ForcePilot/src/agents/deep_agent/graph.py

"""Deep Agent - 基于create_deep_agent的深度分析智能体"""

from deepagents.middleware.filesystem import FilesystemMiddleware
from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware
from deepagents.middleware.subagents import SubAgentMiddleware
from langchain.agents import create_agent
from langchain.agents.middleware import (
    TodoListMiddleware,
    ToolCallLimitMiddleware,
)

from src.agents.common import BaseAgent, load_chat_model
from src.agents.common.backends import create_agent_composite_backend
from src.agents.common.middlewares import RuntimeConfigMiddleware, SummaryOffloadMiddleware, save_attachments_to_fs
from src.agents.common.tools import get_tavily_search
from src.services.mcp_service import get_tools_from_all_servers
from src.utils import logger


from .context import DeepContext


def _create_fs_backend(rt):
    """创建文件存储后端"""
    return create_agent_composite_backend(rt)


def _get_research_sub_agent(search_tools: list) -> dict:
    """Get research sub-agent config with search tools."""
    return {
        "name": "research-agent",
        "description": ("利用搜索工具，用于研究更深入的问题。将调研结果写入到主题研究文件中。"),
        "system_prompt": (
            "你是一位专注的研究员。你的工作是根据用户的问题进行研究。"
            "进行彻底的研究，然后用详细的答案回复用户的问题，只有你的最终答案会被传递给用户。"
            "除了你的最终信息，他们不会知道任何其他事情，所以你的最终报告应该就是你的最终信息！"
            "将调研结果保存到主题研究文件中 /sub_research/xxx.md 中。"
        ),
        "tools": search_tools,
    }


critique_sub_agent = {
    "name": "critique-agent",
    "description": "用于评论最终报告。给这个代理一些关于你希望它如何评论报告的信息。",
    "system_prompt": (
        "你是一位专注的编辑。你的任务是评论一份报告。\n\n"
        "你可以在 `final_report.md` 找到这份报告。\n\n"
        "你可以在 `question.txt` 找到这份报告的问题/主题。\n\n"
        "用户可能会要求评论报告的特定方面。请用详细的评论回复用户，指出报告中可以改进的地方。\n\n"
        "如果有助于你评论报告，你可以使用搜索工具来搜索信息\n\n"
        "不要自己写入 `final_report.md`。\n\n"
        "需要检查的事项：\n"
        "- 检查每个部分的标题是否恰当\n"
        "- 检查报告的写法是否像论文或教科书——它应该是以文本为主，不要只是一个项目符号列表！\n"
        "- 检查报告是否全面。如果任何段落或部分过短，或缺少重要细节，请指出来。\n"
        "- 检查文章是否涵盖了行业的关键领域，确保了整体理解，并且没有遗漏重要部分。\n"
        "- 检查文章是否深入分析了原因、影响和趋势，提供了有价值的见解\n"
        "- 检查文章是否紧扣研究主题并直接回答问题\n"
        "- 检查文章是否结构清晰、语言流畅、易于理解。"
    ),
}


class DeepAgent(BaseAgent):
    name = "深度分析智能体"
    description = "具备规划、深度分析和子智能体协作能力的智能体，可以处理复杂的多步骤任务"
    context_schema = DeepContext
    capabilities = [
        "file_upload",
        "todo",
        "files",
    ]

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.graph = None
        self.checkpointer = None

    async def get_tools(self):
        """返回 Deep Agent 的专用工具"""
        tools = []
        tavily_search = get_tavily_search()
        if tavily_search:
            tools.append(tavily_search)

        # # Assert that search tool is available for DeepAgent
        # assert tools, (
        #     "DeepAgent requires at least one search tool. "
        #     "Please configure TAVILY_API_KEY environment variable to enable web search."
        # )
        if not tools:
            logger.warning("No search tools configured, DeepAgent will work without web search")
            tools = []
        return tools

    async def get_graph(self, **kwargs):
        """构建 Deep Agent 的图"""
        # 获取上下文配置
        context = self.context_schema.from_file(module_name=self.module_name)

        model = load_chat_model(context.model)
        sub_model = load_chat_model(context.subagents_model)
        search_tools = await self.get_tools()
        all_mcp_tools = await get_tools_from_all_servers()
        # 合并搜索工具和 MCP 工具

        # Build subagents with search tools
        research_sub_agent = _get_research_sub_agent(search_tools)

        # 主 Agent 上下文优化：90k tokens 触发压缩（128k context window 的 70%）
        summary_middleware = SummaryOffloadMiddleware(
            model=model,
            trigger=("tokens", 90000),
            trim_tokens_to_summarize=4000,
            summary_offload_threshold=500,
            max_retention_ratio=0.5,
        )

        # 子 Agent 独立的上下文优化：更激进的压缩策略
        sub_summary_middleware = SummaryOffloadMiddleware(
            model=sub_model,
            trigger=("tokens", 50000),
            trim_tokens_to_summarize=2000,
            summary_offload_threshold=300,
            max_retention_ratio=0.4,
        )

        subagents_middleware = SubAgentMiddleware(
            default_model=sub_model,
            default_tools=search_tools,
            subagents=[critique_sub_agent, research_sub_agent],
            default_middleware=[
                RuntimeConfigMiddleware(
                    model_context_name="subagents_model",
                    enable_model_override=True,
                    enable_system_prompt_override=False,
                    enable_tools_override=False,
                ),
                PatchToolCallsMiddleware(),
                sub_summary_middleware,
                # 子 Agent 搜索工具限制：tavily_search 最多 8 次
                ToolCallLimitMiddleware(
                    tool_name="tavily_search",
                    run_limit=8,
                    exit_behavior="continue",
                ),
            ],
            general_purpose_agent=True,
        )

        # 使用 create_deep_agent 创建深度智能体
        graph = create_agent(
            model=model,
            system_prompt=context.system_prompt,
            middleware=[
                FilesystemMiddleware(backend=_create_fs_backend),  # 文件系统后端
                RuntimeConfigMiddleware(extra_tools=all_mcp_tools),
                save_attachments_to_fs,  # 附件注入提示词
                TodoListMiddleware(),
                PatchToolCallsMiddleware(),
                subagents_middleware,
                summary_middleware,
                # 工具调用限制：tavily_search 总调用最多 20 次
                ToolCallLimitMiddleware(
                    tool_name="tavily_search",
                    thread_limit=20,
                    exit_behavior="continue",
                ),
                # 总工具调用轮次限制：防止单次运行无限循环
                ToolCallLimitMiddleware(
                    run_limit=50,
                    exit_behavior="end",
                ),
            ],
            checkpointer=await self._get_checkpointer(),
        )

        return graph