From 2510bd349eb4b7a2e6516e26d65e1ae4b6828611 Mon Sep 17 00:00:00 2001
From: skamenan7 <skamenan@redhat.com>
Date: Wed, 24 Sep 2025 16:07:47 -0400
Subject: [PATCH] Address PR feedback: optimize logging and encapsulate
 document_id access

- Gate debug logging behind isEnabledFor check to avoid unnecessary computation
- Add Chunk.document_id property to safely handle metadata/chunk_metadata extraction
- Simplify RAG memory code using new property
---
 llama_stack/apis/vector_io/vector_io.py          | 16 ++++++++++++++++
 llama_stack/core/routers/vector_io.py            | 14 +++++++-------
 .../providers/inline/tool_runtime/rag/memory.py  |  5 +----
 3 files changed, 24 insertions(+), 11 deletions(-)

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)]],