Initial commit: aitool with 1 tools - src/server/agent_callback_utils.py

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)