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 __future__ import annotations
from pathlib import Path
import uuid import uuid
from collections.abc import Callable, Iterable, Mapping from collections.abc import Callable, Iterable, Mapping
from functools import partial from functools import partial
@ -31,6 +32,8 @@ from langchain_core.messages.utils import (
from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.runtime import Runtime from langgraph.runtime import Runtime
from yuxi.utils.paths import OUTPUTS_DIR_NAME
TokenCounter = Callable[[Iterable[MessageLikeRepresentation]], int] TokenCounter = Callable[[Iterable[MessageLikeRepresentation]], int]
DEFAULT_SUMMARY_PROMPT = """<role> DEFAULT_SUMMARY_PROMPT = """<role>
@ -74,9 +77,7 @@ Messages to summarize:
</messages>""" </messages>"""
_DEFAULT_MESSAGES_TO_KEEP = 20 _DEFAULT_MESSAGES_TO_KEEP = 20
_DEFAULT_TRIM_TOKEN_LIMIT = 4000
_DEFAULT_FALLBACK_MESSAGE_COUNT = 15 _DEFAULT_FALLBACK_MESSAGE_COUNT = 15
_DEFAULT_OFFLOAD_THRESHOLD = 1000 # Token 数阈值,超过此值则卸载到文件系统
_OFFLOAD_DIR = "/summary_offload" # 虚拟文件系统路径 _OFFLOAD_DIR = "/summary_offload" # 虚拟文件系统路径
ContextFraction = tuple[Literal["fraction"], float] ContextFraction = tuple[Literal["fraction"], float]
@ -143,7 +144,7 @@ def _offload_tool_result(msg: ToolMessage, threshold: int, token_counter: TokenC
# 生成文件路径 (工具名称-xxx) # 生成文件路径 (工具名称-xxx)
message_id = msg.id or str(uuid.uuid4())[:8] 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) 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 = [ header_lines = [
@ -226,7 +227,7 @@ class SummaryOffloadMiddleware(AgentMiddleware):
keep: ContextSize = ("messages", _DEFAULT_MESSAGES_TO_KEEP), keep: ContextSize = ("messages", _DEFAULT_MESSAGES_TO_KEEP),
token_counter: TokenCounter = count_tokens_approximately, token_counter: TokenCounter = count_tokens_approximately,
summary_prompt: str = DEFAULT_SUMMARY_PROMPT, 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, summary_offload_threshold: int = 1000,
max_retention_ratio: float = 0.6, max_retention_ratio: float = 0.6,
@ -240,18 +241,12 @@ class SummaryOffloadMiddleware(AgentMiddleware):
keep: 摘要后保留的消息数量/ token 策略 (作为 fallback) keep: 摘要后保留的消息数量/ token 策略 (作为 fallback)
token_counter: token 计数函数 token_counter: token 计数函数
summary_prompt: 生成摘要的提示词模板 summary_prompt: 生成摘要的提示词模板
trim_tokens_to_summarize: 准备摘要消息时的最大 token trim_tokens_to_summarize: Summary 无损保留的消息
summary_offload_threshold: Summary 卸载阈值token 默认 1000 summary_offload_threshold: Summary 工具调用结果超过此 token 数阈值则卸载到文件系统
max_retention_ratio: 触发 Summary 如果不超过此比例相对于 trigger则不删除消息默认 0.6 max_retention_ratio: 触发 Summary 如果不超过此比例相对于 trigger则不删除消息默认 0.6
""" """
super().__init__() 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): if isinstance(model, str):
model = init_chat_model(model) 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.repositories.subagent_repository import SubAgentRepository
from yuxi.storage.postgres.manager import pg_manager from yuxi.storage.postgres.manager import pg_manager
from yuxi.utils import logger from yuxi.utils import logger
from yuxi.utils.paths import OUTPUTS_DIR_NAME
# SubAgent specs cache for get_subagent_specs # SubAgent specs cache for get_subagent_specs
_subagent_specs_cache: list[dict[str, Any]] | None = None _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"], "tools": ["tavily_search"],
"is_builtin": True, "is_builtin": True,