285 lines
7.3 KiB
Python
285 lines
7.3 KiB
Python
|
"""
|
|||
|
|
Agent 基类
|
|||
|
|
定义 Agent 的基本接口和通用功能
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from abc import ABC, abstractmethod
|
|||
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from enum import Enum
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AgentType(Enum):
|
|||
|
|
"""Agent 类型"""
|
|||
|
|
ORCHESTRATOR = "orchestrator"
|
|||
|
|
RECON = "recon"
|
|||
|
|
ANALYSIS = "analysis"
|
|||
|
|
VERIFICATION = "verification"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AgentPattern(Enum):
|
|||
|
|
"""Agent 运行模式"""
|
|||
|
|
REACT = "react" # 反应式：思考-行动-观察循环
|
|||
|
|
PLAN_AND_EXECUTE = "plan_execute" # 计划执行：先规划后执行
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class AgentConfig:
|
|||
|
|
"""Agent 配置"""
|
|||
|
|
name: str
|
|||
|
|
agent_type: AgentType
|
|||
|
|
pattern: AgentPattern = AgentPattern.REACT
|
|||
|
|
|
|||
|
|
# LLM 配置
|
|||
|
|
model: Optional[str] = None
|
|||
|
|
temperature: float = 0.1
|
|||
|
|
max_tokens: int = 4096
|
|||
|
|
|
|||
|
|
# 执行限制
|
|||
|
|
max_iterations: int = 20
|
|||
|
|
timeout_seconds: int = 600
|
|||
|
|
|
|||
|
|
# 工具配置
|
|||
|
|
tools: List[str] = field(default_factory=list)
|
|||
|
|
|
|||
|
|
# 系统提示词
|
|||
|
|
system_prompt: Optional[str] = None
|
|||
|
|
|
|||
|
|
# 元数据
|
|||
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class AgentResult:
|
|||
|
|
"""Agent 执行结果"""
|
|||
|
|
success: bool
|
|||
|
|
data: Any = None
|
|||
|
|
error: Optional[str] = None
|
|||
|
|
|
|||
|
|
# 执行统计
|
|||
|
|
iterations: int = 0
|
|||
|
|
tool_calls: int = 0
|
|||
|
|
tokens_used: int = 0
|
|||
|
|
duration_ms: int = 0
|
|||
|
|
|
|||
|
|
# 中间结果
|
|||
|
|
intermediate_steps: List[Dict[str, Any]] = field(default_factory=list)
|
|||
|
|
|
|||
|
|
# 元数据
|
|||
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict[str, Any]:
|
|||
|
|
return {
|
|||
|
|
"success": self.success,
|
|||
|
|
"data": self.data,
|
|||
|
|
"error": self.error,
|
|||
|
|
"iterations": self.iterations,
|
|||
|
|
"tool_calls": self.tool_calls,
|
|||
|
|
"tokens_used": self.tokens_used,
|
|||
|
|
"duration_ms": self.duration_ms,
|
|||
|
|
"metadata": self.metadata,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BaseAgent(ABC):
|
|||
|
|
"""
|
|||
|
|
Agent 基类
|
|||
|
|
所有 Agent 需要继承此类并实现核心方法
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
config: AgentConfig,
|
|||
|
|
llm_service,
|
|||
|
|
tools: Dict[str, Any],
|
|||
|
|
event_emitter=None,
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
初始化 Agent
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
config: Agent 配置
|
|||
|
|
llm_service: LLM 服务
|
|||
|
|
tools: 可用工具字典
|
|||
|
|
event_emitter: 事件发射器
|
|||
|
|
"""
|
|||
|
|
self.config = config
|
|||
|
|
self.llm_service = llm_service
|
|||
|
|
self.tools = tools
|
|||
|
|
self.event_emitter = event_emitter
|
|||
|
|
|
|||
|
|
# 运行状态
|
|||
|
|
self._iteration = 0
|
|||
|
|
self._total_tokens = 0
|
|||
|
|
self._tool_calls = 0
|
|||
|
|
self._cancelled = False
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def name(self) -> str:
|
|||
|
|
return self.config.name
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def agent_type(self) -> AgentType:
|
|||
|
|
return self.config.agent_type
|
|||
|
|
|
|||
|
|
@abstractmethod
|
|||
|
|
async def run(self, input_data: Dict[str, Any]) -> AgentResult:
|
|||
|
|
"""
|
|||
|
|
执行 Agent 任务
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
input_data: 输入数据
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Agent 执行结果
|
|||
|
|
"""
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
def cancel(self):
|
|||
|
|
"""取消执行"""
|
|||
|
|
self._cancelled = True
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def is_cancelled(self) -> bool:
|
|||
|
|
return self._cancelled
|
|||
|
|
|
|||
|
|
async def emit_event(
|
|||
|
|
self,
|
|||
|
|
event_type: str,
|
|||
|
|
message: str,
|
|||
|
|
**kwargs
|
|||
|
|
):
|
|||
|
|
"""发射事件"""
|
|||
|
|
if self.event_emitter:
|
|||
|
|
from ..event_manager import AgentEventData
|
|||
|
|
await self.event_emitter.emit(AgentEventData(
|
|||
|
|
event_type=event_type,
|
|||
|
|
message=message,
|
|||
|
|
**kwargs
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
async def emit_thinking(self, message: str):
|
|||
|
|
"""发射思考事件"""
|
|||
|
|
await self.emit_event("thinking", f"[{self.name}] {message}")
|
|||
|
|
|
|||
|
|
async def emit_tool_call(self, tool_name: str, tool_input: Dict):
|
|||
|
|
"""发射工具调用事件"""
|
|||
|
|
await self.emit_event(
|
|||
|
|
"tool_call",
|
|||
|
|
f"[{self.name}] 调用工具: {tool_name}",
|
|||
|
|
tool_name=tool_name,
|
|||
|
|
tool_input=tool_input,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
async def emit_tool_result(self, tool_name: str, result: str, duration_ms: int):
|
|||
|
|
"""发射工具结果事件"""
|
|||
|
|
await self.emit_event(
|
|||
|
|
"tool_result",
|
|||
|
|
f"[{self.name}] {tool_name} 完成 ({duration_ms}ms)",
|
|||
|
|
tool_name=tool_name,
|
|||
|
|
tool_duration_ms=duration_ms,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
async def call_tool(self, tool_name: str, **kwargs) -> Any:
|
|||
|
|
"""
|
|||
|
|
调用工具
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
tool_name: 工具名称
|
|||
|
|
**kwargs: 工具参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
工具执行结果
|
|||
|
|
"""
|
|||
|
|
tool = self.tools.get(tool_name)
|
|||
|
|
if not tool:
|
|||
|
|
logger.warning(f"Tool not found: {tool_name}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
self._tool_calls += 1
|
|||
|
|
await self.emit_tool_call(tool_name, kwargs)
|
|||
|
|
|
|||
|
|
import time
|
|||
|
|
start = time.time()
|
|||
|
|
|
|||
|
|
result = await tool.execute(**kwargs)
|
|||
|
|
|
|||
|
|
duration_ms = int((time.time() - start) * 1000)
|
|||
|
|
await self.emit_tool_result(tool_name, str(result.data)[:200], duration_ms)
|
|||
|
|
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
async def call_llm(
|
|||
|
|
self,
|
|||
|
|
messages: List[Dict[str, str]],
|
|||
|
|
tools: Optional[List[Dict]] = None,
|
|||
|
|
) -> Dict[str, Any]:
|
|||
|
|
"""
|
|||
|
|
调用 LLM
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
messages: 消息列表
|
|||
|
|
tools: 可用工具描述
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
LLM 响应
|
|||
|
|
"""
|
|||
|
|
self._iteration += 1
|
|||
|
|
|
|||
|
|
# 这里应该调用实际的 LLM 服务
|
|||
|
|
# 使用 LangChain 或直接调用 API
|
|||
|
|
try:
|
|||
|
|
response = await self.llm_service.chat_completion(
|
|||
|
|
messages=messages,
|
|||
|
|
temperature=self.config.temperature,
|
|||
|
|
max_tokens=self.config.max_tokens,
|
|||
|
|
tools=tools,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if response.get("usage"):
|
|||
|
|
self._total_tokens += response["usage"].get("total_tokens", 0)
|
|||
|
|
|
|||
|
|
return response
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"LLM call failed: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
def get_tool_descriptions(self) -> List[Dict[str, Any]]:
|
|||
|
|
"""获取工具描述（用于 LLM）"""
|
|||
|
|
descriptions = []
|
|||
|
|
|
|||
|
|
for name, tool in self.tools.items():
|
|||
|
|
if name.startswith("_"):
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
desc = {
|
|||
|
|
"type": "function",
|
|||
|
|
"function": {
|
|||
|
|
"name": name,
|
|||
|
|
"description": tool.description,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 添加参数 schema
|
|||
|
|
if hasattr(tool, 'args_schema') and tool.args_schema:
|
|||
|
|
desc["function"]["parameters"] = tool.args_schema.schema()
|
|||
|
|
|
|||
|
|
descriptions.append(desc)
|
|||
|
|
|
|||
|
|
return descriptions
|
|||
|
|
|
|||
|
|
def get_stats(self) -> Dict[str, Any]:
|
|||
|
|
"""获取执行统计"""
|
|||
|
|
return {
|
|||
|
|
"agent": self.name,
|
|||
|
|
"type": self.agent_type.value,
|
|||
|
|
"iterations": self._iteration,
|
|||
|
|
"tool_calls": self._tool_calls,
|
|||
|
|
"tokens_used": self._total_tokens,
|
|||
|
|
}
|
|||
|
|
|