Update megumin_agent/chat.py

megumin_agent/chat.py

@@ -10,6 +10,8 @@ from .bootstrap import bootstrap_environment
 
 bootstrap_environment()
 
+from google.adk.agents.run_config import RunConfig
+from google.adk.agents.run_config import StreamingMode
 from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.genai import types
@@ -170,14 +172,13 @@ def _extract_text(event: Any) -> str:
     return "".join(texts).strip()
 
 
-async def chat_once(
+async def stream_chat(
     user_message: str,
     services: ChatServices,
     session_id: str | None = None,
     user_id: str = "local-user",
-) -> tuple[str, str]:
+) -> AsyncIterator[tuple[str, str]]:
     active_session_id = session_id or str(uuid.uuid4())
-    last_text = ""
     existing_session = await services.session_service.get_session(
         app_name=APP_NAME,
         user_id=user_id,
@@ -190,17 +191,34 @@ async def chat_once(
             session_id=active_session_id,
         )
 
+    streamed_text = ""
+    final_text = ""
+    run_config = RunConfig(streaming_mode=StreamingMode.SSE)
+
     async for event in services.runner.run_async(
         user_id=user_id,
         session_id=active_session_id,
         new_message=types.UserContent(parts=[types.Part(text=user_message)]),
+        run_config=run_config,
     ):
-        if not event.content or not event.content.parts:
+        if getattr(event, "author", None) == "user":
+            continue
+
+        text = _extract_text(event)
+        if not text:
             continue
-        for part in event.content.parts:
-            text = getattr(part, "text", None)
-            if text and event.author != "user":
-                last_text = text
+
+        if getattr(event, "partial", None) is True:
+            streamed_text += text
+            yield streamed_text, active_session_id
+            continue
+
+        if getattr(event, "is_final_response", None) and event.is_final_response():
+            final_text = text
+
+    if final_text and final_text != streamed_text:
+        streamed_text = final_text
+        yield streamed_text, active_session_id
 
     _trim_session_history(
         services,
@@ -208,4 +226,20 @@ async def chat_once(
         session_id=active_session_id,
     )
 
+
+async def chat_once(
+    user_message: str,
+    services: ChatServices,
+    session_id: str | None = None,
+    user_id: str = "local-user",
+) -> tuple[str, str]:
+    last_text = ""
+    active_session_id = session_id or str(uuid.uuid4())
+    async for chunk_text, active_session_id in stream_chat(
+        user_message=user_message,
+        services=services,
+        session_id=active_session_id,
+        user_id=user_id,
+    ):
+        last_text = chunk_text
     return last_text, active_session_id