diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index 3e8065cfb..4a09a5230 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -91,6 +91,22 @@ class Chunk(BaseModel): return generate_chunk_id(str(uuid.uuid4()), str(self.content)) + @property + def document_id(self) -> str | None: + """Returns the document_id from either metadata or chunk_metadata, with metadata taking precedence.""" + # Check metadata first (takes precedence) + doc_id = self.metadata.get("document_id") + if isinstance(doc_id, str): + return doc_id + + # Fall back to chunk_metadata if available + if self.chunk_metadata is not None: + chunk_doc_id = getattr(self.chunk_metadata, "document_id", None) + if isinstance(chunk_doc_id, str): + return chunk_doc_id + + return None + @json_schema_type class QueryChunksResponse(BaseModel): diff --git a/llama_stack/core/routers/vector_io.py b/llama_stack/core/routers/vector_io.py index ebf99bba9..7d4248107 100644 --- a/llama_stack/core/routers/vector_io.py +++ b/llama_stack/core/routers/vector_io.py @@ -5,6 +5,7 @@ # the root directory of this source tree. import asyncio +import logging import uuid from typing import Any @@ -101,13 +102,12 @@ 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, " - f"ttl_seconds={ttl_seconds}, chunk_ids={doc_ids}{' and more...' if len(chunks) > 3 else ''}" - ) + if logger.isEnabledFor(logging.DEBUG): + doc_ids = [chunk.document_id for chunk in chunks[:3]] + logger.debug( + 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) await provider.insert_chunks(vector_db_id, chunks, ttl_seconds) diff --git a/llama_stack/providers/inline/tool_runtime/rag/memory.py b/llama_stack/providers/inline/tool_runtime/rag/memory.py index 80eb47573..29685f5bf 100644 --- a/llama_stack/providers/inline/tool_runtime/rag/memory.py +++ b/llama_stack/providers/inline/tool_runtime/rag/memory.py @@ -279,10 +279,7 @@ class MemoryToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, RAGToolRunti return RAGQueryResult( content=picked, metadata={ - "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)] - ], + "document_ids": [c.document_id 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)]],