From 257df1093cc72f1a01f8ed75da10b1105ced8791 Mon Sep 17 00:00:00 2001 From: zhanggangyong Date: Fri, 30 Jan 2026 17:41:13 +0000 Subject: [PATCH] Initial commit: jina-reader-v4 with 1 tools - src/server/agent_callback_utils.py --- src/server/agent_callback_utils.py | 205 +++++++++++++++++++++++++++++ 1 file changed, 205 insertions(+) create mode 100644 src/server/agent_callback_utils.py diff --git a/src/server/agent_callback_utils.py b/src/server/agent_callback_utils.py new file mode 100644 index 0000000..b6fa9b8 --- /dev/null +++ b/src/server/agent_callback_utils.py @@ -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)