Skip to content

Commit 2cb67a7

Browse files
committed
Embed index name as source attribute in chunk metadata
1 parent a3862ee commit 2cb67a7

2 files changed

Lines changed: 9 additions & 1 deletion

File tree

src/lightspeed_rag_content/document_processor.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -540,10 +540,11 @@ async def upload_file(chunk_indices: list[int]) -> str:
540540
)
541541
embedding = embedding_response.data[0].embedding
542542

543+
metadata = {**doc.get("metadata", {}), "source": index}
543544
chunk = {
544545
"content": doc["content"],
545546
"chunk_id": doc["chunk_id"],
546-
"metadata": doc.get("metadata", {}),
547+
"metadata": metadata,
547548
"chunk_metadata": doc["chunk_metadata"],
548549
"embedding": embedding,
549550
"embedding_model": embedding_model,
@@ -613,6 +614,7 @@ async def _upload_and_process_files( # noqa: C901 # pylint: disable=R0912,R091
613614
attributes = {
614615
**rag_doc.metadata, # type: ignore[union-attr]
615616
"document_id": doc_uuid,
617+
"source": index,
616618
}
617619
vs_file = await client.vector_stores.files.create(
618620
vector_store_id=vector_store.id,

tests/test_document_processor_llama_stack.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,9 @@ def test_save_manual_chunking(self, mocker, llama_stack_processor):
557557
assert call_kwargs["vector_store_id"] == "vs_123"
558558
assert "chunks" in call_kwargs
559559
assert len(call_kwargs["chunks"]) == 2
560+
# Verify index name is embedded in chunk metadata as "source"
561+
for chunk in call_kwargs["chunks"]:
562+
assert chunk["metadata"]["source"] == mock.sentinel.index
560563

561564
def test_save_auto_chunking(self, mocker, llama_stack_processor):
562565
"""Test saving documents with automatic chunking workflow."""
@@ -566,3 +569,6 @@ def test_save_auto_chunking(self, mocker, llama_stack_processor):
566569
# Verify files.create was called for each document (single file upload)
567570
assert client.files.create.await_count == 2
568571
assert client.vector_stores.files.create.await_count == 2
572+
# Verify index name is embedded in file attributes as "source"
573+
for call in client.vector_stores.files.create.await_args_list:
574+
assert call.kwargs["attributes"]["source"] == mock.sentinel.index

0 commit comments

Comments
 (0)