[
{
"path": ".gitignore",
"content": "# Python\n__pycache__/\n*.py[cod]\n*$py.class\n*.so\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# Virtual Environment\nvenv/\nenv/\nENV/\n\n# IDE\n.idea/\n.vscode/\n*.swp\n*.swo\n\n# OS specific\n.DS_Store\nThumbs.db\n\n# LangSmith\n.langchain.db\n.langsmith/\n\n# Logs\n*.log\n\n# Env\n.env\n\n# output\nexampels/logs/\nexampels/output/\nexamples/output/sandbox_test"
},
{
"path": "README.md",
"content": "# Mentis - Agent Development Kit\n\n[](https://www.python.org/downloads/)\n[](https://opensource.org/licenses/MIT) \n## 概述 (Overview)\n\nMentis 是一个基于 LangGraph 构建的、可扩展的多 Agent ADK(Agent Development Kit)。它的核心是一个**状态驱动的规划型 Supervisor Agent**，负责理解用户复杂请求、制定执行计划，并智能地协调一组具有不同专业能力的子 Agent (Specialist Agents) 来共同完成任务。\n\n此框架旨在实现复杂任务的自动化处理，通过 Agent 间的协作提供比单一 Agent 更强大、更灵活的问题解决能力。\n\n## 核心特性 (Core Features)\n\n* **Multi-Agent 架构**: 采用中心化的 Supervisor 协调多个专门的子 Agent (如 Research, Coder, Reporter, Designer, Data Analyst)。\n* **State-Based Planning**: 引入独立的 `Planner` 节点负责初始规划，`Supervisor` 专注于基于计划状态的执行和调度，`Evaluator` 节点负责评估子 Agent 结果并更新状态。计划状态通过 LangGraph 持久化（需配置 Checkpointer）。\n* **模块化 Agent 设计**: 基于 `BaseAgent` 和 `ReactAgent` 构建，易于添加或修改具有不同能力的子 Agent。\n* **工具注册与管理**: 通过 `core/tools/registry.py` 实现工具的集中注册、分类和动态加载。\n* **可配置 LLM**: 支持通过 `LLMManager` (或环境变量) 配置和切换不同的 LLM Provider (OpenAI, DeepSeek, XAI Grok via compatible endpoint) 和模型。\n* **持久化支持**: 基于 LangGraph 的 Checkpointer 机制，可以实现对话状态和计划的持久化。\n* **清晰的执行流程**: Planner -> Supervisor -> (Handoff -> Agent -> Evaluator -> Supervisor 循环) -> 最终输出/Reporter。\n* **A2A 协议支持**: 实现了 Google 的 Agent-to-Agent (A2A) 协议，使 Mentis Agents 能够与其他支持 A2A 协议的系统进行互操作。\n\n## 架构概览 (Architecture Overview)\n\n1. **用户请求 (Input)**: 用户通过入口点 (`main.py` 或 API) 提交任务请求。\n2. **规划节点 (Planner Node)**: 分析请求，生成一个包含任务步骤、建议 Agent 的初始计划 (`Plan`)，并更新到图状态 (`PlanningAgentState`)。\n3. **主管节点 (Supervisor Node)**: 接收带有计划的状态，根据计划状态和消息历史决定下一步行动：\n * 启动新任务 (标记 'in_progress')。\n * 委派 'in_progress' 的任务给合适的子 Agent (通过 Handoff 工具)。\n * 等待子 Agent 完成。\n * 判断计划是否最终完成。\n * 决定最终输出方式（自己总结或调用 Reporter）。\n4. **切换执行器 (Handoff Executor)**: 处理 Supervisor 发出的 `transfer_to_` 工具调用，并将控制权和状态传递给目标子 Agent。\n5. **子 Agent 节点 (Specialist Agent Nodes)**: 继承自 `ReactAgent` 或 `BaseAgent`，执行具体的任务（研究、编码、生成报告/图像、数据分析），可能调用其自身的工具。\n6. **评估节点 (Evaluate Result Node)**: 接收子 Agent 的执行结果，进行确定性评估（成功/失败），更新对应任务的状态和 Plan 的整体状态。\n7. **循环与结束**: 流程在 Evaluator -> Supervisor 之间循环，直到 Supervisor 判断 Plan 完成，然后路由到 `END` 或 `ReporterAgent`。\n\n## 快速开始 (Getting Started)\n\n### 1. 环境设置 (Prerequisites)\n\n* Python 3.11+\n* 使用 `pip` 或 `uv` 等工具管理依赖。\n\n### 2. 安装依赖 (Installation)\n\n在项目根目录运行：\n建议使用 uv 管理\n```bash\nuv venv\nsource .venv/bin/activate\nuv sync\n```\n\n```bash\n# pip install -r requirements.txt \n# 或者 uv pip install -r requirements.txt\n```\n(requirements.txt 我没维护，请确保 `requirements.txt` 文件包含了所有必要的库，如 `langchain`, `langgraph`, `langchain-openai`, `e2b` (如果使用 E2B), `replicate` (如果使用 Replicate), `tavily-python`, `exa-py`, `python-dotenv`, `anyio`, `tiktoken` 等)。\n\n### 3. 配置环境 (Configuration)\n\n* 复制 `.env.example` 文件为 `.env`。\n* 在 `.env` 文件中填入您所需的 API Keys/Tokens：\n * `OPENAI_API_KEY` (如果使用 OpenAI 模型)\n * `DEEPSEEK_API_KEY` (如果使用 DeepSeek 模型)\n * `XAI_API_KEY` (如果使用 XAI Grok，并确认 Base URL)\n * `REPLICATE_API_TOKEN` (如果使用 Replicate 工具)\n * `E2B_API_KEY` (如果使用 E2B Code Interpreter，推荐！)\n * `TAVILY_API_KEY` (如果使用 Tavily 搜索，推荐！)\n * `EXA_API_KEY` (如果使用 Exa 搜索)\n * `LANGCHAIN_TRACING_V2=\"true\"` (强烈推荐，用于 LangSmith 调试)\n * `LANGCHAIN_API_KEY=\"ls_...\"` (您的 LangSmith Key)\n * `LANGCHAIN_PROJECT=\"Your_Project_Name\"` (您在 LangSmith 上的项目名)\n* **LLM 配置**:\n * 如果您使用了 `LLMManager`（如示例所示），请检查并配置其读取的模型配置文件（例如 `config/models.yaml`，路径可能不同）。\n * 如果您在 `tools.py` 中直接根据环境变量初始化 LLM，请确保设置了对应的环境变量，如 `LLM_PROVIDER`, `LLM_MODEL_NAME`, `LLM_BASE_URL` (用于兼容 API)。\n* **工具配置**: 确保 `core/tools/__init__.py` 或 `registry.py` 中的工具预注册逻辑能够正确找到并初始化您需要的工具。\n\n### 4. 运行示例 (Running Examples)\n\n项目包含示例脚本以演示框架的使用：\n```bash\n# 从项目根目录 (mentis/) 运行\npython examples/state_based_supervisor_examples/03_multi_agents.py \n```\n脚本会提示您输入初始请求。您可以进行简单尝试：\n\n* `\"What is the capital of France?\"` (简单测试)\n* `\"Write a short, four-line poem about spring.\"` (测试 Reporter)\n* `\"Generate an image of a cat wearing a top hat, oil painting style.\"` (测试 Designer)\n* `\"Write a Python function to calculate factorial and run it for 5.\"` (测试 Coder)\n\n## 项目结构 (Project Structure)\n\n```\nmentis/\n├── api/ # (可选) API 服务相关代码\n├── core/ # 核心框架代码\n│ ├── a2a/ # A2A 协议的客户端和服务器实现\n│ ├── agents/ # Agent 定义 (base, react, supervisor, sub-agents)\n│ │ ├── base/\n│ │ ├── state_based_supervisor/ # Supervisor 相关 (graph, node, planner, evaluator)\n│ │ ├── sub_agents/ # 具体子 Agent 实现 (research, coder, etc.)\n│ │ └── sb_supervisor_agent.py # SupervisorAgent 类定义\n│ ├── llm/ # (可选) LLM 管理或配置\n│ ├── tools/ # 工具定义和注册表 (registry, e2b, replicate, etc.)\n│ └── utils/ # 通用辅助函数\n├── examples/ # 示例和测试脚本\n│ └── state_based_supervisor_examples/\n│ └── 03_multi_agents.py # 我们使用的测试脚本\n├── super_agents/ # 独立功能型 Agent 实现\n│ └── deep_research/ # DeepResearch Agent 实现\n│ └── a2a_adapter/ # DeepResearch 的 A2A 协议适配器\n├── web/ # (可选) Web 客户端代码\n├── web_for_a2a/ # 基于 A2A 协议的 Web 界面\n├── .env.example # 环境变量示例\n├── requirements.txt # Python 依赖\n└── README.md # 本文件\n```\n\n## Super Agents (独立功能型 Agent)\n\n除了由 Supervisor 协调的、专注于单一技能的 Specialist Agents (如 Coder, Researcher) 之外，本框架也支持构建和集成更复杂的 **\"Super Agents\"**。\n\nSuper Agent 可以理解为一个**独立的、具有端到端能力、能够完成一个相对完整且复杂任务的 Agent 图**。它可以包含自己的规划、执行、甚至内部协调逻辑。\n\n这些 Super Agents 既可以**独立运行**以完成特定的大型任务，也可以被更高层的协调者（例如我们的 Supervisor Agent）**视为一种强大的“能力”或“工具”**来调用，以处理其复杂计划中的某个步骤。\n\n### DeepResearch Agent (第一个实例)\n\n\nhttps://github.com/user-attachments/assets/2a685709-5be0-43a3-9e2d-934ef5fa3315\n\n\n`DeepResearch Agent` 是我们在此框架理念下实现的第一个 Super Agent 实例（其早期版本是我们开发此 Multi-Agent 框架的基础）。\n\n* **核心功能**: 旨在针对用户给定的**任意主题**，自动化地执行一个**深度研究**流程。\n* **内部工作流**: 它包含自己的一套完整的内部步骤，大致如下：\n 1. **研究规划 (Plan Research)**: 分析主题，生成初步的搜索查询和分析点。\n 2. **多源搜索 (Multi-Source Search)**: 调用网页搜索 (Tavily)、学术搜索 (Exa) 等工具获取信息。\n 3. **(可选) 分析执行 (Perform Analysis)**: 对搜索结果进行初步分析（如情感、SWOT 等）。\n 4. **差距分析 (Gap Analysis)**: 评估已有信息，识别知识空白和局限性。\n 5. **(可选) 补充搜索 (Gap Filling)**: 针对知识空白进行额外的、更具针对性的搜索。\n 6. **最终综合 (Final Synthesis)**: 整合所有信息，提炼关键发现和不确定性。\n 7. **报告生成 (Report Generation)**: 将综合结果和上下文信息，撰写成一份详细的、带引用的 Markdown 研究报告。\n* **当前状态**: 该 Agent 的核心逻辑和节点已基本实现，并且现在支持 A2A 协议和专用 Web 界面。\n\n#### A2A 协议支持\n\n我们为 DeepResearch Agent 实现了完整的 A2A 协议适配器，使其能够：\n\n* 作为标准的 A2A 服务被发现和调用\n* 通过 `tasks/send` 和 `tasks/sendSubscribe` 端点接收研究任务\n* 提供实时的流式研究进度更新\n* 返回结构化的研究结果\n* 支持推送通知机制\n\n这使得 DeepResearch Agent 可以轻松地与其他支持 A2A 协议的系统（如 Google Assistant）集成，或者被自定义的前端应用调用。\n\n#### 专用 Web 界面\n\n\nhttps://github.com/user-attachments/assets/640365c7-839b-4765-b9ac-ee0ac961ceb8\n\n\n我们还开发了一个基于 Next.js 的现代 Web 界面，专门用于与 DeepResearch A2A 服务交互：\n\n* 提供直观的用户界面，用于输入研究主题和启动研究任务\n* 实时显示研究进度和中间更新（通过 Server-Sent Events）\n* 美观地展示最终生成的研究报告\n* 演示了如何在前端应用中使用浏览器原生 API 处理 A2A 流式响应\n\n**如何体验 DeepResearch Agent:**\n\n1. **独立运行模式**:\n * 确保环境配置: 确认您的 `.env` 文件中包含了所需的所有 API Keys（例如 `OPENAI_API_KEY`/`DEEPSEEK_API_KEY`, `TAVILY_API_KEY`, `EXA_API_KEY`）。\n * 运行脚本: 在项目根目录执行：\n ```bash\n python super_agents/deep_research/main.py\n ```\n * 输入主题并查看结果: 生成的报告通常会保存在 `output/` 文件夹中。\n\n2. **A2A 服务模式**:\n * 启动 A2A 服务器:\n ```bash\n cd super_agents/deep_research/a2a_adapter\n python run_server.py\n ```\n * 服务器将在默认端口（通常是 8000）启动，并提供符合 A2A 规范的 API 端点。\n\n3. **Web 界面模式**:\n * 确保 A2A 服务器正在运行\n * 启动 Web 界面:\n ```bash\n cd web_for_a2a\n npm install\n npm run dev\n ```\n * 在浏览器中访问 http://localhost:3000/deepresearch 使用图形界面与 DeepResearch Agent 交互。\n\n## 未来工作 (Future Work / Contributing)\n\n* 完善子 Agent 的工具集和 Prompt。\n* 增强 Evaluator Node 的评估逻辑。\n* 添加更复杂的任务依赖处理。\n* 优化长对话历史的管理。\n* 集成持久化 Checkpointer (如 SQLite, Redis)。\n* 欢迎提出 Issue 或 Pull Request！\n* 有问题也可以添加我的微信 brown🩷cony999\n\n\n## 许可证 (License)\n\nThis project is licensed under the MIT License - see the LICENSE file for details.\n"
},
{
"path": "__init__.py",
"content": "# Project package initialization\n"
},
{
"path": "api/__init__.py",
"content": ""
},
{
"path": "api/agent/__init__.py",
"content": ""
},
{
"path": "api/agent/loader.py",
"content": "# Agent Loader Module\n# This module is responsible for loading agents from the web_agents directory\n\nimport importlib\nimport os\nimport sys\nfrom typing import Dict, Optional, Any, List\nfrom langgraph.graph import StateGraph\nfrom langgraph.graph.graph import CompiledGraph # Add this import\n\n# Try to import deep_research_app\ntry:\n # Adjust this import path based on your project structure\n from super_agents.deep_research.reason_graph.graph import web_app as deep_research_app\nexcept ImportError:\n print(\"Warning: Failed to import deep_research_app. DeepResearchAgent will be unavailable.\")\n deep_research_app = None\n\n# Add examples directory to Python path to allow importing web_agents\nexamples_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'examples')\nif examples_path not in sys.path:\n sys.path.append(examples_path)\n\n\ndef list_available_agents() -> Dict[str, str]:\n \"\"\"List all available agents in the web_agents directory\n \n Returns:\n Dict[str, str]: A dictionary mapping agent names to their descriptions\n \"\"\"\n agents = {}\n web_agents_dir = os.path.join(examples_path, 'web_agents')\n \n # Check if web_agents directory exists\n if not os.path.exists(web_agents_dir) or not os.path.isdir(web_agents_dir):\n pass # Continue with empty agents dict\n else:\n # Iterate through subdirectories in web_agents\n for item in os.listdir(web_agents_dir):\n agent_dir = os.path.join(web_agents_dir, item)\n \n # Skip non-directories and special directories\n if not os.path.isdir(agent_dir) or item.startswith('__') or item.startswith('.'):\n continue\n \n # Check if the directory contains an __init__.py file with get_graph function\n init_file = os.path.join(agent_dir, '__init__.py')\n if os.path.exists(init_file):\n # Try to get description from README.md\n readme_file = os.path.join(agent_dir, 'README.md')\n description = item # Default description is the directory name\n \n if os.path.exists(readme_file):\n try:\n with open(readme_file, 'r', encoding='utf-8') as f:\n first_line = f.readline().strip()\n if first_line.startswith('# '):\n description = first_line[2:]\n except Exception:\n pass\n \n agents[item] = description\n \n # Add deep_research to available agents if it's imported successfully\n if deep_research_app is not None:\n agents[\"deep_research\"] = \"Deep Research Agent for in-depth topic exploration\"\n \n return agents\n\n\ndef load_agent(agent_name: str) -> Optional[CompiledGraph]:\n \"\"\"Load an agent from the web_agents directory or special agents\n \n Args:\n agent_name (str): The name of the agent to load\n \n Returns:\n Optional[CompiledGraph]: The compiled graph for the agent, or None if the agent could not be loaded\n \"\"\"\n # Special case for deep_research agent\n if agent_name == \"deep_research\":\n if deep_research_app:\n return deep_research_app\n else:\n print(f\"ERROR: DeepResearchAgent requested but not available.\")\n return None\n \n # Standard agents from web_agents directory\n try:\n # Import the agent module\n module = importlib.import_module(f'web_agents.{agent_name}')\n \n # Check if the module has a get_graph function\n if hasattr(module, 'get_graph'):\n # Call the get_graph function to get the compiled graph\n return module.get_graph()\n else:\n print(f\"Error: Agent '{agent_name}' does not have a get_graph function\")\n return None\n except ImportError as e:\n print(f\"Error importing agent '{agent_name}': {e}\")\n return None\n except Exception as e:\n print(f\"Error loading agent '{agent_name}': {e}\")\n return None\n\n\n# Default agent to use if none is specified\nDEFAULT_AGENT = 'research_assistant'\n# DEFAULT_AGENT = 'weather_agent'\n\n\ndef get_default_agent() -> Optional[CompiledGraph]:\n \"\"\"Get the default agent\n \n Returns:\n Optional[CompiledGraph]: The compiled graph for the default agent, or None if it could not be loaded\n \"\"\"\n return load_agent(DEFAULT_AGENT)"
},
{
"path": "api/server.py",
"content": "import uvicorn\nfrom langgraph.types import Command, Interrupt\nfrom fastapi import FastAPI, Request, HTTPException, Query\nfrom fastapi.middleware.cors import CORSMiddleware\nfrom sse_starlette.sse import EventSourceResponse\nfrom typing import AsyncGenerator, Dict, Optional, Union, Any\nfrom api.utils import message_chunk_event, interrupt_event, custom_event, checkpoint_event, format_state_snapshot, stream_update_event\nimport asyncio\nimport traceback\nimport json\nfrom langchain_core.messages import HumanMessage\nfrom langchain_core.runnables import RunnableConfig\n\n# Import the agent loader\nfrom api.agent.loader import load_agent, list_available_agents, get_default_agent\n\n# Load the default agent\ngraph = get_default_agent()\n\n# Track active connections\nactive_connections: Dict[str, asyncio.Event] = {}\n\napp = FastAPI(\n title=\"LangGraph API\",\n description=\"API for LangGraph interactions\",\n version=\"0.1.0\"\n)\n\n# Configure CORS\napp.add_middleware(\n CORSMiddleware,\n allow_origins=[\"*\"], # In production, replace with specific origins\n allow_credentials=True,\n allow_methods=[\"*\"],\n allow_headers=[\"*\"],\n)\n\n\n@app.get(\"/agents\")\nasync def list_agents():\n \"\"\"Endpoint returning a list of available agents.\"\"\"\n return list_available_agents()\n\n\n@app.get(\"/state\")\nasync def state(thread_id: str | None = None, agent: Optional[str] = Query(None)):\n \"\"\"Endpoint returning current graph state.\"\"\"\n if not thread_id:\n raise HTTPException(status_code=400, detail=\"thread_id is required\")\n \n # Load the specified agent if provided\n current_graph = load_agent(agent) if agent else graph\n if not current_graph:\n raise HTTPException(status_code=404, detail=f\"Agent '{agent}' not found\")\n\n config: RunnableConfig = {\"configurable\": {\"thread_id\": thread_id}}\n\n state = await current_graph.aget_state(config)\n return format_state_snapshot(state)\n\n\n@app.get(\"/history\")\nasync def history(thread_id: str | None = None, agent: Optional[str] = Query(None)):\n \"\"\"Endpoint returning complete state history. Used for restoring graph.\"\"\"\n if not thread_id:\n raise HTTPException(status_code=400, detail=\"thread_id is required\")\n \n # Load the specified agent if provided\n current_graph = load_agent(agent) if agent else graph\n if not current_graph:\n raise HTTPException(status_code=404, detail=f\"Agent '{agent}' not found\")\n\n config: RunnableConfig = {\"configurable\": {\"thread_id\": thread_id}}\n\n records = []\n async for state in current_graph.aget_state_history(config):\n records.append(format_state_snapshot(state))\n return records\n\n\n@app.post(\"/agent/stop\")\nasync def stop_agent(request: Request):\n \"\"\"Endpoint for stopping the running agent.\"\"\"\n body = await request.json()\n thread_id = body.get(\"thread_id\")\n if not thread_id:\n raise HTTPException(status_code=400, detail=\"thread_id is required\")\n\n if thread_id in active_connections:\n active_connections[thread_id].set()\n return {\"status\": \"stopped\", \"thread_id\": thread_id}\n raise HTTPException(status_code=404, detail=\"Thread is not running\")\n\n\n@app.post(\"/agent\")\nasync def agent(request: Request):\n \"\"\"Endpoint for running the agent.\"\"\"\n body = await request.json()\n\n request_type = body.get(\"type\")\n if not request_type:\n raise HTTPException(status_code=400, detail=\"type is required\")\n\n thread_id = body.get(\"thread_id\")\n if not thread_id:\n raise HTTPException(status_code=400, detail=\"thread_id is required\")\n\n # Get the agent name if provided\n agent_name = body.get(\"agent\")\n \n # Load the specified agent if provided\n current_graph = load_agent(agent_name) if agent_name else graph\n if not current_graph:\n raise HTTPException(status_code=404, detail=f\"Agent '{agent_name or 'default'}' not found\")\n\n stop_event = asyncio.Event()\n active_connections[thread_id] = stop_event\n\n config: RunnableConfig = {\"configurable\": {\"thread_id\": thread_id}}\n initial_graph_state: Dict[str, Any] = {}\n input_for_astream: Optional[Union[Dict, Command]] = None # input for astream\n\n # Get initial state or messages from frontend\n initial_state_input = body.get(\"state\", {\"messages\": []})\n if not isinstance(initial_state_input, dict):\n raise HTTPException(status_code=400, detail=\"state must be a dictionary\")\n\n if agent_name == \"deep_research\":\n # --- Prepare state for DeepResearch Agent ---\n print(\"Preparing state for DeepResearchAgent...\")\n # Extract topic from the first message in state['messages']\n first_message_content = \"\"\n try:\n # Ensure initial_state_input['messages'] is a list and not empty\n if isinstance(initial_state_input.get('messages'), list) and initial_state_input['messages']:\n # Assume the first message's content is the topic\n first_message_content = initial_state_input['messages'][0]['content']\n else:\n # Try to get topic from other fields in state (alternative)\n first_message_content = initial_state_input.get('topic', '')\n \n except Exception as e:\n print(f\"Warning: Could not extract topic from initial state input: {e}\")\n\n if not first_message_content or not isinstance(first_message_content, str):\n raise HTTPException(status_code=400, detail=\"A valid 'topic' string is required for deep_research agent, expected in state.messages[0].content or state.topic\")\n\n # Build the ResearchState needed by DeepResearch Agent (at least topic and depth)\n initial_graph_state = {\n \"topic\": first_message_content,\n \"depth\": initial_state_input.get(\"depth\", \"advanced\"), # Optional: allow frontend to specify depth\n \"messages\": [], # DeepResearch manages its own message history\n \"stream_updates\": [], # Initialize stream_updates\n # Initialize other ResearchState fields to None or default values\n \"plan\": None, \"research_plan\": None, \"search_results\": [], \n \"gap_analysis\": None, \"final_synthesis\": None, \n \"final_report_markdown\": None,\n }\n print(f\"Initial ResearchState: {{'topic': '{initial_graph_state['topic']}', 'depth': '{initial_graph_state['depth']}', ...}}\")\n \n # DeepResearch Agent's astream input is the complete initial state\n if request_type == \"run\":\n input_for_astream = initial_graph_state\n elif request_type == \"resume\":\n # DeepResearch Agent might not support or need different resume approach\n print(\"Warning: 'resume' might not be fully supported for DeepResearchAgent yet.\")\n # Assume resume Command can be understood by the graph\n input_for_astream = Command(resume=body.get(\"resume\"))\n config[\"configurable\"][\"checkpoint_id\"] = body.get(\"resume\") # Resume usually needs checkpoint ID\n else: # Fork, Replay typically only need config\n config_from_request = body.get(\"config\")\n if not config_from_request:\n raise HTTPException(status_code=400, detail=\"config is required for fork/replay\")\n config = config_from_request # Use complete config provided in the request\n input_for_astream = None\n\n else: # For Supervisor or other Agents (assume using PlanningAgentState)\n print(\"Preparing state for Supervisor/Other Agent...\")\n # --- Prepare PlanningAgentState ---\n # Ensure messages list contains correct BaseMessage objects (or let BaseAgent preprocess)\n initial_messages = initial_state_input.get(\"messages\", [])\n\n initial_graph_state = {\n \"messages\": initial_messages,\n \"plan\": None, # Planner node will create it\n \"error\": None\n # Add other fields needed by PlanningAgentState and set to None or default values\n }\n \n # --- Set astream input (logic similar to before) ---\n if request_type == \"run\":\n # For PlanningAgentState, initial input typically only contains messages\n input_for_astream = {\"messages\": initial_messages}\n elif request_type == \"resume\":\n resume_val = body.get(\"resume\")\n if not resume_val:\n raise HTTPException(status_code=400, detail=\"resume value is required\")\n input_for_astream = Command(resume=resume_val)\n # Ensure config includes checkpoint_id for resuming\n if \"configurable\" not in config:\n config[\"configurable\"] = {}\n config[\"configurable\"][\"checkpoint_id\"] = resume_val \n elif request_type == \"fork\": \n config_from_request = body.get(\"config\")\n if not config_from_request:\n raise HTTPException(status_code=400, detail=\"config is required for fork\")\n config = config_from_request # Fork uses complete config provided\n # Fork typically starts from specified checkpoint, no extra state dict input needed\n input_for_astream = None \n elif request_type == \"replay\": \n config_from_request = body.get(\"config\")\n if not config_from_request:\n raise HTTPException(status_code=400, detail=\"config is required for replay\")\n config = config_from_request\n input_for_astream = None\n else:\n raise HTTPException(status_code=400, detail=\"invalid request type\")\n \n # Ensure config always has thread_id (important for all agents)\n if \"configurable\" not in config:\n config[\"configurable\"] = {}\n config[\"configurable\"][\"thread_id\"] = thread_id\n\n # --- State and Input preparation complete ---\n\n async def generate_events() -> AsyncGenerator[dict, None]:\n try:\n # 设置recursion_limit为100，解决深度研究时的递归限制问题\n if agent_name == \"deep_research\" and \"recursion_limit\" not in config:\n config[\"recursion_limit\"] = 100\n \n async for chunk in current_graph.astream(\n input_for_astream, # Use prepared input\n config, # Use prepared config\n stream_mode=[\"debug\", \"messages\", \"updates\", \"custom\"],\n ):\n if stop_event.is_set():\n break\n\n chunk_type, chunk_data = chunk\n\n if chunk_type == \"debug\":\n # type can be checkpoint, task, task_result\n if isinstance(chunk_data, dict) and \"type\" in chunk_data:\n debug_type = chunk_data[\"type\"]\n if debug_type == \"checkpoint\":\n yield checkpoint_event(chunk_data)\n elif debug_type == \"task_result\":\n interrupts = chunk_data[\"payload\"].get(\n \"interrupts\", [])\n if interrupts and len(interrupts) > 0:\n yield interrupt_event(interrupts)\n elif chunk_type == \"messages\":\n # 确保chunk_data是一个包含至少两个元素的列表/元组，并且第二个元素是一个包含langgraph_node的字典\n if isinstance(chunk_data, (list, tuple)) and len(chunk_data) > 1 and isinstance(chunk_data[1], dict) and \"langgraph_node\" in chunk_data[1]:\n yield message_chunk_event(chunk_data[1][\"langgraph_node\"], chunk_data[0])\n else:\n print(f\"Warning: Unexpected messages chunk_data format: {chunk_data}\")\n # 尝试使用安全的默认值\n node_name = chunk_data[1].get(\"langgraph_node\", \"unknown\") if isinstance(chunk_data, (list, tuple)) and len(chunk_data) > 1 and isinstance(chunk_data[1], dict) else \"unknown\"\n message = chunk_data[0] if isinstance(chunk_data, (list, tuple)) and len(chunk_data) > 0 else None\n if message is not None:\n yield message_chunk_event(node_name, message)\n elif chunk_type == \"custom\":\n # Check if this is a StreamUpdate\n if isinstance(chunk_data, dict) and all(k in chunk_data for k in ['id', 'type', 'status', 'title']):\n yield stream_update_event(chunk_data)\n else:\n yield custom_event(chunk_data)\n elif chunk_type == \"updates\":\n # Handle state update events (e.g., real-time Plan updates)\n pass # Currently ignore updates events, rely on checkpoint or custom\n \n # --- Loop ended ---\n yield {\"event\": \"end\", \"data\": \"{}\"} # Send an end event to frontend\n\n except Exception as e:\n print(f\"Error during agent execution stream: {e}\")\n traceback.print_exc()\n # Send error event to frontend\n yield {\"event\": \"error\", \"data\": json.dumps({\"message\": f\"Agent execution error: {e}\"})}\n finally:\n if thread_id in active_connections:\n del active_connections[thread_id]\n\n return EventSourceResponse(generate_events())\n\n\ndef main():\n uvicorn.run(\"api.server:app\", host=\"0.0.0.0\", port=8000, reload=True)\n\n\nif __name__ == \"__main__\":\n import sys\n import os\n # 将项目根目录添加到 Python 路径中\n sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))\n main()\n"
},
{
"path": "api/utils.py",
"content": "import json\nfrom typing import Dict, Any, List, Optional\nfrom langchain_core.messages import BaseMessage, AIMessage, HumanMessage, ToolMessage\nfrom langgraph.types import StateSnapshot\n\n\ndef checkpoint_event(value):\n \"\"\"Create a checkpoint event for the client.\"\"\"\n\n def format_values(values: dict):\n formatted_values = values.copy()\n if \"messages\" in formatted_values:\n formatted_values[\"messages\"] = [\n {\n \"type\": msg.get(\"type\") if isinstance(msg, dict) else msg.type,\n \"content\": msg.get(\"content\") if isinstance(msg, dict) else msg.content,\n \"id\": msg.get(\"id\") if isinstance(msg, dict) else msg.id,\n \"tool_calls\": msg.get(\"tool_calls\") if isinstance(msg, dict) else (msg.tool_calls if hasattr(msg, 'tool_calls') else None)\n }\n for msg in formatted_values[\"messages\"]\n ]\n return formatted_values\n\n def format_writes(writes: dict):\n if writes is None:\n return None\n formatted_writes = {}\n for key, value in writes.items():\n if isinstance(value, dict):\n formatted_writes[key] = format_values(value)\n elif isinstance(value, list):\n formatted_writes[key] = [format_values(item) if isinstance(\n item, dict) else item for item in value]\n else:\n formatted_writes[key] = value\n return formatted_writes\n\n configurable = value[\"payload\"][\"config\"][\"configurable\"]\n data = {\n \"next\": value[\"payload\"][\"next\"],\n \"values\": format_values(value[\"payload\"][\"values\"]),\n \"config\": {\n \"configurable\": {\n \"checkpoint_id\": configurable[\"checkpoint_id\"],\n \"checkpoint_ns\": configurable[\"checkpoint_ns\"],\n \"thread_id\": configurable[\"thread_id\"]\n }\n },\n \"metadata\": {\n \"source\": value[\"payload\"][\"metadata\"][\"source\"],\n \"step\": value[\"payload\"][\"metadata\"][\"step\"],\n \"writes\": format_writes(value[\"payload\"][\"metadata\"][\"writes\"]),\n \"parents\": value[\"payload\"][\"metadata\"][\"parents\"]\n }\n }\n return {\n \"event\": \"checkpoint\",\n \"data\": json.dumps(data)\n }\n\n\ndef message_chunk_event(node_name, message_chunk):\n \"\"\"Create a message chunk event for the client.\"\"\"\n\n def format_messages(value):\n \"\"\"Format message chunk into a serializable dictionary. \n This is needed because the message class is not serializable.\n \"\"\"\n return {\n \"content\": value.content,\n \"id\": value.id,\n \"tool_calls\": value.tool_calls if hasattr(value, 'tool_calls') else None,\n \"tool_call_chunks\": value.tool_call_chunks if hasattr(value, 'tool_call_chunks') else None\n }\n\n return {\n \"event\": \"message_chunk\",\n \"data\": json.dumps({\n \"node_name\": node_name,\n \"message_chunk\": format_messages(message_chunk)\n })\n }\n\n\ndef interrupt_event(interrupts):\n \"\"\"Create an interrupt event for the client.\"\"\"\n formatted_interrupts = [{\"value\": interrupt[\"value\"]}\n for interrupt in interrupts]\n return {\n \"event\": \"interrupt\",\n \"data\": json.dumps(formatted_interrupts)\n }\n\n\ndef custom_event(value):\n \"\"\"Create a custom event for the client.\"\"\"\n return {\n \"event\": \"custom\",\n \"data\": json.dumps(value)\n }\n\n\ndef format_state_snapshot(snapshot: StateSnapshot):\n interrupts = []\n for task in snapshot.tasks:\n for interrupt in task.interrupts:\n interrupts.append({\"value\": interrupt.value})\n return {\n \"values\": snapshot.values,\n \"next\": snapshot.next,\n \"config\": snapshot.config,\n \"interrupts\": interrupts,\n \"parent_config\": snapshot.parent_config,\n \"metadata\": snapshot.metadata\n }\n\n\ndef stream_update_event(data: dict):\n \"\"\"为 DeepResearch Agent 的 StreamUpdateData 创建一个 stream_update 事件。\n\n Args:\n data: 从 add_stream_update 产生的、符合 StreamUpdateData 结构的字典。\n\n Returns:\n 符合 SSE EventSourceResponse 格式的字典。\n \"\"\"\n if not isinstance(data, dict):\n # 如果传入的不是字典，返回一个错误事件\n return {\n \"event\": \"error\",\n \"data\": json.dumps({\"message\": \"Internal server error: Invalid stream update data type.\"})\n }\n \n return {\n \"event\": \"stream_update\",\n \"data\": json.dumps(data, default=str)\n }\n"
},
{
"path": "core/__init__.py",
"content": "# Core module initialization"
},
{
"path": "core/a2a/README.md",
"content": "# Mentis A2A (Agent2Agent) 协议集成\n\n本目录 (`core/a2a/`) 包含用于实现 Agent2Agent (A2A) 协议的客户端和服务器实现，使 Mentis Agents 能够与其他支持 A2A 协议的代理系统进行通信和协作。\n\n## 背景\n\nA2A 是由 Google 发起的开放标准，旨在使不同框架（如 LangGraph、CrewAI、Google ADK、Genkit）或不同供应商构建的 AI 代理能够发现彼此的能力，协商交互模式（文本、文件、数据等），并在任务上进行协作。\n\n## 核心组件\n\n### 1. A2A 客户端 (`A2AClient`)\n\n`A2AClient` 类（位于 `client/client.py`）提供了与支持 A2A 协议的服务器进行交互的功能：\n\n* **代理发现:** 支持通过 `.well-known/agent.json` 端点自动发现代理能力（Agent Card）。\n* **任务管理:** 提供发送、获取和取消任务的方法。\n* **推送通知:** 支持设置和获取任务的推送通知配置。\n* **流式响应:** 支持通过流式API接收任务执行的实时更新。\n* **异步架构:** 基于 `asyncio` 和 `httpx` 构建，适合异步应用。\n\n### 2. A2A 服务器 (`A2AServer`)\n\n`A2AServer` 类（位于 `server/server.py`）允许将现有的 Mentis Agent 暴露为支持 A2A 协议的服务：\n\n* **基于 Starlette:** 使用 Starlette 框架提供 HTTP 和 SSE 端点。\n* **任务处理:** 支持任务的创建、执行和状态跟踪。\n* **流式更新:** 通过 Server-Sent Events (SSE) 提供任务执行的实时更新。\n* **Agent Card:** 通过 `.well-known/agent.json` 端点公开代理能力。\n\n### 3. 辅助工具\n\n#### 推送通知认证 (`PushNotificationAuth`)\n\n`PushNotificationAuth` 类（位于 `utils/push_notification_auth.py`）提供了安全的推送通知机制：\n\n* **发送方认证 (`PushNotificationSenderAuth`):** \n - 生成和管理 JWT 密钥对\n - 验证推送通知 URL\n - 签名并发送推送通知\n - 提供 JWKS 端点供接收方获取公钥\n\n* **接收方认证 (`PushNotificationReceiverAuth`):** \n - 从 JWKS URL 加载公钥\n - 验证接收到的推送通知的完整性和时效性\n - 防止重放攻击\n\n#### 内存缓存 (`InMemoryCache`)\n\n`InMemoryCache` 类（位于 `utils/in_memory_cache.py`）提供了线程安全的内存缓存实现：\n\n* **单例模式:** 确保应用中只有一个缓存实例\n* **TTL 支持:** 支持设置缓存项的过期时间\n* **线程安全:** 使用锁机制确保并发安全\n\n## 数据类型\n\nA2A 协议定义了几个关键数据类型（位于 `types.py`）：\n\n* **AgentCard:** 描述代理的元数据，包括名称、描述、URL、能力和技能。\n* **Task:** 表示代理执行的任务，包含状态、内容和产物。\n* **Part:** 内容的一部分，可以是文本、文件或数据。\n* **Artifact:** 代理产生的产物，如结果、生成的文件等。\n* **TaskState:** 任务状态枚举（已提交、进行中、需要输入、已完成、已取消、失败）。\n* **PushNotificationConfig:** 推送通知配置，包含回调URL和认证信息。\n\n## 如何使用\n\n### 1. 创建和使用 A2A 客户端\n\n```python\nimport asyncio\nfrom common.types import AgentCard\nfrom core.a2a.client.client import A2AClient\n\nasync def main():\n # 方式1：直接指定URL创建客户端\n async with A2AClient(url=\"http://localhost:8000/a2a\") as client:\n # 发送任务\n response = await client.send_task({\"text\": \"请帮我研究人工智能\"})\n task_id = response[\"result\"][\"taskId\"]\n \n # 获取任务结果\n task_response = await client.get_task({\"id\": task_id})\n \n # 设置推送通知\n await client.set_task_callback({\n \"taskId\": task_id,\n \"callbackUrl\": \"https://your-callback-url.com/webhook\"\n })\n \n # 方式2：通过Agent Card创建客户端\n agent_card = AgentCard(name=\"Example Agent\", url=\"http://localhost:8000/a2a\")\n async with A2AClient(agent_card=agent_card) as client:\n # 使用流式API接收实时更新\n async for update in client.send_task_streaming({\"text\": \"分析最新的AI趋势\"}):\n print(update)\n\n# 运行\nasyncio.run(main())\n```\n\n### 2. 创建 A2A 服务器\n\n```python\nfrom core.a2a.server.server import A2AServer\nfrom core.a2a.server.task_manager import InMemoryTaskManager\nfrom common.types import AgentCard\n\n# 创建Agent卡片\nagent_card = AgentCard(\n name=\"My Agent\",\n description=\"一个示例代理\",\n url=\"http://localhost:5000\"\n)\n\n# 创建任务管理器\ntask_manager = InMemoryTaskManager()\n\n# 创建服务器\nserver = A2AServer(\n host=\"0.0.0.0\",\n port=5000,\n endpoint=\"/\",\n agent_card=agent_card,\n task_manager=task_manager\n)\n\n# 启动服务器\nserver.start()\n```\n\n### 3. 配置推送通知\n\n#### 发送方配置\n\n```python\nfrom core.a2a.utils.push_notification_auth import PushNotificationSenderAuth\n\n# 创建发送方认证\nsender_auth = PushNotificationSenderAuth()\n\n# 生成密钥对\nsender_auth.generate_jwk()\n\n# 添加JWKS端点到你的服务器\napp.add_route(\"/.well-known/jwks.json\", sender_auth.handle_jwks_endpoint)\n\n# 验证接收方URL\nis_valid = await sender_auth.verify_push_notification_url(\"https://receiver-url.com/webhook\")\n\n# 发送推送通知\nif is_valid:\n await sender_auth.send_push_notification(\n \"https://receiver-url.com/webhook\",\n {\"event\": \"task_completed\", \"taskId\": \"123\"}\n )\n```\n\n#### 接收方配置\n\n```python\nfrom core.a2a.utils.push_notification_auth import PushNotificationReceiverAuth\nfrom starlette.requests import Request\n\n# 创建接收方认证\nreceiver_auth = PushNotificationReceiverAuth()\n\n# 加载发送方的公钥\nawait receiver_auth.load_jwks(\"https://sender-url.com/.well-known/jwks.json\")\n\n# 在webhook处理函数中验证推送通知\nasync def webhook_handler(request: Request):\n is_valid = await receiver_auth.verify_push_notification(request)\n if is_valid:\n # 处理推送通知...\n data = await request.json()\n print(f\"收到有效的推送通知: {data}\")\n```\n\n### 4. 使用内存缓存\n\n```python\nfrom core.a2a.utils.in_memory_cache import InMemoryCache\n\n# 获取缓存实例\ncache = InMemoryCache()\n\n# 设置缓存项（带TTL）\ncache.set(\"api_result\", {\"data\": \"some_value\"}, ttl=300) # 5分钟过期\n\n# 获取缓存项\nresult = cache.get(\"api_result\")\nif result:\n print(f\"从缓存获取结果: {result}\")\nelse:\n print(\"缓存已过期或不存在\")\n \n# 删除缓存项\ncache.delete(\"api_result\")\n\n# 清空所有缓存\ncache.clear()\n```\n\n## 完整示例\n\n查看 `examples/16_a2a_integration_test.py` 获取完整的集成示例，包括：\n\n1. 创建 A2A 服务器，将现有 Agent 暴露为 A2A 服务\n2. 使用 A2A 客户端连接到 A2A 服务器\n3. 创建一个 Agent，使用 A2A 客户端作为工具\n\n运行示例：\n\n```bash\n# 启动 A2A 服务器\npython -m examples.16_a2a_integration_test server\n\n# 运行 A2A 客户端\npython -m examples.16_a2a_integration_test client\n\n# 运行带有 A2A 工具的 Agent\npython -m examples.16_a2a_integration_test agent\n```\n\n## 与 MCP 的关系\n\nMentis 同时支持 MCP（Model Context Protocol）和 A2A（Agent2Agent）协议：\n\n* **MCP:** 专注于 AI 模型与外部工具/服务的交互，主要用于扩展单个 Agent 的能力。\n* **A2A:** 专注于不同 Agent 之间的通信和协作，使多个 Agent 能够协同工作。\n\n这两个协议是互补的，可以同时使用以构建功能强大的 Agent 系统。"
},
{
"path": "core/a2a/__init__.py",
"content": ""
},
{
"path": "core/a2a/agent_task_manager.py",
"content": "import asyncio\nimport logging\nimport traceback\nfrom typing import Dict, Any, Union, AsyncIterable, Optional\nfrom core.a2a.types import (\n TaskState, TaskStatus, Task, Artifact, Message, TextPart,\n SendTaskRequest, SendTaskResponse, GetTaskRequest, GetTaskResponse,\n CancelTaskRequest, CancelTaskResponse, SendTaskStreamingRequest, SendTaskStreamingResponse,\n SetTaskPushNotificationRequest, SetTaskPushNotificationResponse,\n GetTaskPushNotificationRequest, GetTaskPushNotificationResponse,\n TaskResubscriptionRequest, TaskSendParams, JSONRPCResponse, InvalidParamsError,\n TaskNotFoundError, TaskNotCancelableError, PushNotificationNotSupportedError,\n TaskArtifactUpdateEvent, TaskStatusUpdateEvent, InternalError, TaskIdParams,\n PushNotificationConfig\n)\nfrom core.a2a.server.task_manager import TaskManager, InMemoryTaskManager\nfrom core.a2a.server import utils\n\nlogger = logging.getLogger(__name__)\n\nclass AgentTaskManager(InMemoryTaskManager):\n \"\"\"\n AgentTaskManager是连接LangGraph Agent与A2A协议的关键组件。\n 它负责管理任务生命周期、处理流式响应、更新任务状态以及发送推送通知。\n \"\"\"\n def __init__(self, agent, notification_sender_auth=None):\n \"\"\"\n 初始化AgentTaskManager\n \n Args:\n agent: LangGraph Agent实例\n notification_sender_auth: 推送通知认证（可选）\n \"\"\"\n super().__init__()\n self.agent = agent\n self.notification_sender_auth = notification_sender_auth\n \n async def _run_streaming_agent(self, request: SendTaskStreamingRequest):\n \"\"\"\n 运行流式Agent并处理响应\n \n Args:\n request: 流式任务请求\n \"\"\"\n task_send_params: TaskSendParams = request.params\n query = self._get_user_query(task_send_params)\n try:\n async for item in self.agent.stream(query, task_send_params.sessionId):\n is_task_complete = item[\"is_task_complete\"]\n require_user_input = item[\"require_user_input\"]\n artifact = None\n message = None\n parts = [{\"type\": \"text\", \"text\": item[\"content\"]}]\n end_stream = False\n \n if not is_task_complete and not require_user_input:\n task_state = TaskState.WORKING\n message = Message(role=\"agent\", parts=parts)\n elif require_user_input:\n task_state = TaskState.INPUT_REQUIRED\n message = Message(role=\"agent\", parts=parts)\n end_stream = True\n else:\n task_state = TaskState.COMPLETED\n artifact = Artifact(parts=parts, index=0, append=False)\n end_stream = True\n \n task_status = TaskStatus(state=task_state, message=message)\n latest_task = await self.update_store(\n task_send_params.id,\n task_status,\n None if artifact is None else [artifact],\n )\n await self.send_task_notification(latest_task)\n\n if artifact:\n task_artifact_update_event = TaskArtifactUpdateEvent(\n id=task_send_params.id, artifact=artifact\n )\n await self.enqueue_events_for_sse(\n task_send_params.id, task_artifact_update_event\n ) \n task_update_event = TaskStatusUpdateEvent(\n id=task_send_params.id, status=task_status, final=end_stream\n )\n await self.enqueue_events_for_sse(\n task_send_params.id, task_update_event\n )\n except Exception as e:\n logger.error(f\"An error occurred while streaming the response: {e}\")\n await self.enqueue_events_for_sse(\n task_send_params.id,\n InternalError(message=f\"An error occurred while streaming the response: {e}\") \n )\n\n def _get_user_query(self, task_send_params: TaskSendParams) -> str:\n \"\"\"\n 从任务参数中提取用户查询 (采用 Google Demo 的严格方法)\n\n Args:\n task_send_params: 任务发送参数\n\n Returns:\n str: 用户查询文本\n \"\"\"\n if not task_send_params.message or not task_send_params.message.parts:\n logger.warning(f\"[_get_user_query] Message or parts are empty for task {task_send_params.id}\")\n return \"\" # 或者可以抛出错误，取决于你的设计\n\n # 直接获取第一个 part\n part = task_send_params.message.parts[0]\n logger.debug(f\"[_get_user_query] First part: type={type(part)}, value={part!r}\") # 保留调试日志\n\n # 严格检查第一个 part 是否为 TextPart 实例\n if not isinstance(part, TextPart):\n logger.error(f\"[_get_user_query] First part is not a TextPart instance! Type: {type(part)}\")\n # 直接抛出错误，这会中断流程并提供明确信息\n raise ValueError(f\"Expected first message part to be TextPart, but got {type(part)}\")\n\n # 如果检查通过，直接返回文本\n logger.debug(f\"[_get_user_query] Extracted query from TextPart: '{part.text}'\")\n return part.text\n\n\n def _validate_request(\n self, request: Union[SendTaskRequest, SendTaskStreamingRequest]\n ) -> JSONRPCResponse | None:\n \"\"\"\n 验证请求参数\n \n Args:\n request: 任务请求\n \n Returns:\n JSONRPCResponse | None: 错误响应或None\n \"\"\"\n task_send_params: TaskSendParams = request.params\n if not utils.are_modalities_compatible(\n task_send_params.acceptedOutputModes, self.agent.SUPPORTED_CONTENT_TYPES\n ):\n logger.warning(\n \"Unsupported output mode. Received %s, Support %s\",\n task_send_params.acceptedOutputModes,\n self.agent.SUPPORTED_CONTENT_TYPES,\n )\n return utils.new_incompatible_types_error(request.id)\n \n if task_send_params.pushNotification and not task_send_params.pushNotification.url:\n logger.warning(\"Push notification URL is missing\")\n return JSONRPCResponse(id=request.id, error=InvalidParamsError(message=\"Push notification URL is missing\"))\n \n return None\n \n async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:\n \"\"\"\n 处理发送任务请求\n \n Args:\n request: 任务请求\n \n Returns:\n SendTaskResponse: 任务响应\n \"\"\"\n validation_error = self._validate_request(request)\n if validation_error:\n return SendTaskResponse(id=request.id, error=validation_error.error)\n \n if request.params.pushNotification:\n if not await self.set_push_notification_info(request.params.id, request.params.pushNotification):\n return SendTaskResponse(id=request.id, error=InvalidParamsError(message=\"Push notification URL is invalid\"))\n\n await self.upsert_task(request.params)\n task = await self.update_store(\n request.params.id, TaskStatus(state=TaskState.WORKING), None\n )\n await self.send_task_notification(task)\n\n task_send_params: TaskSendParams = request.params\n query = self._get_user_query(task_send_params)\n try:\n agent_response = self.agent.invoke(query, task_send_params.sessionId)\n # 处理Agent响应并更新任务状态\n parts = [{\"type\": \"text\", \"text\": agent_response}]\n artifact = Artifact(parts=parts, index=0, append=False)\n task = await self.update_store(\n task_send_params.id, \n TaskStatus(state=TaskState.COMPLETED), \n [artifact]\n )\n await self.send_task_notification(task)\n return SendTaskResponse(id=request.id, result=task)\n\n except Exception as e:\n # 建议也稍微改进一下异常处理日志和返回信息\n logger.error(f\"Error during agent invocation or task processing: {e}\", exc_info=True)\n # 记录失败状态\n try:\n # 确保即使在异常处理中也能更新状态\n task_failed : Task = await self.update_store(\n task_send_params.id,\n TaskStatus(state=TaskState.FAILED, error={\"message\": str(e)}),\n None\n )\n await self.send_task_notification(task_failed)\n except Exception as update_err:\n # 如果更新状态也失败，记录下来\n logger.error(f\"Failed to update task status to FAILED after initial error: {update_err}\", exc_info=True)\n\n # 返回更合适的错误类型和消息\n # return SendTaskResponse(id=request.id, error=InvalidParamsError(message=f\"Error processing task: {e}\"))\n # InternalError 可能更合适，因为错误发生在服务器内部处理中\n return SendTaskResponse(id=request.id, error=InternalError(message=f\"Error processing task: {str(e) or type(e).__name__}\"))\n\n \n async def on_send_task_subscribe(\n self, request: SendTaskStreamingRequest\n ) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:\n \"\"\"\n 处理流式任务请求\n \n Args:\n request: 流式任务请求\n \n Returns:\n AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse: 流式响应或错误\n \"\"\"\n try:\n error = self._validate_request(request)\n if error:\n return error\n \n await self.upsert_task(request.params)\n \n if request.params.pushNotification:\n if not await self.set_push_notification_info(request.params.id, request.params.pushNotification):\n return JSONRPCResponse(id=request.id, error=InvalidParamsError(message=\"Push notification URL is invalid\"))\n \n task_send_params: TaskSendParams = request.params\n sse_event_queue = await self.setup_sse_consumer(task_send_params.id, False) \n asyncio.create_task(self._run_streaming_agent(request))\n \n return self.dequeue_events_for_sse(\n request.id, task_send_params.id, sse_event_queue\n )\n except Exception as e:\n logger.error(f\"Error in SSE stream: {e}\")\n print(traceback.format_exc())\n return JSONRPCResponse(\n id=request.id,\n error=InternalError(\n message=\"An error occurred while streaming the response\"\n ),\n )\n\n async def _process_agent_response(\n self, request: SendTaskRequest, agent_response: dict\n ) -> SendTaskResponse:\n \"\"\"Processes the agent's response and updates the task store.\"\"\"\n task_send_params: TaskSendParams = request.params\n task_id = task_send_params.id\n history_length = task_send_params.historyLength\n task_status = None\n\n parts = [{\"type\": \"text\", \"text\": agent_response[\"content\"]}]\n artifact = None\n if agent_response[\"require_user_input\"]:\n task_status = TaskStatus(\n state=TaskState.INPUT_REQUIRED,\n message=Message(role=\"agent\", parts=parts),\n )\n else:\n task_status = TaskStatus(state=TaskState.COMPLETED)\n artifact = Artifact(parts=parts)\n task = await self.update_store(\n task_id, task_status, None if artifact is None else [artifact]\n )\n task_result = self.append_task_history(task, history_length)\n await self.send_task_notification(task)\n return SendTaskResponse(id=request.id, result=task_result)\n \n async def on_resubscribe_to_task(\n self, request: TaskResubscriptionRequest\n ) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:\n task_id_params: TaskIdParams = request.params\n try:\n sse_event_queue = await self.setup_sse_consumer(task_id_params.id, True)\n return self.dequeue_events_for_sse(request.id, task_id_params.id, sse_event_queue)\n except Exception as e:\n logger.error(f\"Error while reconnecting to SSE stream: {e}\")\n return JSONRPCResponse(\n id=request.id,\n error=InternalError(\n message=f\"An error occurred while reconnecting to stream: {e}\"\n ),\n )\n \n async def send_task_notification(self, task: Task):\n if not await self.has_push_notification_info(task.id):\n logger.info(f\"No push notification info found for task {task.id}\")\n return\n push_info = await self.get_push_notification_info(task.id)\n\n logger.info(f\"Notifying for task {task.id} => {task.status.state}\")\n await self.notification_sender_auth.send_push_notification(\n push_info.url,\n data=task.model_dump(exclude_none=True)\n )\n\n async def set_push_notification_info(self, task_id: str, push_notification_config: PushNotificationConfig):\n # Verify the ownership of notification URL by issuing a challenge request.\n if self.notification_sender_auth:\n is_verified = await self.notification_sender_auth.verify_push_notification_url(push_notification_config.url)\n if not is_verified:\n return False\n \n await super().set_push_notification_info(task_id, push_notification_config)\n return True"
},
{
"path": "core/a2a/client/__init__.py",
"content": ""
},
{
"path": "core/a2a/client/card_resolver.py",
"content": "import httpx\nfrom core.a2a.types import (\n AgentCard,\n A2AClientJSONError,\n)\nimport json\n\n\nclass A2ACardResolver:\n def __init__(self, base_url, agent_card_path=\"/.well-known/agent.json\"):\n self.base_url = base_url.rstrip(\"/\")\n self.agent_card_path = agent_card_path.lstrip(\"/\")\n\n def get_agent_card(self) -> AgentCard:\n with httpx.Client() as client:\n response = client.get(self.base_url + \"/\" + self.agent_card_path)\n response.raise_for_status()\n try:\n return AgentCard(**response.json())\n except json.JSONDecodeError as e:\n raise A2AClientJSONError(str(e)) from e"
},
{
"path": "core/a2a/client/client.py",
"content": "import httpx\nfrom httpx_sse import connect_sse\nfrom typing import Any, AsyncIterable\nfrom core.a2a.types import (\n AgentCard,\n GetTaskRequest,\n SendTaskRequest,\n SendTaskResponse,\n JSONRPCRequest,\n GetTaskResponse,\n CancelTaskResponse,\n CancelTaskRequest,\n SetTaskPushNotificationRequest,\n SetTaskPushNotificationResponse,\n GetTaskPushNotificationRequest,\n GetTaskPushNotificationResponse,\n A2AClientHTTPError,\n A2AClientJSONError,\n SendTaskStreamingRequest,\n SendTaskStreamingResponse,\n)\nimport json\n\n\nclass A2AClient:\n def __init__(self, agent_card: AgentCard = None, url: str = None):\n if agent_card:\n self.url = agent_card.url\n elif url:\n self.url = url\n else:\n raise ValueError(\"Must provide either agent_card or url\")\n\n async def send_task(self, payload: dict[str, Any]) -> SendTaskResponse:\n request = SendTaskRequest(params=payload)\n return SendTaskResponse(**await self._send_request(request))\n\n async def send_task_streaming(\n self, payload: dict[str, Any]\n ) -> AsyncIterable[SendTaskStreamingResponse]:\n request = SendTaskStreamingRequest(params=payload)\n with httpx.Client(timeout=None) as client:\n with connect_sse(\n client, \"POST\", self.url, json=request.model_dump()\n ) as event_source:\n try:\n for sse in event_source.iter_sse():\n yield SendTaskStreamingResponse(**json.loads(sse.data))\n except json.JSONDecodeError as e:\n raise A2AClientJSONError(str(e)) from e\n except httpx.RequestError as e:\n raise A2AClientHTTPError(400, str(e)) from e\n\n async def _send_request(self, request: JSONRPCRequest) -> dict[str, Any]:\n async with httpx.AsyncClient() as client:\n try:\n # Image generation could take time, adding timeout\n response = await client.post(\n self.url, json=request.model_dump(), timeout=30\n )\n response.raise_for_status()\n return response.json()\n except httpx.HTTPStatusError as e:\n raise A2AClientHTTPError(e.response.status_code, str(e)) from e\n except json.JSONDecodeError as e:\n raise A2AClientJSONError(str(e)) from e\n\n async def get_task(self, payload: dict[str, Any]) -> GetTaskResponse:\n request = GetTaskRequest(params=payload)\n return GetTaskResponse(**await self._send_request(request))\n\n async def cancel_task(self, payload: dict[str, Any]) -> CancelTaskResponse:\n request = CancelTaskRequest(params=payload)\n return CancelTaskResponse(**await self._send_request(request))\n\n async def set_task_callback(\n self, payload: dict[str, Any]\n ) -> SetTaskPushNotificationResponse:\n request = SetTaskPushNotificationRequest(params=payload)\n return SetTaskPushNotificationResponse(**await self._send_request(request))\n\n async def get_task_callback(\n self, payload: dict[str, Any]\n ) -> GetTaskPushNotificationResponse:\n request = GetTaskPushNotificationRequest(params=payload)\n return GetTaskPushNotificationResponse(**await self._send_request(request))"
},
{
"path": "core/a2a/config.json",
"content": "{\n \"local_agent\": {\n \"url\": \"http://127.0.0.1:8000/\",\n \"auth\": {\n \"type\": \"none\"\n }\n }\n}"
},
{
"path": "core/a2a/server/__init__.py",
"content": ""
},
{
"path": "core/a2a/server/server.py",
"content": "# core/a2a/server/server.py\nfrom starlette.applications import Starlette\nfrom starlette.responses import JSONResponse\nfrom sse_starlette.sse import EventSourceResponse\nfrom starlette.requests import Request\nfrom starlette.middleware import Middleware\nfrom starlette.middleware.cors import CORSMiddleware\n\n# --- 添加 Pydantic 的 ValidationError 导入 ---\nfrom pydantic import ValidationError\n# --- 导入结束 ---\n\nfrom core.a2a.types import (\n A2ARequest,\n JSONRPCResponse,\n InvalidRequestError,\n JSONParseError,\n GetTaskRequest,\n CancelTaskRequest,\n SendTaskRequest,\n SetTaskPushNotificationRequest,\n GetTaskPushNotificationRequest,\n InternalError,\n AgentCard,\n TaskResubscriptionRequest,\n SendTaskStreamingRequest,\n MethodNotFoundError,\n # 确保 ValidationError 没有在这里导入\n)\nimport json\nfrom typing import AsyncIterable, Any, Optional, Union\nfrom core.a2a.server.task_manager import TaskManager\n\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n\nclass A2AServer:\n def __init__(\n self,\n host=\"0.0.0.0\",\n port=5000,\n endpoint=\"/\",\n agent_card: AgentCard = None,\n task_manager: TaskManager = None,\n allowed_origins: Optional[list[str]] = None,\n ):\n self.host = host\n self.port = port\n self.endpoint = endpoint\n self.task_manager = task_manager\n self.agent_card = agent_card\n\n if allowed_origins is None:\n # 本地开发时默认只允许 localhost:3000\n allowed_origins = [\"http://localhost:3000\"]\n logger.warning(\"CORS allow_origins set to 'http://localhost:3000' for local development.\")\n else:\n logger.info(f\"CORS allow_origins configured: {allowed_origins}\")\n\n middleware = [\n Middleware(\n CORSMiddleware,\n allow_origins=allowed_origins,\n allow_credentials=True,\n allow_methods=[\"*\"],\n allow_headers=[\"*\"],\n )\n ]\n self.app = Starlette(middleware=middleware, debug=True)\n self.app.add_route(self.endpoint, self._process_request, methods=[\"POST\"])\n self.app.add_route(\n \"/.well-known/agent.json\", self._get_agent_card, methods=[\"GET\"]\n )\n logger.info(f\"A2AServer initialized. Endpoint: {self.endpoint}, Agent Card Endpoint: /.well-known/agent.json\")\n\n def start(self):\n if self.agent_card is None: raise ValueError(\"agent_card must be provided to A2AServer\")\n if self.task_manager is None: raise ValueError(\"task_manager must be provided to A2AServer\")\n import uvicorn\n logger.info(f\"Starting Uvicorn server on {self.host}:{self.port}...\")\n uvicorn.run(self.app, host=self.host, port=self.port)\n\n def _get_agent_card(self, request: Request) -> JSONResponse:\n logger.debug(\"Received request for /.well-known/agent.json\")\n if not self.agent_card:\n logger.error(\"Agent card requested but not configured in A2AServer.\")\n return JSONResponse({\"error\": \"Agent card not configured\"}, status_code=500)\n return JSONResponse(self.agent_card.model_dump(exclude_none=True))\n\n async def _process_request(self, request: Request) -> Union[JSONResponse, EventSourceResponse]:\n result = None; json_rpc_request = None; request_id_for_error = None\n try:\n try: body = await request.json(); logger.debug(f\"Received request body: {body}\")\n except json.JSONDecodeError as e: logger.error(f\"JSON decoding failed: {e}\"); raise JSONParseError()\n\n try:\n json_rpc_request = A2ARequest.validate_python(body); request_id_for_error = getattr(json_rpc_request, 'id', None)\n logger.info(f\"Processing valid A2A request: Method='{json_rpc_request.method}', ID='{request_id_for_error}', TaskID='{getattr(json_rpc_request.params, 'id', 'N/A')}'\")\n except ValidationError as e:\n logger.error(f\"A2A request validation failed: {e}\"); req_id_fallback = body.get('id') if isinstance(body, dict) else None\n # 注意: 这里抛出的 InvalidRequestError 会在下面的 except Exception 中被捕获\n raise InvalidRequestError(data=json.loads(e.json())) from e\n\n # 分发给 TaskManager\n if isinstance(json_rpc_request, GetTaskRequest): result = await self.task_manager.on_get_task(json_rpc_request)\n elif isinstance(json_rpc_request, SendTaskRequest): result = await self.task_manager.on_send_task(json_rpc_request)\n elif isinstance(json_rpc_request, SendTaskStreamingRequest): result = await self.task_manager.on_send_task_subscribe(json_rpc_request)\n elif isinstance(json_rpc_request, CancelTaskRequest): result = await self.task_manager.on_cancel_task(json_rpc_request)\n elif isinstance(json_rpc_request, SetTaskPushNotificationRequest): result = await self.task_manager.on_set_task_push_notification(json_rpc_request)\n elif isinstance(json_rpc_request, GetTaskPushNotificationRequest): result = await self.task_manager.on_get_task_push_notification(json_rpc_request)\n elif isinstance(json_rpc_request, TaskResubscriptionRequest): result = await self.task_manager.on_resubscribe_to_task(json_rpc_request)\n else: logger.warning(f\"Unhandled validated request type: {type(json_rpc_request)}\"); raise MethodNotFoundError(data={\"method\": getattr(json_rpc_request, 'method', 'unknown')})\n\n logger.debug(f\"[A2AServer] Result from TaskManager method '{json_rpc_request.method}': type={type(result)}\")\n return self._create_response(result) # 调用 _create_response\n\n except Exception as e:\n # 统一处理所有在请求处理（包括验证和 task manager 调用）中发生的异常\n logger.error(f\"Exception during request processing: {e}\", exc_info=True)\n return self._handle_exception(e, request_id=request_id_for_error) # 使用 _handle_exception\n\n def _handle_exception(self, e: Exception, request_id: Optional[Union[str, int]] = None) -> JSONResponse:\n status_code = 500; json_rpc_error: Optional[JSONRPCError] = None\n if isinstance(e, JSONParseError): json_rpc_error = e; status_code = 400\n elif isinstance(e, InvalidRequestError): json_rpc_error = e; status_code = 400\n elif isinstance(e, MethodNotFoundError): json_rpc_error = e; status_code = 404 # 或 501\n # --- 现在可以正确捕获 Pydantic 的 ValidationError ---\n elif isinstance(e, ValidationError):\n logger.warning(f\"Pydantic Validation error caught in handler: {e}\")\n error_data = str(e); \n try: error_data = json.loads(e.json()) \n except: pass\n # 通常 Pydantic 验证错误发生在请求处理阶段是 InvalidRequestError 的一种\n # 如果发生在响应创建阶段则更像是 InternalError\n json_rpc_error = InvalidRequestError(message=\"Request/Response data validation failed\", data=error_data)\n status_code = 400 # 认为是客户端请求或服务器返回的数据结构问题\n # --- 捕获结束 ---\n elif isinstance(e, ValueError) and \"Unexpected result type\" in str(e):\n logger.error(f\"Internal error due to unexpected result type: {e}\", exc_info=False)\n json_rpc_error = InternalError(message=\"Server error: Unexpected result type from handler.\")\n status_code = 500\n elif isinstance(e, NotImplementedError):\n logger.error(f\"Method not implemented: {e}\", exc_info=True)\n json_rpc_error = MethodNotFoundError(message=f\"Method not implemented: {e}\")\n status_code = 501\n else:\n logger.error(f\"Unhandled internal exception: {e}\", exc_info=True)\n json_rpc_error = InternalError(message=f\"An internal server error occurred: {type(e).__name__}\")\n status_code = 500\n\n response = JSONRPCResponse(id=request_id, error=json_rpc_error)\n logger.debug(f\"Returning error response: {response.model_dump(exclude_none=True)}\")\n return JSONResponse(response.model_dump(exclude_none=True), status_code=status_code)\n\n def _create_response(self, result: Any) -> Union[JSONResponse, EventSourceResponse]:\n if isinstance(result, AsyncIterable):\n logger.debug(\"[A2AServer] Creating EventSourceResponse (text/event-stream)\")\n async def event_generator(stream_result: AsyncIterable) -> AsyncIterable[dict[str, str]]:\n try:\n async for item in stream_result:\n if hasattr(item, 'model_dump_json'):\n json_data = item.model_dump_json(exclude_none=True)\n logger.debug(f\"A2AServer yielding SSE data: {json_data}\")\n yield {\"data\": json_data}\n else:\n logger.warning(f\"Yielding non-Pydantic object in event stream: {type(item)}\")\n yield {\"data\": json.dumps(str(item))}\n except Exception as gen_err:\n logger.error(f\"Error during SSE event generation: {gen_err}\", exc_info=True)\n try:\n # 尝试 yield 一个标准的 JSON-RPC 错误事件\n error_payload = JSONRPCResponse(id=None, error=InternalError(message=f\"Streaming generation error: {gen_err}\"))\n yield {\"event\": \"error\", \"data\": error_payload.model_dump_json(exclude_none=True)}\n except Exception as yield_err:\n logger.error(f\"Failed to yield error event to SSE stream: {yield_err}\", exc_info=True)\n\n return EventSourceResponse(event_generator(result))\n elif isinstance(result, JSONRPCResponse):\n logger.debug(\"[A2AServer] Creating JSONResponse (application/json)\")\n return JSONResponse(result.model_dump(exclude_none=True))\n else:\n logger.error(f\"Unexpected result type received by _create_response: {type(result)}\")\n raise ValueError(f\"Unexpected result type: {type(result)}\")"
},
{
"path": "core/a2a/server/task_manager.py",
"content": "from abc import ABC, abstractmethod\nfrom typing import Union, AsyncIterable, List\nfrom core.a2a.types import Task\nfrom core.a2a.types import (\n JSONRPCResponse,\n TaskIdParams,\n TaskQueryParams,\n GetTaskRequest,\n TaskNotFoundError,\n SendTaskRequest,\n CancelTaskRequest,\n TaskNotCancelableError,\n SetTaskPushNotificationRequest,\n GetTaskPushNotificationRequest,\n GetTaskResponse,\n CancelTaskResponse,\n SendTaskResponse,\n SetTaskPushNotificationResponse,\n GetTaskPushNotificationResponse,\n PushNotificationNotSupportedError,\n TaskSendParams,\n TaskStatus,\n TaskState,\n TaskResubscriptionRequest,\n SendTaskStreamingRequest,\n SendTaskStreamingResponse,\n Artifact,\n PushNotificationConfig,\n TaskStatusUpdateEvent,\n JSONRPCError,\n TaskPushNotificationConfig,\n InternalError,\n)\nfrom core.a2a.server.utils import new_not_implemented_error\nimport asyncio\nimport logging\n\nlogger = logging.getLogger(__name__)\n\nclass TaskManager(ABC):\n @abstractmethod\n async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:\n pass\n\n @abstractmethod\n async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:\n pass\n\n @abstractmethod\n async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:\n pass\n\n @abstractmethod\n async def on_send_task_subscribe(\n self, request: SendTaskStreamingRequest\n ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:\n pass\n\n @abstractmethod\n async def on_set_task_push_notification(\n self, request: SetTaskPushNotificationRequest\n ) -> SetTaskPushNotificationResponse:\n pass\n\n @abstractmethod\n async def on_get_task_push_notification(\n self, request: GetTaskPushNotificationRequest\n ) -> GetTaskPushNotificationResponse:\n pass\n\n @abstractmethod\n async def on_resubscribe_to_task(\n self, request: TaskResubscriptionRequest\n ) -> Union[AsyncIterable[SendTaskResponse], JSONRPCResponse]:\n pass\n\n\nclass InMemoryTaskManager(TaskManager):\n def __init__(self):\n self.tasks: dict[str, Task] = {}\n self.push_notification_infos: dict[str, PushNotificationConfig] = {}\n self.lock = asyncio.Lock()\n self.task_sse_subscribers: dict[str, List[asyncio.Queue]] = {}\n self.subscriber_lock = asyncio.Lock()\n\n async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:\n logger.info(f\"Getting task {request.params.id}\")\n task_query_params: TaskQueryParams = request.params\n\n async with self.lock:\n task = self.tasks.get(task_query_params.id)\n if task is None:\n return GetTaskResponse(id=request.id, error=TaskNotFoundError())\n\n task_result = self.append_task_history(\n task, task_query_params.historyLength\n )\n\n return GetTaskResponse(id=request.id, result=task_result)\n\n async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:\n logger.info(f\"Cancelling task {request.params.id}\")\n task_id_params: TaskIdParams = request.params\n\n async with self.lock:\n task = self.tasks.get(task_id_params.id)\n if task is None:\n return CancelTaskResponse(id=request.id, error=TaskNotFoundError())\n\n return CancelTaskResponse(id=request.id, error=TaskNotCancelableError())\n\n @abstractmethod\n async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:\n pass\n\n @abstractmethod\n async def on_send_task_subscribe(\n self, request: SendTaskStreamingRequest\n ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:\n pass\n\n async def set_push_notification_info(self, task_id: str, notification_config: PushNotificationConfig):\n async with self.lock:\n task = self.tasks.get(task_id)\n if task is None:\n raise ValueError(f\"Task not found for {task_id}\")\n\n self.push_notification_infos[task_id] = notification_config\n\n return\n \n async def get_push_notification_info(self, task_id: str) -> PushNotificationConfig:\n async with self.lock:\n task = self.tasks.get(task_id)\n if task is None:\n raise ValueError(f\"Task not found for {task_id}\")\n\n return self.push_notification_infos[task_id]\n \n return\n \n async def has_push_notification_info(self, task_id: str) -> bool:\n async with self.lock:\n return task_id in self.push_notification_infos\n \n\n async def on_set_task_push_notification(\n self, request: SetTaskPushNotificationRequest\n ) -> SetTaskPushNotificationResponse:\n logger.info(f\"Setting task push notification {request.params.id}\")\n task_notification_params: TaskPushNotificationConfig = request.params\n\n try:\n await self.set_push_notification_info(task_notification_params.id, task_notification_params.pushNotificationConfig)\n except Exception as e:\n logger.error(f\"Error while setting push notification info: {e}\")\n return JSONRPCResponse(\n id=request.id,\n error=InternalError(\n message=\"An error occurred while setting push notification info\"\n ),\n )\n \n return SetTaskPushNotificationResponse(id=request.id, result=task_notification_params)\n\n async def on_get_task_push_notification(\n self, request: GetTaskPushNotificationRequest\n ) -> GetTaskPushNotificationResponse:\n logger.info(f\"Getting task push notification {request.params.id}\")\n task_params: TaskIdParams = request.params\n\n try:\n notification_info = await self.get_push_notification_info(task_params.id)\n except Exception as e:\n logger.error(f\"Error while getting push notification info: {e}\")\n return GetTaskPushNotificationResponse(\n id=request.id,\n error=InternalError(\n message=\"An error occurred while getting push notification info\"\n ),\n )\n \n return GetTaskPushNotificationResponse(id=request.id, result=TaskPushNotificationConfig(id=task_params.id, pushNotificationConfig=notification_info))\n\n async def upsert_task(self, task_send_params: TaskSendParams) -> Task:\n logger.info(f\"Upserting task {task_send_params.id}\")\n async with self.lock:\n task = self.tasks.get(task_send_params.id)\n if task is None:\n task = Task(\n id=task_send_params.id,\n sessionId = task_send_params.sessionId,\n messages=[task_send_params.message],\n status=TaskStatus(state=TaskState.SUBMITTED),\n history=[task_send_params.message],\n )\n self.tasks[task_send_params.id] = task\n else:\n task.history.append(task_send_params.message)\n\n return task\n\n async def on_resubscribe_to_task(\n self, request: TaskResubscriptionRequest\n ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:\n return new_not_implemented_error(request.id)\n\n async def update_store(\n self, task_id: str, status: TaskStatus, artifacts: list[Artifact]\n ) -> Task:\n async with self.lock:\n try:\n task = self.tasks[task_id]\n except KeyError:\n logger.error(f\"Task {task_id} not found for updating the task\")\n raise ValueError(f\"Task {task_id} not found\")\n\n task.status = status\n\n if status.message is not None:\n task.history.append(status.message)\n\n if artifacts is not None:\n if task.artifacts is None:\n task.artifacts = []\n task.artifacts.extend(artifacts)\n\n return task\n\n def append_task_history(self, task: Task, historyLength: int | None):\n new_task = task.model_copy()\n if historyLength is not None and historyLength > 0:\n new_task.history = new_task.history[-historyLength:]\n else:\n new_task.history = []\n\n return new_task \n\n async def setup_sse_consumer(self, task_id: str, is_resubscribe: bool = False):\n async with self.subscriber_lock:\n if task_id not in self.task_sse_subscribers:\n if is_resubscribe:\n raise ValueError(\"Task not found for resubscription\")\n else:\n self.task_sse_subscribers[task_id] = []\n\n sse_event_queue = asyncio.Queue(maxsize=0) # <=0 is unlimited\n self.task_sse_subscribers[task_id].append(sse_event_queue)\n return sse_event_queue\n\n async def enqueue_events_for_sse(self, task_id, task_update_event):\n async with self.subscriber_lock:\n if task_id not in self.task_sse_subscribers:\n return\n\n current_subscribers = self.task_sse_subscribers[task_id]\n for subscriber in current_subscribers:\n await subscriber.put(task_update_event)\n\n async def dequeue_events_for_sse(\n self, request_id, task_id, sse_event_queue: asyncio.Queue\n ) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:\n try:\n while True: \n event = await sse_event_queue.get()\n if isinstance(event, JSONRPCError):\n yield SendTaskStreamingResponse(id=request_id, error=event)\n break\n \n yield SendTaskStreamingResponse(id=request_id, result=event)\n if isinstance(event, TaskStatusUpdateEvent) and event.final:\n break\n finally:\n async with self.subscriber_lock:\n if task_id in self.task_sse_subscribers:\n self.task_sse_subscribers[task_id].remove(sse_event_queue)\n"
},
{
"path": "core/a2a/server/utils.py",
"content": "from core.a2a.types import (\n JSONRPCResponse,\n ContentTypeNotSupportedError,\n UnsupportedOperationError,\n)\nfrom typing import List\n\n\ndef are_modalities_compatible(\n server_output_modes: List[str], client_output_modes: List[str]\n):\n \"\"\"Modalities are compatible if they are both non-empty\n and there is at least one common element.\"\"\"\n if client_output_modes is None or len(client_output_modes) == 0:\n return True\n\n if server_output_modes is None or len(server_output_modes) == 0:\n return True\n\n return any(x in server_output_modes for x in client_output_modes)\n\n\ndef new_incompatible_types_error(request_id):\n return JSONRPCResponse(id=request_id, error=ContentTypeNotSupportedError())\n\n\ndef new_not_implemented_error(request_id):\n return JSONRPCResponse(id=request_id, error=UnsupportedOperationError())"
},
{
"path": "core/a2a/types.py",
"content": "from typing import Union, Any\nfrom pydantic import BaseModel, Field, TypeAdapter\nfrom typing import Literal, List, Annotated, Optional\nfrom datetime import datetime\nfrom pydantic import model_validator, ConfigDict, field_serializer\nfrom uuid import uuid4\nfrom enum import Enum\nfrom typing_extensions import Self\n\n\nclass TaskState(str, Enum):\n SUBMITTED = \"submitted\"\n WORKING = \"working\"\n INPUT_REQUIRED = \"input-required\"\n COMPLETED = \"completed\"\n CANCELED = \"canceled\"\n FAILED = \"failed\"\n UNKNOWN = \"unknown\"\n\n\nclass TextPart(BaseModel):\n type: Literal[\"text\"] = \"text\"\n text: str\n metadata: dict[str, Any] | None = None\n\n\nclass FileContent(BaseModel):\n name: str | None = None\n mimeType: str | None = None\n bytes: str | None = None\n uri: str | None = None\n\n @model_validator(mode=\"after\")\n def check_content(self) -> Self:\n if not (self.bytes or self.uri):\n raise ValueError(\"Either 'bytes' or 'uri' must be present in the file data\")\n if self.bytes and self.uri:\n raise ValueError(\n \"Only one of 'bytes' or 'uri' can be present in the file data\"\n )\n return self\n\n\nclass FilePart(BaseModel):\n type: Literal[\"file\"] = \"file\"\n file: FileContent\n metadata: dict[str, Any] | None = None\n\n\nclass DataPart(BaseModel):\n type: Literal[\"data\"] = \"data\"\n data: dict[str, Any]\n metadata: dict[str, Any] | None = None\n\n\nPart = Annotated[Union[TextPart, FilePart, DataPart], Field(discriminator=\"type\")]\n\n\nclass Message(BaseModel):\n role: Literal[\"user\", \"agent\"]\n parts: List[Part]\n metadata: dict[str, Any] | None = None\n\n\nclass TaskStatus(BaseModel):\n state: TaskState\n message: Message | None = None\n timestamp: datetime = Field(default_factory=datetime.now)\n\n @field_serializer(\"timestamp\")\n def serialize_dt(self, dt: datetime, _info):\n return dt.isoformat()\n\n\nclass Artifact(BaseModel):\n name: str | None = None\n description: str | None = None\n parts: List[Part]\n metadata: dict[str, Any] | None = None\n index: int = 0\n append: bool | None = None\n lastChunk: bool | None = None\n\n\nclass Task(BaseModel):\n id: str\n sessionId: str | None = None\n status: TaskStatus\n artifacts: List[Artifact] | None = None\n history: List[Message] | None = None\n metadata: dict[str, Any] | None = None\n\n\nclass TaskStatusUpdateEvent(BaseModel):\n id: str\n status: TaskStatus\n final: bool = False\n metadata: dict[str, Any] | None = None\n\n\nclass TaskArtifactUpdateEvent(BaseModel):\n id: str\n artifact: Artifact \n metadata: dict[str, Any] | None = None\n\n\nclass AuthenticationInfo(BaseModel):\n model_config = ConfigDict(extra=\"allow\")\n\n schemes: List[str]\n credentials: str | None = None\n\n\nclass PushNotificationConfig(BaseModel):\n url: str\n token: str | None = None\n authentication: AuthenticationInfo | None = None\n\n\nclass TaskIdParams(BaseModel):\n id: str\n metadata: dict[str, Any] | None = None\n\n\nclass TaskQueryParams(TaskIdParams):\n historyLength: int | None = None\n\n\nclass TaskSendParams(BaseModel):\n id: str\n sessionId: str = Field(default_factory=lambda: uuid4().hex)\n message: Message\n acceptedOutputModes: Optional[List[str]] = None\n pushNotification: PushNotificationConfig | None = None\n historyLength: int | None = None\n metadata: dict[str, Any] | None = None\n\n\nclass TaskPushNotificationConfig(BaseModel):\n id: str\n pushNotificationConfig: PushNotificationConfig\n\n\n## RPC Messages\n\n\nclass JSONRPCMessage(BaseModel):\n jsonrpc: Literal[\"2.0\"] = \"2.0\"\n id: int | str | None = Field(default_factory=lambda: uuid4().hex)\n\n\nclass JSONRPCRequest(JSONRPCMessage):\n method: str\n params: dict[str, Any] | None = None\n\n\nclass JSONRPCError(BaseModel):\n code: int\n message: str\n data: Any | None = None\n\n\nclass JSONRPCResponse(JSONRPCMessage):\n result: Any | None = None\n error: JSONRPCError | None = None\n\n\nclass SendTaskRequest(JSONRPCRequest):\n method: Literal[\"tasks/send\"] = \"tasks/send\"\n params: TaskSendParams\n\n\nclass SendTaskResponse(JSONRPCResponse):\n result: Task | None = None\n\n\nclass SendTaskStreamingRequest(JSONRPCRequest):\n method: Literal[\"tasks/sendSubscribe\"] = \"tasks/sendSubscribe\"\n params: TaskSendParams\n\n\nclass SendTaskStreamingResponse(JSONRPCResponse):\n result: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None\n\n\nclass GetTaskRequest(JSONRPCRequest):\n method: Literal[\"tasks/get\"] = \"tasks/get\"\n params: TaskQueryParams\n\n\nclass GetTaskResponse(JSONRPCResponse):\n result: Task | None = None\n\n\nclass CancelTaskRequest(JSONRPCRequest):\n method: Literal[\"tasks/cancel\",] = \"tasks/cancel\"\n params: TaskIdParams\n\n\nclass CancelTaskResponse(JSONRPCResponse):\n result: Task | None = None\n\n\nclass SetTaskPushNotificationRequest(JSONRPCRequest):\n method: Literal[\"tasks/pushNotification/set\",] = \"tasks/pushNotification/set\"\n params: TaskPushNotificationConfig\n\n\nclass SetTaskPushNotificationResponse(JSONRPCResponse):\n result: TaskPushNotificationConfig | None = None\n\n\nclass GetTaskPushNotificationRequest(JSONRPCRequest):\n method: Literal[\"tasks/pushNotification/get\",] = \"tasks/pushNotification/get\"\n params: TaskIdParams\n\n\nclass GetTaskPushNotificationResponse(JSONRPCResponse):\n result: TaskPushNotificationConfig | None = None\n\n\nclass TaskResubscriptionRequest(JSONRPCRequest):\n method: Literal[\"tasks/resubscribe\",] = \"tasks/resubscribe\"\n params: TaskIdParams\n\n\nA2ARequest = TypeAdapter(\n Annotated[\n Union[\n SendTaskRequest,\n GetTaskRequest,\n CancelTaskRequest,\n SetTaskPushNotificationRequest,\n GetTaskPushNotificationRequest,\n TaskResubscriptionRequest,\n SendTaskStreamingRequest,\n ],\n Field(discriminator=\"method\"),\n ]\n)\n\n## Error types\n\n\nclass JSONParseError(JSONRPCError):\n code: int = -32700\n message: str = \"Invalid JSON payload\"\n data: Any | None = None\n\n\nclass InvalidRequestError(JSONRPCError):\n code: int = -32600\n message: str = \"Request payload validation error\"\n data: Any | None = None\n\n\nclass MethodNotFoundError(JSONRPCError):\n code: int = -32601\n message: str = \"Method not found\"\n data: None = None\n\n\nclass InvalidParamsError(JSONRPCError):\n code: int = -32602\n message: str = \"Invalid parameters\"\n data: Any | None = None\n\n\nclass InternalError(JSONRPCError):\n code: int = -32603\n message: str = \"Internal error\"\n data: Any | None = None\n\n\nclass TaskNotFoundError(JSONRPCError):\n code: int = -32001\n message: str = \"Task not found\"\n data: None = None\n\n\nclass TaskNotCancelableError(JSONRPCError):\n code: int = -32002\n message: str = \"Task cannot be canceled\"\n data: None = None\n\n\nclass PushNotificationNotSupportedError(JSONRPCError):\n code: int = -32003\n message: str = \"Push Notification is not supported\"\n data: None = None\n\n\nclass UnsupportedOperationError(JSONRPCError):\n code: int = -32004\n message: str = \"This operation is not supported\"\n data: None = None\n\n\nclass ContentTypeNotSupportedError(JSONRPCError):\n code: int = -32005\n message: str = \"Incompatible content types\"\n data: None = None\n\n\nclass AgentProvider(BaseModel):\n organization: str\n url: str | None = None\n\n\nclass AgentCapabilities(BaseModel):\n streaming: bool = False\n pushNotifications: bool = False\n stateTransitionHistory: bool = False\n\n\nclass AgentAuthentication(BaseModel):\n schemes: List[str]\n credentials: str | None = None\n\n\nclass AgentSkill(BaseModel):\n id: str\n name: str\n description: str | None = None\n tags: List[str] | None = None\n examples: List[str] | None = None\n inputModes: List[str] | None = None\n outputModes: List[str] | None = None\n\n\nclass AgentCard(BaseModel):\n name: str\n description: str | None = None\n url: str\n provider: AgentProvider | None = None\n version: str\n documentationUrl: str | None = None\n capabilities: AgentCapabilities\n authentication: AgentAuthentication | None = None\n defaultInputModes: List[str] = [\"text\"]\n defaultOutputModes: List[str] = [\"text\"]\n skills: List[AgentSkill]\n\n\nclass A2AClientError(Exception):\n pass\n\n\nclass A2AClientHTTPError(A2AClientError):\n def __init__(self, status_code: int, message: str):\n self.status_code = status_code\n self.message = message\n super().__init__(f\"HTTP Error {status_code}: {message}\")\n\n\nclass A2AClientJSONError(A2AClientError):\n def __init__(self, message: str):\n self.message = message\n super().__init__(f\"JSON Error: {message}\")\n\n\nclass MissingAPIKeyError(Exception):\n \"\"\"Exception for missing API key.\"\"\"\n\n pass"
},
{
"path": "core/a2a/utils/__init__.py",
"content": ""
},
{
"path": "core/a2a/utils/in_memory_cache.py",
"content": "\"\"\"In Memory Cache utility.\"\"\"\n\nimport threading\nimport time\nfrom typing import Any, Dict, Optional\n\n\nclass InMemoryCache:\n \"\"\"A thread-safe Singleton class to manage cache data.\n\n Ensures only one instance of the cache exists across the application.\n \"\"\"\n\n _instance: Optional[\"InMemoryCache\"] = None\n _lock: threading.Lock = threading.Lock()\n _initialized: bool = False\n\n def __new__(cls):\n \"\"\"Override __new__ to control instance creation (Singleton pattern).\n\n Uses a lock to ensure thread safety during the first instantiation.\n\n Returns:\n The singleton instance of InMemoryCache.\n \"\"\"\n if cls._instance is None:\n with cls._lock:\n if cls._instance is None:\n cls._instance = super().__new__(cls)\n return cls._instance\n\n def __init__(self):\n \"\"\"Initialize the cache storage.\n\n Uses a flag (_initialized) to ensure this logic runs only on the very first\n creation of the singleton instance.\n \"\"\"\n if not self._initialized:\n with self._lock:\n if not self._initialized:\n # print(\"Initializing SessionCache storage\")\n self._cache_data: Dict[str, Dict[str, Any]] = {}\n self._ttl: Dict[str, float] = {}\n self._data_lock: threading.Lock = threading.Lock()\n self._initialized = True\n\n def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:\n \"\"\"Set a key-value pair.\n\n Args:\n key: The key for the data.\n value: The data to store.\n ttl: Time to live in seconds. If None, data will not expire.\n \"\"\"\n with self._data_lock:\n self._cache_data[key] = value\n\n if ttl is not None:\n self._ttl[key] = time.time() + ttl\n else:\n if key in self._ttl:\n del self._ttl[key]\n\n def get(self, key: str, default: Any = None) -> Any:\n \"\"\"Get the value associated with a key.\n\n Args:\n key: The key for the data within the session.\n default: The value to return if the session or key is not found.\n\n Returns:\n The cached value, or the default value if not found.\n \"\"\"\n with self._data_lock:\n if key in self._ttl and time.time() > self._ttl[key]:\n del self._cache_data[key]\n del self._ttl[key]\n return default\n return self._cache_data.get(key, default)\n\n def delete(self, key: str) -> None:\n \"\"\"Delete a specific key-value pair from a cache.\n\n Args:\n key: The key to delete.\n\n Returns:\n True if the key was found and deleted, False otherwise.\n \"\"\"\n\n with self._data_lock:\n if key in self._cache_data:\n del self._cache_data[key]\n if key in self._ttl:\n del self._ttl[key]\n return True\n return False\n\n def clear(self) -> bool:\n \"\"\"Remove all data.\n\n Returns:\n True if the data was cleared, False otherwise.\n \"\"\"\n with self._data_lock:\n self._cache_data.clear()\n self._ttl.clear()\n return True\n return False"
},
{
"path": "core/a2a/utils/push_notification_auth.py",
"content": "from jwcrypto import jwk\nimport uuid\nfrom starlette.responses import JSONResponse\nfrom starlette.requests import Request\nfrom typing import Any\n\nimport jwt\nimport time\nimport json\nimport hashlib\nimport httpx\nimport logging\n\nfrom jwt import PyJWK, PyJWKClient\n\nlogger = logging.getLogger(__name__)\nAUTH_HEADER_PREFIX = 'Bearer '\n\nclass PushNotificationAuth:\n def _calculate_request_body_sha256(self, data: dict[str, Any]):\n \"\"\"Calculates the SHA256 hash of a request body.\n\n This logic needs to be same for both the agent who signs the payload and the client verifier.\n \"\"\"\n body_str = json.dumps(\n data,\n ensure_ascii=False,\n allow_nan=False,\n indent=None,\n separators=(\",\", \":\"),\n )\n return hashlib.sha256(body_str.encode()).hexdigest()\n\nclass PushNotificationSenderAuth(PushNotificationAuth):\n def __init__(self):\n self.public_keys = []\n self.private_key_jwk: PyJWK = None\n\n @staticmethod\n async def verify_push_notification_url(url: str) -> bool:\n async with httpx.AsyncClient(timeout=10) as client:\n try:\n validation_token = str(uuid.uuid4())\n response = await client.get(\n url,\n params={\"validationToken\": validation_token}\n )\n response.raise_for_status()\n is_verified = response.text == validation_token\n\n logger.info(f\"Verified push-notification URL: {url} => {is_verified}\") \n return is_verified \n except Exception as e:\n logger.warning(f\"Error during sending push-notification for URL {url}: {e}\")\n\n return False\n\n def generate_jwk(self):\n key = jwk.JWK.generate(kty='RSA', size=2048, kid=str(uuid.uuid4()), use=\"sig\")\n self.public_keys.append(key.export_public(as_dict=True))\n self.private_key_jwk = PyJWK.from_json(key.export_private())\n \n def handle_jwks_endpoint(self, _request: Request):\n \"\"\"Allow clients to fetch public keys.\n \"\"\"\n return JSONResponse({\n \"keys\": self.public_keys\n })\n \n def _generate_jwt(self, data: dict[str, Any]):\n \"\"\"JWT is generated by signing both the request payload SHA digest and time of token generation.\n\n Payload is signed with private key and it ensures the integrity of payload for client.\n Including iat prevents from replay attack.\n \"\"\"\n \n iat = int(time.time())\n\n return jwt.encode(\n {\"iat\": iat, \"request_body_sha256\": self._calculate_request_body_sha256(data)},\n key=self.private_key_jwk,\n headers={\"kid\": self.private_key_jwk.key_id},\n algorithm=\"RS256\"\n )\n\n async def send_push_notification(self, url: str, data: dict[str, Any]):\n jwt_token = self._generate_jwt(data)\n headers = {'Authorization': f\"Bearer {jwt_token}\"}\n async with httpx.AsyncClient(timeout=10) as client: \n try:\n response = await client.post(\n url,\n json=data,\n headers=headers\n )\n response.raise_for_status()\n logger.info(f\"Push-notification sent for URL: {url}\") \n except Exception as e:\n logger.warning(f\"Error during sending push-notification for URL {url}: {e}\")\n\nclass PushNotificationReceiverAuth(PushNotificationAuth):\n def __init__(self):\n self.public_keys_jwks = []\n self.jwks_client = None\n\n async def load_jwks(self, jwks_url: str):\n self.jwks_client = PyJWKClient(jwks_url)\n \n async def verify_push_notification(self, request: Request) -> bool:\n auth_header = request.headers.get(\"Authorization\")\n if not auth_header or not auth_header.startswith(AUTH_HEADER_PREFIX):\n print(\"Invalid authorization header\")\n return False\n \n token = auth_header[len(AUTH_HEADER_PREFIX):]\n signing_key = self.jwks_client.get_signing_key_from_jwt(token)\n\n decode_token = jwt.decode(\n token,\n signing_key,\n options={\"require\": [\"iat\", \"request_body_sha256\"]},\n algorithms=[\"RS256\"],\n )\n\n actual_body_sha256 = self._calculate_request_body_sha256(await request.json())\n if actual_body_sha256 != decode_token[\"request_body_sha256\"]:\n # Payload signature does not match the digest in signed token.\n raise ValueError(\"Invalid request body\")\n \n if time.time() - decode_token[\"iat\"] > 60 * 5:\n # Do not allow push-notifications older than 5 minutes.\n # This is to prevent replay attack.\n raise ValueError(\"Token is expired\")\n \n return True"
},
{
"path": "core/agents/__init__.py",
"content": "# Agents module initialization"
},
{
"path": "core/agents/base/base_agent.py",
"content": "import json\nfrom typing import List, Dict, Any, Optional, Union, Callable, Sequence, TypeVar, cast\nfrom langchain_core.language_models.chat_models import BaseChatModel\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, ToolMessage\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.runnables import RunnableConfig\nfrom langgraph.graph import StateGraph\nfrom langgraph.types import Checkpointer\nfrom langgraph.graph.graph import CompiledGraph\nfrom langgraph.graph.state import CompiledStateGraph\nimport logging\ntry:\n import tiktoken\n TIKTOKEN_AVAILABLE = True\nexcept ImportError:\n TIKTOKEN_AVAILABLE = False\n print(\"Warning: Tiktoken not installed. Using naive token estimation.\")\n\nlogger = logging.getLogger(__name__)\nDEFAULT_MODEL_NAME = \"gpt-4o-mini\"\n\nStateSchema = TypeVar(\"StateSchema\", bound=Union[dict, Any])\n\nclass BaseAgent:\n def __init__(\n self,\n name: str,\n model: Union[BaseChatModel, LanguageModelLike],\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[Union[str, SystemMessage, Callable]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None, # Limit number of recent messages\n max_context_tokens: Optional[int] = None, # Limit total estimated tokens\n model_name: Optional[str] = \"gpt-4o-mini\", # Optional, used for future token estimation improvements\n description: str = \"No description provided.\"\n \n ):\n if max_context_messages and max_context_tokens:\n raise ValueError(\"Only one of max_context_messages or max_context_tokens should be set.\")\n if name is None or name == \"LangGraph\":\n raise ValueError(\"Agent name must be specified.\")\n\n self.name = name\n self.model = model\n self.tools = tools or []\n self.base_prompt = prompt\n self.checkpointer = checkpointer\n self.max_context_messages = max_context_messages\n self.max_context_tokens = max_context_tokens\n self.model_name = model_name or getattr(model, \"model_name\", DEFAULT_MODEL_NAME)\n self.description = description\n \n self._workflow: Optional[StateGraph] = None\n self._compiled_agent: Optional[CompiledGraph] = None # Stores the final compiled graph\n\n self._tokenizer = None\n if TIKTOKEN_AVAILABLE:\n try: self._tokenizer = tiktoken.encoding_for_model(self.model_name)\n except KeyError:\n try:\n self._tokenizer = tiktoken.get_encoding(\"cl100k_base\")\n # print(f\"Warning: Tiktoken encoding for model '{self.model_name}' not found. Using 'cl100k_base'.\")\n except Exception as e: print(f\"Error getting tiktoken encoding 'cl100k_base': {e}.\")\n except Exception as e: print(f\"Error initializing tiktoken for model '{self.model_name}': {e}.\")\n\n\n def _estimate_tokens(self, message: BaseMessage) -> int:\n content_to_encode = \"\"\n if isinstance(message, (HumanMessage, SystemMessage, AIMessage)):\n if isinstance(message.content, str): content_to_encode = message.content\n elif isinstance(message.content, list):\n for block in message.content:\n if isinstance(block, dict) and block.get(\"type\") == \"text\": content_to_encode += block.get(\"text\", \"\") + \"\\n\"\n elif isinstance(message, ToolMessage):\n content_to_encode = message.content if isinstance(message.content, str) else json.dumps(message.content)\n else: content_to_encode = str(message)\n if self._tokenizer:\n try: return len(self._tokenizer.encode(content_to_encode, disallowed_special=()))\n except Exception: pass\n return len(content_to_encode) // 2\n\n \n def _truncate_by_tokens(self, messages: Sequence[BaseMessage]) -> List[BaseMessage]:\n if not self.max_context_tokens: return list(messages)\n truncated_messages: List[BaseMessage] = []\n total_tokens = 0\n preserved_system_message: Optional[SystemMessage] = None\n # Check if the first message is a SystemMessage, preserve it if so\n # Note: This assumes only ONE leading SystemMessage should be preserved.\n if messages and isinstance(messages[0], SystemMessage):\n preserved_system_message = messages[0]\n messages_to_truncate = messages[1:]\n try: \n system_tokens = self._estimate_tokens(preserved_system_message)\n # Only count if it doesn't exceed limit by itself\n if system_tokens <= self.max_context_tokens:\n total_tokens += system_tokens\n else:\n print(f\"Warning: System message alone ({system_tokens} tokens) exceeds token limit ({self.max_context_tokens}). It might be truncated if context grows.\")\n # Don't add to total_tokens yet, let truncation logic handle it.\n preserved_system_message = None # Don't preserve if it's too big initially\n\n except Exception: pass # Ignore errors estimating system message\n else:\n messages_to_truncate = messages\n\n # Iterate backwards from the most recent message\n for msg in reversed(messages_to_truncate):\n try:\n msg_tokens = self._estimate_tokens(msg)\n # Check if adding this message exceeds the limit\n if total_tokens + msg_tokens <= self.max_context_tokens:\n truncated_messages.append(msg)\n total_tokens += msg_tokens\n else:\n print(f\"Context Token Limit ({self.max_context_tokens}) reached. Truncating older messages.\")\n break # Limit reached\n except Exception as e:\n print(f\"Warning: Failed to estimate tokens for message, skipping: {e}\")\n continue\n\n # Re-add the system message at the beginning if it was preserved\n final_list = list(reversed(truncated_messages))\n if preserved_system_message:\n try: system_tokens = self._estimate_tokens(preserved_system_message)\n except Exception: system_tokens = 0\n # Ensure adding system message doesn't push over limit *again* (edge case)\n if total_tokens - (msg_tokens if 'msg_tokens' in locals() and total_tokens + msg_tokens > self.max_context_tokens else 0) + system_tokens <= self.max_context_tokens:\n final_list.insert(0, preserved_system_message)\n elif not final_list: # If only system message fits\n return [preserved_system_message]\n # Else: System message doesn't fit with the truncated history, omit it.\n\n return final_list\n\n\n def _truncate_messages(self, messages: Sequence[BaseMessage]) -> List[BaseMessage]:\n \"\"\"根据配置（优先 token 数，其次消息数）截断消息历史。\"\"\"\n if self.max_context_tokens is not None:\n return self._truncate_by_tokens(messages)\n elif self.max_context_messages is not None:\n if messages and isinstance(messages[0], SystemMessage):\n # Keep system message + last N-1 messages\n keep_count = self.max_context_messages - 1\n return [messages[0]] + list(messages[-keep_count:]) if keep_count > 0 and len(messages) > 1 else [messages[0]]\n else:\n return list(messages[-self.max_context_messages:])\n return list(messages)\n\n def _get_state_value(self, state: StateSchema, key: str, default: Any = None) -> Any:\n return state.get(key, default) if isinstance(state, dict) else getattr(state, key, default)\n \n def _format_tools_for_prompt(self, tools: List[Union[BaseTool, Callable]]) -> str:\n \"\"\"Formats the tool list for inclusion in the prompt.\"\"\"\n if not tools:\n return \"No tools available for use.\"\n # 使用 getattr 安全地访问 name 和 description\n return \"\\n\".join([\n f\"- **{getattr(t, 'name', 'Unnamed Tool')}**: {getattr(t, 'description', 'No description available.')}\"\n for t in tools\n ])\n \n # --- build/compile/get_agent ---\n def build(self) -> Optional[StateGraph]:\n \"\"\"构建 Agent 的 LangGraph 工作流图定义。子类应实现。\"\"\"\n raise NotImplementedError(\"Subclasses must implement build() or override compile() directly.\")\n\n def compile(self) -> CompiledGraph:\n \"\"\"编译 Agent 工作流。\"\"\"\n if self._compiled_agent is not None:\n return self._compiled_agent\n\n # 尝试调用 build() 来获取 StateGraph\n workflow = self.build()\n\n if workflow is None or not isinstance(workflow, StateGraph):\n # 如果 build() 不返回 StateGraph (例如 ReactAgent),\n # 子类的 compile() 需要被覆盖以处理编译\n raise ValueError(\n f\"Agent '{self.name}': build() did not return a valid StateGraph, \"\n \"and compile() was not overridden to handle direct compilation.\"\n )\n\n print(f\"Compiling graph for agent: {self.name}\")\n try:\n # 编译 StateGraph 并存储结果\n self._compiled_agent = workflow.compile(\n checkpointer=self.checkpointer,\n debug=getattr(self, 'debug', False) # 传递 debug 标志\n )\n print(f\"Graph compiled successfully for agent: {self.name}\")\n return self._compiled_agent\n except Exception as e:\n print(f\"!!! Error compiling graph for agent {self.name}: {e}\")\n import traceback\n traceback.print_exc()\n raise e\n\n def get_agent(self) -> CompiledGraph:\n \"\"\"获取编译后的核心图实例，如果未编译则先编译。\"\"\"\n if self._compiled_agent is None:\n print(f\"Agent '{self.name}' not compiled yet. Compiling now.\")\n self.compile()\n if self._compiled_agent is None:\n raise RuntimeError(f\"Failed to get compiled agent for '{self.name}'.\")\n return self._compiled_agent\n \n # --- invoke/ainvoke: 标准入口点，调用编译后的图 ---\n def invoke(self, state: Dict[str, Any], config: Optional[RunnableConfig] = None) -> Dict[str, Any]:\n \"\"\"同步调用编译后的 Agent 图。\"\"\"\n try:\n compiled_agent = self.get_agent() # 获取 (或编译) 图\n print(f\"--- Invoking Agent: {self.name} ---\")\n # 直接调用编译后的图，预处理由图内部的 prompt callable 处理 (如果使用 ReactAgent)\n # 或由 Supervisor 节点逻辑处理 (如果使用自定义 Supervisor)\n result = compiled_agent.invoke(state, config=config)\n print(f\"--- Agent Invocation Complete: {self.name} ---\")\n return cast(Dict[str, Any], result) # 假设返回字典\n except Exception as e:\n print(f\"!!! Error during {self.name} agent invocation: {e}\")\n import traceback\n traceback.print_exc()\n # 返回带错误标记的状态 (可能是输入状态)\n state[\"error\"] = f\"Agent invocation failed: {e}\"\n return state\n\n async def ainvoke(self, state: Dict[str, Any], config: Optional[RunnableConfig] = None) -> Dict[str, Any]:\n \"\"\"异步调用编译后的 Agent 图。\"\"\"\n try:\n compiled_agent = self.get_agent() # 获取 (或编译) 图\n print(f\"--- Invoking Agent Async: {self.name} ---\")\n # 直接调用编译后的图\n result = await compiled_agent.ainvoke(state, config=config)\n print(f\"--- Agent Invocation Complete Async: {self.name} ---\")\n return cast(Dict[str, Any], result) # 假设返回字典\n except Exception as e:\n print(f\"!!! Error during {self.name} agent async invocation: {e}\")\n import traceback\n traceback.print_exc()\n state[\"error\"] = f\"Agent async invocation failed: {e}\"\n return state\n\n def run(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"Run the supervisor workflow synchronously.\n\n Args:\n state: The input state for the workflow\n\n Returns:\n The output state from the workflow\n \"\"\"\n return self.invoke(state)\n \n async def arun(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"Run the supervisor workflow asynchronously.\n Args:\n state: The input state for the workflow\n Returns:\n The output state from the workflow\n \"\"\"\n return await self.ainvoke(state)\n\n def reset(self):\n \"\"\"重置编译状态，强制下次重新编译。\"\"\"\n print(f\"Resetting compiled graph for agent '{self.name}'. Will recompile on next use.\")\n self._compiled_agent = None\n self._workflow = None\n\n def add_tools(self, tools: List[Union[BaseTool, Callable]]) -> None:\n \"\"\"添加工具到 Agent 的工具列表。\"\"\"\n print(f\"Warning: Adding tools to {self.name} post-initialization. Agent needs recompilation.\")\n self.tools.extend(tools)\n self.reset()\n"
},
{
"path": "core/agents/base/create_react_agent_wrapper.py",
"content": "import logging\nfrom typing import Optional, Callable, Dict\nfrom langgraph.utils.runnable import RunnableCallable\nfrom langchain_core.runnables.config import RunnableConfig\n\nlogger = logging.getLogger(__name__)\n\nclass CreateReactAgentWrapper(RunnableCallable):\n def __init__(\n self, \n agent, \n name: str = \"agent\", \n before_invoke: Optional[Callable[[dict], dict]] = None,\n before_ainvoke: Optional[Callable[[dict], dict]] = None,\n after_invoke: Optional[Callable[[dict, dict], None]] = None,\n after_ainvoke: Optional[Callable[[dict, dict], None]] = None\n ):\n \"\"\"\n :param agent: The underlying compiled graph or runnable\n :param name: Unique name for this wrapper (avoid duplicates)\n :param before_invoke: A sync callback that modifies the state before the wrapped agent call\n :param before_ainvoke: An async callback that modifies the state before the wrapped agent call\n :param after_invoke: A sync callback that inspects (state, output) after the wrapped call\n :param after_ainvoke: An async callback that inspects (state, output) after the wrapped call\n \"\"\"\n self._agent = agent\n self.name = name or getattr(agent, \"name\", \"agent\")\n self.before_invoke = before_invoke\n self.after_invoke = after_invoke\n self.before_ainvoke = before_ainvoke\n self.after_ainvoke = after_ainvoke\n\n # We define the sync/async \"call\" functions for RunnableCallable\n def call(state: Dict, config: Optional[RunnableConfig] = None, **kwargs) -> Dict:\n logger.info(f\"[{self.name}] (sync) call() - started. State keys: {list(state.keys())}\")\n # Or use print if you prefer\n # print(f\"🟢 [Sync] Invoking wrapper: {self.name}, state keys: {list(state.keys())}\")\n\n # before_invoke callback\n if self.before_invoke:\n state = self.before_invoke(state)\n\n # Call the underlying agent\n output = self._agent.invoke(state, config, **kwargs)\n\n # after_invoke callback\n if self.after_invoke:\n self.after_invoke(state, output)\n\n logger.info(f\"[{self.name}] (sync) call() - finished. Output keys: {list(output.keys())}\")\n return output\n\n async def acall(state: Dict, config: Optional[RunnableConfig] = None, **kwargs) -> Dict:\n logger.info(f\"[{self.name}] (async) acall() - started. State keys: {list(state.keys())}\")\n # print(f\"🟢 [Async] Invoking wrapper: {self.name}, state keys: {list(state.keys())}\")\n\n if self.before_ainvoke:\n state = await self.before_ainvoke(state)\n\n output = await self._agent.ainvoke(state, config, **kwargs)\n\n if self.after_ainvoke:\n await self.after_ainvoke(state, output)\n\n logger.info(f\"[{self.name}] (async) acall() - finished. Output keys: {list(output.keys())}\")\n return output\n\n # Pass these to RunnableCallable\n super().__init__(call, acall, name=self.name)"
},
{
"path": "core/agents/base/react_agent.py",
"content": "from typing import Any, Callable, Dict, List, Optional, Type, Union, Literal, Sequence\n\nfrom langchain_core.language_models import LanguageModelLike, LanguageModelInput\nfrom langchain_core.tools import BaseTool\nfrom langgraph.graph import StateGraph\nfrom langgraph.graph.graph import CompiledGraph\nfrom langgraph.types import Checkpointer\nfrom langgraph.store.base import BaseStore\nfrom langchain_core.messages import BaseMessage, SystemMessage # 导入 SystemMessage\nfrom langgraph.prebuilt import create_react_agent\nfrom langgraph.prebuilt.chat_agent_executor import (\n AgentState,\n StateSchemaType,\n StructuredResponseSchema,\n)\nfrom core.agents.base.base_agent import BaseAgent\nimport logging\nlogger = logging.getLogger(__name__)\n\nclass ReactAgent(BaseAgent):\n \"\"\"ReAct Agent class for reasoning and acting with tools.\n \n This class provides a high-level interface for creating a ReAct agent workflow\n that can perform multi-step reasoning and tool calling.\n \"\"\"\n \n def __init__(\n self,\n model: LanguageModelLike,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[str] = None,\n response_format: Optional[\n Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]\n ] = None,\n state_schema: StateSchemaType = AgentState,\n config_schema: Type[Any] = None,\n checkpointer: Optional[Checkpointer] = None,\n store: Optional[BaseStore] = None,\n interrupt_before: Optional[List[str]] = None,\n interrupt_after: Optional[List[str]] = None,\n debug: bool = False,\n version: Literal[\"v1\", \"v2\"] = \"v1\",\n name: str = \"react_agent\",\n description: str = \"ReAct agent for reasoning and acting with tools.\",\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = None,\n model_name: Optional[str] = \"gpt-4o-mini\",\n ):\n \"\"\"Initialize a ReAct agent.\n \n Args:\n model: Language model to use for the agent\n tools: Optional list of tools available to the agent\n prompt: Optional prompt to use for the agent\n response_format: Optional schema for structured output\n state_schema: State schema to use for the agent graph\n config_schema: Optional schema for configuration\n interrupt_before: Optional list of nodes to interrupt before execution\n interrupt_after: Optional list of nodes to interrupt after execution\n debug: Whether to enable debug mode\n version: Version of the ReAct agent (\"v1\" or \"v2\")\n name: Name of the agent\n max_context_messages: Optional limit on number of recent messages\n max_context_tokens: Optional limit on total estimated tokens\n model_name: Optional model name for token estimation\n \"\"\"\n # Call BaseAgent's __init__ to initialize parent class attributes\n super().__init__(\n name=name,\n model=model,\n tools=tools or [],\n prompt=prompt,\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n model_name=model_name\n )\n \n # Initialize ReactAgent specific attributes\n self.response_format = response_format\n self.react_state_schema = state_schema\n self.react_config_schema = config_schema\n self.react_store = store\n self.react_interrupt_before = interrupt_before\n self.react_interrupt_after = interrupt_after\n self.react_debug = debug\n self.react_version = version\n\n def _prepare_llm_input(self, state: Dict[str, Any]) -> LanguageModelInput:\n \"\"\"\n 准备 LLM 输入：截断消息历史并添加基础 System Prompt (如果存在)。\n 作为 Callable 传递给 create_react_agent 的 prompt 参数。\n \"\"\"\n # 1. 从状态获取消息 (BaseAgent 的方法)\n messages = self._get_state_value(state, \"messages\", [])\n \n # 2. 截断消息 (BaseAgent 的方法)\n # 注意：这里截断的是进入 LLM 前的列表，checkpointer 中的完整历史不受影响\n # --- 添加 Debug 打印 (截断前) ---\n # print(f\"\\nDEBUG _prepare_llm_input ({self.name}): BEFORE truncation (length {len(messages)}):\")\n # for i, msg in enumerate(messages[-5:]): # 只看最后几条\n # print(f\" Msg {i-5}: Type={type(msg).__name__}, ToolCallID={getattr(msg, 'tool_call_id', 'N/A')}\")\n # ---\n\n truncated_messages = self._truncate_messages(messages)\n\n # --- 添加 Debug 打印 (截断后) ---\n # print(f\"DEBUG _prepare_llm_input ({self.name}): AFTER truncation (length {len(truncated_messages)}):\")\n # for i, msg in enumerate(truncated_messages[-5:]): # 只看最后几条\n # print(f\" Msg {i-5}: Type={type(msg).__name__}, ToolCallID={getattr(msg, 'tool_call_id', 'N/A')}\")\n # ---\n \n # 3. 添加基础 System Prompt (如果存在)\n final_messages: List[BaseMessage] = []\n if self.base_prompt:\n if isinstance(self.base_prompt, str):\n final_messages.append(SystemMessage(content=self.base_prompt))\n elif isinstance(self.base_prompt, SystemMessage):\n final_messages.append(self.base_prompt)\n # 如果 self.base_prompt 是其他 Runnable 或 Callable，需要相应处理\n # 但 create_react_agent 的 prompt 通常是 str 或 SystemMessage\n \n final_messages.extend(truncated_messages)\n \n # print(f\"DEBUG [{self.name}]: Preparing LLM input with {len(final_messages)} messages.\") # Optional debug log\n # 返回最终的消息列表给 LLM\n return final_messages\n \n def build(self) -> Optional[StateGraph]:\n \"\"\"对于 ReactAgent，核心图由 create_react_agent 直接创建，无需 build。\"\"\"\n print(f\"Note: ReactAgent '{self.name}' uses create_react_agent in compile(). Build returns None.\")\n self._workflow = None\n return None\n \n def compile(self) -> CompiledGraph:\n \"\"\"使用 create_react_agent 构建并编译核心 ReAct 工作流，存储在 _compiled_agent。\"\"\"\n if self._compiled_agent is not None:\n return self._compiled_agent\n\n print(f\"[[DEBUG]] Compiling core ReAct agent for: {self.name} using create_react_agent\")\n try:\n # 使用 create_react_agent 创建编译后的图\n # 将 self._prepare_llm_input 作为 prompt callable 传入\n compiled_agent = create_react_agent(\n model=self.model,\n tools=self.tools,\n prompt=self._prepare_llm_input, # <--- 关键改动：传入准备函数\n state_schema=self.react_state_schema,\n config_schema=self.react_config_schema,\n checkpointer=self.checkpointer,\n store=self.react_store,\n interrupt_before=self.react_interrupt_before,\n interrupt_after=self.react_interrupt_after,\n debug=self.react_debug,\n version=self.react_version,\n name=self.name,\n )\n # 存储编译好的图\n self._compiled_agent = compiled_agent\n print(f\"Core ReAct graph compiled successfully for agent: {self.name}\")\n return self._compiled_agent\n except Exception as e:\n print(f\"!!! Error compiling graph for agent {self.name} using create_react_agent: {e}\")\n import traceback\n traceback.print_exc()\n self._compiled_agent = None\n raise e"
},
{
"path": "core/agents/react_based_supervisor/__init__.py",
"content": "# 从当前目录导入create_supervisor函数\nfrom .supervisor import create_supervisor\n\n__all__ = [\"create_supervisor\"]\n"
},
{
"path": "core/agents/react_based_supervisor/agent_name.py",
"content": "import re\nfrom typing import Literal\n\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.messages import AIMessage, BaseMessage\nfrom langchain_core.runnables import RunnableLambda\n\nNAME_PATTERN = re.compile(r\"(.*?)\", re.DOTALL)\nCONTENT_PATTERN = re.compile(r\"(.*?)\", re.DOTALL)\n\nAgentNameMode = Literal[\"inline\"]\n\n\ndef _is_content_blocks_content(content: list[dict] | str) -> bool:\n return (\n isinstance(content, list)\n and len(content) > 0\n and isinstance(content[0], dict)\n and \"type\" in content[0]\n )\n\n\ndef add_inline_agent_name(message: BaseMessage) -> BaseMessage:\n \"\"\"Add name and content XML tags to the message content.\n\n Examples:\n\n >>> add_inline_agent_name(AIMessage(content=\"Hello\", name=\"assistant\"))\n AIMessage(content=\"assistantHello\", name=\"assistant\")\n\n >>> add_inline_agent_name(AIMessage(content=[{\"type\": \"text\", \"text\": \"Hello\"}], name=\"assistant\"))\n AIMessage(content=[{\"type\": \"text\", \"text\": \"assistantHello\"}], name=\"assistant\")\n \"\"\"\n if not isinstance(message, AIMessage) or not message.name:\n return message\n\n formatted_message = message.model_copy()\n if _is_content_blocks_content(formatted_message.content):\n text_blocks = [block for block in message.content if block[\"type\"] == \"text\"]\n non_text_blocks = [block for block in message.content if block[\"type\"] != \"text\"]\n content = text_blocks[0][\"text\"] if text_blocks else \"\"\n formatted_content = f\"{message.name}{content}\"\n formatted_message.content = non_text_blocks + [{\"type\": \"text\", \"text\": formatted_content}]\n else:\n formatted_message.content = (\n f\"{message.name}{formatted_message.content}\"\n )\n return formatted_message\n\n\ndef remove_inline_agent_name(message: BaseMessage) -> BaseMessage:\n \"\"\"Remove explicit name and content XML tags from the AI message content.\n\n Examples:\n\n >>> remove_inline_agent_name(AIMessage(content=\"assistantHello\", name=\"assistant\"))\n AIMessage(content=\"Hello\", name=\"assistant\")\n\n >>> remove_inline_agent_name(AIMessage(content=[{\"type\": \"text\", \"text\": \"assistantHello\"}], name=\"assistant\"))\n AIMessage(content=[{\"type\": \"text\", \"text\": \"Hello\"}], name=\"assistant\")\n \"\"\"\n if not isinstance(message, AIMessage) or not message.name:\n return message\n\n is_content_blocks_content = _is_content_blocks_content(message.content)\n if is_content_blocks_content:\n text_blocks = [block for block in message.content if block[\"type\"] == \"text\"]\n if not text_blocks:\n return message\n\n non_text_blocks = [block for block in message.content if block[\"type\"] != \"text\"]\n content = text_blocks[0][\"text\"]\n else:\n content = message.content\n\n name_match: re.Match | None = NAME_PATTERN.search(content)\n content_match: re.Match | None = CONTENT_PATTERN.search(content)\n if not name_match or not content_match:\n return message\n\n if name_match.group(1) != message.name:\n return message\n\n parsed_content = content_match.group(1)\n parsed_message = message.model_copy()\n if is_content_blocks_content:\n content_blocks = non_text_blocks\n if parsed_content:\n content_blocks.append({\"type\": \"text\", \"text\": parsed_content})\n\n parsed_message.content = content_blocks\n else:\n parsed_message.content = parsed_content\n return parsed_message\n\n\ndef with_agent_name(\n model: LanguageModelLike,\n agent_name_mode: AgentNameMode,\n) -> LanguageModelLike:\n \"\"\"Attach formatted agent names to the messages passed to and from a language model.\n\n This is useful for making a message history with multiple agents more coherent.\n\n NOTE: agent name is consumed from the message.name field.\n If you're using an agent built with create_react_agent, name is automatically set.\n If you're building a custom agent, make sure to set the name on the AI message returned by the LLM.\n\n Args:\n model: Language model to add agent name formatting to.\n agent_name_mode: Use to specify how to expose the agent name to the LLM.\n - \"inline\": Add the agent name directly into the content field of the AI message using XML-style tags.\n Example: \"How can I help you\" -> \"agent_nameHow can I help you?\".\n \"\"\"\n if agent_name_mode == \"inline\":\n process_input_message = add_inline_agent_name\n process_output_message = remove_inline_agent_name\n else:\n raise ValueError(\n f\"Invalid agent name mode: {agent_name_mode}. Needs to be one of: {AgentNameMode.__args__}\"\n )\n\n def process_input_messages(messages: list[BaseMessage]) -> list[BaseMessage]:\n return [process_input_message(message) for message in messages]\n\n model = (\n process_input_messages\n | model\n | RunnableLambda(process_output_message, name=\"process_output_message\")\n )\n return model\n"
},
{
"path": "core/agents/react_based_supervisor/handoff.py",
"content": "import re\nimport uuid\n\nfrom langchain_core.messages import AIMessage, ToolCall, ToolMessage\nfrom langchain_core.tools import BaseTool, InjectedToolCallId, tool\nfrom langgraph.prebuilt import InjectedState\nfrom langgraph.types import Command\nfrom typing_extensions import Annotated\n\nWHITESPACE_RE = re.compile(r\"\\s+\")\n\n\ndef _normalize_agent_name(agent_name: str) -> str:\n \"\"\"Normalize an agent name to be used inside the tool name.\"\"\"\n return WHITESPACE_RE.sub(\"_\", agent_name.strip()).lower()\n\n\ndef create_handoff_tool(*, agent_name: str) -> BaseTool:\n \"\"\"Create a tool that can handoff control to the requested agent.\n\n Args:\n agent_name: The name of the agent to handoff control to, i.e.\n the name of the agent node in the multi-agent graph.\n Agent names should be simple, clear and unique, preferably in snake_case,\n although you are only limited to the names accepted by LangGraph\n nodes as well as the tool names accepted by LLM providers\n (the tool name will look like this: `transfer_to_`).\n \"\"\"\n tool_name = f\"transfer_to_{_normalize_agent_name(agent_name)}\"\n\n @tool(tool_name)\n def handoff_to_agent(\n state: Annotated[dict, InjectedState],\n tool_call_id: Annotated[str, InjectedToolCallId],\n ):\n \"\"\"Ask another agent for help.\"\"\"\n tool_message = ToolMessage(\n content=f\"Successfully transferred to {agent_name}\",\n name=tool_name,\n tool_call_id=tool_call_id,\n )\n return Command(\n goto=agent_name,\n graph=Command.PARENT,\n update={\"messages\": state[\"messages\"] + [tool_message]},\n )\n\n return handoff_to_agent\n\n\ndef create_handoff_back_messages(\n agent_name: str, supervisor_name: str\n) -> tuple[AIMessage, ToolMessage]:\n \"\"\"Create a pair of (AIMessage, ToolMessage) to add to the message history when returning control to the supervisor.\"\"\"\n tool_call_id = str(uuid.uuid4())\n tool_name = f\"transfer_back_to_{_normalize_agent_name(supervisor_name)}\"\n tool_calls = [ToolCall(name=tool_name, args={}, id=tool_call_id)]\n return (\n AIMessage(\n content=f\"Transferring back to {supervisor_name}\",\n tool_calls=tool_calls,\n name=agent_name,\n ),\n ToolMessage(\n content=f\"Successfully transferred back to {supervisor_name}\",\n name=tool_name,\n tool_call_id=tool_call_id,\n ),\n )\n"
},
{
"path": "core/agents/react_based_supervisor/planning_handler.py",
"content": "import uuid\nimport datetime\nfrom typing import List, Dict, Optional\n\nclass PlanningStateHandler:\n \"\"\"\n Manages a project plan.\n A plan is a dict with:\n - title (str)\n - description (str)\n - status (str): \"planning\", \"in_progress\", or \"completed\"\n - tasks (list): each task is a dict with:\n id, description, status, agent, notes, evaluation\n - current_task_id (str or None)\n - created_at (str)\n - updated_at (str)\n \"\"\"\n\n @staticmethod\n def _now() -> str:\n return datetime.datetime.now().isoformat()\n\n @staticmethod\n def _gen_id() -> str:\n return str(uuid.uuid4())\n\n @staticmethod\n def create_plan(title: str, description: str) -> Dict:\n now = PlanningStateHandler._now()\n return {\n \"title\": title,\n \"description\": description,\n \"status\": \"planning\", # initial status\n \"tasks\": [],\n \"current_task_id\": None,\n \"created_at\": now,\n \"updated_at\": now\n }\n\n @staticmethod\n def create_task(description: str,\n status: str = \"pending\",\n agent: str = \"\",\n notes: str = \"\",\n evaluation: str = \"\") -> Dict:\n return {\n \"id\": PlanningStateHandler._gen_id(),\n \"description\": description.strip(),\n \"status\": status.strip() if status else \"pending\",\n \"agent\": agent.strip(),\n \"notes\": notes.strip(),\n \"evaluation\": evaluation.strip()\n }\n\n @staticmethod\n def add_tasks(plan: Dict, tasks_data: List[Dict]) -> Dict:\n for tinfo in tasks_data:\n desc = tinfo.get(\"description\", \"Untitled Task\")\n status = tinfo.get(\"status\", \"pending\")\n agent = tinfo.get(\"agent\", \"\")\n notes = tinfo.get(\"notes\", \"\")\n eval_ = tinfo.get(\"evaluation\", \"\")\n task = PlanningStateHandler.create_task(desc, status, agent, notes, eval_)\n plan[\"tasks\"].append(task)\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan\n\n @staticmethod\n def update_task(plan: Dict,\n by_id: Optional[str] = None,\n new_desc: Optional[str] = None,\n new_status: Optional[str] = None,\n new_agent: Optional[str] = None,\n new_notes: Optional[str] = None,\n new_evaluation: Optional[str] = None) -> Dict:\n \"\"\"\n Update a task identified by by_id.\n \"\"\"\n if not by_id:\n raise ValueError(\"Must provide 'by_id' to update a task.\")\n task = next((t for t in plan[\"tasks\"] if t[\"id\"] == by_id), None)\n if not task:\n raise ValueError(\"No matching task found with the given ID.\")\n\n if new_desc is not None:\n task[\"description\"] = new_desc.strip()\n if new_status is not None:\n task[\"status\"] = new_status.strip()\n if new_agent is not None:\n task[\"agent\"] = new_agent.strip()\n if new_notes is not None:\n task[\"notes\"] = new_notes.strip()\n if new_evaluation is not None:\n task[\"evaluation\"] = new_evaluation.strip()\n\n plan[\"updated_at\"] = PlanningStateHandler._now()\n\n # Determine overall plan status\n if any(t[\"status\"] == \"in_progress\" for t in plan[\"tasks\"]):\n plan[\"status\"] = \"in_progress\"\n if all(t[\"status\"] == \"completed\" for t in plan[\"tasks\"]) and plan[\"tasks\"]:\n plan[\"status\"] = \"completed\"\n\n return plan\n\n @staticmethod\n def set_current_task(plan: Dict, task_id: str) -> Dict:\n found = any(t[\"id\"] == task_id for t in plan[\"tasks\"])\n if not found:\n raise ValueError(\"Task ID not found in plan.\")\n plan[\"current_task_id\"] = task_id\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan\n\n @staticmethod\n def finish_plan(plan: Dict) -> Dict:\n \"\"\"\n Forcefully mark the plan as completed.\n \"\"\"\n plan[\"status\"] = \"completed\"\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan"
},
{
"path": "core/agents/react_based_supervisor/simple_planning_tool.py",
"content": "import json\nfrom typing import Dict, List, Optional\nfrom langchain_core.tools import BaseTool\nfrom core.agents.supervisor.planning_handler import PlanningStateHandler\n\nclass SimplePlanningTool(BaseTool):\n \"\"\"\n A tool that manages a single project plan in memory.\n It supports creating, viewing, adding tasks, updating tasks, setting the current task,\n and finishing the plan. All operations return a JSON string.\n \"\"\"\n name: str = \"planning\"\n description: str = (\"Manage a project plan with actions to create, view, add tasks, update tasks, \"\n \"set current task, and finish the plan. All data is stored in JSON.\")\n\n def __init__(self):\n super().__init__()\n self._plan: Optional[Dict] = None\n\n def _run(self, action: str, **kwargs) -> str:\n try:\n if action == \"create_plan\":\n return self._handle_create_plan(**kwargs)\n elif action == \"view_plan\":\n return self._handle_view_plan()\n elif action == \"add_tasks\":\n return self._handle_add_tasks(**kwargs)\n elif action == \"update_task\":\n return self._handle_update_task(**kwargs)\n elif action == \"set_current_task\":\n return self._handle_set_current_task(**kwargs)\n elif action == \"finish_plan\":\n return self._handle_finish_plan()\n else:\n return self._json_error(f\"Unknown action: {action}\")\n except Exception as e:\n return self._json_error(str(e))\n\n async def _arun(self, action: str, **kwargs) -> str:\n return self._run(action, **kwargs)\n\n def _handle_create_plan(self, **kwargs) -> str:\n title = kwargs.get(\"title\", \"Untitled Plan\")\n description = kwargs.get(\"description\", \"\")\n tasks_data = kwargs.get(\"tasks\", [])\n new_plan = PlanningStateHandler.create_plan(title, description)\n PlanningStateHandler.add_tasks(new_plan, tasks_data)\n self._plan = new_plan\n return self._json_ok(self._plan)\n\n def _handle_view_plan(self) -> str:\n if not self._plan:\n self._plan = PlanningStateHandler.create_plan(\"Untitled\", \"\")\n return self._json_ok(self._plan)\n\n def _handle_add_tasks(self, **kwargs) -> str:\n if not self._plan:\n self._plan = PlanningStateHandler.create_plan(\"Untitled\", \"\")\n tasks_data = kwargs.get(\"tasks\", [])\n PlanningStateHandler.add_tasks(self._plan, tasks_data)\n return self._json_ok(self._plan)\n\n def _handle_update_task(self, **kwargs) -> str:\n if not self._plan:\n raise ValueError(\"No plan exists. Please create a plan first.\")\n # Use 'by_id' instead of 'task_id'\n by_id = kwargs.get(\"by_id\")\n new_desc = kwargs.get(\"description\")\n new_status = kwargs.get(\"status\")\n new_agent = kwargs.get(\"agent\")\n new_notes = kwargs.get(\"notes\")\n new_evaluation = kwargs.get(\"evaluation\")\n updated = PlanningStateHandler.update_task(\n self._plan,\n by_id=by_id,\n new_desc=new_desc,\n new_status=new_status,\n new_agent=new_agent,\n new_notes=new_notes,\n new_evaluation=new_evaluation\n )\n self._plan = updated\n return self._json_ok(self._plan)\n\n def _handle_set_current_task(self, **kwargs) -> str:\n if not self._plan:\n raise ValueError(\"No plan available to set current task.\")\n tid = kwargs.get(\"task_id\")\n if not tid:\n raise ValueError(\"Must provide 'task_id' for set_current_task.\")\n PlanningStateHandler.set_current_task(self._plan, tid)\n return self._json_ok(self._plan)\n\n def _handle_finish_plan(self) -> str:\n if not self._plan:\n raise ValueError(\"No plan exists to finish.\")\n finished_plan = PlanningStateHandler.finish_plan(self._plan)\n self._plan = finished_plan\n return self._json_ok(finished_plan)\n\n def _json_ok(self, plan_data: Dict) -> str:\n return json.dumps({\"ok\": True, \"plan\": plan_data}, ensure_ascii=False, indent=2)\n\n def _json_error(self, message: str) -> str:\n return json.dumps({\"ok\": False, \"error\": message}, ensure_ascii=False, indent=2)"
},
{
"path": "core/agents/react_based_supervisor/state_schema.py",
"content": "from typing import Dict, List, Optional, Any, Literal, TypedDict, Union\nfrom langchain_core.messages import BaseMessage\nfrom langgraph.prebuilt.chat_agent_executor import AgentState\n\n# 定义计划状态类型\nPlanningStatus = Literal[\"not_started\", \"planning\", \"executing\", \"completed\", \"failed\"]\n\n# 定义任务状态类型\nTaskStatus = Literal[\"pending\", \"in_progress\", \"completed\", \"failed\"]\n\n# 定义任务项\nclass Task(TypedDict, total=False):\n \"\"\"任务项定义\n \n 表示计划中的一个任务项，包含任务描述、状态、分配的代理等信息\n \"\"\"\n id: str # 任务唯一标识符\n description: str # 任务描述\n status: TaskStatus # 任务状态\n agent: Optional[str] # 分配的代理名称\n created_at: str # 创建时间\n updated_at: str # 更新时间\n completed_at: Optional[str] # 完成时间\n dependencies: Optional[List[str]] # 依赖的任务ID列表\n notes: Optional[str] # 任务备注\n\n# 定义计划\nclass Plan(TypedDict, total=False):\n \"\"\"计划定义\n \n 表示一个完整的计划，包含计划状态、任务列表等信息\n \"\"\"\n status: PlanningStatus # 计划状态\n tasks: List[Task] # 任务列表\n current_task_id: Optional[str] # 当前执行的任务ID\n created_at: str # 创建时间\n updated_at: str # 更新时间\n completed_at: Optional[str] # 完成时间\n title: Optional[str] # 计划标题\n description: Optional[str] # 计划描述\n\n# 扩展AgentState以支持计划功能\nclass PlanningAgentState(AgentState):\n \"\"\"支持计划功能的代理状态\n \n 扩展了AgentState，添加了plan字段用于存储计划信息\n \"\"\"\n plan: Optional[Plan] = None"
},
{
"path": "core/agents/react_based_supervisor/supervisor.py",
"content": "import inspect\nfrom typing import Any, Callable, Literal, Optional, Type, Union, Dict, Optional\n\nfrom langchain_core.language_models import BaseChatModel, LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langgraph.graph import END, START, StateGraph\nfrom langgraph.prebuilt.chat_agent_executor import (\n create_react_agent,\n AgentState,\n Prompt,\n StateSchemaType,\n StructuredResponseSchema,\n)\nfrom langgraph.pregel import Pregel\nfrom langgraph.utils.runnable import RunnableCallable\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.agents.supervisor.agent_name import AgentNameMode, with_agent_name\nfrom core.agents.supervisor.handoff import (\n create_handoff_back_messages,\n create_handoff_tool,\n)\nOutputMode = Literal[\"full_history\", \"last_message\"]\n\"\"\"Mode for adding agent outputs to the message history in the multi-agent workflow\n\n- `full_history`: add the entire agent message history\n- `last_message`: add only the last message\n\"\"\"\n\n\nMODELS_NO_PARALLEL_TOOL_CALLS = {\"o3-mini\"}\n\n\ndef _supports_disable_parallel_tool_calls(model: LanguageModelLike) -> bool:\n if not isinstance(model, BaseChatModel):\n return False\n\n if hasattr(model, \"model_name\") and model.model_name in MODELS_NO_PARALLEL_TOOL_CALLS:\n return False\n\n if not hasattr(model, \"bind_tools\"):\n return False\n\n if \"parallel_tool_calls\" not in inspect.signature(model.bind_tools).parameters:\n return False\n\n return True\n\n\ndef _make_call_agent(\n agent: Pregel,\n output_mode: OutputMode,\n add_handoff_back_messages: bool,\n supervisor_name: str,\n) -> Callable[[dict], dict] | RunnableCallable:\n if output_mode not in OutputMode.__args__:\n raise ValueError(\n f\"Invalid agent output mode: {output_mode}. Needs to be one of {OutputMode.__args__}\"\n )\n\n def _process_output(output: dict) -> dict:\n messages = output[\"messages\"]\n if output_mode == \"full_history\":\n pass\n elif output_mode == \"last_message\":\n messages = messages[-1:]\n else:\n raise ValueError(\n f\"Invalid agent output mode: {output_mode}. \"\n f\"Needs to be one of {OutputMode.__args__}\"\n )\n\n if add_handoff_back_messages:\n messages.extend(create_handoff_back_messages(agent.name, supervisor_name))\n\n return {\n **output,\n \"messages\": messages,\n }\n\n def call_agent(state: dict) -> dict:\n #print(f\"🟡 [Sync invoke] Handoff to agent '{agent.name}' with state keys: {list(state.keys())}\")\n output = agent.invoke(state)\n #print(f\"✅ [Sync invoke] Agent '{agent.name}' completed.\")\n return _process_output(output)\n\n async def acall_agent(state: dict) -> dict:\n #print(f\"🟡 [Async invoke] Handoff to agent '{agent.name}' with state keys: {list(state.keys())}\")\n output = await agent.ainvoke(state)\n #print(f\"✅ [Async invoke] Agent '{agent.name}' completed.\")\n return _process_output(output)\n\n return RunnableCallable(call_agent, acall_agent)\n\n\ndef create_supervisor(\n agents: list[Pregel],\n *,\n model: LanguageModelLike,\n tools: list[BaseTool | Callable] | None = None,\n prompt: Prompt | None = None,\n response_format: Optional[\n Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]\n ] = None,\n state_schema: StateSchemaType = AgentState,\n config_schema: Type[Any] | None = None,\n output_mode: OutputMode = \"last_message\",\n add_handoff_back_messages: bool = True,\n supervisor_name: str = \"supervisor\",\n include_agent_name: AgentNameMode | None = None,\n) -> StateGraph:\n \"\"\"Create a multi-agent supervisor.\n\n Args:\n agents: List of agents to manage\n model: Language model to use for the supervisor\n tools: Tools to use for the supervisor\n prompt: Optional prompt to use for the supervisor. Can be one of:\n - str: This is converted to a SystemMessage and added to the beginning of the list of messages in state[\"messages\"].\n - SystemMessage: this is added to the beginning of the list of messages in state[\"messages\"].\n - Callable: This function should take in full graph state and the output is then passed to the language model.\n - Runnable: This runnable should take in full graph state and the output is then passed to the language model.\n response_format: An optional schema for the final supervisor output.\n\n If provided, output will be formatted to match the given schema and returned in the 'structured_response' state key.\n If not provided, `structured_response` will not be present in the output state.\n Can be passed in as:\n\n - an OpenAI function/tool schema,\n - a JSON Schema,\n - a TypedDict class,\n - or a Pydantic class.\n - a tuple (prompt, schema), where schema is one of the above.\n The prompt will be used together with the model that is being used to generate the structured response.\n\n !!! Important\n `response_format` requires the model to support `.with_structured_output`\n\n !!! Note\n `response_format` requires `structured_response` key in your state schema.\n You can use the prebuilt `langgraph.prebuilt.chat_agent_executor.AgentStateWithStructuredResponse`.\n state_schema: State schema to use for the supervisor graph.\n config_schema: An optional schema for configuration.\n Use this to expose configurable parameters via supervisor.config_specs.\n output_mode: Mode for adding managed agents' outputs to the message history in the multi-agent workflow.\n Can be one of:\n - `full_history`: add the entire agent message history\n - `last_message`: add only the last message (default)\n add_handoff_back_messages: Whether to add a pair of (AIMessage, ToolMessage) to the message history\n when returning control to the supervisor to indicate that a handoff has occurred.\n supervisor_name: Name of the supervisor node.\n include_agent_name: Use to specify how to expose the agent name to the underlying supervisor LLM.\n\n - None: Relies on the LLM provider using the name attribute on the AI message. Currently, only OpenAI supports this.\n - \"inline\": Add the agent name directly into the content field of the AI message using XML-style tags.\n Example: \"How can I help you\" -> \"agent_nameHow can I help you?\"\n \"\"\"\n agent_names = set()\n for agent in agents:\n if agent.name is None or agent.name == \"LangGraph\":\n raise ValueError(\n \"Please specify a name when you create your agent, either via `create_react_agent(..., name=agent_name)` \"\n \"or via `graph.compile(name=name)`.\"\n )\n\n if agent.name in agent_names:\n raise ValueError(\n f\"Agent with name '{agent.name}' already exists. Agent names must be unique.\"\n )\n\n agent_names.add(agent.name)\n\n handoff_tools = [create_handoff_tool(agent_name=agent.name) for agent in agents]\n all_tools = (tools or []) + handoff_tools\n\n if _supports_disable_parallel_tool_calls(model):\n model = model.bind_tools(all_tools, parallel_tool_calls=False)\n else:\n model = model.bind_tools(all_tools)\n\n if include_agent_name:\n model = with_agent_name(model, include_agent_name)\n \n supervisor = create_react_agent(\n name=supervisor_name,\n model=model,\n tools=all_tools,\n prompt=prompt,\n state_schema=state_schema,\n response_format=response_format,\n debug=False,\n )\n # Build the multi-agent supervisor graph using the langgraph StateGraph setup\n builder = StateGraph(state_schema, config_schema=config_schema)\n builder.add_node(supervisor, destinations=tuple(agent_names) + (END,))\n builder.add_edge(START, supervisor.name)\n for agent in agents:\n # If agent is a \"ReactAgent\" or similar\n if hasattr(agent, \"get_agent\") and callable(agent.get_agent):\n agent = agent.get_agent() # retrieve the compiled subgraph\n \n builder.add_node(\n agent.name,\n _make_call_agent(\n agent,\n output_mode,\n add_handoff_back_messages,\n supervisor_name,\n ),\n )\n builder.add_edge(agent.name, supervisor.name)\n\n return builder\n"
},
{
"path": "core/agents/react_supervisor_agent.py",
"content": "from typing import Any, Callable, Dict, List, Optional, Union\nimport re\n\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langgraph.graph import StateGraph\nfrom langgraph.graph.state import CompiledStateGraph\nfrom langgraph.types import Checkpointer\nfrom langgraph.prebuilt.chat_agent_executor import (\n AgentState,\n StateSchemaType,\n)\nfrom langgraph.utils.runnable import RunnableCallable\nfrom core.agents.react_based_supervisor import create_supervisor\nfrom core.agents.react_based_supervisor.simple_planning_tool import SimplePlanningTool\nfrom core.agents.base.base_agent import BaseAgent\nfrom core.agents.react_based_supervisor.state_schema import PlanningAgentState\nimport logging\n\nlogger = logging.getLogger(__name__)\n\nclass SupervisorAgent(BaseAgent):\n \"\"\"Supervisor class for managing multiple agents with planning capabilities.\n \n This class provides a high-level interface for creating a supervisor workflow\n that can manage and coordinate multiple agents. It also includes planning capabilities\n to create and manage a plan for complex tasks using a state-driven approach.\n \n The planning functionality is implemented using PlanningStateHandler and PlanningTool,\n which provide a more structured and flexible way to manage tasks compared to the\n previous TodolistTool approach.\n \"\"\"\n _PROMPT_TEMPLATE = \"\"\"You are a Supervisor Agent. Your job is to analyze user requests and coordinate multiple agents to complete tasks.\n\n## Task Approach Methodology\n\n### Understanding Requirements\n- Analyzing user requests to identify core needs\n- Asking clarifying questions when requirements are ambiguous\n- Breaking down complex requests into manageable components\n- Identifying potential challenges before beginning work\n\n### Coordination\n- Identifying appropriate agents for each task\n- Delegating tasks to specialized agents\n- Tracking progress and ensuring task completion\n- Synthesizing information from multiple agents\n\nRemember: Effective coordination is essential for successful task completion. Take time to understand the request and delegate appropriately.\n {tools}\n\"\"\"\n\n _PLANNING_PROMPT_TEMPLATE = \"\"\"You are a Supervisor agent. Your role is to analyze user requests, break them down into actionable tasks, and coordinate specialized agents (e.g., research_expert, coder_expert, reporter_expert) to complete them.\n\n# Working with Complex Requests\n1. FIRST, carefully analyze the user's request and break it down into clear, actionable tasks\n2. Identify which agent is best suited for each part of the task\n3. Use the handoff tools to delegate tasks to appropriate agents ONE AT A TIME\n4. WAIT for each agent to COMPLETELY FINISH their assigned task before proceeding\n5. Review the output from each agent before delegating the next task\n6. Maintain a sequential workflow - never delegate multiple tasks simultaneously\n7. Synthesize the results and provide a coherent response to the user\n8. Provide a final summary when all tasks are done\n\"\"\"\n\n _PLANNING_TOOL_TEMPLATE = \"\"\"\n# Planning Tool Instructions\nYou have access to a \"planning\" tool that uses JSON for all operations. Do NOT include any \"state\" field in your calls. Use the following actions exactly as defined:\n\n1. \"create_plan\": Create a new plan.\n - Required fields:\n - title (string)\n - description (string)\n - tasks (list of task objects). Each task object must include:\n \"description\": string,\n \"status\": \"pending\" (all tasks must have \"status\": \"pending\" initially),\n \"agent\": string (empty if not assigned),\n \"notes\": string (empty if none),\n \"evaluation\": string (empty if none)\n - Example:\n {\n \"action\": \"create_plan\",\n \"title\": \"Python Scraper for Tech News\",\n \"description\": \"Build a Python scraper to fetch the latest tech news and save it as CSV\",\n \"tasks\": [\n {\"description\": \"Research Python scraping libraries\", \"status\": \"pending\", \"agent\": \"\", \"notes\": \"\", \"evaluation\": \"\"},\n {\"description\": \"Implement the scraper\", \"status\": \"pending\", \"agent\": \"\", \"notes\": \"\", \"evaluation\": \"\"},\n {\"description\": \"Test the code\", \"status\": \"pending\", \"agent\": \"\", \"notes\": \"\", \"evaluation\": \"\"}\n ]\n }\n\n2. \"view_plan\": Retrieve the current plan.\n - Example:\n {\n \"action\": \"view_plan\"\n }\n\n3. \"add_tasks\": Add additional tasks to the plan.\n - Required:\n - tasks: list of task objects (same format as above)\n - Example:\n {\n \"action\": \"add_tasks\",\n \"tasks\": [\n {\"description\": \"Write documentation\", \"status\": \"pending\", \"agent\": \"\", \"notes\": \"\", \"evaluation\": \"\"}\n ]\n }\n\n4. \"update_task\": Update an existing task.\n - Identify the task by \"by_id\" (the task's unique ID from the plan).\n - You may update any of: \"description\", \"status\", \"agent\", \"notes\", \"evaluation\".\n - Example:\n {\n \"action\": \"update_task\",\n \"by_id\": \"TASK-UUID\",\n \"status\": \"completed\",\n \"evaluation\": \"The scraper works perfectly.\"\n }\n\n5. \"set_current_task\": Set the current task by its ID.\n - Example:\n {\n \"action\": \"set_current_task\",\n \"task_id\": \"TASK-UUID\"\n }\n\n6. \"finish_plan\": Mark the entire plan as completed.\n - Example:\n {\n \"action\": \"finish_plan\"\n }\n\nImportant:\n- Always produce valid JSON for your tool calls.\n- Continuously update and monitor the plan until every task's status is \"completed\" before delivering your final answer.\n- If the plan is not fully completed, do not stop; keep updating the plan with appropriate calls.\n\"\"\"\n def __init__(\n self,\n agents: List[BaseAgent],\n model: LanguageModelLike,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[str] = None,\n state_schema: StateSchemaType = AgentState,\n supervisor_name: str = \"supervisor\",\n checkpointer: Optional[Checkpointer] = None,\n output_mode: str = \"last_message\", # * full_history or last_message *\n enable_planning: bool = True, # * True or False *\n ):\n \"\"\"Initialize a supervisor.\n \n Args:\n agents: List of agents to manage\n model: Language model to use for the supervisor\n tools: Optional list of tools available to the supervisor\n prompt: Optional prompt to use for the supervisor\n state_schema: State schema to use for the supervisor graph\n supervisor_name: Name of the supervisor node\n checkpointer: Optional checkpointer to use for the supervisor\n output_mode: Mode for adding agent outputs to the message history\n (\"full_history\" or \"last_message\")\n enable_planning: Whether to enable planning capabilities\n auto_planning: Whether to automatically generate plans for complex tasks\n \"\"\"\n # 设置规划相关属性\n self._enable_planning = enable_planning\n \n # 如果启用规划功能，设置状态模式为PlanningAgentState\n if self._enable_planning and state_schema == AgentState:\n state_schema = PlanningAgentState\n \n # Store agent-specific attributes before super().__init__\n self.agents = agents\n self.output_mode = output_mode\n self.supervisor_name = supervisor_name\n self.state_schema = state_schema\n self.checkpointer = checkpointer\n self.tools = tools or []\n self._workflow = None\n \n # 生成基础提示词\n # _agents_prompt = self._generate_agents_prompt()\n _final_prompt = self._PLANNING_PROMPT_TEMPLATE + \"/n/n\" + self._PLANNING_TOOL_TEMPLATE if self._enable_planning else self._PROMPT_TEMPLATE\n \n if tools is None:\n tools = []\n # 如果启用规划功能，添加规划提示模板并添加规划工具\n if self._enable_planning:\n tools.append(SimplePlanningTool())\n \n # 初始化BaseAgent父类\n super().__init__(\n name=supervisor_name,\n model=model,\n tools=tools,\n checkpointer=checkpointer,\n prompt=_final_prompt,\n )\n \n def build(self) -> StateGraph:\n \"\"\"Build the supervisor workflow.\n \n Returns:\n The built StateGraph\n \"\"\"\n \n if self._workflow is not None:\n return self._workflow\n \n self._workflow = create_supervisor(\n agents=self.agents,\n model=self.model,\n tools=self.tools,\n prompt=self.base_prompt,\n state_schema=self.state_schema,\n supervisor_name=self.supervisor_name,\n output_mode=self.output_mode,\n )\n \n return self._workflow"
},
{
"path": "core/agents/sb_supervisor_agent.py",
"content": "# reason_graph/supervisor_agent.py\nfrom typing import Callable, List, Optional, Union, cast, Literal\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langgraph.graph import StateGraph\nfrom langgraph.types import Checkpointer\n\n# 内部导入\nfrom core.agents.base.base_agent import BaseAgent\nfrom core.agents.state_based_supervisor.state_schema import PlanningAgentState, StateSchemaType # 导入 PlanningAgentState\n# 导入重构后的 create_supervisor 函数\nfrom core.agents.state_based_supervisor.supervisor_graph import create_supervisor\nfrom core.agents.state_based_supervisor.agent_name import AgentNameMode\n\nimport logging\nlogger = logging.getLogger(__name__)\n\nclass SupervisorAgent(BaseAgent):\n \"\"\"\n Supervisor Agent 类 (最终版)\n 负责协调子 Agent 并管理规划 (使用状态驱动方法)。\n invoke/ainvoke 继承自 BaseAgent，负责完整流程。\n \"\"\"\n\n def __init__(\n self,\n agents: List[BaseAgent], # 子 Agent 实例列表\n model: LanguageModelLike, # Supervisor 使用的 LLM\n tools: Optional[List[Union[BaseTool, Callable]]] = None, # Supervisor 特有工具\n state_schema: StateSchemaType = PlanningAgentState,\n supervisor_name: str = \"supervisor\",\n checkpointer: Optional[Checkpointer] = None,\n output_mode: str = \"last_message\",\n # enable_planning: bool = True, # 不再需要，强制使用 Planning\n include_agent_name: Optional[str] = \"inline\",\n # BaseAgent 参数\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = None,\n model_name: Optional[str] = None,\n ):\n \"\"\"初始化 Supervisor Agent\"\"\"\n if state_schema != PlanningAgentState:\n print(\"Warning: SupervisorAgent forces state_schema to PlanningAgentState.\")\n state_schema = PlanningAgentState\n\n self.sub_agents = agents\n self.output_mode = output_mode\n self.include_agent_name = cast(Optional[AgentNameMode], include_agent_name)\n\n # 初始化 BaseAgent 父类\n super().__init__(\n name=supervisor_name,\n model=model,\n tools=tools or [],\n checkpointer=checkpointer,\n prompt=None, # 核心 Prompt 在 supervisor_node_logic 中处理\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n model_name=model_name,\n )\n # _workflow_definition 和 _executable_agent 由 BaseAgent 管理\n\n def build(self) -> Optional[StateGraph]:\n \"\"\"构建 Supervisor 的 LangGraph 工作流图定义。\"\"\"\n # 调用重构后的 create_supervisor 函数来获取 StateGraph 定义\n # 这个 StateGraph 包含了手写的 supervisor_node_logic\n if self._workflow: return self._workflow\n \n print(f\"Building supervisor graph definition for '{self.name}'...\")\n try:\n graph_definition = create_supervisor(\n model=self.model,\n sub_agents=self.sub_agents,\n state_schema=PlanningAgentState, # 强制使用\n tools=self.tools,\n output_mode=cast(Literal[\"full_history\", \"last_message\"], self.output_mode),\n supervisor_name=self.name,\n include_agent_name=self.include_agent_name,\n )\n self._workflow = graph_definition # 存储图定义\n print(f\"Supervisor graph definition built for '{self.name}'.\")\n return self._workflow\n except Exception as e:\n print(f\"!!! Error building supervisor graph definition '{self.name}': {e}\")\n import traceback\n traceback.print_exc()\n self._workflow = None\n raise e\n\n # compile 方法继承自 BaseAgent\n # 它会调用上面的 build() 获取 StateGraph 定义，然后编译它，\n # 并创建包含预处理步骤的最终 _executable_agent\n\n # invoke, ainvoke, get_agent, reset 继承自 BaseAgent"
},
{
"path": "core/agents/state_based_supervisor/__init__.py",
"content": ""
},
{
"path": "core/agents/state_based_supervisor/agent_name.py",
"content": "import re\nfrom typing import Literal\n\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.messages import AIMessage, BaseMessage\nfrom langchain_core.runnables import RunnableLambda\n\nNAME_PATTERN = re.compile(r\"(.*?)\", re.DOTALL)\nCONTENT_PATTERN = re.compile(r\"(.*?)\", re.DOTALL)\n\nAgentNameMode = Literal[\"inline\"]\n\n\ndef _is_content_blocks_content(content: list[dict] | str) -> bool:\n return (\n isinstance(content, list)\n and len(content) > 0\n and isinstance(content[0], dict)\n and \"type\" in content[0]\n )\n\n\ndef add_inline_agent_name(message: BaseMessage) -> BaseMessage:\n \"\"\"Add name and content XML tags to the message content.\n\n Examples:\n\n >>> add_inline_agent_name(AIMessage(content=\"Hello\", name=\"assistant\"))\n AIMessage(content=\"assistantHello\", name=\"assistant\")\n\n >>> add_inline_agent_name(AIMessage(content=[{\"type\": \"text\", \"text\": \"Hello\"}], name=\"assistant\"))\n AIMessage(content=[{\"type\": \"text\", \"text\": \"assistantHello\"}], name=\"assistant\")\n \"\"\"\n if not isinstance(message, AIMessage) or not message.name:\n return message\n\n formatted_message = message.model_copy()\n if _is_content_blocks_content(formatted_message.content):\n text_blocks = [block for block in message.content if block[\"type\"] == \"text\"]\n non_text_blocks = [block for block in message.content if block[\"type\"] != \"text\"]\n content = text_blocks[0][\"text\"] if text_blocks else \"\"\n formatted_content = f\"{message.name}{content}\"\n formatted_message.content = non_text_blocks + [{\"type\": \"text\", \"text\": formatted_content}]\n else:\n formatted_message.content = (\n f\"{message.name}{formatted_message.content}\"\n )\n return formatted_message\n\n\ndef remove_inline_agent_name(message: BaseMessage) -> BaseMessage:\n \"\"\"Remove explicit name and content XML tags from the AI message content.\n\n Examples:\n\n >>> remove_inline_agent_name(AIMessage(content=\"assistantHello\", name=\"assistant\"))\n AIMessage(content=\"Hello\", name=\"assistant\")\n\n >>> remove_inline_agent_name(AIMessage(content=[{\"type\": \"text\", \"text\": \"assistantHello\"}], name=\"assistant\"))\n AIMessage(content=[{\"type\": \"text\", \"text\": \"Hello\"}], name=\"assistant\")\n \"\"\"\n if not isinstance(message, AIMessage) or not message.name:\n return message\n\n is_content_blocks_content = _is_content_blocks_content(message.content)\n if is_content_blocks_content:\n text_blocks = [block for block in message.content if block[\"type\"] == \"text\"]\n if not text_blocks:\n return message\n\n non_text_blocks = [block for block in message.content if block[\"type\"] != \"text\"]\n content = text_blocks[0][\"text\"]\n else:\n content = message.content\n\n name_match: re.Match | None = NAME_PATTERN.search(content)\n content_match: re.Match | None = CONTENT_PATTERN.search(content)\n if not name_match or not content_match:\n return message\n\n if name_match.group(1) != message.name:\n return message\n\n parsed_content = content_match.group(1)\n parsed_message = message.model_copy()\n if is_content_blocks_content:\n content_blocks = non_text_blocks\n if parsed_content:\n content_blocks.append({\"type\": \"text\", \"text\": parsed_content})\n\n parsed_message.content = content_blocks\n else:\n parsed_message.content = parsed_content\n return parsed_message\n\n\ndef with_agent_name(\n model: LanguageModelLike,\n agent_name_mode: AgentNameMode,\n) -> LanguageModelLike:\n \"\"\"Attach formatted agent names to the messages passed to and from a language model.\n\n This is useful for making a message history with multiple agents more coherent.\n\n NOTE: agent name is consumed from the message.name field.\n If you're using an agent built with create_react_agent, name is automatically set.\n If you're building a custom agent, make sure to set the name on the AI message returned by the LLM.\n\n Args:\n model: Language model to add agent name formatting to.\n agent_name_mode: Use to specify how to expose the agent name to the LLM.\n - \"inline\": Add the agent name directly into the content field of the AI message using XML-style tags.\n Example: \"How can I help you\" -> \"agent_nameHow can I help you?\".\n \"\"\"\n if agent_name_mode == \"inline\":\n process_input_message = add_inline_agent_name\n process_output_message = remove_inline_agent_name\n else:\n raise ValueError(\n f\"Invalid agent name mode: {agent_name_mode}. Needs to be one of: {AgentNameMode.__args__}\"\n )\n\n def process_input_messages(messages: list[BaseMessage]) -> list[BaseMessage]:\n return [process_input_message(message) for message in messages]\n\n model = (\n process_input_messages\n | model\n | RunnableLambda(process_output_message, name=\"process_output_message\")\n )\n return model\n"
},
{
"path": "core/agents/state_based_supervisor/evaluate_result_node.py",
"content": "# reason_graph/evaluate_result_node.py\n\nimport json\nimport time\nimport copy\nimport traceback\nimport anyio \nfrom typing import Dict, Any, List, Optional, Union\nfrom langchain_core.messages import BaseMessage, AIMessage, ToolMessage\nfrom langchain_core.runnables import RunnableConfig\n\n# 内部导入 (确保路径正确)\ntry:\n from .state_schema import PlanningAgentState, TaskStatus, Plan, Task\n from .planning_handler import PlanningStateHandler\nexcept ImportError as e:\n print(f\"Error importing modules in evaluate_result_node.py: {e}\")\n # Fallbacks\n class PlanningAgentState(Dict): pass; \n class Plan(Dict): pass; \n class Task(Dict): pass\n TaskStatus = str \n class PlanningStateHandler: # Dummy\n @staticmethod \n def update_task(plan, by_id, **kwargs): return plan\n @staticmethod\n def set_current_task(plan, task_id): return plan\n @staticmethod\n def get_task(plan, task_id): return None\n @staticmethod\n def update_plan_status(plan): return plan\n\n\nasync def evaluate_result_node_logic(state: PlanningAgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:\n \"\"\"\n 评估子 Agent 返回结果并更新计划状态的节点逻辑 (异步, 优化评估逻辑)。\n \"\"\"\n print(f\"--- Entering Evaluate Result Node ---\")\n messages: List[BaseMessage] = state.get('messages', [])\n plan: Optional[Plan] = state.get('plan')\n last_message = messages[-1] if messages else None\n error_message: Optional[str] = None\n plan_updated: bool = False\n updated_plan: Optional[Plan] = copy.deepcopy(plan) if plan else None \n\n if not updated_plan:\n print(\"Evaluate Result Node: No plan found in state. Skipping.\")\n return {} \n\n current_task_id = updated_plan.get(\"current_task_id\")\n if not current_task_id:\n # Fallback logic for finding current task (不变)\n print(\"Warning: Evaluate Result Node - No current_task_id found in plan...\")\n in_progress_tasks = [t for t in updated_plan.get('tasks', []) if t.get('status') == 'in_progress']\n if len(in_progress_tasks) == 1: current_task_id = in_progress_tasks[0].get('id'); print(f\" Fallback: Found task {current_task_id}\")\n else: error_message = \"Evaluation failed: Cannot determine finished task.\"; print(f\"ERROR: {error_message}\"); return {\"plan\": updated_plan, \"error\": error_message, \"messages\": []}\n\n agent_result_content: Optional[str] = None\n agent_name: Optional[str] = None\n if isinstance(last_message, AIMessage): \n agent_result_content = str(last_message.content) if last_message.content is not None else \"\" # Ensure string\n agent_name = last_message.name or \"SubAgent\"\n print(f\" Evaluating result from: {agent_name} for task ID: {current_task_id}\")\n else:\n agent_result_content = f\"Error: Expected AIMessage result, got {type(last_message).__name__}.\"\n agent_name = \"System/Error\"\n print(f\"Warning: Last message not AIMessage. Assuming task failed for {current_task_id}.\")\n\n\n # --- 优化的评估逻辑 ---\n new_status: TaskStatus = \"completed\" # 默认成功\n evaluation_notes = f\"Result received from {agent_name}.\"\n \n # 1. 检查是否为空内容 (或只有空白符)\n if agent_result_content is None or not agent_result_content.strip():\n new_status = \"failed\"\n evaluation_notes = f\"Task failed: Agent {agent_name} returned empty content.\"\n print(f\" Task {current_task_id} evaluated as FAILED (Empty Result).\")\n # 2. 检查是否以明确的错误标识开头 (需要工具配合)\n # 假设工具出错时会在返回字符串前加上 \"Error: \" 或 \"Execution Failed: \"\n elif agent_result_content.strip().startswith((\"Error:\", \"Execution Failed:\", \"Tool Error:\")):\n new_status = \"failed\"\n evaluation_notes = f\"Task failed: Agent {agent_name} reported an error: {agent_result_content[:150]}...\"\n print(f\" Task {current_task_id} evaluated as FAILED (Explicit Error Signal).\")\n # 3. (可选) 添加其他特定检查，例如检查是否只是\"我不明白\"之类的回复\n elif len(agent_result_content) < 50 and any(kw in agent_result_content.lower() for kw in [\"don't know\", \"cannot fulfill\", \"无法回答\", \"不明白\"]):\n new_status = \"failed\" # 或 \"pending_review\" ? 暂时设为 failed\n evaluation_notes = f\"Task likely failed: Agent {agent_name} indicated inability to fulfill request.\"\n print(f\" Task {current_task_id} evaluated as FAILED (Agent Indicated Inability).\")\n else:\n # 如果以上都不是，则认为是成功\n new_status = \"completed\"\n print(f\" Task {current_task_id} evaluated as COMPLETED.\")\n # --- 评估逻辑结束 ---\n\n\n # --- 更新 Plan 状态 (逻辑不变) ---\n try:\n update_kwargs = {\n \"new_status\": new_status, \n \"new_evaluation\": evaluation_notes,\n \"new_notes\": agent_result_content[:1000] + \"...\" if agent_result_content and len(agent_result_content) > 1000 else agent_result_content \n }\n print(f\" Updating task {current_task_id} with: {{'status': '{new_status}', ...}}\")\n \n if updated_plan and PlanningStateHandler.get_task(updated_plan, current_task_id):\n updated_plan = PlanningStateHandler.update_task(updated_plan, by_id=current_task_id, **update_kwargs)\n updated_plan = PlanningStateHandler.set_current_task(updated_plan, None) \n updated_plan = PlanningStateHandler.update_plan_status(updated_plan)\n print(f\" Plan status after evaluation update: {updated_plan.get('status')}\")\n plan_updated = True\n else:\n raise ValueError(f\"Task ID '{current_task_id}' not found or plan invalid before update.\")\n\n except ValueError as ve: error_message = f\"Error updating plan: {ve}\"; print(f\"ERROR: {error_message}\"); traceback.print_exc()\n except Exception as e: error_message = f\"Unexpected error updating plan: {e}\"; print(f\"ERROR: {error_message}\"); traceback.print_exc()\n\n # --- 准备返回字典 (逻辑不变) ---\n updates: Dict[str, Any] = {}\n if updated_plan is not None: updates[\"plan\"] = updated_plan \n elif plan is not None: updates[\"plan\"] = plan \n \n # 记录本节点错误，或清除旧错误\n current_state_error = state.get(\"error\") \n if error_message: updates[\"error\"] = error_message \n elif current_state_error: updates[\"error\"] = None \n\n updates[\"messages\"] = [] # Evaluator 不添加消息\n\n print(f\"--- Exiting Evaluate Result Node. Plan updated: {plan_updated} ---\")\n return updates\n\n# --- 同步包装器 (保持不变) ---\ndef evaluate_result_node_logic_sync(state: PlanningAgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:\n \"\"\"evaluate_result_node_logic 的同步包装器\"\"\"\n print(f\"--- Entering Evaluate Result Node (Sync Wrapper) ---\")\n try:\n import anyio \n return anyio.run(evaluate_result_node_logic, state, config) # type: ignore\n except Exception as e:\n print(f\"Error running evaluate_result_node_logic synchronously: {e}\")\n traceback.print_exc()\n return {\"error\": f\"Evaluate Result sync execution failed: {e}\", \"plan\": state.get(\"plan\"), \"messages\": []}"
},
{
"path": "core/agents/state_based_supervisor/handoff.py",
"content": "# reason_graph/handoff.py\n# (Paste the code user provided for handoff.py here)\nimport re\nimport uuid\nfrom typing import List, Tuple # Import Tuple\n\nfrom langchain_core.messages import AIMessage, ToolCall, ToolMessage, BaseMessage # Import BaseMessage\nfrom langchain_core.tools import BaseTool, InjectedToolCallId, tool\nfrom langgraph.prebuilt import InjectedState\nfrom langgraph.types import Command\nfrom typing_extensions import Annotated\n\nWHITESPACE_RE = re.compile(r\"\\s+\")\n\ndef _normalize_agent_name(agent_name: str) -> str:\n \"\"\"Normalize an agent name to be used inside the tool name.\"\"\"\n if not agent_name: return \"unknown_agent\"\n return WHITESPACE_RE.sub(\"_\", agent_name.strip()).lower()\n\n# Note: The original code uses @tool decorator which requires function arguments.\n# To inject state, the decorated function needs the Annotated state argument.\n# Let's define the function first and then apply the decorator, or use functools.partial.\n# Using the function approach first for clarity.\n\ndef _handoff_to_agent_implementation(\n state: Annotated[dict, InjectedState], # Inject state here\n tool_call_id: Annotated[str, InjectedToolCallId], # Inject tool_call_id\n target_agent_name: str, # Pass the target agent name\n tool_name: str # Pass the specific tool name for the ToolMessage\n) -> Command:\n \"\"\"Ask another agent for help. This is the core logic.\"\"\"\n # Create the ToolMessage confirming the handoff BEFORE generating the Command\n \"\"\"Handoff 核心逻辑，添加日志\"\"\"\n print(f\"\\n--- DEBUG: Entering _handoff_to_agent_implementation ---\")\n print(f\" - Target Agent: {target_agent_name}\")\n print(f\" - Tool Name: {tool_name}\")\n print(f\" - Tool Call ID: {tool_call_id}\")\n # print(f\" - Current State Keys: {list(state.keys())}\") # 可选：打印状态键\n tool_message = ToolMessage(\n content=f\"Okay, handing off to {target_agent_name}. The current state and task context have been passed.\",\n name=tool_name,\n tool_call_id=tool_call_id,\n )\n print(f\" - Created ToolMessage: ID={tool_message.tool_call_id}, Name={tool_message.name}\")\n # The Command tells LangGraph to route to the target agent node\n # It also includes the ToolMessage in the state update for the next step\n command_obj = Command(\n goto=target_agent_name,\n # graph=Command.PARENT, # PARENT is default, usually not needed unless nested graphs\n update={\"messages\": [tool_message]}, # Return only the NEW message to be added\n )\n print(f\" - Created Command: goto='{command_obj.goto}', update contains {len(command_obj.update.get('messages',[]))} message(s)\")\n print(f\"--- DEBUG: Exiting _handoff_to_agent_implementation ---\")\n return command_obj\n\ndef create_handoff_tool(*, agent_name: str) -> BaseTool:\n \"\"\"Create a tool that can handoff control to the requested agent.\"\"\"\n if not agent_name:\n raise ValueError(\"agent_name cannot be empty for create_handoff_tool\")\n\n normalized_name = _normalize_agent_name(agent_name)\n tool_name = f\"transfer_to_{normalized_name}\"\n\n # Use functools.partial to fix the target_agent_name and tool_name arguments\n import functools\n specific_handoff_logic = functools.partial(\n _handoff_to_agent_implementation,\n target_agent_name=agent_name,\n tool_name=tool_name\n )\n\n # Decorate the partial function\n # The arguments 'state' and 'tool_call_id' will be automatically injected by LangGraph\n # when the tool is called due to the Annotations used in _handoff_to_agent_implementation\n @tool(tool_name)\n def handoff_tool_wrapper(\n state: Annotated[dict, InjectedState],\n tool_call_id: Annotated[str, InjectedToolCallId]\n ) -> Command:\n \"\"\"Dynamically generated tool description: Ask the '{agent_name}' agent for help with the current task or question.\"\"\"\n # --- 添加 Debug 日志 ---\n print(f\"\\n--- DEBUG: Handoff Tool '{tool_name}' (wrapper) CALLED ---\")\n # ---\n return specific_handoff_logic(state=state, tool_call_id=tool_call_id) # type: ignore\n\n # Set a more descriptive description\n handoff_tool_wrapper.description = f\"Use this tool to delegate the current task or ask a question to the '{agent_name}' agent. Pass the necessary context or instructions in your reasoning before calling this tool.\"\n\n return handoff_tool_wrapper\n\n\ndef create_handoff_back_messages(\n agent_name: str, supervisor_name: str\n) -> Tuple[AIMessage, ToolMessage]:\n \"\"\"Create a pair of (AIMessage, ToolMessage) to add to the message history when returning control to the supervisor.\"\"\"\n tool_call_id = str(uuid.uuid4())\n # Although no tool exists for transferring back, we simulate the pattern\n # The AIMessage signals intent, the ToolMessage confirms the transition occurred in the graph logic\n simulated_tool_name = f\"transfer_back_to_{_normalize_agent_name(supervisor_name)}\"\n\n # The AIMessage contains the *final output* of the sub-agent in its content field\n # It should also indicate the intent to hand back, though the graph logic forces this anyway.\n # The content here is just a placeholder - the actual content comes from the agent's final response.\n ai_message_content = f\"Task completed. Transferring back to {supervisor_name}.\"\n\n # We still generate a ToolCall structure for consistency in the AIMessage, even if no real tool is called on supervisor side for hand-back.\n tool_calls = [ToolCall(name=simulated_tool_name, args={}, id=tool_call_id)]\n\n # Create the AIMessage - crucial to include the sub-agent's name\n ai_message = AIMessage(\n content=ai_message_content, # Placeholder - see note below\n tool_calls=tool_calls,\n name=agent_name, # Identify which agent is responding\n )\n\n # The ToolMessage confirms the transition happened from the graph's perspective\n tool_message = ToolMessage(\n content=f\"Successfully transferred back to {supervisor_name} from {agent_name}.\",\n name=simulated_tool_name,\n tool_call_id=tool_call_id,\n )\n\n # IMPORTANT NOTE: The `_make_call_agent` helper function should populate the\n # `ai_message.content` with the *actual* final response message(s) from the sub-agent,\n # replacing the placeholder content above. It keeps the tool_calls structure.\n # The code provided for `_make_call_agent` seems to handle extracting `output['messages']`.\n # We need to ensure it correctly structures the AIMessage part of the tuple returned here.\n # Let's refine create_handoff_back_messages to just create the ToolMessage,\n # as the AIMessage content comes from the sub-agent's actual final output.\n\n # Refined approach: _make_call_agent gets the final AI response, we only need the ToolMessage here?\n # No, the pattern expects both. Let's assume _make_call_agent takes the *last* message from the\n # sub-agent's output and packages it into this AIMessage structure.\n\n return ai_message, tool_message # Return both for the standard pattern"
},
{
"path": "core/agents/state_based_supervisor/planner_node.py",
"content": "import re\nimport json\nimport time\nimport copy\nimport ast\nimport traceback\nimport anyio # <--- 导入 anyio\nfrom typing import Dict, Any, List, Optional, Union\nfrom datetime import datetime\nfrom langchain_core.messages import BaseMessage, AIMessage, SystemMessage, HumanMessage\nfrom langchain_core.runnables import RunnableConfig\n\n# 内部导入\ntry:\n from .state_schema import PlanningAgentState, Plan\n from .planning_handler import PlanningStateHandler\n from .prompt import PLANNER_SYSTEM_PROMPT_TEMPLATE\nexcept ImportError as e:\n print(f\"Error importing modules in planner_node.py: {e}\")\n class PlanningAgentState(Dict): pass; \n class Plan(Dict): pass; \n class PlanningStateHandler: pass\n PLANNER_SYSTEM_PROMPT_TEMPLATE = \"Fallback Planner Prompt: Error loading template. Args: {agent_descriptions}\"\n\n# --- Planner 节点核心逻辑 (异步) ---\nasync def planner_node_logic(\n state: PlanningAgentState,\n config: Optional[RunnableConfig],\n model: Any, # Planner 使用的 LLM\n agent_description_map: Dict[str, str] # 需要 Agent 描述来分配任务\n) -> Dict[str, Any]:\n \"\"\"Planner 节点逻辑：分析请求，生成初始计划\"\"\"\n print(f\"--- Entering Planner Node ---\")\n messages: List[BaseMessage] = state.get('messages', [])\n # Planner 通常在 plan 为空时运行\n plan: Optional[Plan] = state.get('plan')\n if plan:\n print(\"Planner Node: Plan already exists. Skipping plan creation.\")\n # 如果计划已存在，Planner 不应再执行，直接返回当前状态？\n # 或者返回一个空更新，让图流向 Supervisor？\n # 返回空更新更安全，让 Supervisor 继续\n return {} # 返回空字典，状态不变\n\n if not messages:\n print(\"Planner Node: No messages found to create a plan from.\")\n return {\"error\": \"Planner received no messages.\"}\n\n # --- 1. 准备 Planner Prompt ---\n # Planner 只需要 Agent 描述，不需要 plan_json 或 current_date?\n # 可以让它知道日期\n desc_list = [f\"- {name}: {desc}\" for name, desc in agent_description_map.items()]\n agent_descriptions_str = \"\\n\".join(desc_list)\n current_date_str = datetime.now().strftime(\"%a, %b %d, %Y\") # Planner 也可能需要日期\n\n system_prompt_text = \"Error: Planner prompt template could not be loaded/formatted.\"\n try:\n # 加载 Planner 的模板\n from .prompt import PLANNER_SYSTEM_PROMPT_TEMPLATE\n system_prompt_text = PLANNER_SYSTEM_PROMPT_TEMPLATE.format(\n agent_descriptions=agent_descriptions_str,\n # 如果 Planner Prompt 需要日期：\n current_date=current_date_str\n )\n except ImportError: print(\"ERROR: Could not import PLANNER_SYSTEM_PROMPT_TEMPLATE\")\n except KeyError as e: print(f\"ERROR: Missing key in planner prompt formatting: {e}\")\n except Exception as e: print(f\"ERROR: Unexpected error loading/formatting planner prompt: {e}\")\n\n # Planner 的输入只需要 System Prompt 和用户的初始请求（通常是第一条）\n # 或者传递最后几条消息？为了简单，先只用第一条 HumanMessage\n initial_user_request = next((m for m in messages if isinstance(m, HumanMessage)), None)\n if not initial_user_request:\n print(\"Planner Node: No HumanMessage found in initial state.\")\n return {\"error\": \"Planner did not find initial user request.\"}\n\n llm_input_messages = [SystemMessage(content=system_prompt_text), initial_user_request]\n\n # --- 2. 调用 Planner LLM ---\n print(\"--- Calling Planner LLM ---\")\n response: Optional[AIMessage] = None\n llm_error_msg: Optional[str] = None\n try:\n response = await model.ainvoke(llm_input_messages, config=config)\n if not isinstance(response, AIMessage): raise TypeError(\"Planner LLM returned non-AIMessage.\")\n # Planner 的回复主要是指令，可以不设置 name\n print(f\"Planner LLM Raw Response Content: {response.content[:300]}...\")\n # Planner 不应该调用工具\n if response.tool_calls: print(\"Warning: Planner LLM unexpectedly generated tool calls!\")\n messages_to_add: List[BaseMessage] = [response] # 可以选择是否将 Planner 的思考过程加入 history\n except Exception as e:\n print(f\"!!! Error invoking Planner LLM: {e}\"); traceback.print_exc()\n llm_error_msg = f\"Planner LLM invocation failed: {e}\"\n messages_to_add = []\n response = None\n\n # --- 3. 处理 Planner LLM 回复 (解析 CREATE_PLAN) ---\n new_plan: Optional[Plan] = None\n plan_updated: bool = False # 标记计划是否在本节点成功创建\n directive_error_msg: Optional[str] = None\n\n if response and isinstance(response.content, str):\n try:\n plan_match = re.search(r\"PLAN_UPDATE:\\s*CREATE_PLAN\\s*(\\{.*?\\})\\s*$\", response.content, re.IGNORECASE | re.DOTALL | re.MULTILINE)\n if plan_match:\n args_json_str = plan_match.group(1)\n print(f\"Planner directive found: CREATE_PLAN with args: {args_json_str[:100]}...\")\n try:\n args = json.loads(args_json_str)\n if not isinstance(args, dict): raise ValueError(\"Args JSON not a dict.\")\n \n title=args.get(\"title\", \"Plan\"); desc=args.get(\"description\",\"\"); tasks=args.get(\"tasks\",[])\n if isinstance(tasks, list) and all(isinstance(t, dict) and 'description' in t for t in tasks):\n for task_data in tasks: task_data['status'] = 'pending' # 强制状态\n new_plan = PlanningStateHandler.create_plan(title, desc)\n new_plan = PlanningStateHandler.add_tasks(new_plan, tasks); plan_updated = True\n print(\"DEBUG: Plan successfully created by Planner node.\")\n else: raise ValueError(\"Invalid 'tasks' format (must be list of dicts with 'description').\")\n\n except (json.JSONDecodeError, ValueError, KeyError, TypeError) as e:\n err_msg = f\"Error processing CREATE_PLAN directive: {type(e).__name__} - {e}\"\n print(err_msg); traceback.print_exc(); directive_error_msg = err_msg\n except Exception as e:\n err_msg = f\"Unexpected error processing CREATE_PLAN: {type(e).__name__} - {e}\"\n print(err_msg); traceback.print_exc(); directive_error_msg = err_msg\n else:\n directive_error_msg = \"Planner LLM did not output a valid PLAN_UPDATE: CREATE_PLAN directive.\"\n print(f\"Warning: {directive_error_msg}\")\n # 即使没有指令，也可能需要返回 Planner 的回复消息\n # 但如果没有 plan，流程可能无法继续，所以记录错误\n\n except Exception as outer_e:\n directive_error_msg = f\"Error searching for PLAN_UPDATE directive: {outer_e}\"\n print(f\"ERROR: {directive_error_msg}\"); traceback.print_exc()\n\n # --- 4. 准备返回的状态更新 ---\n updates: Dict[str, Any] = {\"messages\": messages_to_add} # 添加 Planner 的回复消息\n if plan_updated and new_plan:\n updates[\"plan\"] = new_plan # 返回新创建的 Plan\n \n final_error = llm_error_msg or directive_error_msg\n if final_error: # 记录 Planner 步骤中遇到的第一个错误\n updates[\"error\"] = final_error\n\n print(f\"--- Exiting Planner Node. Plan created: {plan_updated} ---\")\n return updates\n\n\n# --- Planner 节点的同步包装器 (使用 anyio) ---\ndef planner_node_logic_sync(\n state: PlanningAgentState,\n config: Optional[RunnableConfig],\n model: Any,\n agent_description_map: Dict[str, str]\n) -> Dict[str, Any]:\n \"\"\"planner_node_logic 的同步包装器\"\"\"\n print(f\"--- Entering Planner Node (Sync Wrapper) ---\")\n try:\n # 使用 anyio 在同步函数中运行异步函数\n return anyio.run( # type: ignore\n planner_node_logic, state, config, model, agent_description_map\n )\n except Exception as e:\n print(f\"Error running planner_node_logic synchronously: {e}\")\n traceback.print_exc()\n return {\"error\": f\"Planner sync execution failed: {e}\", \"messages\": state.get(\"messages\",[])}"
},
{
"path": "core/agents/state_based_supervisor/planning_handler.py",
"content": "# reason_graph/planning_handler.py\nimport uuid\nimport datetime\nfrom typing import List, Dict, Optional, Any\nfrom .state_schema import TaskStatus, PlanningStatus, Task, Plan # 从 state_schema 导入类型\n\nclass PlanningStateHandler:\n \"\"\"\n 使用静态方法管理一个表示项目计划的字典。\n 计划现在存储在 LangGraph 的状态中，此类提供操作该字典的函数。\n \"\"\"\n\n @staticmethod\n def _now() -> str:\n return datetime.datetime.now(datetime.timezone.utc).isoformat()\n\n @staticmethod\n def _gen_id() -> str:\n # 生成更易读的任务 ID (可选)\n # return f\"task_{str(uuid.uuid4())[:8]}\"\n return str(uuid.uuid4())\n\n @staticmethod\n def create_plan(title: str, description: str) -> Plan:\n \"\"\"创建一个新的 Plan 字典\"\"\"\n now = PlanningStateHandler._now()\n return Plan(\n title=title,\n description=description,\n status=\"planning\", # 初始状态为规划中\n tasks=[],\n current_task_id=None,\n created_at=now,\n updated_at=now,\n completed_at=None,\n )\n\n @staticmethod\n def create_task(description: str,\n agent: Optional[str] = None,\n dependencies: Optional[List[str]] = None) -> Task:\n \"\"\"创建一个新的 Task 字典\"\"\"\n now = PlanningStateHandler._now()\n return Task(\n id=PlanningStateHandler._gen_id(),\n description=description.strip(),\n status=\"pending\", # 初始状态为待处理\n agent=agent.strip() if agent else None,\n created_at=now,\n updated_at=now,\n completed_at=None,\n dependencies=dependencies or [],\n notes=None,\n evaluation=None,\n result=None,\n )\n\n @staticmethod\n def add_tasks(plan: Plan, tasks_data: List[Dict[str, Any]]) -> Plan:\n \"\"\"向 Plan 字典中添加任务\"\"\"\n if not isinstance(plan, dict) or \"tasks\" not in plan:\n raise ValueError(\"Invalid plan structure provided.\")\n if not isinstance(tasks_data, list):\n raise ValueError(\"tasks_data must be a list of task dictionaries.\")\n\n for tinfo in tasks_data:\n desc = tinfo.get(\"description\")\n if not desc: continue # 跳过没有描述的任务\n agent = tinfo.get(\"agent\")\n deps = tinfo.get(\"dependencies\")\n task = PlanningStateHandler.create_task(desc, agent, deps)\n plan[\"tasks\"].append(task)\n\n # 如果添加任务时计划仍在 planning 阶段，可以转为 ready\n if plan.get(\"status\") == \"planning\":\n plan[\"status\"] = \"ready\"\n\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan\n\n @staticmethod\n def update_task(plan: Plan,\n by_id: Optional[str] = None,\n new_desc: Optional[str] = None,\n new_status: Optional[TaskStatus] = None,\n new_agent: Optional[str] = None,\n new_notes: Optional[str] = None,\n new_evaluation: Optional[str] = None,\n new_result: Optional[Any] = None) -> Plan:\n \"\"\"更新 Plan 字典中指定 ID 的任务\"\"\"\n if not isinstance(plan, dict) or \"tasks\" not in plan:\n raise ValueError(\"Invalid plan structure provided.\")\n if not by_id:\n raise ValueError(\"Must provide 'by_id' to update a task.\")\n\n task = next((t for t in plan[\"tasks\"] if t.get(\"id\") == by_id), None)\n if not task:\n raise ValueError(f\"No matching task found with ID: {by_id}\")\n\n updated = False\n if new_desc is not None and task.get(\"description\") != new_desc.strip():\n task[\"description\"] = new_desc.strip()\n updated = True\n if new_status is not None and task.get(\"status\") != new_status.strip():\n task[\"status\"] = new_status.strip()\n if new_status.strip() == \"completed\":\n task[\"completed_at\"] = PlanningStateHandler._now()\n updated = True\n if new_agent is not None and task.get(\"agent\") != new_agent.strip():\n task[\"agent\"] = new_agent.strip()\n updated = True\n if new_notes is not None and task.get(\"notes\") != new_notes.strip():\n task[\"notes\"] = new_notes.strip()\n updated = True\n if new_evaluation is not None and task.get(\"evaluation\") != new_evaluation.strip():\n task[\"evaluation\"] = new_evaluation.strip()\n updated = True\n if new_result is not None: # 直接更新结果（谨慎使用，可能很大）\n task[\"result\"] = new_result\n updated = True\n\n if updated:\n task[\"updated_at\"] = PlanningStateHandler._now()\n plan[\"updated_at\"] = PlanningStateHandler._now() # 更新整个计划的更新时间\n\n # 检查并更新整个计划的状态\n plan = PlanningStateHandler.update_plan_status(plan)\n\n return plan\n\n @staticmethod\n def update_plan_status(plan: Plan) -> Plan:\n \"\"\"根据任务状态自动更新计划状态\"\"\"\n if not isinstance(plan, dict) or \"tasks\" not in plan:\n return plan # Return as is if invalid\n\n tasks = plan[\"tasks\"]\n if not tasks: # 没有任务\n if plan.get(\"status\") not in [\"completed\", \"failed\", \"error\"]:\n plan[\"status\"] = \"ready\" # 或 \"completed\" 如果没有任务就算完成? 设为 ready 似乎更合理\n return plan\n\n all_completed = all(t.get(\"status\") == \"completed\" for t in tasks)\n any_failed = any(t.get(\"status\") == \"failed\" for t in tasks)\n any_in_progress = any(t.get(\"status\") in [\"in_progress\", \"pending_review\"] for t in tasks)\n any_pending = any(t.get(\"status\") == \"pending\" for t in tasks)\n\n current_status = plan.get(\"status\")\n new_status = current_status\n\n if any_failed:\n new_status = \"failed\" # 或 \"error\"\n elif all_completed:\n new_status = \"completed\"\n plan[\"completed_at\"] = PlanningStateHandler._now()\n elif any_in_progress:\n new_status = \"executing\"\n elif any_pending or not any_in_progress: # 如果还有 pending 或所有任务都结束了但不是 completed/failed\n if current_status not in [\"completed\", \"failed\", \"error\"]: # 避免覆盖最终状态\n new_status = \"ready\" # 准备好执行或等待新任务\n\n if new_status != current_status:\n plan[\"status\"] = new_status\n plan[\"updated_at\"] = PlanningStateHandler._now()\n\n return plan\n\n @staticmethod\n def set_current_task(plan: Plan, task_id: Optional[str]) -> Plan:\n \"\"\"设置 Plan 字典中的当前任务 ID\"\"\"\n if not isinstance(plan, dict):\n raise ValueError(\"Invalid plan structure provided.\")\n\n if task_id is None:\n plan[\"current_task_id\"] = None\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan\n\n found = any(t.get(\"id\") == task_id for t in plan.get(\"tasks\", []))\n if not found:\n raise ValueError(f\"Task ID '{task_id}' not found in plan.\")\n\n if plan.get(\"current_task_id\") != task_id:\n plan[\"current_task_id\"] = task_id\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan\n\n @staticmethod\n def get_task(plan: Plan, task_id: str) -> Optional[Task]:\n \"\"\"根据 ID 获取任务字典\"\"\"\n if not isinstance(plan, dict) or \"tasks\" not in plan:\n return None\n return next((t for t in plan[\"tasks\"] if t.get(\"id\") == task_id), None)\n\n @staticmethod\n def get_next_pending_task(plan: Plan) -> Optional[Task]:\n \"\"\"获取下一个处于 pending 状态且所有依赖已完成的任务\"\"\"\n if not isinstance(plan, dict) or \"tasks\" not in plan:\n return None\n\n completed_task_ids = {t[\"id\"] for t in plan[\"tasks\"] if t.get(\"status\") == \"completed\"}\n\n for task in plan[\"tasks\"]:\n if task.get(\"status\") == \"pending\":\n dependencies = task.get(\"dependencies\", [])\n if not dependencies or all(dep_id in completed_task_ids for dep_id in dependencies):\n return task\n return None # 没有找到合适的下一个任务\n\n @staticmethod\n def finish_plan(plan: Plan) -> Plan:\n \"\"\"强制将 Plan 标记为完成\"\"\"\n if not isinstance(plan, dict):\n raise ValueError(\"Invalid plan structure provided.\")\n if plan.get(\"status\") != \"completed\":\n plan[\"status\"] = \"completed\"\n plan[\"completed_at\"] = PlanningStateHandler._now()\n plan[\"updated_at\"] = PlanningStateHandler._now()\n return plan"
},
{
"path": "core/agents/state_based_supervisor/prompt.py",
"content": "# # --- Planner Agent System Prompt (新增) ---\n# PLANNER_SYSTEM_PROMPT_TEMPLATE = \"\"\"You are an expert planning agent. Your sole responsibility is to analyze a user request and create a detailed, step-by-step plan to fulfill it by coordinating specialized agents.\n\n# The current date is {current_date}.\n\n# ## Agent Descriptions:\n# {agent_descriptions}\n# *(This list includes the capabilities of available specialist agents.)*\n\n# ## Task:\n# Analyze the user request provided in the message history. Break it down into a sequence of logical tasks. For each task, determine the most suitable agent from the descriptions provided.\n\n# ## Output Format:\n# You MUST output **ONLY** a single `PLAN_UPDATE: CREATE_PLAN ` directive in your response content. The JSON arguments MUST be valid and contain:\n# - \"title\": A concise title for the overall plan.\n# - \"description\": A brief description summarizing the user's goal.\n# - \"tasks\": A list of task objects. Each task object MUST contain:\n# - \"description\": A clear and actionable description of the specific sub-task.\n# - \"agent\": The name of the MOST SUITABLE agent from the Agent Descriptions to perform this task. Leave empty (\"\") if unsure or if it's a general task.\n# - \"status\": Set **all** initial tasks to **\"pending\"**.\n# - (Optional) \"dependencies\": A list of task IDs (UUIDs that will be generated later) this task depends on, if any (usually empty for initial plan).\n\n# **Example JSON Args:**\n# `{{\"title\": \"Research and Report on AI Ethics\", \"description\": \"User wants a report on AI ethics, including research and writing.\", \"tasks\": [{{\"description\": \"Research current trends in AI ethics using web search\", \"agent\": \"research_expert\", \"status\": \"pending\"}}, {{\"description\": \"Write a structured report summarizing the findings\", \"agent\": \"reporter_expert\", \"status\": \"pending\", \"dependencies\": [\"\"]}}]}}` \n# *(Note: Actual IDs are UUIDs generated later, dependencies often added via UPDATE_TASK)*\n\n# **CRITICAL**: Output **ONLY** the `PLAN_UPDATE: CREATE_PLAN ` directive and nothing else. Do not add conversational text. Make sure the JSON is valid.\n# \"\"\"\n\n# SUPERVISOR_PLANNING_PROMPT_TEMPLATE = \"\"\"You are a meticulous top-level Supervisor agent responsible for executing an existing plan, coordinating specialist agents, and managing task execution based on the provided state. You rely on an external evaluator node to assess task completion after agents run.\n\n# The current date is {current_date}.\n\n# ## Current Plan State:\n# ```json\n# {plan_json}\n# ```\n# *(Review plan status and individual task statuses and IDs (UUIDs). Your main goal is to drive the plan status to 'completed'.)*\n\n# ## Agent Descriptions:\n# {agent_descriptions}\n\n# ## Your Goal:\n# Execute the **existing plan** strictly step-by-step towards 'completed' status. Make **exactly one** logical primary decision per turn. **Do NOT evaluate agent results or mark tasks 'completed'/'failed' yourself.**\n\n# ## Workflow & Decision Process (Strict Sequence):\n# 1. **Analyze State**: Review the latest messages and the 'Current Plan State'. (Note: If the last message is from a sub-agent, an evaluator node has already processed it and updated the plan state before your turn).\n# 2. **Determine ONE Next Action**: Execute the FIRST matching condition below and **IMMEDIATELY END YOUR TURN**:\n\n# * **A. Initiate Next Task**: If the plan is 'ready' or 'executing', AND no task is currently 'in_progress', AND a 'pending' task is ready (dependencies met):\n# * **Action**: Find the FIRST such task. Output **ONLY** `PLAN_UPDATE: UPDATE_TASK `. **CRITICAL: Use the exact UUID for `by_id`!** JSON Args should be ` {{\"by_id\": \"\", \"status\": \"in_progress\"}}`.\n# * **B. Delegate In-Progress Task**: If a task **currently has status 'in_progress'** (check plan state):\n# * **Action**: Identify the best agent. Output **ONLY** the `transfer_to_` tool call. **CRITICAL**: Tool call args **MUST** include `\"task_id\": \"\"` and clear `\"instructions\"`.\n# * **C. Finish Plan**: If **ALL** tasks in the plan now have status 'completed' AND the plan status is NOT 'completed' yet (check plan state provided):\n# * **Action**: Output **ONLY** `PLAN_UPDATE: FINISH_PLAN {{}}`.\n# * **D. Generate Final Output**: If the **Plan Status IS 'completed'** (check plan state provided):\n# * **Action**: Decide final output format based on original request. EITHER call `transfer_to_reporter_expert` (passing context in args, like relevant task IDs) OR generate the final `AIMessage` content yourself summarizing the overall result.\n# * **E. Waiting/Blocked/Failed**: If no other action is appropriate (e.g., plan status 'failed', or waiting for dependencies):\n# * **Action**: Output a brief waiting or status message explaining the situation.\n\n# ## Output Constraints:\n# - Your response MUST contain exactly ONE primary action (ONE PLAN_UPDATE directive OR ONE transfer_to tool call OR the final answer OR a status message).\n# - `PLAN_UPDATE:` directives MUST be in the text content with **valid JSON arguments**.\n# - **CRITICAL**: `UPDATE_TASK` **MUST** use the correct Task UUID string for `\"by_id\"`.\n\n# ## Planning Directives Format (Mandatory - JSON Args in text):\n# - `PLAN_UPDATE: ADD_TASKS {{\"tasks\": [...]}}` # You can still add tasks if needed mid-plan\n# - `PLAN_UPDATE: UPDATE_TASK {{\"by_id\": \"\", \"status\": \"in_progress\", \"notes\": \"\"}}` (**UUID!** Only use non-terminal statuses).\n# - `PLAN_UPDATE: FINISH_PLAN {{}}`\n\n# ## Tool Usage:\n# - Only `transfer_to_` tools. Args **MUST** include `\"task_id\"` and `\"instructions\"`.\n\n# Now, analyze the current state (which reflects any recent evaluations) and the LAST message, and determine the single next action based strictly on the workflow for **executing the existing plan**. Remember, you do **not** evaluate results or mark tasks complete/failed.\n# \"\"\"\n\n# --- Planner Agent System Prompt ---\nPLANNER_SYSTEM_PROMPT_TEMPLATE = \"\"\"You are an expert planning agent. Your sole responsibility is to analyze a user request and create a detailed, step-by-step plan to fulfill it by coordinating specialized agents.\n\nThe current date is {current_date}.\n\n## Agent Descriptions:\n{agent_descriptions}\n*(This list includes the capabilities of available specialist agents.)*\n\n## Task:\nAnalyze the user request provided in the message history. Break it down into a sequence of logical tasks. For each task, determine the most suitable agent from the descriptions provided.\n\n## Task Granularity Guidelines:\n- **IMPORTANT**: Maintain appropriate task granularity based on complexity:\n - For simple requests, create just 1-2 tasks that can be completed by a single agent\n - For complex requests, break down into 3-5 logical steps\n - Avoid excessive fragmentation of simple tasks\n - Each task should represent a meaningful unit of work\n\n## Output Format:\nYou MUST output **ONLY** a single `PLAN_UPDATE: CREATE_PLAN ` directive in your response content. The JSON arguments MUST be valid and contain:\n- \"title\": A concise title for the overall plan.\n- \"description\": A brief description summarizing the user's goal.\n- \"tasks\": A list of task objects. Each task object MUST contain:\n - \"description\": A clear and actionable description of the specific sub-task.\n - \"agent\": The name of the MOST SUITABLE agent from the Agent Descriptions to perform this task. Leave empty (\"\") if unsure or if it's a general task.\n - \"status\": Set **all** initial tasks to **\"pending\"**.\n - (Optional) \"dependencies\": Usually empty for initial plan.\n\n**Example JSON Args for SIMPLE request:**\n`{{\"title\": \"Answer Question About Python\", \"description\": \"User wants to know how to use list comprehensions in Python\", \"tasks\": [{{\"description\": \"Provide a comprehensive explanation of Python list comprehensions with examples\", \"agent\": \"coder_expert\", \"status\": \"pending\"}}]}}`\n\n**Example JSON Args for COMPLEX request:**\n`{{\"title\": \"Research and Report on AI Ethics\", \"description\": \"User wants a detailed report on AI ethics\", \"tasks\": [{{\"description\": \"Research current trends in AI ethics using web search\", \"agent\": \"research_expert\", \"status\": \"pending\"}}, {{\"description\": \"Write a structured report summarizing the findings\", \"agent\": \"reporter_expert\", \"status\": \"pending\"}}]}}`\n\n**CRITICAL**: Output **ONLY** the `PLAN_UPDATE: CREATE_PLAN ` directive and nothing else. Do not add conversational text. Make sure the JSON is valid.\n\"\"\"\n\n# --- Supervisor Planning Prompt (允许动作组合 + 强制UUID/JSON) ---\nSUPERVISOR_PLANNING_PROMPT_TEMPLATE = \"\"\"You are a meticulous top-level Supervisor agent responsible for executing an existing plan, coordinating specialist agents, and managing task execution based on the provided state.\n\nThe current date is {current_date}.\n\n## Current Plan State:\n```json\n{plan_json}\n```\n*(Review plan status and individual task statuses and IDs (UUIDs). Your main goal is to drive the plan status to 'completed'.)*\n\n## Agent Descriptions:\n{agent_descriptions}\n*(This list includes specialist agents and yourself.)*\n\n## Your Goal:\nExecute the **existing plan** step-by-step towards 'completed' status by making logical decisions and issuing appropriate directives and tool calls.\n\n## Workflow & Decision Guidelines:\n1. **Analyze State**: Review the latest messages (especially agent results) and the 'Current Plan State'.\n2. **Determine Next Action(s)**: Based on the analysis, decide the next logical step(s).\n\n * **If a sub-agent just returned results**:\n a. Evaluate the result against the task.\n b. Issue the `PLAN_UPDATE: UPDATE_TASK `. **CRITICAL: Use the exact Task UUID for `by_id`!** Include `evaluation` and `notes`.\n c. **After** the update directive, **if** more tasks are pending and ready, you **CAN** identify the next task, issue `PLAN_UPDATE: UPDATE_TASK ` (using its UUID), **AND** issue the corresponding `transfer_to_` tool call **in the same response**.\n * **If no agent just returned, AND a 'pending' task is ready**:\n a. Identify the *next* suitable 'pending' task.\n b. Issue `PLAN_UPDATE: UPDATE_TASK ` (using its UUID).\n c. **Immediately following** the directive in the same response, issue the corresponding `transfer_to_` tool call with instructions (including Task UUID).\n * **If ALL tasks are 'completed' AND plan status is NOT 'completed' yet**:\n a. Issue `PLAN_UPDATE: FINISH_PLAN {{}}`.\n b. **In the same response**, decide the final output: EITHER call `transfer_to_reporter_expert` OR generate the final `AIMessage` content yourself.\n * **If Plan Status IS 'completed'**:\n a. Your job is done. Generate the final `AIMessage` content if you didn't call the reporter in the previous step.\n * **If Waiting/Blocked/Failed**: Output a status message explaining the situation.\n\n## Output Constraints:\n- Your response **CAN** contain **both** a `PLAN_UPDATE:` directive (in content) and a `transfer_to_` tool call if logically appropriate (e.g., completing one task and starting the next).\n- Your response **CAN** contain **both** `PLAN_UPDATE: FINISH_PLAN` and the final action (call reporter or final answer).\n- **NEVER** delegate to more than one agent simultaneously (only one `transfer_to_` tool call per response).\n- `PLAN_UPDATE:` directives MUST be in the text content with **valid JSON arguments**.\n- **CRITICAL**: `UPDATE_TASK` **MUST** use the correct Task UUID string for `\"by_id\"`.\n\n## Planning Directives Format (Mandatory - JSON Args in text):\nUse these exact formats **within your response content**. Arguments **MUST** be a valid JSON string.\n- `PLAN_UPDATE: ADD_TASKS {{\"tasks\": [...]}}`\n- `PLAN_UPDATE: UPDATE_TASK {{\"by_id\": \"\", \"status\": \"\", \"evaluation\": \"\", \"notes\": \"\"}}` (**UUID!**)\n- `PLAN_UPDATE: FINISH_PLAN {{}}`\n*(Note: CREATE_PLAN is handled by the Planner Agent)*\n\n## Tool Usage:\n- Only `transfer_to_` tools are callable by you. Args **MUST** include `\"task_id\"` and `\"instructions\"`.\n\nNow, analyze the current state and messages, and determine the necessary action(s) for this turn.\n\"\"\"\n\n\n# **主要调整说明:**\n\n# 1. **允许动作组合**: 修改了 Workflow 和 Output Constraints，明确允许 Supervisor 在一个回合中既更新 Plan 状态（通过 `PLAN_UPDATE:` 指令）又委派任务（通过 `transfer_to_` 工具调用），或者在结束计划的同时进行最终输出操作。这给予 LLM 更大的灵活性，可能更符合它的“思考习惯”。\n# 2. **保留核心要求**: 仍然**强制要求** `PLAN_UPDATE` 的参数必须是有效的 JSON，并且 `UPDATE_TASK` **必须**使用正确的 Task UUID。同时，**仍然禁止**一次委派多个 Agent。\n# 3. **移除了严格的 `STOP` 指令**: 不再强制要求 LLM 在发出 `PLAN_UPDATE` 后必须结束当前回合。\n\n# **预期效果:**\n\n# * Supervisor LLM 在处理完子 Agent 的结果并更新任务状态后，如果发现下一个任务已准备就绪，它可能会在同一个回复中直接发出 `transfer_to_` 指令，从而减少一个交互回合，提高效率。\n# * 在所有任务完成后，它可以一步到位地发出 `FINISH_PLAN` 并同时决定最终输出（调用 Reporter 或自己总结）。\n# * **潜在风险**: 这种灵活性也可能使得 LLM 在复杂情况下更容易出错（例如，忘记更新状态就去委派，或者错误地组合了动作）。但鉴于之前严格分步也遇到了问题，这种方式值得一试。"
},
{
"path": "core/agents/state_based_supervisor/state_schema.py",
"content": "# reason_graph/state_schema.py\nimport operator\nfrom typing import Dict, List, Optional, Any, Literal, TypedDict, Sequence, Annotated, Union\nfrom langchain_core.messages import BaseMessage\nfrom langgraph.graph.message import add_messages\nfrom langgraph.managed import IsLastStep, RemainingSteps\n\n# 定义计划状态类型\nPlanningStatus = Literal[\"not_started\", \"planning\", \"ready\", \"executing\", \"completed\", \"failed\", \"error\"]\n\n# 定义任务状态类型\nTaskStatus = Literal[\"pending\", \"ready\", \"in_progress\", \"completed\", \"failed\", \"skipped\", \"pending_review\", \"revision_needed\"]\n\n# 定义任务项\nclass Task(TypedDict, total=False):\n \"\"\"任务项定义\n\n 表示计划中的一个任务项，包含任务描述、状态、分配的代理等信息\n \"\"\"\n id: str # 任务唯一标识符\n description: str # 任务描述\n status: TaskStatus # 任务状态\n agent: Optional[str] # 分配的代理名称 (建议的执行者)\n created_at: str # 创建时间 (ISO 格式)\n updated_at: str # 更新时间 (ISO 格式)\n completed_at: Optional[str] # 完成时间 (ISO 格式)\n dependencies: Optional[List[str]] # 依赖的任务ID列表\n notes: Optional[str] # 关于任务执行情况的备注 (可由 Agent 或 Supervisor 更新)\n evaluation: Optional[str] # 对任务完成情况的评估 (可由 Supervisor LLM 或 Evaluator Agent 更新)\n result: Optional[Any] # (可选) 存储任务的直接输出结果摘要\n\n# 定义计划\nclass Plan(TypedDict, total=False):\n \"\"\"计划定义\n\n 表示一个完整的计划，包含计划状态、任务列表等信息\n \"\"\"\n status: PlanningStatus # 计划状态\n tasks: List[Task] # 任务列表\n current_task_id: Optional[str] # 当前 Supervisor 关注或正在处理的任务ID\n created_at: str # 创建时间 (ISO 格式)\n updated_at: str # 更新时间 (ISO 格式)\n completed_at: Optional[str] # 完成时间 (ISO 格式)\n title: Optional[str] # 计划标题\n description: Optional[str] # 计划描述 (通常是用户原始请求)\n\n# 扩展基础 AgentState 以支持计划功能\nclass PlanningAgentState(TypedDict):\n \"\"\"支持计划功能的、用于 Supervisor 图的状态定义\"\"\"\n messages: Annotated[Sequence[BaseMessage], add_messages] # 消息历史\n plan: Optional[Plan] = None # 存储计划对象\n # last_agent_result: Optional[Dict[str, Any]] = None # 存储刚结束的子 Agent 的 {name: ..., content: ...}\n is_last_step: IsLastStep # LangGraph 内部状态\n remaining_steps: RemainingSteps # LangGraph 内部状态, 用于防止无限循环\n error: Optional[str] = None # 用于记录执行中发生的错误信息\n # 可以根据需要添加其他全局共享的状态字段\n # 例如: shared_context: Optional[Dict] = None\n\n# 可以为子 Agent 定义一个稍微不同的状态（如果它们不需要 plan）\nclass BasicAgentState(TypedDict):\n \"\"\"基础 Agent 状态，仅包含消息历史\"\"\"\n messages: Annotated[Sequence[BaseMessage], add_messages]\n is_last_step: IsLastStep\n remaining_steps: RemainingSteps\n error: Optional[str] = None\n\n# 方便类型提示\nStateSchemaType = Union[Dict[str, Any], PlanningAgentState, BasicAgentState]"
},
{
"path": "core/agents/state_based_supervisor/supervisor_graph.py",
"content": "# reason_graph/supervisor_graph.py\nimport inspect\nimport re\nimport functools\nimport uuid\nimport asyncio\nimport anyio\nimport traceback \nfrom typing import Any, Callable, List, Optional, Type, Union, Dict, Literal, Sequence, cast # <--- 导入 cast\n\nfrom langchain_core.language_models import BaseChatModel, LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import AIMessage, ToolMessage, BaseMessage, ToolCall, SystemMessage # <--- 导入 SystemMessage\nfrom langchain_core.runnables import RunnableConfig\nfrom langgraph.utils.runnable import RunnableCallable\n\nfrom langgraph.graph import END, START, StateGraph\nfrom langgraph.graph.state import CompiledStateGraph\nfrom langgraph.prebuilt import ToolNode\nfrom langgraph.pregel import Pregel\n\n# 内部导入\ntry:\n from core.agents.base.base_agent import BaseAgent\n from .handoff import create_handoff_tool, _normalize_agent_name # 确保导入 _normalize_agent_name\n from .state_schema import PlanningAgentState, Plan # 导入 PlanningAgentState 和 Plan\n from .supervisor_node import supervisor_node_logic # 导入异步节点逻辑\n from .planner_node import planner_node_logic, planner_node_logic_sync # <--- 导入 Planner 逻辑\n from .evaluate_result_node import evaluate_result_node_logic, evaluate_result_node_logic_sync # <--- 导入 Evaluator 逻辑\n from .agent_name import AgentNameMode, with_agent_name\nexcept ImportError as e:\n print(f\"Error importing modules in supervisor_graph.py: {e}\")\n # Add Dummy classes for type hints if needed\n class BaseAgent: pass\n class PlanningAgentState(Dict): pass\n class Plan(Dict): pass\n class Pregel: pass\n AgentNameMode = Literal[\"inline\"]\n def create_handoff_tool(*args, **kwargs): return None # type: ignore\n def _normalize_agent_name(s: str) -> str: return s\n async def supervisor_node_logic(*args, **kwargs): return {}\n async def planner_node_logic(*args, **kwargs): return {} # <--- 添加 planner_node_logic\n def planner_node_logic_sync(*args, **kwargs): return {} # <--- 添加 planner_node_logic_sync\n async def evaluate_result_node_logic(*args, **kwargs): return {} # 添加 evaluate_result_node_logic \n def evaluate_result_node_logic_sync(*args, **kwargs): return {} # 添加 evaluate_result_node_logic_sync\n def with_agent_name(model, mode): return model\n\n\n# 定义 OutputMode, MODELS_NO_PARALLEL_TOOL_CALLS, _supports_disable_parallel_tool_calls (保持不变)\nOutputMode = Literal[\"full_history\", \"last_message\"]\nMODELS_NO_PARALLEL_TOOL_CALLS = {\"o3-mini\"}\ndef _supports_disable_parallel_tool_calls(model: LanguageModelLike) -> bool:\n if not isinstance(model, BaseChatModel): return False\n if hasattr(model, \"model_name\") and model.model_name in MODELS_NO_PARALLEL_TOOL_CALLS: return False\n if not hasattr(model, \"bind_tools\"): return False\n if \"parallel_tool_calls\" not in inspect.signature(model.bind_tools).parameters: return False\n return True\n\n\n# _make_call_agent (保持不变 - 已支持同步/异步)\ndef _make_call_agent(\n agent_graph: Pregel, \n output_mode: OutputMode,\n add_handoff_back_messages: bool, \n supervisor_name: str,\n) -> RunnableCallable:\n if output_mode not in [\"full_history\", \"last_message\"]: raise ValueError(...)\n\n async def acall_agent(state: Dict, config: Optional[RunnableConfig] = None) -> Dict:\n agent_name = getattr(agent_graph, 'name', 'sub_agent')\n print(f\"🟡 [Async invoke] Handoff to agent '{agent_name}'\")\n sub_agent_input = {\"messages\": state.get(\"messages\", [])}\n output: Dict[str, Any] = {}\n agent_error: Optional[str] = None\n\n try:\n output = await agent_graph.ainvoke(sub_agent_input, config=config)\n print(f\"✅ [Async invoke] Agent '{agent_name}' completed.\")\n except Exception as e:\n print(f\"!!! Error during sub-agent {agent_name} ainvoke: {e}\"); traceback.print_exc()\n agent_error = f\"Error executing agent '{agent_name}': {type(e).__name__}\"\n\n sub_agent_messages: List[BaseMessage] = output.get(\"messages\", [])\n returned_messages: List[BaseMessage] = []\n if not sub_agent_messages and not agent_error:\n returned_messages = [AIMessage(content=\"(No output received from agent)\", name=agent_name)]\n elif output_mode == \"last_message\":\n last_ai_message = next((m for m in reversed(sub_agent_messages) if isinstance(m, AIMessage)), None)\n returned_messages = [last_ai_message] if last_ai_message else sub_agent_messages[-1:]\n else:\n returned_messages = sub_agent_messages\n \n last_content = agent_error\n if not last_content and returned_messages:\n last_content = str(returned_messages[-1].content) if hasattr(returned_messages[-1], 'content') else \"(No textual content)\"\n\n return {\n \"messages\": returned_messages,\n \"last_agent_result\": {\n \"agent_name\": agent_name,\n \"content\": last_content or \"(Agent execution finished without specific output or error)\"\n }\n }\n\n def call_agent(state: Dict, config: Optional[RunnableConfig] = None) -> Dict:\n agent_name = getattr(agent_graph, 'name', 'sub_agent')\n print(f\"🟡 [Sync invoke] Handoff to agent '{agent_name}'\")\n sub_agent_input = {\"messages\": state.get(\"messages\", [])}\n output: Dict[str, Any] = {}\n agent_error: Optional[str] = None\n\n try: output = agent_graph.invoke(sub_agent_input, config=config); print(f\"✅ [Sync invoke] Agent '{agent_name}' completed.\")\n except NotImplementedError: agent_error = f\"Error: Sync invoke not supported by agent '{agent_name}'.\"; print(agent_error)\n except Exception as e: agent_error = f\"Error during sub-agent {agent_name} invoke: {e}\"; print(f\"!!! {agent_error}\")\n \n sub_agent_messages: List[BaseMessage] = output.get(\"messages\", [])\n returned_messages: List[BaseMessage] = []\n if not sub_agent_messages and not agent_error: returned_messages = [AIMessage(content=\"(No output received)\", name=agent_name)]\n elif output_mode == \"last_message\":\n last_ai_message = next((m for m in reversed(sub_agent_messages) if isinstance(m, AIMessage)), None)\n returned_messages = [last_ai_message] if last_ai_message else sub_agent_messages[-1:]\n else: returned_messages = sub_agent_messages\n \n last_content = agent_error\n if not last_content and returned_messages: last_content = str(returned_messages[-1].content) if hasattr(returned_messages[-1], 'content') else \"(No content)\"\n\n return {\n \"messages\": returned_messages,\n \"last_agent_result\": {\n \"agent_name\": agent_name,\n \"content\": last_content or \"(Agent sync execution finished)\"\n }\n }\n\n return RunnableCallable(func=call_agent, afunc=acall_agent, name=f\"Call_{getattr(agent_graph, 'name', 'sub_agent')}\")\n\n\ndef supervisor_node_logic_sync(\n state: PlanningAgentState,\n config: Optional[RunnableConfig],\n model: Any,\n supervisor_name: str,\n agent_description_map: Dict[str, str]\n) -> Dict[str, Any]:\n print(f\"--- Entering Supervisor Node (Sync Wrapper) ---\")\n try:\n return anyio.run(\n supervisor_node_logic, state, config, model, supervisor_name, agent_description_map\n )\n except Exception as e:\n print(f\"Error running supervisor_node_logic synchronously using anyio: {e}\")\n import traceback\n traceback.print_exc()\n return {\"error\": f\"Sync execution wrapper failed: {e}\", \"messages\": state.get(\"messages\",[])}\n\n\ndef create_supervisor(\n model: LanguageModelLike,\n sub_agents: List[BaseAgent],\n state_schema: Type[PlanningAgentState] = PlanningAgentState,\n config_schema: Type[Any] | None = None,\n tools: list[BaseTool | Callable] | None = None,\n output_mode: OutputMode = \"last_message\",\n add_handoff_back_messages: bool = False,\n supervisor_name: str = \"supervisor\",\n planner_node_name: str = \"planner\",\n evaluator_node_name: str = \"evaluate_result\",\n handoff_executor_name: str = \"handoff_executor\",\n include_agent_name: AgentNameMode | None = \"inline\",\n) -> StateGraph:\n agent_graphs: Dict[str, Pregel] = {}\n agent_names: List[str] = []\n agent_description_map: Dict[str, str] = {}\n # --- 1. 提取 Agent 信息 ---\n for agent in sub_agents:\n if not isinstance(agent, BaseAgent): raise TypeError(...)\n if not agent.name or agent.name == \"LangGraph\": raise ValueError(...)\n if agent.name in agent_graphs: raise ValueError(...)\n agent_names.append(agent.name)\n agent_description_map[agent.name] = getattr(agent, 'description', '...')\n try:\n compiled_graph = agent.get_agent()\n if not isinstance(compiled_graph, Pregel): \n core_graph = getattr(compiled_graph, 'last', None)\n if isinstance(core_graph, Pregel):\n compiled_graph = core_graph\n else:\n raise TypeError(f\"Could not retrieve Pregel instance from agent '{agent.name}'.get_agent()\")\n agent_graphs[agent.name] = compiled_graph\n except Exception as e: raise e\n\n # --- 2. 创建 Handoff 工具 ---\n handoff_tools = [create_handoff_tool(agent_name=name) for name in agent_names]\n supervisor_callable_tools = (tools or []) + handoff_tools\n print(f\"Supervisor '{supervisor_name}' bound with tools: {[t.name for t in supervisor_callable_tools]}\")\n\n # --- 3. 绑定工具到 Supervisor 模型 ---\n bound_supervisor_model: LanguageModelLike\n if not supervisor_callable_tools:\n print(f\"Warning: Supervisor '{supervisor_name}' has no tools bound.\")\n bound_supervisor_model = model\n elif _supports_disable_parallel_tool_calls(model):\n bound_supervisor_model = model.bind_tools(supervisor_callable_tools, parallel_tool_calls=False)\n else:\n bound_supervisor_model = model.bind_tools(supervisor_callable_tools)\n if include_agent_name:\n bound_supervisor_model = with_agent_name(bound_supervisor_model, include_agent_name)\n\n # --- 4. 构建 StateGraph ---\n builder = StateGraph(state_schema, config_schema=config_schema)\n \n # --- 5. 添加 Planner 节点 (使用同步/异步包装) ---\n planner_logic_partial_async = functools.partial(\n planner_node_logic,\n model=model,\n agent_description_map=agent_description_map,\n )\n planner_logic_partial_sync = functools.partial(\n planner_node_logic_sync,\n model=model,\n agent_description_map=agent_description_map,\n )\n planner_runnable = RunnableCallable(\n func=planner_logic_partial_sync,\n afunc=planner_logic_partial_async,\n name=planner_node_name\n )\n builder.add_node(planner_node_name, planner_runnable)\n\n # --- 6. 添加 Supervisor 节点 (使用同步/异步包装) ---\n supervisor_logic_partial_async = functools.partial(\n supervisor_node_logic,\n model=bound_supervisor_model,\n supervisor_name=supervisor_name,\n agent_description_map=agent_description_map,\n )\n supervisor_logic_partial_sync = functools.partial(\n supervisor_node_logic_sync,\n model=bound_supervisor_model,\n supervisor_name=supervisor_name,\n agent_description_map=agent_description_map,\n )\n supervisor_runnable = RunnableCallable(\n func=supervisor_logic_partial_sync,\n afunc=supervisor_logic_partial_async,\n name=supervisor_name\n )\n builder.add_node(supervisor_name, supervisor_runnable)\n\n # --- 7. 添加子 Agent 节点 ---\n for name, compiled_graph in agent_graphs.items():\n builder.add_node(name, _make_call_agent(compiled_graph, output_mode, add_handoff_back_messages, supervisor_name))\n builder.add_edge(name, evaluator_node_name)\n\n # --- 8. 添加 Handoff Tool 执行节点 ---\n handoff_executor_node = ToolNode(handoff_tools, name=handoff_executor_name)\n builder.add_node(handoff_executor_name, handoff_executor_node)\n \n # --- 9. 添加 Evaluate Result 节点 ---\n evaluator_runnable = RunnableCallable(func=evaluate_result_node_logic_sync, afunc=evaluate_result_node_logic, name=evaluator_node_name)\n # Evaluator 不需要 model 或 agent descriptions 作为直接参数\n builder.add_node(evaluator_node_name, evaluator_runnable) # type: ignore\n # --- 10. 设置图的入口和边 ---\n builder.set_entry_point(planner_node_name)\n builder.add_edge(planner_node_name, supervisor_name)\n\n def route_from_supervisor(state: PlanningAgentState) -> str:\n messages = state.get('messages', [])\n plan = state.get('plan')\n last_message = messages[-1] if messages else None\n\n if not isinstance(last_message, AIMessage):\n print(\"Routing: Last message not AIMessage, looping supervisor.\")\n return supervisor_name\n\n if last_message.tool_calls:\n tool_call = last_message.tool_calls[0]\n agent_name_match = re.match(r\"transfer_to_(\\w+)\", tool_call[\"name\"])\n if agent_name_match and agent_name_match.group(1) in agent_names: \n extracted_name = agent_name_match.group(1)\n print(f\"DEBUG route_from_supervisor: Tool Call Name = {repr(tool_call['name'])}\") \n print(f\"DEBUG route_from_supervisor: Extracted Target Name = {repr(extracted_name)}\") \n print(f\"DEBUG route_from_supervisor: Available Agent Names = {repr(agent_names)}\") \n print(f\"Routing: Supervisor -> HandoffExecutor (for {extracted_name})\")\n return handoff_executor_name\n else:\n print(f\"DEBUG route_from_supervisor: Membership check failed! ('{extracted_name}' in {repr(agent_names)}) is False.\")\n print(f\"Warning: Supervisor called unknown/invalid tool: {tool_call['name']}. Looping supervisor.\")\n return supervisor_name\n\n if plan and plan.get(\"status\") == \"completed\":\n print(\"Routing: Plan completed -> END\")\n return END\n\n print(f\"Routing: No tool call and plan not completed (status: {plan.get('status') if plan else 'None'}). Looping supervisor.\")\n return supervisor_name\n\n builder.add_conditional_edges(\n supervisor_name,\n route_from_supervisor,\n {\n handoff_executor_name: handoff_executor_name,\n supervisor_name: supervisor_name,\n END: END,\n }\n )\n \n # Handoff Executor 完成后, LangGraph 处理 Command(goto=...) 直接路由到子 Agent\n # 不需要从 Handoff Executor 出发的显式边\n\n # --- 关键修改: 子 Agent 完成后 -> Evaluator ---\n for name in agent_names:\n builder.add_edge(name, evaluator_node_name) # <--- 修改: 指向 Evaluator\n\n # --- 新增: Evaluator 完成后 -> Supervisor ---\n builder.add_edge(evaluator_node_name, supervisor_name) # <--- 新增: Evaluator 指回 Supervisor\n\n print(\"Supervisor graph definition created with Planner and Evaluator nodes.\")\n return builder # 返回 StateGraph 定义"
},
{
"path": "core/agents/state_based_supervisor/supervisor_node.py",
"content": "# reason_graph/supervisor_node.py\n\nimport re\nimport json\nimport time\nimport copy\nimport ast \nimport traceback\nfrom typing import Dict, Any, List, Optional, Union, cast\nfrom datetime import datetime \nfrom langchain_core.messages import BaseMessage, AIMessage, SystemMessage, HumanMessage, ToolMessage\nfrom langchain_core.messages import ToolCall # 确保导入\nfrom langchain_core.runnables import RunnableConfig\nfrom langgraph.graph import END\n\n# 内部导入 (确保路径正确)\ntry:\n from .state_schema import PlanningAgentState, TaskStatus, Plan\n from .planning_handler import PlanningStateHandler\n from .prompt import SUPERVISOR_PLANNING_PROMPT_TEMPLATE\nexcept ImportError as e:\n print(f\"Error importing modules in supervisor_node.py: {e}\")\n # Fallbacks\n class PlanningAgentState(Dict): pass\n class Plan(Dict): pass\n class PlanningStateHandler: \n @staticmethod\n def update_task(*args, **kwargs): return kwargs.get('plan')\n @staticmethod\n def create_plan(*args, **kwargs): return {}\n @staticmethod\n def add_tasks(*args, **kwargs): return kwargs.get('plan')\n @staticmethod\n def finish_plan(*args, **kwargs): return kwargs.get('plan')\n @staticmethod\n def get_task(*args, **kwargs): return None\n @staticmethod\n def update_plan_status(*args, **kwargs): return kwargs.get('plan')\n @staticmethod\n def set_current_task(*args, **kwargs): return kwargs.get('plan')\n SUPERVISOR_PLANNING_PROMPT_TEMPLATE = \"Fallback Prompt: Error loading template.\"\n\n\n# --- 参数解析函数 (使用 JSON / ast.literal_eval) ---\ndef parse_directive_args(directive_str: str) -> Dict[str, Any]:\n \"\"\"从指令字符串中解析 JSON 参数\"\"\"\n args = {}\n # 查找第一个 '{' 到最后一个 '}' 之间的内容作为 JSON 字符串\n json_match = re.search(r\"(\\{.*?\\})\\s*$\", directive_str.split(maxsplit=1)[1] if len(directive_str.split(maxsplit=1)) > 1 else \"\", re.DOTALL)\n if json_match:\n args_json_str = json_match.group(1)\n try:\n args = json.loads(args_json_str)\n if not isinstance(args, dict): raise ValueError(\"Args JSON not a dict.\")\n print(f\"DEBUG: Parsed args via JSON: {args}\")\n return args\n except json.JSONDecodeError as json_err:\n print(f\"Warning: JSON parsing failed ({json_err}), trying ast.literal_eval...\")\n try:\n args = ast.literal_eval(args_json_str)\n if not isinstance(args, dict): raise ValueError(\"ast.literal_eval didn't return dict.\")\n print(f\"DEBUG: Parsed args via ast.literal_eval: {args}\")\n return args\n except Exception as ast_err:\n raise ValueError(f\"Failed to parse args: {ast_err}. Raw: '{args_json_str}'\") from ast_err\n elif directive_str.strip().upper().endswith(\"{}\"): # 处理 FINISH_PLAN {} 的情况\n return {} # 返回空字典\n else:\n # 如果找不到有效的 JSON 参数，但指令需要参数，则抛出错误或返回空字典\n print(f\"Warning: Could not find valid JSON arguments in directive: '{directive_str}'. Returning empty args.\")\n return {}\n\n\n# --- Supervisor 节点核心逻辑 (移除结果处理，增加设置 current_task_id) ---\nasync def supervisor_node_logic(\n state: PlanningAgentState,\n config: Optional[RunnableConfig],\n model: Any,\n supervisor_name: str,\n agent_description_map: Dict[str, str]\n) -> Dict[str, Any]:\n \"\"\"Supervisor 节点核心逻辑 (不再处理 Agent 结果状态更新)\"\"\"\n print(f\"--- Entering Supervisor Node ({supervisor_name}) ---\")\n messages: List[BaseMessage] = state.get('messages', [])\n plan: Optional[Plan] = state.get('plan')\n current_error = state.get('error'); state['error'] = None\n if current_error: print(f\" Supervisor saw previous error: {current_error}\")\n\n # --- 0. 检查 Plan 是否存在 (不变) ---\n if not plan:\n print(\"ERROR: Supervisor node requires a plan, but none found in state.\")\n return {\"error\": \"Plan is missing.\", \"messages\": []}\n\n # --- 1. 准备 Prompt (不变) ---\n plan_json_str = json.dumps(plan, indent=2, ensure_ascii=False)\n desc_list = [f\"- {name}: {desc}\" for name, desc in agent_description_map.items()]\n desc_list.append(f\"- {supervisor_name}: Coordinates tasks...\")\n agent_descriptions_str = \"\\n\".join(desc_list)\n system_prompt_text = \"Error loading/formatting prompt\"\n try:\n current_date_str = datetime.now().strftime(\"%a, %b %d, %Y\")\n system_prompt_text = SUPERVISOR_PLANNING_PROMPT_TEMPLATE.format(\n plan_json=plan_json_str, \n agent_descriptions=agent_descriptions_str, \n current_date=current_date_str\n )\n except Exception as e: print(f\"ERROR loading/formatting prompt: {e}\")\n llm_input_messages = [SystemMessage(content=system_prompt_text)] + messages\n\n # --- 2. 调用 Supervisor LLM (不变) ---\n print(\"--- Calling Supervisor LLM ---\"); response=None; llm_error_msg=None\n try: \n response = await model.ainvoke(llm_input_messages, config=config)\n if not isinstance(response, AIMessage): raise TypeError(f\"LLM returned non-AIMessage: {type(response)}\")\n if not response.name: response.name = supervisor_name\n print(f\"Supervisor LLM Raw Response Content: {response.content[:300]}...\")\n if response.tool_calls: print(f\"Supervisor LLM Tool Calls: {response.tool_calls}\")\n messages_to_add = [response]\n except Exception as e: \n print(f\"!!! Error invoking Supervisor LLM: {e}\"); traceback.print_exc()\n llm_error_msg = f\"LLM failed: {e}\"; messages_to_add = []; response = None\n\n # --- 3. 处理 LLM 回复 ---\n plan_updated: bool = False\n updated_plan: Optional[Plan] = copy.deepcopy(plan) # 从当前 plan 开始\n directive_error_msg: Optional[str] = None\n task_id_to_delegate: Optional[str] = None # <-- 存储本轮要委派的任务 ID\n\n if response and isinstance(response.content, str):\n # --- A. 先解析并执行所有 PLAN_UPDATE 指令 (移除 status='completed/failed' 的处理) ---\n try:\n plan_directives = re.findall(r\"PLAN_UPDATE:\\s*(\\w+)\\s*(\\{.*?\\})\\s*$\", response.content, re.IGNORECASE | re.DOTALL | re.MULTILINE)\n plan_directives.extend(re.findall(r\"PLAN_UPDATE:\\s*(FINISH_PLAN)\\s*(\\{\\})\\s*$\", response.content, re.IGNORECASE | re.DOTALL | re.MULTILINE))\n\n if plan_directives:\n print(f\"Found {len(plan_directives)} PLAN_UPDATE directive(s).\")\n for command, args_json_str in plan_directives:\n command = command.upper(); args_json_str = args_json_str if args_json_str else \"{}\"\n print(f\"Processing directive: {command} with args JSON: {args_json_str[:100]}...\")\n try:\n args = json.loads(args_json_str) # 使用 JSON 解析\n if not isinstance(args, dict): raise ValueError(\"Args not dict.\")\n\n # --- 执行规划指令 ---\n if command == \"ADD_TASKS\":\n if not updated_plan: raise ValueError(\"No plan.\"); tasks=args.get(\"tasks\",[])\n if isinstance(tasks, list): \n # 确保新任务状态是 pending\n for task_data in tasks: task_data['status'] = 'pending'\n updated_plan = PlanningStateHandler.add_tasks(updated_plan, tasks); plan_updated = True\n else: raise ValueError(\"Invalid 'tasks'.\")\n\n elif command == \"UPDATE_TASK\":\n if not updated_plan: raise ValueError(\"No plan.\")\n by_id=args.get(\"by_id\")\n if not by_id or not isinstance(by_id, str): raise ValueError(\"Requires string 'by_id'.\")\n by_id = by_id.strip()\n task_exists = PlanningStateHandler.get_task(updated_plan, by_id)\n if not task_exists: raise ValueError(f\"Task ID '{by_id}' not found!\")\n \n # 只处理状态为 'in_progress' 或 其他非终结状态的更新，以及 notes/evaluation\n new_status=args.get(\"status\"); notes_text=args.get(\"notes\"); eval_text=args.get(\"evaluation\") # 保留 evaluation 用于记录 LLM 的想法\n update_kwargs = {}\n # **不再**设置 \"completed\", \"failed\", \"pending_review\"\n if new_status and new_status == \"in_progress\": \n update_kwargs['new_status'] = \"in_progress\"\n task_id_to_delegate = by_id # 记录这个 ID，将在 Handoff 前设置\n # 总是可以更新 notes 和 evaluation (如果 LLM 提供了)\n if notes_text is not None: update_kwargs['new_notes'] = notes_text\n if eval_text is not None: update_kwargs['new_evaluation'] = eval_text \n\n if update_kwargs: # 只有当确实需要更新时才调用\n print(f\"Updating task {by_id} with: {update_kwargs}\")\n updated_plan = PlanningStateHandler.update_task(updated_plan, by_id=by_id, **update_kwargs); plan_updated = True\n\n elif command == \"FINISH_PLAN\":\n if not updated_plan: raise ValueError(\"No plan.\")\n updated_plan = PlanningStateHandler.finish_plan(updated_plan); plan_updated = True\n \n else: print(f\"Warning: Unknown PLAN_UPDATE command '{command}' ignored by Supervisor.\")\n\n except (json.JSONDecodeError, ValueError, KeyError, TypeError) as e:\n err_msg = f\"Error processing plan directive '{command} {args_json_str}': {type(e).__name__} - {e}\"\n print(err_msg); traceback.print_exc()\n if not directive_error_msg: directive_error_msg = err_msg # 只记录第一个错误\n except Exception as e:\n err_msg = f\"Unexpected error processing directive '{command} {args_json_str}': {type(e).__name__} - {e}\"\n print(err_msg); traceback.print_exc()\n if not directive_error_msg: directive_error_msg = err_msg\n \n # --- 重新计算 Plan 状态 ---\n if plan_updated and updated_plan:\n updated_plan = PlanningStateHandler.update_plan_status(updated_plan)\n print(f\"Plan status after updates by Supervisor: {updated_plan.get('status')}\")\n\n except Exception as outer_e:\n err_msg = f\"Error occurred while searching for PLAN_UPDATE directives: {outer_e}\"\n print(err_msg); traceback.print_exc()\n if not directive_error_msg: directive_error_msg = err_msg\n\n # --- B. 检查 Tool Calls 并设置 Current Task ID ---\n handoff_tool_call: Optional[Dict] = None # 显式初始化\n if response and response.tool_calls:\n for tool_call in response.tool_calls:\n agent_name_match = re.match(r\"transfer_to_(\\w+)\", tool_call[\"name\"])\n # **使用 agent_description_map.keys() 来检查**\n if agent_name_match and agent_name_match.group(1) in agent_description_map.keys():\n handoff_tool_call = cast(Dict, tool_call) # 找到第一个有效的就用它\n break\n\n # 如果决定 Handoff，尝试设置 plan 中的 current_task_id\n if handoff_tool_call and updated_plan:\n # **关键**: 尝试从 Tool Call 的 args 中获取 task_id (Prompt 要求 LLM 必须提供)\n tool_args = handoff_tool_call.get(\"args\", {})\n task_id_from_tool = tool_args.get(\"task_id\") if isinstance(tool_args, dict) else None\n \n # 如果 Tool args 中没有，再使用之前记录的 task_id_to_delegate (标记为 in_progress 的)\n effective_task_id = task_id_from_tool or task_id_to_delegate \n\n if effective_task_id:\n print(f\"Setting current_task_id in plan to: {effective_task_id}\")\n try:\n # 验证 ID 存在\n if PlanningStateHandler.get_task(updated_plan, effective_task_id):\n updated_plan = PlanningStateHandler.set_current_task(updated_plan, effective_task_id)\n # plan_updated 标志可能已经被 Plan Directive 设置，这里不需要重复设置\n else:\n print(f\"Warning: Task ID '{effective_task_id}' provided for delegation not found. Cannot set current_task_id.\")\n # 记录错误，阻止 Handoff? 或者让路由回到 Supervisor?\n directive_error_msg = directive_error_msg or f\"Invalid Task ID '{effective_task_id}' for delegation.\"\n except Exception as e:\n err_msg = f\"Error setting current_task_id to '{effective_task_id}': {e}\"\n print(f\"ERROR: {err_msg}\")\n if not directive_error_msg: directive_error_msg = err_msg\n\n # --- 4. 准备最终返回的状态更新字典 ---\n updates: Dict[str, Any] = {\"messages\": messages_to_add}\n if updated_plan is not None: updates[\"plan\"] = updated_plan\n elif plan is not None: updates[\"plan\"] = plan\n \n final_error = llm_error_msg or directive_error_msg\n if final_error: updates[\"error\"] = final_error\n elif state.get(\"error\"): updates[\"error\"] = None # 清除旧错误\n\n print(f\"--- Exiting Supervisor Node. Plan updated this step: {plan_updated} ---\")\n return updates"
},
{
"path": "core/agents/sub_agents/__init__.py",
"content": ""
},
{
"path": "core/agents/sub_agents/coder_agent.py",
"content": "# Refactored coder_agent.py\nfrom typing import Any, List, Optional, Union, Callable, Type\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import SystemMessage\nfrom langgraph.types import Checkpointer\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_tools_by_category, ToolCategory, get_tool_instance # Import get_tool_instance\n\nimport logging\nlogger = logging.getLogger(__name__)\n\nclass CoderAgent(ReactAgent):\n \"\"\"\n Coder Agent (Refactored)\n - Interacts with a sandboxed Linux environment via code execution tools.\n \"\"\"\n\n def __init__(\n self,\n name: str = \"coder_expert\",\n model: LanguageModelLike = None,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = 100000, # Coding might need more context\n **kwargs\n ):\n # 1. Define Description\n description = \"Writes, executes, tests, and debugs Python code and Linux shell commands within a secure sandboxed environment. Can install packages, manage files, and interact with the network.\"\n\n # 2. Get Tools from Registry\n agent_tools = []\n default_tool_name = \"e2b_code_interpreter\" # Expected tool name\n try:\n code_tools = get_tools_by_category(ToolCategory.CODE_INTERPRETER) + get_tools_by_category(ToolCategory.FILE_SYSTEM)\n\n agent_tools.extend(code_tools)\n # Optionally add file system tools if not included in interpreter tool\n # fs_tools = get_tools_by_category(ToolCategory.FILE_SYSTEM)\n # agent_tools.extend(fs_tools)\n print(f\"[{name}] Loaded tools from registry: {[t.name for t in agent_tools if hasattr(t,'name')]}\")\n # Verify the main execution tool is present\n if not any(getattr(t,'name', None) == default_tool_name for t in agent_tools):\n print(f\"CRITICAL Warning: CoderAgent '{name}' is missing the primary '{default_tool_name}' tool!\")\n # Attempt to get it specifically if missing?\n specific_tool = get_tool_instance(default_tool_name)\n if specific_tool: agent_tools.append(specific_tool)\n\n except Exception as e:\n print(f\"Warning: Failed to get tools from registry for {name}: {e}\")\n\n if tools: # Merge extra tools\n existing_names = {t.name for t in agent_tools if hasattr(t,'name')}\n agent_tools.extend([t for t in tools if getattr(t, 'name', None) not in existing_names])\n\n if not agent_tools:\n print(f\"CRITICAL Warning: CoderAgent '{name}' initialized with NO tools!\")\n\n # 3. Define System Prompt (using the capabilities)\n tool_name_for_prompt = next((t.name for t in agent_tools if hasattr(t, 'name') and 'code' in t.name.lower()), default_tool_name) # Try to get actual tool name\n\n base_prompt = f\"\"\"You are an expert Coder Agent interacting with a secure, sandboxed Linux environment provided by the '{tool_name_for_prompt}' tool. Your goal is to fulfill coding, file manipulation, or shell command requests by generating and executing appropriate code or commands within this sandbox.\n\nAvailable Tools:\n{self._format_tools_for_prompt(agent_tools)}\n- **{tool_name_for_prompt}**: Executes Python code or shell commands within the sandboxed Linux environment. Returns stdout, stderr, execution errors, and potentially file outputs or structured results (like image data). To run shell commands, generate Python code that uses the 'subprocess' module OR if the tool directly supports it, prefix the command with '!'. Always prefer generating Python code for complex shell operations or when needing output capture.\n\nKey Capabilities of the Sandbox Environment (via the tool):\n- Execute Python 3 code.\n- Install Python packages using pip (generate code like `import subprocess; subprocess.run(['pip', 'install', 'requests'], check=True)`).\n- Run standard Linux shell commands (e.g., `ls`, `pwd`, `mkdir`, `curl`, `git`, etc. using Python's subprocess).\n- Access and manipulate a persistent filesystem within the sandbox (typically starting in `/home/user/` or `/`). Create, read, write, delete files and directories.\n- Access the internet from within the sandbox for tasks like cloning repos or fetching data.\n\nWorkflow & Instructions:\n1. **Analyze Request**: Understand the goal, constraints, and required inputs/outputs.\n2. **Plan Steps**: Outline the necessary code or commands. Consider file paths, dependencies, and error handling.\n3. **Generate Code/Command**: Write the Python code or shell command sequence needed. For non-trivial Python, include comments.\n4. **Execute using Tool**: Prepare the arguments for the '{tool_name_for_prompt}' tool (usually the code string or command string) and invoke the tool.\n5. **Analyze Output**: Carefully review the stdout, stderr, errors, and any results returned by the tool.\n6. **Debug/Iterate**: If errors occurred or the output is not as expected, analyze the error, revise the code/command, and execute again using the tool.\n7. **Final Output**: Once the task is successfully completed, provide the final working code (if relevant), a summary of the execution results (stdout/stderr highlights), confirmation of file operations, and any requested explanation. If the task cannot be completed, explain why.\n8. **File Handling**: If generating files (code, data, images), clearly state the full path within the sandbox where the file was saved (e.g., `/home/user/my_script.py`, `/home/user/output.csv`). Do not attempt to display images directly in your response.\n\nFocus strictly on tasks achievable within the sandboxed environment using the provided tool. Be precise and careful with file paths and commands.\n\"\"\"\n\n # 4. Call super().__init__\n super().__init__(\n name=name,\n model=model,\n tools=agent_tools,\n prompt=base_prompt,\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n **kwargs\n )\n print(f\"CoderAgent '{self.name}' initialized with tools: {[t.name for t in self.tools if hasattr(t,'name')]}\")\n\n # Inherits _format_tools_for_prompt and other methods from BaseAgent/ReactAgent"
},
{
"path": "core/agents/sub_agents/data_analyst_agent.py",
"content": "# data_analyst_agent.py (or in main.py)\n\nfrom typing import Any, List, Optional, Union, Callable, Type\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import SystemMessage\nfrom langgraph.types import Checkpointer\n\n# Internal imports - ensure paths are correct\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_tools_by_category, ToolCategory, get_tool_instance # Import necessary functions\n\nimport logging\nlogger = logging.getLogger(__name__)\n\n# Assume ToolCategory.CODE_INTERPRETER exists\n# Assume ToolCategory.FILE_SYSTEM exists if needed\n\nclass DataAnalystAgent(ReactAgent):\n \"\"\"\n Data Analyst Agent (Refactored)\n - Focuses on analyzing structured data using code execution sandbox.\n - Generates insights and saves visualizations to files.\n \"\"\"\n\n def __init__(\n self,\n name: str = \"data_analyst_expert\",\n model: LanguageModelLike = None,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = 120000, # Analysis might need decent context\n debug: bool = False,\n **kwargs\n ):\n # 1. Define Description for Supervisor\n description = \"Analyzes structured data (provided in context or potentially read from sandbox files) using Python (Pandas, NumPy, Matplotlib, Seaborn) within a secure code execution environment. Performs statistical analysis, identifies trends, generates insights, and creates data visualizations (saved as files in the sandbox).\"\n\n # 2. Get Tools from Registry\n agent_tools = []\n default_tool_name = \"e2b_code_interpreter\" # Tool needed for execution\n try:\n # Primarily needs Code Interpreter\n code_tools = get_tools_by_category(ToolCategory.CODE_INTERPRETER) + get_tools_by_category(ToolCategory.FILE_SYSTEM) # 需要代码和文件工具\n\n agent_tools.extend(code_tools)\n # Optionally, add File System tools if needed to read data files\n # fs_tools = get_tools_by_category(ToolCategory.FILE_SYSTEM)\n # agent_tools.extend(fs_tools)\n print(f\"[{name}] Loaded tools from registry: {[t.name for t in agent_tools if hasattr(t,'name')]}\")\n # Verify the main execution tool is present\n if not any(getattr(t,'name', None) == default_tool_name for t in agent_tools):\n print(f\"CRITICAL Warning: DataAnalystAgent '{name}' is missing the primary '{default_tool_name}' tool!\")\n specific_tool = get_tool_instance(default_tool_name)\n if specific_tool: agent_tools.append(specific_tool)\n\n except Exception as e:\n print(f\"Warning: Failed to get tools from registry for {name}: {e}\")\n\n if tools: # Merge extra tools\n existing_names = {t.name for t in agent_tools if hasattr(t,'name')}\n agent_tools.extend([t for t in tools if getattr(t, 'name', None) not in existing_names])\n\n if not agent_tools:\n print(f\"CRITICAL Warning: DataAnalystAgent '{name}' initialized with NO execution tools!\")\n\n # 3. Define System Prompt\n tool_name_for_prompt = next((t.name for t in agent_tools if hasattr(t, 'name') and 'code' in t.name.lower()), default_tool_name)\n\n base_prompt = f\"\"\"You are an expert Data Analyst. Your task is to analyze data using Python code within a secure sandbox environment accessed via the '{tool_name_for_prompt}' tool. Libraries like Pandas, NumPy, Matplotlib, and Seaborn are available (install if needed using pip in your code).\n\nAvailable Tools:\n{self._format_tools_for_prompt(agent_tools)}\n- **{tool_name_for_prompt}**: Executes Python code in the sandbox. Returns stdout, stderr, errors, and potentially structured results.\n\nKey Instructions:\n1. **Understand Data & Goal**: Identify the data source (likely provided in previous messages or mentioned as a sandbox file path like '/home/user/data.csv') and the specific analysis question or goal.\n2. **Plan Analysis**: Briefly outline the Python code steps (e.g., load data into Pandas DataFrame, clean/transform data, perform calculations, generate plot).\n3. **Write Python Code**: Generate the necessary Python code. Use libraries effectively. Import necessary libraries (e.g., `import pandas as pd`, `import matplotlib.pyplot as plt`).\n4. **Handle Files (If Needed)**: If reading/writing files within the sandbox, use standard Python file I/O within your code (e.g., `pd.read_csv('/home/user/data.csv')`, `df.to_csv('/home/user/output.csv')`).\n5. **Handle Visualizations**: If asked to create plots:\n * Generate the plot using Matplotlib/Seaborn.\n * **MUST save the plot to a file** inside the sandbox (e.g., `/home/user/plots/my_plot.png`). Use `plt.savefig('/home/user/plots/my_plot.png')`. Create directories if necessary (`os.makedirs('/home/user/plots', exist_ok=True)`).\n * Use `plt.show()` or `plt.close()` after saving to clear the plot buffer.\n * **DO NOT attempt to return image data directly.** Images cannot be displayed in the response.\n * In your response, **state that the plot was generated and provide the full path** where it was saved in the sandbox (e.g., \"I have generated a scatter plot and saved it to /home/user/plots/scatter_plot.png\").\n6. **Execute Code**: Use the '{tool_name_for_prompt}' tool to run your complete Python script.\n7. **Analyze Results**: Interpret the output (stdout, numerical results, errors) from the tool execution.\n8. **Present Findings**: Summarize your analysis and findings clearly. Use Markdown tables for structured data if helpful. Mention any plots saved and their paths. If errors occurred, explain them.\n9. **Focus**: Concentrate on data analysis using code execution. Do not perform web searches unless specifically instructed and given tools for it.\n\"\"\"\n\n # 4. Call super().__init__\n super().__init__(\n name=name,\n model=model,\n tools=agent_tools,\n prompt=base_prompt,\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n debug=debug,\n **kwargs\n )\n print(f\"DataAnalystAgent '{self.name}' initialized.\")\n\n # Inherits _format_tools_for_prompt and other methods"
},
{
"path": "core/agents/sub_agents/designer_agent.py",
"content": "# 文件路径示例: reason_graph/designer_agent.py\n\nfrom typing import Any, List, Optional, Union, Callable, Type\nfrom langchain_core.language_models import LanguageModelLike # 确保导入正确类型\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import SystemMessage\nfrom langgraph.types import Checkpointer\n\n# 内部导入\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_tools_by_category, ToolCategory # 导入 Registry\n# 假设您的 Flux 工具已注册或在此导入\n# from core.tools.flux_image_tool import FluxImageGeneratorTool \n\nimport logging\nlogger = logging.getLogger(__name__)\n\n# 假设的 ToolCategory.IMAGE_GENERATION\nif not hasattr(ToolCategory, 'IMAGE_GENERATION'):\n ToolCategory.IMAGE_GENERATION = ToolCategory.OTHER\n\nclass DesignerAgent(ReactAgent):\n \"\"\"\n 设计 Agent (重构版)\n - 能够理解图像上下文，并使用工具生成新的视觉内容。\n - 应用设计原则来完成海报、网页等设计任务。\n \"\"\"\n\n def __init__(\n self,\n name: str = \"designer_expert\",\n model: LanguageModelLike = None, # <--- 必须传入多模态模型 (e.g., gpt-4o)\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = 8000, # 调整上下文需求\n debug: bool = False,\n **kwargs\n ):\n # 1. 定义 Agent 描述\n description = \"Understands images provided in context and generates new visual content (images, mockups, diagrams) using specialized image generation tools (like Flux). Can apply design thinking for tasks like poster or web page layout design.\"\n\n # 2. 获取工具 (主要是图像生成工具)\n agent_tools = []\n try:\n # 从 Registry 获取图像生成工具\n img_tools = get_tools_by_category(ToolCategory.IMAGE_GENERATION)\n agent_tools.extend(img_tools)\n # 也可以直接实例化\n # agent_tools.append(FluxImageGeneratorTool()) # 如果不使用 Registry\n print(f\"[{name}] Loaded tools: {[t.name for t in agent_tools if hasattr(t,'name')]}\")\n except Exception as e:\n print(f\"Warning: Failed to get IMAGE_GENERATION tools for {name}: {e}\")\n\n if tools: # 合并额外工具\n existing_names = {t.name for t in agent_tools if hasattr(t,'name')}\n agent_tools.extend([t for t in tools if getattr(t, 'name', None) not in existing_names])\n\n if not agent_tools:\n print(f\"CRITICAL Warning: DesignerAgent '{name}' initialized with NO generation tools!\")\n\n # 3. 定义 System Prompt\n tool_name_for_prompt = next((t.name for t in agent_tools if hasattr(t, 'name') and 'generat' in t.name.lower()), \"image_generator_tool\") # 获取工具名\n\n base_prompt = f\"\"\"You are an expert Visual Designer and Creative Assistant. Your capabilities include understanding images provided in the conversation history and generating new images using available tools based on detailed text prompts.\n\nAvailable Tools:\n{self._format_tools_for_prompt(agent_tools)}\n- **{tool_name_for_prompt}**: Use this tool to generate images. Input requires a detailed 'prompt'.\n\nKey Instructions & Workflow:\n\n1. **Understand Request**: Analyze the user request, paying attention to both text and any images provided in the message history. Identify the core visual goal (e.g., analyze image, generate image, design layout).\n2. **Image Understanding (If Applicable)**: If the request involves analyzing or describing an existing image from the history, provide your analysis directly based on your multimodal understanding.\n3. **Design Thinking (For Generation/Design Tasks)**:\n * **Clarify**: If the request is vague (e.g., \"design a logo\"), think about necessary elements: target audience, brand feeling, key symbols, color preferences, desired style (minimalist, vintage, futuristic, etc.). You might need to state assumptions if details are missing.\n * **Conceptualize**: Describe the visual elements, layout, color palette, and overall composition you plan to generate.\n * **Formulate Prompt for Tool**: Translate your design concept into a **highly detailed and descriptive text prompt** suitable for the `{tool_name_for_prompt}`. Include style, mood, composition, colors, and specific objects.\n4. **Use Generation Tool**: Call the `{tool_name_for_prompt}` with the detailed prompt you formulated.\n5. **Present Result**:\n * State that you have generated the image.\n * Provide the result from the tool (e.g., the image URL or identifier).\n * Briefly describe the generated image and how it matches the design concept or request.\n * **Important**: Do NOT attempt to display the image directly in your text response. Only provide the URL or description.\n6. **Handle Errors**: If the tool fails, report the error clearly.\n\nFocus on visual design and generation tasks. Use your understanding of design principles when conceptualizing visuals for requests like posters or web mockups.\n\"\"\"\n\n # 4. 调用父类 __init__\n super().__init__(\n name=name,\n model=model, # 必须是多模态模型\n tools=agent_tools,\n prompt=base_prompt,\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n debug=debug,\n **kwargs\n )\n print(f\"DesignerAgent '{self.name}' initialized.\")\n\n # 继承 _format_tools_for_prompt 和其他 BaseAgent/ReactAgent 方法"
},
{
"path": "core/agents/sub_agents/reporter_agent.py",
"content": "# 文件路径: reason_graph/reporter_agent.py\n\nimport json\nimport time\nfrom datetime import datetime\nfrom typing import Dict, Any, List, Optional, Union, Type, cast, Sequence\n\n# --- LangChain / LangGraph ---\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import SystemMessage, HumanMessage, BaseMessage, AIMessage\nfrom langchain_core.runnables import RunnableConfig, Runnable\nfrom langgraph.graph import StateGraph, END, START # 导入 StateGraph, END, START\nfrom langgraph.graph.graph import CompiledGraph\nfrom langgraph.types import Checkpointer\n\n# --- 内部导入 ---\nfrom core.agents.base.base_agent import BaseAgent # 导入最终版 BaseAgent\n# 导入最终报告的 Prompt 模板\n\nimport logging\nlogger = logging.getLogger(__name__)\n\nclass ReporterAgent(BaseAgent):\n \"\"\"\n 报告 Agent (最终版)\n - 继承自 BaseAgent。\n - 负责基于完整的消息历史和明确指令生成最终 Markdown 报告。\n - 内部包含一个简单的图用于执行报告生成任务。\n \"\"\"\n\n FINAL_REPORT_SYSTEM_PROMPT_TEMPLATE = \"\"\"You are a professional writer and editor AI assistant. Your primary goal is to generate high-quality, well-structured text content based on the specific instructions provided in the latest message and the relevant information available in the preceding conversation history.\n\nThe current date is {current_date}.\n\n**Your Task Execution Workflow:**\n1. **Identify Instructions:** Carefully read the **last message** you received, which contains the specific writing task assigned to you by the supervisor. Understand the desired output (e.g., summary, report section, full report), format, tone, and any other requirements.\n2. **Gather Context:** Review the preceding messages in the conversation history to find the necessary information, data points, findings, or creative elements needed to complete the assigned task.\n3. **Compose Output:** Write the text according to the instructions.\n * If asked for creative content (like a poem), focus on fulfilling the creative request.\n * If asked for a summary or section, synthesize the relevant information concisely and accurately.\n * If asked to compile a **full report**, structure it logically (e.g., Introduction, Body, Conclusion), use Markdown formatting effectively, and incorporate information/citations from the history as instructed. Adhere to any specified length or style guidelines.\n4. **Final Response:** Your output should be **only** the requested written text. Do not add extra conversational phrases unless necessary for context. Do not include planning directives or attempt to call tools (unless a specific writing/editing tool was provided and instructed for use). If you cannot fulfill the request due to missing information in the history, state that clearly.\n\"\"\"\n\n\n def __init__(\n self,\n name: str = \"reporter_expert\",\n model: LanguageModelLike = None, # 应传入适合长文本生成的模型\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = 16000, # 报告生成可能需要处理长上下文\n debug: bool = False,\n prompt_template: str = FINAL_REPORT_SYSTEM_PROMPT_TEMPLATE, # 使用最终报告模板\n **kwargs # 接收其他 BaseAgent 参数\n ):\n # 1. 定义 Agent 描述 (给 Supervisor 看)\n description = \"Synthesizes information from the complete conversation history and task results into a final, comprehensive, well-structured, and potentially cited Markdown research report, following specific instructions.\"\n\n # 2. 定义工具列表 (Reporter 通常不需要工具)\n agent_tools = []\n\n # 3. 存储基础 Prompt 模板 (将在节点逻辑中使用)\n # 注意：我们将模板本身（或其引用）存储起来，而不是格式化后的 prompt\n self.report_prompt_template = prompt_template\n\n # 4. 调用父类 __init__\n super().__init__(\n name=name,\n model=model, # 传入用于报告生成的 LLM\n tools=agent_tools,\n prompt=None, # BaseAgent 的 prompt 字段不直接用于此 Agent 的核心逻辑\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n # **kwargs 传递 debug 等\n **kwargs\n )\n print(f\"ReporterAgent '{self.name}' initialized.\")\n\n\n async def _generate_report_node_logic(self, state: Dict[str, Any], config: RunnableConfig) -> Dict[str, Any]:\n \"\"\"报告生成节点的核心逻辑\"\"\"\n # 注意：这里的 state 已经是经过 BaseAgent._preprocess_state 处理后的状态\n print(f\"--- Entering Node: {self.name}._generate_report_node_logic ---\")\n\n messages: List[BaseMessage] = state.get(\"messages\", [])\n # 理论上，所有需要的信息都应该在 messages 历史中，\n # 特别是 Supervisor 委派时的最后一条指令消息。\n\n if not messages:\n error_msg = \"Error: No messages found in state for report generation.\"\n print(error_msg)\n return {\"messages\": [AIMessage(content=f\"# Report Generation Failed\\n\\n{error_msg}\", name=self.name)]}\n\n # --- 格式化 System Prompt (包含日期) ---\n try:\n current_date_str = datetime.now().strftime(\"%a, %b %d, %Y\")\n system_prompt = self.report_prompt_template.format(current_date=current_date_str)\n except Exception as e:\n print(f\"Error formatting report system prompt: {e}\")\n system_prompt = \"You are a report writing assistant. Synthesize the provided messages into a final report.\" # Fallback\n\n # --- 准备 LLM 输入 ---\n # 输入是 System Prompt + 完整的、经过预处理（截断）的消息历史\n # BaseAgent 的 _preprocess_state 已经处理了截断\n llm_input_messages = [SystemMessage(content=system_prompt)] + messages\n\n # --- 调用 LLM 生成报告 ---\n final_report_markdown = \"\"\n llm_error = None\n try:\n print(f\"--- Calling LLM for Final Report Generation ({self.name}) ---\")\n # 使用 self.model (初始化时传入的 LLM 实例)\n response = await self.model.ainvoke(llm_input_messages, config=config)\n final_report_markdown = response.content\n print(f\"--- Report Generation LLM Call Successful ({self.name}). Length: {len(final_report_markdown)} chars ---\")\n except Exception as e:\n print(f\"!!! Error during Report Generation LLM call ({self.name}): {e}\")\n llm_error = f\"Report generation failed due to LLM error: {e}\"\n final_report_markdown = f\"# Report Generation Failed\\n\\nError: {str(e)}\"\n # 可以在这里打印更详细的 traceback\n # import traceback\n # traceback.print_exc()\n\n # --- 返回包含报告或错误的状态更新 ---\n # Reporter 的最终输出就是报告本身，放入 messages 中，替换掉历史？\n # 不，应该追加，让调用者（Supervisor 或 main）能看到完整历史和最终报告\n # 使用 AIMessage 返回报告\n return {\n \"messages\": [AIMessage(content=final_report_markdown, name=self.name)],\n \"error\": state.get(\"error\") or llm_error # 保留或记录错误\n }\n\n def build(self) -> Optional[StateGraph]:\n \"\"\"构建 Reporter Agent 的简单工作流： Start -> GenerateReport -> End \"\"\"\n if self._workflow: return self._workflow\n\n print(f\"Building internal graph for ReporterAgent '{self.name}'\")\n # Reporter 通常使用 BasicAgentState，因为它不直接操作 Plan\n # 但为了兼容 Supervisor 可能传递 PlanningAgentState，这里可以暂时用 Any\n # 或者定义一个 ReporterState\n workflow = StateGraph(Dict[str, Any]) # 使用通用字典状态，因为它只关心 messages\n\n # 添加报告生成节点，确保它能访问 self.model\n # functools.partial 不能直接用于异步实例方法，需要包装\n async def node_wrapper(state, config):\n return await self._generate_report_node_logic(state, config)\n\n workflow.add_node(\"generate_report\", node_wrapper) # type: ignore\n workflow.add_edge(START, \"generate_report\")\n workflow.add_edge(\"generate_report\", END)\n\n self._workflow = workflow\n return workflow\n\n # compile 方法继承自 BaseAgent\n # 它会调用上面的 build() 获取 StateGraph 定义，然后编译它，\n # 并创建包含预处理步骤 (_preprocess_state) 的最终 _executable_agent\n\n # invoke, ainvoke, get_agent (get_executable_agent), reset 继承自 BaseAgent"
},
{
"path": "core/agents/sub_agents/research_agent.py",
"content": "# 文件路径示例: reason_graph/research_agent.py\n\nfrom typing import Any, List, Optional, Union, Callable, Type, cast\nfrom langchain_core.language_models import LanguageModelLike\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import SystemMessage\nfrom langgraph.types import Checkpointer\n\n# 内部导入 - 请确保路径正确\nfrom core.agents.base.react_agent import ReactAgent\n# 导入工具 Registry 相关 - 只需要 get_tools_by_category 和 ToolCategory\nfrom core.tools.registry import get_tools_by_category, ToolCategory\n# *** 不再需要导入 get_tool 或 get_registered_tools ***\n\nimport logging\nlogger = logging.getLogger(__name__)\n\n# 假设 ToolCategory 包含 SEARCH 和 WEB_Browse\nif not hasattr(ToolCategory, 'SEARCH'): ToolCategory.SEARCH = ToolCategory.OTHER\nif not hasattr(ToolCategory, 'WEB_Browse'): ToolCategory.WEB_Browse = ToolCategory.OTHER\n\n\nclass ResearchAgent(ReactAgent):\n \"\"\"\n 研究 Agent (重构版)\n - 继承自新的 ReactAgent\n - 专注于定义自身工具和 Prompt\n - 移除了自定义的状态管理和方法\n \"\"\"\n\n def __init__(\n self,\n name: str = \"research_expert\",\n model: LanguageModelLike = None,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = 8000,\n debug: bool = False,\n **kwargs\n ):\n\n # 1. 定义 Agent 描述 (不变)\n description = \"Expert at finding, extracting, and synthesizing the latest information, data, and background knowledge on specific topics using search engines (like Tavily, Google Search) and web Browse tools (like Firecrawl, Arxiv). Capable of providing source links and content summaries.\"\n\n # 2. --- 从 Registry 获取和合并工具 ---\n agent_tools: List[Union[BaseTool, Callable]] = []\n search_tools_loaded: List[Union[BaseTool, Callable]] = [] # 用于后续检查\n Browse_tools_loaded: List[Union[BaseTool, Callable]] = []\n\n try:\n search_tools_loaded = get_tools_by_category(ToolCategory.SEARCH)\n agent_tools.extend(search_tools_loaded)\n try:\n Browse_tools_loaded = get_tools_by_category(ToolCategory.WEB_Browse)\n agent_tools.extend(Browse_tools_loaded)\n except Exception as e:\n if debug: print(f\"[{name}] Info: Failed to get WEB_Browse tools: {e}\")\n print(f\"[{name}] Loaded tools from registry: {[t.name for t in agent_tools if hasattr(t,'name')]}\")\n\n # --- 简化核心工具检查 ---\n if not search_tools_loaded: # 直接检查从 Registry 加载的搜索工具列表是否为空\n print(f\"CRITICAL Warning: ResearchAgent '{name}' initialized without any SEARCH tools from registry!\")\n # ------------------------\n\n except Exception as e:\n print(f\"Warning: Failed to get tools from registry for {name}: {e}\")\n\n # 合并外部传入的 `tools` 参数 (逻辑不变)\n if tools:\n # ... (合并逻辑不变) ...\n existing_tool_names = {t.name for t in agent_tools if hasattr(t, 'name')}\n added_external_count = 0\n for tool in tools:\n tool_name = getattr(tool, 'name', None)\n if tool_name and tool_name not in existing_tool_names:\n agent_tools.append(tool)\n existing_tool_names.add(tool_name)\n added_external_count +=1\n elif not tool_name: \n agent_tools.append(tool)\n added_external_count += 1\n if added_external_count > 0: print(f\"[{name}] Merged {added_external_count} external tool(s).\")\n\n\n # --- 简化最终工具检查 ---\n if not agent_tools:\n print(f\"CRITICAL Warning: ResearchAgent '{name}' initialized with NO tools configured!\")\n # 不再需要那个复杂的 any(...) 检查\n # ----------------------\n\n # 3. 定义 Agent 的 System Prompt (逻辑不变)\n base_prompt = f\"\"\"You are a professional Research Analyst expert...\nAvailable Tools:\n{self._format_tools_for_prompt(agent_tools)} \nInstructions:\n\n- Analyze the request in the message history.\n\n- If the request requires searching for current information, facts, data, or background knowledge, you MUST use one of your search tools (like 'tavily_search_results').\n\n- When using tools, formulate concise and effective search queries based on the request.\n\n- Synthesize the information found from the tools into a clear and informative answer.\n\n- If you use information from a tool, cite the source implicitly in your response (e.g., \"According to [Source Title], ...\").\n\n- If the initial search is insufficient, analyze the results and decide if further searches with refined queries or different tools are needed.\n\n- If you cannot find the information after thorough searching, or if the tools return errors, clearly state the limitations encountered. Do not invent information.\n\"\"\"\n\n # 4. 调用父类 __init__ (逻辑不变)\n super().__init__(\n name=name,\n model=model,\n tools=agent_tools,\n prompt=base_prompt,\n description=description,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n debug=debug,\n **kwargs\n )\n print(f\"ResearchAgent '{self.name}' initialized with final tools: {[t.name for t in self.tools if hasattr(t,'name')]}\")\n"
},
{
"path": "core/llm/llm_manager.py",
"content": "# reason_graph/llm_manager.py\nimport os\nfrom enum import Enum, auto\nfrom typing import Any, Dict, List, Optional, Type, Union, Callable, Tuple\nfrom langchain_core.language_models import BaseChatModel, LanguageModelLike\nfrom langchain_openai import ChatOpenAI\n# (移除 ChatGroq 导入)\n\nfrom dotenv import load_dotenv\n\n# 加载环境变量\nload_dotenv()\n\nclass ModelType(Enum):\n \"\"\"模型提供商类型枚举\"\"\"\n OPENAI = auto()\n XAI = auto()\n DEEPSEEK = auto()\n CUSTOM = auto() # 保持用于其他 OpenAI 兼容 API\n\nclass ModelCapability(Enum):\n \"\"\"模型能力枚举\"\"\"\n GENERAL = auto(); PLANNING = auto(); REASONING = auto()\n CREATIVE = auto(); RESEARCH = auto(); CODE = auto()\n LONG_CONTEXT = auto()\n\nclass LLMManager:\n \"\"\"\n 模型管理器 (融合版 V2)\n - 在初始化时根据配置自动注册模型。\n - 支持按能力获取模型。\n - 支持延迟实例化。\n - 从环境变量加载 API Keys/Base URLs。\n \"\"\"\n\n def __init__(self):\n \"\"\"初始化模型管理器，加载配置并自动注册模型\"\"\"\n self._models_config: Dict[str, Dict[str, Any]] = {}\n self._models_instance: Dict[str, BaseChatModel] = {}\n self._default_model_id: Optional[str] = None\n self._capability_models: Dict[ModelCapability, str] = {}\n\n # 加载 API Keys 和 Base URLs (保持不变)\n self._loaded_api_keys = {\n ModelType.OPENAI: os.getenv(\"OPENAI_API_KEY\"),\n ModelType.XAI: os.getenv(\"XAI_API_KEY\"),\n ModelType.DEEPSEEK: os.getenv(\"DEEPSEEK_API_KEY\"),\n ModelType.CUSTOM: os.getenv(\"LLM_API_KEY\"),\n }\n self._loaded_base_urls = {\n ModelType.OPENAI: os.getenv(\"OPENAI_BASE_URL\"),\n ModelType.XAI: os.getenv(\"XAI_BASE_URL\"),\n ModelType.DEEPSEEK: os.getenv(\"DEEPSEEK_BASE_URL\", \"https://api.deepseek.com/v1\"),\n ModelType.CUSTOM: os.getenv(\"LLM_BASE_URL\"),\n }\n print(\"LLMManager initialized.\")\n print(\"Loaded API Keys for:\", [k.name for k, v in self._loaded_api_keys.items() if v])\n print(\"Loaded Base URLs for:\", {k.name: v for k, v in self._loaded_base_urls.items() if v})\n\n # --- 自动注册模型 ---\n try:\n from .model_config import SUPPORTED_MODELS_CONFIG # 从配置文件导入\n \n print(\"Registering models from config...\")\n for model_id, config in SUPPORTED_MODELS_CONFIG.items():\n # 检查所需 Key/URL 是否存在，如果不存在则跳过注册并警告\n model_type = config.get(\"model_type\")\n api_key = config.get(\"config_override\", {}).get(\"api_key\") or self._loaded_api_keys.get(model_type)\n base_url = config.get(\"config_override\", {}).get(\"base_url\") or self._loaded_base_urls.get(model_type)\n \n # OpenAI 可以只依赖 OPENAI_API_KEY 环境变量\n if model_type == ModelType.OPENAI and not api_key:\n api_key = os.getenv(\"OPENAI_API_KEY\") # 再次检查 OpenAI 专用 Key\n\n # 对于需要 Key 的类型进行检查\n key_required = model_type not in [ModelType.CUSTOM] # 假设 CUSTOM 可能匿名\n url_required = model_type in [ModelType.XAI, ModelType.CUSTOM] # DeepSeek 有默认值\n\n if key_required and not api_key:\n print(f\" Skipping registration for '{model_id}': Required API key for type '{model_type.name}' not found.\")\n continue\n if url_required and not base_url:\n print(f\" Skipping registration for '{model_id}': Required Base URL for type '{model_type.name}' not found.\")\n continue\n\n # 调用内部注册方法\n self._register_model(\n model_id=model_id,\n model_type=config[\"model_type\"],\n model_name=config[\"model_name\"],\n model_class=config.get(\"model_class\"), # 可能为 None\n capabilities=config.get(\"capabilities\", [ModelCapability.GENERAL]),\n set_as_default=config.get(\"is_default\", False),\n config_override=config.get(\"config_override\"),\n **config.get(\"kwargs\", {})\n )\n print(\"Model registration complete.\")\n # 可以在这里设置一个环境变量的默认模型 ID，如果配置中没有 is_default=True\n if not self._default_model_id and self._models_config:\n fallback_default = list(self._models_config.keys())[0]\n print(f\"Warning: No default model marked in config. Falling back to first registered: '{fallback_default}'\")\n self._default_model_id = fallback_default\n\n\n except ImportError:\n print(\"Warning: Could not import model_config.py. No models registered automatically.\")\n except Exception as e:\n print(f\"Error during automatic model registration: {e}\")\n\n print(f\"Default model set to: {self._default_model_id}\")\n print(f\"Capability mapping: {self.list_capabilities()}\")\n print(\"-\" * 20)\n\n\n # register_model 现在是内部方法\n def _register_model(\n self, model_id: str, model_type: ModelType, model_name: str,\n model_class: Optional[Type[BaseChatModel]] = None,\n capabilities: List[ModelCapability] = [ModelCapability.GENERAL],\n set_as_default: bool = False,\n config_override: Optional[Dict[str, Any]] = None,\n **kwargs\n ) -> None:\n \"\"\"(Internal) Registers a model configuration.\"\"\"\n if model_id in self._models_config:\n # Decide on behavior: overwrite or ignore? Let's overwrite with warning.\n print(f\" Overwriting registration for existing model_id: '{model_id}'\")\n # pass # If ignore is preferred\n\n if model_class is None:\n model_class = ChatOpenAI\n\n self._models_config[model_id] = {\n \"type\": model_type, \"name\": model_name, \"class\": model_class,\n \"capabilities\": list(set(capabilities)),\n \"config_override\": config_override or {},\n \"kwargs\": kwargs,\n }\n print(f\" Registered model config: '{model_id}' (Type: {model_type.name}, Class: {model_class.__name__})\")\n\n if set_as_default:\n self._default_model_id = model_id\n print(f\" Set '{model_id}' as default.\")\n\n for capability in capabilities:\n if capability not in self._capability_models:\n self._capability_models[capability] = model_id\n print(f\" Mapped capability '{capability.name}' to '{model_id}'.\")\n\n def set_default_model(self, model_id: str) -> None:\n \"\"\"设置默认模型\"\"\"\n if model_id not in self._models_config: raise ValueError(...)\n self._default_model_id = model_id\n\n def set_capability_model(self, capability: ModelCapability, model_id: str) -> None:\n \"\"\"设置特定能力的模型\"\"\"\n if model_id not in self._models_config: raise ValueError(...)\n model_info = self._models_config[model_id]\n if capability not in model_info.get(\"capabilities\", []):\n print(f\"Warning: Model '{model_id}' not registered with capability '{capability.name}'.\")\n self._capability_models[capability] = model_id\n\n # _get_instance (核心实例化逻辑)\n def _get_instance(self, model_id: str) -> BaseChatModel:\n \"\"\"(Internal) Gets or creates a model instance.\"\"\"\n if model_id in self._models_instance:\n return self._models_instance[model_id]\n\n if model_id not in self._models_config:\n raise ValueError(f\"Model ID '{model_id}' not registered or registration skipped due to missing config.\")\n\n config = self._models_config[model_id]\n model_type = config[\"type\"]\n model_name = config[\"name\"]\n model_class = config[\"class\"]\n config_override = config[\"config_override\"]\n kwargs = config[\"kwargs\"]\n\n # 确定 Key/URL (优先 override, 其次 env)\n api_key = config_override.get(\"api_key\", self._loaded_api_keys.get(model_type))\n base_url = config_override.get(\"base_url\", self._loaded_base_urls.get(model_type))\n\n # OpenAI 特殊 Key 处理\n if model_type == ModelType.OPENAI and not api_key:\n api_key = os.getenv(\"OPENAI_API_KEY\")\n\n # 检查必要配置\n key_required = model_type not in [ModelType.CUSTOM]\n url_required = model_type in [ModelType.XAI, ModelType.DEEPSEEK, ModelType.CUSTOM]\n if key_required and not api_key:\n raise ValueError(f\"API key required but not found for '{model_id}' (Type: {model_type.name}). Set in .env or config_override.\")\n if url_required and not base_url:\n raise ValueError(f\"Base URL required but not found for '{model_id}' (Type: {model_type.name}). Set in .env or config_override.\")\n\n print(f\"Instantiating model: ID='{model_id}', Type='{model_type.name}', Name='{model_name}'\")\n\n # 准备构造函数参数\n init_kwargs = kwargs.copy()\n if model_class == ChatOpenAI:\n init_kwargs['model'] = model_name\n if api_key: init_kwargs['openai_api_key'] = api_key\n if base_url: init_kwargs['openai_api_base'] = base_url\n # elif model_class == ChatGroq: ... # Removed\n else: # 尝试通用参数\n init_kwargs['model'] = model_name # 很多兼容类可能也认 model\n init_kwargs['model_name'] = model_name\n if api_key: init_kwargs['api_key'] = api_key\n if base_url: init_kwargs['base_url'] = base_url\n\n # 移除内部配置键\n for k in [\"config_override\", \"capabilities\", \"type\", \"class\", \"name\", \"instance\"]:\n init_kwargs.pop(k, None)\n \n # 实例化\n try:\n instance = model_class(**init_kwargs)\n self._models_instance[model_id] = instance\n return instance\n except Exception as e:\n print(f\"!!! Failed to instantiate model '{model_id}'\")\n raise e\n\n # get_model 和 get_model_for_capability (保持不变, 调用 _get_instance)\n def get_model(self, model_id: Optional[str] = None) -> BaseChatModel:\n \"\"\"获取模型实例 (通过 ID 或默认)\"\"\"\n target_id = model_id\n if target_id is None:\n if self._default_model_id is None: raise ValueError(\"No default model set.\")\n target_id = self._default_model_id\n if target_id not in self._models_config: raise ValueError(f\"Model ID '{target_id}' not registered.\")\n return self._get_instance(target_id)\n\n def get_model_for_capability(self, capability: ModelCapability) -> BaseChatModel:\n \"\"\"获取具有特定能力的模型实例\"\"\"\n if capability not in self._capability_models:\n print(f\"No preferred model for '{capability.name}'. Falling back to default.\")\n if self._default_model_id is None: raise ValueError(f\"No model for '{capability.name}' and no default set.\")\n model_id = self._default_model_id\n else: model_id = self._capability_models[capability]\n print(f\"Using model '{model_id}' for capability '{capability.name}'.\")\n return self.get_model(model_id)\n\n # list_models 和 list_capabilities (保持不变)\n def list_models(self) -> Dict[str, Dict[str, Any]]:\n \"\"\"列出所有注册的模型及其配置\"\"\"\n result = {}; # ... (populate result) ...\n for model_id, model_info in self._models_config.items():\n result[model_id] = {\n \"type\": model_info[\"type\"].name,\n \"name\": model_info[\"name\"],\n \"class\": model_info[\"class\"].__name__,\n \"capabilities\": [c.name for c in model_info.get(\"capabilities\", [])],\n \"is_default\": model_id == self._default_model_id,\n \"kwargs\": model_info.get(\"kwargs\"),\n \"config_override\": model_info.get(\"config_override\"),\n }\n return result\n\n def list_capabilities(self) -> Dict[str, str]:\n return {capability.name: model_id for capability, model_id in self._capability_models.items()}"
},
{
"path": "core/llm/model_config.py",
"content": "# reason_graph/model_config.py\nfrom langchain_openai import ChatOpenAI\n# from langchain_groq import ChatGroq # 不再需要\n# (如果未来支持其他非 OpenAI 兼容的，在这里 import)\n\nfrom .llm_manager import ModelType, ModelCapability # 从同级 llm_manager 导入枚举\n\n# 定义支持的模型及其配置\n# key 是我们内部使用的 model_id\nSUPPORTED_MODELS_CONFIG = {\n \"openai_gpt4o\": {\n \"model_type\": ModelType.OPENAI,\n \"model_name\": \"gpt-4o\", # API 调用名\n \"model_class\": ChatOpenAI,\n \"capabilities\": [\n ModelCapability.GENERAL, ModelCapability.PLANNING, ModelCapability.REASONING,\n ModelCapability.CREATIVE, ModelCapability.LONG_CONTEXT, ModelCapability.CODE,\n ModelCapability.RESEARCH # GPT-4o 也能做一定研究\n ],\n \"is_default\": False, # 不设为默认\n \"config_override\": {}, # 允许覆盖 env vars, e.g., {'api_key': '...'}\n \"kwargs\": {\"temperature\": 0.1} # 传递给构造函数的额外参数\n },\n \"openai_gpt4o_mini\": {\n \"model_type\": ModelType.OPENAI,\n \"model_name\": \"gpt-4o-mini\",\n \"model_class\": ChatOpenAI,\n \"capabilities\": [ModelCapability.GENERAL, ModelCapability.REASONING, ModelCapability.CREATIVE],\n \"is_default\": True, # <--- 将其设为默认模型\n \"config_override\": {},\n \"kwargs\": {\"temperature\": 0.0}\n },\n \"xai_grok\": { # 假设 ID 命名为 xai_grok\n \"model_type\": ModelType.XAI,\n \"model_name\": \"grok-2-latest\", # 或者是 xAI API 实际接受的模型名\n \"model_class\": ChatOpenAI, # 假设使用兼容 OpenAI 的方式连接\n \"capabilities\": [ModelCapability.GENERAL, ModelCapability.REASONING, ModelCapability.LONG_CONTEXT, ModelCapability.CREATIVE],\n \"is_default\": False,\n \"config_override\": {}, # Key/URL 将从 env (XAI_API_KEY, XAI_BASE_URL) 加载\n \"kwargs\": {\"temperature\": 0.2}\n },\n \"deepseek_v3\": { # 假设 ID 命名为 deepseek_chat\n \"model_type\": ModelType.DEEPSEEK,\n \"model_name\": \"deepseek/deepseek-v3-0324\", # DeepSeek Chat 模型 API 名\n \"model_class\": ChatOpenAI, # 使用兼容 OpenAI 的方式连接\n \"capabilities\": [ModelCapability.GENERAL, ModelCapability.REASONING, ModelCapability.CODE, ModelCapability.LONG_CONTEXT],\n \"is_default\": False,\n \"config_override\": {}, # Key/URL 将从 env (DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL) 加载\n \"kwargs\": {\"temperature\": 0.0}\n },\n # --- 可以继续添加其他模型配置 ---\n # \"groq_llama3_70b\": {\n # \"model_type\": ModelType.GROQ,\n # \"model_name\": \"llama3-70b-8192\",\n # \"model_class\": ChatGroq, # 需要导入 ChatGroq\n # \"capabilities\": [...],\n # \"is_default\": False,\n # \"config_override\": {},\n # \"kwargs\": {\"temperature\": 0.1}\n # },\n}"
},
{
"path": "core/mcp/README.md",
"content": "# Mentis MCP 客户端与配置指南\n\n本目录 (`core/mcp/`) 包含用于与模型上下文协议 (MCP - Model Context Protocol) 服务器进行交互的 Python 客户端实现。\n\n## 背景\n\nMCP 旨在为 AI 模型（如 LLM Agent）提供一个标准的、与外部工具或服务进行交互的协议。本客户端的目标是提供一种灵活、可配置的方式来连接这些 MCP 服务器，并将它们提供的工具集成到 LangChain Agent 中。\n\n## 客户端 (`MCPClient`)\n\n核心实现是 `MCPClient` 类 (位于 `client.py`)，它具备以下特性：\n\n* **配置驱动:** 通过读取一个位于 `core/mcp/config.json` 的 JSON 文件来管理一个或多个服务器的连接/启动信息。兼容 \"Cursor 风格\" 的配置格式。\n* **灵活连接:**\n * **启动本地服务 (stdio):** 如果配置文件中提供了 `command` 和 `args`，客户端会尝试执行该命令启动服务器进程，并通过 **STDIO** 建立通信。这对于使用 `uvx` 或 `python -m` 启动的标准 MCP 服务器很有用。\n * **连接远程服务 (sse):** 如果配置文件中提供了 `url`，客户端会直接通过 **SSE** 连接到该 URL 对应的、已在运行的 MCP 服务器。\n* **异步架构:** 基于 `asyncio` 构建，适合异步应用。\n* **健壮的资源管理:** 使用 `contextlib.AsyncExitStack` 管理连接和会话，旨在提高关闭时的稳定性。\n* **LangChain 集成支持:** 提供了加载 MCP 工具为 LangChain `BaseTool` 对象的基础（尽管存在适配器问题，见下文）。\n\n## 如何使用\n\n### 1. 配置服务器 (`core/mcp/config.json`)\n\n你需要在此目录下创建一个 `config.json` 文件，定义你想要连接的 MCP 服务器。文件是一个 JSON 对象，键是服务器的逻辑名称，值是该服务器的配置详情。\n\n**示例 `config.json` (只包含外部标准服务器):**\n\n```json\n{\n \"fetch_via_uvx\": {\n \"id\": \"fetch-uvx-stdio\",\n \"type\": \"mcp-server\",\n \"description\": \"Fetch Server launched by uvx via stdio\",\n \"connection\": {\n \"transport\": \"stdio\",\n \"command\": \"uvx\",\n \"args\": [ \"mcp-server-fetch\" ],\n \"timeout\": 45\n }\n },\n \"everything\": {\n \"id\": \"everything-stdio\",\n \"type\": \"mcp-server\",\n \"description\": \"Everything Server launched by npx via stdio\",\n \"connection\": {\n \"transport\": \"stdio\",\n \"command\": \"npx\",\n \"args\": [ \"-y\", \"@modelcontextprotocol/server-everything\" ],\n \"env\": {\n // 如果 Everything Server 需要 API Keys, 在此添加\n // 或确保运行客户端脚本的环境变量会被继承\n // \"OPENAI_API_KEY\": \"YOUR_KEY\",\n // \"TAVILY_API_KEY\": \"YOUR_KEY\"\n },\n \"timeout\": 60\n }\n },\n \"external_sse_example\": {\n \"id\": \"external-sse\",\n \"type\": \"mcp-server\",\n \"description\": \"Connect to a pre-running SSE server (Example)\",\n \"connection\": {\n \"transport\": \"sse\",\n \"url\": \"http://localhost:9001/sse\" // 假设有服务器在此运行\n }\n }\n}\n```\n\n**重要:**\n\n* 使用 `command` 启动服务器时，确保 `command` (如 `uvx`, `npx`, `python`) 在你的环境中可用。\n* 如果服务器需要 API Keys，请通过 `env` 字段或系统环境变量提供。\n* `transport: \"stdio\"` 告诉我们的客户端使用 stdio 连接，`transport: \"sse\"` 告诉它使用 sse 连接。\n\n### 2. 客户端代码示例\n\n使用 `config_loader.py` 加载配置，并通过 `async with` 语句使用 `MCPClient`。\n\n```python\nimport asyncio\nimport os\nfrom core.mcp.client import MCPClient\nfrom core.mcp.config_loader import load_config\n# 导入 LangChain 相关 (如果需要 Agent)\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.prebuilt import create_react_agent\nfrom langchain_core.tools import BaseTool, Tool\n# 导入工具的 Pydantic Schema (用于手动创建 Tool)\nfrom pydantic.v1 import BaseModel, Field # 或 v2\n\n# --- Fetch Schema 示例 ---\nclass FetchInputSchema(BaseModel):\n url: str = Field(..., description=\"URL to fetch\")\n # ... 其他字段 ...\n\nasync def main():\n # --- 加载配置 ---\n config_path = os.path.join(os.path.dirname(__file__), \"config.json\") # 假设 config 在同目录\n try:\n all_configs = load_config(config_path)\n # 选择要使用的配置\n server_key = \"fetch_via_uvx\" # 或 \"everything\", \"e2b_stdio\" 等\n mcp_config = all_configs.get(server_key)\n if not mcp_config:\n print(f\"Config '{server_key}' not found.\")\n return\n except Exception as e:\n print(f\"Failed to load config: {e}\")\n return\n\n # --- 使用 MCPClient ---\n async with MCPClient(mcp_config) as client:\n print(f\"Connected to MCP Server '{server_key}'. Session active: {client.session is not None}\")\n if not client.session: return\n\n # --- 获取和使用工具 ---\n\n # 方式一: 标准方式 (但存在已知问题)\n # print(\"\\nAttempting standard tool loading via load_mcp_tools...\")\n # loaded_tools = client.get_tools() # 内部调用 load_mcp_tools\n # print(f\"load_mcp_tools returned {len(loaded_tools)} tools.\")\n # # !! 注意：对于某些服务器实现 (如此处之前的 MentisMCPServer),\n # # !! load_mcp_tools 返回的工具对象的 args_schema 可能是错误的！\n # # !! 这会导致 Agent 调用失败。但对于 Fetch Server 这样的标准服务器，\n # # !! 它加载的 Schema 可能是正确的。需要根据打印的 Schema 判断。\n\n # 方式二: 【当前推荐】手动创建 Tool 对象 (绕过 load_mcp_tools 问题)\n print(\"\\nManually creating Tool object with correct schema...\")\n tool_name = \"fetch\" # 假设测试 Fetch Server\n tool_description = \"Fetches URL content.\" # 可以从服务器获取或手写\n correct_schema = FetchInputSchema # 使用正确的 Pydantic 模型\n\n # 定义调用逻辑\n async def call_mcp_tool_wrapper(**kwargs) -> str:\n # ... (内部使用 client.session.call_tool 发送正确请求) ...\n # 参考 examples/14_mcp_fetch_test.py 中的实现\n if not client or not client.session: return \"ERROR: Session lost.\"\n try:\n req_params = {\"name\": tool_name, \"arguments\": kwargs}\n from mcp.types import CallToolRequest # 需要导入\n request = CallToolRequest(method='tools/call', params=req_params)\n result = await client.session.call_tool(request)\n if hasattr(result, 'result'): return str(result.result)\n elif hasattr(result, 'error'): return f\"Tool Error: {result.error.message}\"\n else: return \"Unknown response\"\n except Exception as e: return f\"Error: {e}\"\n\n # 创建 LangChain Tool\n manual_tool = Tool.from_function(\n name=tool_name,\n description=tool_description,\n args_schema=correct_schema,\n coroutine=call_mcp_tool_wrapper\n )\n tools_for_agent = [manual_tool]\n print(f\"Manual tool '{manual_tool.name}' created.\")\n\n # --- 使用 Agent ---\n try:\n # model = llm_manager.get_model(\"openai_gpt4o_mini\") # 获取 LLM\n # agent = create_react_agent(model, tools_for_agent)\n # response = await agent.ainvoke(...)\n # print(\"Agent Response:\", response)\n print(\"\\nAgent execution part skipped in README example.\")\n print(\"Refer to examples/14_mcp_fetch_test.py for full Agent integration.\")\n except Exception as e:\n print(f\"Agent execution error: {e}\")\n\n# if __name__ == \"__main__\":\n# asyncio.run(main())\n```\n\n## 关于自建 MCP Server (MentisMCPServer)\n\n我们在之前的开发中，尝试在 `core/mcp/server.py` 中构建了一个 `MentisMCPServer` 类，目的是将我们内部工具注册表 (`core/tools/registry.py`) 中的 LangChain `BaseTool` 动态包装成 MCP 工具。\n\n**当前遇到的主要挑战：**\n\n我们发现，当使用 `FastMCP` 库的 `@mcp.tool` 装饰器来动态注册这些包装器时，服务器未能正确地向客户端广播这些工具的**输入模式 (Schema)**。这导致客户端的 `load_mcp_tools` 收到了错误的 Schema 信息，进而使 LangChain Agent 在调用工具时因参数错误而失败。\n\n虽然我们通过重构服务器的注册逻辑（改为在 `run_server.py` 中直接使用 `FastMCP` 实例注册顶层包装函数）**成功解决**了 Schema 广播的问题，使得 `load_mcp_tools` 能够获取到正确的 Schema，但后续测试发现 Agent (`create_react_agent`) 在调用这些工具时仍可能出现内部错误 (`TypeError`)。\n\n**结论与建议：**\n\n由于在结合 LangChain 工具、动态包装、`FastMCP` 和 LangChain Agent 时遇到了较深的库交互和调试障碍，我们**目前不建议**将 `MentisMCPServer` 作为稳定可靠的方案对外提供服务。\n\n**推荐使用以下方式来提供或使用 MCP Server:**\n\n1. **使用社区标准服务器:** 直接使用像 `mcp-server-fetch`, `@modelcontextprotocol/server-everything` 这样由社区或官方提供的、预构建好的 MCP 服务器。通过 `config.json` 配置 `command` (如 `uvx`, `npx`, `python -m`) 或 `url` 来使用它们。\n2. **采用简单服务器模式:** 如果你需要自己实现 MCP Server 来暴露特定功能，建议参考 `modelcontextprotocol/servers` 仓库中的简单示例（如 `math_server`, `time_server`），采用**直接注册工具函数**（用 `@mcp_instance.tool` 装饰顶层 `async def` 函数）的模式，避免复杂的动态包装层。"
},
{
"path": "core/mcp/__init__.py",
"content": "# core/mcp/__init__.py\n\"\"\"\nMCP (Model Context Protocol) 功能模块\n\"\"\""
},
{
"path": "core/mcp/client.py",
"content": "import os\nimport asyncio\nfrom pathlib import Path\nfrom typing import List, Dict, Any, Optional, Union, Type, Literal, TypedDict, cast\nfrom types import TracebackType\nimport re\nimport sys\nimport json\nimport traceback\nfrom contextlib import asynccontextmanager, AsyncExitStack\n\n# --- MCP Imports ---\nfrom mcp import ClientSession, StdioServerParameters\nfrom mcp.client.stdio import stdio_client\nfrom mcp.client.sse import sse_client\n# --- Adapter Import ---\ntry:\n from langchain_mcp_adapters.tools import load_mcp_tools\n LOAD_MCP_TOOLS_AVAILABLE = True\nexcept ImportError:\n print(\"警告: 未找到 langchain-mcp-adapters。 load_mcp_tools 将不可用。\")\n async def load_mcp_tools(session: ClientSession) -> list: return []\n LOAD_MCP_TOOLS_AVAILABLE = False\n# --- LangChain / Pydantic Imports ---\nfrom langchain_core.tools import BaseTool\ntry: from pydantic.v1 import BaseModel as BaseModelV1\nexcept ImportError: from pydantic import BaseModel as BaseModelV1 # Fallback\n# --- Config Loader Import ---\ntry: from .config_loader import MCPConfig, StdioConfig, SSEConfig\nexcept ImportError: print(\"WARNING: Could not import config models from .config_loader.\"); MCPConfig=Any; StdioConfig=Any; SSEConfig=Any # Placeholders\n\nprint(\"--- DEBUG: Loading FINAL client.py (Config-Driven + AsyncExitStack) ---\")\n\nclass MCPClient:\n \"\"\"Config-driven MCP Client using AsyncExitStack.\"\"\"\n def __init__(self, config: MCPConfig):\n self.config = config\n self.session: Optional[ClientSession] = None\n self.tools: List[BaseTool] = []\n self._stack: AsyncExitStack = AsyncExitStack()\n self._server_process: Optional[asyncio.subprocess.Process] = None\n\n async def __aenter__(self) -> \"MCPClient\":\n print(f\"DEBUG: MCPClient entering context for config ID: {getattr(self.config, 'id', 'N/A')}\")\n try:\n connection_config = self.config.connection\n transport_ctx = None\n reader = None\n writer = None\n\n if isinstance(connection_config, SSEConfig) and connection_config.url:\n # --- Direct SSE ---\n print(f\"DEBUG: Connecting via SSE to {connection_config.url}\")\n transport_ctx = sse_client(\n connection_config.url, getattr(connection_config,'headers', None),\n getattr(connection_config,'timeout', 5.0), getattr(connection_config,'sse_read_timeout', 300.0)\n )\n reader, writer = await self._stack.enter_async_context(transport_ctx)\n print(\"DEBUG: SSE transport context entered.\")\n\n elif isinstance(connection_config, StdioConfig) and connection_config.command:\n # --- Launch via Command + STDIO ---\n print(f\"DEBUG: Launching command via STDIO: {connection_config.command} {' '.join(connection_config.args)}\")\n merged_env = os.environ.copy();\n if connection_config.env: merged_env.update(connection_config.env)\n server_params = StdioServerParameters(\n command=connection_config.command, args=connection_config.args, env=merged_env,\n cwd=connection_config.cwd, encoding=connection_config.encoding,\n encoding_error_handler=connection_config.encoding_error_handler,\n startup_timeout=connection_config.timeout\n )\n transport_ctx = stdio_client(server_params)\n reader, writer = await self._stack.enter_async_context(transport_ctx)\n print(\"DEBUG: STDIO transport context entered.\")\n\n else: # Fallback/Error - Handle case where config might be wrong or transport missing\n # Added check for command presence before assuming SSE launch\n if hasattr(connection_config, 'command') and connection_config.command:\n # This is the complex \"launch then connect SSE\" case from the guide\n # Keeping it simple for now - if transport isn't 'stdio', it must be 'sse' with a URL\n raise NotImplementedError(\"Launching command for SSE connection (URL capture) not implemented in this client version. Use direct SSE URL or STDIO command.\")\n else:\n raise ValueError(\"Invalid configuration: must have 'url' for SSE or 'command' for STDIO.\")\n\n\n # --- Establish ClientSession ---\n session_kwargs = getattr(connection_config, 'session_kwargs', None) or {}\n session_ctx = ClientSession(reader, writer, **session_kwargs)\n self.session = await self._stack.enter_async_context(session_ctx)\n print(\"DEBUG: ClientSession context entered.\")\n\n # --- Initialize and Load Tools (with Schema Check) ---\n print(\"Initializing MCP session...\")\n await asyncio.wait_for(self.session.initialize(), timeout=30.0)\n print(\"MCP session initialized.\")\n\n if LOAD_MCP_TOOLS_AVAILABLE:\n print(\"Loading MCP tools (via langchain-mcp-adapters)...\")\n loaded_tools_from_mcp = await load_mcp_tools(self.session)\n print(f\"Successfully loaded {len(loaded_tools_from_mcp)} tool descriptions.\")\n print(\"--- Loaded Tools & Args Schema (Diagnostic) ---\")\n self.tools = []\n for i, tool in enumerate(loaded_tools_from_mcp):\n schema = getattr(tool, 'args_schema', 'N/A'); tool_name = getattr(tool, 'name', f'Tool_{i+1}')\n print(f\"{i+1}. Tool Name: {tool_name}\")\n schema_detail = \"N/A\"\n is_correct = None # Undetermined\n if schema != 'N/A': # Schema printing and basic check\n schema_dict = None\n if isinstance(schema, type) and issubclass(schema, BaseModelV1):\n try: schema_dict = schema.schema(); schema_detail = f\"(PydanticV1): {json.dumps(schema_dict, indent=2)}\"\n except Exception as e_schema: schema_detail = f\"(PydanticV1): Error - {e_schema}\"\n elif hasattr(schema, 'model_json_schema'):\n try: schema_dict = schema.model_json_schema(); schema_detail = f\"(PydanticV2): {json.dumps(schema_dict, indent=2)}\"\n except Exception as e_schema: schema_detail = f\"(PydanticV2): Error - {e_schema}\"\n else: schema_detail = f\"(Unknown Type): {schema}\"\n # Basic check: does it look like the faulty kwargs schema?\n if isinstance(schema_dict, dict):\n props = schema_dict.get('properties', {})\n if list(props.keys()) == ['kwargs'] and props['kwargs'].get('type') == 'string':\n is_correct = False\n schema_detail += \" <-- LOOKS WRONG (kwargs only!)\"\n elif props:\n is_correct = True # Has properties other than just kwargs\n schema_detail += \" <-- Looks structured correctly\"\n else:\n is_correct = True # No properties, might be simple input\n schema_detail += \" <-- No properties defined\"\n else: is_correct = False # No schema is usually wrong\n print(f\" Args Schema: {schema_detail}\")\n print(\"-\" * 15); self.tools.append(tool)\n print(f\"Schema Check Result: {'All schemas look structured correctly.' if all(s is not False for s in [getattr(t, 'args_schema', None) != 'N/A' and 'kwargs' not in str(getattr(t, 'args_schema', '')).lower() for t in self.tools]) else 'One or more schemas look incorrect (kwargs only or missing)!'}\")\n print(\"-------------------------------------------\")\n else: print(\"Warning: load_mcp_tools unavailable.\"); self.tools = []\n print(f\"MCPClient ready. Loaded {len(self.tools)} tools via adapter.\")\n return self\n except Exception as enter_err:\n print(f\"ERROR: Failed during MCPClient __aenter__: {type(enter_err).__name__}: {enter_err}\")\n await self.close(); raise\n\n async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]):\n print(\"DEBUG: MCPClient exiting context...\"); await self.close(); print(\"DEBUG: MCPClient context exited.\")\n\n async def close(self):\n \"\"\"Closes connections and resets state using AsyncExitStack.\"\"\"\n print(\"Closing MCP Client...\");\n if hasattr(self, '_stack') and self._stack:\n print(\" Closing managed async contexts (via AsyncExitStack)...\")\n try: await self._stack.aclose(); print(\" AsyncExitStack closed.\")\n except Exception as e: print(f\"WARNING: Error closing AsyncExitStack: {type(e).__name__}: {e}\")\n finally: self._stack = None\n else: print(\" No active AsyncExitStack.\")\n self.session = None; self.tools = []; self._transport_ctx = None; self._server_process = None\n print(\"MCP Client state reset.\")\n\n def get_tools(self) -> List[BaseTool]:\n \"\"\"Returns the list of tools loaded by load_mcp_tools.\"\"\"\n return self.tools"
},
{
"path": "core/mcp/config_loader.py",
"content": "# core/mcp/config_loader.py (修改 load_config 返回类型)\nimport json\nimport os\nfrom pathlib import Path\nfrom typing import Dict, Any, Optional, List, Literal, Union, Type # 导入 Type\ntry:\n from pydantic.v1 import BaseModel, Field, ValidationError, validator\n PYDANTIC_V = 1\nexcept ImportError:\n try:\n from pydantic import BaseModel, Field, ValidationError, validator # type: ignore\n PYDANTIC_V = 2\n except ImportError: raise ImportError(\"Pydantic (v1 or v2) required.\")\nfrom typing_extensions import TypedDict\n\nEncodingErrorHandler = Literal[\"strict\", \"ignore\", \"replace\"]\n\nclass StdioConfig(BaseModel):\n transport: Literal[\"stdio\"] = \"stdio\"; command: str = Field(...)\n args: List[str] = Field(default_factory=list); env: Optional[Dict[str, str]] = None\n cwd: Optional[Union[str, Path]] = None; encoding: str = Field(default=\"utf-8\")\n encoding_error_handler: EncodingErrorHandler = Field(default=\"strict\")\n timeout: int = Field(default=30, gt=0); session_kwargs: Optional[Dict[str, Any]] = None\n if PYDANTIC_V == 1: \n class Config: extra = 'forbid'\n else: model_config = {'extra': 'forbid'}\n\nclass SSEConfig(BaseModel):\n transport: Literal[\"sse\"] = \"sse\"; url: str = Field(...)\n headers: Optional[Dict[str, Any]] = None; timeout: float = Field(default=5.0, gt=0)\n sse_read_timeout: float = Field(default=300.0, gt=0); session_kwargs: Optional[Dict[str, Any]] = None\n if PYDANTIC_V == 1: \n class Config: extra = 'forbid'\n else: model_config = {'extra': 'forbid'}\n\nclass MCPConfig(BaseModel):\n \"\"\"Represents the structure for a single server configuration.\"\"\"\n id: Optional[str] = Field(default=None)\n type: Literal[\"mcp-server\"] = Field(default=\"mcp-server\")\n description: Optional[str] = Field(default=None)\n connection: Union[StdioConfig, SSEConfig] = Field(..., discriminator='transport')\n if PYDANTIC_V == 1: \n class Config: extra = 'forbid'\n else: model_config = {'extra': 'forbid'}\n\n\n# --- 修改 load_config ---\ndef load_config(config_path: Union[str, Path]) -> Dict[str, MCPConfig]:\n \"\"\"\n Loads the central MCP configuration JSON file and validates each server entry.\n\n Args:\n config_path: Path to the central config.json file.\n\n Returns:\n A dictionary where keys are server names and values are validated MCPConfig objects.\n \"\"\"\n config_p = Path(config_path).resolve()\n if not config_p.is_file():\n raise FileNotFoundError(f\"Configuration file not found at: {config_p}\")\n\n print(f\"DEBUG: Loading central MCP configuration from: {config_p}\")\n validated_configs: Dict[str, MCPConfig] = {}\n try:\n with open(config_p, 'r', encoding='utf-8') as f:\n raw_config_dict = json.load(f)\n\n if not isinstance(raw_config_dict, dict):\n raise TypeError(\"Root configuration must be a JSON object (dictionary).\")\n\n # 遍历字典中的每个服务器配置并验证\n for server_name, config_data in raw_config_dict.items():\n print(f\"DEBUG: Validating config for server: '{server_name}'\")\n if not isinstance(config_data, dict):\n print(f\"WARNING: Entry for '{server_name}' is not a dictionary. Skipping.\")\n continue\n try:\n # 确保 connection 和 transport 存在\n if 'connection' not in config_data: raise ValueError(\"Missing 'connection'\")\n if 'transport' not in config_data.get('connection', {}): raise ValueError(\"Missing 'transport' in connection\")\n\n if PYDANTIC_V == 2:\n validated_config = MCPConfig.model_validate(config_data)\n else: # Pydantic V1\n validated_config = MCPConfig.parse_obj(config_data)\n validated_configs[server_name] = validated_config\n print(f\"DEBUG: Config for '{server_name}' validated successfully.\")\n except (ValidationError, ValueError) as e_val:\n print(f\"ERROR: Validation failed for server '{server_name}' config:\\n{e_val}\\nSkipping this server.\")\n #可以选择继续加载其他配置，或者在这里 raise 让整个加载失败\n\n if not validated_configs:\n print(\"WARNING: No valid server configurations were loaded.\")\n\n print(f\"DEBUG: Central configuration loaded. Found {len(validated_configs)} valid server configs.\")\n return validated_configs\n except json.JSONDecodeError as e:\n print(f\"ERROR: Failed to decode JSON from {config_p}: {e}\"); raise\n except Exception as e:\n print(f\"ERROR: An unexpected error occurred loading config {config_p}: {e}\"); raise"
},
{
"path": "core/mcp/mcp_server_config.json",
"content": "{\n \"fetch_via_uvx\": {\n \"id\": \"fetch-uvx-stdio\",\n \"type\": \"mcp-server\",\n \"description\": \"Fetch Server launched by uvx via stdio\",\n \"connection\": {\n \"transport\": \"stdio\",\n \"command\": \"uvx\",\n \"args\": [\n \"mcp-server-fetch\"\n ],\n \"env\": null,\n \"cwd\": null,\n \"encoding\": \"utf-8\",\n \"encoding_error_handler\": \"strict\",\n \"timeout\": 45\n }\n },\n \"everything\": {\n \"id\": \"everything-stdio\",\n \"type\": \"mcp-server\",\n \"description\": \"Everything Server\",\n \"connection\": {\n \"transport\": \"stdio\",\n \"command\": \"npx\",\n \"args\": [\n \"-y\",\n \"@modelcontextprotocol/server-everything\"\n ],\n \"env\": null,\n \"cwd\": null,\n \"encoding\": \"utf-8\",\n \"encoding_error_handler\": \"strict\",\n \"timeout\": 45\n }\n }\n }"
},
{
"path": "core/mcp/run_server.py",
"content": "# core/mcp/run_server.py (FINAL - Direct FastMCP Registration)\nimport os\nimport sys\nimport argparse\nimport traceback\nimport logging\nfrom typing import List, Dict, Any, Optional, Type\n\n# --- Standard Setup ---\nlogging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')\nlogger = logging.getLogger(\"mcp_server_direct\")\n\ncurrent_dir = os.path.dirname(os.path.abspath(__file__)); \nproject_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir))); sys.path.insert(0, project_root)\n\n# --- Imports ---\nfrom mcp.server.fastmcp import FastMCP # Import FastMCP directly\n# Assume registry is populated correctly by preregister_core_tools\nfrom core.tools.registry import get_registered_tools, get_tool_instance\ntry: \n from core.tools import preregister_core_tools; \n PREREGISTER_AVAILABLE = True\nexcept ImportError: \n print(\"WARNING: preregister_core_tools not found\"); \n def preregister_core_tools(): pass; \n PREREGISTER_AVAILABLE = False\nfrom langchain_core.tools import BaseTool\nimport asyncio\nimport time\nimport json\nimport functools\nimport inspect # Needed for func_metadata potentially\n\nprint(\"--- DEBUG: Loading FINAL run_server.py (Direct FastMCP Registration) ---\")\n\n# --- Tool Wrapper Creation Logic (as a standalone function) ---\ndef create_tool_wrapper(tool_instance: BaseTool):\n \"\"\"\n Creates the async wrapper function for a given tool instance.\n This function will be decorated LATER by the mcp_instance.\n \"\"\"\n tool_name = getattr(tool_instance, 'name', 'unknown_tool')\n print(f\" DEBUG: Defining wrapper function for tool: '{tool_name}'\")\n\n # Define the actual wrapper coroutine\n async def dynamic_tool_wrapper(tool_to_run=tool_instance, **kwargs): # Bind instance\n _tool_name = tool_to_run.name\n log_file = \"/tmp/mcp_wrapper.log\"; \n timestamp = time.strftime(\"%Y-%m-%d %H:%M:%S\"); \n log_prefix = f\"--- {timestamp} WRAPPER for '{_tool_name}' ---\"\n log_lines = [f\"{log_prefix} START\", f\"Received kwargs: {kwargs}\"]\n try: # Main execution block\n result = None\n if hasattr(tool_to_run, '_arun'):\n log_lines.append(f\"Calling await tool_to_run._arun(**kwargs)\")\n result = await tool_to_run._arun(**kwargs)\n log_lines.append(f\"Await _arun completed.\")\n elif hasattr(tool_to_run, '_run'):\n log_lines.append(f\"Calling tool_to_run._run(**kwargs) via run_in_executor\")\n loop = asyncio.get_running_loop()\n sync_func_with_args = functools.partial(tool_to_run._run, **kwargs)\n result = await loop.run_in_executor(None, sync_func_with_args)\n log_lines.append(f\"Executor _run completed.\")\n else: log_lines.append(\"ERROR: Tool no _arun/_run!\"); raise NotImplementedError(f\"Tool {_tool_name} no method.\")\n\n log_lines.append(f\"Raw result type: {type(result)}\"); log_lines.append(f\"Raw value snippet: {str(result)[:500]}...\")\n final_result = result\n try: json.dumps(result); log_lines.append(\"Result JSON serializable.\")\n except TypeError: log_lines.append(f\"WARN: Non-JSON type {type(result)}.->str.\"); final_result = str(result)\n log_lines.append(f\"Returning final (type {type(final_result)}).\"); log_lines.append(f\"{log_prefix} END (Success)\")\n return {\"result\": final_result}\n except Exception as e: # Catch execution errors\n log_lines.append(f\"!!! EXCEPTION in tool exec for '{_tool_name}': {e} !!!\"); tb_lines = traceback.format_exc().splitlines(); log_lines.append(\"--- Traceback ---\"); log_lines.extend(tb_lines); log_lines.append(\"-----------------\"); log_lines.append(f\"{log_prefix} END (Exception)\")\n return f\"ERROR_EXECUTING_TOOL_{_tool_name}: {str(e)}\" # Return error string\n finally: # Ensure logging\n try:\n for line in log_lines: print(line, flush=True, file=sys.stderr)\n with open(log_file, \"a\") as f: f.write(\"\\n\".join(log_lines) + \"\\n\\n\")\n except Exception as log_e: print(f\"!!! Logging Error for tool {_tool_name}: {log_e} !!!\", flush=True, file=sys.stderr)\n\n # Return the created wrapper function AND the original tool's metadata\n return dynamic_tool_wrapper, tool_name, getattr(tool_instance, 'description', f\"Tool {tool_name}\")\n\n# --- Main Execution Logic ---\ndef main():\n parser = argparse.ArgumentParser(description='Start Mentis MCP Server (Direct Registration)')\n parser.add_argument('--transport', type=str, choices=['stdio', 'sse'], default='stdio'); parser.add_argument('--host', type=str, default='0.0.0.0'); parser.add_argument('--port', type=int, default=8000); parser.add_argument('--name', type=str, default='MentisMCP'); parser.add_argument('--tools', nargs='+'); parser.add_argument('--debug', action='store_true')\n args = parser.parse_args()\n\n if args.debug: logger.setLevel(logging.DEBUG); print(\"DEBUG Logging Enabled\")\n\n try:\n # --- 1. Preregister tools into the central registry ---\n if PREREGISTER_AVAILABLE:\n print(\"DEBUG: Calling preregister_core_tools...\")\n preregister_core_tools() # This populates the registry\n print(\"DEBUG: preregister_core_tools finished.\")\n else: print(\"DEBUG: Skipping preregister_core_tools (unavailable).\")\n\n # --- 2. Create FastMCP instance ---\n print(f\"DEBUG: Creating FastMCP instance: name='{args.name}'\")\n fastmcp_kwargs = {}\n if args.transport == 'sse':\n if args.host: fastmcp_kwargs['host'] = args.host\n if args.port: fastmcp_kwargs['port'] = args.port\n mcp_instance = FastMCP(args.name, **fastmcp_kwargs) # Create instance directly\n print(f\"DEBUG: FastMCP instance created.\")\n\n # --- 3. Load tools from registry and register wrappers with FastMCP ---\n registered_count = 0\n target_tools = args.tools # List of names, or None for all\n\n # Get all tools first if needed\n all_tools_dict = get_registered_tools(as_dict=True)\n\n tools_to_register = {}\n if target_tools: # Filter if specific tools requested\n print(f\"DEBUG: Filtering for specific tools: {target_tools}\")\n for name in target_tools:\n if name in all_tools_dict:\n tools_to_register[name] = all_tools_dict[name]\n else:\n print(f\"ERROR: Requested tool '{name}' not found in registry.\")\n else: # Register all tools found in registry\n print(\"DEBUG: Registering all tools found in registry...\")\n tools_to_register = all_tools_dict\n\n # Iterate and register the selected tools\n print(f\"DEBUG: Attempting to register {len(tools_to_register)} tools with FastMCP...\")\n for tool_name, tool_info in tools_to_register.items():\n tool_instance = tool_info.get(\"tool\")\n if isinstance(tool_instance, BaseTool):\n try:\n # Create the wrapper function and get metadata\n wrapper_func, name, description = create_tool_wrapper(tool_instance)\n # Register the wrapper directly using the mcp_instance decorator method\n mcp_instance.tool(name=name, description=description)(wrapper_func)\n print(f\"DEBUG: Successfully registered '{name}' with FastMCP.\")\n registered_count += 1\n except Exception as e_register:\n print(f\"ERROR: Failed to register wrapper for tool '{tool_name}': {e_register}\")\n traceback.print_exc()\n else:\n print(f\"WARNING: Item '{tool_name}' not a BaseTool, skipping.\")\n\n print(f\"DEBUG: Tool registration complete. {registered_count} tools registered with FastMCP.\")\n if registered_count == 0: print(\"WARNING: No tools were registered!\")\n\n # --- 4. Run the FastMCP server ---\n print(f\"Starting MCP Server '{args.name}' (Transport: {args.transport})...\")\n mcp_instance.run(transport=args.transport)\n\n except KeyboardInterrupt: print(\"Server shutting down...\"); sys.exit(0)\n except Exception as e: print(f\"Error starting server: {e}\"); traceback.print_exc(); sys.exit(1)\n\nif __name__ == \"__main__\":\n main()"
},
{
"path": "core/mcp/server.py",
"content": "import os\nimport sys\nimport traceback\nimport asyncio\nimport time\nimport json\nimport functools\nfrom typing import Dict, Any, Optional, List\n\n# mcp & fastmcp\nfrom mcp.server.fastmcp import FastMCP\nfrom mcp.types import CallToolResult, TextContent, ErrorData # <-- 关键导入\n\n# 修正路径，导入你自己的工具与 BaseTool\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))\nfrom core.tools.registry import get_registered_tools, get_tool_instance\nfrom langchain_core.tools import BaseTool\n\nprint(\"--- DEBUG: Loading REFACTORED server.py (Fix InvalidSignature) ---\")\n\nclass MentisMCPServer:\n def __init__(self, name: str = \"MentisMCP\", host: Optional[str] = None, port: Optional[int] = None):\n print(f\"DEBUG: Initializing MentisMCPServer(name='{name}', host={host}, port={port})\")\n fastmcp_kwargs = {}\n if host is not None:\n fastmcp_kwargs['host'] = host\n if port is not None:\n fastmcp_kwargs['port'] = port\n\n try:\n print(f\"DEBUG: Calling FastMCP(name='{name}', **{fastmcp_kwargs})\")\n self.mcp = FastMCP(name, **fastmcp_kwargs)\n print(\"DEBUG: FastMCP initialized successfully.\")\n except Exception as e_fastmcp:\n print(\"ERROR: Failed to initialize FastMCP!\")\n print(traceback.format_exc())\n raise\n\n # 记录注册成功的工具包装器\n self.registered_tools_wrappers = {}\n\n def register_all_tools(self):\n \"\"\"批量注册所有在 registry 中找到的 BaseTool\"\"\"\n tools_dict = get_registered_tools(as_dict=True)\n print(f\"DEBUG: Registering all tools ({len(tools_dict)} found)...\")\n registered_count = 0\n for tool_name, tool_info in tools_dict.items():\n tool_instance = tool_info.get(\"tool\")\n if isinstance(tool_instance, BaseTool):\n if self._register_tool_with_simplified_wrapper(tool_instance):\n registered_count += 1\n else:\n print(f\"WARNING: Item '{tool_name}' not BaseTool, skipping.\")\n print(f\"DEBUG: Finished registering all tools. Registered: {registered_count}\")\n\n def register_single_tool(self, tool_name: str):\n \"\"\"仅注册特定名称的一个工具\"\"\"\n print(f\"DEBUG: Attempting to register single tool: {tool_name}\")\n try:\n tool_instance = get_tool_instance(tool_name)\n if not tool_instance:\n print(f\"ERROR: Tool '{tool_name}' not found in registry.\")\n return\n if isinstance(tool_instance, BaseTool):\n if self._register_tool_with_simplified_wrapper(tool_instance):\n print(f\"DEBUG: Successfully registered single tool: {tool_instance.name}\")\n else:\n print(f\"ERROR: Failed wrapper registration for: {tool_instance.name}\")\n else:\n print(f\"WARNING: Tool '{tool_name}' not BaseTool, skipping.\")\n except Exception as e:\n print(f\"ERROR during register_single_tool for '{tool_name}': {e}\")\n print(traceback.format_exc())\n\n def _register_tool_with_simplified_wrapper(self, tool: BaseTool) -> bool:\n \"\"\"\n 为工具创建并注册一个简化的包装器 (Fix InvalidSignature),\n 并确保返回的数据符合 CallToolResult，以便客户端解析.\n \"\"\"\n try:\n tool_name = getattr(tool, 'name', None)\n tool_description = getattr(tool, 'description', None)\n if not tool_name or not isinstance(tool_name, str):\n print(f\"ERROR: Invalid tool name: {tool_name}. Skip.\")\n return False\n if not tool_description or not isinstance(tool_description, str):\n print(f\"WARNING: Empty/invalid description for '{tool_name}'.\")\n tool_description = f\"Tool {tool_name}\"\n\n print(f\"DEBUG: Defining wrapper for tool: '{tool_name}'\")\n\n @self.mcp.tool(name=tool_name, description=tool_description)\n async def simplified_tool_wrapper(tool_for_wrapper=tool, **kwargs):\n \"\"\"\n 同步或异步地调用 tool_for_wrapper，并将结果包装到\n CallToolResult 中返回给客户端，以匹配 .content 或 .error.\n \"\"\"\n _tool_name = tool_for_wrapper.name\n log_file = \"/tmp/mcp_wrapper.log\"\n timestamp = time.strftime(\"%Y-%m-%d %H:%M:%S\")\n log_prefix = f\"--- {timestamp} WRAPPER for '{_tool_name}' ---\"\n log_lines = [f\"{log_prefix} START\", f\"Received kwargs: {kwargs}\"]\n\n try:\n # 根据工具方法签名决定调用 _arun (异步) 或 _run (同步)\n result = None\n if hasattr(tool_for_wrapper, '_arun'):\n log_lines.append(\"Calling await tool_for_wrapper._arun(**kwargs)\")\n result = await tool_for_wrapper._arun(**kwargs)\n log_lines.append(\"Await _arun completed.\")\n elif hasattr(tool_for_wrapper, '_run'):\n log_lines.append(\"Calling tool_for_wrapper._run(**kwargs) via run_in_executor\")\n loop = asyncio.get_running_loop()\n sync_func_with_args = functools.partial(tool_for_wrapper._run, **kwargs)\n result = await loop.run_in_executor(None, sync_func_with_args)\n log_lines.append(\"Executor _run completed.\")\n else:\n log_lines.append(f\"ERROR: Tool '{_tool_name}' has no _arun/_run!\")\n raise NotImplementedError(f\"Tool '{_tool_name}' cannot be invoked directly.\")\n\n # 记录结果类型和内容片段\n log_lines.append(f\"Raw result type: {type(result)}\")\n log_lines.append(f\"Raw value snippet: {str(result)[:500]}...\")\n\n # 关键：将结果包装成 CallToolResult，让客户端能识别 .content\n call_result = CallToolResult(\n content=[TextContent(text=str(result))]\n )\n log_lines.append(\"Returning standard CallToolResult with .content.\")\n log_lines.append(f\"{log_prefix} END (Success)\")\n return call_result\n\n except Exception as e:\n # 出现异常则使用 .error 返回\n log_lines.append(f\"!!! EXCEPTION in tool exec for '{_tool_name}': {e} !!!\")\n tb_lines = traceback.format_exc().splitlines()\n log_lines.append(\"--- Traceback ---\")\n log_lines.extend(tb_lines)\n log_lines.append(\"-----------------\")\n log_lines.append(f\"{log_prefix} END (Exception)\")\n err_msg = f\"ERROR_EXECUTING_TOOL_{_tool_name}: {str(e)}\"\n return CallToolResult(error=ErrorData(message=err_msg))\n\n finally:\n # 日志记录\n try:\n for line in log_lines:\n print(line, flush=True, file=sys.stderr)\n with open(log_file, \"a\") as f:\n f.write(\"\\n\".join(log_lines) + \"\\n\\n\")\n except Exception as log_e:\n print(f\"!!! Logging Error for '{_tool_name}': {log_e} !!!\",\n flush=True, file=sys.stderr)\n\n # 修正下包装器的名字，避免重复\n simplified_tool_wrapper.__name__ = f\"{tool_name}_simplified_wrapper\"\n self.registered_tools_wrappers[tool_name] = simplified_tool_wrapper\n print(f\"DEBUG: Registered simplified wrapper for tool: '{tool_name}'\")\n return True\n\n except Exception as registration_error:\n failed_tool_name = getattr(tool, 'name', 'unknown')\n print(f\"ERROR: Failed to create/register wrapper for tool '{failed_tool_name}': {registration_error}\")\n print(traceback.format_exc())\n return False\n\n def run(self, transport: str = \"stdio\"):\n \"\"\"运行 MCP 服务器 (签名中移除了 host/port)\"\"\"\n print(f\"DEBUG: MentisMCPServer.run(transport='{transport}') called.\")\n print(f\"正在启动 MCP 服务器，传输方式: {transport}\")\n\n if transport == \"sse\":\n # SSE 方式\n host = 'N/A'\n port = 'N/A'\n if hasattr(self.mcp, 'settings'):\n host = getattr(self.mcp.settings, 'host', 'N/A')\n port = getattr(self.mcp.settings, 'port', 'N/A')\n print(f\"配置 SSE 服务器监听在: http://{host}:{port} (如果 N/A 表示未配置或获取失败)\")\n try:\n import importlib\n try:\n fastmcp_module = importlib.import_module('mcp.server.fastmcp')\n print(f\"FastMCP version: {getattr(fastmcp_module, '__version__', '未知')}\")\n except:\n pass\n import uvicorn\n import fastapi\n print(f\"FastAPI: {fastapi.__version__}, Uvicorn: {uvicorn.__version__}\")\n print(f\"DEBUG: Calling self.mcp.run(transport='{transport}') for SSE\")\n self.mcp.run(transport=transport)\n except Exception as e:\n print(f\"SSE 服务器启动失败: {e}\")\n print(traceback.format_exc())\n raise\n else:\n # 默认 stdio 模式\n print(\"启动 stdio 模式服务器...\")\n try:\n print(f\"DEBUG: Calling self.mcp.run(transport='{transport}') for STDIO\")\n self.mcp.run(transport=transport)\n except Exception as e:\n print(f\"stdio 服务器启动失败: {e}\")\n print(traceback.format_exc())\n raise\n"
},
{
"path": "core/mcp/test/README.md",
"content": "# MCP 测试框架说明\n\n## 概述\n\nMCP（Machine Conversation Protocol）是一个用于机器对话的协议框架，它允许不同的系统通过标准化的接口进行通信。本测试框架提供了一种方式来测试MCP服务器的功能和性能。\n\n## 测试文件结构\n\n测试框架包含以下主要文件：\n\n### 1. minimal_fastmcp_test.py\n\n这是一个最小化的FastMCP服务器实现，用于测试基本功能：\n\n- 创建FastMCP实例\n- 注册简单的工具函数（ping工具）\n- 通过STDIO传输方式运行服务器\n\n该文件可以独立运行，也可以被其他测试脚本作为子进程启动。\n\n### 2. test_minimal_client.py\n\n这个脚本使用MCP客户端库来测试minimal_fastmcp_test.py：\n\n- 导入必要的MCP客户端库（ClientSession, stdio_client等）\n- 连接到minimal_fastmcp_test.py并测试ping工具\n- 展示如何使用客户端API进行工具调用\n\n## 测试方法\n\n### 客户端库测试（test_minimal_client.py）\n\n这种测试方法使用MCP客户端库与MCP服务器通信，展示了如何在实际应用中使用MCP客户端。测试流程如下：\n\n1. 创建ClientSession对象\n2. 连接到MCP服务器\n3. 调用工具并处理结果\n\n## 运行测试\n\n### 运行客户端库测试\n\n```bash\npython core/mcp/test/test_minimal_client.py\n```\n\n## 扩展测试\n\n### 添加新工具\n\n要在minimal_fastmcp_test.py中添加新工具，可以按照以下步骤操作：\n\n1. 定义新的异步工具函数\n2. 使用FastMCP实例的装饰器注册工具\n\n示例：\n```python\nasync def new_tool(param1: str, param2: int = 0) -> str:\n \"\"\"A new tool description.\"\"\"\n # 工具实现\n return f\"Result: {param1}, {param2}\"\n\nmcp_server.tool(name=\"new_tool\", description=\"New tool description.\")(new_tool)\n```\n\n### 创建新的测试脚本\n\n可以参考现有的测试脚本创建新的测试脚本，测试不同的功能或场景。\n\n## 常见问题\n\n### 服务器无响应\n\n- 确保服务器进程正在运行\n- 检查传输方式是否正确（stdio或sse）\n- 检查客户端连接参数是否正确\n\n### 工具调用失败\n\n- 确保工具名称正确\n- 检查参数是否符合工具的要求\n- 查看服务器日志以获取更多信息\n\n## 总结\n\nMCP测试框架提供了使用MCP客户端库测试MCP服务器功能的方法。通过这些测试，可以验证MCP服务器的基本功能和性能，为开发和调试提供支持。"
},
{
"path": "core/mcp/test/__init__.py",
"content": "# MCP测试模块\n# 包含用于测试MCP（Message Control Protocol）功能的各种测试脚本"
},
{
"path": "core/mcp/test/minimal_fastmcp_test.py",
"content": "import asyncio\nfrom mcp.server.fastmcp import FastMCP\nimport logging\n\n# 配置基本日志，看FastMCP内部是否有更多信息\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(\"minimal_test\")\n\nprint(\"--- Minimal FastMCP Server Test ---\")\n\n# 1. 创建 FastMCP 实例\n# (假设 FastMCP 对于 stdio 不需要 host/port in __init__)\nmcp_server = FastMCP(name=\"MinimalServer\")\nprint(\"FastMCP instance created.\")\n\n# 2. 定义一个简单的 async 工具函数\nasync def ping_tool(query: str = \"default ping\") -> str:\n \"\"\"A very basic tool that just returns pong.\"\"\"\n print(f\"\\n--- PING TOOL CALLED! ---\") # 在工具内部打印日志\n print(f\"Received query: {query}\")\n result = f\"pong: {query}\"\n print(f\"Returning: {result}\")\n print(f\"--- PING TOOL END ---\")\n return result\n\n# 3. 直接用 FastMCP 实例的装饰器注册\ntry:\n mcp_server.tool(name=\"ping\", description=\"Returns pong plus the query.\")(ping_tool)\n # 上一行等价于:\n # @mcp_server.tool(name=\"ping\", description=\"Returns pong plus the query.\")\n # async def ping_tool(...) ...\n print(\"Tool 'ping' registered directly with FastMCP.\")\nexcept Exception as e_reg:\n print(f\"Error registering tool directly: {e_reg}\")\n import traceback\n traceback.print_exc()\n exit(1)\n\n# 4. 运行服务器 (使用 STDIO)\ntry:\n print(\"Starting minimal server with STDIO transport...\")\n # 假设 run() 只需 transport 参数对 stdio 有效\n mcp_server.run(transport=\"stdio\")\n print(\"Server finished.\") # 理应不会执行到，除非服务器停止\nexcept Exception as e_run:\n print(f\"Error running minimal server: {e_run}\")\n import traceback\n traceback.print_exc()\n exit(1)"
},
{
"path": "core/mcp/test/test_minimal_client.py",
"content": "# test_minimal_client_fixed.py - 用于测试minimal_fastmcp_test.py的客户端脚本（修复版）\nimport os\nimport sys\nimport asyncio\nimport json\nimport traceback\nfrom typing import Optional, Dict, Any\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))\n\n# 导入必要的MCP客户端库\ntry:\n from mcp import ClientSession\n from mcp.client.stdio import stdio_client, StdioServerParameters\n from mcp.types import CallToolRequest\n DEPS_OK = True\nexcept ImportError as e:\n print(f\"错误: 缺少必要的依赖: {e}\")\n print(\"请确保已安装mcp库: pip install mcp\")\n DEPS_OK = False\n\nasync def main():\n \"\"\"连接到minimal_fastmcp_test.py并测试ping工具\"\"\"\n print(\"=== MCP最小客户端测试（修复版）===\\n\")\n \n if not DEPS_OK:\n print(\"缺少必要的依赖，无法继续。\")\n return\n \n # 准备minimal_fastmcp_test.py的路径\n script_path = os.path.join(os.path.dirname(__file__), \"minimal_fastmcp_test.py\")\n cmd = [sys.executable, script_path]\n print(f\"准备连接到服务器: {script_path}\")\n \n try:\n # 创建StdioServerParameters对象\n server_params = StdioServerParameters(\n command=sys.executable,\n args=[script_path],\n # 可以根据需要添加其他参数，如env, cwd等\n )\n print(\"已创建服务器参数配置。\")\n \n # 创建STDIO客户端连接\n print(\"\\n创建STDIO客户端连接...\")\n async with stdio_client(server_params) as (reader, writer):\n print(\"STDIO连接已建立。创建ClientSession...\")\n async with ClientSession(reader, writer) as session:\n print(\"ClientSession已创建。初始化会话...\")\n await session.initialize()\n print(\"会话已初始化。\")\n \n # 获取服务器支持的工具列表\n print(\"\\n获取服务器支持的工具列表...\")\n tools_result = await session.list_tools()\n print(f\"服务器支持的工具: {tools_result}\")\n \n # 调用ping工具\n print(\"\\n调用ping工具...\")\n ping_request = CallToolRequest(\n method=\"tools/call\",\n params={\n \"name\": \"ping\",\n \"arguments\": {\"query\": \"Hello, MCP!\"}\n }\n )\n \n try:\n print(f\"发送请求: {ping_request}\")\n result = await session.call_tool(\"ping\", {\"query\": \"Hello, MCP!\"})\n print(f\"\\n收到响应: {result}\")\n if hasattr(result, 'result'):\n print(f\"结果: {result.result}\")\n elif hasattr(result, 'error'):\n print(f\"错误: {result.error}\")\n else:\n print(f\"未知响应格式: {result}\")\n except Exception as e:\n print(f\"调用工具时出错: {e}\")\n print(traceback.format_exc())\n \n except Exception as e:\n print(f\"运行测试时出错: {e}\")\n print(traceback.format_exc())\n\nif __name__ == \"__main__\":\n asyncio.run(main())"
},
{
"path": "core/tools/__init__.py",
"content": "# Tools package initialization\nfrom langchain_community.agent_toolkits.load_tools import load_tools\nfrom core.tools.registry import register_tool, ToolCategory, get_registered_tools\nfrom core.tools.firecrawl_tool import FireCrawlTool\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\nimport os\nimport importlib\nimport inspect\nfrom typing import Any, Dict, List, Type, Optional\nfrom langchain_core.tools import BaseTool\n\n# 导入预注册所需的工具\nfrom langchain_community.tools import (\n TavilySearchResults,\n ArxivQueryRun,\n)\nfrom langchain_community.agent_toolkits import FileManagementToolkit\nfrom langchain_community.agent_toolkits.openapi.toolkit import RequestsToolkit,TextRequestsWrapper\nfrom langchain_community.tools.riza.command import ExecPython, ExecJavaScript\n\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\n\n# 预注册核心工具列表 - 定义需要预注册的核心工具\ndef preregister_core_tools():\n \"\"\"预注册核心工具，确保系统启动时这些工具已经可用\"\"\"\n print(\"开始预注册核心工具...\")\n \n # 注册搜索类工具\n try:\n # Tavily搜索工具\n tavily_search = TavilySearchResults()\n register_tool(tavily_search, ToolCategory.SEARCH)\n print(f\"已预注册工具: {tavily_search.name} (类别: {ToolCategory.SEARCH.value})\")\n except Exception as e:\n print(f\"预注册Tavily搜索工具失败: {e}\")\n \n # 注册网页浏览类工具\n try:\n # Arxiv查询工具\n arxiv_tool = ArxivQueryRun()\n register_tool(arxiv_tool, ToolCategory.WEB_BROWSING)\n print(f\"已预注册工具: {arxiv_tool.name} (类别: {ToolCategory.WEB_BROWSING.value})\")\n except Exception as e:\n print(f\"预注册Arxiv查询工具失败: {e}\")\n \n try:\n # RequestoolKit请求工具\n # 创建TextRequestsWrapper实例作为请求包装器\n requests_wrapper = TextRequestsWrapper(headers={})\n # 初始化RequestsToolkit，提供必要的参数\n requests_toolkit = RequestsToolkit(\n requests_wrapper=requests_wrapper,\n allow_dangerous_requests=True # 允许危险请求，使工具可用\n )\n for req_tool in requests_toolkit.get_tools():\n register_tool(req_tool, ToolCategory.WEB_BROWSING)\n print(f\"已预注册工具: {req_tool.name} (类别: {ToolCategory.WEB_BROWSING.value})\")\n except Exception as e:\n print(f\"预注册 RequestoolKit请求工具失败: {e}\")\n \n # 注册文件系统工具\n try:\n # 获取当前目录作为文件系统工具的根目录\n current_dir = os.getcwd()\n # 创建文件系统工具集\n filesystem_toolkit = FileManagementToolkit(\n root_dir=current_dir,\n selected_tools=[\"write_file\", \"read_file\", \"list_directory\"]\n )\n # 获取文件系统工具并注册\n for fs_tool in filesystem_toolkit.get_tools():\n register_tool(fs_tool, ToolCategory.FILE_SYSTEM)\n print(f\"已预注册工具: {fs_tool.name} (类别: {ToolCategory.FILE_SYSTEM.value})\")\n except Exception as e:\n print(f\"预注册文件系统工具失败: {e}\")\n \n # 注册代码解释器工具\n # try:\n # # Python REPL工具\n # python_repl = ExecPython()\n # register_tool(python_repl, ToolCategory.CODE_INTERPRETER)\n # print(f\"已预注册工具: {python_repl.name} (类别: {ToolCategory.CODE_INTERPRETER.value})\")\n # except Exception as e:\n # print(f\"预注册Python REPL工具失败: {e}\")\n\n # # 注册代码解释器工具\n # try:\n # # Python REPL工具\n # javascript_repl = ExecJavaScript()\n # register_tool(javascript_repl, ToolCategory.CODE_INTERPRETER)\n # print(f\"已预注册工具: {javascript_repl.name} (类别: {ToolCategory.CODE_INTERPRETER.value})\")\n # except Exception as e:\n # print(f\"预注册Python REPL工具失败: {e}\")\n \n # 注册自定义工具 - FireCrawl工具\n try:\n firecrawl_tool = FireCrawlTool()\n register_tool(firecrawl_tool, ToolCategory.WEB_BROWSING)\n print(f\"已预注册工具: {firecrawl_tool.name} (类别: {ToolCategory.WEB_BROWSING.value})\")\n except Exception as e:\n print(f\"预注册FireCrawl工具失败: {e}\")\n \n # 注册E2B代码解释器工具\n try:\n e2b_tool = E2BCodeInterpreterTool()\n register_tool(e2b_tool, ToolCategory.CODE_INTERPRETER)\n print(f\"已预注册工具: {e2b_tool.name} (类别: {ToolCategory.CODE_INTERPRETER.value})\")\n except Exception as e:\n print(f\"预注册E2B代码解释器工具失败: {e}\")\n\n\n from .replicate_flux_tool import ReplicateFluxImageTool, category \n try:\n flux_tool = ReplicateFluxImageTool()\n if flux_tool._is_available:\n register_tool(flux_tool, category)\n except Exception as e:\n print(f\"Failed to register ReplicateFluxImageTool: {e}\")\n\nprint(\"核心工具预注册完成\")\n\n# 执行预注册\npreregister_core_tools()\n\n# 注册 LangChain 工具 - 使用load_tools加载的工具列表\ntry:\n langchain_tools = load_tools([\"serpapi\"])\n for tool in langchain_tools:\n register_tool(tool, ToolCategory.SEARCH)\n print(f\"已注册LangChain工具: {tool.name} (类别: {ToolCategory.SEARCH.value})\")\nexcept Exception as e:\n print(f\"加载LangChain工具失败: {e}\")\n\n# 工具类别映射 - 用于自动分类直接导入的工具\ntool_category_mapping = {\n # 搜索类工具\n \"TavilySearchResults\": ToolCategory.SEARCH,\n \"GoogleSearchResults\": ToolCategory.SEARCH,\n \"GoogleSerperResults\": ToolCategory.SEARCH,\n \"WikipediaQueryRun\": ToolCategory.SEARCH,\n \"FireCrawl\": ToolCategory.SEARCH,\n \n # 网页浏览类工具\n \"WebBrowser\": ToolCategory.WEB_BROWSING,\n \"ArxivQueryRun\": ToolCategory.WEB_BROWSING,\n \"RequestsGet\": ToolCategory.WEB_BROWSING,\n \"RequestsPost\": ToolCategory.WEB_BROWSING,\n \n # 文件系统类工具\n \"WriteFile\": ToolCategory.FILE_SYSTEM,\n \"ReadFile\": ToolCategory.FILE_SYSTEM,\n \"ListDirectory\": ToolCategory.FILE_SYSTEM,\n \n # 代码解释器类工具\n \"PythonREPL\": ToolCategory.CODE_INTERPRETER,\n \"ShellTool\": ToolCategory.CODE_INTERPRETER,\n \"E2BCodeInterpreterTool\": ToolCategory.CODE_INTERPRETER,\n \n # 数据库类工具\n \"SQLDatabaseTool\": ToolCategory.DATABASE,\n \n # 默认为其他类别\n \"default\": ToolCategory.OTHER\n}\n\ndef register_direct_tool(tool_instance: BaseTool, category: ToolCategory = None) -> None:\n \"\"\"注册直接从langchain_community.tools导入的工具\n \n Args:\n tool_instance: 工具实例\n category: 工具类别，如果为None则自动根据工具名称判断类别\n \"\"\"\n if not category:\n # 获取工具类名\n tool_class_name = tool_instance.__class__.__name__\n # 根据工具类名自动判断类别\n category = tool_category_mapping.get(tool_class_name, tool_category_mapping[\"default\"])\n \n # 注册工具\n register_tool(tool_instance, category)\n print(f\"已注册工具: {tool_instance.name} (类别: {category.value})\")\n\n# 获取 tools 目录路径\ntools_dir = os.path.dirname(__file__)\n\n# 遍历目录中的所有文件，注册自定义工具\nfor filename in os.listdir(tools_dir):\n # 只处理 .py 文件，且排除 __init__.py 和 registry.py\n if filename.endswith('.py') and filename not in ['__init__.py', 'registry.py']:\n # 提取模块名（去掉 .py 后缀）\n module_name = filename[:-3]\n try:\n # 动态导入模块\n module = importlib.import_module(f'.{module_name}', package='core.tools')\n \n # 查找模块中的工具类（继承自BaseTool的类）\n for name, obj in inspect.getmembers(module):\n # 检查是否是类且是BaseTool的子类\n if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool:\n # 检查该类是否已经被实例化并注册\n tool_name = getattr(obj, 'name', None)\n if tool_name and tool_name not in [info['tool'].name for info in get_registered_tools().values()]:\n # 确定工具类别\n category = getattr(module, 'category', ToolCategory.OTHER)\n # 实例化并注册工具\n try:\n tool_instance = obj()\n register_tool(tool_instance, category)\n print(f\"已注册工具类: {name} (工具名: {tool_instance.name}, 类别: {category.value})\")\n except Exception as e:\n print(f\"实例化工具类 {name} 时出错: {e}\")\n except Exception as e:\n print(f\"导入 {module_name} 时出错: {e}\")"
},
{
"path": "core/tools/e2b_tool.py",
"content": "# core/tools/e2b_tool.py\n\nimport os\nimport json\nimport asyncio\nimport traceback\nfrom typing import Dict, Any, Optional, Type, List # 确保导入 List\nfrom pydantic import BaseModel, Field, PrivateAttr\nfrom langchain_core.tools import BaseTool\n\n\n# --- E2B Imports ---\ntry:\n from e2b_code_interpreter import Sandbox\n from e2b_code_interpreter.exceptions import TimeoutException\n E2B_AVAILABLE = True\nexcept ImportError:\n Sandbox = None # type: ignore\n SandboxException = Exception # type: ignore # Fallback to base Exception\n TimeoutException = TimeoutError # type: ignore # Fallback to base TimeoutError\n E2B_AVAILABLE = False\n print(\"Warning: 'e2b' package not installed (pip install e2b). E2BCodeInterpreterTool will not work.\")\n\n# --- Tool Category ---\ntry:\n from .registry import ToolCategory, register_tool\n if not hasattr(ToolCategory, 'CODE_INTERPRETER'):\n ToolCategory.CODE_INTERPRETER = ToolCategory.OTHER\n category = ToolCategory.CODE_INTERPRETER\nexcept ImportError:\n category = None\n print(\"Tool registry not found.\")\n\n# --- Input Schema (保持不变) ---\nclass E2BCodeInterpreterToolInput(BaseModel):\n code: str = Field(description=\"要执行的Python代码\")\n\n# --- Tool Class (优化版) ---\nclass E2BCodeInterpreterTool(BaseTool):\n \"\"\"\n 使用 E2B SDK 在安全沙箱中执行 Python 代码的工具 (修正异常处理版)。\n 返回执行结果的字符串摘要。\n \"\"\"\n name: str = \"e2b_code_interpreter\"\n description: str = ( # 可以稍微调整描述，强调是 Python 执行环境\n \"Executes Python code in a sandboxed environment. \"\n \"Input MUST be a JSON object with a 'code' key containing the Python code string. \"\n \"Libraries like matplotlib, pandas, numpy, sympy are available. Install others using pip (e.g., `import subprocess; subprocess.run(['pip', 'install', 'requests'])`). \"\n \"Use 'print()' to output results. For plots, save them to a file (e.g., '/home/user/plot.png') and state the path; do not return raw image data. \"\n \"Returns a string summarizing execution status, stdout, stderr, and any errors.\"\n )\n args_schema: Type[BaseModel] = E2BCodeInterpreterToolInput\n\n _sandbox: Optional[Any] = PrivateAttr(default=None)\n _is_available: bool = PrivateAttr(default=False)\n _init_error: Optional[str] = PrivateAttr(default=None)\n # 不再需要 self.ExceptionClass\n\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n self._initialize_sandbox()\n\n def _initialize_sandbox(self):\n \"\"\"初始化沙箱环境\"\"\"\n if not E2B_AVAILABLE:\n self._init_error = \"Package 'e2b' not installed.\"\n print(f\"ERROR: {self._init_error}\")\n return\n\n if \"E2B_API_KEY\" not in os.environ:\n self._init_error = \"Environment variable E2B_API_KEY not set.\"\n print(f\"ERROR: {self._init_error}\")\n return\n\n try:\n print(\"Initializing E2B Sandbox...\")\n # 实例化 Sandbox\n self._sandbox = Sandbox() # 使用导入的 Sandbox 类\n print(\"E2B Sandbox initialized successfully!\")\n self._is_available = True\n self._init_error = None\n except (SandboxException, TimeoutException) as e: # <--- 捕获特定的 E2B 异常\n self._init_error = f\"Failed to initialize E2B Sandbox (E2B Error): {e}\"\n print(f\"ERROR: {self._init_error}\")\n self._is_available = False\n except Exception as e: # 捕获其他意外错误\n self._init_error = f\"An unexpected error occurred during E2B Sandbox initialization: {e}\"\n print(f\"ERROR: {self._init_error}\")\n self._is_available = False\n\n def _run(self, code: str, **kwargs) -> str:\n \"\"\"同步执行 Python 代码并返回结果摘要字符串\"\"\"\n if not self._is_available or self._sandbox is None:\n # ... (返回包含设置指南的错误信息，不变) ...\n error_message = \"E2B Sandbox is not available\"\n if self._init_error: error_message += f\": {self._init_error}\"\n setup_guide = \"\\n\\nSetup: pip install e2b; export E2B_API_KEY='...'\"\n return f\"Execution Failed: {error_message}{setup_guide}\"\n\n output_summary = \"\"\n try:\n print(f\"--- E2B: Executing code synchronously ---\\n{code}\\n--------------------------------------\")\n # 使用 run_python 方法\n execution = self._sandbox.run_code(code)\n\n # 构建结果字符串 (逻辑保持不变)\n if execution.error:\n output_summary += f\"Execution Failed!\\nError Name: {execution.error.name}\\nError Value: {execution.error.value}\\n\"\n if execution.error.traceback:\n traceback_lines = execution.error.traceback.splitlines()\n output_summary += f\"Traceback (last few lines):\\n...\\n\" + \"\\n\".join(traceback_lines[-5:])\n else:\n output_summary += \"Execution Successful.\\n\"\n if execution.logs.stdout: output_summary += f\"\\nSTDOUT:\\n{execution.logs.stdout}\"\n if execution.logs.stderr: output_summary += f\"\\nSTDERR:\\n{execution.logs.stderr}\"\n if execution.results: output_summary += \"\\n\\nNote: Execution produced structured results (e.g., plots saved as files).\"\n if not output_summary.strip() or output_summary.strip() == \"Execution Successful.\": output_summary = \"Code executed successfully with no textual output.\"\n\n print(f\"--- E2B: Execution finished ---\\nResult Summary:\\n{output_summary[:500]}...\\n-----------------------------\")\n return output_summary.strip()\n\n except (SandboxException, TimeoutException) as e: # <--- 捕获特定的 E2B 异常\n error_str = f\"Execution Failed (E2B Error)!\\nError Name: {getattr(e, 'name', type(e).__name__)}\\nDetails: {e}\"\n # TimeoutException 可能没有 traceback 属性，SandboxException 通常有\n tb = getattr(e, 'traceback', traceback.format_exc())\n if tb:\n tb_lines = tb.splitlines()\n error_str += f\"\\nTraceback (last few lines):\\n...\\n\" + \"\\n\".join(tb_lines[-5:])\n print(f\"ERROR during E2B execution: {error_str}\")\n return error_str\n except Exception as e: # 其他错误\n error_str = f\"Execution Failed (Unexpected Error)!\\nError Type: {type(e).__name__}\\nError Details: {str(e)}\\nTraceback:\\n{traceback.format_exc()}\"\n print(f\"ERROR during E2B execution: {error_str}\")\n return error_str\n\n async def _arun(self, code: str, **kwargs) -> str:\n \"\"\"异步执行 Python 代码并返回结果摘要字符串\"\"\"\n if not self._is_available or self._sandbox is None:\n # ... (返回错误信息) ...\n error_message = f\"E2B Sandbox is not available: {self._init_error}\"\n return f\"Execution Failed: {error_message}\"\n\n try:\n loop = asyncio.get_running_loop()\n import functools\n # 注意：传递给 run_in_executor 的函数应该是可调用的\n # 这里 _run 是实例方法，所以直接传递 self._run 即可\n # 但为了确保 code 参数被正确传递，可以用 lambda 或 partial\n sync_run_with_args = functools.partial(self._run, code=code, **kwargs)\n\n print(f\"--- E2B: Executing code asynchronously via executor ---\\n{code}\\n--------------------------------------\")\n result_summary = await loop.run_in_executor(\n None, sync_run_with_args\n )\n print(f\"--- E2B: Async execution finished ---\")\n return result_summary\n except Exception as e: # run_in_executor 或 _run 的异常会在这里捕获\n error_str = f\"Execution Failed (Async Wrapper Error)!\\nError Type: {type(e).__name__}\\nError Details: {str(e)}\"\n # 尝试获取 Traceback\n tb = traceback.format_exc()\n error_str += f\"\\nTraceback:\\n{tb}\"\n print(f\"ERROR during E2B async execution: {error_str}\")\n return error_str\n\n\n def close(self):\n \"\"\"关闭沙箱，释放资源。\"\"\"\n if hasattr(self, \"sandbox\") and self._is_available and self._sandbox is not None:\n try:\n print(\"Attempting to close E2B Sandbox...\")\n self._sandbox.kill()\n print(\"E2B Sandbox closed successfully.\")\n self._is_available = False\n self._sandbox = None\n except (SandboxException, TimeoutException) as e: # 捕获特定异常\n print(f\"Error closing E2B Sandbox (E2B Error): {e}\")\n except Exception as e:\n print(f\"An unexpected error occurred while closing E2B Sandbox: {e}\")\n\n model_config = {\n \"arbitrary_types_allowed\": True\n }\n\n # __del__ 方法用于对象销毁，通常不保证执行，不建议依赖它来关闭资源\n # def __del__(self): self.close()"
},
{
"path": "core/tools/firecrawl_tool.py",
"content": "# 文件路径: core/tools/firecrawl_tool.py (或您存放工具的文件)\n\nimport os\nimport json # 虽然不直接返回 JSON，但可能用于处理 metadata\nfrom typing import Dict, List, Literal, Optional, Tuple, Type, Union, Any # 确保导入\nfrom pydantic import BaseModel, Field, PrivateAttr # 导入 PrivateAttr\nfrom langchain_core.callbacks import (\n AsyncCallbackManagerForToolRun,\n CallbackManagerForToolRun,\n)\nfrom langchain_core.tools import BaseTool\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\n\n\n# 尝试导入 FireCrawlLoader，如果失败则标记\ntry:\n from langchain_community.document_loaders import FireCrawlLoader\n FIRECRAWL_LOADER_AVAILABLE = True\nexcept ImportError:\n FireCrawlLoader = None # type: ignore\n FIRECRAWL_LOADER_AVAILABLE = False\n print(\"Warning: langchain_community or firecrawl-py not installed? FireCrawlLoader unavailable.\")\n print(\"Run: pip install -U langchain-community firecrawl-py\")\n\n# 定义输入 Schema (保持不变)\nclass FireCrawlInput(BaseModel):\n \"\"\"Input for the FireCrawl tool.\"\"\"\n url: str = Field(description=\"URL to crawl or scrape\")\n mode: str = Field(\n default=\"scrape\", # <-- 将默认模式改为 'scrape' 可能更常用\n description=\"Mode: 'scrape' (single page), 'crawl' (multiple pages). Default: 'scrape'\",\n )\n # 可以添加 params 字段如果希望 LLM 控制更多参数\n # params: Optional[Dict[str, Any]] = Field(default=None, description=\"Optional dictionary of additional FireCrawl parameters (e.g., {'pageOptions': {'onlyMainContent': True}})\")\n\n\nclass FireCrawlTool(BaseTool):\n \"\"\"\n Tool that uses FireCrawl API to crawl or scrape web content and return a summary.\n\n Setup:\n pip install -U langchain-community firecrawl-py\n export FIRECRAWL_API_KEY=\"your-api-key\"\n\n Instantiate:\n tool = FireCrawlTool() # Reads API key from env\n # Or explicitly: tool = FireCrawlTool(api_key=\"...\")\n\n Invoke:\n tool.invoke({\"url\": \"https://example.com\", \"mode\": \"scrape\"})\n \"\"\"\n\n name: str = \"firecrawl_web_content\" # 建议用更描述性的名字\n description: str = (\n \"Fetches and extracts the main textual content from a given URL. \"\n \"Use 'scrape' mode (default) for a single page, or 'crawl' mode to follow links (use sparingly). \"\n \"Input should be a URL. Returns a textual summary of the content.\"\n )\n args_schema: Type[BaseModel] = FireCrawlInput\n\n # --- 配置属性 ---\n # API Key 可以通过 __init__ 传入，或者留空让 loader 从环境变量读取\n _api_key: Optional[str] = PrivateAttr(default=None) # 使用 PrivateAttr 避免 Pydantic 验证\n _api_url: Optional[str] = PrivateAttr(default=None)\n # 可以在 __init__ 中设置默认 mode 和 params，或者在 _run/_arun 中处理\n default_mode: str = \"scrape\" # 工具级别的默认模式\n default_params: Dict[str, Any] = Field(default_factory=dict) # 工具级别的默认参数\n\n # 添加 __init__ 以便可以传入 api_key (可选)\n def __init__(self, api_key: Optional[str] = None, api_url: Optional[str] = None,\n mode: str = \"scrape\", params: Optional[Dict[str, Any]] = None, **kwargs):\n super().__init__(**kwargs)\n # Pydantic V2 中，非 model 字段需要用 PrivateAttr 或在 model_config 中设置\n self._api_key = api_key\n self._api_url = api_url\n self.default_mode = mode\n self.default_params = params or {}\n # 检查 Loader 是否可用\n if not FIRECRAWL_LOADER_AVAILABLE:\n print(\"ERROR: FireCrawlLoader is unavailable. Please install required packages.\")\n\n def _run(\n self,\n url: str,\n mode: Optional[str] = None,\n run_manager: Optional[CallbackManagerForToolRun] = None,\n ) -> str: # <--- 返回值必须是字符串\n \"\"\"使用工具同步获取网页内容。\"\"\"\n if not FIRECRAWL_LOADER_AVAILABLE:\n return \"Error: FireCrawlLoader is not available. Required packages might be missing.\"\n \n # 确定使用的 API Key (优先实例属性，其次环境变量)\n key = self._api_key or os.getenv('FIRECRAWL_API_KEY')\n if not key:\n return \"Error: FIRECRAWL_API_KEY not found in environment variables or instantiation.\"\n \n # 打印 Debug 信息 (可选)\n print(f\"DEBUG [FireCrawlTool]: Running for URL='{url}', Mode='{mode or self.default_mode}'\")\n # print(f\"DEBUG [FireCrawlTool]: Effective API Key = {'*' * (len(key) - 4) + key[-4:] if key else 'None'}\")\n\n try:\n current_mode = mode or self.default_mode\n loader = FireCrawlLoader(\n url=url,\n api_key=key, # 传递最终确定的 key\n api_url=self._api_url, # 传递实例属性或 None\n mode=current_mode,\n params=self.default_params, # 传递实例默认参数\n )\n\n print(f\"--- Calling FireCrawl API (Sync) for: '{url}' ---\")\n docs = loader.load()\n print(f\"--- FireCrawl API call successful for: '{url}', received {len(docs)} document(s) ---\")\n\n # --- 格式化结果为字符串 ---\n if not docs:\n return f\"FireCrawl successful but returned no content from {url} (Mode: {current_mode}). The page might be empty or restricted.\"\n\n summary_parts = [f\"Content summary from {url} (Mode: {current_mode}):\"]\n content_limit = 4000 # 限制返回给 LLM 的总字符数 (可调整)\n current_length = len(summary_parts[0])\n doc_count = 0\n\n for doc in docs:\n # 可以考虑只返回第一个文档的内容，如果文档很多\n # if doc_count >= 1 and current_mode == 'scrape': break \n \n source_info = f\"\\n\\n--- Source: {doc.metadata.get('sourceURL', url)} ---\"\n page_content = doc.page_content or \"\"\n \n available_length = content_limit - current_length - len(source_info) - 20 # 预留空间\n if available_length <= 0 and doc_count > 0: # 如果已经有内容且空间不足\n summary_parts.append(\"\\n\\n... (further content truncated)\")\n break\n\n content = source_info + \"\\n\" + page_content\n \n if len(content) > available_length:\n content = content[:available_length] + \"... (truncated)\"\n \n summary_parts.append(content)\n current_length += len(content)\n doc_count += 1\n if current_length >= content_limit: break # 达到总长度限制\n\n return \"\\n\".join(summary_parts).strip()\n # --- 格式化结束 ---\n\n except Exception as e:\n error_msg = f\"Error during FireCrawl for {url} (Mode: {mode or self.default_mode}): {repr(e)}\"\n print(f\"ERROR: {error_msg}\")\n return error_msg # 返回错误信息字符串\n\n async def _arun(\n self,\n url: str,\n mode: Optional[str] = None,\n run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n ) -> str: # <--- 返回值必须是字符串\n \"\"\"使用工具异步获取网页内容。\"\"\"\n if not FIRECRAWL_LOADER_AVAILABLE:\n return \"Error: FireCrawlLoader is not available.\"\n \n key = self.api_key or os.getenv('FIRECRAWL_API_KEY')\n if not key:\n return \"Error: FIRECRAWL_API_KEY not found.\"\n \n print(f\"DEBUG [FireCrawlTool]: Running async for URL='{url}', Mode='{mode or self.default_mode}'\")\n\n try:\n current_mode = mode or self.default_mode\n loader = FireCrawlLoader(\n url=url, api_key=key, api_url=self.api_url,\n mode=current_mode, params=self.default_params,\n )\n\n print(f\"--- Calling FireCrawl API (Async) for: '{url}' ---\")\n # 使用 aload 进行异步加载\n docs = await loader.aload()\n print(f\"--- FireCrawl API call successful for: '{url}', received {len(docs)} document(s) ---\")\n\n # --- 格式化结果为字符串 (与 _run 逻辑相同) ---\n if not docs: return f\"FireCrawl successful but returned no content from {url} (Mode: {current_mode}).\"\n summary_parts = [f\"Content summary from {url} (Mode: {current_mode}):\"]\n content_limit = 4000; current_length = len(summary_parts[0]); doc_count = 0\n for doc in docs:\n # if doc_count >= 1 and current_mode == 'scrape': break\n source_info = f\"\\n\\n--- Source: {doc.metadata.get('sourceURL', url)} ---\"\n page_content = doc.page_content or \"\"\n available_length = content_limit - current_length - len(source_info) - 20\n if available_length <= 0 and doc_count > 0:\n summary_parts.append(\"\\n\\n... (further content truncated)\"); break\n content = source_info + \"\\n\" + page_content\n if len(content) > available_length: content = content[:available_length] + \"... (truncated)\"\n summary_parts.append(content); current_length += len(content); doc_count += 1\n if current_length >= content_limit: break\n return \"\\n\".join(summary_parts).strip()\n # --- 格式化结束 ---\n\n except Exception as e:\n error_msg = f\"Error during Async FireCrawl for {url} (Mode: {mode or self.default_mode}): {repr(e)}\"\n print(f\"ERROR: {error_msg}\")\n return error_msg\n\n # Pydantic V2: 允许额外的私有属性\n model_config = {\n \"arbitrary_types_allowed\": True\n }"
},
{
"path": "core/tools/registry.py",
"content": "from enum import Enum\nfrom typing import List, Dict, Union, Optional\nfrom langchain.tools import Tool\n\n# 定义工具分类枚举\nclass ToolCategory(Enum):\n SEARCH = \"Search\"\n CODE_INTERPRETER = \"Code Interpreter\"\n WEB_BROWSING = \"Web Browsing\"\n DATABASE = \"Database\"\n FILE_SYSTEM = \"FileSystem\"\n IMAGE_GENERATION = \"Image Generation\"\n OTHER = \"Other\"\n\n# 全局工具注册表\n_registered_tools = {}\n\ndef register_tool(tool: Tool, category: ToolCategory) -> None:\n \"\"\"注册一个工具到全局字典中，带有分类信息\n \n 如果工具名已存在，将覆盖现有的工具注册信息\n \"\"\"\n if tool.name in _registered_tools:\n print(f\"警告: 工具名 {tool.name} 已存在，将覆盖现有注册信息\")\n _registered_tools[tool.name] = {\n \"tool\": tool,\n \"category\": category\n }\n\ndef get_registered_tools(as_dict: bool = False) -> Union[List[Tool], Dict[str, Dict]]:\n \"\"\"返回所有已注册的工具\n \n Args:\n as_dict: 如果为True，返回原始字典格式；如果为False，返回工具列表\n \n Returns:\n 如果as_dict为True，返回原始字典格式；否则返回工具列表\n \"\"\"\n if as_dict:\n return _registered_tools\n return [info[\"tool\"] for info in _registered_tools.values()]\n\ndef get_tools_list() -> List[Tool]:\n \"\"\"返回所有已注册的工具列表，直接可用于Agent初始化\n \n Returns:\n 所有已注册工具的列表\n \"\"\"\n return [info[\"tool\"] for info in _registered_tools.values()]\n\ndef get_tools_dict() -> Dict[str, Tool]:\n \"\"\"返回工具名称到工具实例的映射字典\n \n Returns:\n 工具名称到工具实例的映射字典\n \"\"\"\n return {name: info[\"tool\"] for name, info in _registered_tools.items()}\n\ndef get_tool(name: str) -> Optional[Dict]:\n \"\"\"根据名称获取工具及其分类\n \n Args:\n name: 工具名称\n \n Returns:\n 包含工具和分类的字典，如果工具不存在则返回None\n \"\"\"\n tool_info = _registered_tools.get(name)\n if tool_info:\n return {\n \"tool\": tool_info[\"tool\"],\n \"category\": tool_info[\"category\"].value\n }\n return None\n\ndef get_tool_instance(name: str) -> Optional[Tool]:\n \"\"\"根据名称直接获取工具实例\n \n Args:\n name: 工具名称\n \n Returns:\n 工具实例，如果工具不存在则返回None\n \"\"\"\n tool_info = _registered_tools.get(name)\n return tool_info[\"tool\"] if tool_info else None\n\ndef get_tools_by_category(category: ToolCategory, return_instances: bool = True) -> List[Union[str, Tool]]:\n \"\"\"返回指定分类的工具列表\n \n Args:\n category: 工具分类\n return_instances: 如果为True，返回工具实例列表；如果为False，返回工具名称列表\n \n Returns:\n 工具实例列表或工具名称列表\n \"\"\"\n if return_instances:\n return [info[\"tool\"] for name, info in _registered_tools.items() if info[\"category\"] == category]\n return [name for name, info in _registered_tools.items() if info[\"category\"] == category]"
},
{
"path": "core/tools/replicate_flux_tool.py",
"content": "# 文件路径: core/tools/replicate_flux_tool.py (或类似)\n\nimport os\nimport asyncio\nimport json\nfrom typing import Dict, Any, Optional, Type, List, Literal\nfrom pydantic import BaseModel, Field, PrivateAttr\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.callbacks import (\n AsyncCallbackManagerForToolRun,\n CallbackManagerForToolRun,\n)\n\n# --- Replicate Client ---\ntry:\n import replicate\n REPLICATE_AVAILABLE = True\nexcept ImportError:\n replicate = None # type: ignore\n REPLICATE_AVAILABLE = False\n print(\"Warning: 'replicate' package not installed (pip install replicate). ReplicateFluxImageTool will not work.\")\n\n# --- Tool Category (可选, 用于 Registry) ---\ntry:\n from .registry import ToolCategory, register_tool\n if not hasattr(ToolCategory, 'IMAGE_GENERATION'):\n ToolCategory.IMAGE_GENERATION = ToolCategory.OTHER\n category = ToolCategory.IMAGE_GENERATION\nexcept ImportError:\n category = None\n print(\"Tool registry not found. Cannot auto-register ReplicateFluxImageTool.\")\n\n\n# --- Input Schema based on flux-dev ---\nclass ReplicateFluxToolInput(BaseModel):\n \"\"\"Input schema for the Replicate Flux Image Generator Tool.\"\"\"\n prompt: str = Field(description=\"Required. Detailed text description of the image to be generated.\")\n aspect_ratio: Literal[\"1:1\", \"16:9\", \"21:9\", \"3:2\", \"2:3\", \"4:5\", \"5:4\", \"3:4\", \"4:3\", \"9:16\", \"9:21\"] = Field(\n default=\"1:1\", description=\"Aspect ratio for the generated image.\"\n )\n num_outputs: int = Field(\n default=1, description=\"Number of images to generate (1-4).\", ge=1, le=4\n )\n guidance: float = Field(\n default=3.0, description=\"Guidance scale (0-10).\", ge=0, le=10\n )\n num_inference_steps: int = Field(\n default=28, description=\"Number of denoising steps (1-50). Lower is faster, lower quality.\", ge=1, le=50\n )\n seed: Optional[int] = Field(default=None, description=\"Random seed for reproducible generation.\")\n # Add other relevant fields from the schema if needed, e.g., megapixels, output_format\n # megapixels: Literal[\"1\", \"0.25\"] = Field(default=\"1\", description=\"Approximate megapixels for output.\")\n # output_format: Literal[\"webp\", \"jpg\", \"png\"] = Field(default=\"webp\", description=\"Output image format.\")\n\n\n# --- Tool Class (修正返回值处理) ---\nclass ReplicateFluxImageTool(BaseTool):\n \"\"\"Generates images using 'black-forest-labs/flux-dev' on Replicate.\"\"\"\n name: str = \"replicate_flux_image_generator\"\n description: str = (\n \"Generates high-quality images based on a detailed text prompt using the Flux model on Replicate. \"\n \"Specify 'prompt' and optionally other parameters like 'aspect_ratio'. \"\n \"Returns a string containing the URL(s) of the generated image(s).\"\n )\n args_schema: Type[BaseModel] = ReplicateFluxToolInput\n _client: Any = PrivateAttr(default=None)\n _is_available: bool = PrivateAttr(default=False)\n _init_error: Optional[str] = PrivateAttr(default=None)\n model_identifier: str = \"black-forest-labs/flux-dev\"\n\n def __init__(self, api_token: Optional[str] = None, model_id: Optional[str] = None, **kwargs):\n \"\"\"Initialize the Replicate client.\"\"\"\n super().__init__(**kwargs)\n if not REPLICATE_AVAILABLE: self._init_error = \"...\"; print(f\"ERROR: {self._init_error}\"); return\n token = api_token or os.getenv(\"REPLICATE_API_TOKEN\")\n if not token: self._init_error = \"...\"; print(f\"ERROR: {self._init_error}\"); return\n try:\n print(\"Initializing Replicate client...\")\n self._client = replicate.Client(api_token=token)\n print(\"Replicate client initialized successfully.\")\n self._is_available = True; self._init_error = None\n if model_id: self.model_identifier = model_id\n except Exception as e: self._init_error = f\"...: {e}\"; print(f\"ERROR: {self._init_error}\"); self._is_available = False\n\n def _run( self, run_manager: Optional[CallbackManagerForToolRun] = None, **kwargs: Any ) -> str:\n \"\"\"Generates image(s) synchronously.\"\"\"\n if not self._is_available or self._client is None:\n error_message = f\"Replicate client unavailable: {self._init_error}\"\n print(f\"ERROR: {error_message}\"); return f\"Error: {error_message}\"\n\n input_data = {k: v for k, v in kwargs.items() if v is not None and k in self.args_schema.__fields__}\n prompt_short = str(input_data.get('prompt', ''))[:100]\n print(f\"--- TOOL CALL: {self.name} ---\")\n print(f\" Input: Prompt='{prompt_short}...', Args={ {k:v for k,v in input_data.items() if k != 'prompt'} }\")\n\n try:\n # output 现在预期是包含特殊对象 (如 FileOutput 或 URL 字符串) 的列表\n output: List[Any] = self._client.run(self.model_identifier, input=input_data)\n\n if not output or not isinstance(output, list):\n result_str = \"Image generation failed: Replicate API returned no output or unexpected format.\"\n print(f\" Warning: {result_str}\"); return f\"Error: {result_str}\"\n\n # --- 从返回的对象中提取 URL ---\n image_urls: List[str] = []\n for item in output:\n if isinstance(item, str): # 如果直接返回了 URL 字符串\n image_urls.append(item)\n elif hasattr(item, 'url') and isinstance(getattr(item, 'url'), str): # 检查是否有 .url 属性且是字符串\n image_urls.append(getattr(item, 'url'))\n elif hasattr(item, 'read'): # 如果是文件类对象，可能需要其他处理或报错\n print(f\"Warning: Received file-like object from Replicate, cannot directly get URL: {item}\")\n # 或者尝试其他属性？这个需要根据 replicate 库的具体 FileOutput 类型确定\n else:\n print(f\"Warning: Unknown item type in Replicate output list: {type(item)}\")\n\n if not image_urls:\n result_str = \"Image generation succeeded but failed to extract image URLs from the response.\"\n print(f\" Warning: {result_str}\"); return f\"Error: {result_str}\"\n # --- 提取结束 ---\n\n # 格式化 URL 列表为字符串\n url_list_str = \"\\n\".join(image_urls)\n result_str = f\"Successfully generated {len(image_urls)} image(s):\\n{url_list_str}\"\n print(f\" Result: {result_str}\")\n return result_str\n\n except Exception as e: # 捕获 Replicate API 错误等\n # 检查是否是 ReplicateError 并提取更具体的细节\n error_detail = str(e)\n if REPLICATE_AVAILABLE and isinstance(e, replicate.exceptions.ReplicateError):\n error_detail = f\"ReplicateError (Status: {e.status}): {e.title} - {e.detail}\"\n\n error_msg = f\"Error calling Replicate API ({self.model_identifier}): {error_detail}\"\n print(f\" Error: {error_msg}\")\n # traceback.print_exc() # 可以在调试时取消注释\n return f\"Error: {error_msg}\" # 返回错误信息给 LLM\n\n async def _arun( self, run_manager: Optional[AsyncCallbackManagerForToolRun] = None, **kwargs: Any ) -> str:\n \"\"\"Generates image(s) asynchronously using run_in_executor.\"\"\"\n if not self._is_available or self._client is None:\n error_message = f\"Replicate client unavailable: {self._init_error}\"\n print(f\"ERROR: {error_message}\"); return f\"Error: {error_message}\"\n\n input_data = {k: v for k, v in kwargs.items() if v is not None and k in self.args_schema.__fields__}\n prompt_short = str(input_data.get('prompt', ''))[:100]\n print(f\"--- TOOL CALL (Async): {self.name} ---\")\n print(f\" Input: Prompt='{prompt_short}...', Args={ {k:v for k,v in input_data.items() if k != 'prompt'} }\")\n\n try:\n loop = asyncio.get_running_loop()\n import functools\n sync_call_with_args = functools.partial( self._client.run, self.model_identifier, input=input_data )\n output: List[Any] = await loop.run_in_executor( None, sync_call_with_args )\n\n if not output or not isinstance(output, list):\n result_str = \"Async image generation failed: Replicate API returned no output or unexpected format.\"\n print(f\" Warning: {result_str}\"); return f\"Error: {result_str}\"\n\n # --- 从返回的对象中提取 URL (逻辑同 _run) ---\n image_urls: List[str] = []\n for item in output:\n if isinstance(item, str): image_urls.append(item)\n elif hasattr(item, 'url') and isinstance(getattr(item, 'url'), str): image_urls.append(getattr(item, 'url'))\n else: print(f\"Warning: Unknown item type in async Replicate output list: {type(item)}\")\n if not image_urls:\n result_str = \"Async image generation succeeded but failed to extract image URLs.\"\n print(f\" Warning: {result_str}\"); return f\"Error: {result_str}\"\n # --- 提取结束 ---\n\n url_list_str = \"\\n\".join(image_urls)\n result_str = f\"Successfully generated {len(image_urls)} image(s) asynchronously:\\n{url_list_str}\"\n print(f\" Result: {result_str}\")\n return result_str\n\n except Exception as e: # 捕获 Replicate API 错误等\n error_detail = str(e)\n if REPLICATE_AVAILABLE and isinstance(e, replicate.exceptions.ReplicateError):\n error_detail = f\"ReplicateError (Status: {e.status}): {e.title} - {e.detail}\"\n error_msg = f\"Error calling Replicate API asynchronously ({self.model_identifier}): {error_detail}\"\n print(f\" Error: {error_msg}\")\n # traceback.print_exc()\n return f\"Error: {error_msg}\"\n\n\n def close(self):\n \"\"\"关闭沙箱（如果需要的话）。Replicate Client 通常不需要关闭。\"\"\"\n print(f\"Info: Replicate client for '{self.name}' does not require explicit closing.\")\n pass # Replicate client 通常不需要显式关闭\n\n model_config = {\"arbitrary_types_allowed\": True}\n"
},
{
"path": "core/utils/agent_utils.py",
"content": "import os\nfrom typing import Dict, Any, Optional, Literal\nfrom langchain_core.messages import AIMessage, ToolMessage\nimport inspect\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\n \n 这个函数用于在控制台打印Agent的思考过程、工具调用和工具返回结果，\n 便于观察和调试Agent的行为。\n \n Args:\n state: 包含消息历史的状态字典\n \"\"\"\n print(\"\\n\" + \"=\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n # 只打印结果的前500个字符，避免输出过长\n content = latest_message.content\n if len(content) > 500:\n content = content[:500] + \"... (更多内容省略)\"\n print(f\"- 结果: {content}\")\n \n print(\"=\" * 50)\n\ndef save_agent_graph(\n agent, \n caller_file_path: Optional[str] = None,\n output_format: Literal[\"png\", \"svg\", \"mermaid\"] = \"png\",\n custom_filename: Optional[str] = None,\n output_dir: Optional[str] = None\n) -> str:\n \"\"\"保存Agent的图表到指定目录\n \n 这个函数用于生成Agent的图表并保存到指定目录，\n 默认情况下文件名与调用者的文件名保持一致（不含扩展名）。\n \n Args:\n agent: Agent对象，必须有get_graph方法\n caller_file_path: 调用者的文件路径，如果为None则使用调用栈获取\n output_format: 输出格式，可选\"png\"、\"svg\"或\"mermaid\"\n custom_filename: 自定义文件名（不含扩展名），如果提供则使用此名称\n output_dir: 自定义输出目录，如果提供则使用此目录\n \n Returns:\n str: 保存的图表路径\n \"\"\"\n # 如果没有提供调用者文件路径，则从调用栈获取\n if caller_file_path is None:\n # 获取调用者的栈帧\n frame = inspect.currentframe().f_back\n caller_file_path = frame.f_code.co_filename\n \n try:\n # 获取图对象\n graph = agent.get_graph()\n except AttributeError:\n raise ValueError(\"提供的agent对象没有get_graph方法\") \n except Exception as e:\n raise RuntimeError(f\"获取图表时出错: {str(e)}\")\n \n # 确定文件名\n if custom_filename:\n file_name_without_ext = custom_filename\n else:\n # 获取当前文件名（不含路径和扩展名）\n current_file = os.path.basename(caller_file_path)\n file_name_without_ext = os.path.splitext(current_file)[0]\n \n # 确定输出目录\n if output_dir:\n graph_dir = output_dir\n else:\n # 如果调用者在examples目录下，则使用examples/graphs\n # 否则在调用者所在目录创建graphs子目录\n if 'examples' in caller_file_path:\n base_dir = os.path.dirname(os.path.dirname(caller_file_path))\n graph_dir = os.path.join(base_dir, \"examples\", \"graphs\")\n else:\n graph_dir = os.path.join(os.path.dirname(caller_file_path), \"graphs\")\n \n # 确保graphs目录存在\n os.makedirs(graph_dir, exist_ok=True)\n \n # 根据输出格式生成相应文件\n try:\n if output_format == \"png\":\n image_data = graph.draw_mermaid_png()\n graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n with open(graph_path, \"wb\") as f:\n f.write(image_data)\n \n elif output_format == \"svg\":\n image_data = graph.draw_mermaid_svg()\n graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.svg\")\n with open(graph_path, \"wb\") as f:\n f.write(image_data)\n \n elif output_format == \"mermaid\":\n mermaid_code = graph.get_mermaid()\n graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.mmd\")\n with open(graph_path, \"w\") as f:\n f.write(mermaid_code)\n else:\n raise ValueError(f\"不支持的输出格式: {output_format}\")\n \n except Exception as e:\n raise RuntimeError(f\"保存图表时出错: {str(e)}\")\n \n print(f\"图表已保存为 {graph_path}\")\n return graph_path\n\ndef visualize_agent(agent, **kwargs):\n \"\"\"可视化Agent的快捷方法\n \n 这是save_agent_graph的简便包装，用于快速可视化Agent\n \n Args:\n agent: Agent对象\n **kwargs: 传递给save_agent_graph的其他参数\n \n Returns:\n str: 保存的图表路径\n \"\"\"\n # 获取调用者的栈帧\n frame = inspect.currentframe().f_back\n caller_file_path = frame.f_code.co_filename\n \n return save_agent_graph(agent, caller_file_path=caller_file_path, **kwargs)"
},
{
"path": "core/utils/timezone.py",
"content": "from datetime import datetime\nimport os\nfrom typing import Optional\nfrom zoneinfo import ZoneInfo\n\ndef get_timezone() -> str:\n \"\"\"Get timezone from environment variable or use default.\n \n Returns:\n str: Timezone string (e.g. 'Asia/Shanghai')\n \"\"\"\n return os.getenv('TZ', 'UTC')\n\ndef get_formatted_date(timezone: Optional[str] = None) -> str:\n \"\"\"Get formatted date string with timezone awareness.\n \n Args:\n timezone: Optional timezone string. If not provided, uses TZ from env or UTC.\n \n Returns:\n str: Formatted date string (e.g. 'Today's Date: Mon, Jan 01, 2024')\n \"\"\"\n tz = ZoneInfo(timezone or get_timezone())\n now = datetime.now(tz)\n return f\"Today's Date: {now.strftime('%a, %b %d, %Y')}\"\n\ndef get_current_time(timezone: Optional[str] = None) -> datetime:\n \"\"\"Get current time with timezone awareness.\n \n Args:\n timezone: Optional timezone string. If not provided, uses TZ from env or UTC.\n \n Returns:\n datetime: Current time with timezone information\n \"\"\"\n tz = ZoneInfo(timezone or get_timezone())\n return datetime.now(tz)"
},
{
"path": "examples/01_supervisor_test.py",
"content": "from langgraph.prebuilt import create_react_agent\nfrom core.agents.supervisor import create_supervisor\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.func import entrypoint, task\nfrom langgraph.graph import add_messages\nfrom dotenv import load_dotenv\nfrom core.utils.agent_utils import visualize_agent\n\nload_dotenv() # 自动加载 .env 文件\n# 1. 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# Agent 1: Joke Generator (Functional API)\n##############################################################################\n\n@task\ndef generate_joke(messages):\n \"\"\"Generate a short joke (no tool calls).\"\"\"\n system_message = {\n \"role\": \"system\", \n \"content\": \"You are a witty comedian. Write a short joke.\"\n }\n # 直接调用 model.invoke，拼接 system_message + 用户消息\n msg = model.invoke([system_message] + messages)\n return msg\n\n@entrypoint()\ndef joke_agent(state):\n # 调用上面的函数型任务\n joke = generate_joke(state['messages']).result()\n # 将产物插入消息列表\n messages = add_messages(state[\"messages\"], [joke])\n return {\"messages\": messages}\n\njoke_agent.name = \"joke_agent\"\n\n##############################################################################\n# Agent 2: Research Expert (Graph API)\n##############################################################################\n\ndef web_search(query: str) -> str:\n \"\"\"Search the web for information. (Mocked data here)\"\"\"\n return (\n \"Here are the headcounts for each of the FAANG companies in 2024:\\n\"\n \"1. **Facebook (Meta)**: 67,317 employees.\\n\"\n \"2. **Apple**: 164,000 employees.\\n\"\n \"3. **Amazon**: 1,551,000 employees.\\n\"\n \"4. **Netflix**: 14,000 employees.\\n\"\n \"5. **Google (Alphabet)**: 181,269 employees.\"\n )\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[web_search],\n name=\"research_expert\",\n # Prompt 告诉它是一个研究型 Agent，可调用 web_search\n prompt=(\n \"You are a world-class researcher. You have access to a 'web_search(query: str)' tool. \"\n \"Do not do any complicated math, just provide factual info from the web_search if needed.\"\n ),\n)\n\n##############################################################################\n# Supervisor Workflow\n##############################################################################\n\n# 让 Supervisor 在一次对话中可以多轮调用 joke_agent 和 research_expert\n# 这里的 prompt 告诉它：如果用户要“先讲笑话再查信息”，请先调用 joke_agent，再调用 research_expert，\n# 这样可以在同一个用户请求下顺序执行两个 Agent。\n# 这是最简单的示例，只是为了演示 create_supervisor 的基本用法，该方法没有被封装成一个 Agent\n# 也不具备 Planning 能力\nworkflow = create_supervisor(\n [research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions.\\n\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n)\n\n# 编译得到一个可调用的\"App\"\nagent = workflow.compile()\n# 保存为一个可视化的图\n# visualize_agent(agent)\n##############################################################################\n# 测试：单个用户请求想要 \"先讲笑话，再查Apple的2024年人数\" 并合并结果\n##############################################################################\nresult = agent.invoke({\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": (\n \"Hi! I'd like to start with a short joke to lighten the mood, \"\n \"then please check Apple's headcount in 2024. Summarize both.\"\n )\n }\n ]\n})\n\n##############################################################################\n# 打印最终对话消息\n##############################################################################\nfor m in result[\"messages\"]:\n m.pretty_print()"
},
{
"path": "examples/02_supervisor_agent_test.py",
"content": "from langgraph.prebuilt import create_react_agent\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.agents.react_supervisor_agent import SupervisorAgent\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.func import entrypoint, task\nfrom langgraph.graph import add_messages\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\n# 1. 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# Agent 1: Joke Generator (Functional API)\n##############################################################################\n\n@task\ndef generate_joke(messages):\n \"\"\"Generate a short joke (no tool calls).\"\"\"\n system_message = {\n \"role\": \"system\", \n \"content\": \"You are a witty comedian. Write a short joke.\"\n }\n # 直接调用 model.invoke，拼接 system_message + 用户消息\n msg = model.invoke([system_message] + messages)\n return msg\n\n@entrypoint()\ndef joke_agent(state):\n # 调用上面的函数型任务\n joke = generate_joke(state['messages']).result()\n # 将产物插入消息列表\n messages = add_messages(state[\"messages\"], [joke])\n return {\"messages\": messages}\n\njoke_agent.name = \"joke_agent\"\n\n##############################################################################\n# Agent 2: Research Expert (Graph API)\n##############################################################################\n\ndef web_search(query: str) -> str:\n \"\"\"Search the web for information. (Mocked data here)\"\"\"\n return (\n \"Here are the headcounts for each of the FAANG companies in 2024:\\n\"\n \"1. **Facebook (Meta)**: 67,317 employees.\\n\"\n \"2. **Apple**: 164,000 employees.\\n\"\n \"3. **Amazon**: 1,551,000 employees.\\n\"\n \"4. **Netflix**: 14,000 employees.\\n\"\n \"5. **Google (Alphabet)**: 181,269 employees.\"\n )\n\n# research_agent = create_react_agent(\n# model=model,\n# tools=[web_search],\n# name=\"research_expert\",\n# # Prompt 告诉它是一个研究型 Agent，可调用 web_search\n# prompt=(\n# \"You are a world-class researcher. You have access to a 'web_search(query: str)' tool. \"\n# \"Do not do any complicated math, just provide factual info from the web_search if needed.\"\n# ),\n# )\nresearch_agent = ReactAgent(\n model=model,\n tools=[web_search],\n name=\"research_expert\",\n # Prompt 告诉它是一个研究型 Agent，可调用 web_search\n prompt=(\n \"You are a world-class researcher. You have access to a 'web_search(query: str)' tool. \"\n \"Do not do any complicated math, just provide factual info from the web_search if needed.\"\n ),\n)\n\n##############################################################################\n# 使用 SupervisorAgent 类替代直接调用 create_supervisor 函数\n##############################################################################\n\n# 创建 SupervisorAgent 实例\nsupervisor = SupervisorAgent(\n agents=[research_agent],\n model=model,\n # prompt=(\n # \"You are the overall supervisor. You manage two specialized agents:\\n\"\n # \"1) joke_agent: for telling jokes.\\n\"\n # \"2) research_expert: for factual or data-related questions.\\n\\n\"\n # \"If the user wants a joke AND some research data in the same query, \"\n # \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n # \"After both calls, provide a final combined response. \"\n # \"Do not call more than one agent in a single LLM message; do it step by step.\"\n # ),\n)\n##############################################################################\n# 测试：单个用户请求想要 \"先讲笑话，再查Apple的2024年人数\" 并合并结果\n##############################################################################\nresult = supervisor.invoke({\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": (\n \"Hi! I'd like to start with a short joke to lighten the mood, \"\n \"then please check Apple's headcount in 2024. Summarize both.\"\n )\n }\n ]\n})\n\n##############################################################################\n# 打印最终对话消息\n##############################################################################\nfor m in result[\"messages\"]:\n m.pretty_print()"
},
{
"path": "examples/03_tavily_tools_test.py",
"content": "import os\nfrom langgraph.prebuilt import create_react_agent\nfrom core.agents.react_supervisor_agent import SupervisorAgent\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.func import entrypoint, task\nfrom langgraph.graph import add_messages\nfrom langchain_community.tools import TavilySearchResults\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\n# 1. 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# Agent 1: Joke Generator (Functional API)\n##############################################################################\n\n@task\ndef generate_joke(messages):\n \"\"\"Generate a short joke (no tool calls).\"\"\"\n system_message = {\n \"role\": \"system\", \n \"content\": \"You are a witty comedian. Write a short joke.\"\n }\n # 直接调用 model.invoke，拼接 system_message + 用户消息\n msg = model.invoke([system_message] + messages)\n return msg\n\n@entrypoint()\ndef joke_agent(state):\n # 调用上面的函数型任务\n joke = generate_joke(state['messages']).result()\n # 将产物插入消息列表\n messages = add_messages(state[\"messages\"], [joke])\n return {\"messages\": messages}\n\njoke_agent.name = \"joke_agent\"\n\n##############################################################################\n# Agent 2: Research Expert with Tavily Search (Graph API)\n##############################################################################\n\n# 创建Tavily搜索工具\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n)\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[tavily_search],\n name=\"research_expert\",\n # Prompt 告诉它是一个研究型 Agent，可调用 tavily_search\n prompt=(\n \"You are a world-class researcher. You have access to the 'tavily_search_results_json' tool \"\n \"which can search the web for real-time information. \"\n \"When asked a question, use this tool to find accurate and up-to-date information. \"\n \"Summarize the search results in a clear and concise manner. \"\n \"Always cite your sources by including the URLs from the search results.\"\n ),\n)\n\n##############################################################################\n# 使用 SupervisorAgent 类来协调多个智能体\n##############################################################################\n\n# 创建 SupervisorAgent 实例\nsupervisor = SupervisorAgent(\n agents=[research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions using real-time web search.\\n\\n\"\n \"If the user wants a joke, call joke_agent.\\n\"\n \"If the user wants factual information or research data, call research_expert.\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n)\n\n# 编译得到一个可调用的\"App\"\napp = supervisor.compile()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = app.get_graph().draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"Image saved as {graph_path}\")\n\n# 使用示例\nif __name__ == \"__main__\":\n # 示例1：只询问笑话\n result1 = app.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"讲个笑话\"}]})\n print(\"\\n示例1 - 只询问笑话:\")\n for message in result1[\"messages\"]:\n message.pretty_print()\n \n # 示例2：只询问研究数据\n result2 = app.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"谁是现任美国总统？\"}]})\n print(\"\\n示例2 - 只询问研究数据:\")\n for message in result2[\"messages\"]:\n message.pretty_print()\n \n # 示例3：同时询问笑话和研究数据\n result3 = app.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"讲个关于人工智能的笑话，然后告诉我什么是大型语言模型\"}]})\n print(\"\\n示例3 - 同时询问笑话和研究数据:\")\n for message in result3[\"messages\"]:\n message.pretty_print()"
},
{
"path": "examples/04_react_agent_test.py",
"content": "import os\nimport json\nfrom langgraph.prebuilt import create_react_agent\nfrom langchain_openai import ChatOpenAI\nfrom langchain_community.tools import TavilySearchResults\nfrom typing import Dict, Any\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom dotenv import load_dotenv\nfrom core.utils.agent_utils import log_agent_actions, save_agent_graph\nload_dotenv() # 自动加载 .env 文件\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# 创建Tavily搜索工具 - 配置为深度搜索模式\n##############################################################################\n\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=True, # 包含原始内容，便于分析\n include_images=False,\n search_depth=\"advanced\" # 使用高级搜索深度\n)\n\n##############################################################################\n# 创建REACT Agent - 使用更详细的提示词引导多步思考\n##############################################################################\n\nreact_agent = create_react_agent(\n model=model,\n tools=[tavily_search],\n name=\"tesla_research_expert\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"你是一位专业的研究分析师，擅长分析复杂问题并提供深入见解。\\n\"\n \"你有一个强大的工具'tavily_search_results_json'可以搜索网络获取实时信息。\\n\\n\"\n \"当面对复杂问题时，请遵循以下REACT方法论：\\n\"\n \"1. 分解问题：将复杂问题分解为更小的子问题\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及搜索的顺序\\n\"\n \"3. 执行搜索：使用tavily_search_results_json工具执行搜索\\n\"\n \"4. 分析结果：分析搜索结果，确定是否需要进一步搜索\\n\"\n \"5. 综合信息：将所有搜索结果综合成一个连贯的回答\\n\\n\"\n \"重要提示：\\n\"\n \"- 不要一次性搜索过于宽泛的问题\\n\"\n \"- 对于复杂问题，进行多次有针对性的搜索\\n\"\n \"- 每次搜索后评估结果，决定下一步行动\\n\"\n \"- 在最终回答中引用来源，包括搜索结果中的URL\\n\"\n \"- 清晰地展示你的思考过程，包括问题分解和计划制定\\n\"\n ),\n)\n\n# 保存Agent图表\n# save_agent_graph(react_agent)\n\n##############################################################################\n# 测试：查询\"特斯拉2025年的发展预期\"\n##############################################################################\n\nif __name__ == \"__main__\":\n # 复杂查询测试\n print(\"\\n开始测试REACT Agent处理复杂查询的能力...\\n\")\n print(\"查询: 特斯拉2025年的发展预期\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n {\"role\": \"user\", \"content\": \"分析特斯拉2025年的发展预期，包括新车型计划、销量目标、技术创新和市场扩张战略。\"}\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in react_agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用原有的log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print(\"\\n最终回答:\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n message.pretty_print()"
},
{
"path": "examples/05_react_agent_user_input.py",
"content": "import asyncio\nimport os\nfrom typing import Dict, Any\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom langchain_community.tools import TavilySearchResults\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"=\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n # 只打印结果的前200个字符，避免输出过长\n content = latest_message.content\n if len(content) > 200:\n content = content[:200] + \"... (更多内容省略)\"\n print(f\"- 结果: {content}\")\n \n print(\"=\" * 50)\n\n##############################################################################\n# 创建Tavily搜索工具 - 配置为深度搜索模式\n##############################################################################\n\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=True, # 包含原始内容，便于分析\n include_images=False,\n search_depth=\"advanced\" # 使用高级搜索深度\n)\n\n##############################################################################\n# 创建ReactAgent实例\n##############################################################################\n\ndef create_react_agent_instance():\n \"\"\"创建并返回ReactAgent实例\"\"\"\n react_agent = ReactAgent(\n model=model,\n tools=[tavily_search],\n name=\"research_assistant\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"你是一位专业的研究分析师，擅长分析复杂问题并提供深入见解。\\n\"\n \"你有一个强大的工具'tavily_search_results_json'可以搜索网络获取实时信息。\\n\\n\"\n \"当面对复杂问题时，请遵循以下REACT方法论：\\n\"\n \"1. 分解问题：将复杂问题分解为更小的子问题\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及搜索的顺序\\n\"\n \"3. 执行搜索：使用tavily_search_results_json工具执行搜索\\n\"\n \"4. 分析结果：分析搜索结果，确定是否需要进一步搜索\\n\"\n \"5. 综合信息：将所有搜索结果综合成一个连贯的回答\\n\\n\"\n \"重要提示：\\n\"\n \"- 不要一次性搜索过于宽泛的问题\\n\"\n \"- 对于复杂问题，进行多次有针对性的搜索\\n\"\n \"- 每次搜索后评估结果，决定下一步行动\\n\"\n \"- 在最终回答中引用来源，包括搜索结果中的URL\\n\"\n \"- 清晰地展示你的思考过程，包括问题分解和计划制定\\n\"\n ),\n )\n \n # 获取图对象并保存\n agent = react_agent.compile() \n return agent\n\n##############################################################################\n# 主函数 - 处理用户输入\n##############################################################################\n\nasync def main():\n # 创建ReactAgent实例\n react_agent = create_react_agent_instance()\n \n while True:\n # 获取用户输入\n user_input = await asyncio.to_thread(input, \"\\n请输入您想了解的问题 (输入'退出'结束): \")\n \n # 检查是否退出\n if user_input.lower() in ['退出', 'exit', 'quit']:\n print(\"\\n感谢使用，再见！\")\n break\n \n # 准备初始状态\n initial_state = {\n \"messages\": [HumanMessage(content=user_input)]\n }\n \n try:\n print(\"\\n=== 🔍 开始研究 ===\\n\")\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in react_agent.stream(initial_state, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print(\"\\n最终回答:\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(\"\\n\" + \"=\" * 80)\n print(message.content)\n print(\"=\" * 80 + \"\\n\")\n \n except Exception as e:\n print(f\"\\n处理查询时出错: {e}\")\n\n##############################################################################\n# 程序入口\n##############################################################################\n\nif __name__ == \"__main__\":\n print(\"\\n欢迎使用ReactAgent研究助手！\")\n print(\"这个助手可以帮助您研究各种问题，使用Tavily搜索工具获取最新信息。\")\n print(\"您可以输入任何问题，助手将使用REACT方法论进行分析和回答。\")\n \n # 运行主函数\n asyncio.run(main())"
},
{
"path": "examples/06_web_extraction_tools_test.py",
"content": "import os\nimport sys\nfrom langgraph.prebuilt import create_react_agent\nfrom langchain_openai import ChatOpenAI\nimport json\nfrom typing import Dict, Any\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom dotenv import load_dotenv\nfrom langchain_community.tools import JinaSearch\nfrom core.tools.firecrawl_tool import FireCrawlTool\n\n\nload_dotenv() # 自动加载 .env 文件\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"=\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n # 只打印结果的前200个字符，避免输出过长\n content = latest_message.content\n if len(content) > 300:\n content = content[:300] + \"... (更多内容省略)\"\n print(f\"- 结果: {content}\")\n \n print(\"=\" * 50)\n\n##############################################################################\n# 创建Web提取工具 - FireCrawl用于网站结构，Jina用于内容提取\n##############################################################################\n\n# 创建FireCrawl工具 - 用于网站结构分析\nfirecrawl_tool = FireCrawlTool(\n mode=\"crawl\", # 使用爬取模式\n params={\"max_pages\": 10} # 限制爬取页面数量\n)\n\n# 创建Jina Reader工具 - 用于内容提取\njina_reader_tool = JinaSearch()\n\n##############################################################################\n# 创建REACT Agent - 使用更详细的提示词引导多步思考\n##############################################################################\n\nreact_agent = create_react_agent(\n model=model,\n tools=[firecrawl_tool, jina_reader_tool],\n name=\"web_extraction_expert\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"你是一位专业的网页内容分析专家，擅长提取和分析网站结构与内容。\\n\"\n \"你有两个强大的工具:\\n\"\n \"1. 'firecrawl_tool': 用于爬取网站结构和下级页面\\n\"\n \"2. 'jina_reader_tool': 用于从特定URL提取结构化内容，获取干净可读的内容\\n\\n\"\n \"当面对网站分析任务时，请遵循以下方法论:\\n\"\n \"1. 分析任务: 明确需要从网站获取什么信息\\n\"\n \"2. 网站结构分析: 使用firecrawl_tool爬取网站结构，了解可用页面\\n\"\n \"3. 内容提取: 根据网站结构，使用jina_reader_tool从关键页面提取内容\\n\"\n \"4. 信息整合: 将提取的内容整合成有条理的分析结果\\n\\n\"\n \"重要提示:\\n\"\n \"- 先使用firecrawl_tool了解网站结构，再使用jina_reader_tool提取具体内容\\n\"\n \"- 对于大型网站，先分析网站结构，再有针对性地选择重要页面进行内容提取\\n\"\n \"- 每次工具使用后评估结果，决定下一步行动\\n\"\n \"- 在最终回答中提供结构化的分析，包括网站组织方式和关键内容摘要\\n\"\n \"- 清晰地展示你的思考过程，包括为什么选择特定页面进行分析\\n\"\n ),\n)\n\n##############################################################################\n# 测试：分析LangGraph文档网站\n##############################################################################\n\nif __name__ == \"__main__\":\n # 测试网站分析\n print(\"\\n开始测试Web提取Agent分析网站的能力...\\n\")\n print(\"分析目标: LangGraph文档网站\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n {\"role\": \"user\", \"content\": \"爬取LangGraph文档网站的每个章节的内容(https://langchain-ai.github.io/langgraph/how-tos/) \"}\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in react_agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用原有的log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print(\"\\n最终分析结果:\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n message.pretty_print()"
},
{
"path": "examples/07_web_extraction_with_filesystem.py",
"content": "import os\nimport sys\nimport json\nimport asyncio\nfrom datetime import datetime\nfrom typing import Dict, Any, List\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom langchain_community.agent_toolkits import FileManagementToolkit\nfrom langgraph.prebuilt import create_react_agent\nfrom langgraph.checkpoint.memory import MemorySaver\nfrom dotenv import load_dotenv\nfrom langchain_community.tools import TavilySearchResults\nfrom core.agents.react_supervisor_agent import SupervisorAgent\n\nload_dotenv() # 自动加载 .env 文件\n\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"=\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n # 限制内容长度，避免过长输出\n content = latest_message.content\n if len(content) > 500:\n content = content[:250] + \"\\n... (内容过长，已截断) ...\\n\" + content[-250:]\n print(content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n # 限制参数输出长度\n args = str(tool_call['args'])\n if len(args) > 100:\n args = args[:100] + \"... (参数过长，已截断)\"\n print(f\"- 参数: {args}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n # 只打印结果的前200个字符，避免输出过长\n content = latest_message.content\n if len(content) > 200:\n content = content[:100] + \"\\n... (更多内容省略) ...\\n\" + content[-100:]\n print(f\"- 结果: {content}\")\n \n print(\"=\" * 50)\n\n##############################################################################\n# 创建Web提取工具\n##############################################################################\n# 创建Tavily搜索工具\n\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n)\n\n##############################################################################\n# 创建文件系统工具 - 用于保存提取的内容\n##############################################################################\n\n# 设置文件系统工具的根目录为examples/output\noutput_dir = os.path.join(os.path.dirname(__file__), \"output\")\nos.makedirs(output_dir, exist_ok=True)\n\n# 创建文件系统工具集\nfilesystem_toolkit = FileManagementToolkit(\n root_dir=output_dir,\n selected_tools=[\"write_file\", \"read_file\", \"list_directory\"]\n)\n\n# 获取文件系统工具\nfilesystem_tools = filesystem_toolkit.get_tools()\n\n##############################################################################\n# 创建Research Agent - 用于网站内容提取\n##############################################################################\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[tavily_search],\n name=\"research_agent\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"You are a world-class researcher. You have access to the 'tavily_search_results_json' tool \"\n \"which can search the web for real-time information. \"\n \"When asked a question, use this tool to find accurate and up-to-date information. \"\n \"Summarize the search results in a clear and concise manner. \"\n \"Always cite your sources by including the URLs from the search results.\"\n ),\n debug=False)\n\n##############################################################################\n# 创建FileSystem Agent - 用于保存提取的内容\n##############################################################################\n\nfilesystem_agent = create_react_agent(\n model=model,\n tools=filesystem_tools,\n name=\"filesystem_agent\",\n # 提示词强调文件操作和内容保存\n prompt=(\n \"你是一位专业的文件系统管理专家，负责将网页内容保存到本地文件系统。\\n\"\n \"你有以下工具可以使用:\\n\"\n \"1. 'write_file': 用于将内容写入文件\\n\"\n \"2. 'read_file': 用于读取文件内容\\n\"\n \"3. 'list_directory': 用于列出目录内容\\n\\n\"\n \"当接收到保存内容的请求时，请遵循以下方法论:\\n\"\n \"1. 分析内容: 确定内容的类型和结构\\n\"\n \"2. 确定文件名: 根据内容类型和来源创建合适的文件名\\n\"\n \"3. 保存内容: 使用write_file工具将内容保存到文件中\\n\"\n \"4. 验证保存: 使用read_file或list_directory工具验证内容已正确保存\\n\\n\"\n \"重要提示:\\n\"\n \"- 为文件创建有意义的名称，包含日期和内容描述\\n\"\n \"- 对于结构化数据，优先使用JSON格式保存\\n\"\n \"- 对于文本内容，使用TXT或MD格式保存\\n\"\n \"- 确保文件名不包含非法字符\\n\"\n \"- 在保存前，检查是否已存在同名文件，避免覆盖重要内容\\n\"\n ),\n)\n\n##############################################################################\n# 创建Supervisor Agent - 协调Research Agent和FileSystem Agent\n##############################################################################\n# 创建内存存储器用于保存对话状态\nmemory_saver = MemorySaver()\nsupervisor = SupervisorAgent(\n agents=[research_agent, filesystem_agent],\n model=model,\n prompt=(\n \"你是一个智能助手的总协调者，负责管理两个专业智能体:\\n\"\n \"1) research_agent: 网页内容分析专家，可以爬取和分析网站内容\\n\"\n \"2) filesystem_agent: 文件系统管理专家，可以将内容保存到本地文件系统\\n\\n\"\n \"你的工作流程如下:\\n\"\n \"1. 分析用户请求，确定是需要网页内容提取还是文件操作，或两者都需要\\n\"\n \"2. 如果需要网页内容提取，调用research_agent获取网页内容\\n\"\n \"3. 如果需要将提取的内容保存到文件，调用filesystem_agent进行保存\\n\"\n \"4. 如果用户同时需要提取内容并保存，先调用research_agent获取内容，再调用filesystem_agent保存内容\\n\\n\"\n \"重要规则:\\n\"\n \"- 不要在一个消息中同时调用多个智能体，必须一步一步来\\n\"\n \"- 当调用filesystem_agent保存内容时，必须提供完整的内容和建议的文件名\\n\"\n \"- 确保在最终回复中告知用户内容已成功提取和/或保存\\n\"\n \"- 如果用户只想提取内容而不保存，只调用research_agent\\n\"\n \"- 如果用户只想操作文件而不提取新内容，只调用filesystem_agent\\n\\n\"\n \"上下文管理指南:\\n\"\n \"- 当处理大型网站或多个页面时，指导research_agent采用分批处理策略\\n\"\n \"- 对于大型内容提取任务，先让research_agent获取网站结构，再逐步处理各个页面\\n\"\n \"- 当发现research_agent返回的内容过大时，指导它进行内容摘要或分批处理\\n\"\n \"- 如果research_agent一次性尝试处理过多页面导致上下文超限，指导它减少并行处理的页面数量\\n\"\n \"- 对于需要保存的大型内容，考虑将其分割成多个小文件，而不是一个大文件\\n\"\n \"- 在处理多页面内容时，可以采用先保存再处理的策略，减轻上下文负担\\n\"\n ),\n checkpointer=memory_saver\n)\n\n\n\n# 编译得到一个可调用的\"App\"，添加checkpointer实现记忆功能\napp = supervisor.compile()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = app.get_graph().draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"图表已保存为 {graph_path}\")\n\n##############################################################################\n# 主函数 - 处理用户输入\n##############################################################################\n\nasync def main():\n # 创建一个固定的thread_id用于保持对话上下文\n thread_id = \"user_session_1\"\n \n # 创建配置对象，包含thread_id\n config = {\"configurable\": {\"thread_id\": thread_id}}\n \n print(\"\\n当前会话ID:\", thread_id)\n print(\"(所有对话将在同一会话中保持上下文记忆)\")\n \n while True:\n # 获取用户输入\n user_input = await asyncio.to_thread(input, \"\\n请输入您想了解的问题 (输入'退出'结束): \")\n \n # 检查是否退出\n if user_input.lower() in ['退出', 'exit', 'quit']:\n print(\"\\n感谢使用，再见！\")\n break\n \n # 准备初始状态 - 只包含当前用户消息\n initial_state = {\n \"messages\": [HumanMessage(content=user_input)]\n }\n \n try:\n print(\"\\n=== 🔍 开始研究 ===\\n\")\n \n # 使用stream方法逐步获取中间状态，传入config以使用相同的thread_id\n for partial_state in app.stream(initial_state, config, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n except Exception as e:\n print(f\"\\n处理查询时出错: {e}\")\n print(\"可能是由于上下文长度超出限制，请尝试减少查询范围或使用'批处理大小设置为X'命令调整批处理大小(1-5之间)\")\n\n##############################################################################\n# 程序入口\n##############################################################################\n\nif __name__ == \"__main__\":\n print(\"\\n欢迎使用具有记忆功能的网页爬取助手！\")\n print(\"本助手可以记住您之前的对话内容，实现连续对话体验。\")\n print(\"您可以询问之前提到过的内容，助手会根据上下文理解您的问题。\")\n \n # 运行主函数\n asyncio.run(main())"
},
{
"path": "examples/08_react_agent_tool_registry_test.py",
"content": "import os\nimport sys\nimport json\nfrom typing import Dict, Any, List\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom langchain_community.tools import JinaSearch, WikipediaQueryRun\nfrom langchain_community.utilities import WikipediaAPIWrapper\nfrom dotenv import load_dotenv\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools import register_direct_tool\nfrom core.tools.registry import get_registered_tools, ToolCategory\nfrom core.tools.firecrawl_tool import FireCrawlTool\n\nload_dotenv() # 自动加载 .env 文件\n\n##############################################################################\n# 工具注册和ReactAgent测试 - 美联储研究任务\n##############################################################################\n\ndef print_separator(title):\n \"\"\"打印分隔符\"\"\"\n print(\"\\n\" + \"=\" * 80)\n print(f\" {title} \".center(80, \"=\"))\n print(\"=\" * 80)\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"-\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n # 只打印结果的前200个字符，避免输出过长\n content = latest_message.content\n if len(content) > 200:\n content = content[:200] + \"... (更多内容省略)\"\n print(f\"- 结果: {content}\")\n \n print(\"-\" * 50)\n\n##############################################################################\n# 注册工具\n##############################################################################\n\nprint_separator(\"注册搜索工具\")\n\n# 创建JinaSearch工具实例\njina_search = JinaSearch()\n\n# 创建Wikipedia工具实例\n# wiki_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())\n\nfirecrawl_tool = FireCrawlTool()\n\n# 使用register_direct_tool函数注册工具\nregister_direct_tool(jina_search)\nregister_direct_tool(firecrawl_tool)\n# 注册自定义工具 - FireCrawlTool\n\n# 获取所有已注册的工具（以字典格式）\nregistered_tools = get_registered_tools(as_dict=True)\n\n# 打印所有已注册的工具\nprint(\"\\n已注册的工具:\")\nfor name, info in registered_tools.items():\n print(f\"- {name} (类别: {info['category'].value})\")\n\n##############################################################################\n# 创建ReactAgent实例\n##############################################################################\n\nprint_separator(\"创建ReactAgent实例\")\n\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n# 从注册表中只获取搜索类工具列表\nfrom core.tools.registry import get_tools_by_category, ToolCategory\ntools_list = get_tools_by_category(ToolCategory.SEARCH)\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"fed_research_agent\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"你是一位专业的经济研究分析师，擅长分析复杂的经济问题并提供深入见解。\\n\"\n \"你有多个强大的工具可以搜索网络获取实时信息：\\n\"\n \"当面对复杂问题时，请遵循以下方法论：\\n\"\n \"1. 分解问题：将复杂问题分解为更小的子问题\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及使用哪些工具\\n\"\n \"3. 执行搜索：使用适当的工具执行搜索\\n\"\n \"4. 分析结果：分析搜索结果，确定是否需要进一步搜索\\n\"\n \"5. 综合信息：将所有搜索结果综合成一个连贯的回答\\n\\n\"\n \"重要提示：\\n\"\n \"- 每次搜索后评估结果，决定下一步行动\\n\"\n \"- 在最终回答中引用来源\\n\"\n \"- 清晰地展示你的思考过程，包括问题分解和计划制定\\n\"\n ),\n)\n\n# agent = react_agent.compile()\n# 获取图对象\n# graph = agent.get_graph()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = graph.draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"工作流图已保存为 {graph_path}\")\n\n##############################################################################\n# 测试：查询\"美联储的详细介绍和它如何影响全球经济\"\n##############################################################################\n\nif __name__ == \"__main__\":\n print_separator(\"开始测试ReactAgent处理美联储研究任务\")\n print(\"\\n查询: 美联储的详细介绍和它如何影响全球经济\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请提供2025年美联储(Federal Reserve)的详细介绍，包括其历史、结构、职能，以及它如何通过货币政策影响全球经济。\")\n ]\n }\n result = react_agent.run(inputs)\n##############################################################################\n# 打印最终对话消息\n##############################################################################\n for m in result[\"messages\"]:\n m.pretty_print()"
},
{
"path": "examples/09_e2b_code_interpreter_test.py",
"content": "import os\nimport sys\nimport json\nfrom typing import Dict, Any, List\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom dotenv import load_dotenv\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_registered_tools, ToolCategory, get_tools_by_category\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\n\nload_dotenv() # 自动加载 .env 文件\n\n##############################################################################\n# E2B代码解释器工具测试\n##############################################################################\n\ndef print_separator(title):\n \"\"\"打印分隔符\"\"\"\n print(\"\\n\" + \"=\" * 80)\n print(f\" {title} \".center(80, \"=\"))\n print(\"=\" * 80)\n\n##############################################################################\n# 检查E2B代码解释器工具是否已注册\n##############################################################################\n\nprint_separator(\"检查E2B代码解释器工具是否已注册\")\n\n# 获取所有已注册的工具（以字典格式）\nregistered_tools = get_registered_tools(as_dict=True)\n\n# 打印所有已注册的工具\nprint(\"\\n已注册的工具:\")\nfor name, info in registered_tools.items():\n print(f\"- {name} (类别: {info['category'].value})\")\n\n# 检查E2B代码解释器工具是否已注册\ne2b_tool_name = \"e2b_code_interpreter\"\nif e2b_tool_name in registered_tools:\n print(f\"\\nE2B代码解释器工具已成功注册: {e2b_tool_name}\")\nelse:\n print(f\"\\n警告: E2B代码解释器工具未注册\")\n # 手动注册E2B代码解释器工具\n print(\"尝试手动注册E2B代码解释器工具...\")\n try:\n from core.tools.registry import register_tool\n e2b_tool = E2BCodeInterpreterTool()\n register_tool(e2b_tool, ToolCategory.CODE_INTERPRETER)\n print(f\"已手动注册工具: {e2b_tool.name}\")\n except Exception as e:\n print(f\"手动注册E2B代码解释器工具失败: {e}\")\n\n##############################################################################\n# 创建ReactAgent实例\n##############################################################################\n\nprint_separator(\"创建ReactAgent实例\")\n\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n# 从注册表中只获取代码解释器类工具列表\ntools_list = get_tools_by_category(ToolCategory.CODE_INTERPRETER)\n\n# 打印获取到的代码解释器工具\nprint(\"\\n获取到的代码解释器工具:\")\nfor tool in tools_list:\n print(f\"- {tool.name}: {tool.description}\")\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"code_interpreter_agent\",\n # 提示词强调使用代码解释器工具进行数据分析和可视化\n prompt=(\n \"你是一位专业的数据分析师和编程助手，擅长使用Python进行数据分析和可视化。\\n\"\n \"你有多个强大的代码执行工具可以使用：\\n\"\n \"- e2b_code_interpreter: 用于执行Python代码，支持数据分析和可视化\\n\"\n \"当面对编程和数据分析问题时，请遵循以下方法论：\\n\"\n \"1. 分析问题：理解用户的需求和问题本质\\n\"\n \"2. 制定计划：确定解决方案和需要使用的工具\\n\"\n \"3. 编写代码：使用适当的工具编写和执行代码\\n\"\n \"4. 分析结果：解释代码执行结果，提供见解\\n\"\n \"5. 优化方案：如有必要，优化代码或提供改进建议\\n\\n\"\n \"重要提示：\\n\"\n \"- 优先使用e2b_code_interpreter工具执行Python代码\\n\"\n \"- 对于数据分析和可视化任务，确保导入必要的库（如pandas, matplotlib, numpy等）\\n\"\n \"- 对于不存在的库，工具会自动尝试使用pip install进行安装\\n\"\n \"- 在代码中添加详细注释，解释关键步骤\\n\"\n \"- 执行代码后，解释结果含义和见解\\n\"\n ),\n)\n\n# 编译Agent\nagent = react_agent.compile()\n\n# # 获取图对象\n# graph = agent.get_graph()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = graph.draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"工作流图已保存为 {graph_path}\")\n\n##############################################################################\n# 测试：使用E2B代码解释器执行简单的数据分析任务\n##############################################################################\n\nif __name__ == \"__main__\":\n print_separator(\"开始测试ReactAgent使用E2B代码解释器\")\n print(\"\\n查询: 使用Python生成一个简单的正弦波图形\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"使用Python生成一个简单的正弦波图形，如果有找不到的模块，需要自动安装\")\n ]\n }\n result = agent.invoke(inputs)\n\n for m in result[\"messages\"]:\n m.pretty_print()"
},
{
"path": "examples/10_financial_data_analysis.py",
"content": "import os\nimport sys\nimport json\nfrom typing import Dict, Any, List\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom dotenv import load_dotenv\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_registered_tools, ToolCategory, get_tools_by_category\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\n\nload_dotenv() # 自动加载 .env 文件\n\n##############################################################################\n# 财务数据分析报表生成示例\n##############################################################################\n\ndef print_separator(title):\n \"\"\"打印分隔符\"\"\"\n print(\"\\n\" + \"=\" * 80)\n print(f\" {title} \".center(80, \"=\"))\n print(\"=\" * 80)\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"-\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n content = latest_message.content\n print(f\"- 结果: {content}\")\n \n print(\"-\" * 50)\n\n##############################################################################\n# 检查E2B代码解释器工具是否已注册\n##############################################################################\n\nprint_separator(\"检查E2B代码解释器工具是否已注册\")\n\n# 获取所有已注册的工具（以字典格式）\nregistered_tools = get_registered_tools(as_dict=True)\n\n# 打印所有已注册的工具\nprint(\"\\n已注册的工具:\")\nfor name, info in registered_tools.items():\n print(f\"- {name} (类别: {info['category'].value})\")\n\n# 检查E2B代码解释器工具是否已注册\ne2b_tool_name = \"e2b_code_interpreter\"\nif e2b_tool_name in registered_tools:\n print(f\"\\nE2B代码解释器工具已成功注册: {e2b_tool_name}\")\nelse:\n print(f\"\\n警告: E2B代码解释器工具未注册\")\n # 手动注册E2B代码解释器工具\n print(\"尝试手动注册E2B代码解释器工具...\")\n try:\n from core.tools.registry import register_tool\n e2b_tool = E2BCodeInterpreterTool()\n register_tool(e2b_tool, ToolCategory.CODE_INTERPRETER)\n print(f\"已手动注册工具: {e2b_tool.name}\")\n except Exception as e:\n print(f\"手动注册E2B代码解释器工具失败: {e}\")\n\n##############################################################################\n# 创建ReactAgent实例\n##############################################################################\n\nprint_separator(\"创建ReactAgent实例\")\n\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n# 从注册表中只获取代码解释器类工具列表\ntools_list = get_tools_by_category(ToolCategory.CODE_INTERPRETER)\n\n# 打印获取到的代码解释器工具\nprint(\"\\n获取到的代码解释器工具:\")\nfor tool in tools_list:\n print(f\"- {tool.name}: {tool.description}\")\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"financial_data_analyst\",\n # 提示词强调使用代码解释器工具进行财务数据分析和可视化\n prompt=(\n \"你是一位专业的财务数据分析师，擅长使用Python进行财务数据分析和可视化。\\n\"\n \"你有强大的代码执行工具可以使用：\\n\"\n \"- e2b_code_interpreter: 用于执行Python代码，支持数据分析和可视化\\n\\n\"\n \"当面对财务数据分析问题时，请遵循以下方法论：\\n\"\n \"1. 分析问题：理解用户的需求和问题本质\\n\"\n \"2. 制定计划：确定解决方案和需要使用的工具\\n\"\n \"3. 编写代码：使用适当的工具编写和执行代码\\n\"\n \"4. 分析结果：解释代码执行结果，提供财务见解\\n\"\n \"5. 优化方案：如有必要，优化代码或提供改进建议\\n\\n\"\n \"重要提示：\\n\"\n \"- 优先使用e2b_code_interpreter工具执行Python代码\\n\"\n \"- 对于财务数据分析和可视化任务，确保导入必要的库（如pandas, matplotlib, numpy等）\\n\"\n \"- 对于不存在的库，工具会自动尝试使用pip install进行安装\\n\"\n \"- 在代码中添加详细注释，解释关键步骤\\n\"\n \"- 执行代码后，解释结果含义和财务见解\\n\"\n ),\n)\n\n# # 编译Agent\n# agent = react_agent.compile()\n\n# # 获取图对象\n# graph = agent.get_graph()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = graph.draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"工作流图已保存为 {graph_path}\")\n\n##############################################################################\n# 从沙箱下载文件到本地的函数\n##############################################################################\nimport os\n\ndef download_file_from_sandbox(sandbox, sandbox_path, local_path):\n \"\"\"从 e2b 沙箱中下载文件并保存到本地，自动区分文本和二进制文件\"\"\"\n try:\n print(f\"读取文件: {sandbox_path}\")\n\n # 判断是否为常见二进制文件类型（可自行扩展）\n binary_extensions = (\n '.png', '.jpg', '.jpeg', '.gif', '.pdf', '.svg',\n '.xlsx', '.xls', '.zip', '.bin', '.pyc', '.pyd',\n '.pptx', '.docx', '.mp3', '.mp4', '.avi', '.mov',\n )\n is_binary = sandbox_path.lower().endswith(binary_extensions)\n\n # 创建目录\n os.makedirs(os.path.dirname(local_path), exist_ok=True)\n\n if is_binary:\n print(\"📦 识别为二进制文件，使用 sandbox.download_file()\")\n content = sandbox.files.read(sandbox_path) # 返回 bytes\n with open(local_path, 'wb') as f:\n f.write(content)\n else:\n print(\"📄 识别为文本文件，使用 sandbox.files.read()\")\n content = sandbox.files.read(sandbox_path) # 返回 str\n with open(local_path, 'w', encoding='utf-8') as f:\n f.write(content)\n\n print(f\"✅ 文件已保存到本地: {local_path}\")\n return True\n\n except Exception as e:\n print(f\"❌ 下载失败: {e}\")\n return False\n \ndef download_directory_from_sandbox(sandbox, sandbox_dir_path, local_dir_path):\n \"\"\"从沙箱下载整个目录内容到本地\n \n Args:\n sandbox: 沙箱实例\n sandbox_dir_path: 沙箱中的目录路径\n local_dir_path: 本地保存目录路径\n \n Returns:\n bool: 是否成功下载所有文件\n \"\"\"\n try:\n print(f\"尝试下载目录: {sandbox_dir_path} -> {local_dir_path}\")\n \n # 确保本地目录存在\n os.makedirs(local_dir_path, exist_ok=True)\n \n # 列出沙箱中指定目录下的所有文件\n try:\n files = sandbox.files.list(sandbox_dir_path)\n # print(f\"获取到文件列表: {sandbox_dir_path}, 类型: {type(files)}\")\n # if files and len(files) > 0:\n # print(f\"第一个文件类型: {type(files[0])}, 内容: {files[0]}\")\n # # 检查对象属性\n # print(f\"文件对象可用属性: {dir(files[0])}\")\n except Exception as e:\n print(f\"列出文件时出错: {sandbox_dir_path}, 错误: {str(e)}\")\n return False\n \n if not files:\n print(f\"沙箱中目录 {sandbox_dir_path} 为空或不存在\")\n return False\n \n downloaded_count = 0\n # 定义需要跳过的系统文件\n skip_files = {'.bashrc', '.bash_logout', '.profile'}\n \n # 遍历并下载每个文件\n for file_info in files:\n try:\n # 使用dir()查看对象有哪些属性\n print(f\"文件信息对象属性: {dir(file_info)}\")\n \n # 尝试安全获取name和type属性\n file_name = getattr(file_info, \"name\", None)\n if file_name is None:\n print(f\"警告: 无法获取文件名, 跳过此文件\")\n continue\n \n file_type = getattr(file_info, \"type\", \"file\") # 默认为文件类型\n # 如果 file_type 是枚举, 使用其 value 进行判断\n type_value = file_type.value if hasattr(file_type, \"value\") else file_type\n \n # 跳过不需要的系统文件或系统目录（隐藏文件/目录）\n if file_name in skip_files or (file_name.startswith('.') and type_value == 'dir'):\n print(f\"跳过系统文件或目录: {file_name}\")\n continue\n \n print(f\"处理文件: {file_name}, 类型: {type_value}\")\n \n sandbox_file_path = f\"{sandbox_dir_path}/{file_name}\"\n local_file_path = os.path.join(local_dir_path, file_name)\n \n if type_value == 'dir':\n # 递归下载子目录\n print(f\"发现子目录: {sandbox_file_path}\")\n if download_directory_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n else:\n # 下载文件\n print(f\"下载文件: {sandbox_file_path} -> {local_file_path}\")\n if download_file_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n except Exception as e:\n print(f\"处理文件时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n continue\n \n if downloaded_count > 0:\n print(f\"从 {sandbox_dir_path} 下载了 {downloaded_count} 个文件/目录到 {local_dir_path}\")\n return True\n return False\n \n except Exception as e:\n print(f\"从沙箱下载目录时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n return False\n\n##############################################################################\n# 测试：使用E2B代码解释器生成财务数据分析报表\n##############################################################################\n\nif __name__ == \"__main__\":\n print_separator(\"开始测试ReactAgent使用E2B代码解释器进行财务数据分析\")\n print(\"\\n查询: 生成模拟财务数据并进行分析，生成财务报表\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请生成一组模拟的公司财务数据（包括收入、支出、利润等），对数据进行分析，将处理过程（代码）和最终生成的结果保存到本地。\")\n ]\n }\n result = react_agent.run(inputs)\n\n for m in result[\"messages\"]:\n m.pretty_print()\n\n print(\"\\n下载沙盒里的文件\")\n try:\n # 遍历 react_agent.tools 以查找 E2B 相关工具\n sandbox = None\n for tool in react_agent.tools:\n if hasattr(tool, \"sandbox\"):\n sandbox = tool.sandbox\n break # 找到后就退出循环\n\n if sandbox:\n # 设定输出目录\n output_dir = os.path.join(os.getcwd(), \"examples/output/sandbox_files\")\n os.makedirs(output_dir, exist_ok=True)\n\n # 直接下载主要工作目录\n print(\"\\n从沙箱下载文件到本地...\")\n download_directory_from_sandbox(sandbox, \"/home/user\", output_dir)\n\n # 下载临时目录中可能的图表和数据文件\n # download_directory_from_sandbox(sandbox, \"/tmp\", output_dir)\n\n print(f\"\\n文件已保存到目录: {output_dir}\")\n sandbox.close()\n except Exception as e:\n print(f\"从沙箱下载文件时出错: {str(e)}\")"
},
{
"path": "examples/11_e2b_sandbox_test.py",
"content": "import os\nimport sys\nimport json\nfrom typing import Dict, Any, List\nfrom datetime import datetime\n\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom dotenv import load_dotenv\n\nfrom core.agents.base.react_agent import ReactAgent\nfrom core.tools.registry import get_registered_tools, ToolCategory, get_tools_by_category\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\n\nload_dotenv() # 自动加载 .env 文件\n\n##############################################################################\n# E2B沙盒环境测试程序\n##############################################################################\n\ndef print_separator(title):\n \"\"\"打印分隔符\"\"\"\n print(\"\\n\" + \"=\" * 80)\n print(f\" {title} \".center(80, \"=\"))\n print(\"=\" * 80)\n\n##############################################################################\n# 创建一个记录Agent思考过程的函数\n##############################################################################\n\ndef log_agent_actions(state: Dict[str, Any]) -> None:\n \"\"\"记录Agent的思考过程和行动\"\"\"\n print(\"\\n\" + \"-\" * 50)\n print(\"当前状态:\")\n \n # 打印最新消息\n if state.get(\"messages\") and len(state[\"messages\"]) > 0:\n latest_message = state[\"messages\"][-1]\n \n if isinstance(latest_message, AIMessage):\n print(f\"\\nAI思考过程:\")\n print(latest_message.content)\n \n # 如果有工具调用，打印工具调用信息\n if latest_message.tool_calls:\n print(f\"\\n工具调用:\")\n for tool_call in latest_message.tool_calls:\n print(f\"- 工具: {tool_call['name']}\")\n print(f\"- 参数: {tool_call['args']}\")\n \n elif isinstance(latest_message, ToolMessage):\n print(f\"\\n工具返回结果:\")\n print(f\"- 工具: {latest_message.name}\")\n content = latest_message.content\n if len(content) > 500:\n content = content[:250] + \"\\n... (内容过长，已截断) ...\\n\" + content[-250:]\n print(f\"- 结果: {content}\")\n \n print(\"-\" * 50)\n\n##############################################################################\n# 从沙箱下载文件到本地的函数\n##############################################################################\n\ndef download_file_from_sandbox(sandbox, sandbox_path, local_path):\n \"\"\"从 e2b 沙箱中下载文件并保存到本地，自动区分文本和二进制文件\"\"\"\n try:\n print(f\"读取文件: {sandbox_path}\")\n\n # 判断是否为常见二进制文件类型（可自行扩展）\n binary_extensions = (\n '.png', '.jpg', '.jpeg', '.gif', '.pdf', '.svg',\n '.xlsx', '.xls', '.zip', '.bin', '.pyc', '.pyd',\n '.pptx', '.docx', '.mp3', '.mp4', '.avi', '.mov',\n )\n is_binary = sandbox_path.lower().endswith(binary_extensions)\n\n # 创建目录\n os.makedirs(os.path.dirname(local_path), exist_ok=True)\n\n if is_binary:\n print(\"📦 识别为二进制文件，使用 sandbox.download_file()\")\n content = sandbox.files.read(sandbox_path) # 返回 bytes\n with open(local_path, 'wb') as f:\n f.write(content)\n else:\n print(\"📄 识别为文本文件，使用 sandbox.files.read()\")\n content = sandbox.files.read(sandbox_path) # 返回 str\n with open(local_path, 'w', encoding='utf-8') as f:\n f.write(content)\n\n print(f\"✅ 文件已保存到本地: {local_path}\")\n return True\n\n except Exception as e:\n print(f\"❌ 下载失败: {e}\")\n return False\n\ndef run_ai_generated_code(sandbox, code: str, save_results_dir=None):\n \"\"\"在 E2B 沙箱中执行 AI 生成的代码\n \n Args:\n sandbox: 沙箱实例\n code: AI 生成的代码字符串\n save_results_dir: 用于保存结果文件的本地目录路径（可选）\n \n Returns:\n dict: 包含执行结果的字典\n \"\"\"\n try:\n print(\"在沙箱中执行 AI 生成的代码...\")\n # 确保代码是字符串类型\n if not isinstance(code, str):\n code = str(code)\n \n # 执行代码\n execution = sandbox.run_code(code)\n print(\"代码执行完成!\")\n \n # 准备结果字典\n result = {\n \"success\": True,\n \"stdout\": \"\",\n \"results\": []\n }\n \n # 提取标准输出\n if hasattr(execution, \"stdout\"):\n result[\"stdout\"] = execution.stdout\n \n # 检查代码是否执行成功\n if hasattr(execution, \"error\") and execution.error:\n error_name = getattr(execution.error, \"name\", \"Unknown\")\n error_value = getattr(execution.error, \"value\", \"Unknown error\")\n error_traceback = getattr(execution.error, \"traceback\", \"\")\n \n print(\"AI 生成的代码执行出错:\")\n print(f\"错误类型: {error_name}\")\n print(f\"错误信息: {error_value}\")\n if error_traceback:\n print(f\"错误追踪: {error_traceback}\")\n \n result[\"success\"] = False\n result[\"error\"] = {\n \"name\": error_name,\n \"value\": error_value,\n \"traceback\": error_traceback\n }\n return result\n \n # 处理执行结果\n if hasattr(execution, \"results\") and execution.results:\n import base64\n result_idx = 0\n \n for res in execution.results:\n # 默认为文本结果\n result_data = {\"type\": \"text\", \"value\": str(res)}\n \n # 检查是否有PNG图像\n if hasattr(res, \"png\") and res.png:\n result_data[\"type\"] = \"png\"\n result_data[\"value\"] = res.png # base64编码的字符串\n \n # 如果指定了保存目录，保存图像到本地\n if save_results_dir:\n try:\n os.makedirs(save_results_dir, exist_ok=True)\n image_path = os.path.join(save_results_dir, f\"result-{result_idx}.png\")\n \n # 解码并保存图像\n with open(image_path, 'wb') as f:\n f.write(base64.b64decode(res.png))\n print(f\"图像已保存到: {image_path}\")\n result_data[\"local_path\"] = image_path\n except Exception as img_err:\n print(f\"保存图像时出错: {str(img_err)}\")\n \n result[\"results\"].append(result_data)\n result_idx += 1\n \n return result\n \n except Exception as e:\n print(f\"执行AI生成的代码时出错: {str(e)}\")\n import traceback\n print(f\"详细错误: {traceback.format_exc()}\")\n return {\n \"success\": False,\n \"error\": {\n \"name\": type(e).__name__,\n \"value\": str(e),\n \"traceback\": traceback.format_exc()\n }\n }\n\ndef download_directory_from_sandbox(sandbox, sandbox_dir_path, local_dir_path):\n \"\"\"从沙箱下载整个目录内容到本地\n \n Args:\n sandbox: 沙箱实例\n sandbox_dir_path: 沙箱中的目录路径\n local_dir_path: 本地保存目录路径\n \n Returns:\n bool: 是否成功下载所有文件\n \"\"\"\n try:\n print(f\"尝试下载目录: {sandbox_dir_path} -> {local_dir_path}\")\n \n # 确保本地目录存在\n os.makedirs(local_dir_path, exist_ok=True)\n \n # 列出沙箱中指定目录下的所有文件\n try:\n files = sandbox.files.list(sandbox_dir_path)\n # print(f\"获取到文件列表: {sandbox_dir_path}, 类型: {type(files)}\")\n # if files and len(files) > 0:\n # print(f\"第一个文件类型: {type(files[0])}, 内容: {files[0]}\")\n # # 检查对象属性\n # print(f\"文件对象可用属性: {dir(files[0])}\")\n except Exception as e:\n print(f\"列出文件时出错: {sandbox_dir_path}, 错误: {str(e)}\")\n return False\n \n if not files:\n print(f\"沙箱中目录 {sandbox_dir_path} 为空或不存在\")\n return False\n \n downloaded_count = 0\n # 定义需要跳过的系统文件\n skip_files = {'.bashrc', '.bash_logout', '.profile'}\n \n # 遍历并下载每个文件\n for file_info in files:\n try:\n # 使用dir()查看对象有哪些属性\n print(f\"文件信息对象属性: {dir(file_info)}\")\n \n # 尝试安全获取name和type属性\n file_name = getattr(file_info, \"name\", None)\n if file_name is None:\n print(f\"警告: 无法获取文件名, 跳过此文件\")\n continue\n \n file_type = getattr(file_info, \"type\", \"file\") # 默认为文件类型\n # 如果 file_type 是枚举, 使用其 value 进行判断\n type_value = file_type.value if hasattr(file_type, \"value\") else file_type\n \n # 跳过不需要的系统文件或系统目录（隐藏文件/目录）\n if file_name in skip_files or (file_name.startswith('.') and type_value == 'dir'):\n print(f\"跳过系统文件或目录: {file_name}\")\n continue\n \n print(f\"处理文件: {file_name}, 类型: {type_value}\")\n \n sandbox_file_path = f\"{sandbox_dir_path}/{file_name}\"\n local_file_path = os.path.join(local_dir_path, file_name)\n \n if type_value == 'dir':\n # 递归下载子目录\n print(f\"发现子目录: {sandbox_file_path}\")\n if download_directory_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n else:\n # 下载文件\n print(f\"下载文件: {sandbox_file_path} -> {local_file_path}\")\n if download_file_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n except Exception as e:\n print(f\"处理文件时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n continue\n \n if downloaded_count > 0:\n print(f\"从 {sandbox_dir_path} 下载了 {downloaded_count} 个文件/目录到 {local_dir_path}\")\n return True\n return False\n \n except Exception as e:\n print(f\"下载整个目录时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n\n##############################################################################\n# 检查E2B代码解释器工具是否已注册\n##############################################################################\n\nprint_separator(\"检查E2B代码解释器工具是否已注册\")\n\n# 获取所有已注册的工具（以字典格式）\nregistered_tools = get_registered_tools(as_dict=True)\n\n# 打印所有已注册的工具\nprint(\"\\n已注册的工具:\")\nfor name, info in registered_tools.items():\n print(f\"- {name} (类别: {info['category'].value})\")\n\n# 检查E2B代码解释器工具是否已注册\ne2b_tool_name = \"e2b_code_interpreter\"\nif e2b_tool_name in registered_tools:\n print(f\"\\nE2B代码解释器工具已成功注册: {e2b_tool_name}\")\nelse:\n print(f\"\\n警告: E2B代码解释器工具未注册\")\n # 手动注册E2B代码解释器工具\n print(\"尝试手动注册E2B代码解释器工具...\")\n try:\n from core.tools.registry import register_tool\n e2b_tool = E2BCodeInterpreterTool()\n register_tool(e2b_tool, ToolCategory.CODE_INTERPRETER)\n print(f\"已手动注册工具: {e2b_tool.name}\")\n except Exception as e:\n print(f\"手动注册E2B代码解释器工具失败: {e}\")\n\n##############################################################################\n# 创建ReactAgent实例\n##############################################################################\n\nprint_separator(\"创建ReactAgent实例\")\n\n# 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n# 从注册表中只获取代码解释器类工具列表\ntools_list = get_tools_by_category(ToolCategory.CODE_INTERPRETER)\n\n# 打印获取到的代码解释器工具\nprint(\"\\n获取到的代码解释器工具:\")\nfor tool in tools_list:\n print(f\"- {tool.name}: {tool.description}\")\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"sandbox_test_agent\",\n # 提示词强调测试沙箱环境的各种功能\n prompt=(\n \"你是一位专业的沙箱环境测试专家，负责测试E2B代码解释器沙箱环境的各种功能。\\n\"\n \"你有强大的代码执行工具可以使用：\\n\"\n \"- e2b_code_interpreter: 用于在沙箱环境中执行Python代码\\n\\n\"\n \"当进行沙箱环境测试时，请遵循以下方法论：\\n\"\n \"1. 分析测试需求：理解需要测试的沙箱功能\\n\"\n \"2. 设计测试用例：针对特定功能设计测试代码\\n\"\n \"3. 执行测试：使用e2b_code_interpreter工具执行测试代码\\n\"\n \"4. 分析结果：解释测试结果，判断功能是否正常\\n\"\n \"5. 记录问题：如有异常，记录问题并提供详细信息\\n\\n\"\n \"重要提示：\\n\"\n \"- 优先使用e2b_code_interpreter工具执行Python代码\\n\"\n \"- 测试代码应包含详细注释，解释测试目的和预期结果\\n\"\n \"- 所有文件和图片必须保存在沙盒环境中的特定目录，不要直接返回图片\\n\"\n \"- 图片不允许在回复中展示！Images are not allowed in the response!\\n\"\n \"- 测试应覆盖沙箱的各种功能，包括但不限于：\\n\"\n \" * 基本Python代码执行\\n\"\n \" * 文件系统操作（创建、读取、写入文件）\\n\"\n \" * 包管理（安装和使用第三方包）\\n\"\n \" * 系统命令执行（使用!前缀执行shell命令）\\n\"\n \" * 数据处理和可视化\\n\"\n \" * 异常处理和错误恢复\\n\"\n ),\n)\n\n# 添加调试信息，验证工具列表和沙箱实例的初始状态\nprint(\"\\n验证ReactAgent工具列表和沙箱实例初始状态:\")\nprint(f\"react_agent.tools类型: {type(react_agent.tools)}\")\nprint(f\"react_agent.tools长度: {len(react_agent.tools)}\")\n\n# 遍历所有工具，检查是否有sandbox属性\nfor i, tool in enumerate(react_agent.tools):\n print(f\"\\n工具[{i}]类型: {type(tool)}\")\n print(f\"工具[{i}]名称: {getattr(tool, 'name', '未知')}\")\n print(f\"工具[{i}]是否有sandbox属性: {'sandbox' in dir(tool)}\")\n \n # 如果有sandbox属性，打印沙箱实例信息\n if 'sandbox' in dir(tool):\n print(f\"工具[{i}]的sandbox类型: {type(tool.sandbox)}\")\n print(f\"工具[{i}]的sandbox是否可用: {getattr(tool, '_is_available', False)}\")\n print(f\"工具[{i}]的初始化错误: {getattr(tool, '_init_error', None)}\")\n\n# 编译Agent\nagent = react_agent.compile()\n\n# # 获取图对象\n# graph = agent.get_graph()\n\n# # 获取当前文件名（不含路径和扩展名）\n# current_file = os.path.basename(__file__)\n# file_name_without_ext = os.path.splitext(current_file)[0]\n# graph_dir = os.path.join(os.path.dirname(__file__), \"graphs\")\n\n# # 确保 graphs 目录存在\n# os.makedirs(graph_dir, exist_ok=True)\n\n# # 生成与文件名一致的图片名，并保存到 examples/graphs 目录\n# image_data = graph.draw_mermaid_png()\n# graph_path = os.path.join(graph_dir, f\"{file_name_without_ext}.png\")\n\n# # 保存图片（如果已存在则覆盖）\n# with open(graph_path, \"wb\") as f:\n# f.write(image_data)\n\n# print(f\"工作流图已保存为 {graph_path}\")\n\n##############################################################################\n# 测试用例1：基本Python代码执行和环境信息\n##############################################################################\n\ndef run_test_case_1():\n print_separator(\"测试用例1：基本Python代码执行和环境信息\")\n print(\"\\n查询: 测试基本Python代码执行和获取环境信息\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请执行一段Python代码，测试基本的数学运算、字符串操作，并获取沙箱环境的系统信息（Python版本、操作系统信息等）。\")\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例1结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n\n##############################################################################\n# 测试用例2：文件系统操作\n##############################################################################\n\ndef run_test_case_2():\n print_separator(\"测试用例2：文件系统操作\")\n print(\"\\n查询: 测试沙箱环境的文件系统操作\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请测试沙箱环境的文件系统操作，包括创建目录、创建文件、写入内容、读取内容、列出目录内容等。创建一个测试目录结构，并将操作结果保存到文件中。文件保存到 /home/user/test_dir\")\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例2结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n \n # 检查是否有E2B沙箱实例，尝试下载生成的文件\n for msg in final_state[\"messages\"]:\n if isinstance(msg, ToolMessage) and msg.name == \"e2b_code_interpreter\":\n try:\n # 尝试解析工具消息内容\n tool_output = json.loads(msg.content)\n print(f\"\\n工具消息内容解析成功: {type(tool_output)}\")\n \n # 检查是否有原始输出\n if hasattr(msg, 'raw_output') and msg.raw_output:\n print(f\"\\n消息包含raw_output属性: {type(msg.raw_output)}\")\n \n # 打印react_agent.tools的信息\n print(f\"\\nreact_agent.tools类型: {type(react_agent.tools)}\")\n print(f\"react_agent.tools长度: {len(react_agent.tools)}\")\n \n # 遍历所有工具，检查是否有sandbox属性\n for i, tool in enumerate(react_agent.tools):\n print(f\"\\n工具[{i}]类型: {type(tool)}\")\n print(f\"工具[{i}]名称: {getattr(tool, 'name', '未知')}\")\n print(f\"工具[{i}]是否有sandbox属性: {'sandbox' in dir(tool)}\")\n if 'sandbox' in dir(tool):\n print(f\"工具[{i}]的sandbox类型: {type(tool.sandbox)}\")\n \n # 遍历 react_agent.tools 以查找 E2B 相关工具\n sandbox = None\n for tool in react_agent.tools:\n if hasattr(tool, \"sandbox\"):\n sandbox = tool.sandbox\n break # 找到后就退出循环\n \n if sandbox:\n print(\"\\n成功获取沙箱实例!\")\n print(f\"沙箱实例类型: {type(sandbox)}\")\n \n # 从沙箱下载生成的文件\n output_dir = os.path.join(os.path.dirname(__file__), \"output\", \"sandbox_test\")\n os.makedirs(output_dir, exist_ok=True)\n print(f\"输出目录已创建: {output_dir}\")\n \n # 尝试下载测试目录，路径和提示中保持一致\n sandbox_test_path = \"/home/user/test_dir\"\n print(f\"尝试从沙箱下载目录: {sandbox_test_path}\")\n download_directory_from_sandbox(sandbox, sandbox_test_path, os.path.join(output_dir, \"test_dir\"))\n else:\n print(\"\\n错误: 无法获取沙箱实例，没有找到具有sandbox属性的工具\")\n else:\n print(\"\\n错误: 消息没有raw_output属性\")\n except Exception as e:\n print(f\"处理工具消息时出错: {str(e)}\")\n\n##############################################################################\n# 测试用例3：包管理和第三方库使用\n##############################################################################\n\ndef run_test_case_3():\n print_separator(\"测试用例3：包管理和第三方库使用\")\n print(\"\\n查询: 测试沙箱环境的包管理和第三方库使用\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请测试沙箱环境的包管理功能，安装一个不常见的第三方库（如wordcloud、pycountry等），并使用该库编写一个简单的示例程序。验证包安装和使用是否正常。\")\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例3结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n\n##############################################################################\n# 测试用例4：Shell命令执行\n##############################################################################\n\ndef run_test_case_4():\n print_separator(\"测试用例4：Shell命令执行\")\n print(\"\\n查询: 测试沙箱环境的Shell命令执行\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请测试沙箱环境中执行Shell命令的功能，使用!前缀执行一系列Linux命令，包括系统信息查询、目录操作、文件查找等。将命令执行结果保存到文件（/home/user/shell_commands_results.txt）中。\")\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例4结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n \n # 尝试下载生成的文件\n for msg in final_state[\"messages\"]:\n if isinstance(msg, ToolMessage) and msg.name == \"e2b_code_interpreter\":\n try:\n print(f\"\\n测试用例4: 检查工具消息类型: {type(msg)}\")\n print(f\"测试用例4: 工具消息名称: {msg.name}\")\n \n # 检查react_agent.tools的信息\n print(f\"\\n测试用例4: react_agent.tools类型: {type(react_agent.tools)}\")\n print(f\"测试用例4: react_agent.tools长度: {len(react_agent.tools)}\")\n \n # 遍历 react_agent.tools 以查找 E2B 相关工具\n sandbox = None\n for tool in react_agent.tools:\n if hasattr(tool, \"sandbox\"):\n sandbox = tool.sandbox\n break # 找到后就退出循环\n \n if sandbox:\n print(\"\\n测试用例4: 成功获取沙箱实例!\")\n print(f\"测试用例4: 沙箱实例类型: {type(sandbox)}\")\n print(f\"测试用例4: 沙箱实例属性: {dir(sandbox)[:10]}...\")\n \n output_dir = os.path.join(os.path.dirname(__file__), \"output\", \"sandbox_test\")\n os.makedirs(output_dir, exist_ok=True)\n print(f\"测试用例4: 输出目录已创建: {output_dir}\")\n \n # 尝试下载shell命令结果文件，路径和提示中保持一致\n sandbox_file_path = \"/home/user/shell_commands_results.txt\"\n local_file_path = os.path.join(output_dir, \"shell_commands_results.txt\")\n print(f\"测试用例4: 尝试下载文件: {sandbox_file_path} -> {local_file_path}\")\n download_file_from_sandbox(sandbox, sandbox_file_path, local_file_path)\n else:\n print(\"\\n测试用例4: 错误: 无法获取沙箱实例，没有找到具有sandbox属性的工具\")\n print(f\"测试用例4: react_agent.tools的类型和长度: {type(react_agent.tools)}, {len(react_agent.tools)}\")\n except Exception as e:\n print(f\"下载文件时出错: {str(e)}\")\n\n##############################################################################\n# 测试用例5：数据处理和可视化\n##############################################################################\n\ndef run_test_case_5():\n print_separator(\"测试用例5：数据处理和可视化\")\n print(\"\\n查询: 测试沙箱环境的数据处理和可视化功能\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=(\n \"请测试沙箱环境的数据处理和可视化功能，生成一些随机数据，使用pandas进行数据处理，\"\n \"然后使用matplotlib创建多种类型的图表（折线图、柱状图、散点图等）。\\n\"\n \"严格按照以下要求:\\n\"\n \"1. 将所有图表保存到 /home/user/visualizations 目录\\n\"\n \"2. 不要在回复中包含图片 - 图片直接保存到上述目录即可\\n\"\n \"3. Images are not allowed in the response!\\n\"\n \"4. 只需描述你做了什么，创建了哪些图表，并说明它们保存在哪里\\n\"\n \"5. 请确保目录存在后再保存图片\\n\"\n ))\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例5结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n \n # 尝试下载生成的图表文件\n for msg in final_state[\"messages\"]:\n if isinstance(msg, ToolMessage) and msg.name == \"e2b_code_interpreter\":\n try:\n # 遍历 react_agent.tools 以查找 E2B 相关工具\n sandbox = None\n for tool in react_agent.tools:\n if hasattr(tool, \"sandbox\"):\n sandbox = tool.sandbox\n break # 找到后就退出循环\n \n if sandbox:\n output_dir = os.path.join(os.path.dirname(__file__), \"output\", \"sandbox_test\")\n os.makedirs(output_dir, exist_ok=True)\n \n # 针对性地下载可视化目录中的图表\n vis_dir = \"/home/user/visualizations\"\n local_vis_dir = os.path.join(output_dir, \"visualizations\")\n os.makedirs(local_vis_dir, exist_ok=True)\n print(f\"测试用例5: 下载可视化图表目录: {vis_dir} -> {local_vis_dir}\")\n \n # 尝试列出可视化目录中的文件\n try:\n files = sandbox.files.list(vis_dir)\n if files:\n print(f\"找到图表文件:\")\n for file_info in files:\n file_name = getattr(file_info, \"name\", \"未知文件\")\n print(f\"- {file_name}\")\n else:\n print(f\"警告: 可视化目录为空或不存在\")\n except Exception as e:\n print(f\"列出可视化目录文件时出错: {str(e)}\")\n \n # 执行下载\n success = download_directory_from_sandbox(sandbox, vis_dir, local_vis_dir)\n if success:\n print(f\"✅ 成功下载可视化图表\")\n else:\n print(f\"⚠️ 下载可视化图表失败，尝试下载整个用户目录作为备份\")\n download_directory_from_sandbox(sandbox, \"/home/user\", output_dir)\n else:\n print(\"\\n错误: 无法获取沙箱实例，没有找到具有sandbox属性的工具\")\n except Exception as e:\n print(f\"下载文件时出错: {str(e)}\")\n import traceback\n print(f\"错误详情: {traceback.format_exc()}\")\n\n##############################################################################\n# 测试用例6：异常处理和错误恢复\n##############################################################################\n\ndef run_test_case_6():\n print_separator(\"测试用例6：异常处理和错误恢复\")\n print(\"\\n查询: 测试沙箱环境的异常处理和错误恢复能力\")\n \n # 定义输入\n inputs = {\n \"messages\": [\n HumanMessage(content=\"请测试沙箱环境的异常处理和错误恢复能力。编写一段包含各种常见错误的Python代码（如语法错误、除零错误、文件不存在错误等），然后展示如何捕获和处理这些异常。验证沙箱环境是否能正确报告错误并继续执行后续代码。\")\n ]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取消息列表\n messages = partial_state.get(\"messages\", [])\n if not messages:\n continue\n \n # 获取最新消息\n latest_message = messages[-1]\n \n # 使用log_agent_actions函数记录状态\n log_agent_actions({\"messages\": [latest_message]})\n \n # 打印最终回答\n print_separator(\"测试用例6结果\")\n if final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n\n##############################################################################\n# 主函数 - 运行所有测试用例\n##############################################################################\n\nif __name__ == \"__main__\":\n print_separator(\"开始测试E2B沙箱环境\")\n \n try:\n # 确保输出目录存在\n output_dir = os.path.join(os.path.dirname(__file__), \"output\", \"sandbox_test\")\n os.makedirs(output_dir, exist_ok=True)\n print(f\"创建输出目录: {output_dir}\")\n \n # 确保可视化输出目录存在\n vis_output_dir = os.path.join(output_dir, \"visualizations\")\n os.makedirs(vis_output_dir, exist_ok=True)\n print(f\"创建可视化输出目录: {vis_output_dir}\")\n \n # # 运行测试用例\n # # 运行测试用例1：基本Python代码执行和环境信息\n # run_test_case_1()\n \n # # 运行测试用例2：文件系统操作\n # run_test_case_2()\n \n # # 运行测试用例3：包管理和第三方库使用\n # run_test_case_3()\n \n # # 运行测试用例4：Shell命令执行\n # run_test_case_4()\n \n # 运行测试用例5：数据处理和可视化\n run_test_case_5()\n \n # # 运行测试用例6：异常处理和错误恢复\n # run_test_case_6()\n \n print_separator(\"E2B沙箱环境测试完成\")\n print(\"测试结果已保存到 examples/output/sandbox_test 目录\")\n \n except Exception as e:\n print(f\"测试过程中出错: {str(e)}\")\n finally:\n # 关闭E2B沙箱\n print(\"\\n正在关闭E2B沙箱...\")\n for tool in react_agent.tools:\n if hasattr(tool, 'close'):\n tool.close()"
},
{
"path": "examples/12_planning_supervisor_test.py",
"content": "from langgraph.prebuilt import create_react_agent\nfrom core.agents.react_supervisor_agent import SupervisorAgent\nfrom core.agents.research_agent import ResearchAgent\nfrom core.agents.base.react_agent import ReactAgent\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.func import entrypoint, task\nfrom langgraph.graph import add_messages\nfrom dotenv import load_dotenv\nfrom langchain_community.tools import TavilySearchResults\nload_dotenv() # 自动加载 .env 文件\n\n# 1. 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n\n##############################################################################\n# Agent 1: Joke Generator (Functional API)\n##############################################################################\n\n@task\ndef generate_joke(messages):\n \"\"\"Generate a short joke (no tool calls).\"\"\"\n system_message = {\n \"role\": \"system\", \n \"content\": \"You are a witty comedian. Write a short joke.\"\n }\n # 直接调用 model.invoke，拼接 system_message + 用户消息\n msg = model.invoke([system_message] + messages)\n return msg\n\n@entrypoint()\ndef joke_agent(state):\n # 调用上面的函数型任务\n joke = generate_joke(state['messages']).result()\n # 将产物插入消息列表\n messages = add_messages(state[\"messages\"], [joke])\n return {\"messages\": messages}\n\njoke_agent.name = \"joke_agent\"\n\n##############################################################################\n# Agent 2: Research Expert with Tavily Search (Graph API)\n##############################################################################\n\n# 创建Tavily搜索工具\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n)\n\n# 使用我们自定义的ResearchAgent替代create_react_agent创建的agent\nresearch_agent = ResearchAgent(\n name=\"research_expert\",\n model=model,\n max_iterations=5,\n cache_enabled=True,\n debug=False\n)\nresearch_agent_2 = ReactAgent(\n name=\"research_expert\",\n model=model,\n tools=[tavily_search])\n\n##############################################################################\n# 使用带有Planning功能的SupervisorAgent\n##############################################################################\n\n# 创建 SupervisorAgent 实例，启用Planning功能\nsupervisor = SupervisorAgent(\n agents=[joke_agent,research_agent_2],\n model=model,\n)\n##############################################################################\n# 测试：复杂请求需要规划和多个步骤\n##############################################################################\nresult = supervisor.run({\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": (\n \"I'm preparing a presentation about tech companies. I need three things: \"\n \"1) A joke about tech companies to start with, \"\n \"2) The employee count for FANNG, and \"\n \"3) A comparison of which company has more employees.\"\n )\n }\n ]\n})\n\n##############################################################################\n# 打印最终对话消息\n##############################################################################\nfor m in result[\"messages\"]:\n m.pretty_print()\n\n# 打印任务列表\nprint(\"\\n##############################################################################\")\nprint(\"# 最终任务列表\")\nprint(\"##############################################################################\")\nif \"plan\" in result and result[\"plan\"] and \"tasks\" in result[\"plan\"]:\n tasks = result[\"plan\"][\"tasks\"]\n print(f\"总共 {len(tasks)} 个任务:\")\n for i, task in enumerate(tasks):\n print(f\"\\n任务 {i+1}: {task['description']}\")\n print(f\" 状态: {task['status']}\")\n print(f\" 代理: {task['agent'] if task['agent'] else '未分配'}\")\n print(f\" 创建时间: {task['created_at']}\")\n print(f\" 完成时间: {task['completed_at'] if task['completed_at'] else '未完成'}\")\nelse:\n print(\"没有任务列表信息\")\n\n# 打印原始任务列表（如果存在）\nif \"tasks\" in result:\n print(\"\\n原始任务列表:\")\n for t in result[\"tasks\"]:\n t.pretty_print()"
},
{
"path": "examples/13_multi_agent_roles_test.py",
"content": "from langgraph.prebuilt import create_react_agent\nfrom core.agents.react_supervisor_agent import SupervisorAgent\nfrom core.agents.sub_agents.research_agent import ResearchAgent\nfrom core.agents.sub_agents.coder_agent import CoderAgent\nfrom core.agents.sub_agents.reporter_agent import ReporterAgent\nfrom core.agents.sub_agents.designer_agent import DesignerAgent\nfrom core.agents.sub_agents.data_analyst_agent import DataAnalystAgent\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.messages import AIMessage, HumanMessage, ToolMessage\nfrom langgraph.func import entrypoint, task\nfrom langgraph.graph import add_messages\nfrom dotenv import load_dotenv\nfrom langchain_community.tools import TavilySearchResults\nimport os\nimport logging\nimport sys\nimport io\nimport json\nfrom contextlib import redirect_stdout, redirect_stderr\n\nload_dotenv() # 自动加载 .env 文件\n\n# 1. 初始化大模型\nmodel = ChatOpenAI(model=\"gpt-4o-mini\")\n# 设置日志捕获\nclass LogCapture:\n def __init__(self):\n self.log_buffer = io.StringIO()\n self.log_content = []\n \n def start_capture(self):\n self.log_buffer = io.StringIO()\n return self.log_buffer\n \n def stop_capture(self):\n output = self.log_buffer.getvalue()\n self.log_content.append(output)\n return output\n \n def get_content(self):\n return \"\\n\".join(self.log_content)\n\nlog_capture = LogCapture()\n\n##############################################################################\n# 从沙箱下载文件到本地的函数\n##############################################################################\n\ndef download_file_from_sandbox(sandbox, sandbox_path, local_path):\n \"\"\"从 e2b 沙箱中下载文件并保存到本地，自动区分文本和二进制文件\"\"\"\n try:\n print(f\"读取文件: {sandbox_path}\")\n\n # 判断是否为常见二进制文件类型（可自行扩展）\n binary_extensions = (\n '.png', '.jpg', '.jpeg', '.gif', '.pdf', '.svg',\n '.xlsx', '.xls', '.zip', '.bin', '.pyc', '.pyd',\n '.pptx', '.docx', '.mp3', '.mp4', '.avi', '.mov',\n )\n is_binary = sandbox_path.lower().endswith(binary_extensions)\n\n # 创建目录\n os.makedirs(os.path.dirname(local_path), exist_ok=True)\n\n if is_binary:\n print(\"📦 识别为二进制文件，使用 sandbox.download_file()\")\n content = sandbox.files.read(sandbox_path) # 返回 bytes\n with open(local_path, 'wb') as f:\n f.write(content)\n else:\n print(\"📄 识别为文本文件，使用 sandbox.files.read()\")\n content = sandbox.files.read(sandbox_path) # 返回 str\n with open(local_path, 'w', encoding='utf-8') as f:\n f.write(content)\n\n print(f\"✅ 文件已保存到本地: {local_path}\")\n return True\n\n except Exception as e:\n print(f\"❌ 下载失败: {e}\")\n return False\n\ndef download_directory_from_sandbox(sandbox, sandbox_dir_path, local_dir_path):\n \"\"\"从沙箱下载整个目录内容到本地\n \n Args:\n sandbox: 沙箱实例\n sandbox_dir_path: 沙箱中的目录路径\n local_dir_path: 本地保存目录路径\n \n Returns:\n bool: 是否成功下载所有文件\n \"\"\"\n try:\n print(f\"尝试下载目录: {sandbox_dir_path} -> {local_dir_path}\")\n \n # 确保本地目录存在\n os.makedirs(local_dir_path, exist_ok=True)\n \n # 列出沙箱中指定目录下的所有文件\n try:\n files = sandbox.files.list(sandbox_dir_path)\n except Exception as e:\n print(f\"列出文件时出错: {sandbox_dir_path}, 错误: {str(e)}\")\n return False\n \n if not files:\n print(f\"沙箱中目录 {sandbox_dir_path} 为空或不存在\")\n return False\n \n downloaded_count = 0\n # 定义需要跳过的系统文件\n skip_files = {'.bashrc', '.bash_logout', '.profile'}\n \n # 遍历并下载每个文件\n for file_info in files:\n try:\n # 尝试安全获取name和type属性\n file_name = getattr(file_info, \"name\", None)\n if file_name is None:\n print(f\"警告: 无法获取文件名, 跳过此文件\")\n continue\n \n file_type = getattr(file_info, \"type\", \"file\") # 默认为文件类型\n # 如果 file_type 是枚举, 使用其 value 进行判断\n type_value = file_type.value if hasattr(file_type, \"value\") else file_type\n \n # 跳过不需要的系统文件或系统目录（隐藏文件/目录）\n if file_name in skip_files or (file_name.startswith('.') and type_value == 'dir'):\n print(f\"跳过系统文件或目录: {file_name}\")\n continue\n \n print(f\"处理文件: {file_name}, 类型: {type_value}\")\n \n sandbox_file_path = f\"{sandbox_dir_path}/{file_name}\"\n local_file_path = os.path.join(local_dir_path, file_name)\n \n if type_value == 'dir':\n # 递归下载子目录\n print(f\"发现子目录: {sandbox_file_path}\")\n if download_directory_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n else:\n # 下载文件\n print(f\"下载文件: {sandbox_file_path} -> {local_file_path}\")\n if download_file_from_sandbox(sandbox, sandbox_file_path, local_file_path):\n downloaded_count += 1\n except Exception as e:\n print(f\"处理文件时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n continue\n \n if downloaded_count > 0:\n print(f\"从 {sandbox_dir_path} 下载了 {downloaded_count} 个文件/目录到 {local_dir_path}\")\n return True\n return False\n \n except Exception as e:\n print(f\"下载整个目录时出错: {str(e)}\")\n import traceback\n print(f\"详细错误跟踪: {traceback.format_exc()}\")\n\n\n##############################################################################\n# Agent 2: Research Expert - 使用自定义的ResearchAgent\n##############################################################################\n\nresearch_agent = ResearchAgent(\n name=\"research_expert\",\n model=model,\n max_iterations=5,\n cache_enabled=True,\n debug=True\n)\n\n##############################################################################\n# Agent 3: Coder - 使用自定义的CoderAgent\n##############################################################################\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\ne2b_tool = E2BCodeInterpreterTool()\n\ncoder_agent = CoderAgent(\n name=\"coder_expert\",\n model=model,\n tools=[e2b_tool],\n max_iterations=5,\n cache_enabled=True,\n debug=True\n)\n\n##############################################################################\n# Agent 4: Reporter - 使用自定义的ReporterAgent\n##############################################################################\n\nreporter_agent = ReporterAgent(\n name=\"reporter_expert\",\n model=model,\n max_iterations=5,\n cache_enabled=True,\n)\n\n##############################################################################\n# Agent 5: Designer - 使用自定义的DesignerAgent\n##############################################################################\n\ndesigner_agent = DesignerAgent(\n name=\"designer_expert\",\n model=model,\n max_iterations=5,\n cache_enabled=True,\n)\n\n##############################################################################\n# Agent 6: Data Analyst - 使用自定义的DataAnalystAgent\n##############################################################################\n\ndata_analyst_agent = DataAnalystAgent(\n name=\"data_analyst_expert\",\n model=model,\n max_iterations=5,\n cache_enabled=True,\n)\n\n##############################################################################\n# 使用带有Planning功能的SupervisorAgent协调所有角色\n##############################################################################\n\n# 创建 SupervisorAgent 实例，启用Planning功能\nsupervisor = SupervisorAgent(\n agents=[\n research_agent,\n coder_agent,\n reporter_agent,\n designer_agent,\n data_analyst_agent,\n ],\n model=model,\n enable_planning=True,\n output_mode=\"last_message\"\n)\n\n# 获取当前文件名（不含路径和扩展名）\ncurrent_file = os.path.basename(__file__)\nfile_name_without_ext = os.path.splitext(current_file)[0]\nlogs_dir = os.path.join(os.path.dirname(__file__), \"logs\")\n# 创建图表输出文件路径\nos.makedirs(logs_dir, exist_ok=True)\n# 创建Markdown输出文件路径\nmarkdown_path = os.path.join(logs_dir, f\"{file_name_without_ext}.md\")\n\n##############################################################################\n# 测试：复杂请求需要规划和多个步骤\n##############################################################################\n\ndef save_markdown_log():\n \"\"\"将执行结果保存为Markdown文件\"\"\"\n with open(markdown_path, \"w\", encoding=\"utf-8\") as f:\n f.write(f\"# 执行结果: {file_name_without_ext}\\n\\n\")\n f.write(\"## 图表\\n\\n\")\n f.write(\"## 执行日志\\n\\n\")\n f.write(\"```\\n\")\n f.write(log_capture.get_content())\n f.write(\"\\n```\\n\")\n print(f\"执行日志已保存到 {markdown_path}\")\n\nif __name__ == \"__main__\":\n try:\n # 开始捕获输出\n log_buffer = log_capture.start_capture()\n \n with redirect_stdout(log_buffer), redirect_stderr(log_buffer):\n print(f\"开始执行 {current_file} 测试...\")\n \n # 测试1：需要研究和编码的任务\n print(\"\\n## 测试1：需要研究和编码的任务\")\n final_state = supervisor.run({\n \"messages\": [\n {\n \"role\": \"user\",\n \"content\": (\n \"我需要一个Python爬虫来获取 https://www.paulgraham.com/articles.html 所有articles列表，并将结果保存为CSV文件,放在/home/user下面。\"\n \"并将你测试通过的爬虫代码返回给我。\"\n \"请确保你的代码能够正常运行。\"\n \"如果遇到问题，请重试。\"\n )\n }\n ]\n })\n \n print(\"\\n测试1结果:\")\n for m in final_state[\"messages\"]:\n m.pretty_print()\n \n # 遍历 react_agent.tools 以查找 E2B 相关工具\n try:\n # 遍历 react_agent.tools 以查找 E2B 相关工具\n sandbox = None\n for tool in coder_agent.tools:\n if hasattr(tool, \"sandbox\"):\n sandbox = tool.sandbox\n break # 找到后就退出循环\n\n if sandbox:\n # 设定输出目录\n output_dir = os.path.join(os.getcwd(), \"examples/output/sandbox_files\")\n os.makedirs(output_dir, exist_ok=True)\n\n # 直接下载主要工作目录\n print(\"\\n从沙箱下载文件到本地...\")\n download_directory_from_sandbox(sandbox, \"/home/user\", output_dir)\n\n # 下载临时目录中可能的图表和数据文件\n # download_directory_from_sandbox(sandbox, \"/tmp\", output_dir)\n\n print(f\"\\n文件已保存到目录: {output_dir}\")\n sandbox.close()\n except Exception as e:\n print(f\"从沙箱下载文件时出错: {str(e)}\")\n\n \n finally:\n # 停止捕获并保存结果\n log_capture.stop_capture()\n save_markdown_log()"
},
{
"path": "examples/14_mcp_client_fetch_test.py",
"content": "import os\nimport sys\nimport asyncio\nimport traceback\nfrom typing import Dict, Optional, Type\n\nfrom dotenv import load_dotenv\n\n# 在这里添加项目根目录到路径，方便导入\nsys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))\nload_dotenv()\n\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.prebuilt import create_react_agent\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import HumanMessage\n\ntry:\n from pydantic.v1 import BaseModel, Field\nexcept ImportError:\n from pydantic import BaseModel, Field # type: ignore\n\nfrom core.mcp.client import MCPClient\nfrom core.mcp.config_loader import load_config, MCPConfig, StdioConfig\nfrom core.llm.llm_manager import LLMManager\n\ntry:\n from mcp.types import CallToolRequest\n CALL_TOOL_REQ_AVAILABLE = True\nexcept ImportError:\n CallToolRequest = None\n CALL_TOOL_REQ_AVAILABLE = False\n\n# 这是唯一保留的 fetch schema\ntry:\n class FetchInputSchema(BaseModel):\n url: str = Field(..., description=\"URL to fetch\")\n max_length: Optional[int] = Field(default=5000)\n start_index: Optional[int] = Field(default=0)\n raw: Optional[bool] = Field(default=False)\n FETCH_SCHEMA_AVAILABLE = True\nexcept Exception:\n FetchInputSchema = None\n FETCH_SCHEMA_AVAILABLE = False\n\nCENTRAL_CONFIG_PATH = os.path.join(os.path.dirname(__file__), \"..\", \"core\", \"mcp\", \"mcp_server_config.json\")\nLLM_ID_FOR_TESTING = \"openai_gpt4o_mini\"\nllm_manager = LLMManager()\n\nclass MCPToolRunner(BaseTool):\n name: str = \"needs_override\"\n description: str = \"needs_override\"\n args_schema: Optional[Type[BaseModel]] = None\n\n client: MCPClient = Field(exclude=True)\n\n class Config:\n arbitrary_types_allowed = True\n\n async def _arun(self, **kwargs) -> str:\n if not self.client or not self.client.session:\n return f\"ERROR: MCP Client session inactive for {self.name}.\"\n if not CALL_TOOL_REQ_AVAILABLE:\n return \"ERROR: CallToolRequest unavailable.\"\n\n try:\n print(f\" [_arun:{self.name}] Sending MCP request with args: {kwargs}\")\n result_message = await asyncio.wait_for(\n self.client.session.call_tool(self.name, kwargs),\n timeout=120.0\n )\n # 简化: 只检查 result 和 error\n if hasattr(result_message, \"result\"):\n return str(result_message.result)\n elif hasattr(result_message, \"error\"):\n return f\"Tool Error: {result_message.error.message}\"\n else:\n return \"Unknown response\"\n except asyncio.TimeoutError:\n return \"Error: Timeout.\"\n except Exception as e:\n return f\"Error: {e}\\n{traceback.format_exc()}\"\n\n def _run(self, **kwargs) -> str:\n print(f\" [_run:{self.name}] Running async method via asyncio.run()...\")\n try:\n return asyncio.run(self._arun(**kwargs))\n except Exception as e:\n return f\"Error in sync wrapper: {e}\"\n\nasync def run_fetch_test(server_config_key: str, all_configs: Dict[str, MCPConfig]):\n print(f\"\\n=== Running STDIO BaseTool Test for Server '{server_config_key}' (Tool: 'fetch') ===\")\n if not FETCH_SCHEMA_AVAILABLE:\n print(\"ERROR: Fetch Schema missing.\")\n return False\n if not CALL_TOOL_REQ_AVAILABLE:\n print(\"ERROR: CallToolRequest unavailable.\")\n return False\n\n server_config = all_configs.get(server_config_key)\n if not server_config:\n print(f\"ERROR: Config for '{server_config_key}' not found.\")\n return False\n if not isinstance(server_config.connection, StdioConfig):\n print(f\"ERROR: Config '{server_config_key}' not STDIO.\")\n return False\n\n try:\n model = llm_manager.get_model(LLM_ID_FOR_TESTING)\n print(f\"Using LLM: {getattr(model, 'model_name', LLM_ID_FOR_TESTING)}\")\n except ValueError as e:\n print(f\"获取 LLM 出错: {e}.\")\n return False\n\n test_success = False\n async with MCPClient(server_config) as client:\n if not client.session:\n print(\"ERROR: MCP session not established!\")\n return False\n\n try:\n runner = MCPToolRunner(\n client=client,\n name=\"fetch\",\n description=\"Fetches URL content as markdown.\",\n args_schema=FetchInputSchema\n )\n tools = [runner]\n except Exception as e_inst:\n print(f\"ERROR: Failed to instantiate MCPToolRunner: {e_inst}\")\n return False\n\n agent = create_react_agent(model, tools)\n query = (\n \"Use the fetch tool to get the content of https://www.google.com \"\n \"and tell me its title (first 50 chars).\"\n )\n print(f\"\\nQuery: {query}\")\n\n try:\n response = await asyncio.wait_for(\n agent.ainvoke({\"messages\": [{\"role\": \"user\", \"content\": query}]}),\n timeout=180.0\n )\n print(f\"\\nAgent Final Response:\")\n if response and \"messages\" in response and response[\"messages\"]:\n response_content = response[\"messages\"][-1].content\n print(response_content)\n if \"google\" in response_content.lower():\n print(\"\\n✅ Test PASS\")\n test_success = True\n else:\n print(\"\\n❌ Test FAIL (title not found)\")\n test_success = False\n else:\n print(\"No valid response from agent.\")\n test_success = False\n except Exception as e:\n print(f\"Exception: {e}\")\n test_success = False\n\n return test_success\n\nasync def main():\n print(\"Starting a simplified MCP Integration Test for 'fetch_via_uvx' only...\")\n try:\n all_configs = load_config(CENTRAL_CONFIG_PATH)\n print(f\"Loaded {len(all_configs)} server configs.\")\n except Exception as e:\n print(f\"Error loading config: {e}\")\n return\n\n # 只测试 fetch_via_uvx\n result = await run_fetch_test(\"fetch_via_uvx\", all_configs)\n if result:\n print(\"\\nALL GOOD: 'fetch' test passed.\")\n else:\n print(\"\\nTEST FAILED: 'fetch' test didn't pass.\")\n print(\"Done.\")\n\nif __name__ == \"__main__\":\n asyncio.run(main())"
},
{
"path": "examples/15_mcp_agent_test.py",
"content": "# examples/14_mcp_fetch_basetool_test.py (最终版 - BaseTool 子类)\nimport os\nimport sys\nimport asyncio\nimport json\nfrom dotenv import load_dotenv\nimport traceback\nfrom typing import List, Dict, Any, Optional, Type\n\n# --- 前置要求 ---\n# 1. 确保 core/mcp/client.py 和 core/mcp/config_loader.py 是最新版本 (含 AsyncExitStack 和导入修复)。\n# 2. 确保 core/mcp/config.json 文件存在，并包含 \"fetch_via_uvx\" 配置 (使用 uvx + stdio)。\n# 3. 确保已安装 uv (`pip install uv`) 和 mcp-server-fetch。\n# 4. 确保 OpenAI API Key (或其他 LLM Key) 在 .env 或环境变量中设置。\n# 5. 推荐设置 LangSmith 环境变量用于详细追踪 Agent 行为。\n# ---\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))\nload_dotenv()\n\n# --- 核心依赖导入 ---\n# LangChain\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.prebuilt import create_react_agent\nfrom langchain_core.tools import BaseTool\nfrom langchain_core.messages import HumanMessage\ntry:\n # 尝试导入 Pydantic v1 (LangChain 常用的版本)\n from langchain_core.pydantic_v1 import BaseModel, Field\nexcept ImportError:\n try:\n # 如果 V1 不可用，尝试导入 V2\n from pydantic import BaseModel, Field # type: ignore\n except ImportError:\n print(\"CRITICAL ERROR: Pydantic (v1 or v2) not found.\")\n sys.exit(1)\n# MCP Client/Config\ntry: from core.mcp.client import MCPClient\nexcept ImportError: print(\"CRITICAL ERROR: Cannot import MCPClient.\"); sys.exit(1)\ntry: from core.mcp.config_loader import load_config, MCPConfig, StdioConfig\nexcept ImportError: print(\"CRITICAL ERROR: Cannot import config loader.\"); sys.exit(1)\n# LLM\nfrom core.llm.llm_manager import LLMManager\n# MCP Types\ntry: from mcp.types import CallToolRequest; CALL_TOOL_REQ_AVAILABLE = True\nexcept ImportError: CallToolRequest = None; CALL_TOOL_REQ_AVAILABLE = False\n# ---\n\n# --- Fetch Tool Schema 定义 ---\nFETCH_SCHEMA_AVAILABLE = False\nFetchInputSchema = None\ntry:\n class FetchInputSchema(BaseModel): # 使用导入的 BaseModel\n url: str = Field(..., description=\"URL to fetch\")\n max_length: Optional[int] = Field(default=5000, description=\"Maximum number of characters to return\")\n start_index: Optional[int] = Field(default=0, description=\"Start content from this character index\")\n raw: Optional[bool] = Field(default=False, description=\"Get raw content without markdown conversion\")\n FETCH_SCHEMA_AVAILABLE = True\nexcept Exception as e_pyd_fetch: print(f\"ERROR defining FetchInputSchema: {e_pyd_fetch}\")\n# ---\n\n# --- 全局设置 ---\n# **重要**: 确认此路径指向你的中央配置文件\nCENTRAL_CONFIG_PATH = os.path.join(os.path.dirname(__file__), \"..\", \"core\", \"mcp\", \"mcp_server_config.json\")\n# 使用 OpenAI 模型通常更稳定\nLLM_ID_FOR_TESTING = \"openai_gpt4o_mini\"\n# 要测试的服务器在 config.json 中的 key\nSERVER_KEY_TO_TEST = \"fetch_via_uvx\"\n# 要测试的工具名称\nTOOL_NAME_TO_TEST = \"fetch\"\n# 要测试的工具的正确 Schema\nCORRECT_SCHEMA_FOR_TOOL = FetchInputSchema\n# 要测试的工具的描述\nTOOL_DESCRIPTION = \"Fetches web content as markdown. Input requires 'url' (string) and optional 'max_length', 'start_index', 'raw'.\"\n\n# --- Everything MCP 服务器设置 ---\nEVERYTHING_SERVER_KEY = \"everything\"\nEVERYTHING_ECHO_TOOL = \"echo\"\nEVERYTHING_ADD_TOOL = \"add\"\n\n# --- Everything MCP 工具 Schema 定义 ---\nECHO_SCHEMA_AVAILABLE = False\nEchoInputSchema = None\ntry:\n class EchoInputSchema(BaseModel):\n message: str = Field(..., description=\"Message to echo back\")\n ECHO_SCHEMA_AVAILABLE = True\nexcept Exception as e_pyd_echo: print(f\"ERROR defining EchoInputSchema: {e_pyd_echo}\")\n\nADD_SCHEMA_AVAILABLE = False\nAddInputSchema = None\ntry:\n class AddInputSchema(BaseModel):\n a: float = Field(..., description=\"First number\")\n b: float = Field(..., description=\"Second number\")\n ADD_SCHEMA_AVAILABLE = True\nexcept Exception as e_pyd_add: print(f\"ERROR defining AddInputSchema: {e_pyd_add}\")\n\nllm_manager = LLMManager()\n\n# --- 标准 BaseTool 子类定义，用于桥接 MCP 调用 ---\nclass MCPToolRunner(BaseTool):\n \"\"\"\n 通过 MCP 调用服务器上工具的标准 BaseTool 实现。\n \"\"\"\n # --- 类属性 (将在实例化时被覆盖) ---\n name: str = \"mcp_tool_runner\" # Default name\n description: str = \"Runs a tool via MCP\"\n args_schema: Optional[Type[BaseModel]] = None\n\n # --- 实例属性 ---\n client: MCPClient = Field(exclude=True) # 存储客户端引用\n\n # Pydantic 配置 (根据你使用的 BaseModel 版本)\n class Config: arbitrary_types_allowed = True\n\n async def _arun(self, **kwargs) -> str:\n \"\"\"异步执行：构造 MCP 请求并调用 client.session.call_tool\"\"\"\n if not self.client or not self.client.session: return f\"ERROR: MCP Client session inactive for {self.name}.\"\n if not CALL_TOOL_REQ_AVAILABLE: return \"ERROR: CallToolRequest unavailable.\"\n\n try:\n # kwargs 应该是 LangChain 根据 args_schema 验证和准备好的参数\n print(f\" [_arun:{self.name}] Preparing MCP request with args: {kwargs}\")\n # 不再需要构造CallToolRequest对象，直接传递工具名称和参数\n print(f\" [_arun:{self.name}] Calling tool '{self.name}' with args: {kwargs}\")\n\n # 调用 MCP session - 直接传递工具名称和参数\n result_message = await asyncio.wait_for(\n self.client.session.call_tool(self.name, kwargs),\n timeout=120.0 # 给予足够的网络和执行超时\n )\n\n # 处理结果 - 简化处理逻辑，直接检查content属性\n print(f\" [_arun:{self.name}] MCP Response received, type: {type(result_message)}\")\n \n # 直接检查是否有content属性（根据日志显示的响应结构）\n if hasattr(result_message, 'content'):\n content = result_message.content\n print(f\" [_arun:{self.name}] Found content attribute, type: {type(content)}\")\n \n # 如果content是列表且不为空\n if isinstance(content, list) and len(content) > 0:\n first_item = content[0]\n print(f\" [_arun:{self.name}] Content is a list, first item type: {type(first_item)}\")\n \n # 尝试获取text属性\n if hasattr(first_item, 'text'):\n print(f\" [_arun:{self.name}] First item has text attribute, returning text\")\n return first_item.text\n else:\n print(f\" [_arun:{self.name}] First item has no text attribute, converting to string\")\n return str(first_item)\n elif hasattr(content, 'text'):\n print(f\" [_arun:{self.name}] Content has text attribute, returning text\")\n return content.text\n else:\n print(f\" [_arun:{self.name}] Content has no text attribute, converting to string\")\n return str(content)\n # 如果没有content属性，回退到检查result属性\n elif hasattr(result_message, 'result'):\n res_val = result_message.result\n print(f\" [_arun:{self.name}] Found result attribute: {str(res_val)[:500]}...\")\n return str(res_val) if not isinstance(res_val, str) else res_val\n elif hasattr(result_message, 'error'):\n err_msg = result_message.error.message\n print(f\" [_arun:{self.name}] MCP Tool Error: {err_msg}\")\n # 对于 Agent，返回错误通常比抛出异常更好处理\n return f\"Tool Error: {err_msg}\"\n else:\n # 打印完整的响应对象，帮助诊断问题\n print(f\" [_arun:{self.name}] Unknown MCP response format. Full response object: {result_message}\")\n print(f\" [_arun:{self.name}] Response type: {type(result_message)}\")\n print(f\" [_arun:{self.name}] Response dir: {dir(result_message)}\")\n \n # 尝试处理特殊的响应格式\n if hasattr(result_message, 'content'):\n content = result_message.content\n print(f\" [_arun:{self.name}] Found content attribute in response\")\n \n # 处理content是列表的情况\n if isinstance(content, list) and len(content) > 0:\n print(f\" [_arun:{self.name}] Content is a list with {len(content)} items\")\n first_item = content[0]\n if hasattr(first_item, 'text'):\n print(f\" [_arun:{self.name}] First item has text attribute, returning text\")\n return first_item.text\n elif hasattr(first_item, 'type') and hasattr(first_item, 'text'):\n print(f\" [_arun:{self.name}] First item has type and text attributes, returning text\")\n return first_item.text\n else:\n print(f\" [_arun:{self.name}] First item has no text attribute, converting to string\")\n return str(first_item)\n # 处理content是单个对象的情况\n elif hasattr(content, 'text'):\n print(f\" [_arun:{self.name}] Content has text attribute, returning text\")\n return content.text\n else:\n print(f\" [_arun:{self.name}] Content has no text attribute, converting to string\")\n return str(content)\n \n # 尝试提取更多信息\n response_details = \"\"\n for attr in dir(result_message):\n if not attr.startswith('_'):\n try:\n value = getattr(result_message, attr)\n if not callable(value):\n response_details += f\"\\n - {attr}: {value}\"\n except Exception as attr_err:\n response_details += f\"\\n - {attr}: [Error accessing: {attr_err}]\"\n print(f\" [_arun:{self.name}] Response details: {response_details}\")\n return f\"Unknown response from MCP tool {self.name}. Details: {response_details}\"\n except asyncio.TimeoutError:\n print(f\" [_arun:{self.name}] MCP call timeout.\")\n return f\"Error: Timeout calling MCP tool {self.name}.\"\n except Exception as e:\n print(f\" [_arun:{self.name}] Unexpected error during MCP call: {e}\")\n print(traceback.format_exc())\n # 返回包含 Traceback 的错误，方便调试\n return f\"Unexpected Error calling {self.name}: {e}\\n{traceback.format_exc()}\"\n\n def _run(self, **kwargs) -> str:\n \"\"\"同步执行 (简单实现，通过运行异步方法)\"\"\"\n print(f\" [_run:{self.name}] Running async method via asyncio.run()...\")\n try:\n # 注意: 在已运行的事件循环中调用 asyncio.run 会报错\n # 更好的方法是检查当前循环或使用 anyio/nest_asyncio\n # 但为了满足 BaseTool 要求，先用简单方式，如果 Agent 只用 async 就没问题\n # 如果 Agent 强制用 sync，可能需要更复杂的处理\n # return asyncio.run(self._arun(**kwargs))\n # 更安全的方式是提示不支持或使用更复杂的同步转异步\n return \"Synchronous execution not fully supported, please use async.\"\n except Exception as e:\n print(f\" [_run:{self.name}] Error: {e}\")\n return f\"Error in sync wrapper: {e}\"\n# ---\n\n# --- 主要测试逻辑 ---\nasync def run_fetch_test():\n \"\"\"运行 Fetch Server 测试 (使用 BaseTool 子类)\"\"\"\n print(f\"\\n=== Running Fetch Server Test (BaseTool Subclass Method) ===\")\n\n # 检查依赖和 Schema 定义\n if not FETCH_SCHEMA_AVAILABLE: print(\"ERROR: FetchInputSchema not available.\"); return False\n if not CALL_TOOL_REQ_AVAILABLE: print(\"ERROR: CallToolRequest unavailable.\"); return False\n\n # 加载配置\n config: Optional[MCPConfig] = None\n try:\n all_configs = load_config(CENTRAL_CONFIG_PATH)\n config = all_configs.get(SERVER_KEY_TO_TEST)\n if not config: print(f\"ERROR: Config key '{SERVER_KEY_TO_TEST}' not found in '{CENTRAL_CONFIG_PATH}'.\"); return False\n if not isinstance(config.connection, StdioConfig): print(\"ERROR: Config connection is not STDIO.\"); return False\n print(f\"Successfully loaded config for '{SERVER_KEY_TO_TEST}'.\")\n except Exception as e_load: print(f\"ERROR loading config: {e_load}\"); return False\n\n # 获取 LLM\n try: model = llm_manager.get_model(LLM_ID_FOR_TESTING); print(f\"Using LLM: {getattr(model, 'model_name', LLM_ID_FOR_TESTING)}\")\n except ValueError as e: print(f\"获取 LLM 出错: {e}.\"); return False\n\n test_success = False\n # 使用 MCPClient 连接 (它会根据 config 启动服务器)\n async with MCPClient(config) as client:\n print(\"\\nMCPClient context entered.\")\n if not client.session: print(\"ERROR: MCP session not established!\"); return False\n\n # --- 实例化我们定义的 MCPToolRunner ---\n try:\n print(f\"Instantiating MCPToolRunner for '{TOOL_NAME_TO_TEST}'...\")\n mcp_tool_instance = MCPToolRunner(\n client=client, # 注入 client\n name=TOOL_NAME_TO_TEST,\n description=TOOL_DESCRIPTION,\n args_schema=CORRECT_SCHEMA_FOR_TOOL\n )\n tools = [mcp_tool_instance]\n print(f\"Tool instance created successfully.\")\n except Exception as e_inst: print(f\"ERROR instantiating MCPToolRunner: {e_inst}\"); return False\n # ---\n\n # --- Agent 执行 ---\n agent = create_react_agent(model, tools) # Agent 使用这个标准工具\n query = \"Use the fetch tool to get the main content (first 2000 chars) from https://developer.mozilla.org/en-US/docs/Web/HTML\"\n print(f\"\\nRunning Agent Query...\")\n print(f\"Query: {query}\")\n print(\"--- NOTE: Enable LangSmith for detailed tracing! ---\")\n try:\n response = await asyncio.wait_for( agent.ainvoke({\"messages\": [{\"role\": \"user\",\"content\": query}]}), timeout=180.0 )\n print(f\"\\nAgent Final Response:\")\n if response and \"messages\" in response and response[\"messages\"]:\n response_content = response[\"messages\"][-1].content; print(response_content)\n # 检查是否成功获取内容且无报错\n contains_error = \"error\" in response_content.lower() or \"fail\" in response_content.lower() or \"issue\" in response_content.lower() or \"apologi\" in response_content.lower() or \"unable\" in response_content.lower() or \"tool error\" in response_content.lower()\n contains_expected = \"HTML\" in response_content\n\n if not contains_error and contains_expected:\n print(f\"\\n✅ Test PASS: Agent successfully used tool and got expected content.\")\n test_success = True\n else: print(f\"\\n❌ Test FAIL: Agent reported error or didn't get expected content.\"); test_success = False\n else: print(\"Agent returned no valid response.\"); test_success = False\n except asyncio.TimeoutError: print(f\"Agent execution timed out\"); test_success = False\n except Exception as e: print(f\"Agent execution failed: {e}\"); print(f\"Traceback:\\n{traceback.format_exc()}\"); test_success = False\n # ---\n\n # async with 会自动调用 client.close()\n print(f\"\\n--- Fetch Server Test Result: {'PASS' if test_success else 'FAIL'} ---\")\n return test_success\n\nasync def run_everything_test():\n \"\"\"运行 Everything MCP Server 测试 (使用 BaseTool 子类)\"\"\"\n print(f\"\\n=== Running Everything MCP Server Test (BaseTool Subclass Method) ===\")\n\n # 检查依赖和 Schema 定义\n if not ECHO_SCHEMA_AVAILABLE: print(\"ERROR: EchoInputSchema not available.\"); return False\n if not ADD_SCHEMA_AVAILABLE: print(\"ERROR: AddInputSchema not available.\"); return False\n if not CALL_TOOL_REQ_AVAILABLE: print(\"ERROR: CallToolRequest unavailable.\"); return False\n\n # 加载配置\n config: Optional[MCPConfig] = None\n try:\n all_configs = load_config(CENTRAL_CONFIG_PATH)\n config = all_configs.get(EVERYTHING_SERVER_KEY)\n if not config: print(f\"ERROR: Config key '{EVERYTHING_SERVER_KEY}' not found in '{CENTRAL_CONFIG_PATH}'.\"); return False\n if not isinstance(config.connection, StdioConfig): print(\"ERROR: Config connection is not STDIO.\"); return False\n print(f\"Successfully loaded config for '{EVERYTHING_SERVER_KEY}'.\")\n except Exception as e_load: print(f\"ERROR loading config: {e_load}\"); return False\n\n # 获取 LLM\n try: model = llm_manager.get_model(LLM_ID_FOR_TESTING); print(f\"Using LLM: {getattr(model, 'model_name', LLM_ID_FOR_TESTING)}\")\n except ValueError as e: print(f\"获取 LLM 出错: {e}.\"); return False\n\n test_success = False\n # 使用 MCPClient 连接 (它会根据 config 启动服务器)\n async with MCPClient(config) as client:\n print(\"\\nMCPClient context entered for Everything MCP.\")\n if not client.session: print(\"ERROR: MCP session not established!\"); return False\n\n # --- 实例化我们定义的 MCPToolRunner 用于 echo 工具 ---\n try:\n print(f\"Instantiating MCPToolRunner for '{EVERYTHING_ECHO_TOOL}'...\")\n echo_tool = MCPToolRunner(\n client=client, # 注入 client\n name=EVERYTHING_ECHO_TOOL,\n description=\"Echoes back the input message\",\n args_schema=EchoInputSchema\n )\n \n print(f\"Instantiating MCPToolRunner for '{EVERYTHING_ADD_TOOL}'...\")\n add_tool = MCPToolRunner(\n client=client, # 注入 client\n name=EVERYTHING_ADD_TOOL,\n description=\"Adds two numbers together\",\n args_schema=AddInputSchema\n )\n \n tools = [echo_tool, add_tool]\n print(f\"Tool instances created successfully.\")\n except Exception as e_inst: print(f\"ERROR instantiating MCPToolRunner: {e_inst}\"); return False\n # ---\n\n # --- Agent 执行 ---\n agent = create_react_agent(model, tools) # Agent 使用这些工具\n query = \"First, use the echo tool to echo back the message 'Hello from Everything MCP!'. Then, use the add tool to calculate 42 + 58.\"\n print(f\"\\nRunning Agent Query...\")\n print(f\"Query: {query}\")\n print(\"--- NOTE: Enable LangSmith for detailed tracing! ---\")\n try:\n response = await asyncio.wait_for(agent.ainvoke({\"messages\": [{\"role\": \"user\",\"content\": query}]}), timeout=180.0)\n print(f\"\\nAgent Final Response:\")\n if response and \"messages\" in response and response[\"messages\"]:\n response_content = response[\"messages\"][-1].content; print(response_content)\n # 检查是否成功获取内容且无报错\n contains_error = \"error\" in response_content.lower() or \"fail\" in response_content.lower() or \"issue\" in response_content.lower() or \"apologi\" in response_content.lower() or \"unable\" in response_content.lower() or \"tool error\" in response_content.lower()\n contains_echo = \"Hello from Everything MCP!\" in response_content\n contains_add = \"100\" in response_content\n\n if not contains_error and contains_echo and contains_add:\n print(f\"\\n✅ Test PASS: Agent successfully used both tools and got expected content.\")\n test_success = True\n else: \n print(f\"\\n❌ Test FAIL: Agent reported error or didn't get expected content.\")\n print(f\" - Contains error: {contains_error}\")\n print(f\" - Contains echo response: {contains_echo}\")\n print(f\" - Contains add result: {contains_add}\")\n test_success = False\n else: print(\"Agent returned no valid response.\"); test_success = False\n except asyncio.TimeoutError: print(f\"Agent execution timed out\"); test_success = False\n except Exception as e: print(f\"Agent execution failed: {e}\"); print(f\"Traceback:\\n{traceback.format_exc()}\"); test_success = False\n # ---\n\n # async with 会自动调用 client.close()\n print(f\"\\n--- Everything MCP Server Test Result: {'PASS' if test_success else 'FAIL'} ---\")\n return test_success\n\nasync def main():\n \"\"\"主函数 - 运行所有测试\"\"\"\n print(\"Starting MCP Integration Tests...\")\n \n # 运行 Fetch 测试\n fetch_success = await run_fetch_test()\n \n # 运行 Everything MCP 测试\n everything_success = await run_everything_test()\n \n print(\"\\n\" + \"=\"*20 + \" FINAL TEST SUMMARY \" + \"=\"*20);\n print(f\" Fetch Server Test: {'PASS' if fetch_success else 'FAIL'}\")\n print(f\" Everything MCP Test: {'PASS' if everything_success else 'FAIL'}\")\n print(\"=\"*20 + \" MCP Integration Test Finished \" + \"=\"*20)\n\nif __name__ == \"__main__\":\n # 简化依赖检查\n print(\"--- Dependency Check ---\")\n deps_ok = True\n try: import mcp; print(\"mcp available: True\")\n except ImportError: print(\"mcp available: False\"); deps_ok = False\n if CALL_TOOL_REQ_AVAILABLE: print(\"CallToolRequest available: True\")\n else: print(\"CallToolRequest available: False\"); deps_ok = False # 需要它\n try: import langgraph; print(\"langgraph available: True\")\n except ImportError: print(\"langgraph available: False\"); deps_ok = False\n try: import langchain_openai; print(\"langchain_openai available: True\")\n except ImportError: print(\"langchain_openai available: False\"); deps_ok = False\n try: import dotenv; print(\"dotenv available: True\")\n except ImportError: print(\"dotenv available: False\"); deps_ok = False\n try: import pydantic; print(\"pydantic available: True\")\n except ImportError: print(\"pydantic available: False\"); deps_ok = False\n try: from core.mcp.client import MCPClient; print(\"MCPClient available: True\")\n except ImportError: print(\"MCPClient available: False\"); deps_ok = False\n try: from core.mcp.config_loader import load_config; print(\"config_loader available: True\")\n except ImportError: print(\"config_loader available: False\"); deps_ok = False\n if not FETCH_SCHEMA_AVAILABLE: print(\"FetchInputSchema available: False\"); deps_ok=False\n else: print(\"FetchInputSchema available: True\")\n if not ECHO_SCHEMA_AVAILABLE: print(\"EchoInputSchema available: False\"); deps_ok=False\n else: print(\"EchoInputSchema available: True\")\n if not ADD_SCHEMA_AVAILABLE: print(\"AddInputSchema available: False\"); deps_ok=False\n else: print(\"AddInputSchema available: True\")\n print(f\"------------------------\")\n\n if not deps_ok:\n print(\"CRITICAL ERROR: Necessary dependencies missing.\")\n sys.exit(1)\n\n asyncio.run(main())"
},
{
"path": "examples/16_google_a2a/README.md",
"content": "# LangGraph Agent 与 A2A 协议集成框架\n\n## 概述\n\n本项目提供了一个将 **LangGraph Agent**（特别是基于 ReAct 模式并能调用工具的 Agent）与 **A2A (Agent-to-Agent) 协议** 相集成的框架和示例。目标是展示如何将一个用 LangGraph 构建的复杂 Agent 能力，通过标准化的 A2A 接口暴露给外部客户端或其他 Agent。\n\n此框架的核心在于 `AgentTaskManager`，它充当了 A2A 协议层与具体 Agent 实现之间的桥梁。项目包含了一个完整的端到端示例，其中 `CurrencyAgent`（使用 `create_react_agent` 构建，并带有计算器和搜索工具）通过 `A2AServer` 提供服务，并提供了两个不同的客户端示例 (`client_example.py` 和 `currency_agent_test.py`) 来演示如何与之交互。\n\n关键技术栈包括：\n* **A2A 协议:** 定义交互规范。\n* **LangGraph:** 用于构建具备状态管理和工具调用能力的 Agent。\n* **`create_react_agent`:** LangGraph 提供的预构建 ReAct Agent 实现（作为示例）。\n* **Pydantic:** 用于定义和验证 A2A 协议中的数据结构 (`core/a2a/types.py`)。\n* **Starlette/Uvicorn:** 作为底层 Web 框架运行 A2A 服务器 (`core/a2a/server/server.py`)。\n* **OpenAI API:** 作为 LangGraph Agent 使用的后端大语言模型（可替换）。\n\n## 特性\n\n* **A2A 协议兼容:** 提供符合 A2A 规范的服务端点 (`/.well-known/agent.json` 和主任务端点)。\n* **LangGraph Agent 集成:** 可将任意（满足特定接口要求的）LangGraph Agent 作为 A2A 服务的核心处理逻辑。\n* **工具使用:** 集成的 Agent 能够根据需要调用外部工具（示例中为计算器和搜索）。\n* **同步任务处理:** 支持客户端发送任务并等待最终结果。\n* **流式基础:** 包含了处理流式请求和响应的框架（Agent 端流式逻辑需开发者实现）。\n* **类型安全:** 使用 Pydantic 进行严格的数据校验。\n* **环境配置:** 支持通过 `.env` 文件配置 API 密钥等敏感信息。\n* **客户端示例:** 提供了基础和场景化的客户端示例代码。\n\n## 目录结构\n\n```\n.\n├── core/ # 核心 A2A 协议实现\n│ └── a2a/\n│ ├── client/\n│ │ └── client.py # A2AClient 客户端库实现\n│ ├── server/\n│ │ ├── server.py # A2AServer HTTP 服务器实现\n│ │ └── task_manager.py # TaskManager 基础接口 (被 AgentTaskManager 使用)\n│ ├── agent_task_manager.py # AgentTaskManager 实现 (连接 A2A 与 LangGraph)\n│ └── types.py # A2A 协议的 Pydantic 模型定义\n├── examples/ # 示例代码\n│ └── a2a/\n│ ├── langgraph_integration.py # 服务端设置和示例 LangGraph Agent (CurrencyAgent) 定义\n│ ├── client_example.py # 基础 A2A 客户端使用示例脚本\n│ └── currency_agent_test.py # 场景化 A2A 客户端测试脚本\n├── .env # 存储环境变量 (例如 OPENAI_API_KEY) - *需要自行创建*\n├── requirements.txt # Python 依赖项列表 (假设存在)\n└── README.md # 本文档\n```\n\n## 核心组件说明\n\n* **`core/a2a/types.py`:** 定义所有 A2A 数据结构，是协议的基础和校验依据。\n* **`core/a2a/server/server.py` (`A2AServer`):** 基于 Starlette 的 HTTP 服务器，处理 A2A JSON-RPC 请求路由，将请求交给 `AgentTaskManager`。通过 `.start()` 方法启动。\n* **`core/a2a/agent_task_manager.py` (`AgentTaskManager`):** **核心适配器**。连接 A2A 层和 Agent 层。它接收来自 `A2AServer` 的请求，管理任务状态，并调用注入的 Agent 实例的 `invoke` 或 `stream` 方法。\n* **`examples/a2a/langgraph_integration.py`:** 包含 `CurrencyAgent` (使用 `create_react_agent` 的示例 Agent) 的定义，以及如何配置和启动 `A2AServer` 来运行这个 Agent 的完整脚本。\n* **`core/a2a/client/client.py` (`A2AClient`):** 基础 A2A 客户端库。\n* **`examples/a2a/client_example.py`:** 一个简单的脚本，演示如何使用 `A2AClient` 发送基本请求。\n* **`examples/a2a/currency_agent_test.py`:** 一个更复杂的客户端脚本，包含多个测试场景，用于测试服务器端 Agent 的不同交互模式。\n\n## 先决条件\n\n* Python (推荐 3.10 或更高版本)\n* `pip` (Python 包安装器)\n* 虚拟环境 (强烈推荐)\n* 大语言模型 API Key (例如 OpenAI API Key)\n\n## 安装与设置\n\n1. **克隆仓库:**\n ```bash\n git clone \n cd \n ```\n2. **创建并激活虚拟环境:**\n ```bash\n uv venv\n source .venv/bin/activate\n ```\n3. **安装依赖项:**\n ```bash\n uv sync\n ```\n4. **设置环境变量:**\n * 在项目根目录下创建 `.env` 文件。\n * 添加所需的 API Key，例如：\n ```dotenv\n OPENAI_API_KEY=\"sk-...\"\n ```\n\n## 运行示例\n\n1. **启动 A2A 服务器:**\n * 在终端中，激活虚拟环境后运行：\n ```bash\n python -m examples.a2a.langgraph_integration\n ```\n * 服务器将在 `http://127.0.0.1:8000` 启动并监听。\n\n2. **运行 A2A 客户端:**\n * 打开**新的**终端，激活虚拟环境。\n * 你可以选择运行任一客户端示例：\n * **基础示例:**\n ```bash\n python -m examples.a2a.client_example\n ```\n * **场景化测试:**\n ```bash\n python -m examples.a2a.currency_agent_test\n ```\n\n3. **预期输出:**\n * **服务器终端**会显示接收请求、调用 LLM 和工具（如果被触发）的日志。\n * **客户端终端**会显示发送任务、轮询状态（对于同步任务）、接收结果或（模拟的）流式事件的输出。`currency_agent_test.py` 会按场景输出结果。\n\n---\n\n## **重要：集成新的 LangGraph Agent 指南**\n\n如果你创建了一个新的基于 LangGraph 的 Agent，并希望将其接入到这个 A2A 框架中，你需要遵循以下步骤和约定：\n\n### 1. Agent 类必须实现的接口\n\n你的新 Agent 类（例如 `MyNewAgent`）需要被 `AgentTaskManager` 调用。为此，它**必须**实现以下方法和属性：\n\n* **`__init__(self, llm, ...)`:**\n * 构造函数，用于初始化 Agent 所需的资源，例如 LLM 实例、工具列表等。\n * **关键:** 在这里构建或获取你的 LangGraph **Runnable** 实例（例如通过 `create_react_agent` 或手动构建 `StateGraph().compile()`），并将其存储为类的成员（例如 `self.agent_runnable`）。\n\n* **`invoke(self, query: str, session_id: Optional[str] = None) -> str:`**\n * 处理 A2A 的**同步** `tasks/send` 请求。\n * 接收从 `AgentTaskManager` 传递过来的纯文本用户查询 `query` 和可选的 `session_id`。\n * **内部逻辑:**\n * 将 `query` 包装成你的 LangGraph Runnable 所需的输入格式。对于基于 `create_react_agent` 或类似使用消息列表的 Agent，通常是 `{\"messages\": [(\"user\", query)]}`。如果需要 `session_id`，也应包含在内。\n * 调用 LangGraph Runnable 的 `.invoke()` 方法，传入构造好的输入字典。\n * 处理 Runnable 返回的结果字典。对于 ReAct Agent，最终的文本答案通常位于结果字典内 `messages` 列表的最后一条消息的内容中。你需要编写逻辑来提取这个最终答案。\n * **返回值:** **必须**返回一个包含最终答案的**字符串**。\n\n* **`stream(self, query: str, session_id: Optional[str] = None) -> AsyncIterable[Dict[str, Any]]:`**\n * 处理 A2A 的**流式** `tasks/sendSubscribe` 请求。\n * 接收 `query` 和 `session_id`。\n * **必须**是一个**异步生成器** (`async def` 包含 `yield`)。\n * **内部逻辑:**\n * 准备 LangGraph Runnable 流式调用所需的输入（通常与 `invoke` 类似，例如 `{\"messages\": [(\"user\", query)]}`）。\n * 调用 LangGraph Runnable 的流式方法，例如 `self.agent_runnable.astream(...)` 或 `self.agent_runnable.astream_log(...)`。\n * 使用 `async for chunk in ...:` 迭代 LangGraph Runnable 返回的流式数据块 (`chunk`)。\n * **解析 `chunk`**: LangGraph 流式输出的 `chunk` 格式取决于你调用的方法（`astream` vs `astream_log`）和图的结构。你需要解析这些 `chunk`（可能是状态变更、日志补丁等）来获取有意义的中间或最终内容。\n * **`yield` 符合格式的字典**: 对于每个希望发送给客户端的更新，你需要 `yield` 一个字典。这个字典**必须**包含以下键（供 `AgentTaskManager._run_streaming_agent` 使用）：\n * `\"content\"`: `str` - 当前步骤生成的文本内容。\n * `\"is_task_complete\"`: `bool` - 指示这是否是任务的最终产物/结束信号。\n * `\"require_user_input\"`: `bool` - 指示任务是否暂停并需要用户输入。\n * **返回值:** 返回一个异步可迭代对象（由 `async def` + `yield` 自动创建）。\n\n* **`SUPPORTED_CONTENT_TYPES: List[str]` (类属性):**\n * 一个包含 Agent 支持的输出内容类型的列表。对于主要处理文本的 Agent，通常是 `[\"text\"]`。`AgentTaskManager` 会用它来验证客户端请求的 `acceptedOutputModes`。\n\n### 2. `AgentState` 的一致性\n\n如果你手动构建 LangGraph 图，你定义的 `AgentState`（传递给 `StateGraph`）需要与你的 `invoke` 和 `stream` 方法处理输入/输出的方式保持一致。特别是，如果你依赖 `messages` 列表来管理对话历史或传递输入/输出，`AgentState` 中需要正确定义它。\n\n### 3. 集成步骤\n\n1. **创建 Agent 类:**\n * 在你的项目中创建一个新的 Python 文件（例如 `my_new_agent.py`）。\n * 定义你的 Agent 类（例如 `MyNewAgent`），确保它实现了上面描述的 `__init__`, `invoke`, `stream` 方法和 `SUPPORTED_CONTENT_TYPES` 属性。\n * 在 `__init__` 中构建或加载你的 LangGraph Runnable。\n\n2. **修改服务器启动脚本 (例如 `examples/a2a/langgraph_integration.py`):**\n * **导入**你的新 Agent 类：`from my_new_agent import MyNewAgent`。\n * **实例化**你的新 Agent：`my_agent = MyNewAgent(llm)` (确保传递了所需的依赖，如 `llm`)。\n * **更新 `AgentCard`**: 修改 `name`, `description` 和 `skills` 列表以反映新 Agent 的信息。确保 `AgentSkill` 具有唯一的 `id` 和正确的 `name`。\n * **实例化 `AgentTaskManager`**: 使用你的新 Agent 实例：`task_manager = AgentTaskManager(my_agent)`。\n * **实例化 `A2AServer`**: 使用更新后的 `agent_card` 和 `task_manager`。\n\n3. **运行服务器:**\n * 启动修改后的服务器脚本：`python -m examples.a2a.your_server_script`。\n\n4. **测试:**\n * 使用 `client_example.py` 或 `currency_agent_test.py`（可能需要修改发送的查询或 `metadata` 中的 `skill_name`）来向新启动的服务器发送请求，验证你的新 Agent 是否能通过 A2A 协议正常工作。\n\n### 示例 Agent 骨架\n\n```python\n# my_new_agent.py\nimport logging\nfrom typing import List, Optional, AsyncIterable, Dict, Any, Tuple\nfrom langchain_core.language_models import BaseChatModel # 示例 LLM 类型\nfrom langgraph.graph.state import StateGraph # 如果手动构建图\n# from langgraph.prebuilt import create_some_agent # 如果使用预构建\nfrom typing import TypedDict\n\nlogger = logging.getLogger(__name__)\n\n# 1. 定义你 Agent 使用的 State (如果需要)\nclass MyAgentState(TypedDict):\n messages: List[Tuple[str, str]]\n # ... 其他状态字段\n\nclass MyNewAgent:\n SUPPORTED_CONTENT_TYPES: List[str] = [\"text\"]\n\n def __init__(self, llm: BaseChatModel):\n self.llm = llm\n # TODO: 在这里构建或加载你的 LangGraph Runnable\n # 例如: self.agent_runnable = self._build_my_graph()\n # 或者: self.agent_runnable = create_some_agent(llm, tools)\n self.agent_runnable = self._get_placeholder_runnable() # 示例\n logger.info(\"MyNewAgent initialized.\")\n\n def _get_placeholder_runnable(self):\n # 这是一个模拟的 Runnable，你需要替换成真实的 LangGraph Runnable\n class PlaceholderRunnable:\n def invoke(self, input_dict):\n logger.info(f\"PlaceholderRunnable received invoke: {input_dict}\")\n query = input_dict.get(\"messages\", [(\"\", \"\")])[-1][1]\n return {\"messages\": [(\"assistant\", f\"模拟回应 '{query}'\")]}\n async def astream(self, input_dict):\n logger.info(f\"PlaceholderRunnable received astream: {input_dict}\")\n query = input_dict.get(\"messages\", [(\"\", \"\")])[-1][1]\n yield {\"messages\": [(\"assistant\", f\"模拟流式回应1 '{query}' ...\")]}\n await asyncio.sleep(0.5)\n yield {\"messages\": [(\"assistant\", f\"模拟流式回应2 '{query}' 完毕。\")]}\n return PlaceholderRunnable()\n\n # def _build_my_graph(self):\n # # 如果你手动构建图，在这里实现\n # # workflow = StateGraph(MyAgentState)\n # # ... add nodes, edges ...\n # # return workflow.compile()\n # pass\n\n def invoke(self, query: str, session_id: Optional[str] = None) -> str:\n logger.debug(f\"[MyNewAgent.invoke] query: '{query}', session_id: '{session_id}'\")\n # 1. 准备输入\n invoke_input = {\"messages\": [(\"user\", query)]}\n # 2. 调用 Runnable\n try:\n result = self.agent_runnable.invoke(invoke_input)\n logger.debug(f\"[MyNewAgent.invoke] Runnable result: {result}\")\n # 3. 解析结果\n final_output = \"错误：未能解析 Agent 响应。\"\n if isinstance(result, dict) and isinstance(result.get(\"messages\"), list) and result[\"messages\"]:\n last_message = result[\"messages\"][-1]\n if isinstance(last_message, tuple) and len(last_message) == 2:\n final_output = last_message[1]\n elif hasattr(last_message, 'content'):\n final_output = last_message.content\n return str(final_output)\n except Exception as e:\n logger.error(f\"[MyNewAgent.invoke] Error: {e}\", exc_info=True)\n raise # 重新抛出异常，让 TaskManager 处理\n\n async def stream(self, query: str, session_id: Optional[str] = None) -> AsyncIterable[Dict[str, Any]]:\n logger.debug(f\"[MyNewAgent.stream] query: '{query}', session_id: '{session_id}'\")\n # 1. 准备输入\n stream_input = {\"messages\": [(\"user\", query)]}\n # 2. 调用 Runnable 的流式方法\n try:\n # 使用 astream 或 astream_log\n async for chunk in self.agent_runnable.astream(stream_input):\n logger.debug(f\"[MyNewAgent.stream] Received chunk: {chunk}\")\n # 3. 解析 chunk 并 yield 符合格式的字典\n # 这里的解析逻辑高度依赖于你的图和使用的流式方法\n # 你需要根据实际的 chunk 内容提取 content, is_task_complete, require_user_input\n # --- 这是一个 **高度简化** 的示例解析 ---\n content_to_yield = \"\"\n is_complete = False # 你需要根据 chunk 判断任务是否真的结束\n is_input_required = False # 你需要根据 chunk 判断是否需要输入\n\n # 尝试从 chunk 中提取 'messages' 的最新内容作为 content\n if isinstance(chunk, dict) and isinstance(chunk.get(\"messages\"), list) and chunk[\"messages\"]:\n last_message = chunk[\"messages\"][-1]\n if isinstance(last_message, tuple) and len(last_message) == 2:\n content_to_yield = last_message[1]\n elif hasattr(last_message, 'content'):\n content_to_yield = last_message.content\n\n if content_to_yield: # 只在有内容时 yield\n # 在实际应用中，你需要更复杂的逻辑判断 is_task_complete\n # 例如，检查 LangGraph 图是否到达了 END 节点，或者某个特定的最终节点状态\n # is_complete = ???\n yield {\n \"content\": content_to_yield,\n \"is_task_complete\": is_complete, # 需要正确设置\n \"require_user_input\": is_input_required # 需要正确设置\n }\n # --- 简化示例结束 ---\n\n # **重要**: 在循环结束后，如果任务确实完成了，需要再 yield 一个最终状态\n # (除非上面的循环中最后一个 yield 的 is_task_complete 已经是 True)\n # 例如:\n # final_result = await self.agent_runnable.ainvoke(stream_input) # 可能需要再调用一次 invoke 获取最终确认状态\n # final_text = ... # 解析最终文本\n # yield {\"content\": final_text, \"is_task_complete\": True, \"require_user_input\": False}\n\n except Exception as e:\n logger.error(f\"[MyNewAgent.stream] Error: {e}\", exc_info=True)\n # 在流中抛出异常可能会中断 SSE 连接，或者你可以 yield 一个错误信息\n yield {\n \"content\": f\"处理流式请求时出错: {e}\",\n \"is_task_complete\": True, # 标记任务失败并结束\n \"require_user_input\": False\n }\n```\n\n## 当前状态与限制\n\n* 同步任务执行，包括 LangGraph Agent 调用 LLM 和工具，已成功实现并验证。\n* A2A 协议的服务端和客户端基础结构已建立。\n* **Agent 端的流式处理 (`CurrencyAgent.stream`) 目前是模拟的**，并未真正调用 LangGraph 的流式接口。真实的流式更新尚未实现。\n* 当前 Agent 实现 (`CurrencyAgent`) 不支持需要跨请求保持状态的多轮对话澄清。\n* 错误处理可以进一步增强。\n* 任务存储仅在内存中 (`InMemoryTaskManager`)。\n\n## 未来方向\n\n* **实现真实流式输出:** 按照上述指南，在 Agent 类中实现 `stream` 方法，调用 LangGraph 的 `astream` 或 `astream_log`，并正确解析和 `yield` A2A 所需格式的字典。\n* **支持多轮对话:** 修改 `AgentState` 以包含可累加的消息历史 (例如使用 `Annotated[List[BaseMessage], operator.add]`)，并调整 Agent 的 `invoke` 和 `stream` 方法以处理和利用这个历史记录。可能还需要 Agent 能返回 `input-required` 状态。\n* **增强错误处理:** 为网络问题、Agent 执行错误、工具调用失败、类型验证错误等提供更详细、用户友好的错误报告。\n* **持久化任务存储:** 替换 `InMemoryTaskManager`。\n* **配置管理:** 外部化配置。\n* **多技能支持:** 添加路由逻辑。\n"
},
{
"path": "examples/16_google_a2a/__init__.py",
"content": "# examples/a2a/__init__.py\n\n\"\"\"\nA2A协议与LangGraph集成示例\n\n本目录包含了A2A协议与LangGraph Agent集成的示例和文档。\n\"\"\""
},
{
"path": "examples/16_google_a2a/agent_task_manager_test.py",
"content": "# examples/a2a/agent_task_manager_test.py\n\nimport os\nimport sys\nimport asyncio\nimport logging\nfrom typing import TypedDict, Any, List, Optional,Tuple\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))\n\n# 导入环境变量\nfrom dotenv import load_dotenv\nload_dotenv()\n\n# 导入A2A相关组件\nfrom core.a2a.types import (\n TaskState, TaskStatus, Task, Artifact, Message,\n SendTaskRequest, SendTaskResponse, SendTaskStreamingRequest,\n TaskSendParams, JSONRPCResponse\n)\nfrom core.a2a.agent_task_manager import AgentTaskManager\n\n# 导入LangChain和LLM相关组件\nfrom langchain_core.tools import tool\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.graph import END, StateGraph\nfrom langgraph.prebuilt import create_react_agent\n\n# 配置日志\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\n# 定义一个简单的工具\n@tool\ndef search(query: str) -> str:\n \"\"\"搜索互联网获取信息\"\"\"\n return f\"这是关于 '{query}' 的搜索结果。\"\n\n@tool\ndef calculator(expression: str) -> str:\n \"\"\"计算数学表达式\"\"\"\n try:\n result = eval(expression)\n return f\"计算结果: {result}\"\n except Exception as e:\n return f\"计算错误: {e}\"\n\n# 定义一个简单的LangGraph Agent\n\nclass AgentState(TypedDict):\n messages: List[Tuple[str, str]]\n session_id: Optional[str] # 保留 session_id\n\nclass TestAgent:\n \"\"\"测试用Agent\"\"\"\n \n # 支持的内容类型\n SUPPORTED_CONTENT_TYPES = [\"text\"]\n \n def __init__(self, llm=None):\n if llm is None:\n try:\n llm = ChatOpenAI(model=\"gpt-4o-mini\")\n except Exception as e:\n print(f\"警告: 无法创建OpenAI LLM ({e})，使用模拟模式\")\n from langchain.llms.fake import FakeListLLM\n llm = FakeListLLM(responses=[\"这是一个模拟的LLM响应\"])\n \n self.tools = [search, calculator]\n self.agent = create_react_agent(llm, self.tools)\n self.graph = self._build_graph()\n \n def _build_graph(self):\n \"\"\"构建Agent的工作流图\"\"\"\n workflow = StateGraph(AgentState)\n workflow.add_node(\"agent\", self.agent)\n workflow.set_entry_point(\"agent\")\n workflow.add_edge(\"agent\", END)\n return workflow.compile()\n \n def invoke(self, query: str, session_id: str = None) -> str:\n \"\"\"同步调用Agent\"\"\"\n result = self.graph.invoke({\"input\": query, \"session_id\": session_id})\n return result[\"output\"]\n \n async def stream(self, query: str, session_id: str = None):\n \"\"\"流式调用Agent\"\"\"\n # 模拟流式输出\n chunks = [\n \"正在处理您的请求...\",\n \"正在搜索相关信息...\",\n \"找到了一些结果，正在整理...\",\n f\"关于 '{query}' 的信息如下：这是一个模拟的流式响应。\"\n ]\n \n for i, chunk in enumerate(chunks):\n is_last = i == len(chunks) - 1\n yield {\n \"content\": chunk,\n \"is_task_complete\": is_last,\n \"require_user_input\": False\n }\n await asyncio.sleep(0.5) # 模拟延迟\n\n# 测试AgentTaskManager的同步任务处理\nasync def test_sync_task():\n print(\"\\n=== 测试同步任务处理 ===\\n\")\n \n # 创建Agent和AgentTaskManager\n agent = TestAgent()\n task_manager = AgentTaskManager(agent)\n \n # 创建任务请求\n task_id = \"test_sync_task_1\"\n session_id = \"test_session_1\"\n content = [{\"type\": \"text\", \"text\": \"计算 123 + 456 的结果\"}]\n \n task_params = TaskSendParams(\n id=task_id,\n sessionId=session_id,\n message=Message(role=\"user\", parts=content),\n acceptedOutputModes=[\"text\"],\n historyLength=10\n )\n \n request = SendTaskRequest(id=\"req1\", params=task_params)\n \n # 发送任务\n response = await task_manager.on_send_task(request)\n \n # 打印结果\n print(f\"任务ID: {task_id}\")\n print(f\"响应类型: {type(response)}\")\n \n if hasattr(response, \"error\") and response.error:\n print(f\"错误: {response.error}\")\n else:\n print(\"任务成功完成\")\n \n # 获取任务\n task = task_manager.tasks.get(task_id)\n if task:\n print(f\"任务状态: {task.status.state}\")\n if task.artifacts:\n for artifact in task.artifacts:\n for part in artifact.parts:\n if part.get(\"type\") == \"text\":\n print(f\"任务结果: {part.get('text')}\")\n\n# 测试AgentTaskManager的流式任务处理\nasync def test_streaming_task():\n print(\"\\n=== 测试流式任务处理 ===\\n\")\n \n # 创建Agent和AgentTaskManager\n agent = TestAgent()\n task_manager = AgentTaskManager(agent)\n \n # 创建任务请求\n task_id = \"test_stream_task_1\"\n session_id = \"test_session_1\"\n content = [{\"type\": \"text\", \"text\": \"搜索关于人工智能的信息\"}]\n \n task_params = TaskSendParams(\n id=task_id,\n sessionId=session_id,\n message=Message(role=\"user\", parts=content),\n acceptedOutputModes=[\"text\"],\n historyLength=10\n )\n \n request = SendTaskStreamingRequest(id=\"req2\", params=task_params)\n \n # 发送流式任务\n response_generator = await task_manager.on_send_task_subscribe(request)\n \n # 检查响应类型\n if isinstance(response_generator, JSONRPCResponse):\n print(f\"错误: {response_generator.error}\")\n return\n \n # 处理流式响应\n print(\"开始接收流式响应:\")\n async for response in response_generator:\n if hasattr(response, \"error\") and response.error:\n print(f\"流式响应错误: {response.error}\")\n else:\n result = response.result\n if hasattr(result, \"status\") and result.status and result.status.message:\n for part in result.status.message.parts:\n # --- 修改开始 ---\n # 直接访问对象的属性 type 和 text\n if hasattr(part, 'type') and part.type == \"text\":\n text_content = getattr(part, 'text', '') # 安全获取 text\n print(f\"流式更新: {text_content}\")\n # --- 修改结束 ---\n\n if hasattr(result, \"artifact\") and result.artifact:\n for part in result.artifact.parts:\n # --- 修改开始 ---\n # 直接访问对象的属性 type 和 text\n if hasattr(part, 'type') and part.type == \"text\":\n text_content = getattr(part, 'text', '') # 安全获取 text\n print(f\"流式结果: {text_content}\")\n # --- 修改结束 ---\n\n if hasattr(result, \"final\") and result.final:\n print(\"流式响应结束\")\n\n# 主函数\nasync def main():\n print(\"=== AgentTaskManager 测试 ===\\n\")\n \n # 测试同步任务\n await test_sync_task()\n \n # 测试流式任务\n await test_streaming_task()\n\n# 运行测试\nif __name__ == \"__main__\":\n asyncio.run(main())"
},
{
"path": "examples/16_google_a2a/client_example.py",
"content": "# examples/a2a/client_example.py\n\nimport os\nimport sys\nimport asyncio\nimport json\nimport logging # 添加 logging\nfrom typing import Dict, Any, List, Optional\nfrom uuid import uuid4 # 用于生成示例 Task ID\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))\n\n# 导入环境变量\nfrom dotenv import load_dotenv\nload_dotenv()\n\n# 导入A2A客户端和类型\nfrom core.a2a.client.client import A2AClient\n# 导入 Message 和 TextPart 以构建请求，导入响应类型以进行类型提示\nfrom core.a2a.types import (\n Part, TextPart, Message, TaskState, # 添加 TaskState\n SendTaskResponse, GetTaskResponse, SendTaskStreamingResponse, Task, # 添加 Task\n JSONRPCError # 添加 JSONRPCError\n)\n\n# 配置日志\nlogging.basicConfig(level=logging.INFO) # 可以改为 DEBUG 获取更详细客户端日志\nlogger = logging.getLogger(__name__)\n\n# 示例: 使用A2A客户端连接到A2A服务器\nasync def run_a2a_client():\n print(\"\\n=== 运行A2A客户端示例 ===\\n\")\n\n # 创建A2A客户端\n client = A2AClient(url=\"http://127.0.0.1:8000\") # 指向你的服务器地址\n\n # 发送同步任务\n await send_sync_task(client)\n\n # 发送流式任务\n await send_streaming_task(client)\n\n# --- 修正发送同步任务 ---\nasync def send_sync_task(client: A2AClient):\n print(\"\\n=== 发送同步任务 ===\\n\")\n query = \"请计算 123 + 456 的结果\"\n task_id = \"client_sync_\" + uuid4().hex # 生成一个唯一的任务 ID\n try:\n # 1. 构建 Message 对象\n message = Message(role=\"user\", parts=[TextPart(text=query)])\n\n # 2. 构建 TaskSendParams 对应的 payload 字典 (添加 id)\n payload_dict = {\n \"id\": task_id, # --- 添加必需的 id 字段 ---\n \"sessionId\": \"client_session_sync_1\",\n \"message\": message.model_dump(),\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"}\n }\n logger.debug(f\"Sending sync task with payload: {payload_dict}\")\n\n # 3. 调用 send_task，传入 payload 字典\n response: SendTaskResponse = await client.send_task(payload=payload_dict)\n logger.debug(f\"Send task response: {response}\")\n\n # 4. 处理响应\n if response.error:\n # 类型提示帮助访问属性\n error: JSONRPCError = response.error\n print(f\"发送任务时出错: Code={error.code}, Message={error.message}\")\n return\n # SendTaskResponse 的 result 是 Task 对象或 None\n if not response.result:\n print(f\"发送任务成功，但响应中未包含任务详情: {response}\")\n # 我们可以继续使用我们发送的 task_id 来查询状态\n elif response.result.id != task_id:\n # 理论上服务器应该使用或确认客户端提供的 ID\n logger.warning(f\"服务器返回的任务ID '{response.result.id}' 与客户端发送的ID '{task_id}' 不匹配。\")\n task_id = response.result.id # 以服务器返回的为准（如果存在）\n\n\n print(f\"任务已发送，ID: {task_id}\")\n\n # --- 轮询等待任务完成 ---\n print(\"等待任务完成...\")\n task_result: Optional[Task] = None # 用于存储最终的任务对象\n for attempt in range(10): # 最多尝试 10 次\n await asyncio.sleep(2) # 等待 2 秒\n\n # 5. 构建 get_task 的 payload\n get_payload = {\"id\": task_id}\n logger.debug(f\"Getting task with payload: {get_payload} (Attempt {attempt+1})\")\n\n # 6. 获取任务结果 (传入 payload 字典)\n get_response: GetTaskResponse = await client.get_task(payload=get_payload)\n logger.debug(f\"Get task response: {get_response}\")\n\n if get_response.error:\n error: JSONRPCError = get_response.error\n print(f\"获取任务时出错: Code={error.code}, Message={error.message}\")\n return # 出错则停止轮询\n if not get_response.result:\n print(f\"获取任务成功，但未收到任务详情: {get_response}\")\n continue # 继续轮询\n\n task_result = get_response.result # 获取任务对象\n print(f\" 当前任务状态: {task_result.status.state}\")\n # 检查任务是否完成或失败\n if task_result.status.state in [TaskState.COMPLETED, TaskState.FAILED, TaskState.CANCELED, TaskState.INPUT_REQUIRED]:\n break\n else:\n print(\"任务在限定时间内未完成。\")\n return\n\n # 7. 处理最终任务结果 (使用属性访问)\n if task_result.status.state == TaskState.COMPLETED and task_result.artifacts:\n print(\"任务成功完成。结果:\")\n for artifact in task_result.artifacts:\n if artifact.parts:\n for part in artifact.parts:\n if isinstance(part, TextPart):\n print(f\" - {part.text}\")\n elif task_result.status.state == TaskState.FAILED:\n error_msg = \"未知错误\"\n if task_result.status.message and task_result.status.message.parts:\n # 假设错误信息在第一个 TextPart\n if isinstance(task_result.status.message.parts[0], TextPart):\n error_msg = task_result.status.message.parts[0].text\n print(f\"任务失败: {error_msg}\")\n else:\n print(f\"任务最终状态为: {task_result.status.state}\")\n\n except Exception as e:\n logger.error(f\"发送或处理同步任务时发生异常: {e}\", exc_info=True)\n print(f\"发送同步任务失败: {e}\")\n\n# --- 修正发送流式任务 ---\nasync def send_streaming_task(client: A2AClient):\n print(\"\\n=== 发送流式任务 ===\\n\")\n query = \"请搜索关于人工智能的最新进展\"\n task_id = \"client_stream_\" + uuid4().hex # 为流式任务生成 ID\n try:\n # 1. 构建 Message 对象\n message = Message(role=\"user\", parts=[TextPart(text=query)])\n\n # 2. 构建 TaskSendParams 对应的 payload 字典 (添加 id)\n payload_dict = {\n \"id\": task_id, # --- 添加必需的 id 字段 ---\n \"sessionId\": \"client_session_stream_1\",\n \"message\": message.model_dump(),\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"}\n }\n logger.debug(f\"Sending streaming task with payload: {payload_dict}\")\n print(f\"任务已发送，ID: {task_id}\") # 流式任务 ID 在发送时就已知\n\n # 3. 调用 send_task_streaming (不再使用 await)\n # 它返回一个异步生成器\n event_stream_generator = client.send_task_streaming(payload=payload_dict)\n\n # 4. 使用 async for 处理流式事件\n print(\"开始接收流式响应:\")\n async for event_response in event_stream_generator: # 正确迭代异步生成器\n logger.debug(f\"Received stream event: {event_response}\")\n\n # 检查整个响应是否有错误\n if event_response.error:\n error: JSONRPCError = event_response.error\n print(f\"流式传输中出错: Code={error.code}, Message={error.message}\")\n continue # 或 break\n\n # 获取事件具体内容\n event = event_response.result\n if not event:\n logger.warning(\"Received stream response with empty result.\")\n continue\n\n # 处理状态更新事件中的消息部分\n if hasattr(event, \"status\") and event.status and event.status.message:\n if event.status.message.parts:\n for part in event.status.message.parts:\n if isinstance(part, TextPart):\n print(f\" 流式更新: {part.text}\")\n\n # 处理制品更新事件\n if hasattr(event, \"artifact\") and event.artifact:\n print(\" 收到 Artifact:\")\n if event.artifact.parts:\n for part in event.artifact.parts:\n if isinstance(part, TextPart):\n print(f\" 流式结果 (TextPart): {part.text}\")\n\n # 检查流结束标志\n if hasattr(event, \"final\") and event.final:\n print(\"流式响应结束标志收到。\")\n\n print(\"流式任务处理完成。\")\n\n except Exception as e:\n logger.error(f\"发送或处理流式任务时发生异常: {e}\", exc_info=True)\n print(f\"发送流式任务失败: {e}\")\n\n# 主函数\nif __name__ == \"__main__\":\n # 使用 asyncio.run 运行顶层异步函数\n asyncio.run(run_a2a_client())"
},
{
"path": "examples/16_google_a2a/currency_agent_test.py",
"content": "# examples/a2a/currency_agent_test.py\n\nimport os\nimport sys\nimport asyncio\nimport json\nimport logging\nfrom typing import Dict, Any, List, Optional\nfrom uuid import uuid4 # Import uuid\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))\n\n# 导入环境变量\nfrom dotenv import load_dotenv\nload_dotenv()\n\n# 导入A2A客户端和所需类型\nfrom core.a2a.client.client import A2AClient\n# 导入 Message, TextPart, TaskState, SendTaskResponse, GetTaskResponse, Task, JSONRPCError\nfrom core.a2a.types import (\n Part, TextPart, Message, TaskState,\n SendTaskResponse, GetTaskResponse, Task, JSONRPCError,\n SendTaskStreamingResponse # 导入流式响应类型\n)\n\n# 配置日志\nlogging.basicConfig(level=logging.INFO) # 可以改为 DEBUG 获取详细日志\nlogger = logging.getLogger(__name__)\n\n# 测试场景1: 同步请求 - 货币转换查询 (修正)\nasync def test_sync_currency_conversion(client: A2AClient):\n print(\"\\n=== 测试场景1: 同步请求 - Agent 调用 (计算器) ===\")\n # query = \"How much is the exchange rate for 1 USD to INR?\" # 这个查询可能需要搜索工具\n query = \"计算 58 * 34 的结果\" # 使用计算器工具确保能得到结果\n task_id = \"test_sync_\" + uuid4().hex # 客户端生成任务ID\n try:\n # 1. 构建 Message 对象\n message = Message(role=\"user\", parts=[TextPart(text=query)])\n\n # 2. 构建 TaskSendParams 对应的 payload 字典\n payload_dict = {\n \"id\": task_id,\n \"sessionId\": \"test_session_sync_1\",\n \"message\": message.model_dump(), # 序列化为字典\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"} # 与 AgentCard 中的 skill name/id 对应\n }\n logger.debug(f\"Sending sync task with payload: {payload_dict}\")\n\n # 3. 调用 send_task，传入 payload 字典\n response: SendTaskResponse = await client.send_task(payload=payload_dict)\n logger.debug(f\"Send task response: {response}\")\n\n # 4. 处理响应\n if response.error:\n error: JSONRPCError = response.error\n print(f\"发送任务时出错: Code={error.code}, Message={error.message}\")\n return None\n if not response.result:\n print(f\"发送任务成功，但未收到任务详情: {response}\")\n # 继续使用我们发送的 task_id 查询\n elif response.result.id != task_id:\n logger.warning(f\"服务器返回的任务ID '{response.result.id}' 与客户端发送的ID '{task_id}' 不匹配。\")\n task_id = response.result.id # 以服务器返回的为准\n\n print(f\"任务已发送，ID: {task_id}\")\n\n # 5. 轮询等待任务完成\n print(\"等待任务完成...\")\n task_result: Optional[Task] = None\n for attempt in range(10):\n await asyncio.sleep(2)\n get_payload = {\"id\": task_id}\n logger.debug(f\"Getting task with payload: {get_payload} (Attempt {attempt+1})\")\n get_response: GetTaskResponse = await client.get_task(payload=get_payload)\n logger.debug(f\"Get task response: {get_response}\")\n\n if get_response.error:\n error: JSONRPCError = get_response.error\n print(f\"获取任务时出错: Code={error.code}, Message={error.message}\")\n return None\n if not get_response.result:\n print(f\"获取任务成功，但未收到任务详情: {get_response}\")\n continue\n\n task_result = get_response.result\n print(f\" 当前任务状态: {task_result.status.state.value}\") # 使用 .value 获取枚举值\n if task_result.status.state in [TaskState.COMPLETED, TaskState.FAILED, TaskState.CANCELED]:\n break\n else:\n print(\"任务在限定时间内未完成。\")\n return None\n\n # 6. 处理最终任务结果 (使用属性访问)\n if task_result.status.state == TaskState.COMPLETED and task_result.artifacts:\n print(\"任务成功完成。结果:\")\n for artifact in task_result.artifacts:\n if artifact.parts:\n for part in artifact.parts:\n if isinstance(part, TextPart): # 检查类型\n print(f\" - {part.text}\") # 访问属性\n elif task_result.status.state == TaskState.FAILED:\n error_msg = \"未知错误\"\n if task_result.status.message and task_result.status.message.parts:\n if isinstance(task_result.status.message.parts[0], TextPart):\n error_msg = task_result.status.message.parts[0].text\n print(f\"任务失败: {error_msg}\")\n else:\n print(f\"任务最终状态为: {task_result.status.state.value}\")\n\n return task_result\n\n except Exception as e:\n logger.error(f\"处理同步任务时发生异常: {e}\", exc_info=True)\n print(f\"发送同步任务失败: {e}\")\n return None\n\n# 测试场景2: 多轮对话 - 不完整信息 (修正，但有局限性)\nasync def test_multi_turn_conversation(client: A2AClient):\n print(\"\\n=== 测试场景2: 多轮对话 (Agent 可能不支持) ===\")\n print(\"注意：当前服务器端的 Agent 实现可能不支持真正的多轮状态保持。\")\n\n # --- 第一轮对话 ---\n session_id = \"test_session_multi_\" + uuid4().hex # 为多轮对话创建唯一 session ID\n query1 = \"100美元等于多少\" # 故意不指定目标货币\n task_id_1 = \"test_multi_1_\" + uuid4().hex\n\n try:\n print(f\"\\n第一轮对话 (Session: {session_id}): 发送 '{query1}'\")\n # 1a. 构建 Message 和 Payload\n message1 = Message(role=\"user\", parts=[TextPart(text=query1)])\n payload_dict1 = {\n \"id\": task_id_1,\n \"sessionId\": session_id, # 传递 session ID\n \"message\": message1.model_dump(),\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"}\n }\n logger.debug(f\"Sending multi-turn task 1 with payload: {payload_dict1}\")\n\n # 1b. 发送任务\n response1: SendTaskResponse = await client.send_task(payload=payload_dict1)\n logger.debug(f\"Send task 1 response: {response1}\")\n\n if response1.error:\n error: JSONRPCError = response1.error\n print(f\"发送第一轮任务时出错: Code={error.code}, Message={error.message}\")\n return None\n if response1.result:\n task_id_1 = response1.result.id # Use server-confirmed ID\n\n print(f\"第一轮任务已发送，ID: {task_id_1}\")\n\n # 1c. 轮询获取结果\n print(\"等待第一轮任务响应...\")\n task1_result: Optional[Task] = None\n for attempt in range(5): # 减少轮询次数\n await asyncio.sleep(2)\n get_payload1 = {\"id\": task_id_1}\n get_response1: GetTaskResponse = await client.get_task(payload=get_payload1)\n if get_response1.result:\n task1_result = get_response1.result\n print(f\" 当前任务状态: {task1_result.status.state.value}\")\n if task1_result.status.state != TaskState.WORKING:\n break\n else:\n print(\"第一轮任务在限定时间内未完成或未开始。\")\n return None\n\n # 1d. 检查 Agent 是否要求输入 (当前 Agent 可能直接完成或失败)\n if task1_result.status.state == TaskState.INPUT_REQUIRED and task1_result.status.message:\n print(\"Agent 要求更多信息:\")\n for part in task1_result.status.message.parts:\n if isinstance(part, TextPart):\n print(f\" Agent: {part.text}\")\n\n # --- 第二轮对话 ---\n query2 = \"日元\" # 提供目标货币\n task_id_2 = \"test_multi_2_\" + uuid4().hex\n print(f\"\\n第二轮对话 (Session: {session_id}): 发送 '{query2}'\")\n\n # 2a. 构建 Message 和 Payload\n message2 = Message(role=\"user\", parts=[TextPart(text=query2)])\n payload_dict2 = {\n \"id\": task_id_2,\n \"sessionId\": session_id, # 必须使用相同的 session ID\n \"message\": message2.model_dump(),\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"}\n }\n logger.debug(f\"Sending multi-turn task 2 with payload: {payload_dict2}\")\n\n # 2b. 发送任务\n response2: SendTaskResponse = await client.send_task(payload=payload_dict2)\n logger.debug(f\"Send task 2 response: {response2}\")\n\n if response2.error:\n error: JSONRPCError = response2.error\n print(f\"发送第二轮任务时出错: Code={error.code}, Message={error.message}\")\n return None\n if response2.result:\n task_id_2 = response2.result.id\n\n print(f\"第二轮任务已发送，ID: {task_id_2}\")\n\n # 2c. 轮询获取最终结果\n print(\"等待第二轮任务完成...\")\n task2_result: Optional[Task] = None\n for attempt in range(10):\n await asyncio.sleep(2)\n get_payload2 = {\"id\": task_id_2}\n get_response2: GetTaskResponse = await client.get_task(payload=get_payload2)\n if get_response2.result:\n task2_result = get_response2.result\n print(f\" 当前任务状态: {task2_result.status.state.value}\")\n if task2_result.status.state != TaskState.WORKING:\n break\n else:\n print(\"第二轮任务在限定时间内未完成。\")\n return None\n\n # 2d. 处理最终结果\n if task2_result.status.state == TaskState.COMPLETED and task2_result.artifacts:\n print(\"多轮任务成功完成。最终结果:\")\n for artifact in task2_result.artifacts:\n if artifact.parts:\n for part in artifact.parts:\n if isinstance(part, TextPart):\n print(f\" - {part.text}\")\n else:\n print(f\"第二轮任务最终状态为: {task2_result.status.state.value}\")\n\n return task2_result\n\n elif task1_result.status.state == TaskState.COMPLETED:\n print(\"Agent 在第一轮就已完成任务 (可能直接使用了默认货币或无法处理):\")\n if task1_result.artifacts:\n for artifact in task1_result.artifacts:\n if artifact.parts:\n for part in artifact.parts:\n if isinstance(part, TextPart):\n print(f\" - {part.text}\")\n return task1_result\n else:\n print(f\"第一轮任务未要求输入，最终状态为: {task1_result.status.state.value}\")\n return task1_result\n\n except Exception as e:\n logger.error(f\"处理多轮对话时发生异常: {e}\", exc_info=True)\n print(f\"多轮对话测试失败: {e}\")\n return None\n\n\n# 测试场景3: 流式响应 (修正)\nasync def test_streaming_response(client: A2AClient):\n print(\"\\n=== 测试场景3: 流式响应 (Agent 端为模拟) ===\")\n\n # query = \"What are the current exchange rates between USD, EUR, and JPY?\"\n query = \"用中文写一首关于春天的短诗\" # 更适合流式输出的查询\n task_id = \"test_stream_\" + uuid4().hex\n try:\n # 1. 构建 Message 和 Payload\n message = Message(role=\"user\", parts=[TextPart(text=query)])\n payload_dict = {\n \"id\": task_id,\n \"sessionId\": \"test_session_stream_1\",\n \"message\": message.model_dump(),\n \"acceptedOutputModes\": [\"text\"],\n \"metadata\": {\"skill_name\": \"react_query\"}\n }\n logger.debug(f\"Sending streaming task with payload: {payload_dict}\")\n print(f\"任务已发送，ID: {task_id}\")\n\n # 2. 调用 send_task_streaming (不使用 await) 并使用 async for 迭代\n event_stream_generator = client.send_task_streaming(payload=payload_dict)\n\n print(\"开始接收流式响应:\")\n async for event_response in event_stream_generator:\n logger.debug(f\"Received stream event: {event_response}\")\n\n if event_response.error:\n error: JSONRPCError = event_response.error\n print(f\"流式传输中出错: Code={error.code}, Message={error.message}\")\n continue\n\n event = event_response.result\n if not event:\n logger.warning(\"Received stream response with empty result.\")\n continue\n\n # 处理状态更新事件\n if hasattr(event, \"status\") and event.status and event.status.message:\n if event.status.message.parts:\n for part in event.status.message.parts:\n if isinstance(part, TextPart):\n print(f\" 流式更新: {part.text}\")\n\n # 处理 Artifact 事件\n if hasattr(event, \"artifact\") and event.artifact:\n # print(\" 收到 Artifact:\") # 打印多次可能比较干扰，注释掉\n if event.artifact.parts:\n for part in event.artifact.parts:\n if isinstance(part, TextPart):\n print(f\" 流式结果: {part.text}\")\n\n # 检查结束标志\n if hasattr(event, \"final\") and event.final:\n print(\"流式响应结束标志收到。\")\n\n print(\"流式任务处理完成。\")\n return True\n\n except Exception as e:\n logger.error(f\"处理流式任务时发生异常: {e}\", exc_info=True)\n print(f\"发送流式任务失败: {e}\")\n return False\n\n# 主函数 (修正)\nasync def main():\n print(\"=== LangGraph Agent A2A协议测试 ===\\n\")\n # print(\"此测试脚本将测试LangGraph Currency Agent通过A2A协议的三种交互场景:\")\n # print(\"1. 同步请求 - Agent 调用 (计算器)\")\n # print(\"2. 多轮对话 - 处理不完整信息 (Agent 可能不支持)\")\n # print(\"3. 流式响应 - 实时状态更新 (Agent 端为模拟)\")\n\n # 创建A2A客户端\n client = A2AClient(url=\"http://127.0.0.1:8000\")\n\n # --- 移除了 get_agent_info 调用 ---\n # (如果需要验证服务器是否在线，可以尝试发送一个简单的任务)\n print(\"尝试连接到服务器并运行测试...\")\n print(\"-\" * 30)\n\n # 执行测试场景\n await test_sync_currency_conversion(client)\n print(\"-\" * 30)\n # 注意：多轮对话测试依赖于 Agent 对话状态的处理能力\n await test_multi_turn_conversation(client)\n print(\"-\" * 30)\n await test_streaming_response(client)\n print(\"-\" * 30)\n print(\"所有测试场景执行完毕。\")\n\n# 运行测试\nif __name__ == \"__main__\":\n asyncio.run(main())"
},
{
"path": "examples/16_google_a2a/currency_agent_test_README.md",
"content": "# LangGraph Agent A2A协议交互测试\n\n## 概述\n\n本测试脚本 (`examples/a2a/currency_agent_test.py`) 旨在通过具体的交互场景，测试和演示如何使用 A2A 客户端与先前通过 `langgraph_integration.py` 启动的 LangGraph Agent 服务进行通信。它覆盖了同步请求（涉及工具调用）、尝试进行多轮对话以及接收（模拟的）流式响应等场景。\n\n## 测试场景说明\n\n此脚本包含以下三个主要测试场景：\n\n1. **场景 1: 同步请求 - Agent 调用 (涉及工具)**\n * **目的:** 测试发送一个需要 Agent 调用内部工具（如此示例中的计算器）才能完成的请求。\n * **流程:** 客户端发送一个计算任务 -> 服务器端 Agent (LangGraph ReAct) 解析任务 -> 调用 `calculator` 工具 -> 获取结果 -> LLM 整合答案 -> 服务器返回最终结果 -> 客户端轮询获取并显示结果。\n * **预期:** 客户端能成功获取到 Agent 计算后的准确结果。\n\n2. **场景 2: 多轮对话尝试 (Agent 当前实现有限)**\n * **目的:** 测试客户端在需要多步交互时的请求发送方式（使用 `sessionId`），并观察当前 Agent 的响应行为。\n * **流程:**\n * 第一轮：客户端发送一个信息不明确的查询（例如 \"100美元等于多少\"，缺少目标货币），并附带 `sessionId`。\n * 客户端轮询获取结果。\n * **注意:** *根据我们当前的 Agent 实现 (`CurrencyAgent` 使用 `create_react_agent` 且 `invoke` 未特殊处理对话历史)，Agent 很可能不会返回 `input-required` 状态来请求更多信息，而是会直接尝试处理或告知无法处理，然后将任务标记为 `completed` 或 `failed`。*\n * (理想流程中，如果 Agent 返回 `input-required`，客户端会发送第二轮请求补充信息，使用相同的 `sessionId`。)\n * **预期:** 客户端能够正确发送带 `sessionId` 的请求，并能处理 Agent 的最终响应（即使它没有按预期进入多轮澄清状态）。此测试主要验证客户端的多轮请求发送能力和对 Agent 当前行为的观察。\n\n3. **场景 3: 流式响应 (Agent 端模拟)**\n * **目的:** 测试客户端接收 A2A 流式响应 (Server-Sent Events) 的能力。\n * **流程:** 客户端发送一个适合流式输出的查询 -> 服务器端的 `AgentTaskManager` 调用 `CurrencyAgent.stream` 方法 -> **注意:** *`CurrencyAgent.stream` 当前是一个模拟实现，它会发送预设的文本块，而不是真正调用 LangGraph 的流式接口。* -> 客户端接收并打印这些模拟的流式事件。\n * **预期:** 客户端能够成功连接 SSE 端点，并接收、打印服务器发送的（模拟）流式事件。\n\n## 运行测试\n\n### 前提条件\n\n* Python (推荐 3.10 或更高版本)\n* 已根据项目 `requirements.txt` 安装所有必需的 Python 依赖库。\n* 在项目根目录下的 `.env` 文件中配置了有效的 `OPENAI_API_KEY` (或其他所需的 LLM API 密钥)。\n\n### 步骤\n\n1. **启动 A2A 服务器:**\n * 确保你位于项目的根目录。\n * 在终端中运行 (如果尚未运行):\n ```bash\n python -m examples.a2a.langgraph_integration\n ```\n * 服务器应成功启动并监听在 `http://127.0.0.1:8000`。\n\n2. **运行本测试脚本:**\n * 打开 **另一个** 终端。\n * 确保你位于项目的根目录并激活了相同的虚拟环境。\n * 运行测试脚本:\n ```bash\n python -m examples.a2a.currency_agent_test\n ```\n\n## 测试输出示例 (基于实际运行结果)\n\n以下是运行此测试脚本时预期的输出格式，反映了当前 Agent 的实际行为：\n\n### 同步请求示例 (计算器调用)\n\n```\n=== 测试场景1: 同步请求 - Agent 调用 (计算器) ===\n\n任务已发送，ID: test_sync_...\n等待任务完成...\n 当前任务状态: completed\n任务成功完成。结果:\n - 58 * 34 的结果是 1972。\n```\n\n### 多轮对话示例 (Agent 第一轮即完成)\n\n```\n=== 测试场景2: 多轮对话 (Agent 可能不支持) ===\n注意：当前服务器端的 Agent 实现可能不支持真正的多轮状态保持。\n\n第一轮对话 (Session: test_session_multi_...): 发送 '100美元等于多少'\n第一轮任务已发送，ID: test_multi_1_...\n等待第一轮任务响应...\n 当前任务状态: completed\nAgent 在第一轮就已完成任务 (可能直接使用了默认货币或无法处理):\n - 目前无法提供100美元等于多少人民币的具体信息。你可以查阅最新的汇率数据或使用汇率转换工具来获取准确的结果。\n```\n*(注意：Agent 的具体回复可能因 LLM 的不同调用而略有差异)*\n\n### 流式响应示例 (Agent 端模拟)\n\n```\n=== 测试场景3: 流式响应 (Agent 端为模拟) ===\n\n任务已发送，ID: test_stream_...\n开始接收流式响应:\n 流式更新: 正在处理您的请求...\n 流式结果: 关于 '用中文写一首关于春天的短诗' 的信息如下：这是一个模拟的回应，因为真实流未实现。\n流式响应结束标志收到。\n流式任务处理完成。\n```\n\n## 注意事项\n\n* 确保在运行测试前已正确设置 `.env` 文件中的环境变量。\n* 测试脚本默认连接 `http://127.0.0.1:8000`。如果服务器地址或端口不同，请修改脚本中的 `A2AClient` 初始化 URL。\n* 如果连接失败或测试出错，请优先检查 A2A 服务器是否已正确启动且正在运行，并查看服务器端的日志输出。\n\n---\n\n## 两个客户端示例的命名与区别\n\n你项目中有两个客户端示例文件，我们可以为它们命名并说明其侧重点：\n\n1. **`examples/a2a/client_example.py` -> \"基础客户端示例 (Basic Client Example)\"**\n * **目的:** 这个脚本更侧重于**基础演示**，展示了调用 `A2AClient` 库中几个核心方法（`send_task`, `get_task`, `send_task_streaming`）的最基本用法。\n * **特点:** 代码相对简洁，逻辑直接，主要目的是让使用者快速了解如何发起不同类型的 A2A 请求并处理最简单的成功响应。它包含了一个简单的轮询逻辑。\n\n2. **`examples/a2a/currency_agent_test.py` -> \"场景化测试客户端 (Scenario-based Test Client)\"**\n * **目的:** 这个脚本的定位是**功能测试和场景演示**。它针对我们集成的 LangGraph Agent 设计了几个具体的交互场景（同步工具调用、尝试多轮对话、流式接收），以验证端到端的流程和观察 Agent 在特定情况下的行为。\n * **特点:** 结构更清晰地划分为不同的测试函数，包含了更具体的业务逻辑查询（尽管有些是模拟的或揭示了 Agent 的局限性），并且其输出更侧重于展示每个测试场景的结果。它也使用了轮询，并尝试了多轮交互的状态传递（通过 `sessionId`）。\n\n**主要区别总结:**\n\n| 特性 | `client_example.py` (基础示例) | `currency_agent_test.py` (场景化测试) |\n| :----------- | :--------------------------------------------- | :---------------------------------------------------- |\n| **目标** | 演示 Client API 基本用法 | 测试/演示特定交互场景 |\n| **结构** | 简单的顺序调用 | 按测试场景划分函数 |\n| **复杂度** | 较低，核心 API 调用 | 略高，包含场景逻辑（如尝试多轮） |\n| **查询内容** | 通用示例（计算、搜索） | 针对场景设计（计算、不完整查询、适合流式的查询） |\n| **侧重点** | 如何调用 API | Agent 在特定场景下的行为和端到端流程验证 |\n"
},
{
"path": "examples/16_google_a2a/langgraph_integration.py",
"content": "# examples/a2a/langgraph_integration.py\n\nimport os\nimport sys\nimport asyncio # asyncio 仍然可能被依赖库使用，保留导入\nimport logging\n# 确保导入了 List, Tuple, Optional, TypedDict\nfrom typing import Dict, Any, List, Optional, AsyncIterable, Union, TypedDict, Tuple\n\n# 添加项目根目录到路径\nsys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))\n\n# 导入环境变量\nfrom dotenv import load_dotenv\nload_dotenv()\n\n# 导入A2A相关组件\n# 从你的项目结构导入\nfrom core.a2a.types import (\n AgentCard, AgentCapabilities, AgentSkill,\n Task, TaskState, TaskStatus, Artifact, Message, TextPart, # TextPart 可能不再直接使用\n JSONRPCResponse, InvalidParamsError, InternalError,\n SendTaskRequest, SendTaskResponse, TaskSendParams\n)\nfrom core.a2a.server.server import A2AServer\nfrom core.a2a.agent_task_manager import AgentTaskManager\n\n# 导入LangChain和LLM相关组件\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.tools import tool\n# StateGraph 和 END 不再直接使用，但保留导入\nfrom langgraph.graph import END, StateGraph\nfrom langgraph.prebuilt import create_react_agent\n\n# 配置日志\nlogging.basicConfig(level=logging.INFO) # 可以改为 DEBUG 获取更详细日志\nlogger = logging.getLogger(__name__)\n\n# --- 定义工具 (保持不变) ---\n@tool\ndef search(query: str) -> str:\n \"\"\"搜索互联网获取信息\"\"\"\n # 实际应用中应调用真实搜索引擎 API\n logger.info(f\"Tool 'search' called with query: {query}\")\n return f\"这是关于 '{query}' 的模拟搜索结果。\"\n\n@tool\ndef calculator(expression: str) -> str:\n \"\"\"计算数学表达式\"\"\"\n logger.info(f\"Tool 'calculator' called with expression: {expression}\")\n try:\n # 注意：生产环境中使用 eval 非常危险，这里仅作示例\n # 限制 eval 的能力，只允许简单的数学运算\n allowed_names = {\n k: v for k, v in __import__(\"math\").__dict__.items() if not k.startswith(\"_\")\n }\n allowed_names.update({\"abs\": abs, \"int\": int, \"float\": float}) # 添加常用函数\n code = compile(expression, \"\", \"eval\")\n\n for name in code.co_names:\n if name not in allowed_names:\n raise NameError(f\"Use of name '{name}' not allowed\")\n\n result = eval(code, {\"__builtins__\": {}}, allowed_names)\n return f\"计算结果: {result}\"\n except NameError as e:\n logger.error(f\"Calculation error (NameError): {e} in expression '{expression}'\")\n return f\"计算错误: 不允许的名称 '{e.name}'\"\n except Exception as e:\n logger.error(f\"Calculation error: {e} in expression '{expression}'\")\n return f\"计算错误: {e}\"\n\n# --- 修正 AgentState 定义 ---\nclass AgentState(TypedDict):\n # 使用 'messages' 字段来传递对话内容\n # 格式为 (角色, 内容) 的元组列表\n messages: List[Tuple[str, str]]\n # session_id 可以保留，如果Agent内部逻辑需要的话 (create_react_agent 通常不需要)\n # session_id: Optional[str]\n # 注意: ReAct Agent 运行时可能会在状态中添加其他键 (例如 intermediate_steps)\n\n# --- 修正 CurrencyAgent 类 ---\nclass CurrencyAgent:\n \"\"\"一个简单的货币转换和信息查询Agent (已修正)\"\"\"\n\n # 支持的内容类型 (保持不变)\n SUPPORTED_CONTENT_TYPES = [\"text\"]\n\n def __init__(self, llm):\n \"\"\"初始化Agent，直接使用 create_react_agent 创建的 Runnable\"\"\"\n self.tools = [search, calculator]\n # create_react_agent 返回一个可直接调用的 Runnable (图)\n self.agent_runnable = create_react_agent(llm, self.tools)\n logger.info(\"CurrencyAgent initialized with ReAct runnable.\")\n\n def invoke(self, query: str, session_id: str = None) -> str:\n \"\"\"同步调用Agent Runnable\"\"\"\n # (session_id 在此实现中未传递给 agent_runnable，如果需要可以添加)\n logger.debug(f\"[CurrencyAgent.invoke] Received query: '{query}', session_id: '{session_id}'\")\n if not query:\n logger.error(\"[CurrencyAgent.invoke] Query is empty!\")\n return \"错误：输入查询为空。\"\n\n # 准备 ReAct Agent Runnable 所需的输入\n invoke_input = {\"messages\": [(\"user\", query)]}\n\n logger.debug(f\"[CurrencyAgent.invoke] Invoking agent runnable with input: {invoke_input}\")\n try:\n # 直接调用 create_react_agent 返回的 runnable\n result = self.agent_runnable.invoke(invoke_input)\n logger.debug(f\"[CurrencyAgent.invoke] Agent runnable result: {result}\")\n\n # 提取最终响应\n final_output = \"错误：未能从Agent获取有效响应。\"\n if isinstance(result, dict) and isinstance(result.get(\"messages\"), list) and result[\"messages\"]:\n last_message = result[\"messages\"][-1]\n if isinstance(last_message, tuple) and len(last_message) == 2:\n final_output = last_message[1]\n elif hasattr(last_message, 'content'):\n final_output = last_message.content\n else:\n logger.warning(f\"[CurrencyAgent.invoke] Last message format unexpected: {last_message!r}\")\n else:\n logger.warning(f\"[CurrencyAgent.invoke] Could not find 'messages' list in result: {result}\")\n\n logger.debug(f\"[CurrencyAgent.invoke] Returning output: {final_output}\")\n return str(final_output)\n except Exception as e:\n logger.error(f\"[CurrencyAgent.invoke] Exception during agent invocation: {e}\", exc_info=True)\n raise\n\n async def ainvoke(self, inputs: dict) -> dict:\n \"\"\"异步调用Agent Runnable (输入格式也需调整)\"\"\"\n # TODO: 确认这里的输入格式是否也需要转换为 {\"messages\": [...]}\n logger.debug(f\"[CurrencyAgent.ainvoke] Invoking agent runnable async with input: {inputs}\")\n # 假设输入字典已经包含了正确的 \"messages\" 键\n return await self.agent_runnable.ainvoke(inputs)\n\n async def stream(self, query: str, session_id: str = None):\n \"\"\"流式调用Agent (当前为模拟)\"\"\"\n # TODO: 实现真实的流式调用\n logger.warning(\"[CurrencyAgent.stream] Stream method is currently mocked.\")\n # --- 模拟实现 ---\n yield { \"content\": \"正在处理您的请求...\", \"is_task_complete\": False, \"require_user_input\": False }\n await asyncio.sleep(0.5)\n final_simulated_answer = f\"关于 '{query}' 的信息如下：这是一个模拟的回应，因为真实流未实现。\"\n yield { \"content\": final_simulated_answer, \"is_task_complete\": True, \"require_user_input\": False }\n # --- 模拟结束 ---\n\n\n# --- A2A 服务器设置 (修正函数定义和 AgentCard) ---\n# 将函数改为同步定义 (def 而不是 async def)\ndef setup_a2a_server():\n \"\"\"设置并返回 A2A 服务器实例 (同步函数)\"\"\"\n print(\"\\n=== 配置 LangGraph A2A 服务器 ===\\n\")\n\n # 创建LLM\n try:\n llm = ChatOpenAI(model=\"gpt-4o-mini\")\n logger.info(\"Using OpenAI LLM: gpt-4o-mini\")\n except Exception as e:\n print(f\"警告: 无法创建OpenAI LLM ({e})，将使用模拟模式\")\n from langchain.llms.fake import FakeListLLM\n llm = FakeListLLM(responses=[\"这是一个模拟的LLM响应\"])\n logger.info(\"Using FakeListLLM (simulation mode)\")\n\n # 创建 Agent 实例\n agent = CurrencyAgent(llm)\n\n # 创建 Agent 卡片 (添加缺失字段)\n agent_card = AgentCard(\n name=\"LangGraph ReAct Agent\",\n description=\"一个使用LangGraph ReAct处理查询并调用工具的Agent\",\n url=\"http://127.0.0.1:8000/agent\", # Agent 的访问 URL (示例)\n version=\"0.1.0\", # Agent 的版本号\n capabilities=AgentCapabilities( # 设置 Agent 的能力\n streaming=False, # 当前 stream 是模拟的，设为 False\n pushNotifications=False # 假设不支持推送\n ),\n skills=[ # skills 列表在 AgentCard 顶层\n AgentSkill(\n id=\"react_query_skill\", # 技能的唯一 ID\n name=\"react_query\",\n description=\"处理自然语言查询，可使用搜索和计算器工具\",\n inputModes=[\"text\"],\n outputModes=[\"text\"]\n )\n ]\n # 其他可选字段可以按需添加\n )\n\n # 创建 AgentTaskManager\n task_manager = AgentTaskManager(agent)\n\n # 创建A2A服务器实例 (不在此处设置 host/port)\n server = A2AServer(agent_card=agent_card, task_manager=task_manager)\n print(\"A2A服务器实例已创建。\")\n return server # 返回实例\n\n\n# --- 主函数入口 (修正启动逻辑) ---\nif __name__ == \"__main__\":\n try:\n # 调用同步函数来设置服务器\n server_instance = setup_a2a_server()\n\n # 定义 HOST 和 PORT\n HOST = \"127.0.0.1\"\n PORT = 8000\n print(f\"准备启动A2A服务器，监听地址 http://{HOST}:{PORT}\")\n\n # 在调用 start 前设置 host 和 port\n # (或者修改 A2AServer 的 __init__ 让其接受 host/port)\n server_instance.host = HOST\n server_instance.port = PORT\n\n # 启动服务器 (调用同步的 start 方法)\n server_instance.start()\n\n except KeyboardInterrupt:\n print(\"\\n服务器已手动停止。\")\n except Exception as e:\n # 捕获设置或启动过程中的其他异常\n logger.error(f\"启动服务器时发生未处理的异常: {e}\", exc_info=True)"
},
{
"path": "examples/TODO_computer_tool_demo.py",
"content": "from typing import Annotated, Literal\nfrom langchain_core.messages import HumanMessage, AIMessage\nfrom langchain.agents import AgentExecutor, create_openai_tools_agent\nfrom langchain_core.runnables import Runnable, RunnablePassthrough\nfrom langchain_core.runnables.graph import StateGraph, END, START\nfrom langchain.tools.render import render_text_description\nfrom langchain_openai import ChatOpenAI\nfrom langchain_core.tools import Tool\nfrom langchain.agents.format_scratchpad import format_to_openai_tool_messages\nfrom langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\nfrom langchain_core.runnables.config import RunnableConfig\nfrom langgraph.graph import END, StateGraph\nfrom langgraph.prebuilt import ToolNode\nfrom langgraph.graph.message import Command, InjectedState\n\n# Import our custom computer tool\n# TODO: MarinaBox - Import our custom computer tool\nfrom marinabox import mb_start_computer, mb_stop_computer, mb_use_computer_tool\n\n# Set up model with tools\nmodel = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\ntools = [mt_use_computer_tool()]\nmodel_with_tools = model.bind_tools(tools)\n\n# Define workflow nodes\ndef should_continue(state: Annotated[dict, InjectedState()]):\n messages = state[\"messages\"]\n if len(messages) > 0:\n last_message = messages[-1]\n if last_message.tool_calls:\n return Command(goto=\"tool_node\")\n else:\n return Command(goto=\"stop_computer\")\n return Command(goto=\"stop_computer\")\n\ndef call_model(state: Annotated[dict, InjectedState()]):\n input_message = input(\"Enter your message: \")\n if input_message != \"stop_computer\":\n messages = [HumanMessage(content=input_message)]\n response = model_with_tools.invoke(messages)\n return {\"messages\": [response], \"session_id\": state.get(\"session_id\")}\n else:\n return {\"messages\": [], \"session_id\": state.get(\"session_id\")}\n\n# Set up workflow\nworkflow = StateGraph(dict)\nworkflow.add_node(\"start_computer\", mt_start_computer)\nworkflow.add_node(\"agent\", call_model)\nworkflow.add_node(\"tool_node\", ToolNode(tools=tools))\nworkflow.add_node(\"stop_computer\", mt_stop_computer)\nworkflow.add_node(\"should_continue\", should_continue)\n\n# Define workflow edges\nworkflow.add_edge(START, \"start_computer\")\nworkflow.add_edge(\"start_computer\", \"agent\")\nworkflow.add_edge(\"tool_node\", \"agent\")\nworkflow.add_edge(\"agent\", \"should_continue\")\nworkflow.add_edge(\"stop_computer\", END)\n\n# Compile and run workflow\napp = workflow.compile()\nif __name__ == \"__main__\":\n app.invoke({\"messages\": \"\"})\n"
},
{
"path": "examples/__init__.py",
"content": ""
},
{
"path": "examples/state_based_supervisor_examples/01_simple.py",
"content": "import asyncio\nimport json\nimport os\nimport re\nimport time \nfrom datetime import datetime \nfrom typing import Literal, List, Dict, Any, Optional, cast\n\n# --- LangChain / LangGraph ---\ntry:\n # 使用 langchain_openai (或你选择的模型提供商)\n from langchain_openai import ChatOpenAI \nexcept ImportError:\n ChatOpenAI = None \n print(\"Warning: langchain_openai not installed.\")\n\n# 核心消息类型\nfrom langchain_core.messages import HumanMessage, AIMessage, BaseMessage, ToolMessage \n# LangChain 工具相关\nfrom langchain_core.tools import tool, BaseTool \n\n# --- OpenAI 错误处理 ---\ntry:\n from openai import RateLimitError\nexcept ImportError:\n class RateLimitError(Exception): pass\n\n# --- 内部模块导入 (请确保路径正确) ---\ntry:\n # 假设这些是你当前的路径\n from core.agents.sb_supervisor_agent import SupervisorAgent \n from core.agents.supervisor.state_schema import PlanningAgentState\n from core.agents.base.react_agent import ReactAgent # 导入 ReactAgent\n # 导入 StreamUpdate (如果需要在最终状态中检查它，但这里主要关注消息)\n # from core.agents.supervisor.schemas import StreamUpdate \n\nexcept ImportError as e:\n print(f\"Error importing agent components: {e}\")\n print(\"Please ensure paths like 'core.agents.sb_supervisor_agent' are correct.\")\n\nimport traceback\n\n# --- 定义 Web Search 工具 ---\n# 使用 @tool 装饰器明确这是一个工具\n@tool\ndef web_search(query: str) -> str:\n \"\"\"Search the web for current information about a given query. Use this for recent events, data, or facts.\"\"\"\n print(f\"--- TOOL CALLED: web_search(query='{query}') ---\") # 添加日志确认工具被调用\n # Mocked data - 实际使用时会调用 Tavily 或其他搜索引擎\n if \"apple\" in query.lower() and \"headcount\" in query.lower() and \"2024\" in query:\n return (\n \"According to recent (mocked) reports for 2024, Apple's headcount is approximately 164,000 employees globally.\"\n )\n elif \"joke\" in query.lower():\n # 这个工具不适合讲笑话\n return \"I am a web search tool, I cannot tell jokes.\"\n else:\n return f\"Mock search results for query: '{query}'. Found relevant information on various websites.\"\n\n# --- 主执行逻辑 ---\nasync def main():\n # --- 初始化 LLM (确保 API Key 在环境中) ---\n try:\n model_name = os.getenv(\"LLM_MODEL_NAME\", \"gpt-4o\") \n print(f\"Using LLM: {model_name}\")\n if not ChatOpenAI: raise ImportError(\"ChatOpenAI not available.\")\n # 使用温度稍高的模型可能有助于 ReAct 思考和调用工具\n model = ChatOpenAI(model=model_name, temperature=0.2) \n except Exception as e:\n print(f\"Failed to initialize ChatOpenAI model: {e}\")\n return\n\n # --- 实例化 Agents ---\n try:\n # research_agent 现在有了一个明确定义的 web_search 工具\n research_agent = ReactAgent(\n name=\"research_expert\", \n tools=[web_search], # <--- 传入工具列表\n model=model,\n # 添加明确的 Prompt 引导工具使用\n prompt=(\n \"You are a research expert. Use available tools to find information. \"\n \"You have access to 'web_search'. Use it for questions about current data, facts, or events.\"\n ),\n max_context_tokens=8000 \n ) \n \n all_agents = [research_agent]\n\n # --- 实例化 Supervisor ---\n supervisor = SupervisorAgent(\n agents=all_agents,\n model=model, # Supervisor 使用相同的模型\n state_schema=PlanningAgentState, \n include_agent_name=\"inline\" \n # checkpointer=... \n )\n except Exception as e:\n print(f\"Failed to initialize agents or supervisor: {e}\")\n traceback.print_exc()\n return\n\n # --- 准备初始请求 ---\n # 用户请求包含两个意图：讲笑话 + 查信息\n user_request = (\n \"Hi! I'd like to start with a short joke to lighten the mood, \"\n \"then please check Apple's headcount in 2024. Summarize both.\"\n )\n print(f\"Initial Request: '{user_request}'\")\n\n # --- 准备初始状态 ---\n initial_graph_state: PlanningAgentState = {\n \"messages\": [HumanMessage(content=user_request)], # 使用 HumanMessage\n \"plan\": None,\n \"error\": None\n }\n\n # --- 执行 Supervisor (使用 ainvoke) ---\n final_state: Optional[Dict[str, Any]] = None\n error_occurred: Optional[Exception] = None\n config = {\"recursion_limit\": 100} \n\n try:\n print(\"\\n--- Invoking Supervisor Agent (ainvoke) ---\")\n final_state = await supervisor.ainvoke(initial_graph_state, config=config)\n print(\"\\n--- Supervisor Agent Invocation Complete ---\")\n\n # --- 错误处理 ---\n except RateLimitError as e: error_occurred = e; print(f\"\\n!!! OpenAI Quota Error: {e}\")\n except Exception as e: error_occurred = e; print(f\"\\n!!! Error during graph execution: {e}\"); traceback.print_exc()\n\n # --- 处理并打印最终结果 ---\n if error_occurred: print(\"\\n--- Graph Execution INTERRUPTED or FAILED ---\")\n else: print(\"\\n--- Graph Execution Finished ---\")\n\n if not final_state:\n print(\"Error: No final state available.\")\n return\n\n print(\"\\n--- FINAL STATE ---\")\n # 打印错误（如果在状态中记录了）\n if final_state.get(\"error\"): print(f\"\\nERROR RECORDED IN STATE: {final_state['error']}\")\n # 打印计划\n final_plan = final_state.get('plan')\n if final_plan: print(\"\\nFinal Plan State:\", json.dumps(final_plan, indent=2, default=str))\n else: print(\"\\nFinal Plan State: Not available.\")\n # 打印消息历史\n final_messages = final_state.get(\"messages\", [])\n if final_messages:\n print(\"\\nFinal Message History (Last 10):\")\n for m in final_messages[-10:]:\n try:\n if hasattr(m, 'pretty_print'): m.pretty_print()\n else: print(json.dumps(m, indent=2, default=str))\n print(\"-\" * 10)\n except Exception as print_err: print(f\"Error printing final message: {print_err}\")\n else: print(\"\\nFinal Message History: Empty.\")\n\n print(\"\\n--- END OF TEST ---\")\n\n\nif __name__ == \"__main__\":\n try:\n asyncio.run(main())\n except KeyboardInterrupt:\n print(\"\\nExecution interrupted by user.\")\n except Exception as e:\n print(f\"\\nAn unexpected top-level error occurred: {e}\")\n traceback.print_exc()"
},
{
"path": "examples/state_based_supervisor_examples/02_tavily.py",
"content": "# main.py (用于测试 State-Based Supervisor 和 ReactAgent)\n\nimport asyncio\nimport json\nimport os\nfrom typing import Dict, Any, Optional\nfrom langchain_community.tools import TavilySearchResults\n# --- LangChain / LangGraph ---\n# 假设模型直接在此初始化或从别处导入\nfrom dotenv import load_dotenv\nload_dotenv() # 自动加载 .env 文件\ntry:\n from langchain_openai import ChatOpenAI # 或者你使用的其他模型\nexcept ImportError:\n ChatOpenAI = None\n print(\"Warning: langchain_openai not installed.\")\n\n# 核心消息类型\nfrom langchain_core.messages import HumanMessage, AIMessage, BaseMessage, ToolMessage\n\n# --- OpenAI 错误处理 ---\ntry:\n from openai import RateLimitError\nexcept ImportError:\n class RateLimitError(Exception): pass\n\n# --- 内部模块导入 (请确保路径正确) ---\ntry:\n # 从你提供的 core.agents... 路径导入\n from core.agents.sb_supervisor_agent import SupervisorAgent # 你的 Supervisor 实现\n from core.agents.state_based_supervisor.state_schema import PlanningAgentState # 包含 Plan 的状态\n from core.agents.base.react_agent import ReactAgent # 你的 ReactAgent 实现\n from core.llm.llm_manager import LLMManager # LLM 管理器\n # (如果你的子 Agent 有更具体的类，在这里导入它们)\n # 例如:\n # from core.agents.researcher import ResearchAgent\n # from core.agents.coder import CoderAgent\n\n # --- 如果没有具体子 Agent 类，使用 ReactAgent 作为示例 ---\n # (确保 ReactAgent 可以被直接实例化用于测试)\n if not issubclass(ReactAgent, object): # 简单检查 ReactAgent 是否有效\n raise ImportError(\"ReactAgent class not found or invalid.\")\n\nexcept ImportError as e:\n print(f\"Error importing agent components: {e}\")\n print(\"Please ensure paths like 'core.agents.sb_supervisor_agent' are correct relative to your execution path.\")\n\nimport traceback\n\n# --- 主执行函数 (简化版，只关注最终结果) ---\nasync def run_supervisor_test(supervisor_agent: SupervisorAgent, initial_state: Dict[str, Any]):\n \"\"\"Executes the supervisor graph using ainvoke and prints the final state.\"\"\"\n\n print(\"--- Starting Supervisor Graph Test ---\")\n # 获取初始消息列表，检查是否为空\n messages_list = initial_state.get(\"messages\", [])\n initial_query = \"N/A\" # 默认值\n if messages_list:\n first_message = messages_list[0]\n # 检查第一个消息是否有 content 属性 (更健壮)\n if hasattr(first_message, 'content'):\n initial_query = first_message.content\n else:\n # 如果第一个元素不是预期的消息对象，记录一下\n print(f\"Warning: First item in initial messages is not a standard message object: {type(first_message)}\")\n initial_query = str(first_message) # 尝试转换为字符串\n print(f\"Initial Query: '{initial_query}'\")\n print(\"-\" * 30)\n\n config = {\"recursion_limit\": 100} # 使用较高的递归限制\n final_state: Optional[Dict[str, Any]] = None\n error_occurred: Optional[Exception] = None\n\n try:\n print(\"--- Invoking Supervisor Agent (ainvoke) ---\")\n # 直接调用 ainvoke 获取最终状态\n final_state = await supervisor_agent.ainvoke(initial_state, config=config)\n print(\"--- Supervisor Agent Invocation Complete ---\")\n\n # --- 错误处理 ---\n except RateLimitError as e:\n error_occurred = e\n print(\"\\n\" + \"=\"*40 + \"\\n!!! OpenAI API Error: Insufficient Quota !!!\\n\" + \"=\"*40)\n print(\"Execution stopped. Check OpenAI plan/billing.\")\n print(f\"Original error: {e}\")\n except TypeError as e:\n error_occurred = e\n print(\"\\n\" + \"=\"*40 + \"\\n!!! TypeError During Graph Execution !!!\\n\" + \"=\"*40)\n print(f\"Error details: {e}\")\n if \"synchronous function provided\" in str(e):\n print(\"Hint: Ensure all graph nodes support async or run the graph synchronously if needed.\")\n traceback.print_exc()\n except Exception as e:\n error_occurred = e\n print(\"\\n\" + \"=\"*40 + \"\\n!!! An Unexpected Error Occurred !!!\\n\" + \"=\"*40)\n print(f\"Error type: {type(e).__name__}\\nError details: {e}\")\n traceback.print_exc()\n\n # --- Process Final State ---\n if error_occurred: print(\"\\n--- Graph Execution INTERRUPTED or FAILED ---\")\n else: print(\"\\n--- Graph Execution Finished ---\")\n\n if not final_state:\n # 如果 ainvoke 返回 None 或在出错前未赋值 (理论上 ainvoke 会抛错或返回字典)\n print(\"Error: No final state available (Execution might have failed early).\")\n # 尝试从 supervisor agent 获取最后状态 (如果 checkpointer 可用且实现了 get_state)\n if hasattr(supervisor_agent, 'checkpointer') and supervisor_agent.checkpointer and hasattr(supervisor_agent.checkpointer, 'get'):\n try:\n # 需要知道配置中的 thread_id (这里假设是 'test_thread')\n last_checkpoint = supervisor_agent.checkpointer.get({\"configurable\": {\"thread_id\": \"test_thread\"}})\n if last_checkpoint:\n print(\"Attempting to display last known checkpoint state:\")\n final_state = last_checkpoint.get('channel_values', {})\n else:\n print(\"Could not retrieve last checkpoint state.\")\n except Exception as cp_err:\n print(f\"Error retrieving checkpoint state: {cp_err}\")\n\n # 即使出错，也尝试打印 final_state (可能是包含错误信息的状态)\n if final_state and isinstance(final_state, dict):\n print(\"\\n--- FINAL STATE ---\")\n\n # 1. 打印错误信息 (如果存在)\n if final_state.get(\"error\"):\n print(f\"\\nERROR RECORDED IN STATE: {final_state['error']}\")\n\n # 2. 打印最终消息历史 (尝试 pretty_print)\n final_messages = final_state.get(\"messages\", [])\n if final_messages and isinstance(final_messages, list):\n print(\"\\nFinal Message History (Last ~10):\")\n for m in final_messages[-10:]: # 只打印最后一部分\n try:\n if hasattr(m, 'pretty_print'):\n m.pretty_print()\n else: # Fallback for dict or other types\n print(json.dumps(m, indent=2, default=str))\n print(\"-\" * 10)\n except Exception as print_err:\n print(f\"Error printing final message: {print_err}\")\n else:\n print(\"\\nFinal Message History: Not available or empty.\")\n\n # 3. 打印最终计划状态\n final_plan = final_state.get('plan')\n if final_plan and isinstance(final_plan, dict):\n print(\"\\nFinal Plan State:\")\n print(json.dumps(final_plan, indent=2, default=str))\n else:\n print(\"\\nFinal Plan State: Not available or not generated.\")\n\n else:\n print(\"\\n--- No Final State Could Be Displayed ---\")\n\n\n print(\"\\n--- END OF TEST ---\")\n return final_state\n\n# --- Main Execution Block ---\nasync def main():\n # --- 1. 初始化 LLM 管理器 (它会自动注册配置好的模型) ---\n try:\n model_manager = LLMManager()\n # 可以选择打印一下注册了哪些模型\n print(\"Registered Models:\", json.dumps(model_manager.list_models(), indent=2))\n print(\"Capability Mapping:\", model_manager.list_capabilities())\n except Exception as e:\n print(f\"Failed to initialize LLMManager: {e}\")\n return\n\n # --- 2. 实例化 Agents (使用 ModelManager 获取模型) ---\n try:\n # 获取默认模型用于基础任务\n grok = model_manager.get_model(\"xai_grok\") # 获取 ID 由 config 或第一个注册的决定\n deepseek_v3 = model_manager.get_model(\"deepseek_v3\") # 获取 DeepSeek 模型\n # 创建Tavily搜索工具\n tavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n )\n\n # 确保 ReactAgent 使用与 Supervisor 兼容的状态 (例如 BasicAgentState)\n # 或者 Supervisor 能够处理不同类型的子 Agent 状态返回\n researcher_system_prompt = \"\"\"You are a research expert. Use available tools to find the most up-to-date information to answer the user's query. You have access to a 'tavily_search_results_json' tool.\"\"\"\n\n research_agent = ReactAgent(\n name=\"research_expert\", \n tools=[tavily_search],\n description=\"Research expert with access to Tavily search.\",\n model=deepseek_v3,\n prompt=researcher_system_prompt,\n ) \n \n all_agents = [research_agent] # 只包含一个子 Agent 用于测试\n\n # --- 实例化 Supervisor (使用 PlanningAgentState) ---\n supervisor = SupervisorAgent(\n agents=all_agents,\n model=deepseek_v3, # Supervisor 使用的 LLM\n state_schema=PlanningAgentState, # 明确 Supervisor 使用 Planning 状态\n # enable_planning=True, # 不再需要此参数，因为 state_schema 暗示了规划\n include_agent_name=\"inline\" # 推荐\n # checkpointer=... # 添加 Checkpointer 以测试持久化\n )\n except Exception as e:\n print(f\"Failed to initialize agents or supervisor: {e}\")\n import traceback\n traceback.print_exc()\n return\n\n # --- 获取用户输入 ---\n topic = input(\"Please enter the initial request for the supervisor: \")\n if not topic:\n print(\"No request entered. Exiting.\")\n return\n\n # --- 准备初始状态 (使用 PlanningAgentState) ---\n initial_graph_state: PlanningAgentState = {\n \"messages\": [HumanMessage(content=topic)], # 确保是 HumanMessage 对象\n \"plan\": None, # 初始没有计划\n \"error\": None\n }\n\n # --- 运行测试 ---\n await run_supervisor_test(supervisor, initial_graph_state)\n\n\nif __name__ == \"__main__\":\n try:\n asyncio.run(main())\n except KeyboardInterrupt:\n print(\"\\nExecution interrupted by user.\")\n except Exception as e:\n print(f\"\\nAn unexpected top-level error occurred: {e}\")\n traceback.print_exc()"
},
{
"path": "examples/state_based_supervisor_examples/03_multi_agents.py",
"content": "# main.py (Multi-Agent Test with State-Based Supervisor)\n\nimport asyncio\nimport json\nimport os\nimport re\nimport time\nimport traceback # 导入 traceback\nfrom datetime import datetime\nfrom typing import Dict, Any, Optional, List, Literal, cast\n\n# --- LangChain / LangGraph / OpenAI Imports ---\nfrom langchain_core.messages import HumanMessage, AIMessage, BaseMessage, ToolMessage\n\n\n\n# --- Agent 和工具导入 (确保路径正确) ---\ntry:\n from core.agents.sb_supervisor_agent import SupervisorAgent # 替换为你的 SupervisorAgent 类路径\n from core.agents.state_based_supervisor.state_schema import PlanningAgentState\n from core.agents.base.react_agent import ReactAgent # 导入 ReactAgent 基类\n\n # 导入所有重构后的子 Agent 类\n from core.agents.sub_agents.research_agent import ResearchAgent # 假设路径\n from core.agents.sub_agents.coder_agent import CoderAgent # 假设路径\n from core.agents.sub_agents.reporter_agent import ReporterAgent # 假设路径\n from core.agents.sub_agents.designer_agent import DesignerAgent # 假设路径\n from core.agents.sub_agents.data_analyst_agent import DataAnalystAgent # 假设路径\n\n # 导入工具注册表函数和枚举\n from core.tools.registry import get_tools_by_category, ToolCategory, register_tool # 导入 register_tool\n from core.llm.llm_manager import LLMManager # LLM 管理器\n # 导入特定工具实例或类 (如果 Registry 没有预注册所有工具)\n from langchain_community.tools.tavily_search import TavilySearchResults # 示例\n # from core.tools.e2b_tool import E2BCodeInterpreterTool # 示例\n # from core.tools.replicate_flux_tool import ReplicateFluxImageTool # 示例\n\n # --- 确保工具已注册 ---\n # 运行 registry 初始化 (通常在 core/tools/__init__.py 中完成)\n try:\n import core.tools # 尝试导入以触发 __init__.py 中的注册\n print(\"Tool registry potentially initialized.\")\n except ImportError:\n print(\"Warning: Could not import 'core.tools' to initialize registry.\")\n except Exception as reg_err:\n print(f\"Error during tool registry initialization: {reg_err}\")\n \n # (可选) 在这里可以检查或手动注册缺失的核心工具\n # Example: Check and register Tavily if not present\n if not any(getattr(t, 'name', '') == 'tavily_search_results_json' for t in get_tools_by_category(ToolCategory.SEARCH)):\n try:\n print(\"Attempting to register TavilySearchResults...\")\n tavily_tool = TavilySearchResults(max_results=3)\n register_tool(tavily_tool, ToolCategory.SEARCH)\n except Exception as e:\n print(f\"Warning: Failed to register TavilySearchResults manually: {e}\")\n \n # ... 检查并注册其他必要的工具 ...\n\n\nexcept ImportError as e:\n print(f\"Error importing agent/tool components: {e}\")\n print(\"Please ensure all agent/tool class paths and registry setup are correct.\")\n exit(1)\n\n\n# --- 助手函数 ---\ndef slugify(text: str) -> str:\n \"\"\"Converts text to a safe filename part.\"\"\"\n # ... (保持不变) ...\n if not text: return \"no_topic\"\n text = text.lower(); text = re.sub(r'\\s+', '_', text)\n text = re.sub(r'[^\\w\\-]+', '', text); text = text.strip('_')\n return text[:100] if text else \"sanitized_topic\"\n\n# --- 主研究/测试函数 ---\nasync def run_supervisor_test(supervisor_agent: SupervisorAgent, initial_state: Dict[str, Any]):\n \"\"\"Executes the supervisor graph using ainvoke and prints the final state.\"\"\"\n\n print(\"\\n--- Starting Supervisor Graph Execution ---\")\n initial_query = initial_state.get(\"messages\", [{}])[0].content if initial_state.get(\"messages\") and hasattr(initial_state.get(\"messages\")[0], 'content') else \"N/A\"\n print(f\"Initial Query: '{initial_query}'\")\n print(\"-\" * 30)\n\n config = {\"recursion_limit\": 100} # 保持较高的递归限制\n\n final_state: Optional[Dict[str, Any]] = None\n error_occurred: Optional[Exception] = None\n\n try:\n print(\"--- Invoking Supervisor Agent (ainvoke) ---\")\n # 直接调用 ainvoke 获取最终状态\n final_state = await supervisor_agent.ainvoke(initial_state, config=config)\n print(\"--- Supervisor Agent Invocation Complete ---\")\n\n # --- 错误处理 ---\n except Exception as e: error_occurred = e; print(f\"\\n!!! Error during graph execution: {e}\"); traceback.print_exc()\n\n\n # --- 处理最终状态 ---\n if error_occurred: print(\"\\n--- Graph Execution INTERRUPTED or FAILED ---\")\n else: print(\"\\n--- Graph Execution Finished ---\")\n\n if not final_state:\n print(\"Error: No final state available (Execution might have failed early).\")\n # 尝试从 checkpointer 获取最后状态 (如果配置了)\n # ... (checkpoint retrieval logic - optional) ...\n return None\n\n print(\"\\n--- FINAL STATE ---\")\n # 打印错误 (如果在状态中记录了)\n if final_state.get(\"error\"): print(f\"\\nERROR RECORDED IN STATE: {final_state['error']}\")\n\n # 打印计划\n final_plan = final_state.get('plan')\n if final_plan and isinstance(final_plan, dict):\n print(\"\\nFinal Plan State:\")\n print(json.dumps(final_plan, indent=2, default=str))\n else: print(\"\\nFinal Plan State: Not available or not generated.\")\n\n # 打印最终消息历史\n final_messages = final_state.get(\"messages\", [])\n if final_messages and isinstance(final_messages, list):\n print(\"\\nFinal Message History (Last 10):\")\n for m in final_messages[-10:]:\n try:\n if hasattr(m, 'pretty_print'): m.pretty_print()\n else: print(json.dumps(m, indent=2, default=str)) # Fallback\n print(\"-\" * 10)\n except Exception as print_err: print(f\"Error printing final message: {print_err}\")\n else: print(\"\\nFinal Message History: Empty.\")\n\n # --- 保存最终报告 (如果 Reporter Agent 被调用且成功) ---\n # 检查最后一条消息是否来自 Reporter\n final_report_content = None\n if final_messages and isinstance(final_messages[-1], AIMessage) and final_messages[-1].name == \"reporter_expert\":\n final_report_content = final_messages[-1].content\n print(\"\\n--- Final Report Found from Reporter Agent ---\")\n\n if not error_occurred and final_report_content and isinstance(final_report_content, str) and \"Failed\" not in final_report_content:\n print(\"\\n--- Saving Final Output to Markdown ---\")\n try:\n markdown_content = final_report_content\n # 获取原始请求作为文件名基础\n initial_query_text = final_state.get('messages', [{}])[0].content if final_state.get('messages') and hasattr(final_state.get('messages')[0], 'content') else 'unknown_request'\n topic_slug = slugify(initial_query_text)\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n filename = f\"multi_agent_report_{topic_slug}_{timestamp}.md\"\n\n script_dir = os.path.dirname(os.path.abspath(__file__))\n output_dir = os.path.join(script_dir, \"Output\")\n os.makedirs(output_dir, exist_ok=True)\n filepath = os.path.join(output_dir, filename)\n\n with open(filepath, \"w\", encoding=\"utf-8\") as f: f.write(markdown_content)\n print(f\"Successfully saved output to: {filepath}\")\n except Exception as e: print(f\"Error saving output to Markdown: {e}\")\n elif error_occurred: print(\"\\nFinal Report: Not saved due to execution error.\")\n else: print(\"\\nFinal Report: Not generated or not found.\")\n\n print(\"\\n--- END OF TEST ---\")\n return final_state\n\n\n# --- Main Execution Block ---\nasync def main():\n # --- 1. 初始化 LLM 管理器 ---\n try:\n model_manager = LLMManager()\n print(\"Registered Models:\", json.dumps(model_manager.list_models(), indent=2))\n except Exception as e:\n print(f\"Failed to initialize LLMManager: {e}\")\n return\n\n # --- 2. 实例化所有 Agents ---\n try:\n # 获取模型实例\n # 确保 'deepseek_v3' 和 'gpt-4o' 是你 LLMManager 中有效的 ID\n deepseek_model = model_manager.get_model(\"deepseek_v3\")\n gpt4o_model = model_manager.get_model(\"openai_gpt4o\") # 多模态模型\n\n # 实例化 ResearchAgent\n research_agent = ResearchAgent(\n model=deepseek_model,\n )\n\n # 实例化 CoderAgent\n coder_agent = CoderAgent(\n model=deepseek_model,\n )\n\n # 实例化 ReporterAgent\n reporter_agent = ReporterAgent(\n model=deepseek_model\n )\n\n # 实例化 DesignerAgent\n designer_agent = DesignerAgent(\n model=gpt4o_model,\n )\n\n # 实例化 DataAnalystAgent\n data_analyst_agent = DataAnalystAgent(\n model=deepseek_model,\n )\n\n # --- 3. 组合 Agent 列表 ---\n all_agents = [\n research_agent,\n coder_agent,\n reporter_agent,\n designer_agent,\n data_analyst_agent,\n ]\n\n # --- 4. 实例化 Supervisor ---\n supervisor = SupervisorAgent(\n agents=all_agents,\n model=deepseek_model, # Supervisor 自身使用的模型\n # model = gpt4o_model,\n state_schema=PlanningAgentState,\n include_agent_name=\"inline\"\n # checkpointer=... # 可选: 添加 Checkpointer 实现持久化\n )\n\n except Exception as e:\n print(f\"Failed to initialize agents or supervisor: {e}\")\n traceback.print_exc()\n return\n\n # --- 5. 获取用户输入 ---\n topic = input(\"Please enter the initial request for the supervisor: \")\n if not topic:\n print(\"No request entered. Using default topic.\")\n topic = \"\"\"我需要获取法国巴黎当前的实时气温。请按以下步骤操作：\n1. 首先，帮我调研一个可以免费获取巴黎当前天气数据的 API (例如 Open-Meteo, WeatherAPI.com 或其他类似的)，重点是找到获取当前气温的 API 端点(endpoint URL)以及如何构造请求（如果可能，选择不需要 API key 的）。\n2. 然后，编写一个 Python 脚本，使用 'requests' 库来调用上一步找到的 API 端点，并从中提取出巴黎当前的温度（摄氏度）。\n3. 使用你的代码执行工具来运行这个 Python 脚本。\n4. 最后，告诉我你找到的当前巴黎温度是多少。\"\"\"\n\n # --- 6. 准备初始状态 ---\n initial_graph_state: PlanningAgentState = {\n \"messages\": [HumanMessage(content=topic)], \n \"plan\": None,\n \"error\": None\n }\n\n # --- 7. 运行测试 ---\n await run_supervisor_test(supervisor, initial_graph_state)\n\n\nif __name__ == \"__main__\":\n try:\n asyncio.run(main())\n except KeyboardInterrupt:\n print(\"\\nExecution interrupted by user.\")\n except Exception as e:\n print(f\"\\nAn unexpected top-level error occurred: {e}\")\n traceback.print_exc()"
},
{
"path": "examples/web_agents/README.md",
"content": "# Web Agents\n\n这个目录包含可以通过web界面加载的代理示例。每个子目录代表一个独立的代理实现，可以被server.py动态加载。\n\n## 目录结构\n\n每个代理应遵循以下结构：\n\n```\nagent_name/\n __init__.py # 包含get_graph()函数，返回编译好的LangGraph\n README.md # 代理的说明文档\n```\n\n## 接口规范\n\n每个代理必须实现以下接口：\n\n```python\ndef get_graph():\n \"\"\"返回编译好的LangGraph实例\"\"\"\n # 构建并返回图\n return compiled_graph\n```"
},
{
"path": "examples/web_agents/README_SPEC.md",
"content": "# Web Agent 开发规范\n\n## 1. 概述\n\n本规范旨在统一Web Agent的开发流程和命名约定，确保前后端协同工作，避免出现前端组件无法正确显示后端数据的问题。本文档基于实际开发经验，特别强调前后端节点命名一致性的重要性。\n\n## 2. 节点命名规范\n\n## 2. 前后端交互核心机制\n\n### 2.1 关键概念\n\n- **节点名称匹配**: 前端渲染组件时，会根据后端节点的名称来选择对应的组件进行渲染\n- **状态数据结构**: 后端节点生成的状态数据必须符合前端组件期望的结构\n- **渲染函数**: 前端的`renderNode`函数是连接后端节点和前端组件的关键桥梁\n\n### 2.2 渲染流程\n\n1. 后端节点执行并生成状态数据\n2. 前端通过`useLangGraphAgent`钩子接收节点数据\n3. 前端的`renderNode`函数根据节点名称选择对应组件\n4. 组件根据状态数据进行渲染\n\n## 3. 节点命名规范\n\n### 3.1 关键节点命名\n\n所有Web Agent必须在图结构中包含处理消息的节点，这些节点名称必须与前端`renderNode`函数中的case语句匹配：\n\n```python\n# 后端节点命名 - 必须与前端renderNode函数中的case匹配\nbuilder.add_node(\"agent\", agent_function) # 或其他在前端已注册的节点名称\n```\n\n**重要提示**: 前端`page.tsx`中的`renderNode`函数定义了可识别的节点名称。目前支持的节点名称有：\n- `__start__`\n- `agent` (替代了原来的`chatbot`)\n- `weather`\n- `reminder`\n- `research`\n- `search`\n- `report`\n\n如果后端使用了其他节点名称，必须在前端的`renderNode`函数中添加对应的case语句。\n\n### 3.2 状态字段命名\n\n- 状态字段名称应与前端组件期望的字段名称保持一致\n- 使用蛇形命名法（snake_case）命名状态字段\n- 复杂数据结构应使用数组形式，即使只有一个元素\n\n### 3.3 必要的状态字段\n\n每个Web Agent必须在`agent-types.ts`文件中定义其状态接口，并确保后端发送的状态与此接口匹配：\n\n```typescript\nexport interface AgentState extends WithMessages {\n // 定义Agent特有的状态字段\n weather_forecast?: WeatherForecast[];\n research_status?: ResearchStatus[];\n // 其他状态字段\n}\n```\n\n## 4. 前端组件规范\n\n### 4.1 组件结构\n\n- 主组件应根据节点名称渲染不同的子组件\n- 子组件应检查所需状态字段是否存在，并提供合理的默认行为\n\n```typescript\nexport default function renderNode(checkpoint, node) {\n switch (node.name) {\n case '__start__':\n case 'agent': // 注意：这里使用'agent'替代了原来的'chatbot'\n return ;\n case 'weather':\n return ;\n // 其他节点类型\n default:\n return null;\n }\n}\n```\n\n### 4.2 组件注册\n\n所有Web Agent的节点组件必须在`page.tsx`的`renderNode`函数中正确注册：\n\n```typescript\nconst renderNode = (checkpoint, node) => {\n switch (node.name) {\n // 确保这里的节点名称与后端图定义中的节点名称一致\n case '__start__':\n case 'agent': // 注意：这里使用'agent'替代了原来的'chatbot'\n return ;\n case 'weather':\n return ;\n case 'reminder':\n return ;\n case 'research':\n case 'search':\n case 'report':\n return ;\n // 添加新节点类型的渲染逻辑\n default:\n return null;\n }\n}\n```\n\n## 5. 后端图结构规范\n\n### 5.1 节点函数\n\n- 节点函数应使用适当的参数来处理状态\n- 消息处理必须在与前端匹配的节点中进行\n\n```python\nasync def agent(state): # 注意：这里使用'agent'替代了原来的'chatbot'\n # 处理消息并返回结果\n return {\"messages\": [...]} # 必须包含messages字段\n```\n\n### 5.2 图构建\n\n- 图必须包含与前端匹配的节点，用于处理消息\n- 必须实现`get_graph()`函数返回编译好的图实例\n\n```python\ndef get_graph():\n \"\"\"返回编译好的LangGraph实例\"\"\"\n builder = StateGraph()\n builder.add_node(\"agent\", agent) # 注意：这里使用'agent'替代了原来的'chatbot'\n # 添加边和其他节点\n graph = builder.compile(checkpointer=MemorySaver())\n return graph\n```\n\n## 6. 开发流程\n\n### 6.1 新建Web Agent流程\n\n1. 在`examples/web_agents/`下创建新的Agent目录\n2. 创建`graph.py`文件，实现Agent的图结构，确保节点名称与前端`renderNode`函数中的case语句匹配\n3. 在`web/app/chat/[id]/agent-types.ts`中添加Agent所需的状态接口\n4. 在`web/app/chat/[id]/components/`下创建Agent的组件\n5. 在`web/app/chat/[id]/page.tsx`的`renderNode`函数中注册Agent的节点组件（如果使用新的节点名称）\n\n### 6.2 测试验证\n\n在提交代码前，必须进行以下测试：\n\n1. 确认后端图结构中的节点名称与前端`renderNode`函数中的case语句匹配\n2. 验证前端组件能正确渲染不同类型的节点\n3. 检查状态字段名称与前端组件期望的字段名称一致\n\n## 7. 常见问题与解决方案\n\n### 7.1 前端不显示消息问题\n\n如果前端不显示消息内容，请检查：\n\n1. 后端图结构中的节点名称是否与前端`renderNode`函数中的case语句匹配\n2. 前端`renderNode`函数是否正确处理了对应的节点名称\n3. 消息是否正确包含在state的messages字段中\n\n### 7.2 状态更新不生效\n\n确保状态更新时，字段名称与前端期望的字段名称一致，并且数据结构符合前端组件的预期。\n\n### 7.3 添加新节点类型\n\n如果需要添加新的节点类型，必须：\n\n1. 在后端图结构中定义新节点\n2. 在前端`page.tsx`的`renderNode`函数中添加对应的case语句\n3. 创建新节点对应的前端组件\n4. 在`agent-types.ts`中添加新节点所需的状态接口\n\n---\n\n遵循本规范可以有效避免前后端不一致导致的显示问题，提高Web Agent的开发效率和质量。"
},
{
"path": "examples/web_agents/__init__.py",
"content": "# Web Agents Package\n# This package contains web agents that can be loaded by the server"
},
{
"path": "examples/web_agents/research_assistant/README.md",
"content": "# 研究助手\n\n这是一个强大的研究助手代理，可以帮助用户进行在线研究、信息收集和报告生成。\n\n## 功能\n\n- 在线搜索信息\n- 提取和总结网页内容\n- 生成研究报告\n- 实时显示研究进度\n\n## 使用方法\n\n用户可以通过自然语言与代理交互，例如：\n\n- \"帮我研究人工智能在医疗领域的应用\"\n- \"查找关于气候变化的最新研究\"\n- \"总结量子计算的基本原理\"\n\n## 技术实现\n\n该代理使用LangGraph构建，结合了Supervisor和React模式，包含以下节点：\n\n- supervisor: 协调整个研究流程\n- search: 执行在线搜索\n- extract: 提取网页内容\n- analyze: 分析收集的信息\n- report: 生成研究报告\n\n研究过程中会实时更新状态，让用户了解当前进度。"
},
{
"path": "examples/web_agents/research_assistant/__init__.py",
"content": "# Research Assistant Agent\n# This module provides a research assistant agent that can crawl websites and extract content\n\nfrom .graph import get_graph\n\n__all__ = [\"get_graph\"]"
},
{
"path": "examples/web_agents/research_assistant/graph.py",
"content": "from langgraph.prebuilt import create_react_agent\nfrom langchain_openai import ChatOpenAI\nfrom typing import Dict, Any\nfrom dotenv import load_dotenv\nfrom langchain_community.tools import TavilySearchResults\nfrom langgraph.checkpoint.memory import MemorySaver\nfrom core.tools.e2b_tool import E2BCodeInterpreterTool\nfrom core.tools.registry import register_tool, ToolCategory\nfrom core.llm.llm_manager import LLMManager\n\nload_dotenv() # 自动加载 .env 文件\n# 初始化大模型\nmodel = LLMManager().get_model(\"deepseek_v3\")\n\n# 创建Tavily搜索工具\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n)\n\n# 创建E2B代码解释器工具\ne2b_code_interpreter = E2BCodeInterpreterTool()\n\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[tavily_search, e2b_code_interpreter],\n name=\"research_expert\",\n # Prompt 告诉它是一个研究型 Agent，可调用 tavily_search 和 e2b_code_interpreter\n prompt=(\n \"你是一位世界级的研究专家和数据分析师，擅长信息检索和数据分析。你有两个强大的工具可以使用：\\n\"\n \"1. 'tavily_search_results_json'：用于搜索网络获取实时信息\\n\"\n \"2. 'e2b_code_interpreter'：用于执行Python代码，支持数据分析和可视化\\n\\n\"\n \"当面对问题时，请遵循以下方法论：\\n\"\n \"1. 分析问题：理解用户的需求和问题本质\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及是否需要进行数据分析\\n\"\n \"3. 执行搜索：使用tavily_search_results_json工具获取最新信息\\n\"\n \"4. 数据分析：如果需要，使用e2b_code_interpreter工具编写和执行Python代码进行数据分析和可视化\\n\"\n \"5. 综合信息：将搜索结果和数据分析结果综合成一个连贯的回答\\n\\n\"\n \"重要提示：\\n\"\n \"- 对于信息检索任务，使用tavily_search_results_json工具，并在回答中引用来源URL\\n\"\n \"- 对于数据分析和可视化任务，使用e2b_code_interpreter工具执行Python代码\\n\"\n \"- 在使用代码解释器时，确保导入必要的库（如pandas, matplotlib, numpy等）\\n\"\n \"- 在代码中添加详细注释，解释关键步骤\\n\"\n \"- 执行代码后，解释结果含义和见解\"\n ),\n checkpointer=MemorySaver(),\n)\n\n\ndef get_graph():\n return research_agent"
},
{
"path": "examples/web_agents/weather_agent/README.md",
"content": "# 天气代理\n\n这是一个简单的天气查询代理，可以回答用户关于天气的问题，并提供天气预报信息。\n\n## 功能\n\n- 查询当前天气\n- 创建提醒\n\n## 使用方法\n\n用户可以通过自然语言与代理交互，例如：\n\n- \"今天北京的天气怎么样？\"\n- \"帮我设置一个提醒，明天早上8点去开会\"\n\n## 技术实现\n\n该代理使用LangGraph构建，包含以下节点：\n\n- chatbot: 处理用户输入并生成回复\n- weather: 处理天气查询请求\n- reminder: 处理提醒创建请求"
},
{
"path": "examples/web_agents/weather_agent/__init__.py",
"content": "# Weather Agent Example\n# This is a simple weather agent that can be loaded by the server\n\nimport operator\nfrom typing import Literal, TypedDict, Any, Annotated\nfrom dotenv import load_dotenv\nfrom langchain_openai import ChatOpenAI\nfrom langgraph.graph import StateGraph, MessagesState, START, END\nfrom langgraph.checkpoint.memory import MemorySaver\nfrom langgraph.types import StreamWriter, interrupt, Send\nfrom langchain_core.messages import ToolMessage\nfrom langchain_core.tools import tool\nimport random\nimport asyncio\n\nload_dotenv()\n\n\nclass Weather(TypedDict):\n location: str\n search_status: str\n result: str\n\n\nclass State(MessagesState):\n weather_forecast: Annotated[list[Weather], operator.add]\n\n\nclass WeatherInput(TypedDict):\n location: str\n tool_call_id: str\n\n\nclass ToolNodeArgs(TypedDict):\n name: str\n args: dict[str, Any]\n id: str\n\n\n@tool\nasync def weather_tool(query: str) -> str:\n \"\"\"Call to get current weather\"\"\"\n return \"Sunny\"\n\n\n@tool\nasync def create_reminder_tool(reminder_text: str) -> str:\n \"\"\"Call to create a reminder\"\"\"\n return \"Reminder created\"\n\n\nasync def weather(input: WeatherInput, writer: StreamWriter):\n location = input[\"location\"]\n tool_call_id = input[\"tool_call_id\"]\n\n # Send custom event to the client. It will update the state of the last checkpoint and all child nodes.\n # Note: if there are multiple child nodes (e.g. parallel nodes), the state will be updated for all of them.\n writer({\"weather_forecast\": [\n {\"location\": location, \"search_status\": f\"Checking weather in {location}\"}]})\n\n await asyncio.sleep(2)\n weather = random.choice([\"Sunny\", \"Cloudy\", \"Rainy\", \"Snowy\"])\n\n return {\"messages\": [ToolMessage(content=weather, tool_call_id=tool_call_id)], \"weather_forecast\": [{\"location\": location, \"search_status\": \"\", \"result\": weather}]}\n\n\nasync def reminder(input: ToolNodeArgs):\n res = interrupt(input['args']['reminder_text'])\n\n tool_answer = \"Reminder created.\" if res == 'approve' else \"Reminder creation cancelled by user.\"\n\n return {\"messages\": [ToolMessage(content=tool_answer, tool_call_id=input[\"id\"])]}\n\n\nasync def chatbot(state: State):\n llm = ChatOpenAI(\n model=\"gpt-4o-mini\").bind_tools([weather_tool, create_reminder_tool])\n response = await llm.ainvoke(state[\"messages\"])\n return {\"messages\": [response]}\n\n\ndef tool_router(state: State) -> Literal[\"weather\", \"reminder\", \"__end__\"]:\n messages = state[\"messages\"]\n last_message = messages[-1]\n if last_message.tool_calls:\n if last_message.tool_calls[0][\"name\"] == \"weather_tool\":\n return \"weather\"\n elif last_message.tool_calls[0][\"name\"] == \"create_reminder_tool\":\n return \"reminder\"\n return \"__end__\"\n\n\n# Chatbot node router. Based on tool calls, creates the list of the next parallel nodes.\ndef assign_tool(state: State) -> Literal[\"weather\", \"reminder\", \"__end__\"]:\n messages = state[\"messages\"]\n last_message = messages[-1]\n if last_message.tool_calls:\n send_list = []\n for tool in last_message.tool_calls:\n if tool[\"name\"] == 'weather_tool':\n send_list.append(\n Send('weather', {'location': tool['args']['query'], 'tool_call_id': tool['id']}))\n elif tool[\"name\"] == 'create_reminder_tool':\n send_list.append(Send('reminder', tool))\n return send_list if len(send_list) > 0 else \"__end__\"\n return \"__end__\"\n\n\ndef get_graph():\n \"\"\"Return the compiled graph for this agent\"\"\"\n builder = StateGraph(State)\n\n builder.add_node(\"chatbot\", chatbot)\n builder.add_node(\"weather\", weather)\n builder.add_node(\"reminder\", reminder)\n\n builder.add_edge(START, \"chatbot\")\n builder.add_conditional_edges(\"chatbot\", assign_tool)\n builder.add_edge(\"weather\", \"chatbot\")\n builder.add_edge(\"reminder\", \"chatbot\")\n\n builder.add_edge(\"chatbot\", END)\n\n memory = MemorySaver()\n return builder.compile(checkpointer=memory)"
},
{
"path": "instructions/00.Langgraph 和 React Agent.md",
"content": "# 一、LangGraph 的核心思想\n\nLangGraph 是一个可以让开发者以**图（Graph）**的方式来编排对话式AI流程的库，提供了以下能力：\n\n1. **状态驱动**：在传统的对话模型中，我们经常需要维护对话上下文、剩余步骤等各种内部变量。LangGraph 将这些变量统一到一个“状态(State)”里，并约定任何节点的输入/输出都以“状态(State)”的形式表示。\n \n2. **可视化执行流**：LangGraph 将对话/工具调用/自定义逻辑封装成“节点(Node)”与“边(Edge)”。当图被编译后，执行流会在节点之间穿梭，处理对话消息、调用工具、终止或转向某些分支。\n\n3. **可组合**：你可以把一个复杂的对话逻辑拆分为多个可复用的子图，每个子图都可以独立进行单元测试或复用在更大的图中。\n\n4. **多步思考 + 工具调用**：通过 ReAct Agent（一个经典的多步推理+工具调用范式），LangGraph 可以帮助你自动管理**多次**调用语言模型及其衍生工具的过程——只要你把“工具”注册到图里。\n\n在使用时，你基本会经历如下步骤：\n\n1. **定义状态模式（state schema）**：说明 state 中必须包含哪些字段（如：对话消息 `messages`，剩余可用步骤 `remaining_steps`，等等）。\n2. **定义节点（Node）**：比如一个负责调用LLM的节点、一个负责执行特定工具的节点、或者你自定义的Python逻辑节点。\n3. **连接边（Edges）**：决定每个节点之后，下一步走到哪个节点；也可以做条件分支或循环。\n4. **编译图（Compile）**：LangGraph 会把你的“编排逻辑”转换为一个 LangChain 兼容的“可调用对象(CompiledGraph)”。\n5. **执行或流式执行**：可以直接一次性 `graph.invoke(...)` 得到最终结果，也可以使用 `graph.stream(...)` 流式获取每个“阶段性状态（partial state）” 。\n\n---\n\n# 二、LangGraph 核心概念详解\n\nLangGraph 构建的是一个\"流程图\"，每个智能体（agent）或功能模块（tool调用、分支逻辑）都是这个流程图的一个节点（node）。让我们深入理解其中的核心概念：\n\n## 2.1 Graph：有状态的数据流图\n\nGraph 是整个 Agent 系统的执行框架，定义了哪些模块怎么串联、怎么流转。你构建的 graph 是一个有向图：\n\n```python\ngraph = StateGraph(state_schema=MyState)\n\n# 添加节点\ngraph.add_node(\"supervisor\", supervisor_runnable)\ngraph.add_node(\"writer\", writer_runnable)\n\n# 添加边来连接节点\ngraph.add_edge(\"supervisor\", \"writer\")\ngraph.add_edge(\"writer\", \"supervisor\")\n```\n\nLangGraph 根据这些连接关系来控制执行流程，决定在某个节点执行完后下一步应该去哪里。\n\n## 2.2 Node：图中的\"一个执行单元\"\n\n每个 node 是图中的一个处理模块（通常就是一个智能体）。它接受一个输入 state，做点事情，然后返回一个新的 state：\n\n```python\ndef my_node(state: dict) -> dict:\n # 处理 state 中的数据\n new_state = state.copy()\n # 修改状态内容\n new_state[\"some_key\"] = \"new_value\"\n return new_state\n```\n\n节点可以是：\n- 函数（同步或异步）\n- LLM Agent（如 create_react_agent 返回的）\n- 包装后的 Agent（如 MemorySlidingReactAgent）\n\n## 2.3 State：每一轮节点处理的输入/输出\n\n每轮调用，LangGraph 会传递一个 \"state\"（字典类型）给当前节点。这个 state 中可以包含：\n- `messages`: 当前对话历史（主上下文）【默认】\n- `memory`: 你自定义的长期记忆（可以注入系统提示）\n- `todo_list`, `current_task`: 其他任务状态\n- 任何你自定义的字段\n\n每个节点执行后，返回新的 state：\n\n```python\ndef writer(state):\n new_msg = generate_chapter(state[\"current_task\"])\n state[\"messages\"].append({\"role\": \"assistant\", \"content\": new_msg})\n return state\n```\n\n## 2.4 Runnable：Node 的运行接口\n\nLangGraph 要求，每个节点（node）必须是可以运行的，也就是说：你交给 `add_node()` 的对象必须有 `.invoke(state)` 或 `.ainvoke(state)` 方法。\n\n比如：\n- 函数本身（它会自动包装成 Runnable）\n- Agent（React agent 本身就支持 `.invoke`）\n- `RunnableCallable(...)` 是 LangGraph 用来显式包装函数的工具\n\n举个例子：\n\n```python\ndef my_function(state: dict) -> dict:\n # 处理逻辑\n return state\n\nrunnable = RunnableCallable(my_function, async_version)\ngraph.add_node(\"writer\", runnable)\n```\n\n## 2.5 执行流程\n\nLangGraph 的执行流程大致如下：\n\n```\nLangGraph Graph:\n [START]\n ↓\n [Supervisor Node]\n ↓\n [Writer Node]\n ↓\n [Supervisor Node]\n ↓\n [END]\n```\n\n每次节点执行时：\n1. 传入当前 state\n2. `.invoke(state)` 被调用\n3. 返回更新后的 state\n4. 下一节点接着执行\n\n## 2.6 概念类比\n\n| LangGraph 概念 | 类比 |\n|----------------|------|\n| Graph | 工作流程图/数据流图 |\n| Node | 每个处理步骤/智能体 |\n| State | 当前上下文与执行状态（黑匣子） |\n| Runnable | 每个节点\"能被执行\"的接口定义 |\n\n---\n\n# 三、ReAct Agent 与 create_react_agent 概念\n\n## 3.1 什么是 ReAct Agent\n\n“ReAct” 是一种典型的LLM多步推理与工具调用策略。它主要包含两部分：\n\n1. **Reasoning**：先让语言模型（LLM）进行一步推理，产出一个潜在的思考过程以及可能的工具调用。\n2. **Acting**：如果模型说“我要调用某个工具”，则执行该工具，得到结果，再把结果加入对话，然后让模型再次 Reason，看看是否还需要再调用工具，或输出最后的答案。\n\n这个循环可以**多次往返**，直到模型不再调用工具，输出最终结果。\n\n## 3.2 create_react_agent 做了什么\n\n`create_react_agent(...)` 是 LangGraph 中的一个快捷方法，用于**快速创建**一个可执行的“ReAct风格”图（Graph）：\n\n- **自动添加“agent节点”**：用来调用你的语言模型（并在对话中发出可能的工具调用）。\n- **自动添加“tools节点”**：如果 agent 的输出中含有工具调用（tool_calls），则会交给 tools 节点逐个执行，并把执行结果以 `ToolMessage` 的形式返回到对话中。\n- **自动在 agent ↔ tools 之间连线**：只要 agent 产生了工具调用，就进入 tools；tools 执行完返回消息后，再回到 agent；直到不再有工具调用为止。\n- **可选 structured output**：如果你传入了 `response_format` 参数，LangGraph 会在最后一步生成一个结构化输出(“JSON Schema”、“Pydantic”、“OpenAI function schema”等)，以便你获取可解析的最终结果。\n- **控制“剩余步骤”**：Agent 每次回答后会检查是否还可以继续调用工具，或者是否需要中止并返回错误（“抱歉，需要更多步骤”）。\n\n因此，调用 `create_react_agent(...)` 得到的结果，是一个**已经配置好**的 “CompiledGraph”。这个图中带有 “agent” 节点（LLM） 和 “tools” 节点（调用工具），以及检查**是否还有工具要调**的逻辑。你可以直接拿这个对象执行，获得一个 ReAct 流程的多轮对话+工具使用。\n\n---\n\n# 四、执行流程：从输入到输出\n\n创建好 ReAct 图后，你给它一个输入状态（最少包含 `\"messages\"`，如 `{\"messages\": [(\"user\", \"Hello!\")]}`）。执行过程大体是：\n\n1. **entry point: \"agent\"** \n 进入 agent 节点，它会从 state[\"messages\"] 中取出消息，交给 LLM 生成一个 AIMessage。如果 AIMessage 包含 tool_calls，那么 state 会更新多一些字段，比如 `messages` 后面多了这个 AIMessage。\n\n2. **检查是否要调用工具** \n - 如果 `tool_calls` 不为空，则顺着 edges 进入 \"tools\" 节点。\n - 如果没有 tool_calls，则表示 agent 没有想调用任何工具 -> 流程会判断是否要去 “generate_structured_response” 或 “END”。\n \n3. **tools 节点执行** \n \"tools\" 节点会去匹配 agent 要调用的工具，比如： \n ```json\n {\n \"name\": \"search_tool\",\n \"args\": {\"query\": \"something\"},\n \"id\": \"call_abc123\"\n }\n ```\n 然后运行相应的 Python 函数，得到结果后，包装成 `ToolMessage`，附加回 state[\"messages\"] 列表里。\n - 如果 agent 一次性请求了多个工具，在 v1 版本中则会并行执行，再把返回结果依次追加到 messages 里。\n - 在 v2 版本中，LangGraph 会拆分 tool_calls 分批执行。\n\n4. **回到 agent** \n 现在 agent 再次拿到新的 state[\"messages\"]（多了“ToolMessage”），就会针对最新的对话上下文重新进行思考——是否要再调用别的工具、或者是否直接产出最终回答？\n\n5. **循环，直到不再调用工具** \n 只要 AIMessage 继续发出 tool_calls，就进入 Tools 节点；Tools 执行完再回到 Agent 节点。这一过程可能多次往返。（如果你设定了 `remaining_steps`，LangGraph 在每一轮都会减少1，直到不足时终止或报错，避免死循环。）\n\n6. **可选：结构化输出** \n 在最后如果 `response_format` 存在，图会跳到“generate_structured_response”节点，再次对(几乎)所有对话做一次 LLM 调用，要求 LLM 产出符合**你给定schema**的 JSON，并存入 `structured_response` 字段中。然后再返回 END。\n\n7. **结束** \n 整个 ReAct 流程完成后，图会返回一个最终状态，如：\n ```python\n {\n \"messages\": [ # 所有对话消息(包含了Human/AI/Tools等),\n ...,\n AIMessage(content=\"Here is the final answer\", tool_calls=[])\n ],\n \"remaining_steps\": 2,\n \"structured_response\": { ... } # 如果使用了response_format\n }\n ```\n 你可以从中拿到想要的最终 AI 回答。\n\n---\n\n# 五、如何查看“中间推理”或“工具调用”？\n\n从 **langgraph 0.3** 开始，`create_react_agent` 及其返回的 Graph 已经**不再支持** `graph.add_state_change_listener` 或在函数参数里传入 `callbacks`。如果你想**监听**或**打印** Agent 的中间思考、工具调用等过程，最好的方式是 **使用 `graph.stream(...)`**——它会在每一小步执行结束后产出一个“部分状态( partial_state )”，你可以在循环里进行日志记录、可视化或其他操作。示例：\n\n```python\ngraph = create_react_agent(model, tools=[...], prompt=\"...\")\n\ninputs = {\n \"messages\": [\n (\"user\", \"请分析特斯拉2025年的发展预期，包括新车型计划、销量目标、技术创新和市场扩张战略。\")\n ]\n}\n\nfor partial_state in graph.stream(inputs, stream_mode=\"values\"):\n messages = partial_state[\"messages\"]\n last_msg = messages[-1]\n if last_msg.type == \"ai\":\n print(\"[AIMessage] => \", last_msg.content)\n if last_msg.tool_calls:\n print(\"AI wants to call tools:\", last_msg.tool_calls)\n elif last_msg.type == \"tool\":\n print(\"[ToolMessage] => Name:\", last_msg.name, \"Content:\", last_msg.content)\n elif last_msg.type == \"human\":\n print(\"[User] => \", last_msg.content)\n\n# 最后一次迭代时，partial_state 就是最终结果\nfinal_answer = partial_state[\"messages\"][-1].content\nprint(\"最终回答:\", final_answer)\n```\n\n这样就能够**在每一次** Agent 或 Tools 完成后都获取状态，不需要“回调监听器”。\n\n---\n\n# 六、关于一些进阶用法\n\n1. **`interrupt_before` / `interrupt_after`** \n 如果你希望在“agent”节点**执行前**或者**后**打断，可以设置这两个可选参数，比如：\n ```python\n create_react_agent(\n model,\n tools=[...],\n interrupt_before=[\"tools\"],\n interrupt_after=[\"agent\"],\n ...\n )\n ```\n 当执行流程跑到 agent 或 tools 时，会先/后给你一个“交互点”机会，你可以在**流式**执行中察觉到这个点，或者抛出异常提前终止等。但是它比较适合做“用户确认”或“调试介入”，而不是实时日志。\n\n2. **`checkpointer` / `store`** \n - `checkpointer` 主要用来将单个“线程”（单条对话）的状态进行保存、恢复，可以在多回合对话里保留上下文。\n - `store` 提供了更跨线程或跨用户的持久化能力。 \n 通过把 `store` 绑定到 Graph，工具调用里还可以使用 `InjectedStore`，把数据写入或读取到 store 中（如相当于“全局数据库”）。\n\n3. **`response_format`** \n 如果你想让最终输出符合某种 JSON Schema 或 Pydantic 验证，可以这样写：\n ```python\n from langchain_core.prompts import ChatPromptTemplate\n from pydantic import BaseModel, Field\n\n class TeslaPlan(BaseModel):\n new_models: list[str] = Field(..., description=\"新车型列表\")\n sales_target: int = Field(..., description=\"预计销量\")\n technology_innovations: str\n market_strategy: str\n\n my_response_format = TeslaPlan\n\n graph = create_react_agent(\n model, \n tools, \n prompt=\"你是一个专业汽车分析师。\",\n response_format=my_response_format\n )\n ```\n 当 ReAct 流程结束后，LangGraph 会调用一次 LLM 并要求它返回符合 `TeslaPlan` 的 JSON。最终的 `state[\"structured_response\"]` 就是一个 Python 字典或 Pydantic 实例。\n\n4. **`version=\"v1\" / \"v2\"`** \n - **v1**: 工具调用是“把当前 AIMessage 中的所有 tool_calls 一次性并行执行” → tools → 再回到 agent。 \n - **v2**: 更细粒度地把每个 tool_call 拆开，每个都进入一个独立的 ToolNode 实例。如果一个 AIMessage 里有 3 个 tool_calls，就会做 3 次独立的“tools执行→回到agent”循环**（通过 Send API）**。这种方式可以在多工具协作里更灵活，也可以插入更多自定义逻辑，但要做好相应的结构化处理。\n\n---\n\n# 七、常见问题与答疑\n\n1. **Q: 我在旧版本使用 `graph.add_state_change_listener` 或 `callbacks`，现在为什么报错？** \n A: 因为新的 LangGraph 0.3 取消了这种回调API，推荐使用 `graph.stream(...)` 在每一步迭代中自行处理日志或监听逻辑。\n\n2. **Q: 如果不想每次都多轮循环，而只想 LLMC 接受一次输入就结束，怎么做？** \n A: 你可以传递 `tools=[]`（空）到 `create_react_agent`，这样它就生成一个不支持工具调用的图；agent 只会输出一次，然后就结束。此时相当于纯LLM调用。\n\n3. **Q: 要怎么限制调用工具的次数？** \n A: 你可以在输入的 `state` 里设置 `remaining_steps`，或自定义 `AgentState` 包含 `remaining_steps=3` 一类初始值，每次 agent节点执行后，LangGraph 会自动减少1。用完就不会再允许工具调用了。\n\n4. **Q: ReAct 会在同一个消息里多次请求调用工具吗？** \n A: 是可能的。尤其是当 LLM 在回答中生成多个 tool_calls，就会全部执行。你可以在 `v1` 模式下并行运行它们，也可以在 `v2` 模式下逐个执行。\n\n5. **Q: structured response 里的提示是如何工作的？** \n A: 当 `response_format` 是 `(system_prompt, schema)` 这种 tuple 时，LangGraph 会在最后的 LLM 调用里给一个额外的 system_prompt，引导 LLM 返回符合 schema 的 JSON。这样可以做更严格的结构化要求。\n\n---\n\n# 八、总结\n\n- **LangGraph** 是一个以“图”来编排对话和工具调用的框架； \n- **create_react_agent** 是“快捷构造 ReAct 风格图”的核心函数，一次性帮你搭建“agent(LLM) ↔ tools(工具节点) ↔ agent”循环； \n- 执行时默认从 `agent` 开始，如果 `AIMessage` 包含 `tool_calls` 就调用 `tools` 并注入结果，直到不再有工具调用； \n- 可以**流式**(`graph.stream(...)`) 或**一次性**(`graph.invoke(...)`)获取结果； \n- 要想查看中间推理和调用日志，使用 `stream` 在每一步循环里记录； \n- 可选地，你能通过 `interrupt_before` / `interrupt_after` 或 `checkpointer` / `store` 等更高级特性进一步定制执行流程或存储/恢复状态。\n\n这就是从**原理**到**源码**再到**执行流程**的完整解析。希望能帮助你在实际项目里更好地运用 `create_react_agent` 和 LangGraph！"
},
{
"path": "instructions/01.supervisor_pattern.md",
"content": "# Supervisor 模式：多智能体协作的核心实现\n\n## 1. 引言\n\n在人工智能领域，多智能体系统（Multi-Agent System）是一种将复杂任务分解为多个专业智能体协同完成的架构模式。本文将详细介绍我们在 Mentis 项目中实现的 Supervisor（监督者）模式，这是一种高效组织和协调多个智能体的方法。\n\n## 2. 多智能体系统的基本概念\n\n多智能体系统由多个具有不同专业能力的智能体组成，每个智能体负责特定的任务领域。在这种系统中，智能体之间需要有效地协作和通信，以完成复杂的任务。\n\n在我们的实现中，主要包含以下角色：\n\n- **Supervisor（监督者）**：负责任务分发、协调和结果整合的中央控制智能体\n- **Specialized Agents（专业智能体）**：具有特定领域专长的执行智能体\n\n## 3. Supervisor 模式的工作流程\n\n### 3.1 基本工作流程\n\nSupervisor 模式的工作流程如下：\n\n1. 用户向系统提交请求\n2. Supervisor 接收请求并进行任务分析\n3. Supervisor 决定调用哪个专业智能体处理任务\n4. 专业智能体执行任务并返回结果\n5. Supervisor 接收结果，可能进一步调用其他智能体\n6. Supervisor 整合所有结果并返回给用户\n\n### 3.2 控制权转移机制\n\nSupervisor 模式的核心是控制权的转移机制。在我们的实现中，这通过 `handoff` 工具实现：\n\n1. Supervisor 通过调用特定的 `handoff` 工具将控制权转移给目标智能体\n2. 目标智能体完成任务后，通过 `handoff_back_messages` 将控制权返回给 Supervisor\n3. 这种机制确保了在任何时刻只有一个智能体在处理任务，避免了冲突\n\n## 4. Supervisor 的核心实现\n\n### 4.1 核心代码分析\n\n在 `supervisor.py` 中，`create_supervisor` 函数是实现 Supervisor 模式的核心：\n\n```python\ndef create_supervisor(\n agents: list[Pregel],\n *,\n model: LanguageModelLike,\n tools: list[BaseTool | Callable] | None = None,\n prompt: Prompt | None = None,\n # ... 其他参数 ...\n) -> StateGraph:\n # 检查智能体名称唯一性\n agent_names = set()\n for agent in agents:\n if agent.name is None or agent.name == \"LangGraph\":\n raise ValueError(\"Please specify a name when you create your agent...\")\n if agent.name in agent_names:\n raise ValueError(f\"Agent with name '{agent.name}' already exists...\")\n agent_names.add(agent.name)\n \n # 为每个智能体创建 handoff 工具\n handoff_tools = [create_handoff_tool(agent_name=agent.name) for agent in agents]\n all_tools = (tools or []) + handoff_tools\n \n # 绑定工具到模型\n model = model.bind_tools(all_tools)\n \n # 创建 supervisor 智能体\n supervisor_agent = create_react_agent(\n name=supervisor_name,\n model=model,\n tools=all_tools,\n prompt=prompt,\n # ... 其他参数 ...\n )\n \n # 构建状态图\n builder = StateGraph(state_schema, config_schema=config_schema)\n builder.add_node(supervisor_agent, destinations=tuple(agent_names) + (END,))\n builder.add_edge(START, supervisor_agent.name)\n \n # 添加智能体节点和边\n for agent in agents:\n builder.add_node(\n agent.name,\n _make_call_agent(\n agent,\n output_mode,\n add_handoff_back_messages,\n supervisor_name,\n ),\n )\n builder.add_edge(agent.name, supervisor_agent.name)\n \n return builder\n```\n\n### 4.2 智能体调用机制\n\n`_make_call_agent` 函数负责创建智能体调用的包装函数：\n\n```python\ndef _make_call_agent(\n agent: Pregel,\n output_mode: OutputMode,\n add_handoff_back_messages: bool,\n supervisor_name: str,\n) -> Callable[[dict], dict] | RunnableCallable:\n # ... 参数验证 ...\n \n def _process_output(output: dict) -> dict:\n messages = output[\"messages\"]\n # 根据输出模式处理消息\n if output_mode == \"full_history\":\n pass\n elif output_mode == \"last_message\":\n messages = messages[-1:]\n \n # 添加控制权返回消息\n if add_handoff_back_messages:\n messages.extend(create_handoff_back_messages(agent.name, supervisor_name))\n \n return {\n **output,\n \"messages\": messages,\n }\n \n def call_agent(state: dict) -> dict:\n output = agent.invoke(state)\n return _process_output(output)\n \n # ... 异步版本 ...\n \n return RunnableCallable(call_agent, acall_agent)\n```\n\n### 4.3 设计亮点与最佳实践\n\nSupervisor 模式的实现包含了多个多智能体系统设计的黄金经验，以下是关键设计亮点：\n\n#### 4.3.1 自动控制权回传机制\n\n`_make_call_agent` 中的自动 handoff back 机制非常巧妙：\n\n```python\nif add_handoff_back_messages:\n messages.extend(create_handoff_back_messages(agent.name, supervisor_name))\n```\n\n这种设计的优势在于：\n- **隐式交接**：专业智能体无需知道 supervisor 的存在\n- **自动转发**：智能体完成任务后，系统自动将结果打包并转交回 supervisor\n- **消息插入**：在消息历史中自动插入 AIMessage 和 ToolMessage，表明控制权已转移\n- **零侵入性**：对智能体代码没有任何侵入，实现了完全的关注点分离\n\n#### 4.3.2 智能的上下文管理策略\n\n`output_mode` 参数提供了对消息历史的精确控制：\n\n```python\nif output_mode == \"last_message\":\n messages = messages[-1:]\n```\n\n这允许开发者灵活选择：\n- **全量历史模式**（`full_history`）：保留智能体输出的完整历史，提供完整上下文\n- **最后消息模式**（`last_message`）：仅保留最后一条消息，有效节省 token 消耗\n\n这种灵活的上下文压缩策略，在长对话或多轮智能体调用场景中尤为重要，可以有效防止上下文爆炸。\n\n#### 4.3.3 动态工具生成与绑定\n\n系统会自动为每个智能体创建对应的 handoff 工具：\n\n```python\nhandoff_tools = [create_handoff_tool(agent_name=agent.name) for agent in agents]\n```\n\n这些工具允许 supervisor 通过类似 `transfer_to_writer()` 或 `transfer_to_researcher()` 的函数调用来转移控制权，实现了：\n- **声明式调度**：调度逻辑由 LLM 决定，而非硬编码规则\n- **可解释性**：每次转移都有明确的工具调用，便于追踪和调试\n- **灵活性**：可以根据当前状态动态决定下一步调用哪个智能体\n\n#### 4.3.4 统一的 Runnable 接口封装\n\n每个智能体都被统一封装为 `RunnableCallable`：\n\n```python\nbuilder.add_node(agent.name, _make_call_agent(...))\n```\n\n这种封装提供了多种优势：\n- **统一接口**：所有智能体都遵循相同的调用接口\n- **状态管理**：状态由 LangGraph 自动管理，无需手动处理\n- **异步支持**：同时支持同步和异步调用，适应不同场景\n- **自动处理**：输入/输出状态转换自动完成，减少样板代码\n\n#### 4.3.5 灵活的配置选项\n\n系统支持多种配置选项，适应不同需求：\n- **多种提示格式**：支持字符串、SystemMessage 或可调用函数作为提示\n- **结构化输出**：支持 JSON schema、TypedDict 或 Pydantic 类作为输出格式\n- **状态模式**：可自定义状态结构，支持复杂的状态追踪和管理\n- **并行工具调用控制**：可以针对不同模型配置是否支持并行工具调用\n\n## 5. 实践案例：笑话生成与研究专家\n\n在 `01_supervisor_test.py` 中，我们实现了一个包含两个专业智能体的系统：\n\n### 5.1 智能体创建\n\n我们使用了两种不同的方式创建智能体：\n\n#### 5.1.1 功能型 API（Functional API）\n\n笑话生成器使用功能型 API 创建：\n\n```python\n@task\ndef generate_joke(messages):\n \"\"\"Generate a short joke (no tool calls).\"\"\"\n system_message = {\n \"role\": \"system\", \n \"content\": \"You are a witty comedian. Write a short joke.\"\n }\n msg = model.invoke([system_message] + messages)\n return msg\n\n@entrypoint()\ndef joke_agent(state):\n joke = generate_joke(state['messages']).result()\n messages = add_messages(state[\"messages\"], [joke])\n return {\"messages\": messages}\n\njoke_agent.name = \"joke_agent\"\n```\n\n#### 5.1.2 图形 API（Graph API）\n\n研究专家使用图形 API 创建：\n\n```python\ndef web_search(query: str) -> str:\n \"\"\"Search the web for information. (Mocked data here)\"\"\"\n return (\n \"Here are the headcounts for each of the FAANG companies in 2024:\\n\"\n # ... 模拟数据 ...\n )\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[web_search],\n name=\"research_expert\",\n prompt=(\n \"You are a world-class researcher. You have access to a 'web_search(query: str)' tool. \"\n \"Do not do any complicated math, just provide factual info from the web_search if needed.\"\n ),\n)\n```\n\n### 5.2 Supervisor 配置\n\n我们创建了一个 Supervisor 来协调这两个智能体：\n\n```python\nworkflow = create_supervisor(\n [research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions.\\n\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n)\n```\n\n### 5.3 执行流程\n\n当用户请求同时需要笑话和研究数据时，执行流程如下：\n\n1. Supervisor 接收用户请求\n2. Supervisor 分析请求，决定先调用 joke_agent\n3. joke_agent 生成笑话并返回结果\n4. Supervisor 接收笑话，然后调用 research_expert\n5. research_expert 查询数据并返回结果\n6. Supervisor 整合两个结果，生成最终回复\n\n## 6. 可视化与调试\n\n我们使用 LangGraph 的可视化功能生成了工作流图表，保存在 `examples/graphs/1_supervisor_test_01.png`，这有助于理解和调试多智能体系统的工作流程。\n\n## 7. 总结\n\nSupervisor 模式是一种高效组织多智能体系统的方法，它通过中央控制智能体协调专业智能体的工作，实现复杂任务的分解与协作。在我们的实现中，通过精心设计的 handoff 机制实现了智能体之间的控制权转移，确保系统的有序运行。\n\n这种模式的优势在于：\n\n1. **模块化**：每个智能体专注于特定领域，便于开发和维护\n2. **可扩展性**：可以方便地添加新的专业智能体\n3. **灵活性**：Supervisor 可以根据任务需求动态调用不同的智能体\n4. **结果整合**：Supervisor 负责整合各个智能体的结果，提供一致的用户体验\n5. **低耦合**：智能体之间通过消息传递交互，减少直接依赖\n6. **可追踪性**：每次控制权转移都有明确的工具调用记录，便于调试和监控\n7. **资源优化**：通过上下文管理策略，有效控制 token 消耗\n8. **开发便捷**：统一的接口和自动化的状态管理，减少样板代码\n\n通过本文的实践案例和深入分析，我们不仅展示了如何使用 LangGraph 和 LangChain 框架实现 Supervisor 模式，更揭示了背后的设计思想和最佳实践，为构建复杂的多智能体系统提供了宝贵参考。这些设计模式和技巧可以帮助开发者构建更加健壮、可维护和高效的智能体系统。"
},
{
"path": "instructions/02.supervisor_pattern_agent.md",
"content": "# Supervisor 模式：多智能体协作的核心实现 （Agent 封装模式）\n\n## 1. 引言\n\n在人工智能领域，多智能体系统（Multi-Agent System）是一种将复杂任务分解为多个专业智能体协同完成的架构模式。本文将详细介绍我们在 Mentis 项目中实现的 Supervisor（监督者）模式，这是一种高效组织和协调多个智能体的方法。\n\n## 2. 多智能体系统的基本概念\n\n多智能体系统由多个具有不同专业能力的智能体组成，每个智能体负责特定的任务领域。在这种系统中，智能体之间需要有效地协作和通信，以完成复杂的任务。\n\n在我们的实现中，主要包含以下角色：\n\n- **Supervisor（监督者）**：负责任务分发、协调和结果整合的中央控制智能体\n- **Specialized Agents（专业智能体）**：具有特定领域专长的执行智能体\n\n## 3. Supervisor 模式的工作流程\n\n### 3.1 基本工作流程\n\nSupervisor 模式的工作流程如下：\n\n1. 用户向系统提交请求\n2. Supervisor 接收请求并进行任务分析\n3. Supervisor 决定调用哪个专业智能体处理任务\n4. 专业智能体执行任务并返回结果\n5. Supervisor 接收结果，可能进一步调用其他智能体\n6. Supervisor 整合所有结果并返回给用户\n\n### 3.2 控制权转移机制\n\nSupervisor 模式的核心是控制权的转移机制。在我们的实现中，这通过 `handoff` 工具实现：\n\n1. Supervisor 通过调用特定的 `handoff` 工具将控制权转移给目标智能体\n2. 目标智能体完成任务后，通过 `handoff_back_messages` 将控制权返回给 Supervisor\n3. 这种机制确保了在任何时刻只有一个智能体在处理任务，避免了冲突\n\n## 4. 基础架构：BaseAgent 类\n\n在我们的重构中，我们引入了 `BaseAgent` 基类，作为所有智能体的基础。这种设计使得不同类型的智能体可以共享通用功能，同时保持各自的特性。\n\n### 4.1 BaseAgent 核心实现\n\n```python\nclass BaseAgent:\n _PROMPT_TEMPLATE = \"\"\"\n You have access to the following tools:\n {tools}\n Use the above tools to answer the question at the end.\n \"\"\"\n def __init__(\n self,\n name: str,\n model: Union[BaseChatModel, LanguageModelLike],\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[Union[str, SystemMessage, Callable]] = None,\n checkpointer: Optional[Checkpointer] = None,\n max_context_messages: Optional[int] = None, # 限制最近消息数量\n max_context_tokens: Optional[int] = None, # 限制总估计token数\n model_name: Optional[str] = \"gpt-4o-mini\", # 用于未来token估计改进\n ):\n # 初始化基本属性\n self.name = name\n self.model = model\n self.tools = tools or []\n self.prompt = prompt\n self.checkpointer = checkpointer\n self.max_context_messages = max_context_messages\n self.max_context_tokens = max_context_tokens\n self.model_name = model_name\n self._workflow = None\n self._agent = None\n```\n\n### 4.2 上下文管理机制\n\n`BaseAgent` 提供了智能的上下文管理机制，可以根据配置自动截断消息历史：\n\n```python\ndef _inject_context(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"注入记忆并根据配置截断消息。\"\"\"\n memory = state.get(\"memory\") or []\n messages = state.get(\"messages\", [])\n messages = self._truncate_messages(messages)\n memory_messages = [SystemMessage(content=chunk) for chunk in memory]\n state[\"messages\"] = memory_messages + messages\n return state\n```\n\n### 4.3 通用方法接口\n\n`BaseAgent` 定义了所有智能体共享的核心方法接口：\n\n```python\ndef build(self) -> StateGraph:\n \"\"\"构建工作流。\"\"\"\n \ndef compile(self) -> CompiledStateGraph:\n \"\"\"编译工作流。\"\"\"\n \ndef invoke(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"同步调用工作流。\"\"\"\n \nasync def ainvoke(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"异步调用工作流。\"\"\"\n```\n\n## 5. ReactAgent 类实现\n\n`ReactAgent` 是我们实现的基于 ReAct（Reasoning and Acting）模式的智能体，它继承自 `BaseAgent`，专注于推理和工具调用。\n\n### 5.1 ReactAgent 类设计\n\n```python\nclass ReactAgent(BaseAgent):\n \"\"\"ReAct Agent class for reasoning and acting with tools.\n \n This class provides a high-level interface for creating a ReAct agent workflow\n that can perform multi-step reasoning and tool calling.\n \"\"\"\n \n def __init__(\n self,\n model: LanguageModelLike,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[str] = None,\n response_format: Optional[\n Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]\n ] = None,\n state_schema: StateSchemaType = AgentState,\n config_schema: Type[Any] = None,\n checkpointer: Optional[Checkpointer] = None,\n store: Optional[BaseStore] = None,\n interrupt_before: Optional[List[str]] = None,\n interrupt_after: Optional[List[str]] = None,\n debug: bool = False,\n version: Literal[\"v1\", \"v2\"] = \"v1\",\n name: str = \"react_agent\",\n max_context_messages: Optional[int] = None,\n max_context_tokens: Optional[int] = None,\n model_name: Optional[str] = \"gpt-4o-mini\",\n ):\n # 调用父类初始化\n super().__init__(\n name=name,\n model=model,\n tools=tools or [],\n prompt=prompt,\n checkpointer=checkpointer,\n max_context_messages=max_context_messages,\n max_context_tokens=max_context_tokens,\n model_name=model_name\n )\n \n # 初始化ReactAgent特有属性\n self.response_format = response_format\n self.state_schema = state_schema\n self.config_schema = config_schema\n self.store = store\n self.interrupt_before = interrupt_before\n self.interrupt_after = interrupt_after\n self.debug = debug\n self.version = version\n self._agent = None\n```\n\n### 5.2 核心方法实现\n\n#### 5.2.1 compile 方法\n\n`compile` 方法负责编译 ReactAgent 工作流：\n\n```python\ndef compile(self) -> CompiledGraph:\n \"\"\"构建 ReAct agent 工作流。\n \n Returns:\n 编译后的 CompiledGraph\n \"\"\"\n # 如果_agent已经存在，直接返回，避免重复构建\n if self._agent is not None:\n return self._agent\n \n _react_agent = create_react_agent(\n model=self.model,\n tools=self.tools,\n prompt=self.prompt,\n response_format=self.response_format,\n state_schema=self.state_schema,\n config_schema=self.config_schema,\n checkpointer=self.checkpointer,\n store=self.store,\n interrupt_before=self.interrupt_before,\n interrupt_after=self.interrupt_after,\n debug=self.debug,\n version=self.version,\n name=self.name,\n )\n \n self._agent = CreateReactAgentWrapper(_react_agent, \n name=self.name,\n before_invoke=self.invoke,\n before_ainvoke=self.ainvoke)\n return self._agent\n```\n\n#### 5.2.2 invoke 和 ainvoke 方法\n\n`invoke` 和 `ainvoke` 方法负责调用 ReactAgent 处理用户请求，并提供调试信息：\n\n```python\ndef invoke(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"同步调用入口 (真正的 Agent 执行逻辑).\"\"\"\n # 打印调试信息\n messages = state.get(\"messages\", [])\n if messages:\n for i, msg in enumerate(messages, 1):\n type_str = type(msg).__name__\n print(f\"第 {i} 条消息 - {type_str} (Name: {msg.name}):\")\n msg.pretty_print()\n\n # 上下文注入\n state = self._inject_context(state)\n return state\n\nasync def ainvoke(self, state: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"异步调用入口.\"\"\"\n # 上下文注入\n state = await self._inject_context(state)\n return state\n```\n\n## 6. SupervisorAgent 类实现\n\n`SupervisorAgent` 类继承自 `BaseAgent`，专注于协调多个智能体的工作。在重构后，它增加了规划功能，可以更有效地管理复杂任务。\n\n### 6.1 SupervisorAgent 类设计\n\n```python\nclass SupervisorAgent(BaseAgent):\n \"\"\"Supervisor class for managing multiple agents with planning capabilities.\n \n This class provides a high-level interface for creating a supervisor workflow\n that can manage and coordinate multiple agents. It also includes planning capabilities\n to create and manage a plan for complex tasks using a state-driven approach.\n \n The planning functionality is implemented using PlanningStateHandler and PlanningTool,\n which provide a more structured and flexible way to manage tasks compared to the\n previous TodolistTool approach.\n \"\"\"\n \n def __init__(\n self,\n agents: List[BaseAgent],\n model: LanguageModelLike,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[str] = None,\n state_schema: StateSchemaType = AgentState,\n supervisor_name: str = \"supervisor\",\n checkpointer: Optional[Checkpointer] = None,\n output_mode: str = \"last_message\", # * full_history or last_message *\n enable_planning: bool = True, # * True or False *\n ):\n # 设置规划相关属性\n self._enable_planning = enable_planning\n \n # 如果启用规划功能，设置状态模式为PlanningAgentState\n if self._enable_planning and state_schema == AgentState:\n state_schema = PlanningAgentState\n \n # 存储特定于智能体的属性\n self.agents = agents\n self.output_mode = output_mode\n self.supervisor_name = supervisor_name\n self.state_schema = state_schema\n self.checkpointer = checkpointer\n self.tools = tools or []\n self._workflow = None\n self._agent = None\n \n # 生成基础提示词\n _final_prompt = self._PLANNING_PROMPT_TEMPLATE + \"/n/n\" + self._PLANNING_TOOL_TEMPLATE if self._enable_planning else self._PROMPT_TEMPLATE\n \n # 如果启用规划功能，添加规划工具\n if self._enable_planning:\n tools = tools or []\n tools.append(SimplePlanningTool())\n \n # 初始化BaseAgent父类\n super().__init__(\n name=supervisor_name,\n model=model,\n tools=tools,\n checkpointer=checkpointer,\n prompt=_final_prompt,\n )\n```\n\n### 6.2 核心方法实现\n\n#### 6.2.1 build 方法\n\n`build` 方法负责构建 Supervisor 工作流：\n\n```python\ndef build(self) -> StateGraph:\n \"\"\"构建 supervisor 工作流。\n \n Returns:\n 构建的 StateGraph\n \"\"\"\n \n if self._workflow is not None:\n return self._workflow\n \n self._workflow = create_supervisor(\n agents=self.agents,\n model=self.model,\n tools=self.tools,\n prompt=self.prompt,\n state_schema=self.state_schema,\n supervisor_name=self.supervisor_name,\n output_mode=self.output_mode,\n )\n \n return self._workflow\n```\n\n## 7. create_supervisor 函数实现\n\n`create_supervisor` 函数是 SupervisorAgent 的核心依赖，它负责创建多智能体协作的工作流。\n\n```python\ndef create_supervisor(\n agents: list[Pregel],\n *,\n model: LanguageModelLike,\n tools: list[BaseTool | Callable] | None = None,\n prompt: Prompt | None = None,\n response_format: Optional[\n Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]\n ] = None,\n state_schema: StateSchemaType = AgentState,\n config_schema: Type[Any] | None = None,\n output_mode: OutputMode = \"last_message\",\n add_handoff_back_messages: bool = True,\n supervisor_name: str = \"supervisor\",\n include_agent_name: AgentNameMode | None = None,\n) -> StateGraph:\n # 检查智能体名称唯一性\n agent_names = set()\n for agent in agents:\n if agent.name is None or agent.name == \"LangGraph\":\n raise ValueError(\n \"Please specify a name when you create your agent...\"\n )\n\n if agent.name in agent_names:\n raise ValueError(\n f\"Agent with name '{agent.name}' already exists. Agent names must be unique.\"\n )\n\n agent_names.add(agent.name)\n\n # 为每个智能体创建 handoff 工具\n handoff_tools = [create_handoff_tool(agent_name=agent.name) for agent in agents]\n all_tools = (tools or []) + handoff_tools\n\n # 绑定工具到模型\n if _supports_disable_parallel_tool_calls(model):\n model = model.bind_tools(all_tools, parallel_tool_calls=False)\n else:\n model = model.bind_tools(all_tools)\n\n # 处理智能体名称显示方式\n if include_agent_name:\n model = with_agent_name(model, include_agent_name)\n \n # 创建 supervisor 智能体\n _react_agent = ReactAgent(\n name=supervisor_name,\n model=model,\n tools=all_tools,\n prompt=prompt,\n state_schema=state_schema,\n response_format=response_format,\n debug=False,\n )\n supervisor_agent = _react_agent.compile()\n \n # 构建状态图\n builder = StateGraph(state_schema, config_schema=config_schema)\n builder.add_node(supervisor_agent, destinations=tuple(agent_names) + (END,))\n builder.add_edge(START, supervisor_agent.name)\n \n # 添加智能体节点和边\n for agent in agents:\n # 如果智能体是 \"ReactAgent\" 或类似类型\n if hasattr(agent, \"get_agent\") and callable(agent.get_agent):\n agent = agent.get_agent() # 获取编译后的子图\n \n builder.add_node(\n agent.name,\n _make_call_agent(\n agent,\n output_mode,\n add_handoff_back_messages,\n supervisor_name,\n ),\n )\n builder.add_edge(agent.name, supervisor_agent.name)\n\n return builder\n```\n\n## 8. 实践案例\n\n### 8.1 使用 create_supervisor 函数（原始方式）\n\n在 `01_supervisor_test.py` 中，我们使用原始的 `create_supervisor` 函数实现了一个包含两个专业智能体的系统：\n\n```python\nworkflow = create_supervisor(\n [research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions.\\n\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n)\n\n# 编译得到一个可调用的\"App\"\napp = workflow.compile()\n```\n\n### 8.2 使用 SupervisorAgent 类（封装方式）\n\n在 `02_supervisor_agent_test.py` 中，我们使用封装的 `SupervisorAgent` 类实现了相同的功能，但增加了规划能力：\n\n```python\n# 创建 SupervisorAgent 实例\nsupervisor = SupervisorAgent(\n agents=[research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions.\\n\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n enable_planning=True, # 启用规划功能\n)\n\n# 编译得到一个可调用的\"App\"\napp = supervisor.compile()\n```\n\n### 8.3 两种方式的比较\n\n两种实现方式在基本功能上相似，但使用 `SupervisorAgent` 类的方式有以下优势：\n\n1. **更简洁的 API**：封装了复杂的参数和配置，提供了更简洁的接口\n2. **更好的封装性**：将相关功能封装在一个类中，便于维护和扩展\n3. **更好的可读性**：代码结构更清晰，意图更明确\n4. **更好的可重用性**：可以方便地在不同项目中复用\n5. **规划功能**：内置了任务规划能力，可以更有效地管理复杂任务\n6. **上下文管理**：通过 BaseAgent 继承了智能的上下文管理机制\n\n## 9. 总结\n\n在重构后的实现中，我们引入了以下关键改进：\n\n1. **BaseAgent 基类**：提供了所有智能体共享的基础功能，如上下文管理、工作流构建等\n2. **ReactAgent 重构**：现在继承自 BaseAgent，使用 CreateReactAgentWrapper 增强功能\n3. **SupervisorAgent 重构**：现在继承自 BaseAgent，增加了规划功能\n4. **统一的接口**：所有智能体类型现在共享相同的核心方法接口\n5. **智能上下文管理**：可以根据配置自动截断消息历史，优化性能\n\nSupervisor 模式是一种高效组织多智能体系统的方法，它通过中央控制智能体协调专业智能体的工作，实现复杂任务的分解与协作。在我们的重构实现中，通过引入 BaseAgent 基类和增强 SupervisorAgent 的规划能力，使得多智能体系统更加灵活、高效，同时保持了良好的可维护性和可扩展性。\n\n这种模式特别适合以下场景：\n- 需要多种专业知识协作的复杂任务\n- 需要动态决策调用不同专家的场景\n- 需要结果整合和质量控制的任务流程\n- 需要有计划地执行多步骤任务的场景\n\n未来，我们将继续优化 Supervisor 模式的实现，增强其灵活性和可扩展性，并探索更多的应用场景。"
},
{
"path": "instructions/03.tavily_search_integration.md",
"content": "# Tavily搜索工具集成：为多智能体系统提供实时信息能力\n\n## 1. 引言\n\n在多智能体系统中，获取实时、准确的外部信息是提升系统实用性的关键因素。本文将详细介绍我们在 Mentis 项目中集成 Tavily 搜索工具的实现，这使得我们的智能体系统能够获取最新的网络信息，大幅提升了系统的实用价值。\n\n## 2. Tavily 搜索工具概述\n\nTavily 是一个专为 AI 应用设计的搜索 API，它提供了高质量、结构化的网络搜索结果。在我们的实现中，Tavily 工具具有以下特点：\n\n- **实时性**：能够获取最新的网络信息\n- **结构化输出**：返回格式化的搜索结果，便于智能体处理\n- **可配置性**：支持多种参数配置，如搜索深度、结果数量等\n- **多媒体支持**：可选择性地包含图片等多媒体内容\n\n## 3. Tavily 工具的实现\n\n### 3.1 核心代码分析\n\n在 `tavily_tools.py` 中，我们实现了 `TavilySearchResults` 类，它继承自 LangChain 的 `BaseTool`：\n\n```python\nclass TavilySearchResults(BaseTool):\n \"\"\"Tool that queries the Tavily Search API and gets back json.\"\"\"\n \n name: str = \"tavily_search_results_json\"\n description: str = (\n \"A search engine optimized for comprehensive, accurate, and trusted results. \"\n \"Useful for when you need to answer questions about current events. \"\n \"Input should be a search query.\"\n )\n args_schema: Type[BaseModel] = TavilyInput\n \n max_results: int = 5\n \"\"\"Max search results to return, default is 5\"\"\"\n search_depth: str = \"advanced\"\n \"\"\"The depth of the search. It can be \"basic\" or \"advanced\"\"\"\"\n include_domains: List[str] = []\n \"\"\"A list of domains to specifically include in the search results.\"\"\"\n exclude_domains: List[str] = []\n \"\"\"A list of domains to specifically exclude from the search results.\"\"\"\n include_answer: bool = False\n \"\"\"Include a short answer to original query in the search results.\"\"\"\n include_raw_content: bool = False\n \"\"\"Include cleaned and parsed HTML of each site search results.\"\"\"\n include_images: bool = False\n \"\"\"Include a list of query related images in the response.\"\"\"\n \n api_wrapper: TavilySearchAPIWrapper = Field(default_factory=TavilySearchAPIWrapper)\n response_format: Literal[\"content_and_artifact\"] = \"content_and_artifact\"\n```\n\n### 3.2 搜索执行方法\n\n`TavilySearchResults` 类提供了同步和异步两种搜索方法：\n\n```python\ndef _run(\n self,\n query: str,\n run_manager: Optional[CallbackManagerForToolRun] = None,\n) -> Tuple[Union[List[Dict[str, str]], str], Dict]:\n \"\"\"Use the tool.\"\"\"\n try:\n raw_results = self.api_wrapper.raw_results(\n query,\n self.max_results,\n self.search_depth,\n self.include_domains,\n self.exclude_domains,\n self.include_answer,\n self.include_raw_content,\n self.include_images,\n )\n except Exception as e:\n return repr(e), {}\n return self.api_wrapper.clean_results(raw_results[\"results\"]), raw_results\n\nasync def _arun(\n self,\n query: str,\n run_manager: Optional[AsyncCallbackManagerForToolRun] = None,\n) -> Tuple[Union[List[Dict[str, str]], str], Dict]:\n \"\"\"Use the tool asynchronously.\"\"\"\n # 异步实现...\n```\n\n## 4. 在多智能体系统中集成 Tavily 工具\n\n### 4.1 创建研究型智能体\n\n在我们的多智能体系统中，我们创建了一个专门的研究型智能体，它使用 Tavily 搜索工具获取实时信息：\n\n```python\n# 创建Tavily搜索工具\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=False,\n include_images=False,\n search_depth=\"advanced\"\n)\n\nresearch_agent = create_react_agent(\n model=model,\n tools=[tavily_search],\n name=\"research_expert\",\n prompt=(\n \"You are a world-class researcher. You have access to the 'tavily_search_results_json' tool \"\n \"which can search the web for real-time information. \"\n \"When asked a question, use this tool to find accurate and up-to-date information. \"\n \"Summarize the search results in a clear and concise manner. \"\n \"Always cite your sources by including the URLs from the search results.\"\n ),\n)\n```\n\n### 4.2 与 Supervisor 集成\n\n研究型智能体作为专业智能体，被集成到 Supervisor 模式中：\n\n```python\n# 创建 SupervisorAgent 实例\nsupervisor = SupervisorAgent(\n agents=[research_agent, joke_agent],\n model=model,\n prompt=(\n \"You are the overall supervisor. You manage two specialized agents:\\n\"\n \"1) joke_agent: for telling jokes.\\n\"\n \"2) research_expert: for factual or data-related questions using real-time web search.\\n\\n\"\n \"If the user wants a joke, call joke_agent.\\n\"\n \"If the user wants factual information or research data, call research_expert.\\n\"\n \"If the user wants a joke AND some research data in the same query, \"\n \"you MUST call joke_agent first, get the joke, then call research_expert for the data. \"\n \"After both calls, provide a final combined response. \"\n \"Do not call more than one agent in a single LLM message; do it step by step.\"\n ),\n)\n```\n\n## 5. 实践案例\n\n### 5.1 只询问研究数据\n\n当用户只询问研究数据时，Supervisor 会直接调用研究型智能体：\n\n```python\n# 示例2：只询问研究数据\nresult2 = app.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"谁是现任美国总统？\"}]})\n```\n\n在这种情况下，研究型智能体会使用 Tavily 搜索工具获取最新信息，并返回结构化的回答，包括引用的来源。\n\n### 5.2 混合查询\n\n当用户同时需要笑话和研究数据时，Supervisor 会先调用笑话智能体，然后调用研究型智能体：\n\n```python\n# 示例3：同时询问笑话和研究数据\nresult3 = app.invoke({\"messages\": [{\"role\": \"user\", \"content\": \"讲个关于人工智能的笑话，然后告诉我什么是大型语言模型\"}]})\n```\n\n这种情况下，Supervisor 会协调两个智能体的工作，并整合它们的结果。\n\n## 6. 可视化与调试\n\n我们使用 LangGraph 的可视化功能生成了工作流图表，保存在 `examples/graphs/03_tavily_tools_test.png`。这个图表展示了包含 Tavily 搜索工具的多智能体系统的工作流程，有助于理解和调试系统。\n\n## 7. 总结\n\nTavily 搜索工具的集成为我们的多智能体系统带来了以下优势：\n\n1. **实时信息获取**：系统能够获取最新的网络信息，不再局限于模型训练数据的时间范围\n2. **信息准确性提升**：通过引用可靠的网络来源，提高了系统回答的准确性\n3. **功能扩展**：使系统能够回答关于最新事件、数据和信息的问题\n4. **灵活配置**：可以根据需要调整搜索参数，优化搜索结果\n\n通过 Tavily 搜索工具的集成，我们的多智能体系统从一个封闭的知识系统转变为一个能够获取实时信息的开放系统，大大提升了系统的实用价值和应用范围。\n\n未来，我们计划进一步优化搜索工具的使用策略，提高搜索效率和结果质量，并探索更多外部工具的集成，使系统能够处理更复杂的任务。"
},
{
"path": "instructions/04.react_agent.md",
"content": "# ReactAgent：基于ReAct方法论的多步推理与工具调用框架\n\n## 1. 引言\n\nReactAgent是一个基于ReAct方法论的智能体框架，它能够通过多步推理和工具调用来解决复杂问题。本文将详细介绍ReactAgent的核心概念、工作原理、实现方式以及在实际应用中的使用方法。\n\n## 2. ReactAgent的核心概念\n\n### 2.1 什么是ReAct方法论\n\nReAct（Reasoning + Acting）是一种结合推理和行动的AI问题解决方法论，它包含两个核心步骤：\n\n1. **推理（Reasoning）**：让语言模型进行思考，分析问题，并决定下一步行动。\n2. **行动（Acting）**：执行具体的工具调用，获取外部信息或执行特定操作。\n\n这两个步骤可以多次循环往复，直到问题被解决。ReAct方法论特别适合处理需要多步骤、多工具协作的复杂问题。\n\n### 2.2 ReactAgent与LangGraph的关系\n\nReactAgent是基于LangGraph框架实现的，它利用LangGraph的图结构来编排推理和行动的流程。在LangGraph中，ReactAgent被表示为一个包含多个节点和边的有向图：\n\n- **节点（Node）**：包括Agent节点（负责推理）和Tools节点（负责执行工具调用）\n- **边（Edge）**：定义节点之间的转换条件，例如当Agent生成工具调用时，流程转向Tools节点\n\n## 3. ReactAgent的实现\n\n### 3.1 ReactAgent类的设计\n\n在我们的实现中，ReactAgent类继承自LangGraph的Pregel类，提供了一个高级接口来创建和管理ReAct工作流：\n\n```python\nclass ReactAgent(Pregel):\n \"\"\"ReAct Agent class for reasoning and acting with tools.\n \n This class provides a high-level interface for creating a ReAct agent workflow\n that can perform multi-step reasoning and tool calling.\n \"\"\"\n \n def __init__(\n self,\n model: LanguageModelLike,\n tools: Optional[List[Union[BaseTool, Callable]]] = None,\n prompt: Optional[str] = None,\n response_format: Optional[\n Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]\n ] = None,\n state_schema: StateSchemaType = AgentState,\n config_schema: Type[Any] = None,\n interrupt_before: Optional[List[str]] = None,\n interrupt_after: Optional[List[str]] = None,\n debug: bool = False,\n version: Literal[\"v1\", \"v2\"] = \"v1\",\n name: str = \"react_agent\",\n ):\n # 初始化代码...\n```\n\n### 3.2 核心方法\n\nReactAgent类提供了以下核心方法：\n\n1. **build()**: 构建ReAct工作流图\n2. **compile()**: 编译工作流为可执行应用\n3. **invoke()**: 同步执行ReAct工作流\n4. **ainvoke()**: 异步执行ReAct工作流\n5. **stream()**: 流式执行，可以获取中间状态\n6. **get_graph()**: 获取底层图结构，用于可视化或调试\n\n### 3.3 与create_react_agent的关系\n\nReactAgent类内部使用了LangGraph提供的`create_react_agent`函数来构建工作流图。这个函数自动处理了：\n\n- 创建Agent节点（用于调用语言模型）\n- 创建Tools节点（用于执行工具调用）\n- 在节点之间建立连接\n- 处理状态管理和流程控制\n\n## 4. 使用ReactAgent解决复杂问题\n\n### 4.1 基本使用流程\n\n使用ReactAgent的基本流程如下：\n\n1. **初始化ReactAgent**：提供语言模型和工具\n2. **编译工作流**：调用compile()方法\n3. **准备初始状态**：通常包含用户的问题\n4. **执行或流式执行**：使用invoke()或stream()方法\n5. **处理结果**：分析最终状态或中间状态\n\n### 4.2 集成Tavily搜索工具\n\n在实际应用中，我们经常将ReactAgent与Tavily搜索工具集成，使其能够获取实时网络信息：\n\n```python\n# 创建Tavily搜索工具\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=True,\n include_images=False,\n search_depth=\"advanced\"\n)\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=[tavily_search],\n prompt=(\n \"你是一位专业的研究分析师，擅长分析复杂问题并提供深入见解。\\n\"\n \"当面对复杂问题时，请遵循以下REACT方法论：\\n\"\n \"1. 分解问题：将复杂问题分解为更小的子问题\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及搜索的顺序\\n\"\n \"3. 执行搜索：使用tavily_search_results_json工具执行搜索\\n\"\n \"4. 分析结果：分析搜索结果，确定是否需要进一步搜索\\n\"\n \"5. 综合信息：将所有搜索结果综合成一个连贯的回答\\n\"\n ),\n)\n\n# 编译工作流\nagent = react_agent.compile()\n```\n\n### 4.3 处理用户输入\n\n以下是处理用户输入的示例代码：\n\n```python\n# 准备初始状态\ninitial_state = {\n \"messages\": [HumanMessage(content=user_input)]\n}\n\n# 流式执行并获取中间状态\nfor partial_state in react_agent.stream(initial_state, stream_mode=\"values\"):\n # 处理中间状态\n messages = partial_state.get(\"messages\", [])\n if messages:\n latest_message = messages[-1]\n # 记录或显示最新消息\n log_agent_actions({\"messages\": [latest_message]})\n\n# 处理最终结果\nfinal_state = partial_state # 最后一个状态就是最终状态\n```\n\n## 5. ReactAgent的优势与应用场景\n\n### 5.1 优势\n\n- **多步推理**：能够分解复杂问题，逐步解决\n- **工具调用**：可以集成各种外部工具，扩展能力边界\n- **状态管理**：自动管理对话状态和中间结果\n- **可视化**：支持工作流可视化，便于调试和理解\n- **流式执行**：可以获取中间状态，实现更好的用户体验\n\n### 5.2 应用场景\n\n- **研究助手**：帮助用户研究复杂问题，获取最新信息\n- **数据分析**：分步骤处理数据分析任务\n- **决策支持**：通过多步推理和信息收集辅助决策\n- **教育辅导**：分解复杂概念，逐步引导学习\n\n## 6. 实际案例：研究特斯拉2025年发展预期\n\n以下是使用ReactAgent研究特斯拉2025年发展预期的实际案例：\n\n1. **问题分解**：将问题分解为新车型计划、销量目标、技术创新和市场扩张战略\n2. **执行搜索**：针对每个子问题执行Tavily搜索\n3. **分析结果**：分析每个搜索的结果，提取关键信息\n4. **综合信息**：将所有信息整合为一个全面的分析报告\n\n通过这种方式，ReactAgent能够提供比单次查询更全面、更深入的分析结果。\n\n## 7. 总结\n\nReactAgent是一个强大的基于ReAct方法论的智能体框架，它通过多步推理和工具调用来解决复杂问题。在实际应用中，ReactAgent特别适合需要分步骤思考、收集信息和综合分析的任务。通过与Tavily等工具的集成，ReactAgent能够获取实时信息，大幅提升其实用价值。\n\n在未来的开发中，我们将继续优化ReactAgent的性能，增强其推理能力，并集成更多实用工具，使其能够应对更广泛的应用场景。"
},
{
"path": "instructions/05.react_agent_user_input.md",
"content": "# ReactAgent与用户交互：构建交互式研究助手\n\n## 1. 引言\n\n本文将介绍如何使用ReactAgent构建一个能够与用户进行交互的研究助手，该助手能够接收用户输入，使用搜索工具获取信息，并提供深入的分析结果。这种交互式助手特别适合需要实时信息和多轮对话的场景。\n\n## 2. 交互式研究助手的核心概念\n\n### 2.1 用户输入处理\n\n交互式研究助手需要能够处理用户的自然语言输入，理解用户的意图，并将其转化为可执行的搜索查询或其他操作。这涉及到：\n\n1. **输入解析**：分析用户输入，提取关键信息和查询意图\n2. **查询重构**：将用户的自然语言问题转化为更有效的搜索查询\n3. **上下文维护**：在多轮对话中保持对话上下文的连贯性\n\n### 2.2 搜索工具集成\n\n研究助手的核心功能是能够获取和分析信息，这通常通过集成各种搜索工具来实现：\n\n1. **Tavily搜索**：提供高质量的网络搜索结果，支持深度搜索模式\n2. **结果处理**：对搜索结果进行过滤、排序和整合，提取最相关的信息\n3. **多次搜索策略**：对复杂问题进行分解，执行多次有针对性的搜索\n\n## 3. 实现交互式研究助手\n\n### 3.1 基本架构\n\n交互式研究助手的基本架构包括：\n\n```\n用户输入 → ReactAgent → 搜索工具 → 结果分析 → 回复生成 → 用户\n```\n\n这个流程可以多次循环，形成多轮对话。\n\n### 3.2 ReactAgent配置\n\n以下是创建交互式研究助手的核心代码：\n\n```python\ndef create_react_agent_instance():\n \"\"\"创建并返回ReactAgent实例\"\"\"\n react_agent = ReactAgent(\n model=model,\n tools=[tavily_search],\n name=\"research_assistant\",\n # 提示词强调分解问题、多步思考和综合信息\n prompt=(\n \"你是一位专业的研究分析师，擅长分析复杂问题并提供深入见解。\\n\"\n \"你有一个强大的工具'tavily_search_results_json'可以搜索网络获取实时信息。\\n\\n\"\n \"当面对复杂问题时，请遵循以下REACT方法论：\\n\"\n \"1. 分解问题：将复杂问题分解为更小的子问题\\n\"\n \"2. 制定计划：确定需要搜索哪些信息，以及搜索的顺序\\n\"\n \"3. 执行搜索：使用tavily_search_results_json工具执行搜索\\n\"\n \"4. 分析结果：分析搜索结果，确定是否需要进一步搜索\\n\"\n \"5. 综合信息：将所有搜索结果综合成一个连贯的回答\\n\\n\"\n \"重要提示：\\n\"\n \"- 不要一次性搜索过于宽泛的问题\\n\"\n \"- 对于复杂问题，进行多次有针对性的搜索\\n\"\n \"- 每次搜索后评估结果，决定下一步行动\\n\"\n \"- 在最终回答中引用来源，包括搜索结果中的URL\\n\"\n \"- 清晰地展示你的思考过程，包括问题分解和计划制定\\n\"\n ),\n )\n \n return react_agent\n```\n\n### 3.3 Tavily搜索工具配置\n\n```python\ntavily_search = TavilySearchResults(\n max_results=3,\n include_answer=True,\n include_raw_content=True, # 包含原始内容，便于分析\n include_images=False,\n search_depth=\"advanced\" # 使用高级搜索深度\n)\n```\n\n### 3.4 用户交互循环\n\n用户交互循环的核心是通过`stream`方法获取中间状态，并实时显示Agent的思考过程：\n\n```python\ndef process_user_query(query):\n # 创建ReactAgent实例\n react_agent = create_react_agent_instance()\n agent = react_agent.compile()\n \n # 准备输入\n inputs = {\n \"messages\": [HumanMessage(content=query)]\n }\n \n # 使用stream方法逐步获取中间状态\n final_state = None\n for partial_state in react_agent.stream(inputs, stream_mode=\"values\"):\n # 保存最终状态\n final_state = partial_state\n \n # 获取最新消息并记录\n messages = partial_state.get(\"messages\", [])\n if messages:\n latest_message = messages[-1]\n log_agent_actions({\"messages\": [latest_message]})\n \n # 返回最终回答\n return final_state\n```\n\n## 4. 最佳实践与优化策略\n\n### 4.1 提示词优化\n\n提示词对研究助手的性能至关重要，应包含以下要素：\n\n1. **角色定义**：明确助手的专业身份和能力\n2. **方法论指导**：提供结构化的问题解决方法\n3. **工具使用指南**：说明如何有效使用搜索工具\n4. **输出格式要求**：规定回答的结构和引用方式\n\n### 4.2 搜索策略优化\n\n为提高搜索效率和结果质量，可采用以下策略：\n\n1. **渐进式搜索**：从一般到具体，逐步缩小搜索范围\n2. **多角度查询**：使用不同的关键词和表述方式进行搜索\n3. **结果验证**：通过交叉检查多个来源验证信息的准确性\n4. **深度参数调整**：根据问题复杂度调整搜索深度参数\n\n### 4.3 用户体验优化\n\n提升用户体验的关键点包括：\n\n1. **透明的思考过程**：展示Agent的推理过程，增强可信度\n2. **实时反馈**：通过流式输出提供即时反馈\n3. **引用来源**：清晰标注信息来源，便于用户进一步探索\n4. **交互式引导**：在复杂问题上引导用户提供更多上下文或澄清问题\n\n## 5. 应用场景\n\n交互式研究助手适用于多种场景：\n\n1. **学术研究**：帮助研究人员快速获取和分析相关文献\n2. **市场分析**：收集和整合市场趋势、竞争对手信息\n3. **新闻摘要**：汇总和分析最新新闻事件\n4. **技术调研**：探索新技术、框架或工具的特性和评价\n5. **教育辅助**：为学生提供学习资料和解答问题\n\n## 6. 总结\n\nReactAgent结合用户交互和搜索工具，可以构建功能强大的研究助手，能够处理复杂查询并提供深入分析。通过优化提示词、搜索策略和用户体验，可以进一步提升助手的性能和实用性。未来的发展方向包括集成更多专业数据源、增强多模态能力，以及提供更个性化的信息服务。"
},
{
"path": "instructions/06.web_extraction_tools.md",
"content": "# 网页提取工具：FireCrawl与Jina的集成与应用\n\n## 1. 引言\n\n网页内容提取是智能体系统中的重要能力，它使智能体能够从互联网获取、分析和处理结构化和非结构化的网页内容。本文将详细介绍如何在Mentis框架中集成和使用FireCrawl和Jina两种强大的网页提取工具，以实现高效的网站结构分析和内容提取。\n\n## 2. 网页提取工具的核心概念\n\n### 2.1 网页提取的两个关键步骤\n\n高效的网页内容提取通常包含两个关键步骤：\n\n1. **网站结构分析**：了解网站的组织结构、页面之间的链接关系，以及重要页面的位置。\n2. **内容提取**：从特定页面中提取有价值的文本、图像或其他结构化信息。\n\n### 2.2 FireCrawl与Jina的角色分工\n\n在Mentis框架中，我们使用两种工具来分别处理这两个步骤：\n\n1. **FireCrawl**：专注于网站结构分析，能够爬取网站的页面结构和链接关系。\n2. **Jina**：专注于内容提取，能够从特定URL获取干净、结构化的内容。\n\n## 3. FireCrawlTool的实现与使用\n\n### 3.1 FireCrawlTool的基本结构\n\nFireCrawlTool是对FireCrawl API的封装，提供了网站爬取和内容分析的能力：\n\n```python\nclass FireCrawlTool(BaseTool):\n \"\"\"Tool that uses FireCrawl API to crawl or scrape web content.\"\"\"\n\n name: str = \"firecrawl_tool\"\n description: str = (\n \"A web crawler and scraper that extracts content from websites. \"\n \"Useful for when you need to analyze the content of a specific website or webpage. \"\n \"Input should be a URL to crawl or scrape.\"\n )\n args_schema: Type[BaseModel] = FireCrawlInput\n \n api_key: Optional[str] = None\n api_url: Optional[str] = None\n mode: str = \"crawl\"\n params: Dict[str, Any] = Field(default_factory=dict)\n```\n\n### 3.2 FireCrawlTool的配置选项\n\nFireCrawlTool提供了多种配置选项：\n\n1. **mode**：工作模式，可选值包括：\n - `crawl`：爬取网站结构和链接\n - `scrape`：提取特定页面的内容\n - `map`：生成网站地图\n\n2. **params**：额外参数，常用的包括：\n - `max_pages`：限制爬取的最大页面数量\n - `max_depth`：限制爬取的最大深度\n - `follow_links`：是否跟踪页面中的链接\n\n### 3.3 使用FireCrawlTool爬取网站结构\n\n以下是使用FireCrawlTool爬取网站结构的示例代码：\n\n```python\n# 创建FireCrawl工具 - 用于网站结构分析\nfirecrawl_tool = FireCrawlTool(\n mode=\"crawl\", # 使用爬取模式\n params={\n \"max_pages\": 5, # 限制爬取页面数量\n }\n)\n\n# 在Agent中使用该工具\nreact_agent = create_react_agent(\n model=model,\n tools=[firecrawl_tool],\n name=\"web_crawler\",\n prompt=\"你是一位网站结构分析专家...\"\n)\n```\n\n## 4. JinaSearch的实现与使用\n\n### 4.1 JinaSearch的基本功能\n\nJinaSearch是LangChain提供的一个工具，能够从网页中提取干净、可读的内容，去除广告、导航栏等干扰元素：\n\n```python\nfrom langchain_community.tools import JinaSearch\n\n# 创建Jina Reader工具 - 用于内容提取\njina_reader_tool = JinaSearch()\n```\n\n### 4.2 使用JinaSearch提取网页内容\n\nJinaSearch特别适合在确定了目标页面后，提取其中的核心内容：\n\n```python\n# 在Agent中结合FireCrawl和Jina\nreact_agent = create_react_agent(\n model=model,\n tools=[firecrawl_tool, jina_reader_tool],\n name=\"web_extraction_expert\",\n prompt=\"你是一位网页内容分析专家...\"\n)\n```\n\n## 5. 网页提取的最佳实践\n\n### 5.1 两阶段提取策略\n\n为了高效地提取网页内容，建议采用两阶段策略：\n\n1. **第一阶段**：使用FireCrawlTool爬取网站结构，了解网站的组织方式和重要页面。\n2. **第二阶段**：根据第一阶段的结果，使用JinaSearch有针对性地提取重要页面的内容。\n\n### 5.2 提示词优化\n\n为了引导Agent正确使用这两个工具，提示词应该明确指出工具的使用顺序和方法：\n\n```python\nprompt = (\n \"你是一位专业的网页内容分析专家，擅长提取和分析网站结构与内容。\\n\"\n \"你有两个强大的工具:\\n\"\n \"1. 'firecrawl_tool': 用于爬取网站结构和下级页面\\n\"\n \"2. 'jina_reader_tool': 用于从特定URL提取结构化内容\\n\\n\"\n \"当面对网站分析任务时，请遵循以下方法论:\\n\"\n \"1. 先使用firecrawl_tool了解网站结构\\n\"\n \"2. 再使用jina_reader_tool提取关键页面内容\\n\"\n \"3. 最后整合信息提供分析结果\"\n)\n```\n\n### 5.3 处理大型网站的策略\n\n对于大型网站，可以采用以下策略：\n\n1. **限制爬取范围**：设置合理的`max_pages`和`max_depth`参数\n2. **分批处理**：先获取网站结构，然后每次只处理1-3个重要页面\n3. **内容摘要**：对提取的内容进行摘要，减少token消耗\n\n## 6. 实际应用案例\n\n### 6.1 分析LangGraph文档网站\n\n以下是使用FireCrawl和Jina分析LangGraph文档网站的示例：\n\n```python\n# 定义输入\ninputs = {\n \"messages\": [\n {\"role\": \"user\", \"content\": \"爬取LangGraph文档网站的每个章节的内容(https://langchain-ai.github.io/langgraph/how-tos/) \"}\n ]\n}\n\n# 使用stream方法逐步获取中间状态\nfinal_state = None\nfor partial_state in react_agent.stream(inputs, stream_mode=\"values\"):\n # 处理中间状态...\n pass\n```\n\n### 6.2 结果分析与处理\n\nAgent会首先使用FireCrawl获取网站结构，然后使用Jina提取重要页面的内容，最后整合信息提供分析结果：\n\n1. **网站结构分析**：识别主要章节和子页面\n2. **内容提取**：获取每个章节的详细内容\n3. **信息整合**：将内容组织成结构化的文档或摘要\n\n## 7. 总结\n\nFireCrawl和Jina的结合为智能体提供了强大的网页内容提取能力。通过两阶段提取策略，可以高效地分析网站结构并提取有价值的内容。这种能力使智能体能够从互联网获取实时信息，为用户提供更加全面和准确的回答。\n\n未来的发展方向包括增强对JavaScript渲染页面的支持、提高内容提取的准确性，以及集成更多专业领域的内容分析能力。"
},
{
"path": "instructions/07.web_extraction_with_filesystem.md",
"content": "# 网页提取与文件系统集成：构建内容采集与存储系统\n\n## 1. 引言\n\n在智能体系统中，网页内容提取通常需要与文件系统操作相结合，以便将提取的内容持久化存储。本文将详细介绍如何在Mentis框架中集成网页提取工具和文件系统工具，并使用SupervisorAgent协调多个专业智能体，构建一个完整的内容采集与存储系统。\n\n## 2. 系统架构设计\n\n### 2.1 三层架构模式\n\n我们采用三层架构设计，包括：\n\n1. **Supervisor层**：负责协调和管理其他智能体，接收用户指令并分配任务\n2. **Research层**：负责网页内容提取，包括网站结构分析和内容提取\n3. **FileSystem层**：负责文件操作，包括内容保存、读取和目录管理\n\n### 2.2 智能体角色分工\n\n系统中的三个智能体各自承担不同的职责：\n\n1. **SupervisorAgent**：总协调者，负责理解用户需求，并将任务分配给适当的专业智能体\n2. **Research Agent**：网页内容分析专家，负责使用FireCrawl和Jina工具提取网页内容\n3. **FileSystem Agent**：文件系统管理专家，负责将提取的内容保存到本地文件系统\n\n## 3. 组件实现\n\n### 3.1 Research Agent实现\n\nResearch Agent负责网页内容提取，使用FireCrawl和Jina工具：\n\n```python\n# 创建FireCrawl工具 - 用于网站结构分析\nfirecrawl_tool = FireCrawlTool(\n mode=\"crawl\", # 使用爬取模式\n params={\n \"max_pages\": 5, # 限制爬取页面数量\n }\n)\n\n# 创建Jina Reader工具 - 用于内容提取\njina_reader_tool = JinaSearch()\n\n# 创建Research Agent\nresearch_agent = create_react_agent(\n model=model,\n tools=[firecrawl_tool, jina_reader_tool],\n name=\"research_agent\",\n prompt=(\n \"你是一位专业的网页内容分析专家，擅长提取和分析网站结构与内容。\\n\"\n \"你有两个强大的工具...\\n\"\n # 提示词内容\n ),\n)\n```\n\n### 3.2 FileSystem Agent实现\n\nFileSystem Agent负责文件操作，使用LangChain的FileManagementToolkit：\n\n```python\n# 设置文件系统工具的根目录\noutput_dir = os.path.join(os.path.dirname(__file__), \"output\")\nos.makedirs(output_dir, exist_ok=True)\n\n# 创建文件系统工具集\nfilesystem_toolkit = FileManagementToolkit(\n root_dir=output_dir,\n selected_tools=[\"write_file\", \"read_file\", \"list_directory\"]\n)\n\n# 获取文件系统工具\nfilesystem_tools = filesystem_toolkit.get_tools()\n\n# 创建FileSystem Agent\nfilesystem_agent = create_react_agent(\n model=model,\n tools=filesystem_tools,\n name=\"filesystem_agent\",\n prompt=(\n \"你是一位专业的文件系统管理专家，负责将网页内容保存到本地文件系统。\\n\"\n \"你有以下工具可以使用...\\n\"\n # 提示词内容\n ),\n)\n```\n\n### 3.3 SupervisorAgent实现\n\nSupervisorAgent负责协调Research Agent和FileSystem Agent：\n\n```python\n# 创建Supervisor Agent\nsupervisor = SupervisorAgent(\n agents=[research_agent, filesystem_agent],\n model=model,\n prompt=(\n \"你是一个智能助手的总协调者，负责管理两个专业智能体:\\n\"\n \"1) research_agent: 网页内容分析专家，可以爬取和分析网站内容\\n\"\n \"2) filesystem_agent: 文件系统管理专家，可以将内容保存到本地文件系统\\n\\n\"\n # 提示词内容\n ),\n)\n\n# 创建内存存储器用于保存对话状态\nmemory_saver = MemorySaver()\n\n# 编译得到一个可调用的\"App\"，添加checkpointer实现记忆功能\napp = supervisor.compile(checkpointer=memory_saver)\n```\n\n## 4. 工作流程\n\n### 4.1 基本工作流程\n\n系统的基本工作流程如下：\n\n1. **用户请求**：用户提出网页内容提取和保存的请求\n2. **Supervisor分析**：SupervisorAgent分析用户请求，确定需要调用哪个专业智能体\n3. **内容提取**：如果需要提取网页内容，SupervisorAgent调用Research Agent\n4. **内容保存**：如果需要保存内容，SupervisorAgent将Research Agent的结果传递给FileSystem Agent\n5. **结果返回**：SupervisorAgent将最终结果返回给用户\n\n### 4.2 上下文管理策略\n\n为了有效管理上下文长度，系统采用以下策略：\n\n1. **分批处理**：对于大型网站，采用分批处理策略，每次只处理少量页面\n2. **内容摘要**：对于大型内容，进行摘要处理，减少传递的token数量\n3. **先保存再处理**：对于多页面内容，采用先保存再处理的策略，减轻上下文负担\n\n## 5. 提示词设计\n\n### 5.1 SupervisorAgent提示词\n\nSupervisorAgent的提示词强调任务分配和协调：\n\n```\n你是一个智能助手的总协调者，负责管理两个专业智能体:\n1) research_agent: 网页内容分析专家，可以爬取和分析网站内容\n2) filesystem_agent: 文件系统管理专家，可以将内容保存到本地文件系统\n\n你的工作流程如下:\n1. 分析用户请求，确定是需要网页内容提取还是文件操作，或两者都需要\n2. 如果需要网页内容提取，调用research_agent获取网页内容\n3. 如果需要将提取的内容保存到文件，调用filesystem_agent进行保存\n4. 如果用户同时需要提取内容并保存，先调用research_agent获取内容，再调用filesystem_agent保存内容\n\n重要规则:\n- 不要在一个消息中同时调用多个智能体，必须一步一步来\n- 当调用filesystem_agent保存内容时，必须提供完整的内容和建议的文件名\n- 确保在最终回复中告知用户内容已成功提取和/或保存\n```\n\n### 5.2 Research Agent提示词\n\nResearch Agent的提示词强调网页内容提取的方法论：\n\n```\n你是一位专业的网页内容分析专家，擅长提取和分析网站结构与内容。\n你有两个强大的工具:\n1. 'firecrawl_tool': 用于爬取网站结构和下级页面\n2. 'jina_reader_tool': 用于从特定URL提取结构化内容，获取干净可读的内容\n\n当面对网站分析任务时，请遵循以下方法论:\n1. 分析任务: 明确需要从网站获取什么信息\n2. 网站结构分析: 使用firecrawl_tool爬取网站结构，了解可用页面\n3. 内容提取: 根据网站结构，使用jina_reader_tool从关键页面提取内容\n4. 信息整合: 将提取的内容整合成有条理的分析结果\n```\n\n### 5.3 FileSystem Agent提示词\n\nFileSystem Agent的提示词强调文件操作和内容保存：\n\n```\n你是一位专业的文件系统管理专家，负责将网页内容保存到本地文件系统。\n你有以下工具可以使用:\n1. 'write_file': 用于将内容写入文件\n2. 'read_file': 用于读取文件内容\n3. 'list_directory': 用于列出目录内容\n\n当接收到保存内容的请求时，请遵循以下方法论:\n1. 分析内容: 确定内容的类型和结构\n2. 确定文件名: 根据内容类型和来源创建合适的文件名\n3. 保存内容: 使用write_file工具将内容保存到文件中\n4. 验证保存: 使用read_file或list_directory工具验证内容已正确保存\n```\n\n## 6. 记忆功能实现\n\n### 6.1 使用MemorySaver实现记忆\n\n系统使用LangGraph的MemorySaver实现对话状态的持久化：\n\n```python\n# 创建内存存储器用于保存对话状态\nmemory_saver = MemorySaver()\n\n# 编译得到一个可调用的\"App\"，添加checkpointer实现记忆功能\napp = supervisor.compile(checkpointer=memory_saver)\n```\n\n### 6.2 记忆功能的应用场景\n\n记忆功能在以下场景中特别有用：\n\n1. **多轮对话**：在多轮对话中保持上下文连贯性\n2. **长时间任务**：对于需要长时间处理的任务，可以保存中间状态\n3. **断点续传**：支持任务的暂停和恢复\n\n## 7. 应用案例\n\n### 7.1 提取并保存LangGraph文档\n\n以下是一个完整的应用案例，提取并保存LangGraph文档：\n\n```python\n# 用户请求\ninputs = {\n \"messages\": [\n HumanMessage(content=\"请爬取LangGraph文档网站(https://langchain-ai.github.io/langgraph/how-tos/)的内容，并保存为Markdown文件\")\n ]\n}\n\n# 执行工作流\nfinal_state = None\nfor partial_state in app.stream(inputs, stream_mode=\"values\"):\n # 处理中间状态...\n final_state = partial_state\n # 记录状态\n log_agent_actions(partial_state)\n\n# 最终结果\nprint(\"\\n最终结果:\")\nif final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n```\n\n## 8. 总结\n\n网页提取与文件系统集成是构建完整内容采集系统的关键。通过SupervisorAgent协调Research Agent和FileSystem Agent，我们可以实现网页内容的提取、分析和持久化存储。这种多智能体协作模式不仅提高了系统的模块化程度，也使得每个智能体可以专注于自己的专业领域，从而提高整体系统的效率和质量。\n\n未来的发展方向包括增强对复杂网站的处理能力、支持更多文件格式的存储和处理，以及集成数据库存储以支持更大规模的内容管理。"
},
{
"path": "instructions/08.react_agent_tool_registry.md",
"content": "# 工具注册机制与ReactAgent集成：构建可扩展的智能体系统\n\n## 1. 引言\n\n工具注册机制是构建可扩展智能体系统的关键组件，它允许我们以统一的方式管理和使用各种工具，并将这些工具与ReactAgent集成。本文将详细介绍Mentis框架中的工具注册机制，包括工具注册、分类管理以及与ReactAgent的集成方式。\n\n## 2. 工具注册机制的核心概念\n\n### 2.1 工具注册的意义\n\n工具注册机制提供了以下优势：\n\n1. **统一管理**：集中管理所有可用工具，避免重复创建和配置\n2. **分类组织**：按功能和用途对工具进行分类，便于查找和使用\n3. **动态加载**：支持动态注册和加载工具，提高系统的灵活性\n4. **简化集成**：简化工具与Agent的集成过程，只需从注册表中获取工具列表\n\n### 2.2 工具分类体系\n\n在Mentis框架中，我们使用`ToolCategory`枚举定义了工具的分类体系：\n\n```python\nclass ToolCategory(Enum):\n SEARCH = \"Search\"\n CODE_INTERPRETER = \"Code Interpreter\"\n WEB_BROWSING = \"Web Browsing\"\n DATABASE = \"Database\"\n FILE_SYSTEM = \"FileSystem\"\n OTHER = \"Other\"\n```\n\n这种分类体系使我们能够根据任务需求选择特定类别的工具，提高工具使用的针对性和效率。\n\n## 3. 工具注册机制的实现\n\n### 3.1 全局工具注册表\n\n工具注册机制的核心是一个全局工具注册表，它是一个字典，用于存储所有已注册的工具及其分类信息：\n\n```python\n# 全局工具注册表\n_registered_tools = {}\n```\n\n### 3.2 工具注册函数\n\n`register_tool`函数用于将工具注册到全局注册表中：\n\n```python\ndef register_tool(tool: Tool, category: ToolCategory) -> None:\n \"\"\"注册一个工具到全局字典中，带有分类信息\"\"\"\n if tool.name in _registered_tools:\n raise ValueError(f\"工具名 {tool.name} 已存在，请确保工具名唯一\")\n _registered_tools[tool.name] = {\n \"tool\": tool,\n \"category\": category\n }\n```\n\n### 3.3 工具获取函数\n\n框架提供了多种函数来获取已注册的工具：\n\n```python\ndef get_registered_tools(as_dict: bool = False) -> Union[List[Tool], Dict[str, Dict]]:\n \"\"\"返回所有已注册的工具\"\"\"\n if as_dict:\n return _registered_tools\n return [info[\"tool\"] for info in _registered_tools.values()]\n\ndef get_tools_by_category(category: ToolCategory, return_instances: bool = True) -> List[Union[str, Tool]]:\n \"\"\"返回指定分类的工具列表\"\"\"\n if return_instances:\n return [info[\"tool\"] for name, info in _registered_tools.items() if info[\"category\"] == category]\n return [name for name, info in _registered_tools.items() if info[\"category\"] == category]\n```\n\n## 4. 简化工具注册的辅助函数\n\n### 4.1 直接注册工具的函数\n\n为了简化工具注册过程，框架提供了`register_direct_tool`函数，它可以根据工具类名自动判断工具类别：\n\n```python\ndef register_direct_tool(tool_instance: BaseTool, category: ToolCategory = None) -> None:\n \"\"\"注册直接从langchain_community.tools导入的工具\"\"\"\n if not category:\n # 获取工具类名\n tool_class_name = tool_instance.__class__.__name__\n # 根据工具类名自动判断类别\n category = tool_category_mapping.get(tool_class_name, tool_category_mapping[\"default\"])\n \n # 注册工具\n register_tool(tool_instance, category)\n print(f\"已注册工具: {tool_instance.name} (类别: {category.value})\")\n```\n\n### 4.2 自动注册自定义工具\n\n框架还支持自动扫描和注册自定义工具。在`__init__.py`中，我们使用以下代码自动注册自定义工具：\n\n```python\n# 遍历目录中的所有文件，注册自定义工具\nfor filename in os.listdir(tools_dir):\n # 只处理 .py 文件，且排除 __init__.py 和 registry.py\n if filename.endswith('.py') and filename not in ['__init__.py', 'registry.py']:\n # 提取模块名（去掉 .py 后缀）\n module_name = filename[:-3]\n try:\n # 动态导入模块\n module = importlib.import_module(f'.{module_name}', package='core.tools')\n \n # 查找模块中的工具类（继承自BaseTool的类）\n for name, obj in inspect.getmembers(module):\n # 检查是否是类且是BaseTool的子类\n if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool:\n # 检查该类是否已经被实例化并注册\n tool_name = getattr(obj, 'name', None)\n if tool_name and tool_name not in [info['tool'].name for info in get_registered_tools().values()]:\n # 确定工具类别\n category = getattr(module, 'category', ToolCategory.OTHER)\n # 实例化并注册工具\n try:\n tool_instance = obj()\n register_tool(tool_instance, category)\n print(f\"已注册工具类: {name} (工具名: {tool_instance.name}, 类别: {category.value})\")\n except Exception as e:\n print(f\"实例化工具类 {name} 时出错: {e}\")\n except Exception as e:\n print(f\"导入 {module_name} 时出错: {e}\")\n```\n\n这段代码会自动扫描`core/tools`目录中的所有Python文件，查找继承自`BaseTool`的类，并自动实例化和注册这些工具。\n\n## 5. 与ReactAgent的集成\n\n### 5.1 从注册表获取工具列表\n\n在创建ReactAgent实例时，我们可以从注册表中获取工具列表：\n\n```python\n# 从注册表中获取工具列表\ntools_list = [info[\"tool\"] for info in registered_tools.values()]\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"fed_research_agent\",\n prompt=(\n \"你是一位专业的经济研究分析师，擅长分析复杂的经济问题并提供深入见解。\\n\"\n \"你有多个强大的工具可以搜索网络获取实时信息：\\n\"\n \"- jina_search: 用于进行网络搜索获取最新信息\\n\"\n \"- wikipedia_query_run: 用于查询维基百科获取基础知识\\n\"\n \"- firecrawl_tool: 用于抓取和分析特定网页内容\\n\\n\"\n # 提示词内容\n ),\n)\n```\n\n### 5.2 按类别选择工具\n\n在某些场景下，我们可能只需要特定类别的工具。这时，可以使用`get_tools_by_category`函数：\n\n```python\n# 获取所有搜索类工具\nsearch_tools = get_tools_by_category(ToolCategory.SEARCH)\n\n# 创建专注于搜索的ReactAgent\nsearch_agent = ReactAgent(\n model=model,\n tools=search_tools,\n name=\"search_agent\",\n prompt=\"你是一位专业的信息搜索专家...\"\n)\n```\n\n## 6. 实际应用案例\n\n### 6.1 美联储研究任务\n\n以下是一个完整的应用案例，使用工具注册机制和ReactAgent进行美联储研究：\n\n```python\n# 注册搜索工具\njina_search = JinaSearch()\nwiki_tool = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())\n\n# 使用register_direct_tool函数注册工具\nregister_direct_tool(jina_search)\nregister_direct_tool(wiki_tool)\n\n# 注意：FireCrawlTool已经在core/tools/__init__.py中被注册，这里不需要再次注册\n\n# 获取所有已注册的工具（以字典格式）\nregistered_tools = get_registered_tools(as_dict=True)\n\n# 从注册表中获取工具列表\ntools_list = [info[\"tool\"] for info in registered_tools.values()]\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=model,\n tools=tools_list,\n name=\"fed_research_agent\",\n prompt=(\n \"你是一位专业的经济研究分析师，擅长分析复杂的经济问题并提供深入见解。\\n\"\n # 提示词内容\n ),\n)\n\n# 编译Agent\nagent = react_agent.compile()\n\n# 定义输入\ninputs = {\n \"messages\": [\n HumanMessage(content=\"请提供美联储(Federal Reserve)的详细介绍，包括其历史、结构、职能，以及它如何通过货币政策影响全球经济。\")\n ]\n}\n\n# 执行Agent\nfinal_state = None\nfor partial_state in react_agent.stream(inputs, stream_mode=\"values\"):\n # 处理中间状态...\n pass\n```\n\n### 6.2 结果保存\n\n执行完成后，我们可以将结果保存到文件：\n\n```python\n# 打印最终回答\nif final_state and final_state.get(\"messages\"):\n for message in final_state[\"messages\"]:\n if isinstance(message, AIMessage) and not message.tool_calls:\n print(message.content)\n \n # 将结果保存到文件\n output_dir = os.path.join(os.path.dirname(__file__), \"output\")\n os.makedirs(output_dir, exist_ok=True)\n output_file = os.path.join(output_dir, \"fed_research_report.md\")\n \n with open(output_file, \"w\", encoding=\"utf-8\") as f:\n f.write(\"# 美联储研究报告\\n\\n\")\n f.write(message.content)\n \n print(f\"\\n研究报告已保存到: {output_file}\")\n```\n\n## 7. 最佳实践\n\n### 7.1 工具命名规范\n\n为了避免工具名冲突，建议遵循以下命名规范：\n\n1. 使用有意义的名称，反映工具的功能\n2. 对于同一类别的工具，使用统一的前缀或后缀\n3. 避免使用过于通用的名称，如`search`、`get`等\n\n### 7.2 工具分类策略\n\n合理的工具分类策略可以提高工具使用的效率：\n\n1. 根据工具的主要功能进行分类，而不是实现方式\n2. 对于多功能工具，根据其主要功能进行分类\n3. 只有在无法确定主要功能时，才将工具归类为`OTHER`\n\n### 7.3 提示词优化\n\n在提示词中明确说明可用工具及其用途，可以提高Agent的工具使用效率：\n\n```\n你是一位专业的经济研究分析师，擅长分析复杂的经济问题并提供深入见解。\n你有多个强大的工具可以搜索网络获取实时信息：\n- jina_search: 用于进行网络搜索获取最新信息\n- wikipedia_query_run: 用于查询维基百科获取基础知识\n- firecrawl_tool: 用于抓取和分析特定网页内容\n\n当面对复杂问题时，请遵循以下方法论：\n1. 分解问题：将复杂问题分解为更小的子问题\n2. 制定计划：确定需要搜索哪些信息，以及使用哪些工具\n3. 执行搜索：使用适当的工具执行搜索\n4. 分析结果：分析搜索结果，确定是否需要进一步搜索\n5. 综合信息：将所有搜索结果综合成一个连贯的回答\n```\n\n## 8. 总结\n\n工具注册机制为Mentis框架提供了强大的可扩展性，使得智能体系统能够轻松集成各种工具，并根据任务需求灵活选择合适的工具组合。通过分类管理和自动注册，工具注册机制简化了工具的管理和使用流程，提高了开发效率。\n\n结合ReactAgent，工具注册机制使得智能体能够访问丰富的外部功能，从而处理更复杂的任务。未来的发展方向包括支持更多类型的工具、增强工具的自动发现和选择能力，以及提供更细粒度的工具权限控制。"
},
{
"path": "instructions/09.e2b_sandbox_integration.md",
"content": "# E2B沙箱环境与智能代理集成指南\n\n## 1. 引言\n\nE2B沙箱环境是一个强大的代码执行工具，它提供了安全、隔离的环境来运行Python代码和Shell命令。将E2B沙箱与智能代理（如ReactAgent）集成，可以显著增强代理的能力，使其能够执行代码、处理数据、创建可视化，甚至与文件系统交互。本文将详细介绍E2B沙箱的核心概念、工作原理、实现方式以及在智能代理系统中的应用。\n\n## 2. E2B沙箱环境的核心概念\n\n### 2.1 什么是E2B沙箱\n\nE2B（Execution Environment for Bots）是一个专为AI代理设计的代码执行环境，它提供以下核心功能：\n\n1. **安全隔离**：在隔离的容器中执行代码，防止恶意代码影响宿主系统\n2. **多语言支持**：主要支持Python，同时可通过Shell命令执行其他语言代码\n3. **文件系统操作**：允许创建、读取、写入和管理文件\n4. **包管理**：支持安装和使用第三方Python库\n5. **持久化**：可以在会话之间保持状态和文件\n\n### 2.2 E2B沙箱与代码解释器的关系\n\nE2B沙箱是一种特殊的代码解释器实现，它不仅能执行代码，还提供了完整的操作系统环境（基于Debian）。这使得它比简单的代码解释器功能更强大，能够：\n\n- 执行系统命令\n- 管理文件和目录\n- 安装和使用各种软件包\n- 运行网络服务\n- 处理复杂的数据分析和可视化任务\n\n## 3. E2B沙箱的实现\n\n### 3.1 E2BCodeInterpreterTool类的设计\n\n在我们的实现中，`E2BCodeInterpreterTool`类继承自LangChain的`BaseTool`，提供了与E2B沙箱交互的接口：\n\n```python\nclass E2BCodeInterpreterTool(BaseTool):\n \"\"\"使用E2B SDK执行Python代码的工具\n \n 该工具创建一个安全的沙箱环境，用于执行Python代码，并返回执行结果、\n 标准输出、标准错误和任何错误信息。\n \"\"\"\n \n name: str = \"e2b_code_interpreter\"\n description: str = (\n \"在安全的 Debian 基础沙箱环境中执行 Python 代码或 shell 命令，并返回结果。\"\n \"适用于数据分析、可视化、复杂计算以及系统操作。\"\n \"输入应为有效的 Python 代码字符串，或以 '!' 开头的 shell 命令。\"\n \"常见 Python 库（如 numpy、pandas 和 matplotlib）已预装，若需其他库，可通过 pip 安装。\"\n \"沙箱环境充分利用 Debian 系统的强大功能，支持广泛的操作。\"\n )\n```\n\n### 3.2 核心方法\n\n`E2BCodeInterpreterTool`类提供了以下核心方法：\n\n1. **_initialize_sandbox()**: 初始化沙箱环境\n2. **_run()**: 在沙箱中执行代码并返回结果\n3. **close()**: 关闭沙箱并释放资源\n4. **format_to_tool_message()**: 将执行结果格式化为工具消息\n\n### 3.3 沙箱初始化与资源管理\n\n沙箱初始化过程包括：\n\n1. 检查是否安装了`e2b_code_interpreter`包\n2. 验证是否设置了`E2B_API_KEY`环境变量\n3. 创建`Sandbox`实例\n4. 设置沙箱状态标志\n\n资源管理方面，工具提供了`close()`方法来释放沙箱资源：\n\n```python\ndef close(self):\n \"\"\"关闭沙箱，释放资源\"\"\"\n if hasattr(self, \"sandbox\") and self._is_available and self.sandbox is not None:\n try:\n print(\"正在关闭E2B沙箱并释放资源...\")\n self.sandbox.kill()\n print(\"E2B沙箱已成功关闭\")\n except Exception as e:\n print(f\"关闭E2B沙箱时出错: {str(e)}\")\n```\n\n## 4. 将E2B沙箱与ReactAgent集成\n\n### 4.1 基本集成流程\n\n将E2B沙箱与ReactAgent集成的基本流程如下：\n\n1. **注册E2B工具**：将`E2BCodeInterpreterTool`注册到工具注册表中\n2. **创建ReactAgent**：使用包含E2B工具的工具列表初始化ReactAgent\n3. **设计提示词**：编写强调代码执行能力的提示词\n4. **执行工作流**：让Agent使用E2B工具执行代码并处理结果\n\n### 4.2 代码示例\n\n以下是一个基本的集成示例：\n\n```python\n# 导入必要的库\nfrom core.agents.react_agent import ReactAgent\nfrom core.tools.registry import get_tools_by_category, ToolCategory\nfrom langchain_openai import ChatOpenAI\n\n# 获取代码解释器工具\ntools_list = get_tools_by_category(ToolCategory.CODE_INTERPRETER)\n\n# 创建ReactAgent实例\nreact_agent = ReactAgent(\n model=ChatOpenAI(model=\"gpt-4o-mini\"),\n tools=tools_list,\n prompt=(\n \"你是一位专业的数据分析师，可以使用Python代码解决问题。\\n\"\n \"你有强大的代码执行工具可以使用：\\n\"\n \"- e2b_code_interpreter: 用于执行Python代码和shell命令\\n\"\n ),\n)\n\n# 编译Agent\nagent = react_agent.compile()\n\n# 执行任务\nresult = agent.invoke({\"messages\": [HumanMessage(content=\"分析以下数据并创建可视化...\")]})\n```\n\n## 5. E2B沙箱的高级功能\n\n### 5.1 文件系统操作\n\nE2B沙箱提供了完整的文件系统操作能力，可以：\n\n- 创建和管理目录结构\n- 读写文本和二进制文件\n- 列出目录内容\n- 移动和删除文件\n\n示例代码：\n\n```python\n# 在沙箱中创建目录和文件\ncode = \"\"\"\n# 创建目录\nimport os\nos.makedirs('test_dir/subdir', exist_ok=True)\n\n# 创建并写入文件\nwith open('test_dir/example.txt', 'w') as f:\n f.write('Hello from E2B sandbox!')\n \n# 列出目录内容\nprint(os.listdir('test_dir'))\n\n# 读取文件内容\nwith open('test_dir/example.txt', 'r') as f:\n content = f.read()\n print(f'文件内容: {content}')\n\"\"\"\n\n# 执行代码\nresult = e2b_tool.invoke({\"code\": code})\n```\n\n### 5.2 包管理\n\nE2B沙箱允许安装和使用第三方Python库：\n\n```python\n# 安装并使用第三方库\ncode = \"\"\"\n# 安装pandas库\n!pip install pandas matplotlib\n\n# 使用pandas进行数据分析\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\n# 创建示例数据\ndata = {'Category': ['A', 'B', 'C', 'D'], 'Values': [10, 25, 15, 30]}\ndf = pd.DataFrame(data)\n\n# 打印数据\nprint(df)\n\n# 创建可视化\nplt.figure(figsize=(8, 4))\nplt.bar(df['Category'], df['Values'])\nplt.title('Sample Bar Chart')\nplt.savefig('chart.png')\nprint('图表已保存为chart.png')\n\"\"\"\n\n# 执行代码\nresult = e2b_tool.invoke({\"code\": code})\n```\n\n### 5.3 从沙箱下载文件\n\n可以将沙箱中生成的文件下载到本地系统：\n\n```python\ndef download_file_from_sandbox(sandbox, sandbox_path, local_path):\n \"\"\"从沙箱下载文件到本地\"\"\"\n try:\n # 从沙箱读取文件内容\n content = sandbox.files.read(sandbox_path)\n \n # 确保目标目录存在\n os.makedirs(os.path.dirname(local_path), exist_ok=True)\n \n # 写入本地文件\n with open(local_path, 'w', encoding='utf-8') as file:\n file.write(content)\n \n print(f\"文件已从沙箱下载到本地: {local_path}\")\n return True\n except Exception as e:\n print(f\"从沙箱下载文件时出错: {str(e)}\")\n return False\n```\n\n## 6. 实际应用案例\n\n### 6.1 数据分析与可视化\n\nE2B沙箱特别适合数据分析和可视化任务，可以：\n\n- 加载和处理各种格式的数据（CSV、JSON、Excel等）\n- 使用pandas进行数据清洗和转换\n- 使用matplotlib、seaborn等创建可视化\n- 生成分析报告\n\n### 6.2 文件处理与转换\n\nE2B沙箱可以处理各种文件格式的转换和处理：\n\n- 文本文件处理（如日志分析）\n- 图像处理和转换\n- 数据格式转换（如CSV到JSON）\n- 文档生成（如生成HTML或PDF报告）\n\n### 6.3 Web爬虫与API调用\n\nE2B沙箱可以执行网络相关任务：\n\n- 使用requests或BeautifulSoup进行网页爬取\n- 调用各种API并处理响应\n- 下载和处理网络资源\n\n## 7. 最佳实践与注意事项\n\n### 7.1 安全考虑\n\n虽然E2B沙箱提供了隔离环境，但在使用时仍需注意：\n\n- 不要在沙箱中处理敏感数据\n- 避免执行未经验证的用户输入代码\n- 限制沙箱的网络访问权限\n- 定期关闭和重新创建沙箱实例\n\n### 7.2 资源管理\n\nE2B沙箱会消耗系统资源，因此：\n\n- 在不需要时关闭沙箱（使用`close()`方法）\n- 避免在单个沙箱中运行过多或过大的任务\n- 监控沙箱的内存和CPU使用情况\n\n### 7.3 错误处理\n\n在与E2B沙箱交互时，应当实施健壮的错误处理：\n\n- 捕获并处理代码执行异常\n- 验证沙箱初始化是否成功\n- 提供有意义的错误消息给用户\n- 实现重试机制处理临时故障\n\n## 8. 总结\n\nE2B沙箱为智能代理提供了强大的代码执行能力，使其能够处理各种复杂任务。通过将E2B沙箱与ReactAgent集成，我们可以创建能够执行代码、处理数据、创建可视化，甚至与文件系统交互的智能系统。\n\n正确使用E2B沙箱需要理解其核心概念、实现方式和最佳实践。通过本文的指导，开发者应能够有效地将E2B沙箱集成到自己的智能代理系统中，并充分利用其强大功能。\n\n## 9. 参考资源\n\n- [E2B官方文档](https://e2b.dev/docs)\n- [E2B Code Interpreter SDK](https://github.com/e2b-dev/code-interpreter)\n- [LangChain工具集成指南](https://python.langchain.com/docs/integrations/tools)\n- [ReactAgent文档](https://python.langchain.com/docs/modules/agents/agent_types/react)"
},
{
"path": "log_analyzer.py",
"content": "import re\nimport sys\nimport argparse\nfrom collections import defaultdict\nimport json\n\ndef parse_log_file(file_path):\n \"\"\"Parse the execution log file and extract agent interactions.\"\"\"\n with open(file_path, 'r', encoding='utf-8') as f:\n content = f.read()\n \n # Extract different sections of the log\n sections = content.split(\"================================ Human Message =================================\")\n if len(sections) > 1:\n main_content = sections[1] # Skip header\n else:\n main_content = content\n \n # Extract messages\n messages = []\n \n # Pattern for AI messages\n ai_pattern = r\"================================== Ai Message ==================================\\nName: (\\w+)\\n\\n(.*?)(?=(==================================|$))\"\n ai_matches = re.finditer(ai_pattern, main_content, re.DOTALL)\n \n for match in ai_matches:\n agent_name = match.group(1)\n message_content = match.group(2).strip()\n \n # Check if message has tool calls\n tool_calls = []\n tool_call_pattern = r\"Tool Calls:\\n(.*?)(?=\\n==================================|$)\"\n tool_call_match = re.search(tool_call_pattern, message_content, re.DOTALL)\n if tool_call_match:\n # Extract tool calls\n tool_calls_text = tool_call_match.group(1)\n tool_call_entries = re.findall(r\" (\\w+) \\(([^)]+)\\)\", tool_calls_text)\n tool_calls = [{\"name\": name, \"id\": call_id} for name, call_id in tool_call_entries]\n \n # Remove tool calls from the message content\n message_content = re.sub(r\"Tool Calls:.*?(?=\\n==================================|$)\", \"\", message_content, flags=re.DOTALL).strip()\n \n messages.append({\n \"role\": \"agent\",\n \"agent\": agent_name,\n \"content\": message_content,\n \"tool_calls\": tool_calls\n })\n \n # Pattern for Tool messages\n tool_pattern = r\"================================= Tool Message =================================\\nName: (\\w+)\\n\\n(.*?)(?=(==================================|$))\"\n tool_matches = re.finditer(tool_pattern, main_content, re.DOTALL)\n \n for match in tool_matches:\n tool_name = match.group(1)\n tool_content = match.group(2).strip()\n \n messages.append({\n \"role\": \"tool\",\n \"tool\": tool_name,\n \"content\": tool_content\n })\n \n # Sort messages by their position in the log\n messages.sort(key=lambda x: main_content.find(x[\"content\"]))\n \n return messages\n\ndef analyze_agent_interactions(messages):\n \"\"\"Analyze the interactions between agents.\"\"\"\n interactions = []\n current_sender = None\n tool_call_map = {}\n \n for i, msg in enumerate(messages):\n if msg[\"role\"] == \"agent\":\n current_sender = msg[\"agent\"]\n # Check if this agent is using tool calls\n for tool_call in msg.get(\"tool_calls\", []):\n tool_name = tool_call[\"name\"]\n tool_id = tool_call[\"id\"]\n tool_call_map[tool_id] = {\n \"sender\": current_sender,\n \"tool\": tool_name\n }\n interactions.append({\n \"step\": i,\n \"from\": current_sender,\n \"to\": f\"SYSTEM ({tool_name})\",\n \"action\": f\"Called tool {tool_name}\",\n \"content\": f\"Tool call ID: {tool_id}\"\n })\n elif msg[\"role\"] == \"tool\":\n # Find which agent invoked this tool\n for prev_msg in reversed(messages[:i]):\n if prev_msg[\"role\"] == \"agent\" and any(tc[\"name\"] == msg[\"tool\"] for tc in prev_msg.get(\"tool_calls\", [])):\n sender = prev_msg[\"agent\"]\n break\n else:\n sender = \"SYSTEM\"\n \n interactions.append({\n \"step\": i,\n \"from\": f\"SYSTEM ({msg['tool']})\",\n \"to\": sender,\n \"action\": f\"Tool response\",\n \"content\": msg[\"content\"]\n })\n \n return interactions\n\ndef visualize_interactions(interactions):\n \"\"\"Visualize the interactions between agents.\"\"\"\n print(\"\\n\" + \"=\"*100)\n print(\" \"*40 + \"AGENT INTERACTIONS SUMMARY\")\n print(\"=\"*100 + \"\\n\")\n \n for idx, interaction in enumerate(interactions):\n print(f\"[{idx+1}] {interaction['from']} → {interaction['to']}\")\n print(f\" Action: {interaction['action']}\")\n content = interaction['content']\n if len(content) > 100:\n content = content[:97] + \"...\"\n print(f\" Content: {content}\\n\")\n\ndef visualize_conversation_flow(messages):\n \"\"\"Visualize the conversation flow between agents.\"\"\"\n print(\"\\n\" + \"=\"*100)\n print(\" \"*40 + \"CONVERSATION FLOW\")\n print(\"=\"*100 + \"\\n\")\n \n for idx, message in enumerate(messages):\n if message[\"role\"] == \"agent\":\n agent_name = message[\"agent\"]\n print(f\"[{idx+1}] Agent: {agent_name}\")\n content = message[\"content\"]\n if len(content) > 150:\n content = content[:147] + \"...\"\n print(f\" Content: {content}\")\n \n if message.get(\"tool_calls\"):\n tools = \", \".join([tc[\"name\"] for tc in message[\"tool_calls\"]])\n print(f\" Tools Called: {tools}\")\n else:\n print(f\"[{idx+1}] Tool: {message['tool']}\")\n content = message[\"content\"]\n if len(content) > 100:\n content = content[:97] + \"...\"\n print(f\" Response: {content}\")\n print()\n\ndef main():\n parser = argparse.ArgumentParser(description='Analyze Mentis execution logs.')\n parser.add_argument('log_file', help='Path to the log file')\n parser.add_argument('--format', choices=['interactions', 'flow', 'all'], default='all',\n help='Output format: interactions, flow, or all')\n \n args = parser.parse_args()\n \n try:\n messages = parse_log_file(args.log_file)\n interactions = analyze_agent_interactions(messages)\n \n if args.format in ['interactions', 'all']:\n visualize_interactions(interactions)\n \n if args.format in ['flow', 'all']:\n visualize_conversation_flow(messages)\n \n except Exception as e:\n print(f\"Error: {e}\")\n sys.exit(1)\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\"setuptools>=42\", \"wheel\"]\nbuild-backend = \"setuptools.build_meta\"\nreadme = \"README.md\"\nrequires-python = \">=3.11\"\n\n[project]\nname = \"mentis\"\nversion = \"0.1.0\"\ndescription = \"A Multi-Agents project based on langgraph\"\nrequires-python = \">=3.11\"\ndependencies = [\n \"dotenv>=0.9.9\",\n \"langchain-community>=0.3.19\",\n \"langchain-core>=0.3.45\",\n \"langchain-openai>=0.3.8\",\n \"langgraph>=0.3.11\",\n \"pydantic>=2.10.6\",\n \"typing-extensions>=4.12.2\",\n \"python-dotenv>=1.0.0\",\n \"firecrawl-py\",\n \"wikipedia>=1.4.0\",\n \"serpapi>=0.1.5\",\n \"google-search-results>=2.4.2\",\n \"duckduckgo-search>=7.5.2\",\n \"arxiv>=2.1.3\",\n \"rizaio>=0.9.0\",\n \"e2b-code-interpreter>=1.1.0\",\n \"fastapi>=0.115.11\",\n \"uvicorn>=0.34.0\",\n \"sse-starlette>=2.2.1\",\n \"exa-py>=1.9.1\",\n \"tavily-python>=0.5.1\",\n \"replicate>=1.0.4\",\n \"langchain-mcp-adapters>=0.0.7\",\n \"mcp>=1.6.0\",\n \"playwright>=1.51.0\",\n \"pillow>=11.2.1\",\n \"yfinance>=0.2.55\",\n]\n\n[tool.setuptools]\npackages = [\"core\"]\n"
},
{
"path": "requirements.txt",
"content": "# This file was autogenerated by uv via the following command:\n# uv pip compile pyproject.toml -o requirements.txt\naiohappyeyeballs==2.6.1\n # via aiohttp\naiohttp==3.11.14\n # via langchain-community\naiosignal==1.3.2\n # via aiohttp\nannotated-types==0.7.0\n # via pydantic\nanyio==4.9.0\n # via\n # httpx\n # openai\n # rizaio\narxiv==2.1.3\n # via mentis (pyproject.toml)\nattrs==25.3.0\n # via\n # aiohttp\n # e2b\n # e2b-code-interpreter\nbeautifulsoup4==4.13.3\n # via wikipedia\ncertifi==2025.1.31\n # via\n # httpcore\n # httpx\n # requests\ncharset-normalizer==3.4.1\n # via requests\nclick==8.1.8\n # via duckduckgo-search\ndataclasses-json==0.6.7\n # via langchain-community\ndistro==1.9.0\n # via\n # openai\n # rizaio\ndotenv==0.9.9\n # via mentis (pyproject.toml)\nduckduckgo-search==7.5.2\n # via mentis (pyproject.toml)\ne2b==1.1.0\n # via e2b-code-interpreter\ne2b-code-interpreter==1.1.0\n # via mentis (pyproject.toml)\nfeedparser==6.0.11\n # via arxiv\nfirecrawl-py==1.14.1\n # via mentis (pyproject.toml)\nfrozenlist==1.5.0\n # via\n # aiohttp\n # aiosignal\ngoogle-search-results==2.4.2\n # via mentis (pyproject.toml)\ngreenlet==3.1.1\n # via sqlalchemy\nh11==0.14.0\n # via httpcore\nhttpcore==1.0.7\n # via\n # e2b\n # httpx\nhttpx==0.28.1\n # via\n # e2b\n # e2b-code-interpreter\n # langgraph-sdk\n # langsmith\n # openai\n # rizaio\nhttpx-sse==0.4.0\n # via langchain-community\nidna==3.10\n # via\n # anyio\n # httpx\n # requests\n # yarl\njiter==0.9.0\n # via openai\njsonpatch==1.33\n # via langchain-core\njsonpointer==3.0.0\n # via jsonpatch\nlangchain==0.3.20\n # via langchain-community\nlangchain-community==0.3.19\n # via mentis (pyproject.toml)\nlangchain-core==0.3.45\n # via\n # mentis (pyproject.toml)\n # langchain\n # langchain-community\n # langchain-openai\n # langchain-text-splitters\n # langgraph\n # langgraph-checkpoint\n # langgraph-prebuilt\nlangchain-openai==0.3.9\n # via mentis (pyproject.toml)\nlangchain-text-splitters==0.3.6\n # via langchain\nlanggraph==0.3.11\n # via mentis (pyproject.toml)\nlanggraph-checkpoint==2.0.20\n # via\n # langgraph\n # langgraph-prebuilt\nlanggraph-prebuilt==0.1.3\n # via langgraph\nlanggraph-sdk==0.1.57\n # via langgraph\nlangsmith==0.3.15\n # via\n # langchain\n # langchain-community\n # langchain-core\nlxml==5.3.1\n # via duckduckgo-search\nmarshmallow==3.26.1\n # via dataclasses-json\nmsgpack==1.1.0\n # via langgraph-checkpoint\nmultidict==6.2.0\n # via\n # aiohttp\n # yarl\nmypy-extensions==1.0.0\n # via typing-inspect\nnest-asyncio==1.6.0\n # via firecrawl-py\nnumpy==2.2.4\n # via langchain-community\nopenai==1.66.3\n # via langchain-openai\norjson==3.10.15\n # via\n # langgraph-sdk\n # langsmith\npackaging==24.2\n # via\n # e2b\n # langchain-core\n # langsmith\n # marshmallow\nprimp==0.14.0\n # via duckduckgo-search\npropcache==0.3.0\n # via\n # aiohttp\n # yarl\nprotobuf==5.29.3\n # via e2b\npydantic==2.10.6\n # via\n # mentis (pyproject.toml)\n # firecrawl-py\n # langchain\n # langchain-core\n # langsmith\n # openai\n # pydantic-settings\n # rizaio\npydantic-core==2.27.2\n # via pydantic\npydantic-settings==2.8.1\n # via langchain-community\npython-dateutil==2.9.0.post0\n # via e2b\npython-dotenv==1.0.1\n # via\n # mentis (pyproject.toml)\n # dotenv\n # firecrawl-py\n # pydantic-settings\npyyaml==6.0.2\n # via\n # langchain\n # langchain-community\n # langchain-core\nregex==2024.11.6\n # via tiktoken\nrequests==2.32.3\n # via\n # arxiv\n # firecrawl-py\n # google-search-results\n # langchain\n # langchain-community\n # langsmith\n # requests-toolbelt\n # serpapi\n # tiktoken\n # wikipedia\nrequests-toolbelt==1.0.0\n # via langsmith\nrizaio==0.9.0\n # via mentis (pyproject.toml)\nserpapi==0.1.5\n # via mentis (pyproject.toml)\nsgmllib3k==1.0.0\n # via feedparser\nsix==1.17.0\n # via python-dateutil\nsniffio==1.3.1\n # via\n # anyio\n # openai\n # rizaio\nsoupsieve==2.6\n # via beautifulsoup4\nsqlalchemy==2.0.39\n # via\n # langchain\n # langchain-community\ntenacity==9.0.0\n # via\n # langchain-community\n # langchain-core\ntiktoken==0.9.0\n # via langchain-openai\ntqdm==4.67.1\n # via openai\ntyping-extensions==4.12.2\n # via\n # mentis (pyproject.toml)\n # anyio\n # beautifulsoup4\n # e2b\n # langchain-core\n # openai\n # pydantic\n # pydantic-core\n # rizaio\n # sqlalchemy\n # typing-inspect\ntyping-inspect==0.9.0\n # via dataclasses-json\nurllib3==2.3.0\n # via requests\nwebsockets==15.0.1\n # via firecrawl-py\nwikipedia==1.4.0\n # via mentis (pyproject.toml)\nyarl==1.18.3\n # via aiohttp\nzstandard==0.23.0\n # via langsmith\n"
},
{
"path": "setup.py",
"content": "from setuptools import setup\n\nsetup()\n"
},
{
"path": "super_agents/__init__.py",
"content": ""
},
{
"path": "super_agents/browser_use/README.md",
"content": "n# Browser Agent (基于 LangGraph) - super_agents/browser_use\n\n## 概述\n\n本项目实现了一个基于 LangGraph 框架的 Web 浏览和交互 Agent。其核心目标是让一个大型语言模型 (LLM) 能够像人一样理解任务指令，自主地控制浏览器（通过 Playwright）来访问网页、分析内容、与页面元素交互（点击、输入、滚动等），并最终完成用户指定的任务，例如信息提取、表单填写、在线搜索等。\n\n该 Agent 采用了多模态感知的设计思路，结合了传统的 DOM/Accessibility Tree 分析和可选的视觉语言模型 (VLM) 分析，以期在复杂网页上获得更鲁棒的理解和定位能力。\n\n## 核心技术栈\n\n* **流程编排:** LangGraph (LangChain 的状态图编排库)\n* **浏览器自动化:** Playwright (异步 Python 版本)\n* **模型调用:** LangChain ChatModels (`langchain-openai`, `langchain-community` 等)\n* **语言模型 (LLM/VLM):**\n * **规划/决策 LLM:** 可配置，支持 OpenAI, Groq, xAI (Grok), 及其他 OpenAI 兼容 API (通过 `llm.py` 和 `.env` 配置)。\n * **视觉分析 VLM:** 可选，通过 OpenRouter 调用支持 Vision 的模型 (如 Qwen-VL, GPT-4o, Claude 3.5 Sonnet 等) (通过 `detector.py` 和 `.env` 配置)。\n* **依赖管理:** `uv` (或 `pip`)\n* **配置:** `.env` 文件\n\n## 项目架构\n\n项目主要文件和目录结构如下：\n\n```\nsuper_agents/\n└── browser_use/ # Agent 根目录\n ├── agent/ # LangGraph 核心实现\n │ ├── __init__.py\n │ ├── graph.py # 定义 LangGraph 图结构、节点连接、条件边\n │ ├── nodes.py # 定义图中各节点 (Node) 的具体执行逻辑 (AgentNodes 类)\n │ ├── state.py # 定义 Agent 在图中流转的状态 (AgentState)\n │ ├── schemas.py # 定义数据模型 (如动作指令 Action Schema, VLM 输出 Schema)\n │ └── prompts.py # 管理发送给规划 LLM 和 VLM 的 Prompt 模板\n │\n ├── browser/ # 浏览器交互底层实现 (基于原始项目代码)\n │ ├── __init__.py\n │ ├── browser.py # 核心 Browser 类，封装 Playwright 操作、感知方法 (get_content, update_state)\n │ ├── detector.py # 视觉检测器类，实现 VLM 调用逻辑\n │ ├── models.py # 定义浏览器状态、元素等 Pydantic 模型\n │ ├── utils.py # 浏览器相关的工具函数\n │ └── findVisibleInteractiveElements.js # 用于 DOM 元素检测的 JS 脚本\n │\n ├── llm/ # LLM 相关实现\n │ ├── __init__.py\n │ └── llm.py # 定义 ChatOpenRouter (VLM 调用), initialize_llms (规划 LLM 初始化), generate_structured_output\n │\n ├── main.py # Agent 的主入口脚本\n ├── requirements.txt # Python 依赖列表\n ├── README.md # 本文件\n └── .env # 环境变量配置文件 (需要手动创建)\n```\n\n## 核心概念与设计\n\n### 1. LangGraph 状态机\n\nAgent 的核心控制流由 LangGraph 管理。它被实现为一个状态机 (`StateGraph`)：\n\n* **状态 (State):** `agent/state.py` 中的 `AgentState` (TypedDict) 定义了在节点间传递的数据，包含当前任务、浏览器内容/状态、LLM 解析出的动作、历史记录、错误信息等。\n* **节点 (Nodes):** `agent/nodes.py` 中的 `AgentNodes` 类定义了主要的处理步骤，作为图的节点：\n * `get_browser_state`: 调用 `Browser` 类的感知方法 (当前是 `get_content`) 获取页面信息。\n * `plan_action`: 将感知信息和任务包装成 Prompt，调用**规划 LLM** (通过 `llm.py` 的 `generate_structured_output`) 获取结构化的下一步动作 JSON。\n * `execute_action`: 解析 `plan_action` 返回的动作 JSON，并调用 `Browser` 类中相应的交互方法 (如 `Maps_to`, `click`, `type`, `scroll`, `wait`) 执行操作。\n* **边 (Edges):** `agent/graph.py` 定义了节点间的固定跳转（如 `get_browser_state` -> `plan_action`）和条件跳转（如 `execute_action` 后根据 `should_end` 函数判断是结束 `END` 还是回到 `get_browser_state`）。\n\n### 2. 感知 (Perception)\n\nAgent 通过 `browser.py` 中的 `Browser.get_content()` 方法（被 `get_browser_state` 节点调用）来理解当前网页状态。该方法整合了多种信息源，旨在为 LLM 提供丰富且相对简洁的页面表示：\n\n* **简化 DOM:** 通过注入并执行 `SIMPLIFY_PAGE_SCRIPT` JavaScript，移除无关标签（脚本、样式等），提取关键交互元素及其属性，并为这些元素添加 `x-pw-id` 唯一标识。结果以伪 HTML 字符串形式返回。\n* **可访问性树 (AX Tree):** (当前实现中暂时禁用/存在错误) 理论上通过 `page.accessibility.snapshot()` 获取页面的语义结构信息（角色、名称等），以 JSON 字符串形式返回。\n* **视觉元素 (VLM):** (可选，需配置)\n * 如果 `.env` 文件中配置了 VLM (`OPENROUTER_API_KEY`, `VLM_API_MODEL`)，`get_content` 会调用 `Detector` 实例。\n * `Detector` (在 `browser/detector.py` 中) 使用 LangChain 的 `ChatOpenRouter` (在 `llm.py` 中定义) 调用配置的 VLM API。\n * 通过精心设计的 Prompt (`VLM_PROMPT_TEMPLATE`) 请求 VLM 返回页面交互元素的**描述、类型和边界框百分比坐标** (JSON 格式)。\n * `Detector` 解析 VLM 返回的 JSON，创建 `InteractiveElement` 对象列表（目前坐标是占位符）。\n * `get_content` 将这些视觉元素信息格式化为**文本摘要** (包含 VLM 分配的 ID 和边界框信息)。\n* **合并与截断:** `get_content` 将 URL、简化 DOM、AX Tree (如果成功)、视觉元素摘要合并为一个长的文本字符串，并在超过 `max_length` 时进行截断，最后返回给 `plan_action` 节点。\n\n### 3. 规划 (Planning)\n\n* `plan_action` 节点接收 `get_content` 返回的**混合文本字符串**。\n* `agent/prompts.py` 中的 `create_agent_prompt` 函数将任务描述、历史记录、错误信息（如果有）和这段混合文本整合成一个 Prompt。\n* 该 Prompt 被发送给**规划 LLM**（通过 `llm.py` 中的 `generate_structured_output` 函数，该函数使用 LangChain 的 `.with_structured_output()` 功能）。\n* LLM 被要求分析输入信息，决定下一步动作，并**严格按照 `agent/schemas.py` 中定义的 `LLMResponse` Pydantic 模型返回一个包含具体动作指令的 JSON**。Prompt 中包含了对生成**健壮 CSS 选择器**（优先使用稳定 ID、aria-label、文本内容，结合 `x-pw-id`）的明确指导。\n\n### 4. 行动 (Action Execution)\n\n* `execute_action` 节点接收规划 LLM 返回的结构化动作 JSON (存储在 `state['parsed_action']`)。\n* 它解析出动作类型 (`type`) 和参数 (`selector`, `url`, `text`, `direction` 等)。\n* 根据动作类型，调用 `browser/browser.py` 中 `Browser` 类对应的**简单交互方法**:\n * `Maps_to(url)`\n * `click(selector)`\n * `type(selector, text)`\n * `scroll(direction)`\n * `wait(milliseconds)`\n* 这些方法内部使用 Playwright 的 `page.goto`, `page.locator(...).click`, `page.locator(...).fill`, `page.evaluate(...)` 等函数执行实际的浏览器操作。\n* 如果动作是 `finish` 或 `error`，图流程会根据 `graph.py` 中的 `should_end` 函数判断并终止。\n\n## 安装与配置\n\n1. **环境:** 推荐使用 Python 3.10+。\n2. **依赖安装:**\n * 克隆项目。\n * 进入 `super_agents/browser_use/` 目录。\n * 创建并激活虚拟环境 (使用 uv):\n ```bash\n uv venv\n source .venv/bin/activate # Linux/macOS\n # 或者 .venv\\Scripts\\activate # Windows\n ```\n * 安装依赖项 (使用 uv):\n ```bash\n uv sync\n ```\n3. **Playwright 浏览器:** 运行 `playwright install` (至少需要 `playwright install chromium`) 来下载浏览器驱动。\n4. **环境变量 (`.env` 文件):**\n * 在 `super_agents/browser_use/` 目录下创建一个名为 `.env` 的文件。\n * 参考我们之前讨论的 `.env` 示例，**至少需要配置**：\n * **规划 LLM:** 选择一个 Provider (如 `openai`), 设置 `LLM_PROVIDER`, `LLM_MODEL_NAME`, 以及对应的 API Key (如 `OPENAI_API_KEY`)。\n * **VLM (可选):** 如果要启用视觉分析，设置 `OPENROUTER_API_KEY` 和 `VLM_API_MODEL` (设置为 OpenRouter 上支持视觉的模型 ID，如 `openai/gpt-4.1`等)。\n * 确保 `.env` 文件被正确加载（`main.py` 和 `llm.py` 中包含 `load_dotenv()`）。\n\n## 如何运行\n\n1. 确保已完成安装和配置。\n2. 激活虚拟环境。\n3. 从 `super_agents/` 目录（即 `browser_use` 的**上级**目录）运行 `main.py`：\n\n ```bash\n # 基本运行\n python -m browser_use.main \"您的任务描述\"\n\n # 示例：访问 Hacker News 并获取导航栏信息\n python -m browser_use.main \"访问 news.ycombinator.com，返回页面导航栏信息\"\n\n # 示例：使用其他命令行参数（如果有定义，如下面的最大步骤数）\n python -m browser_use.main \"您的任务描述\" --max-steps 30\n ```\n\n## 当前状态、局限性与未来工作\n\n* **核心流程:** Agent 的基本 LangGraph 流程（感知-规划-行动循环）、浏览器操作（导航、点击、输入、滚动、等待）、规划 LLM 调用、可选的 VLM 调用**已经跑通**，能够完成一些多步骤的 Web 任务。\n* **视觉集成 (部分):** VLM 调用流程已集成到 `Detector` 类并通过 `get_content` 触发（需配置 API Key 和 Model）。VLM 能够返回 JSON 格式的检测结果，并且可以被成功解析为内部数据结构 (`InteractiveElement`)。\n* **局限性 & 待完善:**\n 1. **VLM 坐标处理:** VLM 返回的是百分比坐标，但在解析时 (`_parse_vlm_detections`) 目前使用的是**占位符像素坐标**。需要获取截图的实际尺寸，实现准确的百分比到像素的转换，才能真正利用视觉信息进行定位。\n 2. **动作执行方式:** 当前 `execute_action` 仍然**完全依赖规划 LLM 生成的 CSS 选择器**。尚未实现基于 VLM 的元素 ID 或坐标进行点击/输入的操作，这限制了视觉能力的实际应用，特别是在 CSS 选择器不可靠的复杂页面上。\n 3. **感知信息完整性:**\n * **内容截断:** `get_content` 方法返回的内容会因为 `max_length` 限制而被截断，影响需要完整页面信息的任务（如“摘录全文”）。需要增大 `max_length` 或实现更智能的内容提取/滚动策略。\n * **AX Tree 缺失:** 获取 Accessibility Tree 的代码目前被注释或存在错误，导致缺少重要的语义信息。需要修复 `page.accessibility.snapshot()` 调用。\n 4. **滚动策略:** 当前依靠 Prompt 指示 LLM 进行滚动。可能需要更鲁棒的机制来处理长页面，例如 Agent 内部判断是否需要滚动，或者让 LLM 能获取滚动状态信息。\n 5. **Pydantic V1 警告:** 调用规划 LLM 的 `with_structured_output` 时仍然出现 Pydantic V1 警告，建议保持 LangChain 相关库和 Pydantic 为最新版本。\n 6. **错误处理:** 当前错误处理相对简单（例如 VLM 解析失败直接跳过，执行错误直接终止图），可以增加更复杂的重试、回退或用户介入机制。\n 7. **VLM 稳定性:** VLM 能否稳定、准确地返回所需的 JSON 格式和边界框，高度依赖所选模型和 Prompt，可能需要进一步调优。\n\n* **未来工作:**\n * 修复 AX Tree 获取。\n * 实现 VLM 百分比坐标到像素坐标的准确转换。\n * 增强 `execute_action` 和 `Browser` 类以支持基于坐标的交互。\n * 优化 Prompt，指导 LLM 输出 VLM 元素 ID 或在 CSS 选择器失败时提供坐标作为备选。\n * 实现更智能的滚动策略以处理长页面和完整内容提取。\n * 持续更新依赖库，解决 Pydantic 警告。\n * 增强错误处理和恢复能力。"
},
{
"path": "super_agents/browser_use/__init__.py",
"content": ""
},
{
"path": "super_agents/browser_use/agent/__init__.py",
"content": "# super_agents/browser_use/agent/__init__.py\n\"\"\"\nBrowser agent module that handles browser automation using LLM guidance.\n\"\"\"\n"
},
{
"path": "super_agents/browser_use/agent/graph.py",
"content": "# super_agents/browser_use/agent/graph.py\nimport logging\nfrom typing import Dict, Any\n\nfrom langchain_core.runnables.base import RunnableSerializable\nfrom langgraph.graph import StateGraph, END\n\nfrom .state import AgentState\nfrom .nodes import AgentNodes\nfrom ..browser.browser import Browser\n\nlogger = logging.getLogger(__name__)\n\nNODE_GET_BROWSER_STATE = \"get_browser_state\"\nNODE_PLAN_ACTION = \"plan_action\"\nNODE_EXECUTE_ACTION = \"execute_action\"\n\n# --- UPDATED Conditional Edge Logic ---\ndef should_end(state: AgentState) -> bool:\n \"\"\"Determines if the graph should end.\"\"\"\n action = state.get(\"parsed_action\", {})\n action_type = action.get(\"type\")\n error_occurred = state.get(\"error\") is not None # Check if execute_action reported an error\n\n # End if the LLM planned action is 'finish' or 'error'\n if action_type == \"finish\":\n logger.info(\"Graph execution: 'finish' action planned. Ending.\")\n return True\n if action_type == \"error\":\n # Log the error message from the action payload\n logger.error(f\"Graph execution: 'error' action planned by LLM: {action.get('message', 'Unknown error')}. Ending.\")\n return True\n\n # End if the execute_action node reported an error in the state\n # Note: Depending on desired behavior, you might want to retry instead of ending on execution errors\n # if error_occurred:\n # logger.error(f\"Graph execution: Error occurred during execution: {state['error']}. Ending.\")\n # return True # Uncomment this line if ANY execution error should terminate the graph\n\n return False # Continue otherwise\n\ndef create_graph_app(browser: Browser, llm: RunnableSerializable):\n \"\"\"\n Creates the LangGraph application using class-based nodes.\n \"\"\"\n agent_nodes = AgentNodes(browser=browser, llm=llm)\n workflow = StateGraph(AgentState)\n\n workflow.add_node(NODE_GET_BROWSER_STATE, agent_nodes.get_browser_state)\n workflow.add_node(NODE_PLAN_ACTION, agent_nodes.plan_action)\n workflow.add_node(NODE_EXECUTE_ACTION, agent_nodes.execute_action)\n\n workflow.set_entry_point(NODE_GET_BROWSER_STATE)\n workflow.add_edge(NODE_GET_BROWSER_STATE, NODE_PLAN_ACTION)\n workflow.add_edge(NODE_PLAN_ACTION, NODE_EXECUTE_ACTION)\n\n # After executing action, decide whether to end or loop back\n workflow.add_conditional_edges(\n NODE_EXECUTE_ACTION,\n # Function to decide the next step based on the state *after* execution\n lambda state: END if should_end(state) else NODE_GET_BROWSER_STATE,\n {\n END: END,\n NODE_GET_BROWSER_STATE: NODE_GET_BROWSER_STATE\n }\n )\n\n logger.info(\"Compiling LangGraph workflow...\")\n app = workflow.compile()\n logger.info(\"LangGraph workflow compiled successfully.\")\n return app"
},
{
"path": "super_agents/browser_use/agent/nodes.py",
"content": "# super_agents/browser_use/agent/nodes.py\nimport asyncio\nimport logging\nfrom typing import Dict, Any, Optional\n\n# --- LangChain Core Import for Type Hint ---\nfrom langchain_core.runnables.base import RunnableSerializable # <--- Import this\n\nfrom .state import AgentState\nfrom .schemas import (\n BaseAction, LLMResponse\n)\nfrom .prompts import create_agent_prompt\n# --- CORRECTED LLM IMPORT ---\n# Import only the necessary functions/classes that actually exist in llm.py\nfrom ..llm import generate_structured_output\n\n# Import the correct Browser from the browser subdirectory\nfrom ..browser.browser import Browser\n\nlogger = logging.getLogger(__name__)\n\n# --- Class to hold nodes and dependencies ---\nclass AgentNodes:\n \"\"\"Encapsulates agent nodes and their dependencies (browser, llm).\"\"\"\n # --- CORRECTED TYPE HINT for llm ---\n def __init__(self, browser: Browser, llm: RunnableSerializable): # <--- Use RunnableSerializable\n if not isinstance(llm, RunnableSerializable):\n logger.warning(f\"LLM instance provided to AgentNodes is not of type RunnableSerializable (actual type: {type(llm)}).\")\n self.browser = browser\n self.llm = llm\n logger.info(\"AgentNodes initialized with browser and llm instances.\")\n\n # --- Node method implementations remain the same ---\n async def get_browser_state(self, state: AgentState) -> Dict[str, Any]:\n \"\"\"Node method to get the current state of the browser page.\"\"\"\n logger.info(\"Node: get_browser_state\")\n try:\n content = await self.browser.get_content()\n return {\"browser_content\": content, \"error\": None}\n except Exception as e:\n logger.error(f\"Error getting browser state: {e}\", exc_info=True)\n return {\"error\": f\"Failed to get browser state: {e}\"}\n\n async def plan_action(self, state: AgentState) -> Dict[str, Any]:\n \"\"\"Node method to decide the next action using the LLM's structured output.\"\"\"\n logger.info(\"Node: plan_action\")\n if state.get(\"error\"):\n logger.warning(f\"Planning action with existing error: {state['error']}\")\n\n prompt = create_agent_prompt(\n task=state[\"task\"],\n current_browser_content=state[\"browser_content\"],\n history=state.get(\"history\", []),\n error_message=state.get(\"error\")\n )\n system_message = \"You are an AI agent controlling a web browser. Respond with the single next action formatted as JSON matching the required schema.\"\n\n try:\n llm_response: Optional[LLMResponse] = await generate_structured_output(\n model=self.llm, # Pass the llm instance\n schema=LLMResponse,\n prompt=prompt,\n system_message=system_message\n )\n\n if llm_response and isinstance(llm_response, LLMResponse):\n parsed_action_model: BaseAction = llm_response.action\n parsed_action_dict = parsed_action_model.dict()\n logger.info(f\"LLM proposed action: {parsed_action_dict.get('type', 'unknown')}\")\n return {\"parsed_action\": parsed_action_dict, \"error\": None}\n else:\n logger.error(\"Failed to get valid structured output from LLM.\")\n error_action_dict = {\"type\": \"error\", \"message\": \"Failed to get valid structured output from LLM.\"}\n return {\"parsed_action\": error_action_dict, \"error\": \"LLM did not return valid structured output.\"}\n\n except Exception as e:\n logger.error(f\"Error during structured action planning: {e}\", exc_info=True)\n error_action_dict = {\"type\": \"error\", \"message\": f\"LLM planning exception: {e}\"}\n return {\"parsed_action\": error_action_dict, \"error\": f\"LLM planning exception: {e}\"}\n\n\n async def execute_action(self, state: AgentState) -> Dict[str, Any]:\n \"\"\"Node method to execute the action dictionary from the state.\"\"\"\n logger.info(\"Node: execute_action\")\n action_dict = state.get(\"parsed_action\")\n history = state.get(\"history\", [])\n\n if not action_dict or not isinstance(action_dict, dict) or \"type\" not in action_dict:\n error_msg = \"No valid action dictionary provided to execute.\"\n logger.error(error_msg)\n return {\"error\": error_msg}\n\n action_type = action_dict.get(\"type\")\n action_repr = f\"Action: {action_type}, Details: { {k:v for k,v in action_dict.items() if k != 'type'} }\"\n logger.info(f\"Executing {action_repr}\")\n\n new_history = history + [action_repr]\n\n try:\n if action_type == \"navigate\":\n await self.browser.navigate_to(action_dict[\"url\"]) # Check if method name/args match Browser class\n elif action_type == \"click\":\n await self.browser.click(action_dict[\"selector\"]) # Check Browser class for click method/args\n elif action_type == \"type\":\n await self.browser.type(action_dict[\"selector\"], action_dict[\"text\"]) # Check Browser class for type method/args\n elif action_type == \"scroll\":\n await self.browser.scroll(action_dict[\"direction\"]) # Check Browser class for scroll method/args\n elif action_type == \"wait\":\n await self.browser.wait(action_dict[\"milliseconds\"]) # Check Browser class for wait method/args\n elif action_type == \"get_content\":\n logger.info(\"Action 'get_content' requested (will be handled by next cycle)\")\n pass\n elif action_type == \"finish\":\n logger.info(f\"Action 'finish' received. Result: {action_dict.get('result')}\")\n pass\n elif action_type == \"error\":\n error_msg = action_dict.get(\"message\", \"LLM signaled an error.\")\n logger.error(f\"Executing 'error' action from LLM: {error_msg}\")\n return {\"error\": error_msg, \"history\": new_history}\n else:\n error_msg = f\"Attempted to execute unknown/unhandled action type: {action_type}\"\n logger.error(error_msg)\n return {\"error\": error_msg, \"history\": new_history}\n\n return {\"error\": None, \"history\": new_history}\n\n except Exception as e:\n logger.error(f\"Error executing action '{action_type}': {e}\", exc_info=True)\n return {\"error\": f\"Failed to execute action '{action_type}': {e}\", \"history\": new_history}"
},
{
"path": "super_agents/browser_use/agent/prompts.py",
"content": "from typing import List\n\ndef create_agent_prompt(\n task: str,\n current_browser_content: str, # This string now potentially contains URL, DOM, AX Tree, and Visual Elements\n history: List[str],\n error_message: str = None\n) -> str:\n \"\"\"\n Generates the prompt to be sent to the LLM based on the current state.\n Includes sections for Simplified DOM, Accessibility Tree, and Visual Elements.\n \"\"\"\n prompt_parts = []\n prompt_parts.append(\"You are an AI agent controlling a web browser to complete a task.\")\n prompt_parts.append(f\"Your current task is: {task}\")\n\n if error_message:\n prompt_parts.append(f\"\\nAn error occurred in the previous step: {error_message}\")\n prompt_parts.append(\"Please analyze the error and the current browser state, then decide the next best action.\")\n\n prompt_parts.append(\"\\n\\n# Current Browser Perception:\")\n # The browser_content string now contains multiple sections, as generated by get_content\n prompt_parts.append(current_browser_content)\n\n if history:\n prompt_parts.append(\"\\n\\n# History of Previous Actions:\")\n for i, item in enumerate(history[-5:], 1):\n prompt_parts.append(f\"{i}. {item}\")\n\n # --- Instructions with guidance on using all perception data ---\n instructions = \"\"\"\n\n# Instructions:\nAnalyze the **Current Browser Perception** section above, which includes:\n1. **Page URL:** The current web address.\n2. **Simplified DOM:** A structural view of the page with interactive elements marked with `x-pw-id` attributes.\n3. **Accessibility Tree:** Semantic information about elements (roles, names).\n4. **Visual Elements:** Elements detected visually via Computer Vision (CV), including their bounding boxes `[L:left, T:top, R:right, B:bottom]` and IDs (e.g., `cv-0`, `cv-1`).\n\nBased on the task and ALL available perception information, decide the single next action to take.\nYour response MUST be a JSON object with a single top-level key named \"action\".\nThe value of the \"action\" key MUST be an object matching one of the following action schemas:\n\n- Navigate: {{\"type\": \"navigate\", \"url\": \"\"}}\n- Click: {{\"type\": \"click\", \"selector\": \"\", \"description\": \"\"}}\n- Type: {{\"type\": \"type\", \"selector\": \"\", \"text\": \"\", \"description\": \"\"}}\n- Scroll: {{\"type\": \"scroll\", \"direction\": \"\"}}\n- Finish: {{\"type\": \"finish\", \"result\": \"\"}}\n- Error: {{\"type\": \"error\", \"message\": \"\"}} (Use if you detect an unrecoverable error or loop)\n- GetContent: {{\"type\": \"get_content\", \"description\": \"\"}}\n\n**Important Task Handling Guidance:**\n1. **Identify elements** using the DOM, AX Tree (if available), and Visual Elements. Use robust selectors as previously guided.\n2. **If the task requires reading or extracting content that might extend beyond the current view (e.g., '摘录全文', 'find all items', 'read the article'), and you haven't finished scrolling, your next action should likely be to SCROLL DOWN.** Use: `{{\"action\": {{\"type\": \"scroll\", \"direction\": \"down\"}}}}`\n3. Only use `get_content` if you believe scrolling will not help or if you need to re-analyze after a non-scroll action.\n4. Once you believe you have scrolled enough and have all necessary information visible in the content provided, proceed with the extraction or final action.\n5. If the task is complete, use the 'finish' action.\n\nExample Response:\n```json\n{{\n \"action\": {{\n \"type\": \"click\",\n \"selector\": \"a[x-pw-id='pw-16']:has-text('new')\",\n \"description\": \"Click the 'new' link, corresponds to visual element cv-3\"\n }}\n}}\n{{\n \"action\": {{\n \"type\": \"scroll\",\n \"direction\": \"down\"\n }}\n}}\n```\n\nProvide ONLY the JSON object containing the 'action' key in a ```json ... ``` block.\nThink step-by-step. Correlate information from the DOM, AX Tree, and Visual Elements if possible. Choose the most precise and stable selector.\nIf the task is complete, use the 'finish' action.\n\"\"\"\n prompt_parts.append(instructions)\n # --- End Instructions ---\n\n final_prompt = \"\\n\".join(prompt_parts)\n return final_prompt"
},
{
"path": "super_agents/browser_use/agent/schemas.py",
"content": "# super_agents/browser_use/agent/schemas.py\nfrom typing import Literal, Optional, Union, List, Dict, Any, Type\n# Use Pydantic V2+ if installed, otherwise V1 syntax\ntry:\n from pydantic.v1 import BaseModel, Field\nexcept ImportError:\n from pydantic import BaseModel, Field # Fallback to V2\n\n# --- Action Type ---\nActionTypeLiteral = Literal[\n \"navigate\",\n \"click\",\n \"type\",\n \"scroll\",\n \"wait\",\n \"get_content\",\n \"finish\",\n \"error\"\n]\n\n# --- Pydantic Schemas for Actions ---\n# Using Pydantic allows for better validation and compatibility\n# with LangChain's structured output features.\n\nclass BaseAction(BaseModel):\n \"\"\"Base schema for all actions, containing the type.\"\"\"\n type: ActionTypeLiteral = Field(..., description=\"The type of action to perform.\")\n\nclass NavigateAction(BaseAction):\n type: Literal[\"navigate\"] = \"navigate\"\n url: str = Field(..., description=\"The URL to navigate to.\")\n\nclass ClickAction(BaseAction):\n type: Literal[\"click\"] = \"click\"\n selector: str = Field(..., description=\"CSS selector for the element to click.\")\n description: Optional[str] = Field(None, description=\"Optional description of the element being clicked.\")\n\nclass TypeAction(BaseAction):\n type: Literal[\"type\"] = \"type\"\n selector: str = Field(..., description=\"CSS selector for the input field.\")\n text: str = Field(..., description=\"The text to type into the field.\")\n description: Optional[str] = Field(None, description=\"Optional description of the element being typed into.\")\n\nclass ScrollAction(BaseAction):\n type: Literal[\"scroll\"] = \"scroll\"\n direction: Literal[\"up\", \"down\", \"left\", \"right\"] = Field(..., description=\"The direction to scroll the page.\")\n # selector: Optional[str] = Field(None, description=\"Optional CSS selector of element to scroll within.\") # Add if needed\n\nclass WaitAction(BaseAction):\n type: Literal[\"wait\"] = \"wait\"\n milliseconds: int = Field(..., description=\"Duration to wait in milliseconds.\")\n\nclass GetContentAction(BaseAction):\n type: Literal[\"get_content\"] = \"get_content\"\n # No extra fields needed, just signifies intent to refresh state\n description: Optional[str] = Field(\"Requesting updated browser content\", description=\"Reason for requesting content.\")\n\nclass FinishAction(BaseAction):\n type: Literal[\"finish\"] = \"finish\"\n result: str = Field(..., description=\"The final answer or summary of the completed task.\")\n\nclass ErrorAction(BaseAction):\n type: Literal[\"error\"] = \"error\"\n message: str = Field(..., description=\"Description of the error encountered or signaled by the LLM.\")\n\n# --- Union for Parsing ---\n# LangChain's with_structured_output often works best when targeting a single Pydantic model\n# that uses discriminated unions (if available in your Pydantic version) or by prompting\n# the LLM clearly to only output ONE type of action JSON matching the base structure.\n# For simplicity here, we define the *expected output structure* the LLM should generate.\n# The parsing function might need refinement based on how the LLM structures the output.\n\n# Define the overall structure the LLM should output, which includes one of the actions.\n# This structure helps `with_structured_output`.\nclass LLMResponse(BaseModel):\n action: Union[\n NavigateAction,\n ClickAction,\n TypeAction,\n ScrollAction,\n WaitAction,\n GetContentAction,\n FinishAction,\n ErrorAction\n ] = Field(..., description=\"The specific action determined by the LLM.\")\n\n# --- Parsing Function (Placeholder/Example) ---\n# The `generate_structured_output` function in llm.py now handles the parsing\n# directly into the Pydantic schema (LLMResponse).\n# So, we might not need a separate manual parsing function here if using that.\n\n# If you need manual parsing from raw text (less reliable):\n# def parse_llm_response_manual(response: str) -> Optional[BaseAction]:\n# # ... (complex logic using regex or JSON parsing as in previous example)\n# # This would return one of the action models (NavigateAction, ClickAction, etc.)\n# pass\n"
},
{
"path": "super_agents/browser_use/agent/state.py",
"content": "# super_agents/browser_use/agent/state.py\nfrom typing import Dict, List, Optional, Any, TypedDict\n\n# Define the state structure using TypedDict for type hinting\nclass AgentState(TypedDict, total=False):\n \"\"\"\n TypedDict representing the state of the browser agent during execution.\n \n Attributes:\n task: The user task description\n browser_content: The current HTML content of the browser\n parsed_action: The last action parsed from LLM response\n history: List of previous actions taken\n error: Any error message from the last operation\n \"\"\"\n task: str\n browser_content: str\n parsed_action: Dict[str, Any]\n history: List[str]\n error: Optional[str]\n"
},
{
"path": "super_agents/browser_use/agent/tools.py",
"content": ""
},
{
"path": "super_agents/browser_use/agent.py",
"content": "# super_agents/browser_use/agent.py\n\"\"\"\nAgent API for browser-based task execution.\nProvides a simplified interface similar to the original implementation.\n\"\"\"\n\nimport asyncio\nimport logging\nfrom typing import Any, Dict, Optional\n\nfrom .agent.graph import create_graph_app\nfrom .agent.state import AgentState\nfrom .browser.browser import Browser\nfrom .browser.config import BrowserConfig\nfrom .llm import initialize_llms\n\nlogger = logging.getLogger(__name__)\n\nclass Agent:\n \"\"\"\n Agent class that provides a simple interface for browser automation with LLM.\n \n This implementation is similar to the original API but uses the current\n browser automation stack with LangGraph.\n \"\"\"\n \n def __init__(\n self, \n llm=None,\n browser_config: Optional[BrowserConfig] = None,\n max_steps: int = 50\n ):\n \"\"\"\n Initialize the Agent with optional LLM and browser configuration.\n \n Args:\n llm: LLM instance to use (if None, will initialize from environment)\n browser_config: Browser configuration options\n max_steps: Maximum number of steps the agent can take\n \"\"\"\n self.browser_config = browser_config or BrowserConfig()\n self.llm = llm\n self.max_steps = max_steps\n self.browser = None\n self._app = None\n \n async def _initialize(self):\n \"\"\"Initialize the browser and LLM if not already initialized.\"\"\"\n # Initialize LLM if not provided\n if self.llm is None:\n logger.info(\"Initializing LLM from environment variables\")\n self.llm, _ = initialize_llms()\n \n if self.llm is None:\n raise ValueError(\"Failed to initialize LLM. Check API keys and .env settings.\")\n \n # Initialize browser\n self.browser = Browser(config=self.browser_config)\n await self.browser.initialize()\n \n # Initialize LangGraph app\n self._app = create_graph_app(browser=self.browser, llm=self.llm)\n \n async def run(self, prompt: str) -> Dict[str, Any]:\n \"\"\"\n Run the agent with the given prompt/task.\n \n Args:\n prompt: The task description or prompt for the agent\n \n Returns:\n Dictionary containing the execution result\n \"\"\"\n # Ensure initialization\n if self.browser is None or self._app is None:\n await self._initialize()\n \n # Define the initial state\n initial_state = AgentState(\n task=prompt,\n browser_content=\"\",\n parsed_action={},\n history=[],\n error=None\n )\n \n # Run the graph\n logger.info(f\"Starting agent execution for task: {prompt}\")\n try:\n final_state = await self._app.ainvoke(\n initial_state, \n config={\"recursion_limit\": self.max_steps}\n )\n \n # Process result\n if final_state.get(\"error\"):\n logger.error(f\"Agent finished with error: {final_state['error']}\")\n return {\"result\": f\"Error: {final_state['error']}\", \"success\": False}\n elif final_state.get(\"parsed_action\", {}).get(\"type\") == \"finish\":\n result = final_state[\"parsed_action\"].get(\"result\", \"Task finished, but no result extracted.\")\n logger.info(f\"Agent finished successfully. Result: {result}\")\n return {\"result\": result, \"success\": True}\n else:\n logger.warning(\"Agent finished without a 'finish' action or error.\")\n return {\n \"result\": \"Agent stopped without producing a final answer.\", \n \"success\": False,\n \"state\": final_state\n }\n \n except Exception as e:\n logger.error(f\"Agent execution failed: {e}\", exc_info=True)\n return {\"result\": f\"Error during execution: {str(e)}\", \"success\": False}\n finally:\n # Clean up resources\n if self.browser:\n await self.browser.close()\n self.browser = None\n self._app = None\n \n def __del__(self):\n \"\"\"Ensure resources are cleaned up.\"\"\"\n if self.browser:\n asyncio.create_task(self.browser.close())\n\n\n# Provider classes for compatibility with original API\nclass OpenAIProvider:\n \"\"\"OpenAI provider compatible with the interface\"\"\"\n \n def __init__(self, model=\"gpt-4o-mini\", api_key=None, temperature=0.1):\n \"\"\"\n Initialize OpenAI provider.\n \n Args:\n model: Model name to use\n api_key: OpenAI API key (if None, will use from environment)\n temperature: Temperature for generation\n \"\"\"\n self.model = model\n self.api_key = api_key\n self.temperature = temperature\n \n # These parameters will be used by initialize_llms() internally\n import os\n if api_key:\n os.environ[\"OPENAI_API_KEY\"] = api_key\n os.environ[\"LLM_PROVIDER\"] = \"openai\"\n os.environ[\"LLM_MODEL_NAME\"] = model\n os.environ[\"LLM_TEMPERATURE\"] = str(temperature)\n\n\nclass AnthropicProvider:\n \"\"\"Anthropic provider compatible with the interface\"\"\"\n \n def __init__(self, model=\"claude-3-opus-20240229\", api_key=None, temperature=0.1, \n enable_thinking=False, thinking_token_budget=None):\n \"\"\"\n Initialize Anthropic provider.\n \n Args:\n model: Model name to use\n api_key: Anthropic API key (if None, will use from environment)\n temperature: Temperature for generation\n enable_thinking: Enable thinking step (not fully supported in current implementation)\n thinking_token_budget: Tokens for thinking (not fully supported)\n \"\"\"\n self.model = model\n self.api_key = api_key\n self.temperature = temperature\n self.enable_thinking = enable_thinking\n self.thinking_token_budget = thinking_token_budget\n \n # These parameters will be used by initialize_llms() internally\n import os\n if api_key:\n os.environ[\"ANTHROPIC_API_KEY\"] = api_key\n os.environ[\"LLM_PROVIDER\"] = \"anthropic\" \n os.environ[\"LLM_MODEL_NAME\"] = model\n os.environ[\"LLM_TEMPERATURE\"] = str(temperature)\n\n\n# Add convenience imports to __init__.py\n# This will allow: from super_agents.browser_use import Agent, OpenAIProvider, BrowserConfig\n"
},
{
"path": "super_agents/browser_use/browser/browser.py",
"content": "# super_agents/browser_use/browser/browser.py\n\"\"\"\nStreamlined Playwright browser implementation with integrated perception capabilities.\nIncludes DOM/AX Tree/Visual analysis and basic interaction methods.\n\"\"\"\n\nimport asyncio\nimport json\nimport logging\nimport functools \nimport base64\nimport os\nfrom dataclasses import dataclass, field\n# from importlib import resources # Not used\nfrom typing import Any, Optional, TypedDict, List, Dict # Added List, Dict\n\n# --- Local Imports (Ensure these files exist in the same directory) ---\ntry:\n from .observe_helper import observe\nexcept ImportError:\n def observe(name, ignore_input=False, ignore_output=False):\n def decorator(func): return func\n return decorator\n logging.basicConfig(level=logging.WARNING) # Setup basic logging if needed\n logger_observe = logging.getLogger(__name__)\n logger_observe.warning(\"observe_helper not found, using dummy decorator.\")\n\ntry:\n from .detector import Detector\n from .models import (\n BrowserError,\n BrowserState,\n InteractiveElementsData,\n TabInfo,\n InteractiveElement,\n )\n from .utils import (\n combine_and_filter_elements,\n put_highlight_elements_on_screenshot,\n )\nexcept ImportError as e:\n logging.basicConfig(level=logging.ERROR)\n logger_import = logging.getLogger(__name__)\n logger_import.error(f\"Failed to import local browser dependencies (detector, models, utils): {e}. Browser class may not function correctly.\", exc_info=True)\n # Define dummy classes to allow file loading, but functionality will be broken\n class Detector: enabled=False\n class BrowserError(Exception): pass\n class BrowserState: pass\n class InteractiveElementsData: elements=[]; viewport={}\n class TabInfo: pass\n class InteractiveElement: pass\n def combine_and_filter_elements(a, b): return []\n def put_highlight_elements_on_screenshot(a, b): return None\n# --- End Local Imports ---\n\n# --- Playwright Imports ---\nfrom playwright.async_api import (\n Browser as PlaywrightBrowser,\n BrowserContext as PlaywrightBrowserContext,\n Page,\n Playwright,\n StorageState,\n async_playwright,\n Error as PlaywrightError\n)\n# --- Tenacity Import ---\nfrom tenacity import (\n retry,\n retry_if_exception_type,\n stop_after_attempt,\n wait_exponential,\n)\n\nlogger = logging.getLogger(__name__)\n# Ensure basic logging is configured if not done elsewhere\nif not logger.hasHandlers():\n logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')\n\n\n# --- Load JavaScript Files ---\nINTERACTIVE_ELEMENTS_JS_CODE = \"\"\nSIMPLIFY_PAGE_SCRIPT = \"\"\ntry:\n current_dir = os.path.dirname(os.path.abspath(__file__))\n # JS for DOM-based interactive elements used in get_interactive_elements_data\n js_file_path_interactive = os.path.join(current_dir, 'findVisibleInteractiveElements.js')\n with open(js_file_path_interactive, 'r', encoding='utf-8') as js_file:\n INTERACTIVE_ELEMENTS_JS_CODE = js_file.read()\n\n # JS for DOM simplification used in get_content\n # (Re-paste the script here for completeness)\n SIMPLIFY_PAGE_SCRIPT = \"\"\"\n (() => {\n const MAX_ELEMENTS = 250; const MAX_TEXT_LENGTH = 200;\n const INTERACTIVE_TAGS = ['a', 'button', 'input', 'textarea', 'select', 'option', 'details', 'summary', 'label'];\n const EXCLUDED_TAGS = ['script', 'style', 'noscript', 'svg', 'link', 'meta', 'head', 'embed', 'object', 'path', 'canvas', 'iframe', 'video', 'audio'];\n let elementCount = 0; let uniqueIdCounter = 0;\n function isVisible(el) { if (!el || !el.checkVisibility) return false; return el.checkVisibility({checkOpacity: true, checkVisibilityCSS: true}); }\n function truncateText(text, maxLength = MAX_TEXT_LENGTH) { if (typeof text !== 'string') return text; return text.length > maxLength ? text.substring(0, maxLength) + '...' : text; }\n function getElementData(el) {\n const data = { tag: el.tagName.toLowerCase(), attributes: {}, text: '', children: [], pw_id: `pw-${uniqueIdCounter++}` };\n try { if (document.body.contains(el)) el.setAttribute('x-pw-id', data.pw_id); } catch(e){}\n const attrsToKeep = ['id', 'class', 'role', 'aria-label', 'aria-labelledby', 'aria-describedby', 'aria-hidden', 'aria-invalid', 'aria-required', 'placeholder', 'title', 'alt', 'for', 'name', 'type', 'href', 'value', 'selected', 'checked', 'disabled', 'readonly', 'open'];\n for (const attr of attrsToKeep) {\n if (el.hasAttribute(attr)) {\n let value = el.getAttribute(attr);\n if (attr === 'class' && value) value = value.split(' ').filter(c => c && c.length > 1 && c.length < 30 && !/^[0-9]+$/.test(c)).slice(0, 5).join(' ');\n if (value !== null && value !== '') data.attributes[attr] = truncateText(String(value), 80);\n }\n }\n if (['button', 'a', 'label', 'summary'].includes(data.tag) && !data.attributes['aria-label'] && el.textContent) data.attributes['aria-label'] = truncateText(el.textContent.trim(), 80);\n try {\n if (el.tagName.toLowerCase() === 'input' && !data.attributes.value && el.value) data.attributes.value = truncateText(el.value);\n else if (el.tagName.toLowerCase() === 'textarea' && !data.attributes.value && el.value) data.attributes.value = truncateText(el.value);\n else if (el.tagName.toLowerCase() === 'select' && el.options && el.selectedIndex !== -1 && !data.attributes.value) data.attributes.value = truncateText(el.options[el.selectedIndex].text);\n } catch (e) {}\n try {\n const directText = Array.from(el.childNodes).filter(node => node.nodeType === Node.TEXT_NODE && node.textContent.trim().length > 0).map(node => node.textContent.trim()).join(' ').replace(/\\s+/g, ' ');\n if (directText) data.text = truncateText(directText);\n } catch (e) {}\n return data;\n }\n function simplifyNode(node) {\n if (elementCount >= MAX_ELEMENTS) return null;\n if (node.nodeType !== Node.ELEMENT_NODE || EXCLUDED_TAGS.includes(node.tagName.toLowerCase())) { if(node.nodeType === Node.TEXT_NODE && node.textContent.trim().length === 0) return null; return null; }\n elementCount++; const elementData = getElementData(node);\n if (node.hasChildNodes()) {\n Array.from(node.childNodes).forEach(child => {\n if (INTERACTIVE_TAGS.includes(node.tagName.toLowerCase()) && child.nodeType === Node.ELEMENT_NODE) return;\n const simplifiedChild = simplifyNode(child); if (simplifiedChild) elementData.children.push(simplifiedChild);\n });\n }\n const isInteractive = INTERACTIVE_TAGS.includes(elementData.tag); const hasMeaningfulAttrs = Object.keys(elementData.attributes).some(k => k !== 'x-pw-id');\n if (!isInteractive && !hasMeaningfulAttrs && elementData.children.length === 0 && !elementData.text) { try { if (document.body.contains(node)) node.removeAttribute('x-pw-id'); } catch(e){} return null; }\n return elementData;\n }\n if (!document.body) return \" element not found.\"; const simplifiedBody = simplifyNode(document.body);\n function convertToPseudoHTML(node) {\n if (!node) return ''; let attrs = `x-pw-id=\"${node.pw_id}\"`;\n for (const [key, value] of Object.entries(node.attributes)) attrs += ` ${key}=\"${String(value).replace(/\"/g, '"')}\"`;\n let childrenHTML = node.children.map(convertToPseudoHTML).join('');\n let textContent = node.text ? String(node.text).replace(//g, '>') : '';\n if (['input', 'img', 'br', 'hr'].includes(node.tag)) return `<${node.tag} ${attrs} />`;\n else return `<${node.tag} ${attrs}>${textContent}${childrenHTML}`;\n }\n return convertToPseudoHTML(simplifiedBody);\n })()\n \"\"\"\nexcept FileNotFoundError:\n logger.error(f\"JavaScript file 'findVisibleInteractiveElements.js' not found in {current_dir}. Interactive element detection (JS based) will fail.\")\n INTERACTIVE_ELEMENTS_JS_CODE = \"() => ({ viewport: { width: window.innerWidth, height: window.innerHeight }, elements: [] });\" # Provide fallback\nexcept Exception as e:\n logger.error(f\"Error loading JavaScript file(s): {e}\", exc_info=True)\n INTERACTIVE_ELEMENTS_JS_CODE = \"() => ({ viewport: { width: window.innerWidth, height: window.innerHeight }, elements: [] });\"\n SIMPLIFY_PAGE_SCRIPT = \"() => 'Error loading simplification script.';\"\n\n\n# --- TypedDict for Viewport Size ---\nclass ViewportSize(TypedDict):\n width: int\n height: int\n\n# --- BrowserConfig Dataclass (Corrected: No CV Endpoints) ---\n@dataclass\nclass BrowserConfig:\n \"\"\"\n Configuration for the Browser.\n \"\"\"\n cdp_url: Optional[str] = None\n viewport_size: ViewportSize = field(default_factory=lambda: {\"width\": 1200, \"height\": 900})\n storage_state: Optional[StorageState] = None\n # CV/Sheets Endpoints Removed\n\n# --- Main Browser Class ---\nclass Browser:\n \"\"\"\n Unified Browser responsible for interacting with the browser via Playwright.\n Includes methods for navigation, simple actions, perception (DOM, AX Tree, optional VLM),\n and state management. Initializes its own VLM detector based on environment variables.\n \"\"\"\n def __init__(self, config: BrowserConfig = BrowserConfig(), close_context: bool = True):\n \"\"\"\n Initializes the Browser instance.\n \"\"\"\n logger.debug('Initializing browser')\n self.config = config\n self.close_context = close_context\n # Playwright attributes\n self.playwright: Optional[Playwright] = None\n self.playwright_browser: Optional[PlaywrightBrowser] = None\n self.context: Optional[PlaywrightBrowserContext] = None\n # Page and state management\n self.current_page: Optional[Page] = None\n self._state: Optional[BrowserState] = None # This holds the rich state from update_state\n self._cdp_session = None\n # Initialize Detector internally\n try:\n self.detector: Optional[Detector] = Detector()\n if not self.detector.enabled:\n self.detector = None\n logger.warning(\"Detector initialized but disabled due to missing config/errors.\")\n else:\n logger.info(\"Detector initialized successfully.\")\n except NameError:\n logger.error(\"Detector class not found (likely due to import errors). Vision disabled.\")\n self.detector = None\n except Exception as e:\n logger.error(f\"Unexpected error initializing Detector: {e}\", exc_info=True)\n self.detector = None\n # REMOVED self._init_state() call as method doesn't exist / state init is implicit\n\n # --- Context Management Methods ---\n async def __aenter__(self):\n await self.initialize()\n return self\n\n async def __aexit__(self, exc_type, exc_val, exc_tb):\n if self.close_context:\n await self.close()\n\n # --- Public Initialization and Closing ---\n async def initialize(self):\n \"\"\"Initializes browser, context, page if not already done.\"\"\"\n if self.current_page and self.context and self.playwright_browser and self.playwright:\n logger.debug(\"Browser already initialized.\")\n return self\n logger.info(\"Initializing browser instance via initialize()\") # Changed level\n await self._init_browser()\n return self\n\n async def close(self):\n \"\"\"Closes the browser and cleans up Playwright resources.\"\"\"\n if not self.playwright: return\n logger.info('Closing browser...')\n try:\n self._cdp_session = None\n if self.context:\n try: await self.context.close()\n except Exception as e: logger.warning(f'Failed to close context: {e}')\n if self.playwright_browser and not self.config.cdp_url:\n try: await self.playwright_browser.close()\n except Exception as e: logger.warning(f'Failed to close browser: {e}')\n if self.playwright:\n try: await self.playwright.stop()\n except Exception as e: logger.warning(f'Failed to stop Playwright: {e}')\n except Exception as e:\n logger.error(f'Error during browser cleanup: {e}', exc_info=True)\n finally: # Ensure attributes are cleared\n self.context = None; self.current_page = None; self._state = None\n self.playwright_browser = None; self.playwright = None; self._cdp_session = None\n logger.info(\"Browser closed.\")\n\n # --- Internal Initialization Helper ---\n async def _init_browser(self):\n \"\"\"Internal method to initialize Playwright components.\"\"\"\n if self.current_page and self.context: return # Avoid re-init if basics exist\n logger.debug('Running internal browser context initialization _init_browser()')\n try:\n if self.playwright is None: self.playwright = await async_playwright().start()\n if self.playwright_browser is None:\n if self.config.cdp_url:\n logger.info(f'Connecting to remote browser via CDP {self.config.cdp_url}')\n self.playwright_browser = await self.playwright.chromium.connect_over_cdp(self.config.cdp_url, timeout=5000)\n else:\n logger.info(f'Launching new browser instance (headless=False assumed)')\n # Note: Headless mode might need to be configurable via BrowserConfig again if needed\n self.playwright_browser = await self.playwright.chromium.launch(\n headless=False,\n args=[ # Common args for stability/anti-detection\n '--no-sandbox', '--disable-setuid-sandbox', '--disable-infobars',\n '--disable-blink-features=AutomationControlled',\n '--disable-dev-shm-usage', '--disable-gpu', '--window-size=1200,900', # Use configured size later\n # '--disable-web-security', # Use with caution\n # '--disable-site-isolation-trials',\n # '--disable-features=IsolateOrigins,site-per-process',\n ]\n )\n if self.context is None:\n existing_contexts = self.playwright_browser.contexts\n if existing_contexts and not self.config.cdp_url: # Reuse only if we launched it? Be careful.\n self.context = existing_contexts[0]\n logger.info(\"Reusing existing browser context.\")\n else:\n logger.info(\"Creating new browser context.\")\n self.context = await self.playwright_browser.new_context(\n viewport=self.config.viewport_size,\n user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36',\n java_script_enabled=True, bypass_csp=True, ignore_https_errors=True,\n storage_state=self.config.storage_state if self.config.storage_state else None\n )\n await self._apply_anti_detection_scripts() # Apply only to new contexts\n self.context.on('page', self._on_page_change) # Attach listener\n\n if self.current_page is None:\n if len(self.context.pages) > 0:\n self.current_page = self.context.pages[-1] # Default to last open page\n logger.info(f\"Using existing page: {self.current_page.url}\")\n else:\n self.current_page = await self.context.new_page()\n logger.info(\"Created new page.\")\n # Ensure viewport is applied regardless\n try: await self.current_page.set_viewport_size(self.config.viewport_size)\n except Exception as vp_err: logger.warning(f\"Failed to set viewport: {vp_err}\")\n\n if not self.current_page: raise BrowserError(\"Failed to get or create a page.\")\n await self.get_cdp_session() # Initialize CDP session for current page\n\n except PlaywrightError as pe:\n logger.error(f\"Playwright Error during browser init: {pe}\", exc_info=True)\n await self.close(); raise BrowserError(f\"Playwright initialization failed: {pe}\") from pe\n except Exception as e:\n logger.error(f\"Unexpected error during browser init: {e}\", exc_info=True)\n await self.close(); raise BrowserError(f\"Unexpected browser initialization failed: {e}\") from e\n\n # --- Method Implementations (Ensure ALL referenced methods are defined) ---\n\n async def _apply_anti_detection_scripts(self):\n \"\"\"Apply scripts to avoid detection as automation\"\"\"\n if self.context is None: return # Should not happen if called from _init_browser correctly\n try:\n await self.context.add_init_script(\n \"\"\"\n Object.defineProperty(navigator, 'webdriver', { get: () => undefined });\n Object.defineProperty(navigator, 'languages', { get: () => ['en-US', 'en'] });\n Object.defineProperty(navigator, 'plugins', { get: () => [] }); // Empty is safer\n // ... other scripts from previous version ...\n const originalQuery = window.navigator.permissions.query;\n window.navigator.permissions.query = (parameters) => (\n parameters.name === 'notifications' ?\n Promise.resolve({ state: Notification.permission }) :\n originalQuery(parameters)\n );\n \"\"\"\n )\n logger.debug(\"Applied anti-detection init script.\")\n except Exception as e:\n logger.error(f\"Failed to add anti-detection init script: {e}\", exc_info=True)\n\n async def _on_page_change(self, page: Page):\n \"\"\"Handle page creation/popup events.\"\"\"\n # Don't automatically switch current page, just log\n logger.info(f'Page event detected. New/Popup URL: {page.url}')\n self._cdp_session = None # Invalidate CDP session as context changed\n\n async def get_current_page(self) -> Page:\n \"\"\"Get the current page, ensuring browser is initialized.\"\"\"\n if self.current_page is None or self.current_page.is_closed():\n logger.warning(\"Current page is None or closed, re-initializing.\")\n await self._init_browser()\n if self.current_page is None: raise BrowserError(\"Unable to get a valid page.\")\n return self.current_page\n\n # Inside Browser class in browser.py\n async def get_cdp_session(self):\n \"\"\"Get or create a CDP session for the *current* page.\"\"\"\n page = await self.get_current_page()\n session_invalid = True # Assume invalid\n if self._cdp_session:\n # More robust check: try a simple CDP command to see if session is active\n try:\n # Example: Get cookies via CDP (relatively harmless check)\n await self._cdp_session.send(\"Network.getAllCookies\")\n # Check if session page matches current page (using internal attr is risky)\n if hasattr(self._cdp_session, '_client') and hasattr(self._cdp_session._client, '_page') and self._cdp_session._client._page == page:\n session_invalid = False # Session seems alive and for the correct page\n else:\n logger.debug(\"CDP session page mismatch or internals unclear, recreating.\")\n except Exception as session_check_err:\n logger.debug(f\"Existing CDP session check failed ({session_check_err}), recreating.\")\n session_invalid = True\n\n if session_invalid:\n try:\n if self.context is None: await self._init_browser()\n logger.debug(f\"Attempting to create new CDP session for page: {page.url}\")\n self._cdp_session = await self.context.new_cdp_session(page)\n logger.debug(f\"Created new CDP session successfully.\")\n except Exception as e:\n logger.error(f\"Failed to create CDP session: {e}\", exc_info=True)\n self._cdp_session = None\n raise BrowserError(f\"Failed to create CDP session: {e}\") from e\n return self._cdp_session\n\n @observe(name='browser.fast_screenshot', ignore_output=True)\n async def fast_screenshot(self) -> str:\n \"\"\"Returns a base64 encoded screenshot using CDP.\"\"\"\n cdp_session = await self.get_cdp_session()\n try:\n screenshot_data = await cdp_session.send(\"Page.captureScreenshot\", {\"format\": \"png\", \"fromSurface\": False, \"captureBeyondViewport\": False})\n return screenshot_data[\"data\"]\n except Exception as e:\n logger.error(f\"Failed to capture screenshot via CDP: {e}\")\n # Fallback to playwright's screenshot? Or raise error?\n page = await self.get_current_page()\n try:\n logger.warning(\"CDP screenshot failed, falling back to Playwright screenshot.\")\n buffer = await page.screenshot()\n return base64.b64encode(buffer).decode()\n except Exception as pw_e:\n logger.error(f\"Fallback Playwright screenshot also failed: {pw_e}\")\n raise BrowserError(f\"Failed to take screenshot: {e}\") from e\n\n # --- Simple Action Methods ---\n @observe(name='browser.navigate_to')\n async def navigate_to(self, url: str):\n page = await self.get_current_page()\n logger.info(f\"Navigating to: {url}\")\n try:\n await page.goto(url, wait_until='domcontentloaded', timeout=60000)\n logger.info(f\"Navigation successful. Current URL: {page.url}\")\n except PlaywrightError as e: raise BrowserError(f\"Navigation failed: {e}\") from e\n except Exception as e: raise BrowserError(f\"Navigation failed unexpectedly: {e}\") from e\n\n @observe(name='browser.click')\n async def click(self, selector: str):\n page = await self.get_current_page()\n logger.info(f\"Attempting to click element: '{selector}'\")\n try:\n element = page.locator(selector).first\n await element.wait_for(state=\"visible\", timeout=15000)\n await element.scroll_into_view_if_needed(timeout=10000)\n await element.click(timeout=15000, delay=50)\n logger.info(f\"Successfully clicked element: '{selector}'\")\n except PlaywrightError as e: raise BrowserError(f\"Click action failed: {e}\") from e\n except Exception as e: raise BrowserError(f\"Click action failed unexpectedly: {e}\") from e\n\n @observe(name='browser.type')\n async def type(self, selector: str, text: str):\n page = await self.get_current_page()\n log_text = '***' if 'password' in selector.lower() else text\n logger.info(f\"Attempting to type into element: '{selector}', Text: '{log_text}'\")\n try:\n element = page.locator(selector).first\n await element.wait_for(state=\"visible\", timeout=15000)\n await element.scroll_into_view_if_needed(timeout=10000)\n await element.fill(text, timeout=15000)\n logger.info(f\"Successfully typed into element: '{selector}'\")\n except PlaywrightError as e: raise BrowserError(f\"Type action failed: {e}\") from e\n except Exception as e: raise BrowserError(f\"Type action failed unexpectedly: {e}\") from e\n\n @observe(name='browser.scroll')\n async def scroll(self, direction: str):\n page = await self.get_current_page()\n logger.info(f\"Scrolling page {direction}\")\n try:\n if direction == \"down\": await page.evaluate(\"window.scrollBy(0, window.innerHeight)\")\n elif direction == \"up\": await page.evaluate(\"window.scrollBy(0, -window.innerHeight)\")\n elif direction == \"left\": await page.evaluate(\"window.scrollBy(-window.innerWidth, 0)\")\n elif direction == \"right\": await page.evaluate(\"window.scrollBy(window.innerWidth, 0)\")\n else: logger.warning(f\"Unknown scroll direction: {direction}\"); return\n await asyncio.sleep(0.3)\n logger.info(f\"Scrolled page {direction}\")\n except PlaywrightError as e: raise BrowserError(f\"Scroll action failed: {e}\") from e\n except Exception as e: raise BrowserError(f\"Scroll action failed unexpectedly: {e}\") from e\n\n async def wait(self, milliseconds: int):\n logger.info(f\"Waiting for {milliseconds} ms\")\n if milliseconds <= 0: return\n await asyncio.sleep(milliseconds / 1000.0)\n logger.info(\"Wait finished\")\n\n # --- Perception & State Methods ---\n async def get_content(self, max_length: int = 120000) -> str:\n \"\"\"Gets comprehensive text representation: URL, DOM, AX Tree, VLM Elements.\"\"\"\n page = await self.get_current_page()\n logger.info(\"Getting comprehensive page content with vision...\")\n combined_content = \"\"\n error_messages = []\n current_url = \"Unknown\"\n screenshot_b64 = None\n try:\n current_url = page.url\n combined_content += f\"# Page URL:\\n{current_url}\\n\\n\"\n try:\n screenshot_b64 = await self.fast_screenshot()\n logger.debug(f\"Screenshot captured (size: {len(screenshot_b64) if screenshot_b64 else 0})\")\n except Exception as ss_err: error_messages.append(f\"Screenshot Error: {ss_err}\"); logger.error(\"Screenshot error\", exc_info=False); combined_content += \"# Screenshot Error\\n\"\n try:\n if SIMPLIFY_PAGE_SCRIPT:\n simplified_dom = await page.evaluate(SIMPLIFY_PAGE_SCRIPT)\n if simplified_dom: combined_content += f\"# Simplified DOM:\\n```html\\n{simplified_dom}\\n```\\n\\n\"; logger.debug(f\"DOM length: {len(simplified_dom)}\")\n else: combined_content += \"# Simplified DOM:\\n(Empty)\\n\\n\"; logger.warning(\"JS simplification empty.\")\n else: combined_content += \"# Simplified DOM:\\n(JS Script Error)\\n\\n\"; logger.error(\"SIMPLIFY_PAGE_SCRIPT empty.\")\n except Exception as js_err: error_messages.append(f\"JS Error: {js_err}\"); logger.error(\"JS Simp. Error\", exc_info=False); combined_content += f\"# Simplified DOM Error: {js_err}\\n\"\n try:\n ax_tree = await page.accessibility.snapshot(interesting_only=False) # No root arg\n if ax_tree:\n try:\n ax_tree_str = json.dumps(ax_tree, separators=(',', ':')) # Compact\n ax_max_len = 2000\n if len(ax_tree_str) > ax_max_len: ax_tree_str = ax_tree_str[:ax_max_len] + \"...(AX Tree truncated)\"\n combined_content += f\"# Accessibility Tree (JSON, Partial):\\n```json\\n{ax_tree_str}\\n```\\n\\n\"; logger.debug(f\"AX Tree length: {len(ax_tree_str)}\")\n except Exception as json_err: error_messages.append(f\"AX JSON Error: {json_err}\"); logger.error(\"AX JSON Error\", exc_info=False); combined_content += \"# AX Tree Error (JSON)\\n\"\n else: combined_content += \"# Accessibility Tree:\\n(Empty)\\n\\n\"; logger.warning(\"AX snapshot empty.\")\n except Exception as ax_err: error_messages.append(f\"AX Tree Error: {ax_err}\"); logger.error(\"AX Tree Error\", exc_info=False); combined_content += f\"# Accessibility Tree Error: {ax_err}\\n\"\n\n if self.detector and screenshot_b64:\n logger.info(\"Attempting visual detection via Detector...\")\n try:\n detect_sheets = 'docs.google.com/spreadsheets/d' in current_url\n visual_elements = await self.detector.detect_from_image(screenshot_b64, detect_sheets)\n if visual_elements:\n formatted = [f\"- ID: {el.browser_agent_id}, Box: [L:{el.rect.get('left',0)}, T:{el.rect.get('top',0)}, R:{el.rect.get('right',0)}, B:{el.rect.get('bottom',0)}] (Tag: {el.tag_name})\" for el in visual_elements[:20]]\n combined_content += f\"# Visual Elements (Detected via CV, Max 20):\\n{chr(10).join(formatted)}\\n\\n\"; logger.info(f\"Added {len(formatted)} visual elements.\") # Use chr(10) for newline\n else: combined_content += \"# Visual Elements:\\n(None detected or VLM error)\\n\\n\"; logger.info(\"No visual elements detected.\")\n except Exception as cv_err: error_messages.append(f\"CV Error: {cv_err}\"); logger.error(\"CV Detector Error\", exc_info=True); combined_content += f\"# Visual Elements Error: {cv_err}\\n\"\n else:\n if not self.detector: logger.info(\"CV Detector not available.\")\n if not screenshot_b64: logger.info(\"Screenshot missing.\")\n combined_content += \"# Visual Elements:\\n(Not Run)\\n\\n\"\n\n if len(combined_content) > max_length:\n logger.warning(f\"Combined content ({len(combined_content)}) exceeds limit ({max_length}). Truncating.\")\n reserve = len(\"\\n\\n# Content Retrieval Errors:\\n- \") + sum(len(str(e)) + 4 for e in error_messages) + 50\n trunc_len = max(0, max_length - reserve); combined_content = combined_content[:trunc_len].rstrip() + \"\\n\\n... (Content truncated)\"\n if error_messages: combined_content += \"\\n\\n# Content Retrieval Errors:\\n- \" + \"\\n- \".join(map(str, error_messages))\n logger.info(f\"Finished getting content (final length: {len(combined_content)})\")\n return combined_content\n except Exception as e: logger.error(f\"General error in get_content: {e}\", exc_info=True); return f\"# Page URL:\\n{current_url}\\n# Error:\\nFailed to get content: {e}\"\n\n # --- Other Methods from Original Code ---\n\n async def get_cookies(self) -> list[dict[str, Any]]:\n \"\"\"Get cookies from the current browser context.\"\"\"\n if self.context:\n try: return await self.context.cookies()\n except Exception as e: logger.error(f\"Failed to get cookies: {e}\"); return []\n return []\n\n async def get_storage_state(self) -> dict[str, Any]:\n \"\"\"Get storage state (currently only cookies) from the browser.\"\"\"\n # Playwright's get_storage_state includes local/session storage too,\n # but might require more careful handling or filtering if large.\n # Sticking to cookies for simplicity based on original user code structure.\n if self.context:\n try:\n # cookies = await self.context.cookies() # Redundant if get_cookies exists\n # return {'cookies': cookies}\n # Or use the full state function if available and needed\n state = await self.context.storage_state()\n return state\n except Exception as e:\n logger.error(f\"Failed to get storage state: {e}\")\n return {}\n return {}\n\n async def get_tabs_info(self) -> list[TabInfo]:\n \"\"\"Get information about all open tabs in the current context.\"\"\"\n tabs_info = []\n if not self.context: return []\n try:\n # Ensure pages list is accessed correctly\n pages = self.context.pages\n for i, page in enumerate(pages):\n if not page.is_closed(): # Check if page is open\n try:\n url = page.url\n title = await page.title()\n # Ensure TabInfo model is available\n tabs_info.append(TabInfo(page_id=i, url=url, title=title))\n except Exception as page_err:\n logger.warning(f\"Failed to get info for tab {i}: {page_err}\")\n # Add placeholder if needed?\n tabs_info.append(TabInfo(page_id=i, url=\"Error\", title=\"Error retrieving info\"))\n\n except Exception as e:\n logger.error(f\"Failed to get tabs info: {e}\")\n return tabs_info\n\n async def switch_to_tab(self, page_id: int) -> None:\n \"\"\"Switch focus to a specific tab by its index.\"\"\"\n if self.context is None: await self._init_browser()\n pages = self.context.pages\n if not 0 <= page_id < len(pages):\n raise BrowserError(f'Invalid page_id: {page_id}. Available pages: {len(pages)}')\n if pages[page_id].is_closed():\n raise BrowserError(f'Page with page_id {page_id} is closed.')\n\n logger.info(f\"Switching to tab (page_id): {page_id}\")\n self.current_page = pages[page_id]\n try:\n await self.current_page.bring_to_front()\n # Wait briefly for potential state changes after switch\n await self.current_page.wait_for_load_state('domcontentloaded', timeout=5000)\n except Exception as e:\n logger.warning(f\"Error during tab switch finalization for page {page_id}: {e}\")\n # Continue anyway, page is switched internally\n\n async def create_new_tab(self, url: str | None = None) -> None:\n \"\"\"Create a new tab, optionally navigating to a URL, and switch to it.\"\"\"\n if self.context is None: await self._init_browser()\n logger.info(f\"Creating new tab. Navigate to: {url if url else 'about:blank'}\")\n try:\n new_page = await self.context.new_page()\n self.current_page = new_page # Switch focus to the new page\n if url:\n await self.navigate_to(url) # Reuse navigate method\n else:\n await new_page.wait_for_load_state('domcontentloaded') # Wait for about:blank load\n logger.info(f\"Switched to new tab. URL: {self.current_page.url}\")\n except Exception as e:\n logger.error(f\"Failed to create new tab: {e}\")\n raise BrowserError(f\"Failed to create new tab: {e}\") from e\n\n\n async def close_current_tab(self):\n \"\"\"Close the currently focused tab.\"\"\"\n if self.current_page is None: logger.warning(\"No current page to close.\"); return\n if len(self.context.pages) <= 1: logger.warning(\"Cannot close the last remaining tab.\"); return # Prevent closing last tab? Or allow context close?\n\n logger.info(f\"Closing current tab: {self.current_page.url}\")\n page_to_close = self.current_page\n # Find index to switch to after closing (e.g., previous or first)\n pages = self.context.pages\n current_index = pages.index(page_to_close) if page_to_close in pages else -1\n switch_to_index = 0 if current_index != 0 else 1 # Switch to first unless closing first\n if switch_to_index >= len(pages): switch_to_index = 0 # Fallback\n\n try:\n await page_to_close.close()\n logger.info(\"Tab closed.\")\n # Need to wait briefly for context.pages to update sometimes\n await asyncio.sleep(0.1)\n # Switch to another tab if possible\n if self.context and self.context.pages:\n new_current_page = self.context.pages[min(switch_to_index, len(self.context.pages)-1)]\n self.current_page = new_current_page\n await self.current_page.bring_to_front()\n logger.info(f\"Switched to tab index {min(switch_to_index, len(self.context.pages)-1)} after closing.\")\n else:\n self.current_page = None # No pages left\n logger.info(\"Closed the last tab.\")\n\n except Exception as e:\n logger.error(f\"Error closing tab or switching: {e}\")\n # Attempt to recover current page if possible\n if self.context and self.context.pages: self.current_page = self.context.pages[0]\n else: self.current_page = None\n\n async def refresh_page(self):\n \"\"\"Refresh the current page.\"\"\"\n page = await self.get_current_page()\n logger.info(f\"Refreshing page: {page.url}\")\n try:\n await page.reload(wait_until='domcontentloaded')\n logger.info(\"Page refreshed.\")\n except Exception as e:\n logger.error(f\"Failed to refresh page: {e}\")\n raise BrowserError(f\"Failed to refresh page: {e}\") from e\n\n async def go_forward(self):\n \"\"\"Navigate forward in the current page's history.\"\"\"\n page = await self.get_current_page()\n logger.info(f\"Going forward in history for: {page.url}\")\n try:\n await page.go_forward(wait_until='domcontentloaded', timeout=10000) # Added timeout\n logger.info(f\"Navigated forward. New URL: {page.url}\")\n except Exception as e:\n # Often fails if no forward history exists, log as warning\n logger.warning(f'Failed to go forward (might be end of history): {e}')\n # raise BrowserError(f\"Failed to go forward: {e}\") from e # Option: re-raise if needed\n\n # --- State Update Methods (using CV potentially) ---\n def get_state(self) -> Optional[BrowserState]:\n \"\"\"Get the last updated internal browser state.\"\"\"\n # Returns the state cached from the last update_state call\n logger.debug(f\"Returning cached browser state (URL: {self._state.url if self._state else 'None'})\")\n return self._state\n\n @observe(name='browser.update_state', ignore_output=True)\n async def update_state(self) -> BrowserState:\n \"\"\"Update the internal browser state by re-evaluating the page (incl. CV if enabled).\"\"\"\n logger.info(\"Updating browser state...\")\n try:\n self._state = await self._update_state()\n logger.info(\"Browser state updated successfully.\")\n if not self._state: raise BrowserError(\"State update returned None unexpectedly.\") # Should not happen if _update_state raises\n return self._state\n except Exception as e:\n logger.error(f\"Failed to update browser state: {e}\", exc_info=True)\n # Decide whether to return old state or raise error\n # Raising error seems more appropriate if update fails\n raise BrowserError(f\"Failed to update state: {e}\") from e\n\n\n @observe(name='browser._update_state', ignore_output=True)\n async def _update_state(self) -> BrowserState:\n \"\"\"Internal method to get comprehensive state with retry logic.\"\"\"\n @retry(\n stop=stop_after_attempt(3),\n wait=wait_exponential(multiplier=0.5, min=0.5, max=2),\n retry=retry_if_exception_type((Exception)), # Retry on any exception during state fetch\n reraise=True # Re-raise the exception after retries fail\n )\n async def get_stable_state():\n page = await self.get_current_page() # Ensures page exists\n url = page.url\n detect_sheets = 'docs.google.com/spreadsheets/d' in url\n screenshot_b64 = await self.fast_screenshot() # Get screenshot\n\n interactive_elements_data: Optional[InteractiveElementsData] = None\n # Get combined elements using CV if detector is enabled\n if self.detector and screenshot_b64:\n logger.debug(\"Getting interactive elements with CV...\")\n interactive_elements_data = await self.get_interactive_elements_with_cv(screenshot_b64, detect_sheets)\n # Fallback to browser-only if detector disabled or screenshot failed\n elif INTERACTIVE_ELEMENTS_JS_CODE: # Ensure JS code loaded\n logger.debug(\"Getting interactive elements with browser JS only...\")\n interactive_elements_data = await self.get_interactive_elements_data()\n else:\n logger.error(\"Cannot get interactive elements: Detector disabled/failed and JS code missing.\")\n interactive_elements_data = InteractiveElementsData(viewport={\"width\":0,\"height\":0}, elements=[]) # Return empty state\n\n\n # Check if interactive_elements_data is valid before proceeding\n if interactive_elements_data is None or not hasattr(interactive_elements_data, 'elements'):\n raise BrowserError(\"Failed to retrieve valid interactive elements data.\")\n\n # Process elements into dictionary for state\n interactive_elements = {element.browser_agent_id: element for element in interactive_elements_data.elements}\n\n # Generate highlighted screenshot\n screenshot_with_highlights = None\n if screenshot_b64 and 'put_highlight_elements_on_screenshot' in globals():\n try:\n screenshot_with_highlights = put_highlight_elements_on_screenshot(\n list(interactive_elements.values()), # Pass list of elements\n screenshot_b64\n )\n except Exception as high_err:\n logger.warning(f\"Failed to generate highlighted screenshot: {high_err}\")\n\n # Get tab info\n tabs = await self.get_tabs_info()\n\n # Ensure BrowserState model is available\n if 'BrowserState' not in globals() or 'BrowserState' not in locals():\n raise ImportError(\"BrowserState model is not defined or imported.\")\n\n # Create and return the state object\n return BrowserState(\n url=url,\n tabs=tabs,\n screenshot_with_highlights=screenshot_with_highlights,\n screenshot=screenshot_b64,\n viewport=interactive_elements_data.viewport, # Use viewport from data\n interactive_elements=interactive_elements,\n )\n\n # Execute the retry logic\n try:\n new_state = await get_stable_state()\n self._state = new_state # Cache the new state\n return new_state\n except Exception as e:\n logger.error(f'Failed to update state after multiple attempts: {e}', exc_info=True)\n # Don't return potentially stale state, let error propagate\n raise BrowserError(f\"Failed to update state definitively: {e}\") from e\n\n @observe(name='browser.get_interactive_elements')\n async def get_interactive_elements_data(self) -> InteractiveElementsData:\n \"\"\"Gets interactive elements using only in-browser JavaScript.\"\"\"\n page = await self.get_current_page()\n if not INTERACTIVE_ELEMENTS_JS_CODE:\n logger.error(\"INTERACTIVE_ELEMENTS_JS_CODE is empty. Cannot get elements.\")\n # Return default empty structure\n vp = await page.viewport_size() or {\"width\":0, \"height\":0}\n return InteractiveElementsData(viewport=vp, elements=[])\n try:\n result = await page.evaluate(INTERACTIVE_ELEMENTS_JS_CODE)\n # Validate result basic structure\n if not isinstance(result, dict) or 'viewport' not in result or 'elements' not in result:\n logger.error(f\"JS evaluation returned unexpected structure: {type(result)}\")\n vp = await page.viewport_size() or {\"width\":0, \"height\":0}\n return InteractiveElementsData(viewport=vp, elements=[])\n # Parse using Pydantic model if available\n if 'InteractiveElementsData' in globals() and 'InteractiveElementsData' in locals():\n return InteractiveElementsData(**result)\n else:\n # Fallback if model missing (though this indicates setup error)\n logger.error(\"InteractiveElementsData model missing, returning raw dict.\")\n return result # type: ignore\n except Exception as e:\n logger.error(f\"Error evaluating INTERACTIVE_ELEMENTS_JS_CODE: {e}\", exc_info=True)\n vp = await page.viewport_size() or {\"width\":0, \"height\":0}\n return InteractiveElementsData(viewport=vp, elements=[])\n\n\n @observe(name='browser.get_interactive_elements_with_cv')\n async def get_interactive_elements_with_cv(self, screenshot_b64: Optional[str] = None, detect_sheets: bool = False) -> InteractiveElementsData:\n \"\"\"Combines browser JS element detection with VLM detection.\"\"\"\n if self.detector is None:\n logger.warning(\"CV detector not available. Falling back to browser-only detection.\")\n return await self.get_interactive_elements_data()\n\n # Ensure screenshot exists\n current_screenshot_b64 = screenshot_b64 or await self.fast_screenshot()\n if not current_screenshot_b64:\n logger.error(\"Screenshot unavailable for CV detection.\")\n return await self.get_interactive_elements_data() # Fallback\n\n logger.debug(\"Getting combined browser + CV elements...\")\n try:\n # Run browser JS detection and VLM detection concurrently\n browser_elements_data_task = asyncio.create_task(self.get_interactive_elements_data())\n cv_elements_task = asyncio.create_task(self.detector.detect_from_image(current_screenshot_b64, detect_sheets))\n\n browser_elements_data = await browser_elements_data_task\n cv_elements = await cv_elements_task\n\n # Ensure results are valid before combining\n if not browser_elements_data or not hasattr(browser_elements_data, 'elements'):\n logger.warning(\"Browser element data invalid or missing for combine step.\")\n browser_elements = []\n viewport = await self.get_current_page().viewport_size() or {\"width\":0,\"height\":0}\n else:\n browser_elements = browser_elements_data.elements\n viewport = browser_elements_data.viewport # Use viewport from browser data\n\n if not isinstance(cv_elements, list):\n logger.warning(\"CV elements result is not a list.\")\n cv_elements = []\n\n # Combine results using utility function\n if 'combine_and_filter_elements' in globals():\n combined_elements = combine_and_filter_elements(browser_elements, cv_elements)\n logger.info(f\"Combined browser ({len(browser_elements)}) and CV ({len(cv_elements)}) elements into {len(combined_elements)}.\")\n else:\n logger.error(\"combine_and_filter_elements utility function not found. Returning only browser elements.\")\n combined_elements = browser_elements # Fallback\n\n # Return combined data in the expected structure\n if 'InteractiveElementsData' in globals() and 'InteractiveElementsData' in locals():\n return InteractiveElementsData(viewport=viewport, elements=combined_elements)\n else:\n logger.error(\"InteractiveElementsData model missing, returning raw combined list.\")\n # This fallback is problematic, structure is needed downstream\n return {\"viewport\": viewport, \"elements\": combined_elements} # type: ignore\n\n except Exception as e:\n logger.error(f\"Error during combined CV+Browser element detection: {e}\", exc_info=True)\n # Fallback gracefully to browser-only if possible\n try: return await self.get_interactive_elements_data()\n except Exception: return InteractiveElementsData(viewport={\"width\":0,\"height\":0}, elements=[]) # Final fallback"
},
{
"path": "super_agents/browser_use/browser/detector.py",
"content": "# super_agents/browser_use/browser/detector.py\nimport os\nimport json\nimport logging\nimport base64\nfrom typing import List, Optional, Dict, Any\n\n# LangChain Core Imports\nfrom langchain_core.messages import HumanMessage, SystemMessage\nfrom langchain_core.runnables.base import RunnableSerializable\n# Pydantic for schema\ntry:\n from pydantic.v1 import BaseModel\nexcept ImportError:\n from pydantic import BaseModel\n\nfrom tenacity import (\n retry,\n retry_if_exception_type,\n stop_after_attempt,\n wait_exponential,\n)\n\n# Local imports (ensure they exist)\ntry:\n from .observe_helper import observe\nexcept ImportError:\n def observe(name, ignore_input=False, ignore_output=False):\n def decorator(func): return func\n return decorator\n # Setup basic logger if not configured by main app yet\n logging.basicConfig(level=logging.WARNING)\n logger = logging.getLogger(__name__)\n logger.warning(\"observe_helper not found, using dummy decorator.\")\ntry:\n from .models import InteractiveElement\n # Define the expected VLM output schema here or import from agent.schemas\n # Let's define it here for clarity in this step\n class VLMJsonOutput(BaseModel):\n detected_elements: List[Dict[str, Any]] = []\nexcept ImportError:\n class InteractiveElement: pass\n class VLMJsonOutput(BaseModel): detected_elements: List = []\n # Setup basic logger if not configured by main app yet\n logging.basicConfig(level=logging.WARNING)\n logger = logging.getLogger(__name__)\n logger.error(\"Failed to import InteractiveElement or define VLMJsonOutput! Detector parsing will fail.\")\n\n# Import the specific ChatOpenRouter class from the updated llm.py\n# Adjust path if llm.py is elsewhere relative to detector.py\ntry:\n from ..llm import ChatOpenRouter # Assumes llm.py is one level up\nexcept ImportError:\n logger.error(\"Failed to import ChatOpenRouter from ..llm. Ensure llm.py is in the parent directory.\")\n # Define a dummy class to allow loading, but it won't work\n class ChatOpenRouter: pass\n\nlogger = logging.getLogger(__name__)\n\n# --- VLM Configuration (Read by Detector's __init__ via ChatOpenRouter) ---\nVLM_API_MODEL = os.getenv(\"VLM_API_MODEL\", \"openai/gpt-4o\") # Read desired VLM model from .env\n\n# --- VLM Prompt Template ---\nVLM_PROMPT_TEMPLATE = \"\"\"\nAnalyze the provided screenshot of a webpage. Your task is to identify all significant interactive elements visible on the screen. Interactive elements include: buttons, links ( tags), text input fields (, , etc.), password fields (), text areas (