zcbot/core/context_fold.py

"""§8.8 Phase 2:阈值结构化摘要 —— 长会话窗口中段折叠。

Phase 1(context_base_idx 软重置)解决"跨话题累积";单个分段内持续对话,窗口仍只增
不减 —— 50% 门槛的旧 tool 压缩(core/context.py)只摘工具正文,user/assistant 对话
本体一字不少,最终仍会撞 reliable_context。本模块在 run 起点检测窗口体量,达
reliable_context × FOLD_TRIGGER_RATIO 时把窗口中段(保留尾部约 FOLD_TAIL_RATIO)
折叠为一份固定模板的结构化摘要。

设计要点(取舍详 DESIGN §8.8):
- **边界而非删除**:原文全量留 messages 表,只写 tasks.context_summary + 推进
  context_base_idx;web 端照旧翻完整历史。
- **增量更新**:旧摘要作为注入消息本就在被折叠的前缀里,摘要指令要求合并 ——
  新摘要 = f(旧摘要, 新折叠段),不重读全史。
- **前缀缓存友好**:摘要调用复用会话自身的消息前缀(prepare 同款压缩后取 [:切点]),
  末尾追加一条摘要指令 user 消息 —— 前缀与上一轮 chat 的缓存字节一致,几乎全程
  cache hit;折叠后的窗口(system + 摘要 + 尾段)在下次折叠前字节稳定,缓存恢复全暖。
- **精确值原文保留**:模板硬性要求路径/ID/数值不得改写(mem0 实测自由摘要会静默丢
  精确值),守 §8.2 质量边界(确认过的需求/结论不删)。
- **失败零阻塞**:摘要调用失败只 warn + 跳过本次折叠,run 照常跑(85% 距硬上限有垫)。
"""
from __future__ import annotations

from typing import Any, Callable, List, Optional
from uuid import UUID

from sqlalchemy import update

from .context import (
    CHARS_PER_TOKEN,
    COMPACT_CONTEXT_RATIO,
    _message_chars,
    prepare_messages_with_stats,
)
from .storage import session_scope
from .storage.models import Task
from .storage.usage import record_chat_usage

# 窗口体量达 reliable_context 的此比例 → 触发折叠(50% 压缩、85% 折叠双层门槛)。
FOLD_TRIGGER_RATIO = 0.85
# 折叠后保留的尾部体量占比(近期对话原文,续聊手感;其余折进摘要)。
FOLD_TAIL_RATIO = 0.30
# 可折叠段少于这么多条消息就不折 —— 摘要调用本身有成本,段太短折了省不了几个字。
MIN_FOLD_MESSAGES = 8

SUMMARY_NOTE_HEADER = "[前情摘要]"

# 摘要指令:作为最后一条 user 消息追加在会话前缀之后(前缀缓存复用,见模块注释)。
FOLD_INSTRUCTION = (
    "请暂停手头任务,现在只做一件事:把以上全部对话内容压缩成一份结构化摘要,"
    "它将替代这些消息作为后续对话的上下文。若开头已有「[前情摘要]」,把它合并进来。\n"
    "严格按以下模板输出,除模板内容外不要输出任何别的文字:\n"
    "## 目标\n## 已确认的约束与决定\n## 进展\n## 待办\n## 关键路径 / ID / 数值\n"
    "硬性要求:文件路径、URL、ID、数字、代码符号、命令等精确值必须原文保留,"
    "不得改写、四舍五入或省略;用户明确确认过的需求与结论必须保留;"
    "不确定的内容标注[待确认]。"
)


def build_summary_note(summary: str) -> dict:
    """把摘要包成注入消息(user 角色,仅内存不落库;Session.load 与折叠现场共用)。"""
    return {
        "role": "user",
        "content": (
            f"{SUMMARY_NOTE_HEADER} 以下是本对话更早内容的结构化摘要"
            "(原始消息已归档,不在当前上下文中);请把它当作已发生的事实,"
            "直接继续对话,不必复述:\n\n" + summary
        ),
    }


def find_cutoff(messages: List[dict], head_len: int) -> Optional[int]:
    """选折叠切点:返回内存 index(切点本身保留在窗口里),不可折返回 None。

    从尾部向前累计体量,第一条使尾部 ≥ FOLD_TAIL_RATIO 的 **user 消息**即切点 ——
    切在 user 消息上天然不劈开 assistant tool_calls / tool 配对,与 Phase 1 软重置
    的锚点语义一致(窗口永远以 user 消息开头)。user 消息稀疏时尾部占比可能超过
    FOLD_TAIL_RATIO,保守方向,可接受。
    """
    n = len(messages)
    if n <= head_len:
        return None
    total = sum(_message_chars(m) for m in messages[head_len:])
    tail_target = total * FOLD_TAIL_RATIO
    acc = 0
    cutoff: Optional[int] = None
    for i in range(n - 1, head_len - 1, -1):
        acc += _message_chars(messages[i])
        if acc >= tail_target and messages[i].get("role") == "user":
            cutoff = i
            break
    if cutoff is None or cutoff - head_len < MIN_FOLD_MESSAGES:
        return None
    return cutoff


def _user_ordinal(messages: List[dict], idx: int) -> int:
    """messages[idx] 是从头数第几条 user 消息(1-based)。"""
    return sum(1 for m in messages[: idx + 1] if m.get("role") == "user")

def _kth_user_index(messages: List[dict], k: int) -> Optional[int]:
    """从头数第 k 条(1-based)user 消息的 index;不存在返回 None。"""
    seen = 0
    for i, m in enumerate(messages):
        if m.get("role") == "user":
            seen += 1
            if seen == k:
                return i
    return None


def persist_fold(task_id: UUID, new_base: int, summary: str) -> None:
    """折叠结果落库:摘要 + 窗口起点,一个事务。先于内存变更执行 —— DB 写失败时
    内存未动,本次折叠白跑一次摘要调用但状态零污染,下次 run 重试。"""
    with session_scope() as s:
        s.execute(
            update(Task)
            .where(Task.task_id == task_id)
            .values(context_base_idx=new_base, context_summary=summary)
        )


def maybe_fold(
    session: Any,
    llm: Any,
    caps: Any,
    *,
    user_id: UUID,
    emit: Callable[[dict], None],
) -> Optional[dict]:
    """run 起点调用:体量达阈值则折叠,返回统计 dict;未触发/不可折返回 None。

    session 鸭子类型,需要:messages / task_id / context_head_len / context_base /
    apply_fold(cutoff, summary)。异常向上抛,由 loop 包 try 兜(失败零阻塞纪律)。
    """
    msgs = session.messages
    before_chars = sum(_message_chars(m) for m in msgs)
    threshold = int(caps.reliable_context * FOLD_TRIGGER_RATIO * CHARS_PER_TOKEN)
    if before_chars < threshold:
        return None
    head_len = session.context_head_len
    cutoff = find_cutoff(msgs, head_len)
    if cutoff is None:
        return None
    folded_count = cutoff - head_len
    emit({"type": "context_fold", "phase": "start", "folding_messages": folded_count})

    # 复用 loop 同款 prepare(50% 门槛旧 tool 压缩)再取前缀 —— 字节与上一轮 chat 的
    # 缓存前缀一致,摘要调用吃到 cache hit。切点在两个列表间用「第 k 条 user 消息」对齐:
    # prepare 只改写 tool 内容 / 补占位 tool 消息,user 消息原样且相对顺序不变。
    compact_threshold = int(caps.reliable_context * COMPACT_CONTEXT_RATIO * CHARS_PER_TOKEN)
    prepared, _ = prepare_messages_with_stats(msgs, compact_threshold_chars=compact_threshold)
    p_cut = _kth_user_index(prepared, _user_ordinal(msgs, cutoff))
    if p_cut is None:  # 理论不可达,防御
        return None
    fold_input = prepared[:p_cut] + [{"role": "user", "content": FOLD_INSTRUCTION}]

    response = llm.chat(messages=fold_input)
    choices = getattr(response, "choices", None) or []
    summary = ((choices[0].message.content if choices else "") or "").strip()
    if not summary:
        emit({"type": "warn", "msg": "context fold: 摘要调用返回空,本轮跳过折叠"})
        return None

    usage = getattr(response, "usage", None)
    try:
        record_chat_usage(
            task_id=session.task_id,
            user_id=user_id,
            message_id=None,  # 摘要不入 messages,event 挂 task 维度
            model_profile=f"{caps.family}.{caps.variant}",
            prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
            completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
            input_cny_per_mtoken=caps.input_cny_per_mtoken,
            output_cny_per_mtoken=caps.output_cny_per_mtoken,
            response=response,
            # 独立 kind:按 kind 对账"折叠烧多少钱",且折叠次数 = count(该 kind)
            # (任务详情 context_folds 字段的数据源,免加计数列)
            kind="context_fold",
        )
    except Exception:
        pass  # 记账失败不阻塞折叠(与 loop 记账同纪律)

    new_base = session.context_base + folded_count
    persist_fold(session.task_id, new_base, summary)
    session.apply_fold(cutoff, summary)
    after_chars = sum(_message_chars(m) for m in session.messages)
    stats = {
        "type": "context_fold",
        "phase": "done",
        "folded_messages": folded_count,
        "before_chars": before_chars,
        "after_chars": after_chars,
        "new_base": new_base,
    }
    emit(stats)
    return stats