e4dcfc49aa
Tests / Import Check (Python 3.13) (push) Has been cancelled
Tests / Import Check (Python 3.14) (push) Has been cancelled
Tests / Python Tests (Python 3.11) (push) Has been cancelled
Tests / Python Tests (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.14) (push) Has been cancelled
Tests / Test Summary (push) Has been cancelled
Tests / Lint and Format (push) Has been cancelled
Tests / Web Node Tests (push) Has been cancelled
Tests / Import Check (Python 3.11) (push) Has been cancelled
Tests / Import Check (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.13) (push) Has been cancelled
479 lines
19 KiB
Python
479 lines
19 KiB
Python
"""
|
|
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 deeptutor.services.llm.context_window import resolve_effective_context_window
|
|
|
|
from .protocol import SessionStoreProtocol
|
|
|
|
#: When the summarizer's output lands within this fraction of its hard token
|
|
#: cap, assume the provider cut it mid-sentence and trim the partial tail.
|
|
TRUNCATION_GUARD_RATIO = 0.95
|
|
|
|
|
|
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 trim_incomplete_tail(text: str) -> str:
|
|
"""Drop the trailing partial line from output that hit a hard token cap.
|
|
|
|
A summary cut mid-sentence would otherwise be persisted as-is; losing the
|
|
last line is cheaper than carrying a corrupted entry forward.
|
|
"""
|
|
lines = text.rstrip().split("\n")
|
|
if len(lines) > 1:
|
|
return "\n".join(lines[:-1]).rstrip()
|
|
return text.rstrip()
|
|
|
|
|
|
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: SessionStoreProtocol,
|
|
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 _effective_context_window(self, llm_config: LLMConfig) -> int:
|
|
return resolve_effective_context_window(
|
|
context_window=getattr(llm_config, "context_window", None),
|
|
model=str(getattr(llm_config, "model", "") or ""),
|
|
max_tokens=getattr(llm_config, "max_tokens", None),
|
|
)
|
|
|
|
def _history_budget(self, llm_config: LLMConfig) -> int:
|
|
effective_context_window = self._effective_context_window(llm_config)
|
|
return max(256, int(effective_context_window * 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 _rebuild_source_budget(self, llm_config: LLMConfig) -> int:
|
|
# Raw-rebuild input may use up to half the effective context window;
|
|
# beyond that we degrade to fold-in (existing summary + new turns).
|
|
return max(1024, self._effective_context_window(llm_config) // 2)
|
|
|
|
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"}
|
|
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,
|
|
)
|
|
# The instruction targets ~80% of the hard cap so the model's own
|
|
# length control — not the max_tokens cut — is the binding limit.
|
|
target_tokens = max(96, int(summary_budget * 0.8))
|
|
system_prompt = (
|
|
"You maintain a running summary of a conversation so future turns can "
|
|
"continue seamlessly. Rewrite the summary from the material provided, "
|
|
"organized under these headings (omit any heading with no content):\n"
|
|
"- Goals: what the user wants to accomplish, and why if stated\n"
|
|
"- Key facts & context: stable facts, definitions, data points, names, "
|
|
"references (files, links, IDs)\n"
|
|
"- Decisions & preferences: choices made, options rejected, style or "
|
|
"format preferences, capability/mode switches\n"
|
|
"- Progress: what has been produced or completed so far\n"
|
|
"- Open items: unanswered questions, pending tasks, known blockers\n"
|
|
"Carry forward still-relevant entries from the existing summary unchanged "
|
|
"unless new information contradicts them; drop only what is obsolete. "
|
|
"Prefer concrete details (numbers, identifiers, exact terms) over "
|
|
"abstract restatement. Never invent information."
|
|
)
|
|
if language.startswith("zh"):
|
|
system_prompt = (
|
|
"你负责维护一份对话的滚动摘要,供后续轮次无缝衔接。请基于给定材料重写摘要,"
|
|
"按以下小节组织(无内容的小节直接省略):\n"
|
|
"- 目标:用户想完成什么,以及(如有说明)原因\n"
|
|
"- 关键事实与上下文:稳定的事实、定义、数据、名称、引用(文件、链接、ID)\n"
|
|
"- 决定与偏好:已做的选择、被否决的方案、风格/格式偏好、能力或模式切换\n"
|
|
"- 进展:目前已经产出或完成的内容\n"
|
|
"- 待办事项:未回答的问题、未完成的任务、已知阻塞\n"
|
|
"已有摘要中仍然有效的条目应原样保留,仅在新信息与之矛盾时修改,只删除确已过时"
|
|
"的内容。优先保留具体细节(数字、标识符、确切措辞),不要抽象转述,绝不虚构。"
|
|
)
|
|
user_prompt = (
|
|
f"Update the summary using the material below. "
|
|
f"Keep the total under {target_tokens} tokens.\n\n{source_text}"
|
|
)
|
|
if language.startswith("zh"):
|
|
user_prompt = (
|
|
f"请基于下面的材料更新摘要,总长度不超过 {target_tokens} tokens。\n\n{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).strip()
|
|
if count_tokens(summary) >= int(summary_budget * TRUNCATION_GUARD_RATIO):
|
|
summary = trim_incomplete_tail(summary)
|
|
return summary, 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,
|
|
leaf_message_id: int | None = None,
|
|
) -> ContextBuildResult:
|
|
session = await self.store.get_session(session_id)
|
|
# When ``leaf_message_id`` is given (edit-branch turn), only the
|
|
# ancestor path of that message is included in context — sibling
|
|
# branches at any depth are excluded.
|
|
messages = await self.store.get_messages_for_context(
|
|
session_id, leaf_message_id=leaf_message_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)
|
|
# Branch guard: the watermark must sit on this turn's ancestor chain.
|
|
# After an edit-branch switch it may point into a sibling branch — the
|
|
# stored summary would then carry content this branch never saw.
|
|
# Discard both and rebuild from this branch's own messages.
|
|
if summary_up_to_msg_id > 0 and not any(
|
|
int(item.get("id", 0) or 0) == summary_up_to_msg_id for item in messages
|
|
):
|
|
stored_summary = ""
|
|
summary_up_to_msg_id = 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
|
|
)
|
|
# Everything not retained verbatim: previously summarized messages
|
|
# plus the older unsummarized turns.
|
|
prefix_messages = messages[: len(messages) - len(recent_messages)]
|
|
prefix_transcript = format_messages_as_transcript(prefix_messages)
|
|
|
|
# Anti-drift: while the raw prefix still fits the rebuild budget,
|
|
# re-summarize from the original messages instead of folding the
|
|
# previous summary into itself — summary-of-summary loses detail
|
|
# monotonically. Only beyond that budget degrade to fold-in.
|
|
rebuild_from_raw = bool(prefix_transcript) and count_tokens(
|
|
prefix_transcript
|
|
) <= self._rebuild_source_budget(llm_config)
|
|
merge_parts: list[str] = []
|
|
if rebuild_from_raw:
|
|
merge_parts.append(f"Conversation history to summarize:\n{prefix_transcript}")
|
|
else:
|
|
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))
|
|
|
|
summarize_ok = True
|
|
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:
|
|
summarize_ok = False
|
|
new_summary = ""
|
|
events = []
|
|
|
|
if summarize_ok and new_summary:
|
|
# Advance the watermark only on a successful summarize — never
|
|
# past turns that were not actually folded in.
|
|
up_to_msg_id = summary_up_to_msg_id
|
|
if prefix_messages:
|
|
up_to_msg_id = max(summary_up_to_msg_id, int(prefix_messages[-1].get("id", 0) or 0))
|
|
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)
|
|
else:
|
|
# Degrade for this turn only: keep the stale summary and as many
|
|
# unsummarized turns as fit; nothing is marked as summarized, so
|
|
# the next turn retries with the full material.
|
|
final_history = self._build_history(stored_summary, unsummarized)
|
|
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",
|
|
"TRUNCATION_GUARD_RATIO",
|
|
"build_history_text",
|
|
"count_tokens",
|
|
"format_messages_as_transcript",
|
|
"trim_incomplete_tail",
|
|
]
|