Update megumin_agent/chat.py

megumin_agent/chat.py

@@ -2,12 +2,16 @@ from __future__ import annotations
 
 import uuid
 from dataclasses import dataclass
+from typing import Any
+from typing import AsyncIterator
 from typing import Iterable
 
 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
@@ -20,6 +24,9 @@ from .agent import root_agent
 APP_NAME = "megumin_rag_app"
 MAX_TURNS_IN_CONTEXT = 6
 SUMMARY_MAX_CHARS = 800
+SUMMARY_USER_LIMIT = 3
+SUMMARY_ASSISTANT_LIMIT = 2
+SUMMARY_ITEM_CHARS = 42
 
 
 @dataclass
@@ -45,14 +52,78 @@ def _event_texts(events: Iterable) -> list[str]:
     return lines
 
 
-def _compress_summary(previous_summary: str, new_lines: list[str]) -> str:
-    pieces = [previous_summary.strip()] if previous_summary.strip() else []
-    if new_lines:
-        pieces.append(" / ".join(new_lines))
-    summary = " | ".join(piece for piece in pieces if piece).strip()
-    if len(summary) <= SUMMARY_MAX_CHARS:
-        return summary
-    return "..." + summary[-(SUMMARY_MAX_CHARS - 3) :]
+def _compact_summary_item(text: str, limit: int = SUMMARY_ITEM_CHARS) -> str:
+    compact = " ".join(str(text or "").split()).strip()
+    if len(compact) <= limit:
+        return compact
+    return compact[: limit - 3].rstrip() + "..."
+
+
+def _parse_summary_map(value: Any) -> dict[str, list[str]]:
+    if not isinstance(value, dict):
+        return {
+            "user_topics": [],
+            "assistant_points": [],
+        }
+    return {
+        "user_topics": [
+            str(item) for item in value.get("user_topics", []) if str(item).strip()
+        ],
+        "assistant_points": [
+            str(item)
+            for item in value.get("assistant_points", [])
+            if str(item).strip()
+        ],
+    }
+
+
+def _merge_unique_tail(previous: list[str], additions: list[str], limit: int) -> list[str]:
+    merged: list[str] = []
+    for item in [*previous, *additions]:
+        if not item or item in merged:
+            continue
+        merged.append(item)
+    return merged[-limit:]
+
+
+def _compress_summary(
+    previous_summary_map: Any,
+    new_lines: list[str],
+) -> dict[str, list[str]]:
+    summary_map = _parse_summary_map(previous_summary_map)
+    user_lines = [
+        _compact_summary_item(line.removeprefix("user:").strip())
+        for line in new_lines
+        if line.startswith("user:")
+    ]
+    assistant_lines = [
+        _compact_summary_item(line.removeprefix("assistant:").strip())
+        for line in new_lines
+        if line.startswith("assistant:")
+    ]
+    summary_map["user_topics"] = _merge_unique_tail(
+        summary_map["user_topics"],
+        user_lines,
+        SUMMARY_USER_LIMIT,
+    )
+    summary_map["assistant_points"] = _merge_unique_tail(
+        summary_map["assistant_points"],
+        assistant_lines,
+        SUMMARY_ASSISTANT_LIMIT,
+    )
+    return summary_map
+
+
+def _render_summary(summary_map: dict[str, list[str]]) -> str:
+    chunks: list[str] = []
+    if summary_map.get("user_topics"):
+        chunks.append("user_topics: " + " ; ".join(summary_map["user_topics"]))
+    if summary_map.get("assistant_points"):
+        chunks.append("assistant_points: " + " ; ".join(summary_map["assistant_points"]))
+    rendered = "\n".join(chunks).strip()
+    if len(rendered) <= SUMMARY_MAX_CHARS:
+        return rendered
+    return rendered[: SUMMARY_MAX_CHARS - 3].rstrip() + "..."
 
 
 def _trim_session_history(
@@ -62,11 +133,7 @@ def _trim_session_history(
     session_id: str,
 ) -> None:
     session_store = services.session_service.sessions
-    storage_session = (
-        session_store.get(APP_NAME, {})
-        .get(user_id, {})
-        .get(session_id)
-    )
+    storage_session = session_store.get(APP_NAME, {}).get(user_id, {}).get(session_id)
     if storage_session is None:
         return
 
@@ -76,11 +143,10 @@ def _trim_session_history(
 
     overflow = storage_session.events[:-max_events]
     storage_session.events = storage_session.events[-max_events:]
-    previous_summary = str(storage_session.state.get("conversation_summary", ""))
-    storage_session.state["conversation_summary"] = _compress_summary(
-        previous_summary,
-        _event_texts(overflow),
-    )
+    previous_summary_map = storage_session.state.get("conversation_summary_map", {})
+    summary_map = _compress_summary(previous_summary_map, _event_texts(overflow))
+    storage_session.state["conversation_summary_map"] = summary_map
+    storage_session.state["conversation_summary"] = _render_summary(summary_map)
 
 
 def create_chat_services() -> ChatServices:
@@ -95,14 +161,24 @@ def create_chat_services() -> ChatServices:
     return ChatServices(runner=runner, session_service=session_service)
 
 
-async def chat_once(
+def _extract_text(event: Any) -> str:
+    if not getattr(event, "content", None) or not getattr(event.content, "parts", None):
+        return ""
+    texts = [
+        getattr(part, "text", "")
+        for part in event.content.parts
+        if getattr(part, "text", "")
+    ]
+    return "".join(texts).strip()
+
+
+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,
@@ -115,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,
@@ -133,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