""" Build bounded conversation history for unified chat sessions. """ from __future__ import annotations from dataclasses import dataclass from typing import Any, Awaitable, Callable from deeptutor.agents.base_agent import BaseAgent from deeptutor.core.stream import StreamEvent, StreamEventType from deeptutor.core.trace import build_trace_metadata, merge_trace_metadata, new_call_id from deeptutor.services.llm.config import LLMConfig from .sqlite_store import SQLiteSessionStore def count_tokens(text: str) -> int: """Estimate token count with tiktoken when available.""" if not text: return 0 try: import tiktoken encoding = tiktoken.get_encoding("cl100k_base") return len(encoding.encode(text)) except Exception: return max(1, len(text) // 4) def format_messages_as_transcript(messages: list[dict[str, Any]]) -> str: lines: list[str] = [] role_map = { "user": "User", "assistant": "Assistant", "system": "System", } for item in messages: content = str(item.get("content", "") or "").strip() if not content: continue role = role_map.get(str(item.get("role", "user")), "User") lines.append(f"{role}: {content}") return "\n\n".join(lines) def build_history_text(history: list[dict[str, Any]]) -> str: lines: list[str] = [] for item in history: role = str(item.get("role", "user")) content = str(item.get("content", "") or "").strip() if not content: continue if role == "system": lines.append(f"Conversation summary:\n{content}") elif role == "assistant": lines.append(f"Assistant: {content}") else: lines.append(f"User: {content}") return "\n\n".join(lines) @dataclass class ContextBuildResult: conversation_history: list[dict[str, Any]] conversation_summary: str context_text: str events: list[StreamEvent] token_count: int budget: int class _ContextSummaryAgent(BaseAgent): """Small helper agent for compressing older conversation turns.""" def __init__(self, language: str = "en") -> None: super().__init__( module_name="chat", agent_name="context_summary_agent", language=language, ) async def process(self, *_args, **_kwargs) -> dict[str, Any]: raise NotImplementedError class ContextBuilder: """Construct a bounded conversation history plus optional summary trace.""" def __init__( self, store: SQLiteSessionStore, history_budget_ratio: float = 0.35, summary_target_ratio: float = 0.40, ) -> None: self.store = store self.history_budget_ratio = history_budget_ratio self.summary_target_ratio = summary_target_ratio def _history_budget(self, llm_config: LLMConfig) -> int: configured = int(getattr(llm_config, "max_tokens", 4096) or 4096) return max(256, int(configured * self.history_budget_ratio)) def _summary_budget(self, budget: int) -> int: return max(96, int(budget * self.summary_target_ratio)) def _recent_budget(self, budget: int) -> int: return max(128, budget - self._summary_budget(budget)) def _build_history(self, summary: str, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: history: list[dict[str, Any]] = [] cleaned_summary = summary.strip() if cleaned_summary: history.append({"role": "system", "content": cleaned_summary}) history.extend( { "role": item.get("role", "user"), "content": str(item.get("content", "") or ""), } for item in messages if item.get("role") in {"user", "assistant", "system"} and str(item.get("content", "") or "").strip() ) return history async def _append_event( self, events: list[StreamEvent], event: StreamEvent, on_event: Callable[[StreamEvent], Awaitable[None]] | None = None, ) -> None: events.append(event) if on_event is not None: await on_event(event) def _select_recent_messages( self, messages: list[dict[str, Any]], recent_budget: int, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: selected: list[dict[str, Any]] = [] total = 0 for item in reversed(messages): content = str(item.get("content", "") or "") tokens = count_tokens(content) if selected and total + tokens > recent_budget: break selected.insert(0, item) total += tokens cutoff = len(messages) - len(selected) return messages[:cutoff], selected async def _summarize( self, *, session_id: str, language: str, source_text: str, summary_budget: int, on_event: Callable[[StreamEvent], Awaitable[None]] | None = None, ) -> tuple[str, list[StreamEvent]]: events: list[StreamEvent] = [] if not source_text.strip(): return "", events agent = _ContextSummaryAgent(language=language) trace_meta = build_trace_metadata( call_id=new_call_id("context-summary"), phase="summarize_context", label="Summarize context", call_kind="llm_summarization", trace_id=session_id, ) async def _trace_bridge(update: dict[str, Any]) -> None: if str(update.get("event", "")) != "llm_call": return state = str(update.get("state", "running")) metadata = { key: value for key, value in update.items() if key not in {"event", "state", "response", "chunk"} } if state == "running": await self._append_event( events, StreamEvent( type=StreamEventType.PROGRESS, source="context_builder", stage="summarize_context", content="Compressing conversation history...", metadata=merge_trace_metadata( metadata, {"trace_kind": "call_status", "call_state": "running"}, ), ), on_event, ) elif state == "complete": response = str(update.get("response", "") or "") if response: await self._append_event( events, StreamEvent( type=StreamEventType.CONTENT, source="context_builder", stage="summarize_context", content=response, metadata=merge_trace_metadata( metadata, {"trace_kind": "llm_output"}, ), ), on_event, ) await self._append_event( events, StreamEvent( type=StreamEventType.PROGRESS, source="context_builder", stage="summarize_context", content="", metadata=merge_trace_metadata( metadata, {"trace_kind": "call_status", "call_state": "complete"}, ), ), on_event, ) elif state == "error": await self._append_event( events, StreamEvent( type=StreamEventType.ERROR, source="context_builder", stage="summarize_context", content=str(update.get("response", "") or "Context summarization failed."), metadata=merge_trace_metadata(metadata, {"call_state": "error"}), ), on_event, ) agent.set_trace_callback(_trace_bridge) await self._append_event( events, StreamEvent( type=StreamEventType.STAGE_START, source="context_builder", stage="summarize_context", metadata=trace_meta, ), on_event, ) system_prompt = ( "You compress conversation history for future turns. Preserve factual context, " "user goals, constraints, decisions, unresolved questions, and any capability " "switches. Keep the summary concise and faithful. Use bullet points only if useful." ) if language.startswith("zh"): system_prompt = ( "你负责把对话历史压缩成后续轮次可直接使用的上下文。保留用户目标、约束、已做决定、" "未解决问题，以及能力切换带来的关键信息。总结要忠实、紧凑，不要虚构。" ) user_prompt = ( f"Compress the following conversation history into <= {summary_budget} tokens.\n\n" f"{source_text}" ) if language.startswith("zh"): user_prompt = ( f"请把下面的对话历史压缩到不超过 {summary_budget} tokens 的长度，" "供后续对话直接继承上下文。\n\n" f"{source_text}" ) try: _chunks: list[str] = [] async for _c in agent.stream_llm( user_prompt=user_prompt, system_prompt=system_prompt, max_tokens=summary_budget, stage="summarize_context", trace_meta=trace_meta, ): _chunks.append(_c) summary = "".join(_chunks) return summary.strip(), events finally: await self._append_event( events, StreamEvent( type=StreamEventType.STAGE_END, source="context_builder", stage="summarize_context", metadata=trace_meta, ), on_event, ) async def build( self, *, session_id: str, llm_config: LLMConfig, language: str = "en", on_event: Callable[[StreamEvent], Awaitable[None]] | None = None, ) -> ContextBuildResult: session = await self.store.get_session(session_id) messages = await self.store.get_messages_for_context(session_id) if session is None: return ContextBuildResult([], "", "", [], 0, self._history_budget(llm_config)) budget = self._history_budget(llm_config) summary_budget = self._summary_budget(budget) recent_budget = self._recent_budget(budget) stored_summary = str(session.get("compressed_summary", "") or "").strip() summary_up_to_msg_id = int(session.get("summary_up_to_msg_id", 0) or 0) unsummarized = [item for item in messages if int(item.get("id", 0) or 0) > summary_up_to_msg_id] current_history = self._build_history(stored_summary, unsummarized) current_tokens = count_tokens(build_history_text(current_history)) if current_tokens <= budget: return ContextBuildResult( conversation_history=current_history, conversation_summary=stored_summary, context_text=build_history_text(current_history), events=[], token_count=current_tokens, budget=budget, ) older_unsummarized, recent_messages = self._select_recent_messages(unsummarized, recent_budget) merge_parts: list[str] = [] if stored_summary: merge_parts.append(f"Existing summary:\n{stored_summary}") older_transcript = format_messages_as_transcript(older_unsummarized) if older_transcript: merge_parts.append(f"Older turns to fold in:\n{older_transcript}") if not merge_parts and recent_messages: merge_parts.append(format_messages_as_transcript(recent_messages)) try: new_summary, events = await self._summarize( session_id=session_id, language=language, source_text="\n\n".join(part for part in merge_parts if part.strip()), summary_budget=summary_budget, on_event=on_event, ) except Exception: new_summary = stored_summary events = [] up_to_msg_id = summary_up_to_msg_id if older_unsummarized: up_to_msg_id = int(older_unsummarized[-1]["id"]) elif stored_summary: up_to_msg_id = summary_up_to_msg_id if new_summary: await self.store.update_summary(session_id, new_summary, up_to_msg_id) stored_summary = new_summary final_history = self._build_history(stored_summary, recent_messages) while len(final_history) > 1 and count_tokens(build_history_text(final_history)) > budget: summary_prefix = 1 if final_history and final_history[0].get("role") == "system" else 0 if len(final_history) <= summary_prefix + 1: break final_history.pop(summary_prefix) final_text = build_history_text(final_history) return ContextBuildResult( conversation_history=final_history, conversation_summary=stored_summary, context_text=final_text, events=events, token_count=count_tokens(final_text), budget=budget, ) __all__ = [ "ContextBuildResult", "ContextBuilder", "build_history_text", "count_tokens", "format_messages_as_transcript", ]