Added merging of chunks from inline and tool-based RAG in streaming query

Author: asimurka

src/app/endpoints/streaming_query.py

@@ -243,6 +243,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
         moderation_result=moderation_result,
         vector_store_ids=extract_vector_store_ids_from_tools(responses_params.tools),
         rag_id_mapping=configuration.rag_id_mapping,
+        inline_rag_context=inline_rag_context,
     )
 
     # Update metrics for the LLM call
@@ -254,7 +255,6 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     generator, turn_summary = await retrieve_response_generator(
         responses_params=responses_params,
         context=context,
-        inline_rag_docs=inline_rag_context.referenced_documents,
     )
 
     # Combine inline RAG results (BYOK + Solr) with tool-based results
@@ -283,7 +283,6 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 async def retrieve_response_generator(
     responses_params: ResponsesApiParams,
     context: ResponseGeneratorContext,
-    inline_rag_docs: list[ReferencedDocument],
 ) -> tuple[AsyncIterator[str], TurnSummary]:
     """
     Retrieve the appropriate response generator.
@@ -295,7 +294,6 @@ async def retrieve_response_generator(
     Args:
         responses_params: The Responses API parameters
         context: The response generator context
-        inline_rag_docs: Inline RAG (BYOK + Solr) documents
     Returns:
         tuple[AsyncIterator[str], TurnSummary]: The response generator and turn summary
 
@@ -328,7 +326,6 @@ async def retrieve_response_generator(
                 response,
                 context,
                 turn_summary,
-                inline_rag_docs,
             ),
             turn_summary,
         )
@@ -582,7 +579,6 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
     turn_response: AsyncIterator[OpenAIResponseObjectStream],
     context: ResponseGeneratorContext,
     turn_summary: TurnSummary,
-    inline_rag_docs: list[ReferencedDocument],
 ) -> AsyncIterator[str]:
     """Generate SSE formatted streaming response.
 
@@ -594,7 +590,6 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
         turn_response: The streaming response from Llama Stack
         context: The response generator context
         turn_summary: TurnSummary to populate during streaming
-        inline_rag_docs: Inline RAG (BYOK + Solr) documents
     Yields:
         SSE-formatted strings for tokens, tool calls, tool results,
         turn completion, and error events.
@@ -773,7 +768,11 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
     )
     # Combine inline RAG results (BYOK + Solr) with tool-based results
     turn_summary.referenced_documents = deduplicate_referenced_documents(
-        inline_rag_docs + tool_rag_docs
+        context.inline_rag_context.referenced_documents + tool_rag_docs
+    )
+    # Combine inline RAG chunks (BYOK + Solr) with tool-based chunks
+    turn_summary.rag_chunks = (
+        context.inline_rag_context.rag_chunks + turn_summary.rag_chunks
     )
 
 

src/models/context.py

@@ -4,7 +4,7 @@
 from llama_stack_client import AsyncLlamaStackClient
 
 from models.requests import QueryRequest
-from utils.types import ShieldModerationResult
+from utils.types import RAGContext, ShieldModerationResult
 
 
 @dataclass
@@ -25,6 +25,7 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
         started_at: Timestamp when the request started (ISO 8601 format)
         client: The Llama Stack client for API interactions
         moderation_result: The moderation result
+        inline_rag_context: Inline RAG context
         vector_store_ids: Vector store IDs used in the query for source resolution.
         rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
     """
@@ -47,5 +48,6 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
     moderation_result: ShieldModerationResult
 
     # RAG index identification
+    inline_rag_context: RAGContext
     vector_store_ids: list[str] = field(default_factory=list)
     rag_id_mapping: dict[str, str] = field(default_factory=dict)

tests/unit/app/endpoints/test_query.py

@@ -17,6 +17,9 @@
 from models.responses import QueryResponse
 from utils.token_counter import TokenCounter
 from utils.types import (
+    RAGChunk,
+    RAGContext,
+    ReferencedDocument,
     ResponsesApiParams,
     ShieldModerationPassed,
     ToolCallSummary,
@@ -174,6 +177,93 @@ async def mock_retrieve_response(*_args: Any, **_kwargs: Any) -> TurnSummary:
         assert response.conversation_id == "123"
         assert response.response == "Kubernetes is a container orchestration platform"
 
+    @pytest.mark.asyncio
+    async def test_query_merges_inline_and_tool_rag_chunks_and_documents(
+        self,
+        dummy_request: Request,
+        setup_configuration: AppConfig,
+        mocker: MockerFixture,
+    ) -> None:
+        """Test that inline RAG and tool-based RAG chunks/docs are correctly merged."""
+        query_request = QueryRequest(
+            query="What is Kubernetes?"
+        )  # pyright: ignore[reportCallIssue]
+
+        mocker.patch("app.endpoints.query.configuration", setup_configuration)
+        mocker.patch("app.endpoints.query.check_configuration_loaded")
+        mocker.patch("app.endpoints.query.check_tokens_available")
+        mocker.patch("app.endpoints.query.validate_model_provider_override")
+
+        mock_client = mocker.AsyncMock(spec=AsyncLlamaStackClient)
+        mock_response_obj = mocker.Mock()
+        mock_response_obj.output = []
+        mock_client.responses = mocker.Mock()
+        mock_client.responses.create = mocker.AsyncMock(return_value=mock_response_obj)
+        mock_client_holder = mocker.Mock()
+        mock_client_holder.get_client.return_value = mock_client
+        mocker.patch(
+            "app.endpoints.query.AsyncLlamaStackClientHolder",
+            return_value=mock_client_holder,
+        )
+        mocker.patch(
+            "app.endpoints.query.run_shield_moderation",
+            new=mocker.AsyncMock(return_value=ShieldModerationPassed()),
+        )
+
+        inline_chunk = RAGChunk(content="inline chunk content", source="byok")
+        inline_doc = ReferencedDocument(doc_title="Inline Doc")
+        inline_rag = RAGContext(
+            context_text="",
+            rag_chunks=[inline_chunk],
+            referenced_documents=[inline_doc],
+        )
+        mocker.patch(
+            "app.endpoints.query.build_rag_context",
+            new=mocker.AsyncMock(return_value=inline_rag),
+        )
+
+        mock_responses_params = mocker.Mock(spec=ResponsesApiParams)
+        mock_responses_params.model = "provider1/model1"
+        mock_responses_params.conversation = "conv_123"
+        mock_responses_params.tools = None
+        mock_responses_params.model_dump.return_value = {
+            "input": "test",
+            "model": "provider1/model1",
+        }
+        mocker.patch(
+            "app.endpoints.query.prepare_responses_params",
+            new=mocker.AsyncMock(return_value=mock_responses_params),
+        )
+
+        tool_chunk = RAGChunk(content="tool chunk content", source="vs-1")
+        tool_doc = ReferencedDocument(doc_title="Tool Doc")
+        mock_turn_summary = TurnSummary()
+        mock_turn_summary.rag_chunks = [tool_chunk]
+        mock_turn_summary.referenced_documents = [tool_doc]
+
+        mocker.patch(
+            "app.endpoints.query.retrieve_response",
+            new=mocker.AsyncMock(return_value=mock_turn_summary),
+        )
+        mocker.patch("app.endpoints.query.store_query_results")
+        mocker.patch("app.endpoints.query.consume_query_tokens")
+        mocker.patch("app.endpoints.query.get_available_quotas", return_value={})
+
+        response = await query_endpoint_handler(
+            request=dummy_request,
+            query_request=query_request,
+            auth=MOCK_AUTH,
+            mcp_headers={},
+        )
+
+        assert isinstance(response, QueryResponse)
+        assert len(response.rag_chunks) == 2
+        assert response.rag_chunks[0].content == "inline chunk content"
+        assert response.rag_chunks[1].content == "tool chunk content"
+        assert len(response.referenced_documents) == 2
+        assert response.referenced_documents[0].doc_title == "Inline Doc"
+        assert response.referenced_documents[1].doc_title == "Tool Doc"
+
     @pytest.mark.asyncio
     async def test_successful_query_with_conversation(
         self,