kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

refactor(middleware): 简化 SummaryOffloadMiddleware 卸载逻辑

Browse Source 
- 移除 immediate_offload_threshold 参数，仅在触发 Summary 时卸载
- 优化工具结果卸载文件命名（工具名称-id）并添加头部信息
- 调整 FilesystemMiddleware 后端工厂函数接收 runtime 参数
- 优化 DeepAgent 提示词，调整报告格式要求
This commit is contained in:
Wenjie Zhang 2026-02-03 12:13:13 +08:00
parent 66cfb50d3f
commit cab244227f
3 changed files with 86 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

122
src/agents/common/middlewares/summary_middleware.py
View File

@ -135,17 +135,32 @@ def _offload_tool_result(msg: ToolMessage, threshold: int, token_counter: TokenC
if token_counter([msg]) <= threshold:
return None
# 生成文件路径
message_id = msg.id or str(uuid.uuid4())
file_path = f"{_OFFLOAD_DIR}/{message_id}"
# 获取工具名称和参数
tool_name = msg.name or "unknown"
tool_call_id = msg.tool_call_id or ""
# 保存到 files 格式
# 生成文件路径 (工具名称-xxx)
message_id = msg.id or str(uuid.uuid4())[:8]
safe_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in tool_name)
file_path = f"{_OFFLOAD_DIR}/{safe_name}-{message_id}"
# 构建文件头部信息
header_lines = [
f"=== Tool Invocation ===",
f"Tool: {tool_name}",
f"Tool Call ID: {tool_call_id}",
"=" * 40,
"",
]
header = "\n".join(header_lines)
# 保存到 files 格式（包含头部信息）
from datetime import datetime
timestamp = datetime.now().isoformat()
files_update = {
file_path: {
"content": content_str.splitlines(keepends=True),
"content": [header + content_str],
"created_at": timestamp,
"modified_at": timestamp,
}
@ -195,12 +210,11 @@ class SummaryOffloadMiddleware(AgentMiddleware):
基于 LangChain SummarizationMiddleware，额外功能：
- 保留原有的 summary 历史记录逻辑
- 将 ToolMessage 的 results 卸载到虚拟文件系统
1. 立即卸载：超过 immediate_offload_threshold (默认 4k tokens)
2. Summary 时卸载：超过 summary_offload_threshold (默认 1k tokens)
- 智能保留策略：
1. 触发 Summary (达到 trigger) 时，首先进行 Summary 卸载
2. 只有当总 Token 数超过 max_retention_ratio * trigger 时，才进行消息清理(Summary)
3. 始终保留 System Message
1. 触发 Summary 时，卸载超过阈值的工具结果
2. 智能保留策略：
- 触发 Summary 时，首先进行卸载
- 只有当总 Token 数超过 max_retention_ratio * trigger 时，才进行消息清理(Summary)
- 始终保留 System Message
"""
def __init__(
@ -212,8 +226,7 @@ class SummaryOffloadMiddleware(AgentMiddleware):
token_counter: TokenCounter = count_tokens_approximately,
summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
trim_tokens_to_summarize: int | None = _DEFAULT_TRIM_TOKEN_LIMIT,
# 新增/修改：工具结果卸载参数
immediate_offload_threshold: int = 4000,
# 工具结果卸载参数
summary_offload_threshold: int = 1000,
max_retention_ratio: float = 0.6,
**deprecated_kwargs: Any,
@ -227,7 +240,6 @@ class SummaryOffloadMiddleware(AgentMiddleware):
token_counter: token 计数函数
summary_prompt: 生成摘要的提示词模板
trim_tokens_to_summarize: 准备摘要消息时的最大 token 数
immediate_offload_threshold: 立即卸载阈值（token 数），默认 4000
summary_offload_threshold: Summary 时卸载阈值（token 数），默认 1000
max_retention_ratio: 触发 Summary 后，如果不超过此比例（相对于 trigger），则不删除消息。默认 0.6
"""
@ -265,7 +277,6 @@ class SummaryOffloadMiddleware(AgentMiddleware):
self.trim_tokens_to_summarize = trim_tokens_to_summarize
# 工具结果卸载配置
self.immediate_offload_threshold = immediate_offload_threshold
self.summary_offload_threshold = summary_offload_threshold
self.max_retention_ratio = max_retention_ratio
@ -305,53 +316,34 @@ class SummaryOffloadMiddleware(AgentMiddleware):
self._ensure_message_ids(messages)
# 1. 立即卸载 (Scenario 1: > 4k)
files_update, modified_messages = _offload_tool_results(
messages, self.immediate_offload_threshold, self.token_counter
)
# 注意: modified_messages 指向 state['messages'] 中的对象 (in-place modification)
total_tokens = self.token_counter(messages)
# 2. 检查是否触发 Summary
# 1. 检查是否触发 Summary
if not self._should_summarize(messages, total_tokens):
if files_update:
return {"files": files_update, "messages": modified_messages}
return None
# 3. 触发 Summary (Scenario 2)
# Phase A: Aggressive Offload (> 1k)
# 2. 触发 Summary：卸载超阈值的工具结果
files_update: dict[str, Any] = {}
modified_messages: list[AnyMessage] = []
agg_files, agg_msgs = _offload_tool_results(messages, self.summary_offload_threshold, self.token_counter)
files_update = agg_files
modified_messages = agg_msgs
files_update.update(agg_files)
# 合并 modified messages 列表 (注意去重或直接使用最新的)
all_modified_ids = {m.id for m in modified_messages}
for m in agg_msgs:
if m.id not in all_modified_ids:
modified_messages.append(m)
all_modified_ids.add(m.id)
# Phase B: 检查 Retention Ratio
# 3. 检查 Retention Ratio
current_tokens = self.token_counter(messages)
trigger_value = self._get_token_trigger_value()
# 如果没有明确的 token trigger，或者当前 tokens 已经小于等于 (trigger * ratio)
# 则不进行删除/Summary，只返回 offload 的结果
retention_limit = float("inf")
if trigger_value:
retention_limit = trigger_value * self.max_retention_ratio
if current_tokens <= retention_limit:
# 不删除消息，只卸载文件
if files_update:
return {"files": files_update, "messages": modified_messages}
return None
# Phase C: 超过 limit，需要 Eviction (Summary)
# 保护 System Message
# 4. 超过 limit，需要 Eviction (Summary)
system_msg_count = 0
messages_to_process = messages
@ -359,22 +351,15 @@ class SummaryOffloadMiddleware(AgentMiddleware):
system_msg_count = 1
messages_to_process = messages[1:]
# 找到 cutoff，使得保留的部分 <= retention_limit
# 注意: 这里简化处理，未扣除 system message 的 tokens，
# 实际保留的可能会略多于 limit (即 limit + system_msg_tokens)
cutoff_relative = self._find_cutoff_by_token_limit(messages_to_process, int(retention_limit))
cutoff_index = system_msg_count + cutoff_relative
if cutoff_index <= system_msg_count:
# 不需要截断或截断点无效
if files_update:
return {"files": files_update, "messages": modified_messages}
return None
messages_to_summarize, preserved_messages = self._partition_messages(messages, cutoff_index)
# 再次确保 preserved_messages 里的 tool results 是卸载过的 (理论上 Phase A 已经做了)
summary = self._create_summary(messages_to_summarize)
new_messages = self._build_new_messages(summary)
@ -402,50 +387,34 @@ class SummaryOffloadMiddleware(AgentMiddleware):
self._ensure_message_ids(messages)
# 1. 立即卸载 (Scenario 1: > 4k)
files_update, modified_messages = _offload_tool_results(
messages, self.immediate_offload_threshold, self.token_counter
)
total_tokens = self.token_counter(messages)
# 2. 检查是否触发 Summary
# 1. 检查是否触发 Summary
if not self._should_summarize(messages, total_tokens):
if files_update:
return {"files": files_update, "messages": modified_messages}
return None
# 3. 触发 Summary (Scenario 2)
# Phase A: Aggressive Offload (> 1k)
# 2. 触发 Summary：卸载超阈值的工具结果
files_update: dict[str, Any] = {}
modified_messages: list[AnyMessage] = []
agg_files, agg_msgs = _offload_tool_results(messages, self.summary_offload_threshold, self.token_counter)
files_update = agg_files
modified_messages = agg_msgs
files_update.update(agg_files)
# 合并 modified messages
all_modified_ids = {m.id for m in modified_messages}
for m in agg_msgs:
if m.id not in all_modified_ids:
modified_messages.append(m)
all_modified_ids.add(m.id)
# Phase B: 检查 Retention Ratio
# 3. 检查 Retention Ratio
current_tokens = self.token_counter(messages)
trigger_value = self._get_token_trigger_value()
retention_limit = float("inf")
if trigger_value:
retention_limit = trigger_value * self.max_retention_ratio
if current_tokens <= retention_limit:
# 不删除消息，只卸载文件
if files_update:
return {"files": files_update, "messages": modified_messages}
return None
# Phase C: 超过 limit，需要 Eviction (Summary)
# 保护 System Message
# 4. 超过 limit，需要 Eviction (Summary)
system_msg_count = 0
messages_to_process = messages
@ -454,7 +423,6 @@ class SummaryOffloadMiddleware(AgentMiddleware):
messages_to_process = messages[1:]
cutoff_relative = self._find_cutoff_by_token_limit(messages_to_process, int(retention_limit))
cutoff_index = system_msg_count + cutoff_relative
if cutoff_index <= system_msg_count:
@ -762,7 +730,6 @@ def create_summary_offload_middleware(
*,
trigger: ContextSize | list[ContextSize] | None = None,
keep: ContextSize = ("messages", _DEFAULT_MESSAGES_TO_KEEP),
immediate_offload_threshold: int = 4000,
summary_offload_threshold: int = 1000,
max_retention_ratio: float = 0.6,
) -> SummaryOffloadMiddleware:
@ -771,7 +738,6 @@ def create_summary_offload_middleware(
model=model,
trigger=trigger,
keep=keep,
immediate_offload_threshold=immediate_offload_threshold,
summary_offload_threshold=summary_offload_threshold,
max_retention_ratio=max_retention_ratio,
)

23
src/agents/deep_agent/context.py
View File

@ -12,6 +12,7 @@ DEEP_PROMPT = """你是一位专家级研究员。你的工作是进行彻底的
首先，你应该并行使用 research-agent 进行深入研究。，当你认为有足够的信息来撰写最终报告时，就把它写入 `final_report.md`
其次（如果有必要），你可以调用 critique-agent 来获取对最终报告（文件）的评价。
之后（如果需要）你可以做更多的研究并编辑 `final_report.md`
最终通知用户，已生成完毕，可以在状态工作台中下载最终报告。
你可以根据需要重复这个过程，直到你对结果满意为止。
@ -32,8 +33,8 @@ DEEP_PROMPT = """你是一位专家级研究员。你的工作是进行彻底的
2. 包含研究中的具体事实和见解
3. 使用 [标题](URL) 格式引用相关来源
4. 图片应该使用 ![描述](图片URL) 格式引用
4. 提供平衡、透彻的分析。尽可能全面，并包含与整体研究问题相关的所有信息。使用你进行深入研究，并期望得到详细、全面的答案
5. 在末尾包含一个“来源”部分，列出所有引用的链接
5. 提供平衡、透彻的分析。尽可能全面，并包含与整体研究问题相关的所有信息。使用你进行深入研究，并期望得到详细、全面的答案
6. 在末尾包含一个“来源”部分，列出所有引用的链接
你可以用多种不同的方式来组织你的报告。以下是一些例子：
@ -58,34 +59,26 @@ DEEP_PROMPT = """你是一位专家级研究员。你的工作是进行彻底的
4/ 概念3
5/ 结论
如果你认为你可以用一个部分来回答问题，你也可以这样做！
1/ 答案
请记住：章节是一个非常灵活和松散的概念。你可以按照你认为最好的方式来组织你的报告，包括上面没有列出的方式！
确保你的各个部分是连贯的，并且对读者来说是有意义的。
对于报告的每个部分，请执行以下操作：
- 使用简单、清晰的语言
- 使用简单、清晰的语言，报告内容要详实。
- 格式符合学术论文、技术报告、公文文件等格式，不要过于随意，段落不可过于简短。
- 对报告的每个部分使用 ## 作为章节标题（Markdown 格式）
- 绝不要将自己称为报告的作者。这应该是一份专业的报告，不含任何自我指涉的语言。
- 不要在报告中说你正在做什么。只需撰写报告，不要添加任何你自己的评论。
- 每个部分的长度应足以用你收集到的信息。预计各部分会长且详尽。你正在撰写一份深入的研究报告，用户会期望得到透彻的答案。
- 在适当的时候使用项目符号来列出信息，但默认情况下，请以段落形式撰写。
请记住：
简报和研究可能是英文的，但在撰写最终答案时，你需要将这些信息翻译成正确的语言。
确保最终答案报告的语言与消息历史中的人类信息语言相同。
用清晰的 markdown 格式化报告，结构合理，并在适当的地方包含来源引用。
<引用规则>
- 在你的文本中为每个唯一的 URL 分配一个引文编号
- 在你的文本中为每个唯一的 URL/文件路径 分配一个引文编号
- 以 ### 来源 结尾，列出每个来源及其对应的编号
- 重要提示：无论你选择哪些来源，最终列表中的来源编号都应连续无间断（1,2,3,4...）
- 每个来源都应该是列表中的一个独立行项目，这样在 markdown 中它会被渲染成一个列表。
- 示例格式：
[1] 来源标题: URL
[2] 来源标题: URL
[1] 来源标题: URL/文件路径
[2] 来源标题: URL/文件路径
- 引用非常重要。请确保包含这些内容，并特别注意确保其正确性。用户通常会使用这些引文来查找更多信息。
</引用规则>
</report_instructions>

74
src/agents/deep_agent/graph.py
View File

@ -21,17 +21,16 @@ from src.services.mcp_service import get_tools_from_all_servers
from .context import DeepContext
def _create_filesystem_backend_factory() -> CompositeBackend:
"""创建混合文件存储后端工厂函数。
def _create_filesystem_backend_factory(rt) -> CompositeBackend:
"""创建混合文件存储后端工厂函数（供 FilesystemMiddleware 使用）。
/attachments/* 路由到真实文件系统（供附件中间件使用）
其他路径使用 StateBackend（内存存储，用于临时文件和大结果卸载）
返回 CompositeBackend 实例，供 FilesystemMiddleware 使用。
注意：StateBackend 会在 middleware 初始化时由 runtime 注入。
注意：rt (runtime) 由 FilesystemMiddleware 在调用时自动传入。
"""
return CompositeBackend(
default=StateBackend, # 使用类而非实例，FilesystemMiddleware 会自动注入 runtime
default=StateBackend(rt), # 传入 runtime 创建实例
routes={
"/attachments/": FilesystemBackend(root_dir=".", virtual_mode=False),
},
@ -118,49 +117,44 @@ class DeepAgent(BaseAgent):
# Build subagents with search tools
research_sub_agent = _get_research_sub_agent(search_tools)
summary_middleware = SummaryOffloadMiddleware(
model=model,
trigger=("tokens", 160000),
trim_tokens_to_summarize=None,
summary_offload_threshold=1000,
max_retention_ratio=0.6,
)
subagents_middleware = SubAgentMiddleware(
default_model=sub_model,
default_tools=search_tools,
subagents=[critique_sub_agent, research_sub_agent],
default_middleware=[
FilesystemMiddleware(backend=_create_filesystem_backend_factory),
RuntimeConfigMiddleware(
model_context_name="subagents_model",
enable_model_override=True,
enable_system_prompt_override=False,
enable_tools_override=False,
),
PatchToolCallsMiddleware(),
summary_middleware,
],
general_purpose_agent=True
)
# 使用 create_deep_agent 创建深度智能体
fs_backend_factory = _create_filesystem_backend_factory()
graph = create_agent(
model=model,
system_prompt=context.system_prompt,
middleware=[
save_attachments_to_fs, # 附件保存到文件系统
FilesystemMiddleware(backend=_create_filesystem_backend_factory),
RuntimeConfigMiddleware(extra_tools=all_mcp_tools),
save_attachments_to_fs, # 附件保存到文件系统
TodoListMiddleware(),
FilesystemMiddleware(backend=fs_backend_factory, tool_token_limit_before_evict=5000),
SubAgentMiddleware(
default_model=sub_model,
default_tools=search_tools,
subagents=[critique_sub_agent, research_sub_agent],
default_middleware=[
FilesystemMiddleware(backend=fs_backend_factory),
RuntimeConfigMiddleware(
model_context_name="subagents_model",
enable_model_override=True,
enable_system_prompt_override=False,
enable_tools_override=False,
),
SummaryOffloadMiddleware(
model=sub_model,
trigger=("tokens", 120000),
trim_tokens_to_summarize=None,
immediate_offload_threshold=4000,
summary_offload_threshold=1000,
max_retention_ratio=0.6,
),
PatchToolCallsMiddleware(),
],
general_purpose_agent=True,
),
SummaryOffloadMiddleware(
model=model,
trigger=("tokens", 120000),
trim_tokens_to_summarize=None,
immediate_offload_threshold=4000,
summary_offload_threshold=1000,
max_retention_ratio=0.6,
),
PatchToolCallsMiddleware(),
subagents_middleware,
summary_middleware,
],
checkpointer=await self._get_checkpointer(),
)