Added agent streaming query utils

Author: asimurka

src/app/endpoints/streaming_query.py

@@ -74,7 +74,10 @@
 from models.common.responses.types import ResponseInput
 from models.common.turn_summary import TurnSummary
 from models.config import Action
-from utils.agents.streaming import generate_agent_response, retrieve_agent_response_generator
+from utils.agents.streaming import (
+    generate_agent_response,
+    retrieve_agent_response_generator,
+)
 from utils.conversation_compaction import (
     CompactionResult,
     CompactionStartedEvent,
@@ -348,6 +351,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
             context=context,
             responses_params=responses_params,
             turn_summary=turn_summary,
+            background_topic_summary_tasks=_background_topic_summary_tasks,
         ),
         media_type=response_media_type,
     )

src/pydantic_ai_lightspeed/llamastack/_model.py

@@ -1,3 +1,5 @@
+# pylint: disable=too-many-branches, too-many-locals, too-many-statements
+
 """Custom OpenAI Responses model that works around Llama Stack streaming quirks.
 
 Llama Stack's Responses API emits MCP tool argument events *before*
@@ -190,15 +192,15 @@ async def _filtered_iter(
                 yield buffer.pending_output_done
             self._pre_args_buffers.pop(item_id, None)
 
-        for item_id, buffer in list(self._list_tools_buffers.items()):
-            if not buffer.has_content():
+        for item_id, list_tools_buffer in list(self._list_tools_buffers.items()):
+            if not list_tools_buffer.has_content():
                 continue
             logger.warning(
                 "Flushing buffered mcp_list_tools event(s) without output_item.added "
                 "for item_id=%s",
                 item_id,
             )
-            for replayed in self._replay_list_tools_events(buffer):
+            for replayed in self._replay_list_tools_events(list_tools_buffer):
                 yield replayed
             self._list_tools_buffers.pop(item_id, None)
 

src/utils/agents/streaming.py

@@ -292,9 +292,7 @@ async def agent_response_generator(
     logger.debug("Starting agent streaming response processing")
     async with agent.run_stream_events(prompt) as stream:
         async for event in stream:
-            logger.error(f"event: {event.event_kind} {repr(event)}")
             if payload := dispatch_stream_event(event, dispatch_state):
-                # print(f"payload: {payload.serialize_json()}")
                 yield serialize_event(payload, media_type)
 
     if dispatch_state.run_result is None:
@@ -428,7 +426,6 @@ def _(
 
     if isinstance(part, NativeToolReturnPart):
         if tool_result := process_native_tool_result(state, part):
-            # print(f"Tool result summarized: {tool_result.model_dump_json()}")
             return ToolResultStreamPayload(data=tool_result)
         return None