Trae Agent中的agent模块

12 阅读5分钟

Agent核心逻辑全面解析

概述

Agent层是Trae Agent的核心,负责接收任务、与LLM交互、调用工具、执行任务并返回结果。本文档详细解析Agent的架构和执行流程。


Agent类架构

类的层级关系

┌─────────────────────────────────────────────────────────────┐
│                        Agent (agent.py)                     │
│                   入口类,工厂模式                            │
│  - 负责创建具体Agent实例                                    │
│  - 协调CLIConsole和TrajectoryRecorder                      │
│  - 管理Agent生命周期                                        │
└─────────────────────────────────────────────────────────────┘
                            │
                            ▼
┌─────────────────────────────────────────────────────────────┐
│                   BaseAgent (base_agent.py)                │
│                    抽象基类                                  │
│  - 定义Agent核心接口和执行流程                               │
│  - LLM调用和工具调用逻辑                                     │
│  - 轨迹记录和状态管理                                       │
└─────────────────────────────────────────────────────────────┘
                            │
                            ▼
┌─────────────────────────────────────────────────────────────┐
│                  TraeAgent (trae_agent.py)                 │
│                    具体实现                                  │
│  - 实现new_task方法                                        │
│  - MCP工具初始化和发现                                      │
│  - 特定业务逻辑                                            │
└─────────────────────────────────────────────────────────────┘

Agent入口类

文件: trae_agent/agent/agent.py

核心职责

  1. 创建具体Agent实例
  2. 协调各组件
  3. 管理生命周期

初始化流程

class Agent:
    def __init__(
        self,
        agent_type: AgentType | str,       # Agent类型
        config: Config,                     # 配置对象
        trajectory_file: str | None = None, # 轨迹文件
        cli_console: CLIConsole | None = None, # CLI控制台
        docker_config: dict | None = None,  # Docker配置
        docker_keep: bool = True,
    ):
        # 1. 转换Agent类型
        if isinstance(agent_type, str):
            agent_type = AgentType(agent_type)
        
        # 2. 创建轨迹记录器
        self.trajectory_recorder = TrajectoryRecorder(trajectory_file)
        
        # 3. 根据类型创建具体Agent (工厂模式)
        match self.agent_type:
            case AgentType.TraeAgent:
                self.agent = TraeAgent(
                    self.agent_config, 
                    docker_config=docker_config, 
                    docker_keep=docker_keep
                )
        
        # 4. 设置CLI控制台
        self.agent.set_cli_console(cli_console)
        
        # 5. 设置轨迹记录器
        self.agent.set_trajectory_recorder(self.trajectory_recorder)

run方法 - 任务执行入口

async def run(
    self,
    task: str,
    extra_args: dict[str, str] | None = None,
    tool_names: list[str] | None = None,
):
    # 1. 创建新任务
    self.agent.new_task(task, extra_args, tool_names)
    
    # 2. 初始化MCP工具
    if self.agent.allow_mcp_servers:
        await self.agent.initialise_mcp()
    
    # 3. 显示任务详情
    if self.agent.cli_console:
        self.agent.cli_console.print_task_details(task_details)
    
    # 4. 启动CLI控制台任务
    cli_console_task = asyncio.create_task(
        self.agent.cli_console.start()
    ) if self.agent.cli_console else None
    
    # 5. 执行任务
    try:
        execution = await self.agent.execute_task()
    finally:
        # 确保MCP清理
        await self.agent.cleanup_mcp_clients()
    
    # 6. 等待CLI控制台任务完成
    if cli_console_task:
        await cli_console_task
    
    return execution

BaseAgent抽象基类

文件: trae_agent/agent/base_agent.py

核心属性

class BaseAgent(ABC):
    def __init__(self, agent_config: AgentConfig, ...):
        # LLM客户端
        self._llm_client = LLMClient(agent_config.model)
        
        # 模型配置
        self._model_config = agent_config.model
        
        # 最大步数限制
        self._max_steps = agent_config.max_steps
        
        # 初始消息
        self._initial_messages: list[LLMMessage] = []
        
        # 当前任务
        self._task: str = ""
        
        # 可用工具列表
        self._tools: list[Tool] = [
            tools_registry[tool_name](...)
            for tool_name in agent_config.tools
        ]
        
        # 工具执行器
        self._tool_caller: Union[ToolExecutor, DockerToolExecutor]
        
        # CLI控制台
        self._cli_console: CLIConsole | None = None
        
        # 轨迹记录器
        self._trajectory_recorder: TrajectoryRecorder | None = None
        
        # Docker管理器
        self.docker_manager: DockerManager | None = None

核心执行流程

execute_task方法 - 主循环

async def execute_task(self) -> AgentExecution:
    # 1. 启动Docker容器
    if self.docker_manager:
        self.docker_manager.start()
    
    # 2. 初始化执行对象
    execution = AgentExecution(task=self._task, steps=[])
    execution.agent_state = AgentState.RUNNING
    start_time = time.time()
    
    messages = self._initial_messages
    step_number = 1
    
    # 3. 主循环 - 按步执行
    while step_number <= self._max_steps:
        # 创建步骤
        step = AgentStep(step_number=step_number, state=AgentStepState.THINKING)
        
        try:
            # 执行一步LLM交互
            messages = await self._run_llm_step(step, messages, execution)
            
            # 完成步骤 (记录轨迹,更新CLI)
            await self._finalize_step(step, messages, execution)
            
            # 检查是否完成
            if execution.agent_state == AgentState.COMPLETED:
                break
                
            step_number += 1
            
        except Exception as error:
            # 错误处理
            execution.agent_state = AgentState.ERROR
            step.state = AgentStepState.ERROR
            step.error = str(error)
            await self._finalize_step(step, messages, execution)
            break
    
    # 4. 超时处理
    if step_number > self._max_steps and not execution.success:
        execution.final_result = "Task execution exceeded maximum steps..."
        execution.agent_state = AgentState.ERROR
    
    # 5. 清理
    finally:
        if self.docker_manager and not self.docker_keep:
            self.docker_manager.stop()
    
    # 6. 释放工具资源
    await self._close_tools()
    
    # 7. 清理MCP客户端
    await self.cleanup_mcp_clients()
    
    execution.execution_time = time.time() - start_time
    
    return execution

LLM调用流程

_run_llm_step方法

async def _run_llm_step(
    self, 
    step: AgentStep, 
    messages: list[LLMMessage], 
    execution: AgentExecution
) -> list[LLMMessage]:
    
    # 1. 显示思考状态
    step.state = AgentStepState.THINKING
    self._update_cli_console(step, execution)
    
    # 2. 调用LLM获取响应
    llm_response = self._llm_client.chat(
        messages, 
        self._model_config, 
        self._tools
    )
    step.llm_response = llm_response
    
    # 3. 更新显示
    self._update_cli_console(step, execution)
    
    # 4. 更新token使用统计
    self._update_llm_usage(llm_response, execution)
    
    # 5. 检查LLM是否表示任务完成
    if self.llm_indicates_task_completed(llm_response):
        if self._is_task_completed(llm_response):
            # 任务完成
            execution.agent_state = AgentState.COMPLETED
            execution.final_result = llm_response.content
            execution.success = True
            return messages
        else:
            # 任务未完成,发送提示
            execution.agent_state = AgentState.RUNNING
            return [LLMMessage(
                role="user", 
                content=self.task_incomplete_message()
            )]
    else:
        # 6. 处理工具调用
        tool_calls = llm_response.tool_calls
        return await self._tool_call_handler(tool_calls, step)

工具调用流程

_tool_call_handler方法

async def _tool_call_handler(
    self, 
    tool_calls: list[ToolCall] | None, 
    step: AgentStep
) -> list[LLMMessage]:
    
    # 1. 无工具调用时返回提示
    if not tool_calls or len(tool_calls) <= 0:
        return [LLMMessage(
            role="user",
            content="It seems that you have not completed the task."
        )]
    
    # 2. 设置状态为调用工具中
    step.state = AgentStepState.CALLING_TOOL
    step.tool_calls = tool_calls
    self._update_cli_console(step)
    
    # 3. 执行工具调用 (并行或串行)
    if self._model_config.parallel_tool_calls:
        tool_results = await self._tool_caller.parallel_tool_call(tool_calls)
    else:
        tool_results = await self._tool_caller.sequential_tool_call(tool_calls)
    
    step.tool_results = tool_results
    self._update_cli_console(step)
    
    # 4. 将工具结果转换为消息
    messages = []
    for tool_result in tool_results:
        message = LLMMessage(role="user", tool_result=tool_result)
        messages.append(message)
    
    # 5. 反思工具执行结果
    reflection = self.reflect_on_result(tool_results)
    if reflection:
        step.state = AgentStepState.REFLECTING
        step.reflection = reflection
        self._update_cli_console(step)
        messages.append(LLMMessage(role="assistant", content=reflection))
    
    return messages

步骤状态机

AgentStepState枚举

┌─────────────────────────────────────────────────────────────┐
│                     步骤状态机                               │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  THINKING ──▶ CALLING_TOOL ──▶ REFLECTING ──▶ COMPLETED   │
│     │              │                │                      │
│     │              │                │                      │
│     ▼              ▼                ▼                      │
│  ERROR          (返回LLM)        (返回LLM)                 │
│                                                             │
└─────────────────────────────────────────────────────────────┘
状态说明
THINKINGLLM正在思考
CALLING_TOOL正在调用工具
REFLECTING正在反思工具结果
COMPLETED步骤完成
ERROR发生错误

Agent执行状态

AgentState枚举

状态说明
RUNNING正在执行
COMPLETED成功完成
ERROR执行出错

TraeAgent具体实现

文件: trae_agent/agent/trae_agent.py

核心职责

  1. 实现new_task方法 - 创建新任务并初始化消息
  2. MCP工具管理 - 发现和初始化MCP工具
  3. 业务逻辑 - 特定于TraeAgent的逻辑

new_task方法

async def new_task(
    self,
    task: str,
    extra_args: dict[str, str] | None = None,
    tool_names: list[str] | None = None,
):
    # 1. 设置任务
    self._task = task
    
    # 2. 构建系统提示词
    system_prompt = TRAE_AGENT_SYSTEM_PROMPT
    
    # 3. 构建初始消息
    self._initial_messages = [
        LLMMessage(role="system", content=system_prompt),
        LLMMessage(
            role="user", 
            content=f"[Project root path]: {project_root}\n\n{issue_text}"
        ),
    ]
    
    # 4. 更新工具列表
    if tool_names:
        self._tools = [
            tools_registry[tool_name](...)
            for tool_name in tool_names
        ]

MCP初始化

async def initialise_mcp(self):
    """初始化MCP工具"""
    await self.discover_mcp_tools()
    
    # 将MCP工具添加到工具列表
    if self.mcp_tools:
        self._tools.extend(self.mcp_tools)

async def discover_mcp_tools(self):
    """发现MCP服务器提供的工具"""
    if self.mcp_servers_config:
        for mcp_server_name, mcp_server_config in self.mcp_servers_config.items():
            # 检查是否允许
            if mcp_server_name not in self.allow_mcp_servers:
                continue
            
            # 连接并发现工具
            mcp_client = MCPClient()
            await mcp_client.connect_and_discover(...)
            
            self.mcp_clients.append(mcp_client)

完整执行流程图

┌─────────────────────────────────────────────────────────────────────┐
│                        CLI层                                         │
│  trae-cli run "修复bug"                                             │
└────────────────────────────┬──────────────────────────────────────────┘
                           │
                           ▼
┌─────────────────────────────────────────────────────────────────────┐
│                     Agent.run()                                       │
│  1. new_task() - 创建任务                                           │
│  2. initialise_mcp() - 初始化MCP                                   │
│  3. execute_task() - 执行任务                                       │
└────────────────────────────┬──────────────────────────────────────────┘
                           │
                           ▼
┌─────────────────────────────────────────────────────────────────────┐
│                   BaseAgent.execute_task()                            │
│                                                                     │
│  ┌──────────────────────────────────────────────────────────────┐  │
│  │                     主循环 (while step <= max_steps)          │  │
│  │                                                               │  │
│  │  1. 创建AgentStep (THINKING)                                  │  │
│  │         │                                                     │  │
│  │         ▼                                                     │  │
│  │  2. _run_llm_step()                                          │  │
│  │         │                                                     │  │
│  │         ├── LLM调用 (llm_client.chat)                       │  │
│  │         │                                                     │  │
│  │         ├── 检查完成标志                                      │  │
│  │         │                                                     │  │
│  │         └── _tool_call_handler() (如果有工具调用)            │  │
│  │                  │                                            │  │
│  │                  ├── 串行/并行执行工具                        │  │
│  │                  │                                             │  │
│  │                  └── 反思结果                                  │  │
│  │                                                               │  │
│  │  3. _finalize_step()                                          │  │
│  │         │                                                     │  │
│  │         ├── 记录轨迹                                          │  │
│  │         │                                                     │  │
│  │         └── 更新CLI显示                                        │  │
│  │                                                               │  │
│  └──────────────────────────────────────────────────────────────┘  │
│                           │                                         │
│                           ▼                                         │
│                    任务完成/超时/错误                                 │
└────────────────────────────┬──────────────────────────────────────────┘
                           │
                           ▼
┌─────────────────────────────────────────────────────────────────────┐
│                     返回AgentExecution                               │
│  - execution.success                                               │
│  - execution.final_result                                         │
│  - execution.steps[]                                              │
│  - execution.execution_time                                       │
└─────────────────────────────────────────────────────────────────────┘

核心类说明

类名文件职责
Agentagent.py入口类,工厂模式创建具体Agent
BaseAgentbase_agent.py抽象基类,定义核心执行逻辑
TraeAgenttrae_agent.py具体实现,MCP支持
LLMClientllm_client.pyLLM客户端封装
ToolExecutorbase.py工具执行器
CLIConsolecli_console.pyCLI输出接口

关键流程总结

  1. 入口: CLI调用 Agent.run()
  2. 初始化: Agent 创建 TraeAgent 实例
  3. 任务创建: new_task() 构建初始消息
  4. MCP初始化: initialise_mcp() 发现工具
  5. 主循环: execute_task() 按步执行
  6. LLM交互: _run_llm_step() 调用LLM
  7. 工具调用: _tool_call_handler() 执行工具
  8. 状态更新: 每个步骤更新状态和CLI显示
  9. 轨迹记录: _record_handler() 记录每步
  10. 资源清理: 清理MCP和工具资源

关键文件索引

文件行数职责
agent/agent.py1-100Agent入口类
agent/base_agent.py1-350BaseAgent抽象基类
agent/trae_agent.py1-260TraeAgent实现
agent/agent_basics.py全部数据类型定义
utils/llm_client.py全部LLM客户端
tools/base.py全部工具基类

最后更新: 2025-03-12