feat：记忆&上下文管理

2026-10-14 08:19:52 -07:00
parent 41b769adde
commit 58604e0f1c
15 changed files with 357 additions and 49 deletions
@@ -1,21 +1,89 @@
-"""会话历史：数据库保存完整对话，每次请求按 token 预算只带最近的完整轮次
+"""会话历史：数据库保存完整对话，每次请求按 token 预算带最近的轮次，最近一轮总是带上

 裁剪在服务层组装请求时做一次，意图判断、问答和 screen agent 看到的是同一段历史
+每轮是 {user, assistant, outcome}，有轮次记录时另带 changes、queries、clarifications（字段同 Turn）：
+用户原话和回复转成对话消息，记录由 history_records 放进本次请求的 JSON，两者共用一个预算
 """

-from langchain_core.messages import (
-    AIMessage,
-    BaseMessage,
-    HumanMessage,
-    trim_messages,
-)
+import json
+from functools import cache
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
 from langchain_core.messages.utils import count_tokens_approximately
+from loguru import logger

 from app.schemas.base import JsonObject
+from app.schemas.screen import Screen
+
+RECORD_KEYS = ("changes", "queries", "clarifications")
+# 放进请求的单条 SQL 上限（字符）；更长的不带，原文仍在会话记录里
+RECORD_SQL_MAX_CHARS = 2000
+# 每条消息在请求里的固定开销（角色、分隔），计入历史预算
+MESSAGE_OVERHEAD = 4
+# 最近一轮超过预算被截短时，回复和用户原话各自至少保留的 token
+MIN_KEPT_TOKENS = 200
+
+
+@cache
+def _encoding():
+    """o200k 近似各家模型的 token；编码表第一次使用时由 tiktoken 下载并缓存，拿不到时退回按字符估算"""
+    try:
+        import tiktoken
+
+        return tiktoken.get_encoding("o200k_base")
+    except Exception:  # noqa: BLE001 计数只用来裁剪历史，不能让请求失败
+        logger.exception("tiktoken 编码表加载失败，历史预算改按字符估算，中文会低估")
+        return None
+
+
+def count_tokens(text: str) -> int:
+    encoding = _encoding()
+    if encoding is None:
+        return count_tokens_approximately([HumanMessage(content=text)])
+    return len(encoding.encode(text, disallowed_special=()))
+
+
+def _record(turn: JsonObject, ago: int, present: set[str] | None) -> JsonObject | None:
+    """一轮放进请求的记录，没有的字段不出现；present 是当前大屏的组件 id，None 时按都不在算（只用于计数）
+
+    过长的 SQL 截断了就不能执行，这条查询整条不带，原文仍在会话记录里
+    """
+    record: JsonObject = {key: turn[key] for key in RECORD_KEYS if turn.get(key)}
+    if "queries" in record:
+        record["queries"] = [
+            query
+            for query in record["queries"]
+            if len(query["sql"]) <= RECORD_SQL_MAX_CHARS
+        ] or None
+    if "changes" in record:
+        record["changes"] = [
+            {**change, "on_screen": _on_screen(change, present or set())}
+            for change in record["changes"]
+        ]
+    record = {key: value for key, value in record.items() if value}
+    return {"ago": ago, **record} if record else None
+
+
+def _on_screen(change: JsonObject, present: set[str]) -> bool:
+    if change["type"] == "canvas":
+        return True
+    return change["node_id"] in present
+
+
+def turn_tokens(turn: JsonObject) -> int:
+    """一轮在请求里的大小：两条消息（含每条的固定开销）和放进请求的记录"""
+    tokens = (
+        count_tokens(turn["user"])
+        + count_tokens(turn["assistant"])
+        + MESSAGE_OVERHEAD * 2
+    )
+    if record := _record(turn, ago=1, present=None):
+        tokens += count_tokens(json.dumps(record, ensure_ascii=False))
+    return tokens


 def history_messages(history: list[JsonObject]) -> list[BaseMessage]:
-    """每轮转成一条用户消息和一条助手消息，按时间从早到晚排列"""
+    """每轮转成一条用户消息和一条助手消息，按时间从早到晚排列；记录不进消息"""
    messages: list[BaseMessage] = []
    for turn in history:
        messages.append(HumanMessage(content=turn["user"]))
@@ -24,16 +92,72 @@ def history_messages(history: list[JsonObject]) -> list[BaseMessage]:


 def trim_history(history: list[JsonObject], max_tokens: int) -> list[JsonObject]:
-    """从最近一轮往前保留，总量不超过 max_tokens；只保留完整的轮次，不从一段回复中间截断
+    """从最近一轮往前保留，消息和记录合计不超过 max_tokens；更早的轮次放不下就整轮不带

-    单独一轮就超过预算时这一轮也不保留
+    最近一轮总是保留：它单独就超过预算时截短再带上（见 _fit），不能因为一条长回答把整段历史丢掉
    """
-    kept = trim_messages(
-        history_messages(history),
-        max_tokens=max_tokens,
-        token_counter=count_tokens_approximately,
-        strategy="last",
-        # 从用户消息开始，助手回复放不下时连同它的用户消息一起去掉
-        start_on="human",
-    )
-    return history[len(history) - len(kept) // 2 :]
+    if not history:
+        return []
+    latest = _fit(history[-1], max_tokens)
+    used = turn_tokens(latest)
+    kept = [latest]
+    for turn in reversed(history[:-1]):
+        used += turn_tokens(turn)
+        if used > max_tokens:
+            break
+        kept.append(turn)
+    return kept[::-1]
+
+
+def _fit(turn: JsonObject, max_tokens: int) -> JsonObject:
+    """一轮超过预算时先截回复、再截用户原话，都是去掉中间、保留开头和结尾（提议和结论通常在回复末尾）
+
+    记录不截：查询最多 10 条、每条 SQL 不超过 RECORD_SQL_MAX_CHARS，默认预算下放得下
+    """
+    for key in ("assistant", "user"):
+        over = turn_tokens(turn) - max_tokens
+        if over <= 0:
+            break
+        keep = max(count_tokens(turn[key]) - over, MIN_KEPT_TOKENS)
+        turn = {**turn, key: truncate_middle(turn[key], keep)}
+    return turn
+
+
+def truncate_middle(text: str, max_tokens: int) -> str:
+    """去掉中间、保留开头和结尾各一半，中间写明省略了多少字；结果不超过 max_tokens"""
+    total = count_tokens(text)
+    if total <= max_tokens:
+        return text
+    keep = len(text) * max_tokens // total
+    while keep > 0:
+        head, tail = text[: keep // 2], text[len(text) - (keep - keep // 2) :]
+        shortened = f"{head}\n…（中间省略 {len(text) - keep} 字）…\n{tail}"
+        if count_tokens(shortened) <= max_tokens:
+            return shortened
+        keep = keep * 9 // 10
+    return f"…（省略 {len(text)} 字）…"
+
+
+def intent_history(history: list[JsonObject]) -> list[JsonObject]:
+    """意图判断看的历史：原话、回复、结果和澄清问答；SQL 和修改记录不给"""
+    projected = []
+    for turn in history:
+        item = {key: turn[key] for key in ("user", "assistant", "outcome")}
+        if turn.get("clarifications"):
+            item["clarifications"] = turn["clarifications"]
+        projected.append(item)
+    return projected
+
+
+def history_records(history: list[JsonObject], screen: Screen) -> list[JsonObject]:
+    """保留下来的各轮的记录，最近的在前；ago 是往前数第几轮，1 是上一轮
+
+    on_screen 按当前大屏现算：只表示这个 id 现在还在大屏上，不代表那次修改已经生效
+    澄清恢复时可能换了大屏，所以在节点里调用，不在组装请求时算好
+    """
+    present = {node.id for node in screen.nodes}
+    records = []
+    for ago, turn in enumerate(reversed(history), start=1):
+        if record := _record(turn, ago, present):
+            records.append(record)
+    return records