diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index 786b0e391..ebf99bba9 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -101,11 +101,15 @@ class VectorIORouter(VectorIO): chunks: list[Chunk], ttl_seconds: int | None = None, ) -> None: + doc_ids = [ + getattr(chunk.chunk_metadata, "document_id", None) if chunk.chunk_metadata else None for chunk in chunks[:3] + ] logger.debug( - f"VectorIORouter.insert_chunks: {vector_db_id}, {len(chunks)} chunks, ttl_seconds={ttl_seconds}, chunk_ids={[chunk.metadata['document_id'] for chunk in chunks[:3]]}{' and more...' if len(chunks) > 3 else ''}", + f"VectorIORouter.insert_chunks: {vector_db_id}, {len(chunks)} chunks, " + f"ttl_seconds={ttl_seconds}, chunk_ids={doc_ids}{' and more...' if len(chunks) > 3 else ''}" ) provider = await self.routing_table.get_provider_impl(vector_db_id) - return await provider.insert_chunks(vector_db_id, chunks, ttl_seconds) + await provider.insert_chunks(vector_db_id, chunks, ttl_seconds) async def query_chunks( self, diff --git a/llama_stack/providers/inline/tool_runtime/rag/memory.py b/llama_stack/providers/inline/tool_runtime/rag/memory.py index bc68f198d..80eb47573 100644 --- a/llama_stack/providers/inline/tool_runtime/rag/memory.py +++ b/llama_stack/providers/inline/tool_runtime/rag/memory.py @@ -279,7 +279,10 @@ class MemoryToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, RAGToolRunti return RAGQueryResult( content=picked, metadata={ - "document_ids": [c.metadata["document_id"] for c in chunks[: len(picked)]], + "document_ids": [ + c.metadata.get("document_id") or (c.chunk_metadata.document_id if c.chunk_metadata else None) + for c in chunks[: len(picked)] + ], "chunks": [c.content for c in chunks[: len(picked)]], "scores": scores[: len(picked)], "vector_db_ids": [c.metadata["vector_db_id"] for c in chunks[: len(picked)]], diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index 98889f38e..12f1fb946 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -113,6 +113,25 @@ async def test_insert_chunks_missing_db_raises(vector_io_adapter): await vector_io_adapter.insert_chunks("db_not_exist", []) +async def test_insert_chunks_with_missing_document_id(vector_io_adapter): + """Ensure no KeyError when document_id is missing or in different places.""" + from llama_stack.apis.vector_io import Chunk, ChunkMetadata + + fake_index = AsyncMock() + vector_io_adapter.cache["db1"] = fake_index + + # Various document_id scenarios that shouldn't crash + chunks = [ + Chunk(content="has doc_id in metadata", metadata={"document_id": "doc-1"}), + Chunk(content="no doc_id anywhere", metadata={"source": "test"}), + Chunk(content="doc_id in chunk_metadata", chunk_metadata=ChunkMetadata(document_id="doc-3")), + ] + + # Should work without KeyError + await vector_io_adapter.insert_chunks("db1", chunks) + fake_index.insert_chunks.assert_awaited_once() + + async def test_query_chunks_calls_underlying_index_and_returns(vector_io_adapter): expected = QueryChunksResponse(chunks=[Chunk(content="c1")], scores=[0.1]) fake_index = AsyncMock(query_chunks=AsyncMock(return_value=expected))