""" 消息遥测记录器。 从 api/handlers/base/base_handler.py 迁移到 services 层, 消除 services→api 的反向依赖。 """ from __future__ import annotations from typing import Any from sqlalchemy.orm import Session from src.core.logger import logger from src.services.system.audit import audit_service from src.services.usage.service import UsageService class MessageTelemetry: """ 负责记录 Usage/Audit,避免处理器里重复代码。 """ def __init__( self, db: Session, user: Any, api_key: Any, request_id: str, client_ip: str ) -> None: self.db = db self.user = user self.api_key = api_key self.request_id = request_id self.client_ip = client_ip def _build_usage_metadata( self, *, request_metadata: dict[str, Any] | None = None, response_metadata: dict[str, Any] | None = None, ) -> dict[str, Any] | None: metadata: dict[str, Any] | None = None if request_metadata: metadata = dict(request_metadata) if response_metadata: metadata.setdefault("response", response_metadata) elif response_metadata: metadata = dict(response_metadata) return metadata async def calculate_cost( self, provider: str, model: str, *, input_tokens: int, output_tokens: int, cache_creation_tokens: int = 0, cache_read_tokens: int = 0, ) -> float: input_price, output_price = await UsageService.get_model_price_async( self.db, provider, model ) _, _, _, _, _, _, total_cost = UsageService.calculate_cost( input_tokens, output_tokens, input_price, output_price, cache_creation_tokens, cache_read_tokens, *await UsageService.get_cache_prices_async(self.db, provider, model, input_price), ) return total_cost async def record_success( self, *, provider: str, model: str, input_tokens: int, output_tokens: int, response_time_ms: int, status_code: int, request_body: dict[str, Any], request_headers: dict[str, Any], response_body: Any, response_headers: dict[str, Any], client_response_headers: dict[str, Any] | None = None, cache_creation_tokens: int = 0, cache_read_tokens: int = 0, cache_creation_tokens_5m: int = 0, cache_creation_tokens_1h: int = 0, is_stream: bool = False, provider_request_headers: dict[str, Any] | None = None, provider_request_body: Any | None = None, client_response_body: Any | None = None, # 时间指标 first_byte_time_ms: int | None = None, # 首字时间/TTFB # Provider 侧追踪信息(用于记录真实成本) provider_id: str | None = None, provider_endpoint_id: str | None = None, provider_api_key_id: str | None = None, api_format: str | None = None, # 结构化格式维度(从 Adapter 层透传) api_family: str | None = None, endpoint_kind: str | None = None, # 格式转换追踪 endpoint_api_format: str | None = None, # 端点原生 API 格式 has_format_conversion: bool = False, # 是否发生了格式转换 # 模型映射信息 target_model: str | None = None, # Provider 响应元数据(如 Gemini 的 modelVersion) response_metadata: dict[str, Any] | None = None, # 请求元数据(用于性能与调试记录) request_metadata: dict[str, Any] | None = None, ) -> float: metadata = self._build_usage_metadata( request_metadata=request_metadata, response_metadata=response_metadata, ) usage = await UsageService.record_usage( db=self.db, user=self.user, api_key=self.api_key, provider=provider, model=model, input_tokens=input_tokens, output_tokens=output_tokens, cache_creation_input_tokens=cache_creation_tokens, cache_read_input_tokens=cache_read_tokens, cache_creation_input_tokens_5m=cache_creation_tokens_5m, cache_creation_input_tokens_1h=cache_creation_tokens_1h, request_type="chat", api_format=api_format, api_family=api_family, endpoint_kind=endpoint_kind, endpoint_api_format=endpoint_api_format, has_format_conversion=has_format_conversion, is_stream=is_stream, response_time_ms=response_time_ms, first_byte_time_ms=first_byte_time_ms, # 传递首字时间 status_code=status_code, request_headers=request_headers, request_body=request_body, provider_request_headers=provider_request_headers or {}, provider_request_body=provider_request_body, response_headers=response_headers, client_response_headers=client_response_headers, response_body=response_body, client_response_body=client_response_body, request_id=self.request_id, # Provider 侧追踪信息(用于记录真实成本) provider_id=provider_id, provider_endpoint_id=provider_endpoint_id, provider_api_key_id=provider_api_key_id, # 模型映射信息 target_model=target_model, # Provider 响应元数据/请求元数据 metadata=metadata, ) total_cost = float(getattr(usage, "total_cost_usd", 0.0) or 0.0) if self.user and self.api_key: audit_service.log_api_request( db=self.db, user_id=self.user.id, api_key_id=self.api_key.id, request_id=self.request_id, model=model, provider=provider, success=True, ip_address=self.client_ip, status_code=status_code, input_tokens=getattr(usage, "input_tokens", input_tokens), output_tokens=getattr(usage, "output_tokens", output_tokens), cost_usd=total_cost, ) return total_cost async def record_failure( self, *, provider: str, model: str, response_time_ms: int, status_code: int, error_message: str, request_body: dict[str, Any], request_headers: dict[str, Any], is_stream: bool, api_format: str | None = None, api_family: str | None = None, endpoint_kind: str | None = None, provider_request_headers: dict[str, Any] | None = None, provider_request_body: Any | None = None, # 预估 token 信息(来自 message_start 事件,用于中断请求的成本估算) input_tokens: int = 0, output_tokens: int = 0, cache_creation_tokens: int = 0, cache_read_tokens: int = 0, cache_creation_tokens_5m: int = 0, cache_creation_tokens_1h: int = 0, response_body: dict[str, Any] | None = None, response_headers: dict[str, Any] | None = None, client_response_headers: dict[str, Any] | None = None, client_response_body: Any | None = None, # Provider 侧追踪信息(用于 curl 复现等场景) provider_id: str | None = None, provider_endpoint_id: str | None = None, provider_api_key_id: str | None = None, # 格式转换追踪 endpoint_api_format: str | None = None, has_format_conversion: bool = False, # 模型映射信息 target_model: str | None = None, # 请求元数据(用于性能与调试记录) request_metadata: dict[str, Any] | None = None, ) -> None: """ 记录失败请求 Args: input_tokens: 预估输入 tokens(来自 message_start,用于中断请求的成本估算) output_tokens: 预估输出 tokens(来自已收到的内容) cache_creation_tokens: 缓存创建 tokens cache_read_tokens: 缓存读取 tokens response_body: 响应体(如果有部分响应) response_headers: 响应头(Provider 返回的原始响应头) client_response_headers: 返回给客户端的响应头 target_model: 映射后的目标模型名(如果发生了映射) """ provider_name = provider or "unknown" if provider_name == "unknown": logger.warning( "[Telemetry] Recording failure with unknown provider (request_id={})", self.request_id, ) metadata = self._build_usage_metadata( request_metadata=request_metadata, ) await UsageService.record_usage( db=self.db, user=self.user, api_key=self.api_key, provider=provider_name, model=model, input_tokens=input_tokens, output_tokens=output_tokens, cache_creation_input_tokens=cache_creation_tokens, cache_read_input_tokens=cache_read_tokens, cache_creation_input_tokens_5m=cache_creation_tokens_5m, cache_creation_input_tokens_1h=cache_creation_tokens_1h, request_type="chat", api_format=api_format, api_family=api_family, endpoint_kind=endpoint_kind, endpoint_api_format=endpoint_api_format, has_format_conversion=has_format_conversion, is_stream=is_stream, response_time_ms=response_time_ms, status_code=status_code, error_message=error_message, request_headers=request_headers, request_body=request_body, provider_request_headers=provider_request_headers or {}, provider_request_body=provider_request_body, response_headers=response_headers or {}, client_response_headers=client_response_headers, response_body=response_body or {"error": error_message}, client_response_body=client_response_body, request_id=self.request_id, # Provider 侧追踪信息 provider_id=provider_id, provider_endpoint_id=provider_endpoint_id, provider_api_key_id=provider_api_key_id, # 模型映射信息 target_model=target_model, # 请求元数据 metadata=metadata, ) async def record_cancelled( self, *, provider: str, model: str, response_time_ms: int, first_byte_time_ms: int | None, status_code: int, request_body: dict[str, Any], request_headers: dict[str, Any], is_stream: bool, api_format: str | None = None, api_family: str | None = None, endpoint_kind: str | None = None, provider_request_headers: dict[str, Any] | None = None, provider_request_body: Any | None = None, input_tokens: int = 0, output_tokens: int = 0, cache_creation_tokens: int = 0, cache_read_tokens: int = 0, cache_creation_tokens_5m: int = 0, cache_creation_tokens_1h: int = 0, response_body: dict[str, Any] | None = None, response_headers: dict[str, Any] | None = None, client_response_headers: dict[str, Any] | None = None, client_response_body: Any | None = None, # Provider 侧追踪信息 provider_id: str | None = None, provider_endpoint_id: str | None = None, provider_api_key_id: str | None = None, # 格式转换追踪 endpoint_api_format: str | None = None, has_format_conversion: bool = False, target_model: str | None = None, # 请求元数据(用于性能与调试记录) request_metadata: dict[str, Any] | None = None, ) -> None: """ 记录客户端取消的请求 客户端主动断开连接不算系统失败,使用 cancelled 状态。 """ provider_name = provider or "unknown" metadata = self._build_usage_metadata( request_metadata=request_metadata, ) await UsageService.record_usage( db=self.db, user=self.user, api_key=self.api_key, provider=provider_name, model=model, input_tokens=input_tokens, output_tokens=output_tokens, cache_creation_input_tokens=cache_creation_tokens, cache_read_input_tokens=cache_read_tokens, cache_creation_input_tokens_5m=cache_creation_tokens_5m, cache_creation_input_tokens_1h=cache_creation_tokens_1h, request_type="chat", api_format=api_format, api_family=api_family, endpoint_kind=endpoint_kind, endpoint_api_format=endpoint_api_format, has_format_conversion=has_format_conversion, is_stream=is_stream, response_time_ms=response_time_ms, first_byte_time_ms=first_byte_time_ms, status_code=status_code, status="cancelled", request_headers=request_headers, request_body=request_body, provider_request_headers=provider_request_headers or {}, provider_request_body=provider_request_body, response_headers=response_headers or {}, client_response_headers=client_response_headers, response_body=response_body or {}, client_response_body=client_response_body, request_id=self.request_id, # Provider 侧追踪信息 provider_id=provider_id, provider_endpoint_id=provider_endpoint_id, provider_api_key_id=provider_api_key_id, target_model=target_model, metadata=metadata, )