feat: 优化 summary 组件

This commit is contained in:
Wenjie Zhang 2026-04-02 09:05:56 +08:00
parent 4ee15e377f
commit e05eacfae3
2 changed files with 9 additions and 13 deletions

View File

@ -7,6 +7,7 @@
from __future__ import annotations
from pathlib import Path
import uuid
from collections.abc import Callable, Iterable, Mapping
from functools import partial
@ -31,6 +32,8 @@ from langchain_core.messages.utils import (
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.runtime import Runtime
from yuxi.utils.paths import OUTPUTS_DIR_NAME
TokenCounter = Callable[[Iterable[MessageLikeRepresentation]], int]
DEFAULT_SUMMARY_PROMPT = """<role>
@ -74,9 +77,7 @@ Messages to summarize:
</messages>"""
_DEFAULT_MESSAGES_TO_KEEP = 20
_DEFAULT_TRIM_TOKEN_LIMIT = 4000
_DEFAULT_FALLBACK_MESSAGE_COUNT = 15
_DEFAULT_OFFLOAD_THRESHOLD = 1000 # Token 数阈值,超过此值则卸载到文件系统
_OFFLOAD_DIR = "/summary_offload" # 虚拟文件系统路径
ContextFraction = tuple[Literal["fraction"], float]
@ -143,7 +144,7 @@ def _offload_tool_result(msg: ToolMessage, threshold: int, token_counter: TokenC
# 生成文件路径 (工具名称-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}"
file_path = (Path(OUTPUTS_DIR_NAME) / f"{_OFFLOAD_DIR}/{safe_name}-{message_id}").as_posix()
# 构建文件头部信息
header_lines = [
@ -226,7 +227,7 @@ class SummaryOffloadMiddleware(AgentMiddleware):
keep: ContextSize = ("messages", _DEFAULT_MESSAGES_TO_KEEP),
token_counter: TokenCounter = count_tokens_approximately,
summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
trim_tokens_to_summarize: int | None = _DEFAULT_TRIM_TOKEN_LIMIT,
trim_tokens_to_summarize: int | None = 4000,
# 工具结果卸载参数
summary_offload_threshold: int = 1000,
max_retention_ratio: float = 0.6,
@ -240,18 +241,12 @@ class SummaryOffloadMiddleware(AgentMiddleware):
keep: 摘要后保留的消息数量/ token 策略 (作为 fallback)
token_counter: token 计数函数
summary_prompt: 生成摘要的提示词模板
trim_tokens_to_summarize: 准备摘要消息时的最大 token
summary_offload_threshold: Summary 卸载阈值token 默认 1000
trim_tokens_to_summarize: Summary 无损保留的消息
summary_offload_threshold: Summary 工具调用结果超过此 token 数阈值则卸载到文件系统
max_retention_ratio: 触发 Summary 如果不超过此比例相对于 trigger则不删除消息默认 0.6
"""
super().__init__()
# Handle renamed argument for backward compatibility if needed,
# but since we are refactoring, we map deprecated 'result_offload_threshold'
# to 'summary_offload_threshold' if present in kwargs.
if "result_offload_threshold" in deprecated_kwargs:
summary_offload_threshold = deprecated_kwargs.pop("result_offload_threshold")
if isinstance(model, str):
model = init_chat_model(model)

View File

@ -9,6 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from yuxi.repositories.subagent_repository import SubAgentRepository
from yuxi.storage.postgres.manager import pg_manager
from yuxi.utils import logger
from yuxi.utils.paths import OUTPUTS_DIR_NAME
# SubAgent specs cache for get_subagent_specs
_subagent_specs_cache: list[dict[str, Any]] | None = None
@ -34,7 +35,7 @@ _DEFAULT_SUBAGENTS = [
"你是一位专注的研究员。你的工作是根据用户的问题进行研究。"
"进行彻底的研究,然后用详细的答案回复用户的问题,只有你的最终答案会被传递给用户。"
"除了你的最终信息,他们不会知道任何其他事情,所以你的最终报告应该就是你的最终信息!"
"将调研结果保存到主题研究文件中 sub_research/xxx.md 中。"
f"将调研结果保存到主题研究文件中 {OUTPUTS_DIR_NAME}/sub_research/xxx.md 中。"
),
"tools": ["tavily_search"],
"is_builtin": True,