populate rag_chunks

src/app/endpoints/query.py

@@ -3,10 +3,11 @@
 import ast
 import json
 import logging
-import traceback
 import re
+import traceback
 from datetime import UTC, datetime
 from typing import Annotated, Any, Optional, cast
+from urllib.parse import urljoin
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from llama_stack_client import (
@@ -17,15 +18,14 @@
 from llama_stack_client.types import Shield, UserMessage  # type: ignore
 from llama_stack_client.types.agents.turn import Turn
 from llama_stack_client.types.agents.turn_create_params import (
+    Document,
     Toolgroup,
     ToolgroupAgentToolGroupWithArgs,
-    Document,
 )
 from llama_stack_client.types.model_list_response import ModelListResponse
 from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
 from llama_stack_client.types.tool_execution_step import ToolExecutionStep
 
-
 import constants
 import metrics
 from app.database import get_session
@@ -41,33 +41,39 @@
 from models.responses import (
     ForbiddenResponse,
     QueryResponse,
+    RAGChunk,
     ReferencedDocument,
     ToolCall,
     UnauthorizedResponse,
 )
 from utils.endpoints import (
     check_configuration_loaded,
     get_agent,
-    get_topic_summary_system_prompt,
-    get_temp_agent,
     get_system_prompt,
+    get_temp_agent,
+    get_topic_summary_system_prompt,
     store_conversation_into_cache,
     validate_conversation_ownership,
     validate_model_provider_override,
 )
+from utils.mcp_headers import handle_mcp_headers_with_toolgroups, mcp_headers_dependency
 from utils.quota import (
-    get_available_quotas,
     check_tokens_available,
     consume_tokens,
+    get_available_quotas,
 )
-from utils.mcp_headers import handle_mcp_headers_with_toolgroups, mcp_headers_dependency
+from utils.token_counter import TokenCounter, extract_and_update_token_metrics
 from utils.transcripts import store_transcript
 from utils.types import TurnSummary
-from utils.token_counter import extract_and_update_token_metrics, TokenCounter
 
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["query"])
 
+# When OFFLINE is False, use reference_url for chunk source
+# When OFFLINE is True, use parent_id for chunk source
+# TODO: move this setting to a higher level configuration
+OFFLINE = True
+
 query_response: dict[int | str, dict[str, Any]] = {
     200: {
         "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
@@ -296,15 +302,18 @@ async def query_endpoint_handler_base(  # pylint: disable=R0914
                 user_conversation=user_conversation, query_request=query_request
             ),
         )
-        summary, conversation_id, referenced_documents, token_usage = (
-            await retrieve_response_func(
-                client,
-                llama_stack_model_id,
-                query_request,
-                token,
-                mcp_headers=mcp_headers,
-                provider_id=provider_id,
-            )
+        (
+            summary,
+            conversation_id,
+            referenced_documents,
+            token_usage,
+        ) = await retrieve_response_func(
+            client,
+            llama_stack_model_id,
+            query_request,
+            token,
+            mcp_headers=mcp_headers,
+            provider_id=provider_id,
         )
 
         # Get the initial topic summary for the conversation
@@ -755,6 +764,8 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
     vector_db_ids = [vdb.identifier for vdb in vector_dbs]
 
     rag_context = ""
+    retrieved_chunks = []
+    retrieved_scores = []
     try:
         # Use the first available vector database if any exist
         if vector_db_ids:
@@ -763,7 +774,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
             query_response = await client.vector_io.query(
                 vector_db_id=vector_db_id,
                 query=query_request.query,
-                params={"k": 5, "score_threshold": 0.0 }
+                params={"k": 5, "score_threshold": 0.0},
             )
             logger.info(f"The query response total payload:{query_response}")
 
@@ -772,6 +783,12 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
                 for i, chunk in enumerate(query_response.chunks[:3], 1):
                     rag_context += f"{i}. {chunk.content}\n"
 
+                # Store chunks and scores for later inclusion in TurnSummary
+                retrieved_chunks = query_response.chunks
+                retrieved_scores = (
+                    query_response.scores if hasattr(query_response, "scores") else []
+                )
+
                 logger.info(
                     f"Retrieved {len(query_response.chunks)} chunks from vector DB"
                 )
@@ -795,6 +812,30 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
     )
     response = cast(Turn, response)
 
+    # Convert retrieved chunks to RAGChunk format
+    rag_chunks = []
+    for i, chunk in enumerate(retrieved_chunks):
+        # Extract source from chunk metadata based on OFFLINE flag
+        source = None
+        if chunk.metadata:
+            if OFFLINE:
+                parent_id = chunk.metadata.get("parent_id")
+                if parent_id:
+                    source = urljoin("https://mimir.corp.redhat.com", parent_id)
+            else:
+                source = chunk.metadata.get("reference_url")
+
+        # Get score from retrieved_scores list if available
+        score = retrieved_scores[i] if i < len(retrieved_scores) else None
+
+        rag_chunks.append(
+            RAGChunk(
+                content=chunk.content,
+                source=source,
+                score=score,
+            )
+        )
+
     summary = TurnSummary(
         llm_response=(
             interleaved_content_as_str(response.output_message.content)
@@ -805,6 +846,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
             else ""
         ),
         tool_calls=[],
+        rag_chunks=rag_chunks,
     )
 
     referenced_documents = parse_referenced_documents(response)