6.3 KiB
6.3 KiB
Agent 回调功能实现总结
概述
根据 /home/taiji/tools/agent-manager/plans/LiteLLM和AgentManager回调接口文档.md 的要求，已为所有 agent 模板添加了运行时长回调功能。
修改内容
1. 新增工具模块
文件: agent_callback_utils.py
提供两个核心类：
- AgentCallbackHandler: 回调处理器，负责记录请求开始/结束时间、使用的工具、并发送回调到 Agent Manager
- CallbackContextManager: 上下文管理器，支持
with语句自动处理回调的开始和结束
关键功能:
- 自动记录运行时长（Pod running time）
- 记录使用的工具列表
- 向回调接口发送 POST 请求：
http://mcp-server:8002/api/v1/billing/agent-callback - 支持用户ID、请求ID追踪
2. API Key 传递方式改进
之前: 所有配置包括 API key 都从环境变量获取
现在:
- API key: 从用户请求中传入（每次请求携带）
- 其他配置: 依然从环境变量获取（数据库连接信息、服务端口等）
3. 修改的 Agent 文件
3.1 search_agent.py
请求模型修改:
class SearchRequest(BaseModel):
query: str
llm_api_key: str # 新增：从请求传入
user_id: Optional[str] # 新增：用于计费回调
auto_configure: bool
回调集成:
- 使用
CallbackContextManager自动处理回调 - 记录使用的工具：
web_search,content_reader - 临时更新 API key 后执行搜索，完成后恢复原值
3.2 jina_search_agent.py
请求模型修改:
class SearchRequest(BaseModel):
url: str
jina_api_key: str # 新增：从请求传入
user_id: Optional[str] # 新增
timeout: int
回调集成:
- 使用
CallbackContextManager自动处理回调 - 记录使用的工具：
jina_reader - 移除了全局
JINA_API_KEY环境变量依赖
3.3 mysql_agent.py
重大改动: 从循环示例查询模式改为 FastAPI HTTP 服务
请求模型:
class QueryRequest(BaseModel):
query: str
openai_api_key: str # 从请求传入
user_id: Optional[str] # 用于计费回调
model: str = "gpt-3.5-turbo"
新增端点:
GET /health- 健康检查POST /query- 执行SQL查询（带回调）GET /- 服务信息
回调集成:
- 使用
CallbackContextManager - 记录使用的工具：
sql_database
3.4 postgresql_agent.py
修改内容与 mysql_agent.py 类似
请求模型:
class QueryRequest(BaseModel):
query: str
openai_api_key: str # 从请求传入
user_id: Optional[str] # 用于计费回调
model: str = "gpt-3.5-turbo"
回调集成: 同 MySQL Agent
回调接口规范
根据文档，每次请求结束后自动发送以下回调：
{
"agentName": "pod-name",
"userId": "user-123",
"podRunningTimeSeconds": 120,
"toolsUsed": ["web_search", "content_reader"],
"startTime": "2026-01-15T10:00:00Z",
"endTime": "2026-01-15T10:02:00Z",
"requestId": "search-1737801600"
}
回调地址:
- 默认:
http://mcp-server:8002/api/v1/billing/agent-callback - 可通过环境变量
AGENT_CALLBACK_URL覆盖
环境变量配置
必需环境变量（所有 Agent）
POD_NAME=agent-name # Pod 名称（用于回调）
USER_ID=default-user-id # 默认用户ID（可被请求中的user_id覆盖）
可选环境变量
AGENT_CALLBACK_URL=http://mcp-server:8002/api/v1/billing/agent-callback # 回调地址
Agent 特定环境变量
Search Agent:
SERPER_API_KEY=xxx # Serper API（从环境变量）
JINA_API_KEY=xxx # Jina API（从环境变量，可被请求覆盖）
LLM_BASE_URL=xxx # LLM API基础URL
Jina Search Agent:
SERVICE_HOST=0.0.0.0
SERVICE_PORT=8080
MySQL Agent:
MYSQL_HOST=localhost
MYSQL_PORT=3306
MYSQL_USER=root
MYSQL_PASSWORD=password
MYSQL_DATABASE=test
SERVICE_HOST=0.0.0.0
SERVICE_PORT=8080
PostgreSQL Agent:
POSTGRES_HOST=localhost
POSTGRES_PORT=5432
POSTGRES_USER=postgres
POSTGRES_PASSWORD=password
POSTGRES_DATABASE=postgres
SERVICE_HOST=0.0.0.0
SERVICE_PORT=8080
使用示例
Search Agent
curl -X POST http://agent-ip:8080/search \
-H "Content-Type: application/json" \
-d '{
"query": "什么是人工智能？",
"llm_api_key": "sk-xxx",
"user_id": "user-123"
}'
Jina Search Agent
curl -X POST http://agent-ip:8080/search \
-H "Content-Type: application/json" \
-d '{
"url": "https://example.com",
"jina_api_key": "jina_xxx",
"user_id": "user-123"
}'
MySQL/PostgreSQL Agent
curl -X POST http://agent-ip:8080/query \
-H "Content-Type: application/json" \
-d '{
"query": "列出所有表",
"openai_api_key": "sk-xxx",
"user_id": "user-123",
"model": "gpt-3.5-turbo"
}'
回调流程
-
请求开始:
- 创建
CallbackContextManager上下文 - 记录开始时间
- 设置 user_id 和 request_id
- 创建
-
执行过程:
- 临时更新 API key（如需要）
- 执行 Agent 逻辑
- 记录使用的工具（通过
ctx.add_tool()）
-
请求结束:
- 自动计算运行时长
- 发送 POST 请求到回调接口
- 包含所有必需字段（agentName, userId, podRunningTimeSeconds, toolsUsed, startTime, endTime, requestId）
-
错误处理:
- 回调失败不影响主流程
- 错误日志记录
注意事项
- API Key 安全: API key 仅在请求期间临时使用，不持久化
- 回调可选: 如果
user_id未提供，回调处理器会跳过发送 - 幂等性: 使用
request_id确保回调幂等性 - 时区: 所有时间戳使用 UTC 时区
- 兼容性: 保持向后兼容，不强制要求 user_id
依赖要求
所有 Agent 需要添加以下 Python 包依赖：
fastapi
uvicorn
requests
pydantic
已有依赖的 Agent 无需额外安装。
测试建议
- 单元测试: 测试回调函数的正确性
- 集成测试: 验证回调接口能正常接收数据
- 压力测试: 确保回调不影响 Agent 性能
- 错误测试: 验证回调失败时 Agent 仍能正常工作
修改完成时间: 2026-01-15
修改人: GitHub Copilot
版本: v1.0