Initial commit: jina-reader-agent-v3 with 1 tools - src/server/agent_callback_utils.py
This commit is contained in:
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
Agent回调工具 - 用于向Agent Manager回调运行时长记录（计费）
|
||||
"""
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import requests
|
||||
from typing import Optional, List
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentCallbackHandler:
|
||||
"""Agent回调处理器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_name: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
callback_url: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
初始化回调处理器
|
||||
|
||||
Args:
|
||||
agent_name: Agent名称，默认从环境变量 POD_NAME 获取
|
||||
user_id: 用户ID，默认从环境变量 USER_ID 获取
|
||||
callback_url: 回调URL，默认从环境变量 AGENT_CALLBACK_URL 获取
|
||||
"""
|
||||
self.agent_name = agent_name or os.getenv("POD_NAME", "unknown-agent")
|
||||
self.user_id = user_id or os.getenv("USER_ID", "")
|
||||
self.callback_url = callback_url or os.getenv(
|
||||
"AGENT_CALLBACK_URL",
|
||||
"http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback"
|
||||
)
|
||||
|
||||
self.start_time: Optional[datetime] = None
|
||||
self.tools_used: List[str] = []
|
||||
self.request_id: Optional[str] = None
|
||||
|
||||
logger.info(f"AgentCallbackHandler initialized: agent={self.agent_name}, callback_url={self.callback_url}")
|
||||
|
||||
def start_request(self, request_id: Optional[str] = None, user_id: Optional[str] = None):
|
||||
"""
|
||||
开始一次请求处理
|
||||
|
||||
Args:
|
||||
request_id: 请求ID
|
||||
user_id: 用户ID（如果提供则覆盖默认值）
|
||||
"""
|
||||
self.start_time = datetime.now(timezone.utc)
|
||||
self.tools_used = []
|
||||
self.request_id = request_id or f"req-{int(time.time())}"
|
||||
|
||||
if user_id:
|
||||
self.user_id = user_id
|
||||
|
||||
logger.info(f"Request started: request_id={self.request_id}, user_id={self.user_id}")
|
||||
|
||||
def add_tool_used(self, tool_name: str):
|
||||
"""
|
||||
记录使用的工具
|
||||
|
||||
Args:
|
||||
tool_name: 工具名称
|
||||
"""
|
||||
if tool_name not in self.tools_used:
|
||||
self.tools_used.append(tool_name)
|
||||
logger.debug(f"Tool used: {tool_name}")
|
||||
|
||||
def end_request(self, tools_used: Optional[List[str]] = None) -> bool:
|
||||
"""
|
||||
结束请求并发送回调
|
||||
|
||||
Args:
|
||||
tools_used: 使用的工具列表（可选，如果提供则覆盖内部记录）
|
||||
|
||||
Returns:
|
||||
是否成功发送回调
|
||||
"""
|
||||
if not self.start_time:
|
||||
logger.warning("Cannot end request: no start time recorded")
|
||||
return False
|
||||
|
||||
if not self.user_id:
|
||||
logger.warning("Cannot send callback: user_id not set")
|
||||
return False
|
||||
|
||||
end_time = datetime.now(timezone.utc)
|
||||
running_time = (end_time - self.start_time).total_seconds()
|
||||
|
||||
# 使用提供的工具列表或内部记录
|
||||
final_tools_used = tools_used if tools_used is not None else self.tools_used
|
||||
|
||||
# 发送回调
|
||||
success = self._send_callback(
|
||||
running_time_seconds=int(running_time),
|
||||
start_time=self.start_time,
|
||||
end_time=end_time,
|
||||
tools_used=final_tools_used
|
||||
)
|
||||
|
||||
# 重置状态
|
||||
self.start_time = None
|
||||
self.tools_used = []
|
||||
self.request_id = None
|
||||
|
||||
return success
|
||||
|
||||
def _send_callback(
|
||||
self,
|
||||
running_time_seconds: int,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
tools_used: List[str]
|
||||
) -> bool:
|
||||
"""
|
||||
发送回调到Agent Manager
|
||||
|
||||
Args:
|
||||
running_time_seconds: 运行时长（秒）
|
||||
start_time: 开始时间
|
||||
end_time: 结束时间
|
||||
tools_used: 使用的工具列表
|
||||
|
||||
Returns:
|
||||
是否成功发送
|
||||
"""
|
||||
try:
|
||||
payload = {
|
||||
"agentName": self.agent_name,
|
||||
"userId": self.user_id,
|
||||
"podRunningTimeSeconds": running_time_seconds,
|
||||
"toolsUsed": tools_used,
|
||||
"startTime": start_time.isoformat(),
|
||||
"endTime": end_time.isoformat(),
|
||||
"requestId": self.request_id
|
||||
}
|
||||
|
||||
logger.info(f"Sending callback: {payload}")
|
||||
|
||||
response = requests.post(
|
||||
self.callback_url,
|
||||
json=payload,
|
||||
timeout=5
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
logger.info(f"Callback sent successfully: {response.json()}")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Callback failed with status {response.status_code}: {response.text}")
|
||||
return False
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Failed to send callback: {str(e)}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error sending callback: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
class CallbackContextManager:
|
||||
"""回调上下文管理器 - 使用with语句自动处理开始和结束"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
handler: AgentCallbackHandler,
|
||||
request_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
tools_used: Optional[List[str]] = None
|
||||
):
|
||||
"""
|
||||
初始化上下文管理器
|
||||
|
||||
Args:
|
||||
handler: AgentCallbackHandler实例
|
||||
request_id: 请求ID
|
||||
user_id: 用户ID
|
||||
tools_used: 使用的工具列表（可选）
|
||||
"""
|
||||
self.handler = handler
|
||||
self.request_id = request_id
|
||||
self.user_id = user_id
|
||||
self.tools_used = tools_used or []
|
||||
|
||||
def __enter__(self):
|
||||
"""进入上下文时开始计时"""
|
||||
self.handler.start_request(
|
||||
request_id=self.request_id,
|
||||
user_id=self.user_id
|
||||
)
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""退出上下文时发送回调"""
|
||||
self.handler.end_request(tools_used=self.tools_used)
|
||||
return False # 不抑制异常
|
||||
|
||||
def add_tool(self, tool_name: str):
|
||||
"""添加使用的工具"""
|
||||
self.handler.add_tool_used(tool_name)
|
||||
if tool_name not in self.tools_used:
|
||||
self.tools_used.append(tool_name)
|
||||
Reference in New Issue
Block a user