xiamuceer
187 lines
6.7 KiB
Python
187 lines
6.7 KiB
Python
"""AI 调用统计与中文日志格式化工具"""
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
|
||
@dataclass
|
||
class TokenUsage:
|
||
"""Token 使用量统计"""
|
||
|
||
prompt_tokens: Optional[int] = None
|
||
completion_tokens: Optional[int] = None
|
||
total_tokens: Optional[int] = None
|
||
|
||
@classmethod
|
||
def from_response(cls, response: Optional[Dict[str, Any]]) -> "TokenUsage":
|
||
"""从响应中提取 usage 信息"""
|
||
if not response:
|
||
return cls()
|
||
|
||
usage = response.get("usage") or {}
|
||
prompt_tokens = cls._to_int(usage.get("prompt_tokens"))
|
||
completion_tokens = cls._to_int(usage.get("completion_tokens"))
|
||
total_tokens = cls._to_int(usage.get("total_tokens"))
|
||
|
||
return cls(
|
||
prompt_tokens=prompt_tokens,
|
||
completion_tokens=completion_tokens,
|
||
total_tokens=total_tokens,
|
||
)
|
||
|
||
@staticmethod
|
||
def _to_int(value: Any) -> Optional[int]:
|
||
if value is None:
|
||
return None
|
||
try:
|
||
return int(value)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
def add(self, other: "TokenUsage") -> None:
|
||
"""累加另一个 usage"""
|
||
self.prompt_tokens = self._sum_optional(self.prompt_tokens, other.prompt_tokens)
|
||
self.completion_tokens = self._sum_optional(self.completion_tokens, other.completion_tokens)
|
||
self.total_tokens = self._sum_optional(self.total_tokens, other.total_tokens)
|
||
|
||
@staticmethod
|
||
def _sum_optional(left: Optional[int], right: Optional[int]) -> Optional[int]:
|
||
if left is None and right is None:
|
||
return None
|
||
return (left or 0) + (right or 0)
|
||
|
||
|
||
@dataclass
|
||
class ToolCallMetrics:
|
||
"""MCP 工具调用统计"""
|
||
|
||
tool_calls_count: int = 0
|
||
mcp_rounds: int = 0
|
||
tool_error_count: int = 0
|
||
tool_names: List[str] = field(default_factory=list)
|
||
usage: TokenUsage = field(default_factory=TokenUsage)
|
||
|
||
def add_tool_name(self, tool_name: str) -> None:
|
||
if tool_name and tool_name not in self.tool_names:
|
||
self.tool_names.append(tool_name)
|
||
|
||
|
||
@dataclass
|
||
class AICallMetrics:
|
||
"""单次 AI 调用统计"""
|
||
|
||
request_mode: str
|
||
provider: str
|
||
model: str
|
||
user_id: Optional[str] = None
|
||
stream: bool = False
|
||
auto_mcp: bool = False
|
||
tools_count: int = 0
|
||
prompt_length: int = 0
|
||
response_length: int = 0
|
||
chunk_count: int = 0
|
||
retry_count: int = 0
|
||
json_parse_success: Optional[bool] = None
|
||
finish_reason: Optional[str] = None
|
||
success: bool = False
|
||
error_type: Optional[str] = None
|
||
error_message: Optional[str] = None
|
||
ttft_ms: Optional[int] = None
|
||
duration_ms: Optional[int] = None
|
||
has_output: bool = False
|
||
usage: TokenUsage = field(default_factory=TokenUsage)
|
||
tool_metrics: ToolCallMetrics = field(default_factory=ToolCallMetrics)
|
||
started_at: float = field(default_factory=time.perf_counter)
|
||
first_chunk_at: Optional[float] = None
|
||
|
||
def mark_first_chunk(self) -> None:
|
||
if self.first_chunk_at is None:
|
||
self.first_chunk_at = time.perf_counter()
|
||
self.ttft_ms = int((self.first_chunk_at - self.started_at) * 1000)
|
||
|
||
def finish(
|
||
self,
|
||
*,
|
||
success: bool,
|
||
response_length: Optional[int] = None,
|
||
finish_reason: Optional[str] = None,
|
||
usage: Optional[TokenUsage] = None,
|
||
error: Optional[BaseException] = None,
|
||
) -> None:
|
||
self.success = success
|
||
self.duration_ms = int((time.perf_counter() - self.started_at) * 1000)
|
||
if response_length is not None:
|
||
self.response_length = response_length
|
||
self.has_output = self.response_length > 0
|
||
if finish_reason is not None:
|
||
self.finish_reason = finish_reason
|
||
if usage is not None:
|
||
self.usage = usage
|
||
if error is not None:
|
||
self.error_type = type(error).__name__
|
||
self.error_message = self._truncate(str(error), 180)
|
||
|
||
def merge_tool_metrics(self, tool_metrics: ToolCallMetrics) -> None:
|
||
self.tool_metrics = tool_metrics
|
||
self.usage.add(tool_metrics.usage)
|
||
|
||
def to_log_message(self, title: str) -> str:
|
||
fields = [
|
||
("请求类型", self.request_mode),
|
||
("提供商", self.provider),
|
||
("模型", self.model),
|
||
("状态", "成功" if self.success else "失败"),
|
||
("首字耗时", self._format_latency(self.ttft_ms, allow_empty=True)),
|
||
("总耗时", self._format_latency(self.duration_ms, allow_empty=False)),
|
||
("输入字符数", str(self.prompt_length)),
|
||
("输出字符数", str(self.response_length)),
|
||
("输入Token", self._format_optional_number(self.usage.prompt_tokens)),
|
||
("输出Token", self._format_optional_number(self.usage.completion_tokens)),
|
||
("总Token", self._format_optional_number(self.usage.total_tokens)),
|
||
("流式块数", str(self.chunk_count) if self.stream else "不适用"),
|
||
("启用MCP", "是" if self.auto_mcp else "否"),
|
||
("工具数", str(self.tools_count)),
|
||
("工具调用次数", str(self.tool_metrics.tool_calls_count)),
|
||
("MCP轮次", str(self.tool_metrics.mcp_rounds)),
|
||
("重试次数", str(self.retry_count) if self.retry_count else "0"),
|
||
("JSON解析", self._format_json_parse_result()),
|
||
("结束原因", self.finish_reason or "未知"),
|
||
]
|
||
|
||
if self.user_id:
|
||
fields.append(("用户ID", self.user_id))
|
||
if self.tool_metrics.tool_names:
|
||
fields.append(("工具名称", ",".join(self.tool_metrics.tool_names)))
|
||
if self.error_type:
|
||
fields.append(("异常类型", self.error_type))
|
||
if self.error_message:
|
||
fields.append(("异常摘要", self.error_message))
|
||
|
||
formatted = "｜".join(f"{key}={value}" for key, value in fields)
|
||
return f"{title}｜{formatted}"
|
||
|
||
def _format_json_parse_result(self) -> str:
|
||
if self.json_parse_success is None:
|
||
return "不适用"
|
||
return "成功" if self.json_parse_success else "失败"
|
||
|
||
@staticmethod
|
||
def _format_optional_number(value: Optional[int]) -> str:
|
||
return str(value) if value is not None else "未知"
|
||
|
||
@staticmethod
|
||
def _format_latency(value: Optional[int], allow_empty: bool) -> str:
|
||
if value is None:
|
||
return "无" if allow_empty else "未知"
|
||
if value < 1000:
|
||
return f"{value}ms"
|
||
return f"{value / 1000:.2f}s"
|
||
|
||
@staticmethod
|
||
def _truncate(text: str, limit: int) -> str:
|
||
if len(text) <= limit:
|
||
return text
|
||
return f"{text[:limit]}..."
|