Fixes from PR review

Author: luis5tb

src/app/endpoints/streaming_query.py

@@ -47,6 +47,7 @@
 from metrics.utils import update_llm_token_count_from_turn
 from models.cache_entry import CacheEntry
 from models.config import Action
+from models.context import ResponseGeneratorContext
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
 from models.responses import ForbiddenResponse, UnauthorizedResponse
@@ -696,17 +697,8 @@ def _handle_heartbeat_event(
     )
 
 
-def create_agent_response_generator(  # pylint: disable=too-many-arguments,too-many-locals
-    conversation_id: str,
-    user_id: str,
-    model_id: str,
-    provider_id: str,
-    query_request: QueryRequest,
-    metadata_map: dict[str, dict[str, Any]],
-    client: AsyncLlamaStackClient,
-    llama_stack_model_id: str,
-    started_at: str,
-    _skip_userid_check: bool,
+def create_agent_response_generator(  # pylint: disable=too-many-locals
+    context: ResponseGeneratorContext,
 ) -> Any:
     """
     Create a response generator function for Agent API streaming.
@@ -715,16 +707,7 @@ def create_agent_response_generator(  # pylint: disable=too-many-arguments,too-m
     responses from the Agent API and yields Server-Sent Events (SSE).
 
     Args:
-        conversation_id: The conversation identifier
-        user_id: The user identifier
-        model_id: The model identifier
-        provider_id: The provider identifier
-        query_request: The query request object
-        metadata_map: Dictionary for storing metadata from tool responses
-        client: The Llama Stack client
-        llama_stack_model_id: The full llama stack model ID
-        started_at: Timestamp when the request started
-        _skip_userid_check: Whether to skip user ID validation
+        context: Context object containing all necessary parameters for response generation
 
     Returns:
         An async generator function that yields SSE-formatted strings
@@ -748,10 +731,10 @@ async def response_generator(
         summary = TurnSummary(llm_response="No response from the model", tool_calls=[])
 
         # Determine media type for response formatting
-        media_type = query_request.media_type or MEDIA_TYPE_JSON
+        media_type = context.query_request.media_type or MEDIA_TYPE_JSON
 
         # Send start event at the beginning of the stream
-        yield stream_start_event(conversation_id)
+        yield stream_start_event(context.conversation_id)
 
         latest_turn: Any | None = None
 
@@ -764,10 +747,10 @@ async def response_generator(
                     p.turn.output_message.content
                 )
                 latest_turn = p.turn
-                system_prompt = get_system_prompt(query_request, configuration)
+                system_prompt = get_system_prompt(context.query_request, configuration)
                 try:
                     update_llm_token_count_from_turn(
-                        p.turn, model_id, provider_id, system_prompt
+                        p.turn, context.model_id, context.provider_id, system_prompt
                     )
                 except Exception:  # pylint: disable=broad-except
                     logger.exception("Failed to update token usage metrics")
@@ -776,7 +759,11 @@ async def response_generator(
                     summary.append_tool_calls_from_llama(p.step_details)
 
             for event in stream_build_event(
-                chunk, chunk_id, metadata_map, media_type, conversation_id
+                chunk,
+                chunk_id,
+                context.metadata_map,
+                media_type,
+                context.conversation_id,
             ):
                 chunk_id += 1
                 yield event
@@ -788,49 +775,53 @@ async def response_generator(
             else TokenCounter()
         )
 
-        yield stream_end_event(metadata_map, summary, token_usage, media_type)
+        yield stream_end_event(context.metadata_map, summary, token_usage, media_type)
 
         if not is_transcripts_enabled():
             logger.debug("Transcript collection is disabled in the configuration")
         else:
             store_transcript(
-                user_id=user_id,
-                conversation_id=conversation_id,
-                model_id=model_id,
-                provider_id=provider_id,
+                user_id=context.user_id,
+                conversation_id=context.conversation_id,
+                model_id=context.model_id,
+                provider_id=context.provider_id,
                 query_is_valid=True,  # TODO(lucasagomes): implement as part of query validation
-                query=query_request.query,
-                query_request=query_request,
+                query=context.query_request.query,
+                query_request=context.query_request,
                 summary=summary,
                 rag_chunks=create_rag_chunks_dict(summary),
                 truncated=False,  # TODO(lucasagomes): implement truncation as part
                 # of quota work
-                attachments=query_request.attachments or [],
+                attachments=context.query_request.attachments or [],
             )
 
         # Get the initial topic summary for the conversation
         topic_summary = None
         with get_session() as session:
             existing_conversation = (
-                session.query(UserConversation).filter_by(id=conversation_id).first()
+                session.query(UserConversation)
+                .filter_by(id=context.conversation_id)
+                .first()
             )
             if not existing_conversation:
                 topic_summary = await get_topic_summary(
-                    query_request.query, client, llama_stack_model_id
+                    context.query_request.query,
+                    context.client,
+                    context.llama_stack_model_id,
                 )
 
         completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
 
         referenced_documents = create_referenced_documents_with_metadata(
-            summary, metadata_map
+            summary, context.metadata_map
         )
 
         cache_entry = CacheEntry(
-            query=query_request.query,
+            query=context.query_request.query,
             response=summary.llm_response,
-            provider=provider_id,
-            model=model_id,
-            started_at=started_at,
+            provider=context.provider_id,
+            model=context.model_id,
+            started_at=context.started_at,
             completed_at=completed_at,
             referenced_documents=(
                 referenced_documents if referenced_documents else None
@@ -839,25 +830,25 @@ async def response_generator(
 
         store_conversation_into_cache(
             configuration,
-            user_id,
-            conversation_id,
+            context.user_id,
+            context.conversation_id,
             cache_entry,
-            _skip_userid_check,
+            context.skip_userid_check,
             topic_summary,
         )
 
         persist_user_conversation_details(
-            user_id=user_id,
-            conversation_id=conversation_id,
-            model=model_id,
-            provider_id=provider_id,
+            user_id=context.user_id,
+            conversation_id=context.conversation_id,
+            model=context.model_id,
+            provider_id=context.provider_id,
             topic_summary=topic_summary,
         )
 
     return response_generator
 
 
-async def streaming_query_endpoint_handler_base(  # pylint: disable=too-many-locals,too-many-statements,too-many-arguments
+async def streaming_query_endpoint_handler_base(  # pylint: disable=too-many-locals,too-many-statements,too-many-arguments,too-many-positional-arguments
     request: Request,
     query_request: QueryRequest,
     auth: AuthTuple,
@@ -866,7 +857,7 @@ async def streaming_query_endpoint_handler_base(  # pylint: disable=too-many-loc
     create_response_generator_func: Callable[..., Any],
 ) -> StreamingResponse:
     """
-    Base handler for streaming query endpoints.
+    Handle streaming query endpoints with common logic.
 
     This base handler contains all the common logic for streaming query endpoints
     and accepts functions for API-specific behavior (Agent API vs Responses API).
@@ -937,20 +928,23 @@ async def streaming_query_endpoint_handler_base(  # pylint: disable=too-many-loc
         )
         metadata_map: dict[str, dict[str, Any]] = {}
 
-        # Create the response generator using the provided factory function
-        response_generator = create_response_generator_func(
+        # Create context object for response generator
+        context = ResponseGeneratorContext(
             conversation_id=conversation_id,
             user_id=user_id,
+            skip_userid_check=_skip_userid_check,
             model_id=model_id,
             provider_id=provider_id,
-            query_request=query_request,
-            metadata_map=metadata_map,
-            client=client,
             llama_stack_model_id=llama_stack_model_id,
+            query_request=query_request,
             started_at=started_at,
-            _skip_userid_check=_skip_userid_check,
+            client=client,
+            metadata_map=metadata_map,
         )
 
+        # Create the response generator using the provided factory function
+        response_generator = create_response_generator_func(context)
+
         # Update metrics for the LLM call
         metrics.llm_calls_total.labels(provider_id, model_id).inc()