Update megumin_agent/chat.py

megumin_agent/chat.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import uuid
 from dataclasses import dataclass
+from typing import Iterable
 
 from .bootstrap import bootstrap_environment
 
@@ -17,6 +18,8 @@ from .retrieval import JsonQaRetriever
 
 
 APP_NAME = "megumin_rag_app"
+MAX_TURNS_IN_CONTEXT = 6
+SUMMARY_MAX_CHARS = 800
 
 
 @dataclass
@@ -25,6 +28,61 @@ class ChatServices:
     session_service: InMemorySessionService
 
 
+def _event_texts(events: Iterable) -> list[str]:
+    lines: list[str] = []
+    for event in events:
+        if not getattr(event, "content", None) or not event.content.parts:
+            continue
+        text_parts = [
+            getattr(part, "text", None)
+            for part in event.content.parts
+            if getattr(part, "text", None)
+        ]
+        if not text_parts:
+            continue
+        author = "user" if event.author == "user" else "assistant"
+        lines.append(f"{author}: {' '.join(text_parts).strip()}")
+    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 _trim_session_history(
+    services: ChatServices,
+    *,
+    user_id: str,
+    session_id: str,
+) -> None:
+    session_store = services.session_service.sessions
+    storage_session = (
+        session_store.get(APP_NAME, {})
+        .get(user_id, {})
+        .get(session_id)
+    )
+    if storage_session is None:
+        return
+
+    max_events = MAX_TURNS_IN_CONTEXT * 2
+    if len(storage_session.events) <= max_events:
+        return
+
+    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),
+    )
+
+
 def create_chat_services() -> ChatServices:
     JsonQaRetriever(DATASET_DIR).warmup()
     session_service = InMemorySessionService()
@@ -68,4 +126,10 @@ async def chat_once(
             if text and event.author != "user":
                 last_text = text
 
+    _trim_session_history(
+        services,
+        user_id=user_id,
+        session_id=active_session_id,
+    )
+
     return last_text, active_session_id