152 lines
4.7 KiB
Python
152 lines
4.7 KiB
Python
"""
|
|
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
|
|
):
|
|
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(
|
|
"AgentCallbackHandler initialized: agent=%s, callback_url=%s",
|
|
self.agent_name,
|
|
self.callback_url,
|
|
)
|
|
|
|
def start_request(self, request_id: Optional[str] = None, user_id: Optional[str] = None):
|
|
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("Request started: request_id=%s, user_id=%s", self.request_id, self.user_id)
|
|
|
|
def add_tool_used(self, tool_name: str):
|
|
if tool_name not in self.tools_used:
|
|
self.tools_used.append(tool_name)
|
|
logger.debug("Tool used: %s", tool_name)
|
|
|
|
def end_request(self, tools_used: Optional[List[str]] = None) -> bool:
|
|
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:
|
|
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("Sending callback: %s", payload)
|
|
|
|
response = requests.post(
|
|
self.callback_url,
|
|
json=payload,
|
|
timeout=5
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
logger.info("Callback sent successfully: %s", response.json())
|
|
return True
|
|
|
|
logger.error("Callback failed with status %s: %s", response.status_code, response.text)
|
|
return False
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
logger.error("Failed to send callback: %s", str(e))
|
|
return False
|
|
except Exception as e:
|
|
logger.error("Unexpected error sending callback: %s", 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
|
|
):
|
|
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)
|