172 lines
5.6 KiB
Python
172 lines
5.6 KiB
Python
"""Agent 基类 - 参考 OpenManus BaseAgent 设计"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from contextlib import asynccontextmanager
|
|
from typing import List, Optional, Dict, Any
|
|
|
|
from pydantic import BaseModel, Field, model_validator
|
|
|
|
from app.llm import LLM
|
|
from app.logger import logger
|
|
from app.schema import AgentState, Memory, Message, ROLE_TYPE
|
|
|
|
|
|
class BaseAgent(BaseModel, ABC):
|
|
"""Agent 基类
|
|
|
|
提供基础的状态管理、记忆管理和执行循环。
|
|
子类需要实现 step() 方法来定义具体行为。
|
|
"""
|
|
|
|
# 核心属性
|
|
name: str = Field(default="BaseAgent", description="Agent 名称")
|
|
description: Optional[str] = Field(default=None, description="Agent 描述")
|
|
|
|
# 提示词
|
|
system_prompt: Optional[str] = Field(default=None, description="系统提示词")
|
|
next_step_prompt: Optional[str] = Field(default=None, description="下一步提示词")
|
|
|
|
# 依赖
|
|
llm: Optional[LLM] = Field(default=None, description="LLM 实例")
|
|
memory: Memory = Field(default_factory=Memory, description="Agent 记忆")
|
|
state: AgentState = Field(default=AgentState.IDLE, description="当前状态")
|
|
|
|
# 执行控制
|
|
max_steps: int = Field(default=10, description="最大执行步数")
|
|
current_step: int = Field(default=0, description="当前步数")
|
|
|
|
# 重复检测阈值
|
|
duplicate_threshold: int = 2
|
|
|
|
class Config:
|
|
arbitrary_types_allowed = True
|
|
extra = "allow"
|
|
|
|
@model_validator(mode="after")
|
|
def initialize_agent(self) -> "BaseAgent":
|
|
"""初始化 Agent"""
|
|
if self.llm is None:
|
|
try:
|
|
self.llm = LLM()
|
|
except Exception as e:
|
|
logger.warning(f"无法初始化 LLM: {e}")
|
|
if not isinstance(self.memory, Memory):
|
|
self.memory = Memory()
|
|
return self
|
|
|
|
@asynccontextmanager
|
|
async def state_context(self, new_state: AgentState):
|
|
"""状态上下文管理器，用于安全的状态转换"""
|
|
if not isinstance(new_state, AgentState):
|
|
raise ValueError(f"无效状态: {new_state}")
|
|
|
|
previous_state = self.state
|
|
self.state = new_state
|
|
try:
|
|
yield
|
|
except Exception as e:
|
|
self.state = AgentState.ERROR
|
|
raise e
|
|
finally:
|
|
self.state = previous_state
|
|
|
|
def update_memory(
|
|
self,
|
|
role: ROLE_TYPE,
|
|
content: str,
|
|
base64_image: Optional[str] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
"""添加消息到记忆"""
|
|
message_map = {
|
|
"user": Message.user_message,
|
|
"system": Message.system_message,
|
|
"assistant": Message.assistant_message,
|
|
"tool": lambda content, **kw: Message.tool_message(content, **kw),
|
|
}
|
|
|
|
if role not in message_map:
|
|
raise ValueError(f"不支持的消息角色: {role}")
|
|
|
|
if role == "tool":
|
|
self.memory.add_message(message_map[role](content, **kwargs))
|
|
else:
|
|
self.memory.add_message(message_map[role](content, base64_image=base64_image))
|
|
|
|
async def run(self, request: Optional[str] = None) -> str:
|
|
"""执行 Agent 主循环
|
|
|
|
Args:
|
|
request: 可选的初始用户请求
|
|
|
|
Returns:
|
|
执行结果摘要
|
|
"""
|
|
if self.state != AgentState.IDLE:
|
|
raise RuntimeError(f"无法从状态 {self.state} 启动 Agent")
|
|
|
|
if request:
|
|
self.update_memory("user", request)
|
|
|
|
results: List[str] = []
|
|
async with self.state_context(AgentState.RUNNING):
|
|
while (
|
|
self.current_step < self.max_steps
|
|
and self.state != AgentState.FINISHED
|
|
):
|
|
self.current_step += 1
|
|
logger.info(f"执行步骤 {self.current_step}/{self.max_steps}")
|
|
step_result = await self.step()
|
|
|
|
# 检测卡住状态
|
|
if self.is_stuck():
|
|
self.handle_stuck_state()
|
|
|
|
results.append(f"Step {self.current_step}: {step_result}")
|
|
|
|
if self.current_step >= self.max_steps:
|
|
self.current_step = 0
|
|
self.state = AgentState.IDLE
|
|
results.append(f"已终止: 达到最大步数 ({self.max_steps})")
|
|
|
|
return "\n".join(results) if results else "未执行任何步骤"
|
|
|
|
@abstractmethod
|
|
async def step(self) -> str:
|
|
"""执行单步操作 - 子类必须实现"""
|
|
pass
|
|
|
|
def handle_stuck_state(self):
|
|
"""处理卡住状态"""
|
|
stuck_prompt = "检测到重复响应。请考虑新策略，避免重复已尝试过的无效路径。"
|
|
self.next_step_prompt = f"{stuck_prompt}\n{self.next_step_prompt or ''}"
|
|
logger.warning(f"Agent 检测到卡住状态，已添加提示")
|
|
|
|
def is_stuck(self) -> bool:
|
|
"""检测是否陷入循环"""
|
|
if len(self.memory.messages) < 2:
|
|
return False
|
|
|
|
last_message = self.memory.messages[-1]
|
|
if not last_message.content:
|
|
return False
|
|
|
|
# 统计相同内容出现次数
|
|
duplicate_count = sum(
|
|
1
|
|
for msg in reversed(self.memory.messages[:-1])
|
|
if msg.role == "assistant" and msg.content == last_message.content
|
|
)
|
|
|
|
return duplicate_count >= self.duplicate_threshold
|
|
|
|
@property
|
|
def messages(self) -> List[Message]:
|
|
"""获取记忆中的消息列表"""
|
|
return self.memory.messages
|
|
|
|
@messages.setter
|
|
def messages(self, value: List[Message]):
|
|
"""设置记忆中的消息列表"""
|
|
self.memory.messages = value
|