第 21 章 · 记忆与上下文管理

def build_context_full(state) -> list[Message]:
    return state.messages          # 全部历史

def build_context_window(state, keep: int = 5, keep_head: int = 2) -> list[Message]:
    """保留最前面的（system + 初始目标）+ 最近的 N 轮。"""
    msgs = state.messages
    if len(msgs) <= keep + keep_head:
        return msgs
    return msgs[:keep_head] + [Message.system(
        f"（此处省略了 {len(msgs) - keep - keep_head} 轮中间对话）"
    )] + msgs[-keep:]

# clinic/agent_memory.py（节选）
class SummarizingMemory:
    """把旧对话压缩成结构化摘要，而不是直接丢弃。"""

    SUMMARY_PROMPT = """把下面的对话压缩成简明的结构化摘要。
必须保留：① 已经确认的事实与数字（含出处）② 已排除的可能性 ③ 未解决的问题
不要保留：寒暄、重复的推理过程、失败尝试的细节

输出格式：
## 已确认
- ...
## 已排除
- ...
## 待解决
- ...
"""

    def __init__(self, llm, keep_recent: int = 4, max_summary_chars: int = 1500):
        self.llm = llm
        self.keep_recent = keep_recent
        self.max_summary_chars = max_summary_chars
        self.summary: str = ""

    def build(self, state) -> list[Message]:
        msgs = state.messages
        if len(msgs) <= self.keep_recent + 1:
            return msgs
        old, recent = msgs[1:-self.keep_recent], msgs[-self.keep_recent:]
        if old:
            self.summary = self._compress(self.summary, old)
        head = Message.system(f"## 早前对话摘要\n{self.summary}") if self.summary else None
        return [msgs[0]] + ([head] if head else []) + recent

# ❌ 靠对话记忆
# 轮次 3：assistant 说"ADSL 共 254 行"
# 轮次 9：模型要引用这个数字 → 可能记错、记串

# ✅ 靠结构化状态
@dataclass
class AgentState:
    ...
    artifacts: dict[str, Any] = field(default_factory=dict)
    """工作记忆：工具产出的事实，按 key 存放。每轮渲染进上下文。"""

    facts: list[Fact] = field(default_factory=list)

@dataclass
class Fact:
    key: str            # "adsl.n_rows"
    value: Any          # 254
    source: str         # "describe_dataset(dataset='adsl')"
    step: int           # 3

# clinic/agent_memory.py（节选）
@dataclass
class Chunk:
    doc: str            # 文件名
    section: str        # 章节号，如 "6.3.2"
    title: str
    text: str

class KeywordRetriever:
    """基于关键词打分的检索器。零依赖，对结构化文档效果好。

    打分 = 命中词数 / 查询词数 + 标题命中加权
    """

    def __init__(self, chunks: list[Chunk]):
        self.chunks = chunks
        self._index = [self._tokenize(c.text + " " + c.title) for c in chunks]

    @staticmethod
    def _tokenize(s: str) -> set[str]:
        # 中文按 2-gram + 英文按单词。对"缺失值""受试者"这类词足够有效
        s = s.lower()
        words = set(re.findall(r"[a-z_][a-z0-9_]{2,}", s))
        cjk = re.findall(r"[\u4e00-\u9fff]", s)
        words |= {"".join(cjk[i:i + 2]) for i in range(len(cjk) - 1)}
        return words

    def search(self, query: str, top_k: int = 3) -> list[tuple[Chunk, float]]:
        q = self._tokenize(query)
        if not q:
            return []
        scored = []
        for chunk, toks in zip(self.chunks, self._index):
            hit = len(q & toks)
            if not hit:
                continue
            score = hit / len(q)
            if q & self._tokenize(chunk.title):
                score += 0.3                     # 标题命中更相关
            scored.append((chunk, score))
        scored.sort(key=lambda x: -x[1])
        return scored[:top_k]

# ❌ 不可核查
"根据 SAP 的规定，主要疗效分析应使用 ITT 人群。"
# 监管问：SAP 第几节说的？

# ✅ 可核查
"根据 SAP 第 6.3.2 节（引用自 SAP_v2.1.pdf，第 34 页）：
 主要疗效分析基于 ITT 人群。"

@dataclass
class Chunk:
    doc: str
    section: str        # "6.3.2"
    title: str
    page: int | None    # PDF 页码
    text: str

def answer_with_citation(query: str, retriever, llm) -> str:
    hits = retriever.search(query, top_k=3)
    if not hits or hits[0][1] < 0.25:               # ★ 低分阈值
        return ("在提供的文档中没有找到与该问题相关的规定。"
                "请确认文档范围，或提供相关章节。")
    ...

{
    "name": "search_sap",
    "description": (
        "在项目 SAP 文档中检索相关条款，返回原文片段与章节号。"
        "当你需要确认某个分析口径时使用本工具，不要凭记忆回答。"
        "如果返回空结果，说明文档中没有相关规定，请如实说明。"
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "用中文描述你要确认的问题"},
            "top_k": {"type": "integer", "minimum": 1, "maximum": 5, "default": 3},
        },
        "required": ["query"],
        "additionalProperties": False,
    },
}

def rough_tokens(text: str) -> int:
    """粗略估算 token 数。够用即可，不必精确。"""
    cjk = len(re.findall(r"[\u4e00-\u9fff]", text))
    other = len(text) - cjk
    return int(cjk * 0.8 + other * 0.3)

def render_tool_result(data: Any, max_chars: int = 1200) -> str:
    """把工具结果渲染进上下文，超长则摘要化。"""
    text = json.dumps(data, ensure_ascii=False, indent=None)
    if len(text) <= max_chars:
        return text
    # 保留结构骨架 + 前若干行，中间明确省略
    if isinstance(data, dict) and "行" in data:
        rows = data["行"]
        keep = max(3, max_chars // 90)
        trimmed = dict(data, 行=rows[:keep],
                       说明=f"（此处省略 {len(rows) - keep} 行，"
                            f"总数 {data.get('总数', len(rows))}）")
        return json.dumps(trimmed, ensure_ascii=False)
    return text[:max_chars] + f"…（已截断，原长 {len(text)} 字符）"