642 lines
27 KiB
Python
642 lines
27 KiB
Python
"""
|
||
core/agent/compaction.py
|
||
========================
|
||
🌟 pi 上下文压缩算法的 Python 1:1 移植(harness 原版算法)
|
||
|
||
对照 pi-main 源码(逐函数对应):
|
||
packages/agent/src/harness/compaction/compaction.ts
|
||
DEFAULT_COMPACTION_SETTINGS -> CompactionSettings / DEFAULT_COMPACTION_SETTINGS
|
||
shouldCompact -> 见 context.py: should_compact(同一公式)
|
||
estimateTokens / estimateContextTokens -> context.py(usage 锚定,CJK 感知见说明)
|
||
findValidCutPoints / findCutPoint -> find_valid_cut_points / find_cut_point
|
||
findTurnStartIndex -> find_turn_start
|
||
prepareCompaction -> prepare_compaction
|
||
SUMMARIZATION_SYSTEM_PROMPT -> 同名(逐字移植)
|
||
SUMMARIZATION_PROMPT -> 同名(逐字移植)
|
||
UPDATE_SUMMARIZATION_PROMPT -> 同名(逐字移植,迭代更新用)
|
||
TURN_PREFIX_SUMMARIZATION_PROMPT -> 同名(逐字移植,断轮前缀用)
|
||
generateSummaryWithUsage -> generate_summary(maxTokens = 0.8×reserve)
|
||
generateTurnPrefixSummary -> generate_turn_prefix_summary(0.5×reserve)
|
||
compact -> compact_context(断轮双摘要 + 拼接格式 1:1)
|
||
packages/agent/src/harness/compaction/utils.ts
|
||
serializeConversation -> serialize_conversation(1:1,含 2000 字符截断)
|
||
extractFileOpsFromMessage -> extract_file_ops_from_message
|
||
computeFileLists -> compute_file_lists
|
||
formatFileOperations -> format_file_operations
|
||
TOOL_RESULT_MAX_CHARS = 2000 -> 同名常量
|
||
|
||
摘要 LLM 调用由上层注入:
|
||
summarize_fn(prompt_text: str, system_prompt: str, max_tokens: int) -> str
|
||
(pi 里是 models.completeSimple + retry;haocode 用 OpenAI 客户端非流式调用,
|
||
由 llm_engine.AgentWorker 实现并注入。)
|
||
|
||
已声明的偏差(仅 2 处,见 context.py 头注):
|
||
1. 单条消息 token 估算用 CJK 感知启发式(pi 是 chars/4)——对中文会话更安全
|
||
2. 摘要 LLM 调用失败时降级为机械摘录(pi 返回 CompactionError)——桌面应用优先不丢上下文
|
||
其余全部 1:1:触发公式、usage 锚定、token 预算切点、有效切点规则、断轮双摘要、
|
||
迭代式 previousSummary 更新、摘要提示词逐字、文件操作附录、拼接格式。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable, List, Optional, Set, Tuple
|
||
|
||
from .types import AgentMessage, ModelConfig
|
||
|
||
|
||
# ======================================================================
|
||
# 压缩设置 —— 1:1 对照 DEFAULT_COMPACTION_SETTINGS
|
||
# ======================================================================
|
||
@dataclass
|
||
class CompactionSettings:
|
||
"""pi: interface CompactionSettings { enabled; reserveTokens; keepRecentTokens }"""
|
||
enabled: bool = True
|
||
reserve_tokens: int = 16384 # pi: 16384(摘要提示词与输出预留)
|
||
keep_recent_tokens: int = 20000 # pi: 20000(压缩后保留的近期上下文预算)
|
||
|
||
|
||
DEFAULT_COMPACTION_SETTINGS = CompactionSettings()
|
||
|
||
# pi utils.ts: const TOOL_RESULT_MAX_CHARS = 2000
|
||
TOOL_RESULT_MAX_CHARS = 2000
|
||
|
||
|
||
# ======================================================================
|
||
# 摘要提示词 —— 从 pi compaction.ts 逐字移植(不得改写,摘要质量依赖它)
|
||
# ======================================================================
|
||
SUMMARIZATION_SYSTEM_PROMPT = \
|
||
"You are a context summarization assistant. Your task is to read a conversation " \
|
||
"between a user and an AI assistant, then produce a structured summary following " \
|
||
"the exact format specified.\n\n" \
|
||
"Do NOT continue the conversation. Do NOT respond to any questions in the " \
|
||
"conversation. ONLY output the structured summary."
|
||
|
||
SUMMARIZATION_PROMPT = """The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||
|
||
Use this EXACT format:
|
||
|
||
## Goal
|
||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||
|
||
## Constraints & Preferences
|
||
- [Any constraints, preferences, or requirements mentioned by user]
|
||
- [Or "(none)" if none were mentioned]
|
||
|
||
## Progress
|
||
### Done
|
||
- [x] [Completed tasks/changes]
|
||
|
||
### In Progress
|
||
- [ ] [Current work]
|
||
|
||
### Blocked
|
||
- [Issues preventing progress, if any]
|
||
|
||
## Key Decisions
|
||
- **[Decision]**: [Brief rationale]
|
||
|
||
## Next Steps
|
||
1. [Ordered list of what should happen next]
|
||
|
||
## Critical Context
|
||
- [Any data, examples, or references needed to continue]
|
||
- [Or "(none)" if not applicable]
|
||
|
||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||
|
||
UPDATE_SUMMARIZATION_PROMPT = """The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||
|
||
Update the existing structured summary with new information. RULES:
|
||
- PRESERVE all existing information from the previous summary
|
||
- ADD new progress, decisions, and context from the new messages
|
||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||
- UPDATE "Next Steps" based on what was accomplished
|
||
- PRESERVE exact file paths, function names, and error messages
|
||
- If something is no longer relevant, you may remove it
|
||
|
||
Use this EXACT format:
|
||
|
||
## Goal
|
||
[Preserve existing goals, add new ones if the task expanded]
|
||
|
||
## Constraints & Preferences
|
||
- [Preserve existing, add new ones discovered]
|
||
|
||
## Progress
|
||
### Done
|
||
- [x] [Include previously done items AND newly completed items]
|
||
|
||
### In Progress
|
||
- [ ] [Current work - update based on progress]
|
||
|
||
### Blocked
|
||
- [Current blockers - remove if resolved]
|
||
|
||
## Key Decisions
|
||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||
|
||
## Next Steps
|
||
1. [Update based on current state]
|
||
|
||
## Critical Context
|
||
- [Preserve important context, add new if needed]
|
||
|
||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||
|
||
TURN_PREFIX_SUMMARIZATION_PROMPT = """This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||
|
||
Summarize the prefix to provide context for the retained suffix:
|
||
|
||
## Original Request
|
||
[What did the user ask for in this turn?]
|
||
|
||
## Early Progress
|
||
- [Key decisions and work done in the prefix]
|
||
|
||
## Context for Suffix
|
||
- [Information needed to understand the retained recent work]
|
||
|
||
Be concise. Focus on what's needed to understand the kept suffix."""
|
||
|
||
|
||
# ======================================================================
|
||
# 对话序列化 —— 1:1 对照 utils.ts serializeConversation
|
||
# ======================================================================
|
||
def _safe_json(value) -> str:
|
||
try:
|
||
s = json.dumps(value, ensure_ascii=False)
|
||
return s if s is not None else "undefined"
|
||
except Exception:
|
||
return "[unserializable]"
|
||
|
||
|
||
def _content_text(content) -> str:
|
||
"""pi contentText: str 或 [{type:"text",text}] 列表取文本拼接"""
|
||
if content is None:
|
||
return ""
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts = []
|
||
for b in content:
|
||
if isinstance(b, dict) and b.get("type") == "text":
|
||
parts.append(str(b.get("text", "")))
|
||
return "\n".join(p for p in parts if p)
|
||
return str(content)
|
||
|
||
|
||
def _truncate_for_summary(text: str, max_chars: int) -> str:
|
||
"""pi utils.ts truncateForSummary(逐字逻辑)"""
|
||
if len(text) <= max_chars:
|
||
return text
|
||
truncated = len(text) - max_chars
|
||
return f"{text[:max_chars]}\n\n[... {truncated} more characters truncated]"
|
||
|
||
|
||
def serialize_conversation(messages: List[AgentMessage]) -> str:
|
||
"""
|
||
1:1 对照 utils.ts serializeConversation(输出格式逐字一致):
|
||
[User]: ...
|
||
[Assistant thinking]: ...
|
||
[Assistant]: ...
|
||
[Assistant tool calls]: name(k=v, k2=v2); name2(...)
|
||
[Tool result]: ...(超 2000 字符截断)
|
||
"""
|
||
parts: List[str] = []
|
||
for msg in messages:
|
||
if msg.role == "user":
|
||
content = _content_text(msg.content)
|
||
if content:
|
||
parts.append(f"[User]: {content}")
|
||
elif msg.role == "assistant":
|
||
thinking_parts = []
|
||
tool_calls = []
|
||
if msg.reasoning:
|
||
thinking_parts.append(msg.reasoning)
|
||
text = _content_text(msg.content)
|
||
for tc in (msg.tool_calls or []):
|
||
args_str = ", ".join(f"{k}={_safe_json(v)}"
|
||
for k, v in (tc.arguments or {}).items())
|
||
tool_calls.append(f"{tc.name}({args_str})")
|
||
if thinking_parts:
|
||
parts.append(f"[Assistant thinking]: {chr(10).join(thinking_parts)}")
|
||
if text:
|
||
parts.append(f"[Assistant]: {text}")
|
||
if tool_calls:
|
||
parts.append(f"[Assistant tool calls]: {'; '.join(tool_calls)}")
|
||
elif msg.role == "toolResult":
|
||
content = _content_text(msg.content)
|
||
if content:
|
||
parts.append(f"[Tool result]: "
|
||
f"{_truncate_for_summary(content, TOOL_RESULT_MAX_CHARS)}")
|
||
return "\n\n".join(parts)
|
||
|
||
|
||
# ======================================================================
|
||
# 文件操作提取 —— 1:1 对照 utils.ts extractFileOps*/computeFileLists/formatFileOperations
|
||
# ======================================================================
|
||
class FileOperations:
|
||
def __init__(self):
|
||
self.read: Set[str] = set()
|
||
self.written: Set[str] = set()
|
||
self.edited: Set[str] = set()
|
||
|
||
|
||
def extract_file_ops_from_message(message: AgentMessage, file_ops: FileOperations):
|
||
"""pi: assistant 的 toolCall 参数里 path 字段 → read/write/edit 归类"""
|
||
if message.role != "assistant":
|
||
return
|
||
for tc in (message.tool_calls or []):
|
||
args = tc.arguments or {}
|
||
path = args.get("path")
|
||
if not isinstance(path, str) or not path:
|
||
continue
|
||
if tc.name == "read":
|
||
file_ops.read.add(path)
|
||
elif tc.name == "write":
|
||
file_ops.written.add(path)
|
||
elif tc.name == "edit":
|
||
file_ops.edited.add(path)
|
||
|
||
|
||
def compute_file_lists(file_ops: FileOperations) -> Tuple[List[str], List[str]]:
|
||
"""pi computeFileLists: modified = edited|written;readOnly = read-modified;均排序"""
|
||
modified = file_ops.edited | file_ops.written
|
||
read_only = sorted(f for f in file_ops.read if f not in modified)
|
||
return read_only, sorted(modified)
|
||
|
||
|
||
def format_file_operations(read_files: List[str], modified_files: List[str]) -> str:
|
||
"""pi formatFileOperations: <read-files>/<modified-files> 标签拼接"""
|
||
sections = []
|
||
if read_files:
|
||
sections.append("<read-files>\n" + "\n".join(read_files) + "\n</read-files>")
|
||
if modified_files:
|
||
sections.append("<modified-files>\n" + "\n".join(modified_files) + "\n</modified-files>")
|
||
if not sections:
|
||
return ""
|
||
return "\n\n" + "\n\n".join(sections)
|
||
|
||
|
||
# ======================================================================
|
||
# 切分点 —— 1:1 对照 findValidCutPoints / findTurnStartIndex / findCutPoint
|
||
# ======================================================================
|
||
@dataclass
|
||
class CutPointResult:
|
||
"""pi: interface CutPointResult"""
|
||
first_kept_index: int # 保留段首条在 compactable 列表里的下标
|
||
turn_start_index: int = -1 # 断轮时:该轮起点(user 消息)下标;否则 -1
|
||
is_split_turn: bool = False
|
||
|
||
|
||
def find_valid_cut_points(messages: List[AgentMessage],
|
||
start_index: int, end_index: int) -> List[int]:
|
||
"""
|
||
pi findValidCutPoints:消息角色为 user/assistant 的位置是有效切点
|
||
(toolResult 不能做切点——它会与前面的 toolCall 分离)。
|
||
pi 里的 bashExecution/branchSummary/compactionSummary 等角色
|
||
在 haocode 消息模型中不存在,等价规则即 role in (user, assistant)。
|
||
"""
|
||
cut_points = []
|
||
for i in range(start_index, end_index):
|
||
if messages[i].role in ("user", "assistant"):
|
||
cut_points.append(i)
|
||
return cut_points
|
||
|
||
|
||
def find_turn_start(messages: List[AgentMessage], entry_index: int,
|
||
start_index: int) -> int:
|
||
"""pi findTurnStartIndex:向前找本轮起点(user 消息 / branch_summary)"""
|
||
for i in range(entry_index, start_index - 1, -1):
|
||
if messages[i].role == "user":
|
||
return i
|
||
return -1
|
||
|
||
|
||
def _estimate(msg: AgentMessage) -> int:
|
||
from .context import estimate_message_tokens
|
||
return estimate_message_tokens(msg)
|
||
|
||
|
||
def find_cut_point(messages: List[AgentMessage], start_index: int,
|
||
end_index: int, keep_recent_tokens: int) -> CutPointResult:
|
||
"""
|
||
1:1 对照 pi findCutPoint:
|
||
1. 从尾部向前累计 token,直到累计 >= keep_recent_tokens
|
||
2. 取该位置(含)之后的第一个有效切点
|
||
3. 切点不是 user 消息 → 断轮:找本轮起点,前缀单独摘要
|
||
(pi 的「回退跳过状态条目」循环针对 session 状态条目;
|
||
haocode 消息列表没有状态条目,等价省略。)
|
||
"""
|
||
cut_points = find_valid_cut_points(messages, start_index, end_index)
|
||
if not cut_points:
|
||
return CutPointResult(first_kept_index=start_index)
|
||
|
||
accumulated = 0
|
||
cut_index = cut_points[0]
|
||
for i in range(end_index - 1, start_index - 1, -1):
|
||
accumulated += _estimate(messages[i])
|
||
if accumulated >= keep_recent_tokens:
|
||
for c in cut_points:
|
||
if c >= i:
|
||
cut_index = c
|
||
break
|
||
break
|
||
|
||
is_user = messages[cut_index].role == "user"
|
||
turn_start = -1 if is_user else find_turn_start(messages, cut_index, start_index)
|
||
is_split = (not is_user) and turn_start != -1
|
||
return CutPointResult(first_kept_index=cut_index,
|
||
turn_start_index=turn_start,
|
||
is_split_turn=is_split)
|
||
|
||
|
||
# ======================================================================
|
||
# 压缩准备 —— 1:1 对照 prepareCompaction
|
||
# ======================================================================
|
||
@dataclass
|
||
class CompactionPreparation:
|
||
"""pi: interface CompactionPreparation"""
|
||
messages_to_summarize: List[AgentMessage] = field(default_factory=list)
|
||
turn_prefix_messages: List[AgentMessage] = field(default_factory=list)
|
||
retained_tail: List[AgentMessage] = field(default_factory=list)
|
||
is_split_turn: bool = False
|
||
tokens_before: int = 0
|
||
previous_summary: Optional[str] = None
|
||
file_ops: FileOperations = field(default_factory=FileOperations)
|
||
# 注意(P2-04 跨平台修复):Python 3.11+ 的 dataclass 禁止不可哈希默认值
|
||
# (eq dataclass 实例 __hash__=None);Ubuntu 24.04 自带 3.12,
|
||
# 直接写 `= DEFAULT_COMPACTION_SETTINGS` 在 3.11+ 会 ValueError。
|
||
# default_factory 返回同一共享实例,与 3.10 行为完全一致。
|
||
settings: CompactionSettings = field(default_factory=lambda: DEFAULT_COMPACTION_SETTINGS)
|
||
|
||
|
||
def _compaction_diag(line: str) -> None:
|
||
"""压缩自诊断日志(与 recovery.compact_diag_log 同一文件/规范:
|
||
print + 追加 compaction_diag.log + 失败静默)。"""
|
||
import os as _os
|
||
import time as _time
|
||
t = _time.time()
|
||
stamp = (f"[{_time.strftime('%H:%M:%S', _time.localtime(t))}"
|
||
f".{int(t * 1000) % 1000:03d}]")
|
||
print(f"{stamp} {line}", flush=True)
|
||
try:
|
||
path = _os.path.join(
|
||
_os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))),
|
||
"compaction_diag.log")
|
||
with open(path, "a", encoding="utf-8") as f:
|
||
f.write(f"{stamp} {line}\n")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _nothing_to_summarize_diag(messages: List[AgentMessage],
|
||
compactable: List[AgentMessage],
|
||
cut: CutPointResult,
|
||
settings: CompactionSettings) -> str:
|
||
"""prepare_compaction 返回 None(无可摘要内容)时的自诊断行:
|
||
记录分支子类型 + 关键量,下次失败可直接从日志定位原因。
|
||
|
||
子类型:
|
||
no_valid_cut_points —— 可压缩范围内没有任何 user/assistant 条目
|
||
total_below_keep_recent —— 可压缩范围总量 < keep_recent(通常意味着
|
||
上下文主体是旧摘要本身,无新内容)
|
||
cut_pinned_at_zero —— 预算从尾部累加只在 i=0 才达标,即首条
|
||
条目独占 ≥ (总量-keep_recent) 的 token
|
||
(单条超长消息/巨型工具输出主导上下文)
|
||
"""
|
||
from .context import estimate_context_tokens, estimate_message_tokens
|
||
try:
|
||
total = estimate_context_tokens(messages).tokens
|
||
n_user = sum(1 for m in compactable if m.role == "user")
|
||
n_asst = sum(1 for m in compactable if m.role == "assistant")
|
||
ests = [(estimate_message_tokens(m), i) for i, m in enumerate(compactable)]
|
||
e0 = ests[0][0] if ests else 0
|
||
rest = sum(e for e, _ in ests[1:])
|
||
top3 = sorted(ests, reverse=True)[:3]
|
||
roles = " ".join(("T" if m.role == "toolResult" else m.role[0].upper())
|
||
for m in compactable[:10])
|
||
if not find_valid_cut_points(compactable, 0, len(compactable)):
|
||
sub = "no_valid_cut_points"
|
||
elif sum(e for e, _ in ests) < settings.keep_recent_tokens:
|
||
sub = "total_below_keep_recent"
|
||
else:
|
||
sub = "cut_pinned_at_zero"
|
||
big = "; ".join(f"idx{i}={e}" for e, i in top3)
|
||
return (f"[COMPACT_NONE] branch=nothing_to_summarize sub={sub} "
|
||
f"msgs={len(messages)} est_total={total} compactable={len(compactable)} "
|
||
f"keep_recent={settings.keep_recent_tokens} "
|
||
f"first_kept={cut.first_kept_index} turn_start={cut.turn_start_index} "
|
||
f"split={cut.is_split_turn} user={n_user} asst={n_asst} "
|
||
f"est_first={e0} est_rest={rest} top3=[{big}] head_roles=[{roles}]")
|
||
except Exception as ex: # 诊断本身失败不能影响主流程
|
||
return f"[COMPACT_NONE] branch=nothing_to_summarize (diag failed: {ex})"
|
||
|
||
|
||
def prepare_compaction(messages: List[AgentMessage],
|
||
settings: Optional[CompactionSettings] = None
|
||
) -> Optional[CompactionPreparation]:
|
||
"""
|
||
1:1 对照 pi prepareCompaction:
|
||
- 上一条压缩摘要(messages[0].kind == "compaction_summary")不重复摘要,
|
||
其内容作为 previousSummary 走迭代更新提示词
|
||
- 可压缩范围 = 摘要之后的全部消息(pi 里等价于「上次保留尾 + 新消息」)
|
||
- 切点在可压缩范围内选;tokens_before 按完整上下文(含摘要消息)估算
|
||
不可压缩(空/无摘要对象)时返回 None(对照 pi 返回 ok(undefined))。
|
||
"""
|
||
from .context import estimate_context_tokens
|
||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||
if not messages:
|
||
_compaction_diag("[COMPACT_NONE] branch=empty_messages")
|
||
return None
|
||
if messages[-1].kind == "compaction_summary":
|
||
_compaction_diag(
|
||
f"[COMPACT_NONE] branch=tail_is_summary msgs={len(messages)} "
|
||
f"tail_head={_content_text(messages[-1].content)[:60]!r}")
|
||
return None
|
||
|
||
previous_summary = None
|
||
if messages[0].kind == "compaction_summary":
|
||
previous_summary = _content_text(messages[0].content)
|
||
compactable = messages[1:]
|
||
else:
|
||
compactable = messages
|
||
if not compactable:
|
||
_compaction_diag("[COMPACT_NONE] branch=compactable_empty (上下文只剩旧摘要)")
|
||
return None
|
||
|
||
tokens_before = estimate_context_tokens(messages).tokens
|
||
cut = find_cut_point(compactable, 0, len(compactable),
|
||
settings.keep_recent_tokens)
|
||
|
||
history_end = cut.turn_start_index if cut.is_split_turn else cut.first_kept_index
|
||
messages_to_summarize = compactable[:history_end]
|
||
turn_prefix_messages = []
|
||
if cut.is_split_turn:
|
||
turn_prefix_messages = compactable[cut.turn_start_index:cut.first_kept_index]
|
||
retained_tail = compactable[cut.first_kept_index:]
|
||
|
||
if not messages_to_summarize and not turn_prefix_messages:
|
||
# 🆕 自诊断:记录是哪种子条件导致无东西可摘要(见 _nothing_to_summarize_diag)
|
||
_compaction_diag(_nothing_to_summarize_diag(messages, compactable, cut,
|
||
settings))
|
||
return None # 没有可摘要内容
|
||
|
||
file_ops = FileOperations()
|
||
for m in messages_to_summarize:
|
||
extract_file_ops_from_message(m, file_ops)
|
||
if cut.is_split_turn:
|
||
for m in turn_prefix_messages:
|
||
extract_file_ops_from_message(m, file_ops)
|
||
|
||
return CompactionPreparation(
|
||
messages_to_summarize=messages_to_summarize,
|
||
turn_prefix_messages=turn_prefix_messages,
|
||
retained_tail=retained_tail,
|
||
is_split_turn=cut.is_split_turn,
|
||
tokens_before=tokens_before,
|
||
previous_summary=previous_summary,
|
||
file_ops=file_ops,
|
||
settings=settings,
|
||
)
|
||
|
||
|
||
# ======================================================================
|
||
# 摘要生成 —— 1:1 对照 generateSummaryWithUsage / generateTurnPrefixSummary
|
||
# ======================================================================
|
||
def _build_summary_prompt(conversation_text: str, previous_summary: Optional[str]) -> str:
|
||
"""pi generateSummaryWithUsage 的 prompt 组装(逐字结构)"""
|
||
base = UPDATE_SUMMARIZATION_PROMPT if previous_summary else SUMMARIZATION_PROMPT
|
||
prompt = f"<conversation>\n{conversation_text}\n</conversation>\n\n"
|
||
if previous_summary:
|
||
prompt += f"<previous-summary>\n{previous_summary}\n</previous-summary>\n\n"
|
||
prompt += base
|
||
return prompt
|
||
|
||
|
||
def generate_summary(messages: List[AgentMessage],
|
||
summarize_fn: Callable[[str, str, int], str],
|
||
reserve_tokens: int,
|
||
model_max_tokens: int,
|
||
previous_summary: Optional[str] = None
|
||
) -> Tuple[Optional[str], Optional[str]]:
|
||
"""
|
||
返回 (summary_text, error)。
|
||
maxTokens = min(0.8 × reserveTokens, model.maxTokens) —— 1:1 对照。
|
||
"""
|
||
max_tokens = min(
|
||
int(0.8 * reserve_tokens),
|
||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||
)
|
||
conversation = serialize_conversation(messages)
|
||
prompt = _build_summary_prompt(conversation, previous_summary)
|
||
try:
|
||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||
except Exception as e:
|
||
return None, f"Summarization failed: {e}"
|
||
if not text or not text.strip():
|
||
return None, "Summarization failed: empty response"
|
||
return text.strip(), None
|
||
|
||
|
||
def generate_turn_prefix_summary(messages: List[AgentMessage],
|
||
summarize_fn: Callable[[str, str, int], str],
|
||
reserve_tokens: int,
|
||
model_max_tokens: int
|
||
) -> Tuple[Optional[str], Optional[str]]:
|
||
"""maxTokens = min(0.5 × reserveTokens, model.maxTokens) —— 1:1 对照"""
|
||
max_tokens = min(
|
||
int(0.5 * reserve_tokens),
|
||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||
)
|
||
conversation = serialize_conversation(messages)
|
||
prompt = f"<conversation>\n{conversation}\n</conversation>\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
|
||
try:
|
||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||
except Exception as e:
|
||
return None, f"Turn prefix summarization failed: {e}"
|
||
if not text or not text.strip():
|
||
return None, "Turn prefix summarization failed: empty response"
|
||
return text.strip(), None
|
||
|
||
|
||
# ======================================================================
|
||
# 主入口 —— 1:1 对照 compact()
|
||
# ======================================================================
|
||
def compact_context(messages: List[AgentMessage],
|
||
model: ModelConfig,
|
||
summarize_fn: Callable[[str, str, int], str],
|
||
settings: Optional[CompactionSettings] = None
|
||
) -> Optional[List[AgentMessage]]:
|
||
"""
|
||
执行压缩。返回 [压缩摘要消息] + 保留尾巴;不可压缩时返回 None。
|
||
|
||
摘要消息: role="user", kind="compaction_summary",content 为纯摘要文本
|
||
(pi 的 compaction 条目;下次压缩时自动走迭代更新提示词)。
|
||
|
||
断轮(切点落在某轮中间)时 1:1 对照 pi compact():
|
||
历史摘要 与 轮前缀摘要 分两次 LLM 调用,拼接为
|
||
{history}\n\n---\n\n**Turn Context (split turn):**\n\n{prefix}
|
||
最后追加文件操作附录(<read-files>/<modified-files>)。
|
||
"""
|
||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||
prep = prepare_compaction(messages, settings)
|
||
if prep is None:
|
||
return None
|
||
|
||
history_text: Optional[str] = None
|
||
prefix_error: Optional[str] = None
|
||
history_error: Optional[str] = None
|
||
|
||
if prep.is_split_turn and prep.turn_prefix_messages:
|
||
if prep.messages_to_summarize:
|
||
history_text, history_error = generate_summary(
|
||
prep.messages_to_summarize, summarize_fn,
|
||
prep.settings.reserve_tokens, model.max_tokens,
|
||
prep.previous_summary)
|
||
if history_error:
|
||
return _degraded_compact(prep, history_error)
|
||
else:
|
||
history_text = "No prior history."
|
||
prefix_text, prefix_error = generate_turn_prefix_summary(
|
||
prep.turn_prefix_messages, summarize_fn,
|
||
prep.settings.reserve_tokens, model.max_tokens)
|
||
if prefix_error:
|
||
return _degraded_compact(prep, prefix_error)
|
||
summary = (f"{history_text}\n\n---\n\n"
|
||
f"**Turn Context (split turn):**\n\n{prefix_text}")
|
||
else:
|
||
if not prep.messages_to_summarize:
|
||
return None
|
||
summary, history_error = generate_summary(
|
||
prep.messages_to_summarize, summarize_fn,
|
||
prep.settings.reserve_tokens, model.max_tokens,
|
||
prep.previous_summary)
|
||
if history_error:
|
||
return _degraded_compact(prep, history_error)
|
||
|
||
read_files, modified_files = compute_file_lists(prep.file_ops)
|
||
summary += format_file_operations(read_files, modified_files)
|
||
|
||
summary_msg = AgentMessage(role="user", content=summary,
|
||
kind="compaction_summary")
|
||
return [summary_msg] + prep.retained_tail
|
||
|
||
|
||
def _degraded_compact(prep: CompactionPreparation,
|
||
error: str) -> Optional[List[AgentMessage]]:
|
||
"""
|
||
已声明偏差(对照 pi: 直接返回 CompactionError):
|
||
桌面应用优先「不丢上下文」——摘要失败时降级为机械摘录,
|
||
保留尾巴原样不动。
|
||
"""
|
||
old = prep.messages_to_summarize + prep.turn_prefix_messages
|
||
if not old:
|
||
return None
|
||
excerpt = serialize_conversation(old)[-500:]
|
||
summary = (f"(自动压缩:摘要生成失败 [{error}],以下为旧对话尾部摘录)\n\n{excerpt}")
|
||
summary_msg = AgentMessage(role="user", content=summary,
|
||
kind="compaction_summary")
|
||
return [summary_msg] + prep.retained_tail
|