mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: 移除 Python 后端源码，全面迁移至 Rust gateway 架构
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
fawney19
15
_deprecated_py_src/services/usage/__init__.py
Normal file
15
_deprecated_py_src/services/usage/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
使用量服务模块
|
||||
|
||||
包含使用量追踪、流式使用量、配额调度等功能。
|
||||
"""
|
||||
|
||||
from src.services.usage.quota_scheduler import QuotaScheduler
|
||||
from src.services.usage.service import UsageService
|
||||
from src.services.usage.stream import StreamUsageTracker
|
||||
|
||||
__all__ = [
|
||||
"UsageService",
|
||||
"StreamUsageTracker",
|
||||
"QuotaScheduler",
|
||||
]
|
||||
262
_deprecated_py_src/services/usage/_billing_integration.py
Normal file
262
_deprecated_py_src/services/usage/_billing_integration.py
Normal file
@@ -0,0 +1,262 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
|
||||
from src.services.usage._recording_helpers import (
|
||||
build_usage_params,
|
||||
deserialize_body_if_json,
|
||||
sanitize_request_metadata,
|
||||
)
|
||||
from src.services.usage._types import UsageCostInfo, UsageRecordParams
|
||||
|
||||
|
||||
class UsageBillingIntegrationMixin:
|
||||
"""计费集成方法 -- 准备用量记录的共享逻辑"""
|
||||
|
||||
@classmethod
|
||||
async def _prepare_usage_record(
|
||||
cls,
|
||||
params: UsageRecordParams,
|
||||
) -> tuple[dict[str, Any], float]:
|
||||
"""准备用量记录的共享逻辑
|
||||
|
||||
此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑：
|
||||
- 获取费率倍数
|
||||
- 计算成本
|
||||
- 构建 Usage 参数
|
||||
|
||||
Args:
|
||||
params: 用量记录参数数据类
|
||||
|
||||
Returns:
|
||||
(usage_params 字典, total_cost 总成本)
|
||||
"""
|
||||
# 计费口径以 Provider 为准（优先 endpoint_api_format）
|
||||
billing_api_format: str | None = None
|
||||
if params.endpoint_api_format:
|
||||
try:
|
||||
billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
|
||||
except Exception:
|
||||
billing_api_format = None
|
||||
if billing_api_format is None and params.api_format:
|
||||
try:
|
||||
billing_api_format = normalize_signature_key(str(params.api_format))
|
||||
except Exception:
|
||||
billing_api_format = None
|
||||
|
||||
input_tokens_for_billing = normalize_input_tokens_for_billing(
|
||||
billing_api_format,
|
||||
params.input_tokens,
|
||||
params.cache_read_input_tokens,
|
||||
)
|
||||
|
||||
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
|
||||
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
|
||||
params.db, params.provider_api_key_id, params.provider_id, billing_api_format
|
||||
)
|
||||
|
||||
metadata = dict(params.metadata or {})
|
||||
is_failed_request = params.status_code >= 400 or params.error_message is not None
|
||||
|
||||
# Helper: compute billing task_type (billing domain)
|
||||
billing_task_type = (params.request_type or "").lower()
|
||||
if billing_task_type not in {"chat", "cli", "video", "image", "audio"}:
|
||||
billing_task_type = "chat"
|
||||
|
||||
# 使用新计费系统计算费用
|
||||
from src.services.billing.service import BillingService
|
||||
|
||||
request_count = 0 if is_failed_request else 1
|
||||
has_cache_tokens = bool(
|
||||
params.cache_creation_input_tokens > 0 or params.cache_read_input_tokens > 0
|
||||
)
|
||||
effective_cache_ttl_minutes = params.cache_ttl_minutes
|
||||
|
||||
# 主链路很多场景不会显式传 cache_ttl_minutes，这里补全以确保 1h/5m TTL 差异化计价生效。
|
||||
if effective_cache_ttl_minutes is None and has_cache_tokens and params.provider_api_key_id:
|
||||
try:
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
key_ttl = (
|
||||
params.db.query(ProviderAPIKey.cache_ttl_minutes)
|
||||
.filter(ProviderAPIKey.id == params.provider_api_key_id)
|
||||
.scalar()
|
||||
)
|
||||
if key_ttl is not None:
|
||||
key_ttl_int = int(key_ttl)
|
||||
if key_ttl_int >= 0:
|
||||
effective_cache_ttl_minutes = key_ttl_int
|
||||
except Exception:
|
||||
# Best-effort fallback below.
|
||||
pass
|
||||
|
||||
# 无法从 key 获取时，尽量从 5m/1h 细分回推（主要覆盖 Claude cache_creation）。
|
||||
if effective_cache_ttl_minutes is None and has_cache_tokens:
|
||||
t5m = int(params.cache_creation_input_tokens_5m or 0)
|
||||
t1h = int(params.cache_creation_input_tokens_1h or 0)
|
||||
if t1h > 0 and t5m == 0:
|
||||
effective_cache_ttl_minutes = 60
|
||||
elif t5m > 0 and t1h == 0:
|
||||
effective_cache_ttl_minutes = 5
|
||||
elif t1h > 0:
|
||||
# 混合场景优先按长 TTL 计，避免 1h 缓存被按 5m 误计。
|
||||
effective_cache_ttl_minutes = 60
|
||||
|
||||
dims: dict[str, Any] = {
|
||||
"input_tokens": input_tokens_for_billing,
|
||||
"output_tokens": params.output_tokens,
|
||||
"cache_creation_input_tokens": params.cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": params.cache_read_input_tokens,
|
||||
"request_count": request_count,
|
||||
}
|
||||
if effective_cache_ttl_minutes is not None:
|
||||
dims["cache_ttl_minutes"] = effective_cache_ttl_minutes
|
||||
# If tiered pricing is disabled, force first tier by using tier-key=0.
|
||||
if not params.use_tiered_pricing:
|
||||
dims["total_input_context"] = 0
|
||||
|
||||
billing = BillingService(params.db)
|
||||
result = billing.calculate(
|
||||
task_type=billing_task_type,
|
||||
model=params.model,
|
||||
provider_id=params.provider_id or "",
|
||||
dimensions=dims,
|
||||
strict_mode=None,
|
||||
)
|
||||
snap = result.snapshot
|
||||
|
||||
breakdown = snap.cost_breakdown or {}
|
||||
input_cost = float(breakdown.get("input_cost", 0.0))
|
||||
output_cost = float(breakdown.get("output_cost", 0.0))
|
||||
cache_creation_cost = float(breakdown.get("cache_creation_cost", 0.0))
|
||||
cache_read_cost = float(breakdown.get("cache_read_cost", 0.0))
|
||||
request_cost = float(breakdown.get("request_cost", 0.0))
|
||||
cache_cost = cache_creation_cost + cache_read_cost
|
||||
total_cost = float(snap.total_cost or 0.0)
|
||||
|
||||
rv = snap.resolved_variables or {}
|
||||
|
||||
def _as_float(v: Any, d: float | None) -> float | None:
|
||||
try:
|
||||
if v is None:
|
||||
return d
|
||||
return float(v)
|
||||
except Exception:
|
||||
return d
|
||||
|
||||
input_price = _as_float(rv.get("input_price_per_1m"), 0.0) or 0.0
|
||||
output_price = _as_float(rv.get("output_price_per_1m"), 0.0) or 0.0
|
||||
cache_creation_price = _as_float(rv.get("cache_creation_price_per_1m"), None)
|
||||
cache_read_price = _as_float(rv.get("cache_read_price_per_1m"), None)
|
||||
request_price = _as_float(rv.get("price_per_request"), None)
|
||||
|
||||
# Audit snapshot (pruned later by sanitize_request_metadata)
|
||||
metadata["billing_snapshot"] = snap.to_dict()
|
||||
|
||||
# Best-effort prune metadata to reduce DB/memory pressure.
|
||||
metadata = sanitize_request_metadata(metadata)
|
||||
|
||||
# 构建 Usage 参数
|
||||
request_body = deserialize_body_if_json(params.request_body)
|
||||
provider_request_body = deserialize_body_if_json(params.provider_request_body)
|
||||
response_body = deserialize_body_if_json(params.response_body)
|
||||
client_response_body = deserialize_body_if_json(params.client_response_body)
|
||||
usage_params = build_usage_params(
|
||||
db=params.db,
|
||||
user=params.user,
|
||||
api_key=params.api_key,
|
||||
provider=params.provider,
|
||||
model=params.model,
|
||||
input_tokens=input_tokens_for_billing,
|
||||
output_tokens=params.output_tokens,
|
||||
cache_creation_input_tokens=params.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=params.cache_read_input_tokens,
|
||||
cache_creation_input_tokens_5m=params.cache_creation_input_tokens_5m,
|
||||
cache_creation_input_tokens_1h=params.cache_creation_input_tokens_1h,
|
||||
request_type=params.request_type,
|
||||
api_format=params.api_format,
|
||||
api_family=params.api_family,
|
||||
endpoint_kind=params.endpoint_kind,
|
||||
endpoint_api_format=params.endpoint_api_format,
|
||||
has_format_conversion=params.has_format_conversion,
|
||||
is_stream=params.is_stream,
|
||||
response_time_ms=params.response_time_ms,
|
||||
first_byte_time_ms=params.first_byte_time_ms,
|
||||
status_code=params.status_code,
|
||||
error_message=params.error_message,
|
||||
metadata=metadata,
|
||||
request_headers=params.request_headers,
|
||||
request_body=request_body,
|
||||
provider_request_headers=params.provider_request_headers,
|
||||
provider_request_body=provider_request_body,
|
||||
response_headers=params.response_headers,
|
||||
client_response_headers=params.client_response_headers,
|
||||
response_body=response_body,
|
||||
client_response_body=client_response_body,
|
||||
request_id=params.request_id,
|
||||
provider_id=params.provider_id,
|
||||
provider_endpoint_id=params.provider_endpoint_id,
|
||||
provider_api_key_id=params.provider_api_key_id,
|
||||
status=params.status,
|
||||
target_model=params.target_model,
|
||||
cost=UsageCostInfo(
|
||||
input_cost=input_cost,
|
||||
output_cost=output_cost,
|
||||
cache_creation_cost=cache_creation_cost,
|
||||
cache_read_cost=cache_read_cost,
|
||||
cache_cost=cache_cost,
|
||||
request_cost=request_cost,
|
||||
total_cost=total_cost,
|
||||
input_price=input_price,
|
||||
output_price=output_price,
|
||||
cache_creation_price=cache_creation_price,
|
||||
cache_read_price=cache_read_price,
|
||||
request_price=request_price,
|
||||
actual_rate_multiplier=actual_rate_multiplier,
|
||||
is_free_tier=is_free_tier,
|
||||
),
|
||||
)
|
||||
|
||||
return usage_params, total_cost
|
||||
|
||||
@classmethod
|
||||
async def _prepare_usage_records_batch(
|
||||
cls,
|
||||
params_list: list[UsageRecordParams],
|
||||
) -> list[tuple[dict[str, Any], float, Exception | None]]:
|
||||
"""批量并行准备用量记录（性能优化）
|
||||
|
||||
并行调用 _prepare_usage_record，提高批量处理效率。
|
||||
|
||||
Args:
|
||||
params_list: 用量记录参数列表
|
||||
|
||||
Returns:
|
||||
列表，每项为 (usage_params, total_cost, exception)
|
||||
如果处理成功，exception 为 None
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
async def prepare_single(
|
||||
params: UsageRecordParams,
|
||||
) -> tuple[dict[str, Any], float, Exception | None]:
|
||||
try:
|
||||
usage_params, total_cost = await cls._prepare_usage_record(params)
|
||||
return (usage_params, total_cost, None)
|
||||
except Exception as e:
|
||||
return ({}, 0.0, e)
|
||||
|
||||
if not params_list:
|
||||
return []
|
||||
|
||||
# 避免一次性创建过多 task（并且 _prepare_usage_record 内部也可能包含并行调用）
|
||||
# 这里采用分批 gather 来限制并发量。
|
||||
chunk_size = 50
|
||||
results: list[tuple[dict[str, Any], float, Exception | None]] = []
|
||||
for i in range(0, len(params_list), chunk_size):
|
||||
chunk = params_list[i : i + chunk_size]
|
||||
chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
|
||||
results.extend(chunk_results)
|
||||
return results
|
||||
397
_deprecated_py_src/services/usage/_recording_helpers.py
Normal file
397
_deprecated_py_src/services/usage/_recording_helpers.py
Normal file
@@ -0,0 +1,397 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.database import ApiKey, Usage, User
|
||||
from src.services.billing.precision import to_money_decimal
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.usage._types import UsageCostInfo
|
||||
from src.services.usage.error_classifier import classify_error
|
||||
|
||||
|
||||
def _parse_format_dimensions(api_format: str | None) -> tuple[str | None, str | None]:
|
||||
"""从 api_format (如 'claude:chat') 解析出 (api_family, endpoint_kind)"""
|
||||
if not api_format:
|
||||
return None, None
|
||||
parts = api_format.lower().split(":", 1)
|
||||
if len(parts) == 2:
|
||||
return parts[0], parts[1]
|
||||
return parts[0], None
|
||||
|
||||
|
||||
# Metadata pruning configuration (ordered by priority - drop first to last)
|
||||
METADATA_PRUNE_KEYS: tuple[str, ...] = (
|
||||
"raw_response_ref",
|
||||
"poll_raw_response",
|
||||
"trace",
|
||||
"debug",
|
||||
"dimensions",
|
||||
"provider_response_headers",
|
||||
"client_response_headers",
|
||||
)
|
||||
|
||||
# Keys to preserve even under aggressive pruning
|
||||
METADATA_KEEP_KEYS: frozenset[str] = frozenset(
|
||||
{
|
||||
"billing_snapshot",
|
||||
"billing_updated_at",
|
||||
"perf",
|
||||
"pool_summary",
|
||||
"scheduling_audit",
|
||||
"_metadata_truncated",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def deserialize_body_if_json(value: Any) -> Any:
|
||||
"""写库前按需反序列化 body JSON 字符串。
|
||||
|
||||
仅对 JSON object/array 字符串做 json.loads，其他值保持原样。
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
stripped = value.lstrip()
|
||||
if not stripped or stripped[0] not in "{[":
|
||||
return value
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return value
|
||||
if isinstance(parsed, (dict, list)):
|
||||
return parsed
|
||||
return value
|
||||
|
||||
|
||||
def build_usage_params(
|
||||
*,
|
||||
db: Session,
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
provider: str,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int,
|
||||
cache_read_input_tokens: int,
|
||||
cache_creation_input_tokens_5m: int = 0,
|
||||
cache_creation_input_tokens_1h: int = 0,
|
||||
request_type: str,
|
||||
api_format: str | None,
|
||||
api_family: str | None = None,
|
||||
endpoint_kind: str | None = None,
|
||||
endpoint_api_format: str | None,
|
||||
provider_api_family: str | None = None,
|
||||
provider_endpoint_kind: str | None = None,
|
||||
has_format_conversion: bool,
|
||||
is_stream: bool,
|
||||
response_time_ms: int | None,
|
||||
first_byte_time_ms: int | None,
|
||||
status_code: int,
|
||||
error_message: str | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
request_headers: dict[str, Any] | None,
|
||||
request_body: Any | None,
|
||||
provider_request_headers: dict[str, Any] | None,
|
||||
provider_request_body: Any | None,
|
||||
response_headers: dict[str, Any] | None,
|
||||
client_response_headers: dict[str, Any] | None,
|
||||
response_body: Any | None,
|
||||
client_response_body: Any | None,
|
||||
request_id: str,
|
||||
provider_id: str | None,
|
||||
provider_endpoint_id: str | None,
|
||||
provider_api_key_id: str | None,
|
||||
status: str,
|
||||
target_model: str | None,
|
||||
cost: UsageCostInfo,
|
||||
) -> dict[str, Any]:
|
||||
"""构建 Usage 记录的参数字典（内部方法，避免代码重复）"""
|
||||
|
||||
# 展开成本信息
|
||||
input_cost = cost.input_cost
|
||||
output_cost = cost.output_cost
|
||||
cache_creation_cost = cost.cache_creation_cost
|
||||
cache_read_cost = cost.cache_read_cost
|
||||
cache_cost = cost.cache_cost
|
||||
request_cost = cost.request_cost
|
||||
total_cost = cost.total_cost
|
||||
input_price = cost.input_price
|
||||
output_price = cost.output_price
|
||||
cache_creation_price = cost.cache_creation_price
|
||||
cache_read_price = cost.cache_read_price
|
||||
request_price = cost.request_price
|
||||
actual_rate_multiplier = cost.actual_rate_multiplier
|
||||
is_free_tier = cost.is_free_tier
|
||||
|
||||
# 根据配置决定是否记录请求详情
|
||||
should_log_headers = SystemConfigService.should_log_headers(db)
|
||||
should_log_body = SystemConfigService.should_log_body(db)
|
||||
|
||||
# 处理请求头（可能需要脱敏）
|
||||
processed_request_headers = None
|
||||
if should_log_headers and request_headers is not None:
|
||||
processed_request_headers = SystemConfigService.mask_sensitive_headers(db, request_headers)
|
||||
|
||||
# 处理提供商请求头（可能需要脱敏）
|
||||
processed_provider_request_headers = None
|
||||
if should_log_headers and provider_request_headers is not None:
|
||||
processed_provider_request_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, provider_request_headers
|
||||
)
|
||||
|
||||
# 处理请求体和响应体（可能需要截断）
|
||||
processed_request_body = None
|
||||
processed_provider_request_body = None
|
||||
processed_response_body = None
|
||||
processed_client_response_body = None
|
||||
if should_log_body:
|
||||
if request_body is not None:
|
||||
processed_request_body = SystemConfigService.truncate_body(
|
||||
db, request_body, is_request=True
|
||||
)
|
||||
if provider_request_body is not None:
|
||||
processed_provider_request_body = SystemConfigService.truncate_body(
|
||||
db, provider_request_body, is_request=True
|
||||
)
|
||||
if response_body is not None:
|
||||
processed_response_body = SystemConfigService.truncate_body(
|
||||
db, response_body, is_request=False
|
||||
)
|
||||
if client_response_body is not None:
|
||||
processed_client_response_body = SystemConfigService.truncate_body(
|
||||
db, client_response_body, is_request=False
|
||||
)
|
||||
|
||||
# 处理响应头
|
||||
processed_response_headers = None
|
||||
if should_log_headers and response_headers is not None:
|
||||
processed_response_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, response_headers
|
||||
)
|
||||
|
||||
# 处理返回给客户端的响应头
|
||||
processed_client_response_headers = None
|
||||
if should_log_headers and client_response_headers is not None:
|
||||
processed_client_response_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, client_response_headers
|
||||
)
|
||||
|
||||
# 计算真实成本（表面成本 * 倍率），免费套餐实际费用为 0
|
||||
if is_free_tier:
|
||||
actual_input_cost = 0.0
|
||||
actual_output_cost = 0.0
|
||||
actual_cache_creation_cost = 0.0
|
||||
actual_cache_read_cost = 0.0
|
||||
actual_request_cost = 0.0
|
||||
actual_total_cost = 0.0
|
||||
else:
|
||||
actual_input_cost = input_cost * actual_rate_multiplier
|
||||
actual_output_cost = output_cost * actual_rate_multiplier
|
||||
actual_cache_creation_cost = cache_creation_cost * actual_rate_multiplier
|
||||
actual_cache_read_cost = cache_read_cost * actual_rate_multiplier
|
||||
actual_request_cost = request_cost * actual_rate_multiplier
|
||||
actual_total_cost = total_cost * actual_rate_multiplier
|
||||
|
||||
error_category = None
|
||||
if status_code >= 400 or error_message or status in {"failed", "cancelled"}:
|
||||
error_category = classify_error(status_code, error_message, status).value
|
||||
|
||||
# 从 api_format / endpoint_api_format 解析 api_family + endpoint_kind
|
||||
# 优先使用透传值，fallback 到字符串解析
|
||||
parsed_family, parsed_kind = _parse_format_dimensions(api_format)
|
||||
client_family = api_family or parsed_family
|
||||
client_kind = endpoint_kind or parsed_kind
|
||||
|
||||
parsed_ep_family, parsed_ep_kind = _parse_format_dimensions(endpoint_api_format)
|
||||
ep_family = provider_api_family or parsed_ep_family
|
||||
ep_kind = provider_endpoint_kind or parsed_ep_kind
|
||||
|
||||
return {
|
||||
"user_id": user.id if user else None,
|
||||
"api_key_id": api_key.id if api_key else None,
|
||||
"username": user.username if user else None,
|
||||
"api_key_name": api_key.name if api_key else None,
|
||||
"request_id": request_id,
|
||||
"provider_name": provider,
|
||||
"model": model,
|
||||
"target_model": target_model,
|
||||
"provider_id": provider_id,
|
||||
"provider_endpoint_id": provider_endpoint_id,
|
||||
"provider_api_key_id": provider_api_key_id,
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": input_tokens + output_tokens,
|
||||
"cache_creation_input_tokens": cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": cache_read_input_tokens,
|
||||
"cache_creation_input_tokens_5m": cache_creation_input_tokens_5m,
|
||||
"cache_creation_input_tokens_1h": cache_creation_input_tokens_1h,
|
||||
"input_cost_usd": to_money_decimal(input_cost),
|
||||
"output_cost_usd": to_money_decimal(output_cost),
|
||||
"cache_cost_usd": to_money_decimal(cache_cost),
|
||||
"cache_creation_cost_usd": to_money_decimal(cache_creation_cost),
|
||||
"cache_read_cost_usd": to_money_decimal(cache_read_cost),
|
||||
"request_cost_usd": to_money_decimal(request_cost),
|
||||
"total_cost_usd": to_money_decimal(total_cost),
|
||||
"actual_input_cost_usd": to_money_decimal(actual_input_cost),
|
||||
"actual_output_cost_usd": to_money_decimal(actual_output_cost),
|
||||
"actual_cache_creation_cost_usd": to_money_decimal(actual_cache_creation_cost),
|
||||
"actual_cache_read_cost_usd": to_money_decimal(actual_cache_read_cost),
|
||||
"actual_request_cost_usd": to_money_decimal(actual_request_cost),
|
||||
"actual_total_cost_usd": to_money_decimal(actual_total_cost),
|
||||
"rate_multiplier": actual_rate_multiplier,
|
||||
"input_price_per_1m": input_price,
|
||||
"output_price_per_1m": output_price,
|
||||
"cache_creation_price_per_1m": cache_creation_price,
|
||||
"cache_read_price_per_1m": cache_read_price,
|
||||
"price_per_request": request_price,
|
||||
"request_type": request_type,
|
||||
"api_format": api_format,
|
||||
"api_family": client_family,
|
||||
"endpoint_kind": client_kind,
|
||||
"endpoint_api_format": endpoint_api_format,
|
||||
"provider_api_family": ep_family,
|
||||
"provider_endpoint_kind": ep_kind,
|
||||
"has_format_conversion": has_format_conversion,
|
||||
"is_stream": is_stream,
|
||||
"status_code": status_code,
|
||||
"error_message": error_message,
|
||||
"error_category": error_category,
|
||||
"response_time_ms": response_time_ms,
|
||||
"first_byte_time_ms": first_byte_time_ms,
|
||||
"status": status,
|
||||
"request_metadata": metadata,
|
||||
"request_headers": processed_request_headers,
|
||||
"request_body": processed_request_body,
|
||||
"provider_request_headers": processed_provider_request_headers,
|
||||
"provider_request_body": processed_provider_request_body,
|
||||
"response_headers": processed_response_headers,
|
||||
"client_response_headers": processed_client_response_headers,
|
||||
"response_body": processed_response_body,
|
||||
"client_response_body": processed_client_response_body,
|
||||
}
|
||||
|
||||
|
||||
def update_existing_usage(
|
||||
existing_usage: Usage,
|
||||
usage_params: dict[str, Any],
|
||||
target_model: str | None,
|
||||
) -> None:
|
||||
"""更新已存在的 Usage 记录（内部方法）"""
|
||||
# 更新关键字段
|
||||
existing_usage.provider_name = usage_params["provider_name"]
|
||||
existing_usage.model = usage_params["model"]
|
||||
existing_usage.request_type = usage_params["request_type"]
|
||||
existing_usage.api_format = usage_params["api_format"]
|
||||
existing_usage.api_family = usage_params.get("api_family")
|
||||
existing_usage.endpoint_kind = usage_params.get("endpoint_kind")
|
||||
existing_usage.endpoint_api_format = usage_params["endpoint_api_format"]
|
||||
existing_usage.provider_api_family = usage_params.get("provider_api_family")
|
||||
existing_usage.provider_endpoint_kind = usage_params.get("provider_endpoint_kind")
|
||||
existing_usage.has_format_conversion = usage_params["has_format_conversion"]
|
||||
existing_usage.is_stream = usage_params["is_stream"]
|
||||
existing_usage.status = usage_params["status"]
|
||||
existing_usage.status_code = usage_params["status_code"]
|
||||
existing_usage.error_message = usage_params["error_message"]
|
||||
existing_usage.error_category = usage_params.get("error_category")
|
||||
existing_usage.response_time_ms = usage_params["response_time_ms"]
|
||||
existing_usage.first_byte_time_ms = usage_params["first_byte_time_ms"]
|
||||
|
||||
# 更新请求头和请求体（如果有新值）
|
||||
if usage_params["request_headers"] is not None:
|
||||
existing_usage.request_headers = usage_params["request_headers"]
|
||||
if usage_params["request_body"] is not None:
|
||||
existing_usage.request_body = usage_params["request_body"]
|
||||
if usage_params["provider_request_headers"] is not None:
|
||||
existing_usage.provider_request_headers = usage_params["provider_request_headers"]
|
||||
if usage_params["provider_request_body"] is not None:
|
||||
existing_usage.provider_request_body = usage_params["provider_request_body"]
|
||||
existing_usage.response_body = usage_params["response_body"]
|
||||
existing_usage.response_headers = usage_params["response_headers"]
|
||||
existing_usage.client_response_headers = usage_params["client_response_headers"]
|
||||
existing_usage.client_response_body = usage_params["client_response_body"]
|
||||
|
||||
# 更新 token 和费用信息
|
||||
existing_usage.input_tokens = usage_params["input_tokens"]
|
||||
existing_usage.output_tokens = usage_params["output_tokens"]
|
||||
existing_usage.total_tokens = usage_params["total_tokens"]
|
||||
existing_usage.cache_creation_input_tokens = usage_params["cache_creation_input_tokens"]
|
||||
existing_usage.cache_read_input_tokens = usage_params["cache_read_input_tokens"]
|
||||
existing_usage.cache_creation_input_tokens_5m = usage_params.get(
|
||||
"cache_creation_input_tokens_5m", 0
|
||||
)
|
||||
existing_usage.cache_creation_input_tokens_1h = usage_params.get(
|
||||
"cache_creation_input_tokens_1h", 0
|
||||
)
|
||||
existing_usage.input_cost_usd = usage_params["input_cost_usd"]
|
||||
existing_usage.output_cost_usd = usage_params["output_cost_usd"]
|
||||
existing_usage.cache_cost_usd = usage_params["cache_cost_usd"]
|
||||
existing_usage.cache_creation_cost_usd = usage_params["cache_creation_cost_usd"]
|
||||
existing_usage.cache_read_cost_usd = usage_params["cache_read_cost_usd"]
|
||||
existing_usage.request_cost_usd = usage_params["request_cost_usd"]
|
||||
existing_usage.total_cost_usd = usage_params["total_cost_usd"]
|
||||
existing_usage.actual_input_cost_usd = usage_params["actual_input_cost_usd"]
|
||||
existing_usage.actual_output_cost_usd = usage_params["actual_output_cost_usd"]
|
||||
existing_usage.actual_cache_creation_cost_usd = usage_params["actual_cache_creation_cost_usd"]
|
||||
existing_usage.actual_cache_read_cost_usd = usage_params["actual_cache_read_cost_usd"]
|
||||
existing_usage.actual_request_cost_usd = usage_params["actual_request_cost_usd"]
|
||||
existing_usage.actual_total_cost_usd = usage_params["actual_total_cost_usd"]
|
||||
existing_usage.rate_multiplier = usage_params["rate_multiplier"]
|
||||
|
||||
# 更新 Provider 侧追踪信息（仅在有新值时更新，避免覆盖已有数据）
|
||||
if usage_params.get("provider_id"):
|
||||
existing_usage.provider_id = usage_params["provider_id"]
|
||||
if usage_params.get("provider_endpoint_id"):
|
||||
existing_usage.provider_endpoint_id = usage_params["provider_endpoint_id"]
|
||||
if usage_params.get("provider_api_key_id"):
|
||||
existing_usage.provider_api_key_id = usage_params["provider_api_key_id"]
|
||||
|
||||
# 更新元数据（如 billing_snapshot/dimensions 等）
|
||||
if usage_params.get("request_metadata") is not None:
|
||||
existing_usage.request_metadata = usage_params["request_metadata"]
|
||||
|
||||
# 更新模型映射信息
|
||||
if target_model is not None:
|
||||
existing_usage.target_model = target_model
|
||||
|
||||
|
||||
def sanitize_request_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Best-effort metadata pruning to reduce DB/CPU/memory pressure.
|
||||
|
||||
This is called right before persisting Usage rows (or updating request_metadata).
|
||||
Pruning order is defined by `METADATA_PRUNE_KEYS` (first key is dropped first).
|
||||
"""
|
||||
if not isinstance(metadata, dict) or not metadata:
|
||||
return {}
|
||||
|
||||
from src.config.settings import config
|
||||
|
||||
# Enforce global metadata size limit (best-effort)
|
||||
max_bytes = int(getattr(config, "usage_metadata_max_bytes", 0) or 0)
|
||||
if max_bytes <= 0:
|
||||
return metadata
|
||||
|
||||
def _size(d: dict[str, Any]) -> int:
|
||||
try:
|
||||
return len(json.dumps(d, ensure_ascii=False, default=str))
|
||||
except Exception:
|
||||
return len(str(d))
|
||||
|
||||
if _size(metadata) <= max_bytes:
|
||||
return metadata
|
||||
|
||||
# Progressive pruning (configurable order)
|
||||
metadata["_metadata_truncated"] = True
|
||||
|
||||
for k in METADATA_PRUNE_KEYS:
|
||||
if k in metadata:
|
||||
metadata.pop(k, None)
|
||||
if _size(metadata) <= max_bytes:
|
||||
return metadata
|
||||
|
||||
# Fallback: keep only billing-related metadata
|
||||
reduced = {k: metadata.get(k) for k in METADATA_KEEP_KEYS if k in metadata}
|
||||
return reduced
|
||||
110
_deprecated_py_src/services/usage/_types.py
Normal file
110
_deprecated_py_src/services/usage/_types.py
Normal file
@@ -0,0 +1,110 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.database import ApiKey, User
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageRecordParams:
|
||||
"""用量记录参数数据类，用于在内部方法间传递数据"""
|
||||
|
||||
db: Session
|
||||
user: User | None
|
||||
api_key: ApiKey | None
|
||||
provider: str
|
||||
model: str
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
request_type: str
|
||||
api_format: str | None
|
||||
api_family: str | None # 协议族（从 Adapter 层透传）
|
||||
endpoint_kind: str | None # 端点类型（从 Adapter 层透传）
|
||||
endpoint_api_format: str | None # 端点原生 API 格式
|
||||
has_format_conversion: bool # 是否发生了格式转换
|
||||
is_stream: bool
|
||||
response_time_ms: int | None
|
||||
first_byte_time_ms: int | None
|
||||
status_code: int
|
||||
error_message: str | None
|
||||
metadata: dict[str, Any] | None
|
||||
request_headers: dict[str, Any] | None
|
||||
request_body: Any | None
|
||||
provider_request_headers: dict[str, Any] | None
|
||||
provider_request_body: Any | None
|
||||
response_headers: dict[str, Any] | None
|
||||
client_response_headers: dict[str, Any] | None
|
||||
response_body: Any | None
|
||||
client_response_body: Any | None
|
||||
request_id: str
|
||||
provider_id: str | None
|
||||
provider_endpoint_id: str | None
|
||||
provider_api_key_id: str | None
|
||||
status: str
|
||||
cache_ttl_minutes: int | None
|
||||
use_tiered_pricing: bool
|
||||
target_model: str | None
|
||||
cache_creation_input_tokens_5m: int = 0
|
||||
cache_creation_input_tokens_1h: int = 0
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""验证关键字段，确保数据完整性"""
|
||||
# Token 数量不能为负数
|
||||
if self.input_tokens < 0:
|
||||
raise ValueError(f"input_tokens 不能为负数: {self.input_tokens}")
|
||||
if self.output_tokens < 0:
|
||||
raise ValueError(f"output_tokens 不能为负数: {self.output_tokens}")
|
||||
if self.cache_creation_input_tokens < 0:
|
||||
raise ValueError(
|
||||
f"cache_creation_input_tokens 不能为负数: {self.cache_creation_input_tokens}"
|
||||
)
|
||||
if self.cache_read_input_tokens < 0:
|
||||
raise ValueError(f"cache_read_input_tokens 不能为负数: {self.cache_read_input_tokens}")
|
||||
|
||||
# 响应时间不能为负数
|
||||
if self.response_time_ms is not None and self.response_time_ms < 0:
|
||||
raise ValueError(f"response_time_ms 不能为负数: {self.response_time_ms}")
|
||||
if self.first_byte_time_ms is not None and self.first_byte_time_ms < 0:
|
||||
raise ValueError(f"first_byte_time_ms 不能为负数: {self.first_byte_time_ms}")
|
||||
|
||||
# HTTP 状态码范围校验
|
||||
if not (100 <= self.status_code <= 599):
|
||||
raise ValueError(f"无效的 HTTP 状态码: {self.status_code}")
|
||||
|
||||
# 状态值校验
|
||||
# - pending: 请求已创建，等待处理
|
||||
# - streaming: 流式响应进行中
|
||||
# - completed: 请求成功完成
|
||||
# - failed: 请求失败（上游错误、超时等）
|
||||
# - cancelled: 客户端主动断开连接
|
||||
valid_statuses = {"pending", "streaming", "completed", "failed", "cancelled"}
|
||||
if self.status not in valid_statuses:
|
||||
raise ValueError(f"无效的状态值: {self.status}，有效值: {valid_statuses}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageCostInfo:
|
||||
"""成本与价格信息，用于 _build_usage_params 参数封装"""
|
||||
|
||||
# 成本计算结果
|
||||
input_cost: float = 0.0
|
||||
output_cost: float = 0.0
|
||||
cache_creation_cost: float = 0.0
|
||||
cache_read_cost: float = 0.0
|
||||
cache_cost: float = 0.0
|
||||
request_cost: float = 0.0
|
||||
total_cost: float = 0.0
|
||||
# 价格信息
|
||||
input_price: float | None = None
|
||||
output_price: float | None = None
|
||||
cache_creation_price: float | None = None
|
||||
cache_read_price: float | None = None
|
||||
request_price: float | None = None
|
||||
# 倍率
|
||||
actual_rate_multiplier: float = 1.0
|
||||
is_free_tier: bool = False
|
||||
433
_deprecated_py_src/services/usage/active_requests.py
Normal file
433
_deprecated_py_src/services/usage/active_requests.py
Normal file
@@ -0,0 +1,433 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.api_format.metadata import can_passthrough_endpoint
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
from src.core.logger import logger
|
||||
from src.models.database import RequestCandidate, Usage
|
||||
|
||||
|
||||
class UsageActiveRequestsMixin:
|
||||
"""活跃请求管理方法"""
|
||||
|
||||
@staticmethod
|
||||
def _find_completed_request_ids(
|
||||
db: Session,
|
||||
request_ids: list[str],
|
||||
) -> set[str]:
|
||||
"""查询已成功完成的 request_id 集合
|
||||
|
||||
通过 RequestCandidate 表判断哪些请求实际已成功完成：
|
||||
1. status='success' 且 stream_completed=True（正常完成）
|
||||
2. status='streaming'（Provider 已返回成功响应头，流因重启中断）
|
||||
"""
|
||||
if not request_ids:
|
||||
return set()
|
||||
|
||||
from sqlalchemy import or_
|
||||
|
||||
candidates = (
|
||||
db.query(
|
||||
RequestCandidate.request_id,
|
||||
RequestCandidate.status,
|
||||
RequestCandidate.extra_data,
|
||||
)
|
||||
.filter(
|
||||
RequestCandidate.request_id.in_(request_ids),
|
||||
or_(
|
||||
RequestCandidate.status == "success",
|
||||
RequestCandidate.status == "streaming",
|
||||
),
|
||||
)
|
||||
.all()
|
||||
)
|
||||
completed: set[str] = set()
|
||||
for c in candidates:
|
||||
extra_data = c.extra_data or {}
|
||||
if c.status == "success" and extra_data.get("stream_completed", False):
|
||||
completed.add(c.request_id)
|
||||
elif c.status == "streaming":
|
||||
completed.add(c.request_id)
|
||||
return completed
|
||||
|
||||
@staticmethod
|
||||
def _sync_candidate_status_to_success(
|
||||
db: Session,
|
||||
request_ids: list[str],
|
||||
now: datetime | None = None,
|
||||
) -> None:
|
||||
"""将指定请求的 streaming candidate 同步更新为 success"""
|
||||
if not request_ids:
|
||||
return
|
||||
if now is None:
|
||||
now = datetime.now(timezone.utc)
|
||||
db.query(RequestCandidate).filter(
|
||||
RequestCandidate.request_id.in_(request_ids),
|
||||
RequestCandidate.status == "streaming",
|
||||
).update(
|
||||
{"status": "success", "finished_at": now},
|
||||
synchronize_session=False,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_active_requests(
|
||||
cls,
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
limit: int = 50,
|
||||
) -> list[Usage]:
|
||||
"""
|
||||
获取活跃的请求（pending 或 streaming 状态）
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID（可选，用于过滤）
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
活跃请求的 Usage 列表
|
||||
"""
|
||||
query = db.query(Usage).filter(Usage.status.in_(["pending", "streaming"]))
|
||||
|
||||
if user_id:
|
||||
query = query.filter(Usage.user_id == user_id)
|
||||
|
||||
return query.order_by(Usage.created_at.desc()).limit(limit).all()
|
||||
|
||||
@classmethod
|
||||
def cleanup_stale_pending_requests(
|
||||
cls,
|
||||
db: Session,
|
||||
timeout_minutes: int = 10,
|
||||
batch_size: int = 200,
|
||||
) -> int:
|
||||
"""
|
||||
清理超时的 pending/streaming 请求
|
||||
|
||||
将超过指定时间仍处于 pending 或 streaming 状态的请求标记为 failed 或恢复为 completed。
|
||||
会检查 RequestCandidate 表，如果 Provider 已返回成功响应（status=streaming 或 stream_completed），
|
||||
则恢复为 completed 而非标记为 failed，同时同步更新 candidate 状态。
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
timeout_minutes: 超时时间（分钟），默认 10 分钟
|
||||
batch_size: 每次处理的记录数，限制在 1-200 之间
|
||||
|
||||
Returns:
|
||||
清理的记录数
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
cutoff_time = now - timedelta(minutes=timeout_minutes)
|
||||
|
||||
batch_size = max(1, batch_size)
|
||||
failed_count = 0
|
||||
recovered_count = 0
|
||||
|
||||
while True:
|
||||
stale_requests = (
|
||||
db.query(Usage.id, Usage.request_id, Usage.status, Usage.billing_status)
|
||||
.filter(
|
||||
Usage.status.in_(["pending", "streaming"]),
|
||||
Usage.created_at < cutoff_time,
|
||||
)
|
||||
.order_by(Usage.created_at.asc(), Usage.id.asc())
|
||||
.limit(batch_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not stale_requests:
|
||||
break
|
||||
|
||||
stale_request_ids = [request_id for _, request_id, _, _ in stale_requests if request_id]
|
||||
completed_request_ids = cls._find_completed_request_ids(db, stale_request_ids)
|
||||
|
||||
usage_updates = []
|
||||
failed_request_ids: list[str] = []
|
||||
|
||||
for usage_id, request_id, old_status, billing_status in stale_requests:
|
||||
if request_id and request_id in completed_request_ids:
|
||||
usage_updates.append(
|
||||
{
|
||||
"id": usage_id,
|
||||
"status": "completed",
|
||||
"status_code": 200,
|
||||
"error_message": None,
|
||||
}
|
||||
)
|
||||
recovered_count += 1
|
||||
else:
|
||||
entry: dict[str, Any] = {
|
||||
"id": usage_id,
|
||||
"status": "failed",
|
||||
"status_code": 504,
|
||||
"error_message": (
|
||||
f"请求超时: 状态 '{old_status}' 超过 {timeout_minutes} 分钟未完成"
|
||||
),
|
||||
}
|
||||
if billing_status == "pending":
|
||||
entry["billing_status"] = "void"
|
||||
entry["finalized_at"] = now
|
||||
entry["total_cost_usd"] = 0.0
|
||||
entry["request_cost_usd"] = 0.0
|
||||
entry["actual_total_cost_usd"] = 0.0
|
||||
entry["actual_request_cost_usd"] = 0.0
|
||||
usage_updates.append(entry)
|
||||
failed_count += 1
|
||||
if request_id:
|
||||
failed_request_ids.append(request_id)
|
||||
|
||||
if usage_updates:
|
||||
db.bulk_update_mappings(Usage, usage_updates)
|
||||
|
||||
cls._sync_candidate_status_to_success(db, list(completed_request_ids), now)
|
||||
|
||||
if failed_request_ids:
|
||||
db.query(RequestCandidate).filter(
|
||||
RequestCandidate.request_id.in_(failed_request_ids),
|
||||
RequestCandidate.status.in_(["streaming", "pending"]),
|
||||
).update(
|
||||
{
|
||||
"status": "failed",
|
||||
"finished_at": now,
|
||||
"error_message": "请求超时（服务器可能已重启）",
|
||||
},
|
||||
synchronize_session=False,
|
||||
)
|
||||
|
||||
db.commit()
|
||||
db.expunge_all()
|
||||
|
||||
total = failed_count + recovered_count
|
||||
if total > 0:
|
||||
parts = []
|
||||
if failed_count:
|
||||
parts.append(f"{failed_count} 条标记为 failed")
|
||||
if recovered_count:
|
||||
parts.append(f"{recovered_count} 条恢复为 completed")
|
||||
logger.info(
|
||||
f"清理超时请求: 超过 {timeout_minutes} 分钟的 pending/streaming 请求 - "
|
||||
+ ", ".join(parts)
|
||||
)
|
||||
|
||||
return total
|
||||
|
||||
@classmethod
|
||||
def get_stale_pending_count(
|
||||
cls,
|
||||
db: Session,
|
||||
timeout_minutes: int = 10,
|
||||
) -> int:
|
||||
"""
|
||||
获取超时的 pending/streaming 请求数量（用于监控）
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
timeout_minutes: 超时时间（分钟）
|
||||
|
||||
Returns:
|
||||
超时请求数量
|
||||
"""
|
||||
cutoff_time = datetime.now(timezone.utc) - timedelta(minutes=timeout_minutes)
|
||||
|
||||
return int(
|
||||
db.query(func.count(Usage.id))
|
||||
.filter(
|
||||
Usage.status.in_(["pending", "streaming"]),
|
||||
Usage.created_at < cutoff_time,
|
||||
)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_active_requests_status(
|
||||
cls,
|
||||
db: Session,
|
||||
ids: list[str] | None = None,
|
||||
user_id: str | None = None,
|
||||
default_timeout_seconds: int = 300,
|
||||
*,
|
||||
include_admin_fields: bool = False,
|
||||
maintain_status: bool | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
获取活跃请求状态（用于前端轮询）。
|
||||
|
||||
与 get_active_requests 不同，此方法：
|
||||
1. 返回轻量级的状态字典而非完整 Usage 对象
|
||||
2. 可选地检测并清理超时的 pending/streaming 请求
|
||||
3. 支持按 ID 列表查询特定请求
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
ids: 指定要查询的请求 ID 列表（可选）
|
||||
user_id: 限制只查询该用户的请求（可选，用于普通用户接口）
|
||||
default_timeout_seconds: 默认超时时间（秒），当端点未配置时使用
|
||||
maintain_status: 是否执行超时修复与状态回写；默认仅在全量活跃请求查询时执行
|
||||
|
||||
Returns:
|
||||
请求状态列表
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# 构建基础查询
|
||||
query = db.query(
|
||||
Usage.id,
|
||||
Usage.status,
|
||||
Usage.input_tokens,
|
||||
Usage.output_tokens,
|
||||
Usage.cache_creation_input_tokens,
|
||||
Usage.cache_read_input_tokens,
|
||||
Usage.total_cost_usd,
|
||||
Usage.actual_total_cost_usd,
|
||||
Usage.rate_multiplier,
|
||||
Usage.response_time_ms,
|
||||
Usage.first_byte_time_ms, # 首字时间 (TTFB)
|
||||
Usage.created_at,
|
||||
Usage.provider_endpoint_id,
|
||||
# API 格式 / 格式转换（streaming 状态时已可确定）
|
||||
Usage.api_format,
|
||||
Usage.endpoint_api_format,
|
||||
Usage.has_format_conversion,
|
||||
# 模型映射（streaming 时已可确定）
|
||||
Usage.target_model,
|
||||
)
|
||||
|
||||
# 管理员轮询：可附带 provider 与上游 key 名称（注意：不要在普通用户接口暴露上游 key 信息）
|
||||
if include_admin_fields:
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
query = query.add_columns(
|
||||
Usage.provider_name,
|
||||
ProviderAPIKey.name.label("api_key_name"),
|
||||
).outerjoin(ProviderAPIKey, Usage.provider_api_key_id == ProviderAPIKey.id)
|
||||
|
||||
if ids:
|
||||
query = query.filter(Usage.id.in_(ids))
|
||||
if user_id:
|
||||
query = query.filter(Usage.user_id == user_id)
|
||||
else:
|
||||
# 查询所有活跃请求
|
||||
query = query.filter(Usage.status.in_(["pending", "streaming"]))
|
||||
if user_id:
|
||||
query = query.filter(Usage.user_id == user_id)
|
||||
query = query.order_by(Usage.created_at.desc()).limit(50)
|
||||
|
||||
records = query.all()
|
||||
should_maintain_status = maintain_status if maintain_status is not None else not ids
|
||||
|
||||
# 检查超时的 pending/streaming 请求
|
||||
# 收集可能超时的 usage_id 列表
|
||||
timeout_candidates: list[str] = []
|
||||
if should_maintain_status:
|
||||
for r in records:
|
||||
if r.status in ("pending", "streaming") and r.created_at:
|
||||
timeout_seconds = default_timeout_seconds
|
||||
|
||||
created_at = r.created_at
|
||||
if created_at.tzinfo is None:
|
||||
created_at = created_at.replace(tzinfo=timezone.utc)
|
||||
elapsed = (now - created_at).total_seconds()
|
||||
if elapsed > timeout_seconds:
|
||||
timeout_candidates.append(r.id)
|
||||
|
||||
# 批量更新超时的请求（排除已有成功完成记录的请求）
|
||||
timeout_ids = []
|
||||
if should_maintain_status and timeout_candidates:
|
||||
# 先获取这些 Usage 的 request_id
|
||||
usage_request_ids = (
|
||||
db.query(Usage.id, Usage.request_id).filter(Usage.id.in_(timeout_candidates)).all()
|
||||
)
|
||||
usage_id_to_request_id = {u.id: u.request_id for u in usage_request_ids}
|
||||
request_id_to_usage_id = {u.request_id: u.id for u in usage_request_ids}
|
||||
request_ids = list(request_id_to_usage_id.keys())
|
||||
|
||||
# 查询已成功完成的 request_id
|
||||
completed_rids = cls._find_completed_request_ids(db, request_ids)
|
||||
completed_usage_ids = {
|
||||
request_id_to_usage_id[rid]
|
||||
for rid in completed_rids
|
||||
if rid in request_id_to_usage_id
|
||||
}
|
||||
|
||||
# 只对没有成功完成记录的请求标记超时
|
||||
timeout_ids = [uid for uid in timeout_candidates if uid not in completed_usage_ids]
|
||||
|
||||
if timeout_ids:
|
||||
db.query(Usage).filter(Usage.id.in_(timeout_ids)).update(
|
||||
{"status": "failed", "error_message": "请求超时（服务器可能已重启）"},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.commit()
|
||||
|
||||
# 对于已完成但状态未更新的请求，主动恢复状态为 completed
|
||||
if completed_usage_ids:
|
||||
db.query(Usage).filter(Usage.id.in_(list(completed_usage_ids))).update(
|
||||
{"status": "completed"},
|
||||
synchronize_session=False,
|
||||
)
|
||||
# 同步更新 candidate 状态：streaming -> success
|
||||
completed_request_ids = [
|
||||
usage_id_to_request_id[uid]
|
||||
for uid in completed_usage_ids
|
||||
if uid in usage_id_to_request_id
|
||||
]
|
||||
cls._sync_candidate_status_to_success(db, completed_request_ids)
|
||||
db.commit()
|
||||
logger.info(
|
||||
"[Usage] 恢复 {} 个已完成请求的状态（遥测回调丢失）",
|
||||
len(completed_usage_ids),
|
||||
)
|
||||
|
||||
result: list[dict[str, Any]] = []
|
||||
for r in records:
|
||||
api_format = getattr(r, "api_format", None)
|
||||
endpoint_api_format = getattr(r, "endpoint_api_format", None)
|
||||
has_format_conversion = getattr(r, "has_format_conversion", None)
|
||||
|
||||
# 兼容历史数据：当 streaming 状态已拿到两个格式但 has_format_conversion 为空时，回填推断结果
|
||||
if has_format_conversion is None and api_format and endpoint_api_format:
|
||||
client_raw = str(api_format).strip()
|
||||
endpoint_raw = str(endpoint_api_format).strip()
|
||||
if ":" in client_raw and ":" in endpoint_raw:
|
||||
client_fmt = normalize_signature_key(client_raw)
|
||||
endpoint_fmt = normalize_signature_key(endpoint_raw)
|
||||
has_format_conversion = not can_passthrough_endpoint(client_fmt, endpoint_fmt)
|
||||
|
||||
item: dict[str, Any] = {
|
||||
"id": r.id,
|
||||
"status": "failed" if r.id in timeout_ids else r.status,
|
||||
"input_tokens": r.input_tokens,
|
||||
"output_tokens": r.output_tokens,
|
||||
"cache_creation_input_tokens": r.cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": r.cache_read_input_tokens,
|
||||
"cost": float(r.total_cost_usd) if r.total_cost_usd else 0,
|
||||
"actual_cost": (
|
||||
float(r.actual_total_cost_usd) if r.actual_total_cost_usd is not None else None
|
||||
),
|
||||
"rate_multiplier": (
|
||||
float(r.rate_multiplier) if r.rate_multiplier is not None else None
|
||||
),
|
||||
"response_time_ms": r.response_time_ms,
|
||||
"first_byte_time_ms": r.first_byte_time_ms, # 首字时间 (TTFB)
|
||||
}
|
||||
if api_format:
|
||||
item["api_format"] = api_format
|
||||
if endpoint_api_format:
|
||||
item["endpoint_api_format"] = endpoint_api_format
|
||||
if has_format_conversion is not None:
|
||||
item["has_format_conversion"] = bool(has_format_conversion)
|
||||
# 模型映射（streaming 时已可确定）
|
||||
if r.target_model:
|
||||
item["target_model"] = r.target_model
|
||||
if include_admin_fields:
|
||||
item["provider"] = r.provider_name
|
||||
item["api_key_name"] = r.api_key_name
|
||||
result.append(item)
|
||||
|
||||
return result
|
||||
529
_deprecated_py_src/services/usage/cache_analysis.py
Normal file
529
_deprecated_py_src/services/usage/cache_analysis.py
Normal file
@@ -0,0 +1,529 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.database import Usage, User
|
||||
|
||||
|
||||
class UsageCacheAnalysisMixin:
|
||||
"""缓存分析方法"""
|
||||
|
||||
@staticmethod
|
||||
def analyze_cache_affinity_ttl(
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
hours: int = 168,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
分析用户请求间隔分布，推荐合适的缓存亲和性 TTL
|
||||
|
||||
通过分析同一用户连续请求之间的时间间隔，判断用户的使用模式：
|
||||
- 高频用户（间隔短）：5 分钟 TTL 足够
|
||||
- 中频用户：15-30 分钟 TTL
|
||||
- 低频用户（间隔长）：需要 60 分钟 TTL
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 指定用户 ID（可选，为空则分析所有用户）
|
||||
api_key_id: 指定 API Key ID（可选）
|
||||
hours: 分析最近多少小时的数据
|
||||
|
||||
Returns:
|
||||
包含分析结果的字典
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
# 计算时间范围
|
||||
start_date = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
|
||||
# 构建 SQL 查询 - 使用窗口函数计算请求间隔
|
||||
# 按 user_id 或 api_key_id 分组，计算同一组内连续请求的时间差
|
||||
group_by_field = "api_key_id" if api_key_id else "user_id"
|
||||
|
||||
# 构建过滤条件
|
||||
filter_clause = ""
|
||||
if user_id or api_key_id:
|
||||
filter_clause = f"AND {group_by_field} = :filter_id"
|
||||
|
||||
sql = text(f"""
|
||||
WITH user_requests AS (
|
||||
SELECT
|
||||
{group_by_field} as group_id,
|
||||
created_at,
|
||||
LAG(created_at) OVER (
|
||||
PARTITION BY {group_by_field}
|
||||
ORDER BY created_at
|
||||
) as prev_request_at
|
||||
FROM usage
|
||||
WHERE status = 'completed'
|
||||
AND created_at > :start_date
|
||||
AND {group_by_field} IS NOT NULL
|
||||
{filter_clause}
|
||||
),
|
||||
intervals AS (
|
||||
SELECT
|
||||
group_id,
|
||||
EXTRACT(EPOCH FROM (created_at - prev_request_at)) / 60.0 as interval_minutes
|
||||
FROM user_requests
|
||||
WHERE prev_request_at IS NOT NULL
|
||||
),
|
||||
user_stats AS (
|
||||
SELECT
|
||||
group_id,
|
||||
COUNT(*) as request_count,
|
||||
COUNT(*) FILTER (WHERE interval_minutes <= 5) as within_5min,
|
||||
COUNT(*) FILTER (WHERE interval_minutes > 5 AND interval_minutes <= 15) as within_15min,
|
||||
COUNT(*) FILTER (WHERE interval_minutes > 15 AND interval_minutes <= 30) as within_30min,
|
||||
COUNT(*) FILTER (WHERE interval_minutes > 30 AND interval_minutes <= 60) as within_60min,
|
||||
COUNT(*) FILTER (WHERE interval_minutes > 60) as over_60min,
|
||||
PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY interval_minutes) as median_interval,
|
||||
PERCENTILE_CONT(0.75) WITHIN GROUP (ORDER BY interval_minutes) as p75_interval,
|
||||
PERCENTILE_CONT(0.90) WITHIN GROUP (ORDER BY interval_minutes) as p90_interval,
|
||||
AVG(interval_minutes) as avg_interval,
|
||||
MIN(interval_minutes) as min_interval,
|
||||
MAX(interval_minutes) as max_interval
|
||||
FROM intervals
|
||||
GROUP BY group_id
|
||||
HAVING COUNT(*) >= 2
|
||||
)
|
||||
SELECT * FROM user_stats
|
||||
ORDER BY request_count DESC
|
||||
""")
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"start_date": start_date,
|
||||
}
|
||||
if user_id:
|
||||
params["filter_id"] = user_id
|
||||
elif api_key_id:
|
||||
params["filter_id"] = api_key_id
|
||||
|
||||
result = db.execute(sql, params)
|
||||
rows = result.fetchall()
|
||||
|
||||
# 收集所有 user_id 以便批量查询用户信息
|
||||
group_ids = [row[0] for row in rows]
|
||||
|
||||
# 如果是按 user_id 分组，查询用户信息
|
||||
user_info_map: dict[str, dict[str, str]] = {}
|
||||
if group_by_field == "user_id" and group_ids:
|
||||
users = db.query(User).filter(User.id.in_(group_ids)).all()
|
||||
for user in users:
|
||||
user_info_map[str(user.id)] = {
|
||||
"username": str(user.username),
|
||||
"email": str(user.email) if user.email else "",
|
||||
}
|
||||
|
||||
# 处理结果
|
||||
users_analysis = []
|
||||
for row in rows:
|
||||
# row 是一个 tuple，按查询顺序访问
|
||||
(
|
||||
group_id,
|
||||
request_count,
|
||||
within_5min,
|
||||
within_15min,
|
||||
within_30min,
|
||||
within_60min,
|
||||
over_60min,
|
||||
median_interval,
|
||||
p75_interval,
|
||||
p90_interval,
|
||||
avg_interval,
|
||||
min_interval,
|
||||
max_interval,
|
||||
) = row
|
||||
|
||||
# 计算推荐 TTL
|
||||
recommended_ttl = UsageCacheAnalysisMixin._calculate_recommended_ttl(
|
||||
p75_interval, p90_interval
|
||||
)
|
||||
|
||||
# 获取用户信息
|
||||
user_info = user_info_map.get(str(group_id), {})
|
||||
|
||||
# 计算各区间占比
|
||||
total_intervals = request_count
|
||||
users_analysis.append(
|
||||
{
|
||||
"group_id": group_id,
|
||||
"username": user_info.get("username"),
|
||||
"email": user_info.get("email"),
|
||||
"request_count": request_count,
|
||||
"interval_distribution": {
|
||||
"within_5min": within_5min,
|
||||
"within_15min": within_15min,
|
||||
"within_30min": within_30min,
|
||||
"within_60min": within_60min,
|
||||
"over_60min": over_60min,
|
||||
},
|
||||
"interval_percentages": {
|
||||
"within_5min": round(within_5min / total_intervals * 100, 1),
|
||||
"within_15min": round(within_15min / total_intervals * 100, 1),
|
||||
"within_30min": round(within_30min / total_intervals * 100, 1),
|
||||
"within_60min": round(within_60min / total_intervals * 100, 1),
|
||||
"over_60min": round(over_60min / total_intervals * 100, 1),
|
||||
},
|
||||
"percentiles": {
|
||||
"p50": round(float(median_interval), 2) if median_interval else None,
|
||||
"p75": round(float(p75_interval), 2) if p75_interval else None,
|
||||
"p90": round(float(p90_interval), 2) if p90_interval else None,
|
||||
},
|
||||
"avg_interval_minutes": (
|
||||
round(float(avg_interval), 2) if avg_interval else None
|
||||
),
|
||||
"min_interval_minutes": (
|
||||
round(float(min_interval), 2) if min_interval else None
|
||||
),
|
||||
"max_interval_minutes": (
|
||||
round(float(max_interval), 2) if max_interval else None
|
||||
),
|
||||
"recommended_ttl_minutes": recommended_ttl,
|
||||
"recommendation_reason": UsageCacheAnalysisMixin._get_ttl_recommendation_reason(
|
||||
recommended_ttl, p75_interval, p90_interval
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
# 汇总统计
|
||||
ttl_distribution = {"5min": 0, "15min": 0, "30min": 0, "60min": 0}
|
||||
for analysis in users_analysis:
|
||||
ttl = analysis["recommended_ttl_minutes"]
|
||||
if ttl <= 5:
|
||||
ttl_distribution["5min"] += 1
|
||||
elif ttl <= 15:
|
||||
ttl_distribution["15min"] += 1
|
||||
elif ttl <= 30:
|
||||
ttl_distribution["30min"] += 1
|
||||
else:
|
||||
ttl_distribution["60min"] += 1
|
||||
|
||||
return {
|
||||
"analysis_period_hours": hours,
|
||||
"total_users_analyzed": len(users_analysis),
|
||||
"ttl_distribution": ttl_distribution,
|
||||
"users": users_analysis,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _calculate_recommended_ttl(
|
||||
p75_interval: float | None,
|
||||
p90_interval: float | None,
|
||||
) -> int:
|
||||
"""
|
||||
根据请求间隔分布计算推荐的缓存 TTL
|
||||
|
||||
策略：
|
||||
- 如果 90% 的请求间隔都在 5 分钟内 -> 5 分钟 TTL
|
||||
- 如果 75% 的请求间隔在 15 分钟内 -> 15 分钟 TTL
|
||||
- 如果 75% 的请求间隔在 30 分钟内 -> 30 分钟 TTL
|
||||
- 否则 -> 60 分钟 TTL
|
||||
"""
|
||||
if p90_interval is None or p75_interval is None:
|
||||
return 5 # 默认值
|
||||
|
||||
# 如果 90% 的间隔都在 5 分钟内
|
||||
if p90_interval <= 5:
|
||||
return 5
|
||||
|
||||
# 如果 75% 的间隔在 15 分钟内
|
||||
if p75_interval <= 15:
|
||||
return 15
|
||||
|
||||
# 如果 75% 的间隔在 30 分钟内
|
||||
if p75_interval <= 30:
|
||||
return 30
|
||||
|
||||
# 低频用户，需要更长的 TTL
|
||||
return 60
|
||||
|
||||
@staticmethod
|
||||
def _get_ttl_recommendation_reason(
|
||||
ttl: int,
|
||||
p75_interval: float | None,
|
||||
p90_interval: float | None,
|
||||
) -> str:
|
||||
"""生成 TTL 推荐理由"""
|
||||
if p75_interval is None or p90_interval is None:
|
||||
return "数据不足，使用默认值"
|
||||
|
||||
if ttl == 5:
|
||||
return f"高频用户：90% 的请求间隔在 {p90_interval:.1f} 分钟内"
|
||||
elif ttl == 15:
|
||||
return f"中高频用户：75% 的请求间隔在 {p75_interval:.1f} 分钟内"
|
||||
elif ttl == 30:
|
||||
return f"中频用户：75% 的请求间隔在 {p75_interval:.1f} 分钟内"
|
||||
else:
|
||||
return f"低频用户：75% 的请求间隔为 {p75_interval:.1f} 分钟，建议使用长 TTL"
|
||||
|
||||
@staticmethod
|
||||
def get_cache_hit_analysis(
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
hours: int = 168,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
分析缓存命中情况
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 指定用户 ID（可选）
|
||||
api_key_id: 指定 API Key ID（可选）
|
||||
hours: 分析最近多少小时的数据
|
||||
|
||||
Returns:
|
||||
缓存命中分析结果
|
||||
"""
|
||||
start_date = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
|
||||
# 基础查询
|
||||
query = db.query(
|
||||
func.count(Usage.id).label("total_requests"),
|
||||
func.sum(Usage.input_tokens).label("total_input_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("total_cache_read_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("total_cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_cost_usd).label("total_cache_read_cost"),
|
||||
func.sum(Usage.cache_creation_cost_usd).label("total_cache_creation_cost"),
|
||||
).filter(
|
||||
Usage.status == "completed",
|
||||
Usage.created_at >= start_date,
|
||||
)
|
||||
|
||||
if user_id:
|
||||
query = query.filter(Usage.user_id == user_id)
|
||||
if api_key_id:
|
||||
query = query.filter(Usage.api_key_id == api_key_id)
|
||||
|
||||
result = query.first()
|
||||
|
||||
if result is None:
|
||||
total_requests = 0
|
||||
total_input_tokens = 0
|
||||
total_cache_read_tokens = 0
|
||||
total_cache_creation_tokens = 0
|
||||
total_cache_read_cost = 0.0
|
||||
total_cache_creation_cost = 0.0
|
||||
else:
|
||||
total_requests = result.total_requests or 0
|
||||
total_input_tokens = result.total_input_tokens or 0
|
||||
total_cache_read_tokens = result.total_cache_read_tokens or 0
|
||||
total_cache_creation_tokens = result.total_cache_creation_tokens or 0
|
||||
total_cache_read_cost = float(result.total_cache_read_cost or 0)
|
||||
total_cache_creation_cost = float(result.total_cache_creation_cost or 0)
|
||||
|
||||
# 计算缓存命中率（按 token 数）
|
||||
# 总输入上下文 = input_tokens + cache_read_tokens（因为 input_tokens 不含 cache_read）
|
||||
# 或者如果 input_tokens 已经包含 cache_read，则直接用 input_tokens
|
||||
# 这里假设 cache_read_tokens 是额外的，命中率 = cache_read / (input + cache_read)
|
||||
total_context_tokens = total_input_tokens + total_cache_read_tokens
|
||||
cache_hit_rate = 0.0
|
||||
if total_context_tokens > 0:
|
||||
cache_hit_rate = total_cache_read_tokens / total_context_tokens * 100
|
||||
|
||||
# 计算节省的费用
|
||||
# 缓存读取价格是正常输入价格的 10%，所以节省了 90%
|
||||
# 节省 = cache_read_tokens * (正常价格 - 缓存价格) = cache_read_cost * 9
|
||||
# 因为 cache_read_cost 是按 10% 价格算的，如果按 100% 算就是 10 倍
|
||||
estimated_savings = total_cache_read_cost * 9 # 节省了 90%
|
||||
|
||||
# 统计有缓存命中的请求数
|
||||
requests_with_cache_hit = db.query(func.count(Usage.id)).filter(
|
||||
Usage.status == "completed",
|
||||
Usage.created_at >= start_date,
|
||||
Usage.cache_read_input_tokens > 0,
|
||||
)
|
||||
if user_id:
|
||||
requests_with_cache_hit = requests_with_cache_hit.filter(Usage.user_id == user_id)
|
||||
if api_key_id:
|
||||
requests_with_cache_hit = requests_with_cache_hit.filter(Usage.api_key_id == api_key_id)
|
||||
requests_with_cache_hit_count = int(requests_with_cache_hit.scalar() or 0)
|
||||
|
||||
return {
|
||||
"analysis_period_hours": hours,
|
||||
"total_requests": total_requests,
|
||||
"requests_with_cache_hit": requests_with_cache_hit_count,
|
||||
"request_cache_hit_rate": (
|
||||
round(requests_with_cache_hit_count / total_requests * 100, 2)
|
||||
if total_requests > 0
|
||||
else 0
|
||||
),
|
||||
"total_input_tokens": total_input_tokens,
|
||||
"total_cache_read_tokens": total_cache_read_tokens,
|
||||
"total_cache_creation_tokens": total_cache_creation_tokens,
|
||||
"token_cache_hit_rate": round(cache_hit_rate, 2),
|
||||
"total_cache_read_cost_usd": round(total_cache_read_cost, 4),
|
||||
"total_cache_creation_cost_usd": round(total_cache_creation_cost, 4),
|
||||
"estimated_savings_usd": round(estimated_savings, 4),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_interval_timeline(
|
||||
db: Session,
|
||||
hours: int = 24,
|
||||
limit: int = 10000,
|
||||
user_id: str | None = None,
|
||||
include_user_info: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取请求间隔时间线数据，用于散点图展示
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
hours: 分析最近多少小时的数据（默认24小时）
|
||||
limit: 最大返回数据点数量（默认10000）
|
||||
user_id: 指定用户 ID（可选，为空则返回所有用户）
|
||||
include_user_info: 是否包含用户信息（用于管理员多用户视图）
|
||||
|
||||
Returns:
|
||||
包含时间线数据点的字典，每个数据点包含 model 字段用于按模型区分颜色
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
start_date = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
|
||||
# 构建用户过滤条件
|
||||
user_filter = "AND u.user_id = :user_id" if user_id else ""
|
||||
|
||||
# 根据是否需要用户信息选择不同的查询
|
||||
if include_user_info and not user_id:
|
||||
# 管理员视图：返回带用户信息的数据点
|
||||
# 使用按比例采样，保持每个用户的数据量比例不变
|
||||
sql = text(f"""
|
||||
WITH request_intervals AS (
|
||||
SELECT
|
||||
u.created_at,
|
||||
u.user_id,
|
||||
u.model,
|
||||
usr.username,
|
||||
LAG(u.created_at) OVER (
|
||||
PARTITION BY u.user_id
|
||||
ORDER BY u.created_at
|
||||
) as prev_request_at
|
||||
FROM usage u
|
||||
LEFT JOIN users usr ON u.user_id = usr.id
|
||||
WHERE u.status = 'completed'
|
||||
AND u.created_at > :start_date
|
||||
AND u.user_id IS NOT NULL
|
||||
{user_filter}
|
||||
),
|
||||
filtered_intervals AS (
|
||||
SELECT
|
||||
created_at,
|
||||
user_id,
|
||||
model,
|
||||
username,
|
||||
EXTRACT(EPOCH FROM (created_at - prev_request_at)) / 60.0 as interval_minutes,
|
||||
ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY created_at) as rn
|
||||
FROM request_intervals
|
||||
WHERE prev_request_at IS NOT NULL
|
||||
AND EXTRACT(EPOCH FROM (created_at - prev_request_at)) / 60.0 <= 120
|
||||
),
|
||||
total_count AS (
|
||||
SELECT COUNT(*) as cnt FROM filtered_intervals
|
||||
),
|
||||
user_totals AS (
|
||||
SELECT user_id, COUNT(*) as user_cnt FROM filtered_intervals GROUP BY user_id
|
||||
),
|
||||
user_limits AS (
|
||||
SELECT
|
||||
ut.user_id,
|
||||
CASE WHEN tc.cnt <= :limit THEN ut.user_cnt
|
||||
ELSE GREATEST(CEIL(ut.user_cnt::float * :limit / tc.cnt), 1)::int
|
||||
END as user_limit
|
||||
FROM user_totals ut, total_count tc
|
||||
)
|
||||
SELECT
|
||||
fi.created_at,
|
||||
fi.user_id,
|
||||
fi.model,
|
||||
fi.username,
|
||||
fi.interval_minutes
|
||||
FROM filtered_intervals fi
|
||||
JOIN user_limits ul ON fi.user_id = ul.user_id
|
||||
WHERE fi.rn <= ul.user_limit
|
||||
ORDER BY fi.created_at
|
||||
""")
|
||||
else:
|
||||
# 普通视图：返回时间、间隔和模型信息
|
||||
sql = text(f"""
|
||||
WITH request_intervals AS (
|
||||
SELECT
|
||||
u.created_at,
|
||||
u.user_id,
|
||||
u.model,
|
||||
LAG(u.created_at) OVER (
|
||||
PARTITION BY u.user_id
|
||||
ORDER BY u.created_at
|
||||
) as prev_request_at
|
||||
FROM usage u
|
||||
WHERE u.status = 'completed'
|
||||
AND u.created_at > :start_date
|
||||
AND u.user_id IS NOT NULL
|
||||
{user_filter}
|
||||
)
|
||||
SELECT
|
||||
created_at,
|
||||
model,
|
||||
EXTRACT(EPOCH FROM (created_at - prev_request_at)) / 60.0 as interval_minutes
|
||||
FROM request_intervals
|
||||
WHERE prev_request_at IS NOT NULL
|
||||
AND EXTRACT(EPOCH FROM (created_at - prev_request_at)) / 60.0 <= 120
|
||||
ORDER BY created_at
|
||||
LIMIT :limit
|
||||
""")
|
||||
|
||||
params: dict[str, Any] = {"start_date": start_date, "limit": limit}
|
||||
if user_id:
|
||||
params["user_id"] = user_id
|
||||
|
||||
result = db.execute(sql, params)
|
||||
rows = result.fetchall()
|
||||
|
||||
# 转换为时间线数据点
|
||||
points = []
|
||||
users_map: dict[str, str] = {} # user_id -> username
|
||||
models_set: set = set() # 收集所有出现的模型
|
||||
|
||||
if include_user_info and not user_id:
|
||||
for row in rows:
|
||||
created_at, row_user_id, model, username, interval_minutes = row
|
||||
point_data: dict[str, Any] = {
|
||||
"x": created_at.isoformat(),
|
||||
"y": round(float(interval_minutes), 2),
|
||||
"user_id": str(row_user_id),
|
||||
}
|
||||
if model:
|
||||
point_data["model"] = model
|
||||
models_set.add(model)
|
||||
points.append(point_data)
|
||||
if row_user_id and username:
|
||||
users_map[str(row_user_id)] = username
|
||||
else:
|
||||
for row in rows:
|
||||
created_at, model, interval_minutes = row
|
||||
point_data = {"x": created_at.isoformat(), "y": round(float(interval_minutes), 2)}
|
||||
if model:
|
||||
point_data["model"] = model
|
||||
models_set.add(model)
|
||||
points.append(point_data)
|
||||
|
||||
response: dict[str, Any] = {
|
||||
"analysis_period_hours": hours,
|
||||
"total_points": len(points),
|
||||
"points": points,
|
||||
}
|
||||
|
||||
if include_user_info and not user_id:
|
||||
response["users"] = users_map
|
||||
|
||||
# 如果有模型信息，返回模型列表
|
||||
if models_set:
|
||||
response["models"] = sorted(models_set)
|
||||
|
||||
return response
|
||||
555
_deprecated_py_src/services/usage/consumer_streams.py
Normal file
555
_deprecated_py_src/services/usage/consumer_streams.py
Normal file
@@ -0,0 +1,555 @@
|
||||
"""
|
||||
Usage Redis Streams consumer.
|
||||
|
||||
高性能消费者实现，支持批量处理和单次提交多条记录。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import ResponseError
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from src.clients.redis_client import get_usage_queue_redis_client as get_redis_client
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.database.database import create_session
|
||||
from src.services.usage.events import UsageEvent, UsageEventType
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
|
||||
def _consumer_name() -> str:
|
||||
host = socket.gethostname() or "unknown"
|
||||
return f"{host}:{os.getpid()}"
|
||||
|
||||
|
||||
def _parse_body(value: Any) -> Any:
|
||||
"""消费者阶段保留原始 body，反序列化延迟到写库阶段。"""
|
||||
return value
|
||||
|
||||
|
||||
def _event_to_record(event: UsageEvent) -> dict[str, Any]:
|
||||
"""将 UsageEvent 转换为 record_usage_batch 所需的字典格式"""
|
||||
data = event.data
|
||||
status = "completed"
|
||||
if event.event_type == UsageEventType.FAILED:
|
||||
status = "failed"
|
||||
elif event.event_type == UsageEventType.CANCELLED:
|
||||
status = "cancelled"
|
||||
|
||||
finalized_at = None
|
||||
if event.timestamp_ms > 0:
|
||||
finalized_at = datetime.fromtimestamp(event.timestamp_ms / 1000, tz=timezone.utc)
|
||||
|
||||
return {
|
||||
"request_id": event.request_id,
|
||||
"user_id": data.get("user_id"),
|
||||
"api_key_id": data.get("api_key_id"),
|
||||
"provider": data.get("provider") or "unknown",
|
||||
"model": data.get("model") or "unknown",
|
||||
"input_tokens": data.get("input_tokens") or 0,
|
||||
"output_tokens": data.get("output_tokens") or 0,
|
||||
"cache_creation_input_tokens": data.get("cache_creation_input_tokens") or 0,
|
||||
"cache_read_input_tokens": data.get("cache_read_input_tokens") or 0,
|
||||
"cache_creation_input_tokens_5m": data.get("cache_creation_input_tokens_5m") or 0,
|
||||
"cache_creation_input_tokens_1h": data.get("cache_creation_input_tokens_1h") or 0,
|
||||
"request_type": data.get("request_type") or "chat",
|
||||
"api_format": data.get("api_format"),
|
||||
"api_family": data.get("api_family"),
|
||||
"endpoint_kind": data.get("endpoint_kind"),
|
||||
"endpoint_api_format": data.get("endpoint_api_format"),
|
||||
"has_format_conversion": data.get("has_format_conversion"),
|
||||
"is_stream": data.get("is_stream", True),
|
||||
"response_time_ms": data.get("response_time_ms"),
|
||||
"first_byte_time_ms": data.get("first_byte_time_ms"),
|
||||
"status_code": data.get("status_code") or 200,
|
||||
"error_message": data.get("error_message"),
|
||||
"metadata": data.get("metadata"),
|
||||
"request_headers": data.get("request_headers"),
|
||||
"request_body": _parse_body(data.get("request_body")),
|
||||
"provider_request_headers": data.get("provider_request_headers"),
|
||||
"provider_request_body": _parse_body(data.get("provider_request_body")),
|
||||
"response_headers": data.get("response_headers"),
|
||||
"client_response_headers": data.get("client_response_headers"),
|
||||
"response_body": _parse_body(data.get("response_body")),
|
||||
"client_response_body": _parse_body(data.get("client_response_body")),
|
||||
"provider_id": data.get("provider_id"),
|
||||
"provider_endpoint_id": data.get("provider_endpoint_id"),
|
||||
"provider_api_key_id": data.get("provider_api_key_id"),
|
||||
"status": status,
|
||||
"target_model": data.get("target_model"),
|
||||
"finalized_at": finalized_at,
|
||||
}
|
||||
|
||||
|
||||
async def ensure_usage_stream_group() -> None:
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if not redis_client:
|
||||
return
|
||||
try:
|
||||
await redis_client.xgroup_create(
|
||||
config.usage_queue_stream_key,
|
||||
config.usage_queue_stream_group,
|
||||
id="0-0",
|
||||
mkstream=True,
|
||||
)
|
||||
logger.info("[usage-queue] Created consumer group {}", config.usage_queue_stream_group)
|
||||
except ResponseError as exc:
|
||||
if "BUSYGROUP" in str(exc):
|
||||
return
|
||||
raise
|
||||
|
||||
|
||||
class UsageQueueConsumer:
|
||||
"""Usage 队列消费者
|
||||
|
||||
性能优化：
|
||||
- 缓存配置值避免重复属性访问
|
||||
- STREAMING 事件使用 pipeline 批量 ACK
|
||||
- 记录事件批量写入数据库
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._running = False
|
||||
self._task: asyncio.Task | None = None
|
||||
self._consumer = _consumer_name()
|
||||
self._last_claim = 0.0
|
||||
self._last_metrics_log = 0.0
|
||||
# 缓存配置值，避免热路径上的属性访问开销
|
||||
self._stream_key = config.usage_queue_stream_key
|
||||
self._stream_group = config.usage_queue_stream_group
|
||||
self._batch_size = config.usage_queue_consumer_batch
|
||||
self._block_ms = config.usage_queue_consumer_block_ms
|
||||
self._claim_idle_ms = config.usage_queue_claim_idle_ms
|
||||
self._claim_interval = config.usage_queue_claim_interval_seconds
|
||||
self._max_retries = config.usage_queue_max_retries
|
||||
self._dlq_key = config.usage_queue_dlq_key
|
||||
self._dlq_maxlen = config.usage_queue_dlq_maxlen
|
||||
self._metrics_interval = config.usage_queue_metrics_interval_seconds
|
||||
# 清理长期闲置的旧 consumer，避免 Redis consumer group 元数据持续累积。
|
||||
# 仅清理 pending=0 且空闲时间足够长的 consumer，不影响正常重投递。
|
||||
self._stale_consumer_idle_ms = max(self._claim_idle_ms * 10, 60 * 60 * 1000)
|
||||
|
||||
@staticmethod
|
||||
def _is_duplicate_key_error(exc: IntegrityError) -> bool:
|
||||
"""判断是否为重复键错误（唯一约束冲突）"""
|
||||
err_str = str(exc).lower()
|
||||
return "unique" in err_str or "duplicate" in err_str
|
||||
|
||||
async def _record_usage_batch(self, records: list[dict[str, Any]]) -> None:
|
||||
"""批量写库。
|
||||
|
||||
record_usage_batch 内部包含 async 准备阶段（费率查询等），
|
||||
必须在当前事件循环中 await，不能用 asyncio.run 在子线程创建新循环，
|
||||
否则会导致 Redis 连接泄漏（每次 asyncio.run 都会在 _redis_by_loop 中
|
||||
注册一个短命循环的连接，且永远不会被清理）。
|
||||
"""
|
||||
db = create_session()
|
||||
try:
|
||||
await UsageService.record_usage_batch(db, records)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._running:
|
||||
return
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if redis_client:
|
||||
await self._cleanup_stale_consumers(redis_client)
|
||||
self._running = True
|
||||
self._task = asyncio.create_task(self._run(), name="usage-queue-consumer")
|
||||
logger.info("[usage-queue] Consumer started: {}", self._consumer)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if not self._running:
|
||||
return
|
||||
self._running = False
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
try:
|
||||
await self._task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if redis_client:
|
||||
await self._delete_consumer(redis_client, self._consumer)
|
||||
logger.info("[usage-queue] Consumer stopped: {}", self._consumer)
|
||||
|
||||
async def _delete_consumer(self, redis_client: Any, consumer_name: str) -> None:
|
||||
try:
|
||||
await redis_client.xgroup_delconsumer(
|
||||
self._stream_key,
|
||||
self._stream_group,
|
||||
consumer_name,
|
||||
)
|
||||
except ResponseError as exc:
|
||||
# Group/stream may already be gone during shutdown; ignore in that case.
|
||||
if "NOGROUP" in str(exc) or "ERR no such key" in str(exc):
|
||||
return
|
||||
logger.debug("[usage-queue] DELCONSUMER failed for {}: {}", consumer_name, exc)
|
||||
except Exception as exc:
|
||||
logger.debug("[usage-queue] DELCONSUMER failed for {}: {}", consumer_name, exc)
|
||||
|
||||
async def _cleanup_stale_consumers(self, redis_client: Any) -> None:
|
||||
try:
|
||||
consumers = await redis_client.xinfo_consumers(
|
||||
self._stream_key,
|
||||
self._stream_group,
|
||||
)
|
||||
except ResponseError as exc:
|
||||
if "NOGROUP" in str(exc):
|
||||
return
|
||||
logger.debug("[usage-queue] XINFO CONSUMERS failed: {}", exc)
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.debug("[usage-queue] XINFO CONSUMERS failed: {}", exc)
|
||||
return
|
||||
|
||||
deleted = 0
|
||||
for consumer in consumers or []:
|
||||
if not isinstance(consumer, dict):
|
||||
continue
|
||||
consumer_name = str(consumer.get("name") or "").strip()
|
||||
if not consumer_name or consumer_name == self._consumer:
|
||||
continue
|
||||
|
||||
pending = int(consumer.get("pending", 0) or 0)
|
||||
idle_ms = int(consumer.get("idle", 0) or 0)
|
||||
if pending > 0 or idle_ms < self._stale_consumer_idle_ms:
|
||||
continue
|
||||
|
||||
await self._delete_consumer(redis_client, consumer_name)
|
||||
deleted += 1
|
||||
|
||||
if deleted:
|
||||
logger.info(
|
||||
"[usage-queue] Cleaned up {} stale consumers from group {}",
|
||||
deleted,
|
||||
self._stream_group,
|
||||
)
|
||||
|
||||
async def _ack_and_delete_messages(self, redis_client: Any, message_ids: list[str]) -> None:
|
||||
"""ACK messages and immediately delete them from the main stream.
|
||||
|
||||
usage:events is only meant to be a short-lived buffer. Once an event is
|
||||
successfully persisted (or moved to DLQ), keeping it in Redis only
|
||||
retains duplicate history and inflates memory.
|
||||
"""
|
||||
if not message_ids:
|
||||
return
|
||||
|
||||
pipe = redis_client.pipeline()
|
||||
for message_id in message_ids:
|
||||
pipe.xack(self._stream_key, self._stream_group, message_id)
|
||||
for message_id in message_ids:
|
||||
pipe.xdel(self._stream_key, message_id)
|
||||
await pipe.execute()
|
||||
|
||||
async def _run(self) -> None:
|
||||
while self._running:
|
||||
try:
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if not redis_client:
|
||||
await asyncio.sleep(1)
|
||||
continue
|
||||
|
||||
await self._maybe_claim_pending(redis_client)
|
||||
await self._read_new(redis_client)
|
||||
await self._log_metrics(redis_client)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except (RedisTimeoutError, RedisConnectionError) as exc:
|
||||
logger.warning("[usage-queue] Redis connection issue: {}", exc)
|
||||
await asyncio.sleep(1)
|
||||
except Exception as exc:
|
||||
logger.exception("[usage-queue] Consumer loop error: {}", exc)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def _maybe_claim_pending(self, redis_client: Any) -> None:
|
||||
now = time.time()
|
||||
if now - self._last_claim < self._claim_interval:
|
||||
return
|
||||
self._last_claim = now
|
||||
try:
|
||||
result = await redis_client.xautoclaim(
|
||||
self._stream_key,
|
||||
self._stream_group,
|
||||
self._consumer,
|
||||
min_idle_time=self._claim_idle_ms,
|
||||
start_id="0-0",
|
||||
count=self._batch_size,
|
||||
)
|
||||
except ResponseError as exc:
|
||||
logger.warning("[usage-queue] XAUTOCLAIM failed: {}", exc)
|
||||
return
|
||||
if not result:
|
||||
return
|
||||
_, messages = result[:2]
|
||||
await self._process_messages(redis_client, messages)
|
||||
|
||||
async def _read_new(self, redis_client: Any) -> None:
|
||||
try:
|
||||
result = await redis_client.xreadgroup(
|
||||
groupname=self._stream_group,
|
||||
consumername=self._consumer,
|
||||
streams={self._stream_key: ">"},
|
||||
count=self._batch_size,
|
||||
block=self._block_ms,
|
||||
)
|
||||
except ResponseError as exc:
|
||||
if "NOGROUP" in str(exc):
|
||||
await ensure_usage_stream_group()
|
||||
return
|
||||
raise
|
||||
if not result:
|
||||
return
|
||||
for _stream, messages in result:
|
||||
await self._process_messages(redis_client, messages)
|
||||
|
||||
async def _process_messages(self, redis_client: Any, messages: list) -> None:
|
||||
"""批量处理消息，区分 STREAMING（状态更新）和其他事件（记录写入）"""
|
||||
if not messages:
|
||||
return
|
||||
|
||||
# 分类消息
|
||||
streaming_messages: list[tuple[str, UsageEvent]] = []
|
||||
record_messages: list[tuple[str, dict[str, Any], UsageEvent]] = []
|
||||
failed_messages: list[tuple[str, dict[str, Any], Exception]] = []
|
||||
|
||||
for message_id, fields in messages:
|
||||
try:
|
||||
event = UsageEvent.from_stream_fields(fields)
|
||||
if event.event_type == UsageEventType.STREAMING:
|
||||
streaming_messages.append((message_id, event))
|
||||
else:
|
||||
record_messages.append((message_id, fields, event))
|
||||
except Exception as exc:
|
||||
failed_messages.append((message_id, fields, exc))
|
||||
|
||||
# 批量处理 STREAMING 事件（状态更新）
|
||||
if streaming_messages:
|
||||
await self._process_streaming_batch(redis_client, streaming_messages)
|
||||
|
||||
# 批量处理记录事件
|
||||
if record_messages:
|
||||
await self._process_record_batch(redis_client, record_messages)
|
||||
|
||||
# 处理解析失败的消息
|
||||
for message_id, fields, exc in failed_messages:
|
||||
await self._handle_processing_error(redis_client, message_id, fields, exc)
|
||||
|
||||
async def _process_streaming_batch(
|
||||
self,
|
||||
redis_client: Any,
|
||||
messages: list[tuple[str, UsageEvent]],
|
||||
) -> None:
|
||||
"""批量处理 STREAMING 事件（状态更新）"""
|
||||
success_ids: list[str] = []
|
||||
|
||||
for message_id, event in messages:
|
||||
try:
|
||||
await self._apply_streaming_event(event)
|
||||
success_ids.append(message_id)
|
||||
except Exception as exc:
|
||||
await self._handle_processing_error(redis_client, message_id, {}, exc)
|
||||
|
||||
# 使用 pipeline 批量 ACK 成功处理的消息
|
||||
if success_ids:
|
||||
await self._ack_and_delete_messages(redis_client, success_ids)
|
||||
|
||||
async def _process_record_batch(
|
||||
self,
|
||||
redis_client: Any,
|
||||
messages: list[tuple[str, dict[str, Any], UsageEvent]],
|
||||
) -> None:
|
||||
"""批量处理记录类型的事件"""
|
||||
try:
|
||||
# 准备批量记录数据
|
||||
records: list[dict[str, Any]] = []
|
||||
message_ids: list[str] = []
|
||||
|
||||
for message_id, fields, event in messages:
|
||||
records.append(_event_to_record(event))
|
||||
message_ids.append(message_id)
|
||||
|
||||
# 批量写入
|
||||
await self._record_usage_batch(records)
|
||||
|
||||
# 写库成功后立即从主队列删除，避免 Redis 保留已入库历史。
|
||||
await self._ack_and_delete_messages(redis_client, message_ids)
|
||||
|
||||
logger.debug("[usage-queue] Batch processed {} records", len(records))
|
||||
|
||||
except Exception as exc:
|
||||
# 批量处理失败，回退到逐条处理，确保每条消息在线程内独立写库
|
||||
logger.warning(
|
||||
"[usage-queue] Batch processing failed, falling back to individual: {}", exc
|
||||
)
|
||||
success_ids: list[str] = []
|
||||
for message_id, fields, event in messages:
|
||||
try:
|
||||
await self._apply_record_event(event)
|
||||
success_ids.append(message_id)
|
||||
except IntegrityError as ie:
|
||||
# 重复 request_id 导致的唯一约束冲突，视为成功（记录已存在）
|
||||
if self._is_duplicate_key_error(ie):
|
||||
logger.debug(
|
||||
"[usage-queue] Duplicate request_id, skipping: {}", event.request_id
|
||||
)
|
||||
success_ids.append(message_id)
|
||||
else:
|
||||
await self._handle_processing_error(redis_client, message_id, fields, ie)
|
||||
except Exception as individual_exc:
|
||||
await self._handle_processing_error(
|
||||
redis_client, message_id, fields, individual_exc
|
||||
)
|
||||
# 批量 ACK 成功处理的消息
|
||||
if success_ids:
|
||||
await self._ack_and_delete_messages(redis_client, success_ids)
|
||||
|
||||
async def _handle_processing_error(
|
||||
self,
|
||||
redis_client: Any,
|
||||
message_id: str,
|
||||
fields: dict[str, Any],
|
||||
error: Exception,
|
||||
) -> None:
|
||||
retries = await self._get_delivery_count(redis_client, message_id)
|
||||
if retries >= self._max_retries:
|
||||
try:
|
||||
dlq_fields = dict(fields)
|
||||
dlq_fields["source_id"] = message_id
|
||||
dlq_fields["error"] = str(error)[:200]
|
||||
if self._dlq_maxlen > 0:
|
||||
await redis_client.xadd(
|
||||
self._dlq_key,
|
||||
dlq_fields,
|
||||
maxlen=self._dlq_maxlen,
|
||||
approximate=True,
|
||||
)
|
||||
else:
|
||||
await redis_client.xadd(self._dlq_key, dlq_fields)
|
||||
await self._ack_and_delete_messages(redis_client, [message_id])
|
||||
logger.error(
|
||||
"[usage-queue] Message moved to DLQ after {} attempts: {}", retries, message_id
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("[usage-queue] Failed to move message to DLQ: {}", exc)
|
||||
else:
|
||||
logger.warning(
|
||||
"[usage-queue] Processing failed (attempt {}): {} error={}",
|
||||
retries,
|
||||
message_id,
|
||||
error,
|
||||
)
|
||||
|
||||
async def _get_delivery_count(self, redis_client: Any, message_id: str) -> int:
|
||||
try:
|
||||
pending = await redis_client.xpending_range(
|
||||
self._stream_key,
|
||||
self._stream_group,
|
||||
min=message_id,
|
||||
max=message_id,
|
||||
count=1,
|
||||
)
|
||||
if not pending:
|
||||
return 0
|
||||
info = pending[0]
|
||||
if isinstance(info, dict):
|
||||
return int(info.get("times_delivered", 0))
|
||||
if isinstance(info, (list, tuple)) and len(info) >= 4:
|
||||
return int(info[3])
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
async def _apply_streaming_event(self, event: UsageEvent) -> None:
|
||||
"""处理 STREAMING 事件（状态更新）"""
|
||||
data = event.data
|
||||
|
||||
def _run_update() -> None:
|
||||
db = create_session()
|
||||
try:
|
||||
UsageService.update_usage_status(
|
||||
db=db,
|
||||
request_id=event.request_id,
|
||||
status="streaming",
|
||||
provider=data.get("provider"),
|
||||
target_model=data.get("target_model"),
|
||||
first_byte_time_ms=data.get("first_byte_time_ms"),
|
||||
provider_id=data.get("provider_id"),
|
||||
provider_endpoint_id=data.get("provider_endpoint_id"),
|
||||
provider_api_key_id=data.get("provider_api_key_id"),
|
||||
api_format=data.get("api_format"),
|
||||
endpoint_api_format=data.get("endpoint_api_format"),
|
||||
has_format_conversion=data.get("has_format_conversion"),
|
||||
request_headers=data.get("request_headers"),
|
||||
request_body=data.get("request_body"),
|
||||
provider_request_headers=data.get("provider_request_headers"),
|
||||
provider_request_body=data.get("provider_request_body"),
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
await asyncio.to_thread(_run_update)
|
||||
|
||||
async def _apply_record_event(self, event: UsageEvent) -> None:
|
||||
"""处理记录类型事件（逐条写入，用于 fallback）"""
|
||||
await self._record_usage_batch([_event_to_record(event)])
|
||||
|
||||
async def _apply_event(self, event: UsageEvent) -> None:
|
||||
"""处理单个事件（兼容旧接口，用于测试）"""
|
||||
if event.event_type == UsageEventType.STREAMING:
|
||||
await self._apply_streaming_event(event)
|
||||
else:
|
||||
await self._apply_record_event(event)
|
||||
|
||||
async def _log_metrics(self, redis_client: Any) -> None:
|
||||
now = time.time()
|
||||
if now - self._last_metrics_log < self._metrics_interval:
|
||||
return
|
||||
self._last_metrics_log = now
|
||||
try:
|
||||
# 使用 XINFO GROUPS 获取更准确的 lag（未处理消息数）
|
||||
groups_info = await redis_client.xinfo_groups(self._stream_key)
|
||||
lag = 0
|
||||
pending_count = 0
|
||||
for group in groups_info:
|
||||
if isinstance(group, dict) and group.get("name") == self._stream_group:
|
||||
lag = group.get("lag", 0) or 0
|
||||
pending_count = group.get("pending", 0) or 0
|
||||
break
|
||||
# lag=未读消息数, pending=已读但未ACK的消息数
|
||||
if lag > 0 or pending_count > 0:
|
||||
logger.info("[usage-queue] lag={} pending={}", lag, pending_count)
|
||||
except Exception as exc:
|
||||
logger.debug("[usage-queue] metrics log failed: {}", exc)
|
||||
|
||||
|
||||
_consumer_instance: UsageQueueConsumer | None = None
|
||||
|
||||
|
||||
async def start_usage_queue_consumer() -> UsageQueueConsumer | None:
|
||||
global _consumer_instance
|
||||
if not config.usage_queue_enabled or not config.usage_queue_python_consumer_enabled:
|
||||
return None
|
||||
await ensure_usage_stream_group()
|
||||
if _consumer_instance is None:
|
||||
_consumer_instance = UsageQueueConsumer()
|
||||
await _consumer_instance.start()
|
||||
return _consumer_instance
|
||||
|
||||
|
||||
async def stop_usage_queue_consumer() -> None:
|
||||
global _consumer_instance
|
||||
if _consumer_instance:
|
||||
await _consumer_instance.stop()
|
||||
_consumer_instance = None
|
||||
70
_deprecated_py_src/services/usage/error_classifier.py
Normal file
70
_deprecated_py_src/services/usage/error_classifier.py
Normal file
@@ -0,0 +1,70 @@
|
||||
"""Error classification helpers for Usage records."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.core.enums import ErrorCategory
|
||||
|
||||
_STATUS_CODE_MAP: dict[int, ErrorCategory] = {
|
||||
400: ErrorCategory.INVALID_REQUEST,
|
||||
401: ErrorCategory.AUTH,
|
||||
403: ErrorCategory.AUTH,
|
||||
404: ErrorCategory.NOT_FOUND,
|
||||
408: ErrorCategory.TIMEOUT,
|
||||
429: ErrorCategory.RATE_LIMIT,
|
||||
500: ErrorCategory.SERVER_ERROR,
|
||||
502: ErrorCategory.SERVER_ERROR,
|
||||
503: ErrorCategory.SERVER_ERROR,
|
||||
504: ErrorCategory.TIMEOUT,
|
||||
}
|
||||
|
||||
_CONTEXT_LENGTH_PATTERNS = (
|
||||
"context_length_exceeded",
|
||||
"maximum context length",
|
||||
"too many tokens",
|
||||
"input is too long",
|
||||
)
|
||||
|
||||
_CONTENT_FILTER_PATTERNS = (
|
||||
"content_filter",
|
||||
"content_policy",
|
||||
"safety_block",
|
||||
"blocked by content",
|
||||
)
|
||||
|
||||
_NETWORK_PATTERNS = ("connection", "network", "dns", "socket")
|
||||
|
||||
|
||||
def classify_error(
|
||||
status_code: int | None,
|
||||
error_message: str | None,
|
||||
status: str | None = None,
|
||||
) -> ErrorCategory:
|
||||
"""Map provider errors to ErrorCategory."""
|
||||
if status and status.lower() == "cancelled":
|
||||
return ErrorCategory.CANCELLED
|
||||
|
||||
if status_code is not None:
|
||||
mapped = _STATUS_CODE_MAP.get(status_code)
|
||||
if mapped:
|
||||
return mapped
|
||||
|
||||
if error_message:
|
||||
msg_lower = error_message.lower()
|
||||
if any(p in msg_lower for p in _CONTEXT_LENGTH_PATTERNS):
|
||||
return ErrorCategory.CONTEXT_LENGTH
|
||||
if any(p in msg_lower for p in _CONTENT_FILTER_PATTERNS):
|
||||
return ErrorCategory.CONTENT_FILTER
|
||||
if "rate limit" in msg_lower or "rate_limit" in msg_lower:
|
||||
return ErrorCategory.RATE_LIMIT
|
||||
if "timeout" in msg_lower or "timed out" in msg_lower:
|
||||
return ErrorCategory.TIMEOUT
|
||||
if any(p in msg_lower for p in _NETWORK_PATTERNS):
|
||||
return ErrorCategory.NETWORK
|
||||
|
||||
if status_code is not None:
|
||||
if status_code >= 500:
|
||||
return ErrorCategory.SERVER_ERROR
|
||||
if status_code >= 400:
|
||||
return ErrorCategory.INVALID_REQUEST
|
||||
|
||||
return ErrorCategory.UNKNOWN
|
||||
125
_deprecated_py_src/services/usage/events.py
Normal file
125
_deprecated_py_src/services/usage/events.py
Normal file
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
Usage 事件定义与序列化工具（用于 Redis Streams）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import msgpack
|
||||
from msgpack.exceptions import OutOfData
|
||||
|
||||
USAGE_EVENT_VERSION = 1
|
||||
|
||||
|
||||
class UsageEventType(str, Enum):
|
||||
STREAMING = "streaming"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
def now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def _sanitize_value(value: Any) -> Any:
|
||||
if value is None or isinstance(value, (str, int, float, bool)):
|
||||
return value
|
||||
if isinstance(value, dict):
|
||||
return {str(k): _sanitize_value(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_sanitize_value(item) for item in value]
|
||||
return str(value)
|
||||
|
||||
|
||||
def sanitize_payload(data: dict[str, Any]) -> dict[str, Any]:
|
||||
return {str(k): _sanitize_value(v) for k, v in data.items()}
|
||||
|
||||
|
||||
def _decode_payload(raw: Any) -> dict[str, Any]:
|
||||
"""兼容解码：优先 msgpack，回退旧 JSON。
|
||||
|
||||
统一将输入归一化为 bytes 后走单一解码路径：msgpack → JSON fallback。
|
||||
str 输入来自 decode_responses=True + surrogateescape 的 Redis 客户端，
|
||||
通过 surrogateescape 可无损还原回原始 bytes。
|
||||
"""
|
||||
if isinstance(raw, str):
|
||||
raw = raw.encode("utf-8", errors="surrogateescape")
|
||||
|
||||
if not isinstance(raw, (bytes, bytearray, memoryview)):
|
||||
raise ValueError("Invalid payload field in usage event")
|
||||
|
||||
payload_bytes = bytes(raw)
|
||||
|
||||
# 新格式：msgpack
|
||||
try:
|
||||
payload = msgpack.unpackb(payload_bytes, raw=False)
|
||||
except (ValueError, OutOfData, TypeError):
|
||||
# 兼容旧格式：JSON bytes（含 surrogateescape 还原后的纯 UTF-8 JSON）
|
||||
try:
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError, TypeError) as exc:
|
||||
raise ValueError("Invalid payload field in usage event") from exc
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Invalid payload field in usage event")
|
||||
return payload
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageEvent:
|
||||
event_type: UsageEventType
|
||||
request_id: str
|
||||
timestamp_ms: int
|
||||
data: dict[str, Any]
|
||||
|
||||
def to_stream_fields(self) -> dict[str, bytes]:
|
||||
"""序列化为 Redis Stream 字段。
|
||||
|
||||
该函数返回 bytes payload，要求读写 usage queue 的 Redis 客户端使用
|
||||
decode_responses=True，并以 surrogateescape 做 UTF-8 编解码，
|
||||
以保证 bytes <-> str 往返无损。
|
||||
"""
|
||||
payload = {
|
||||
"v": USAGE_EVENT_VERSION,
|
||||
"type": self.event_type.value,
|
||||
"request_id": self.request_id,
|
||||
"timestamp_ms": self.timestamp_ms,
|
||||
# 兜底清洗，避免 metadata 中混入非 JSON 类型导致队列写入失败。
|
||||
"data": sanitize_payload(self.data),
|
||||
}
|
||||
return {"payload": msgpack.packb(payload, use_bin_type=True)}
|
||||
|
||||
@classmethod
|
||||
def from_stream_fields(cls, fields: dict[str, Any]) -> UsageEvent:
|
||||
raw = fields.get("payload")
|
||||
if not raw:
|
||||
raise ValueError("Missing payload field in usage event")
|
||||
payload = _decode_payload(raw)
|
||||
event_type = UsageEventType(payload["type"])
|
||||
return cls(
|
||||
event_type=event_type,
|
||||
request_id=payload["request_id"],
|
||||
timestamp_ms=int(payload.get("timestamp_ms", 0)),
|
||||
data=payload.get("data", {}) or {},
|
||||
)
|
||||
|
||||
|
||||
def build_usage_event(
|
||||
*,
|
||||
event_type: UsageEventType,
|
||||
request_id: str,
|
||||
data: dict[str, Any],
|
||||
timestamp_ms: int | None = None,
|
||||
) -> UsageEvent:
|
||||
return UsageEvent(
|
||||
event_type=event_type,
|
||||
request_id=request_id,
|
||||
timestamp_ms=timestamp_ms or now_ms(),
|
||||
data=data,
|
||||
)
|
||||
539
_deprecated_py_src/services/usage/lifecycle.py
Normal file
539
_deprecated_py_src/services/usage/lifecycle.py
Normal file
@@ -0,0 +1,539 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Usage, User
|
||||
from src.services.billing.precision import to_money_decimal
|
||||
from src.services.provider_keys.codex_quota_sync_dispatcher import (
|
||||
dispatch_codex_quota_sync_from_response_headers,
|
||||
)
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.wallet import WalletService
|
||||
|
||||
|
||||
class UsageLifecycleMixin:
|
||||
"""使用记录生命周期管理方法"""
|
||||
|
||||
@staticmethod
|
||||
def _is_billing_terminal(usage: Usage | None) -> bool:
|
||||
return bool(
|
||||
usage is not None and getattr(usage, "billing_status", None) in {"settled", "void"}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def begin_pending_usage(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
model: str,
|
||||
*,
|
||||
is_stream: bool = False,
|
||||
request_type: str = "chat",
|
||||
api_format: str | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
) -> Usage:
|
||||
"""
|
||||
创建（或返回已有）pending Usage 记录，但**不提交事务**。
|
||||
|
||||
适用场景：
|
||||
- ApplicationService 在同一事务内创建 pending usage + task + candidates
|
||||
- submit 幂等：重复调用同一 request_id 时返回已有记录
|
||||
"""
|
||||
existing = db.query(Usage).filter(Usage.request_id == request_id).first()
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
# 根据配置决定是否记录请求详情
|
||||
should_log_headers = SystemConfigService.should_log_headers(db)
|
||||
should_log_body = SystemConfigService.should_log_body(db)
|
||||
|
||||
# 处理请求头
|
||||
processed_request_headers = None
|
||||
if should_log_headers and request_headers is not None:
|
||||
processed_request_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, request_headers
|
||||
)
|
||||
|
||||
# 处理请求体
|
||||
processed_request_body = None
|
||||
if should_log_body and request_body is not None:
|
||||
processed_request_body = SystemConfigService.truncate_body(
|
||||
db, request_body, is_request=True
|
||||
)
|
||||
|
||||
usage = Usage(
|
||||
user_id=user.id if user else None,
|
||||
api_key_id=api_key.id if api_key else None,
|
||||
username=user.username if user else None,
|
||||
api_key_name=api_key.name if api_key else None,
|
||||
request_id=request_id,
|
||||
provider_name="pending", # 尚未确定 provider
|
||||
model=model,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
total_tokens=0,
|
||||
total_cost_usd=0.0,
|
||||
request_type=request_type,
|
||||
api_format=api_format,
|
||||
is_stream=is_stream,
|
||||
status="pending",
|
||||
billing_status="pending",
|
||||
request_headers=processed_request_headers,
|
||||
request_body=processed_request_body,
|
||||
)
|
||||
|
||||
db.add(usage)
|
||||
db.flush()
|
||||
return usage
|
||||
|
||||
@classmethod
|
||||
def create_pending_usage(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
model: str,
|
||||
is_stream: bool = False,
|
||||
request_type: str = "chat",
|
||||
api_format: str | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
) -> Usage:
|
||||
"""
|
||||
创建 pending 状态的使用记录（在请求开始时调用）
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
request_id: 请求ID
|
||||
user: 用户对象
|
||||
api_key: API Key 对象
|
||||
model: 模型名称
|
||||
is_stream: 是否流式请求
|
||||
api_format: API 格式
|
||||
request_headers: 请求头
|
||||
request_body: 请求体
|
||||
|
||||
Returns:
|
||||
创建的 Usage 记录
|
||||
"""
|
||||
usage = cls.begin_pending_usage(
|
||||
db,
|
||||
request_id=request_id,
|
||||
user=user,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
is_stream=is_stream,
|
||||
request_type=request_type,
|
||||
api_format=api_format,
|
||||
request_headers=request_headers,
|
||||
request_body=request_body,
|
||||
)
|
||||
db.commit()
|
||||
|
||||
logger.debug("创建 pending 使用记录: request_id={}, model={}", request_id, model)
|
||||
|
||||
return usage
|
||||
|
||||
# ========== billing_status 并发幂等 finalize ==========
|
||||
|
||||
@classmethod
|
||||
def finalize_settled(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
*,
|
||||
total_cost_usd: float,
|
||||
request_cost_usd: float | None = None,
|
||||
status: str = "completed",
|
||||
status_code: int = 200,
|
||||
error_message: str | None = None,
|
||||
response_time_ms: int | None = None,
|
||||
billing_snapshot: dict[str, Any] | None = None,
|
||||
extra_metadata: dict[str, Any] | None = None,
|
||||
finalized_at: datetime | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
并发安全的幂等 finalize（settled）。
|
||||
|
||||
约定：
|
||||
- 仅当 billing_status='pending' 时才会生效（rowcount==1）
|
||||
- 不在本方法内 commit，由调用方决定事务提交时机
|
||||
"""
|
||||
now = finalized_at or datetime.now(timezone.utc)
|
||||
cost = to_money_decimal(total_cost_usd)
|
||||
request_cost = to_money_decimal(request_cost_usd) if request_cost_usd is not None else cost
|
||||
|
||||
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
|
||||
if not usage or usage.billing_status != "pending":
|
||||
return False
|
||||
|
||||
usage.billing_status = "settled"
|
||||
usage.finalized_at = now
|
||||
usage.total_cost_usd = cost
|
||||
usage.request_cost_usd = request_cost
|
||||
usage.status = status
|
||||
usage.status_code = status_code
|
||||
usage.error_message = error_message
|
||||
usage.response_time_ms = response_time_ms
|
||||
if cost > 0:
|
||||
WalletService.apply_usage_charge(db, usage=usage, amount_usd=cost)
|
||||
|
||||
# 写入审计快照（只在本次 finalize 生效时执行）
|
||||
metadata = usage.request_metadata or {}
|
||||
if billing_snapshot is not None:
|
||||
metadata["billing_snapshot"] = billing_snapshot
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
usage.request_metadata = cls._sanitize_request_metadata(metadata)
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def finalize_void(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
*,
|
||||
reason: str | None = None,
|
||||
status_code: int = 499,
|
||||
finalized_at: datetime | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
并发安全的幂等 finalize（void，不收费）。
|
||||
|
||||
约定：
|
||||
- 仅当 billing_status='pending' 时才会生效（rowcount==1）
|
||||
- 不在本方法内 commit，由调用方决定事务提交时机
|
||||
"""
|
||||
now = finalized_at or datetime.now(timezone.utc)
|
||||
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
|
||||
if not usage or usage.billing_status != "pending":
|
||||
return False
|
||||
|
||||
usage.billing_status = "void"
|
||||
usage.finalized_at = now
|
||||
usage.total_cost_usd = to_money_decimal(0)
|
||||
usage.request_cost_usd = to_money_decimal(0)
|
||||
usage.actual_total_cost_usd = to_money_decimal(0)
|
||||
usage.actual_request_cost_usd = to_money_decimal(0)
|
||||
usage.status = "cancelled"
|
||||
usage.status_code = status_code
|
||||
usage.error_message = reason
|
||||
usage.response_time_ms = None
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def finalize_submitted(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
*,
|
||||
provider_name: str,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
response_time_ms: int | None = None,
|
||||
status_code: int = 200,
|
||||
endpoint_api_format: str | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
response_headers: dict[str, Any] | None = None,
|
||||
response_body: Any | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
异步任务提交成功时的幂等结算。
|
||||
|
||||
将 pending 使用记录保留为 pending，仅补齐已知的 provider/响应信息。
|
||||
后续轮询完成后通过 update_settled_billing 一次性写入实际费用并扣钱包。
|
||||
|
||||
约定：
|
||||
- 仅当 billing_status='pending' 时才会生效（rowcount==1）
|
||||
- 不在本方法内 commit，由调用方决定事务提交时机
|
||||
"""
|
||||
# 处理响应头和响应体
|
||||
should_log_headers = SystemConfigService.should_log_headers(db)
|
||||
should_log_body = SystemConfigService.should_log_body(db)
|
||||
|
||||
processed_provider_headers = None
|
||||
if should_log_headers and provider_request_headers is not None:
|
||||
processed_provider_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, provider_request_headers
|
||||
)
|
||||
|
||||
processed_response_headers = None
|
||||
if should_log_headers and response_headers is not None:
|
||||
processed_response_headers = dict(response_headers)
|
||||
|
||||
processed_response_body = None
|
||||
if should_log_body and response_body is not None:
|
||||
processed_response_body = SystemConfigService.truncate_body(
|
||||
db, response_body, is_request=False
|
||||
)
|
||||
|
||||
values: dict[str, Any] = {
|
||||
"status": "pending",
|
||||
"status_code": status_code,
|
||||
"response_time_ms": response_time_ms,
|
||||
"provider_name": provider_name,
|
||||
"provider_id": provider_id,
|
||||
"provider_endpoint_id": provider_endpoint_id,
|
||||
"provider_api_key_id": provider_api_key_id,
|
||||
"endpoint_api_format": endpoint_api_format,
|
||||
}
|
||||
|
||||
if processed_provider_headers is not None:
|
||||
values["provider_request_headers"] = processed_provider_headers
|
||||
if processed_response_headers is not None:
|
||||
values["response_headers"] = processed_response_headers
|
||||
if processed_response_body is not None:
|
||||
values["response_body"] = processed_response_body
|
||||
|
||||
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
|
||||
if not usage or usage.billing_status != "pending":
|
||||
return False
|
||||
for key, value in values.items():
|
||||
setattr(usage, key, value)
|
||||
finalized = True
|
||||
if finalized:
|
||||
dispatch_codex_quota_sync_from_response_headers(
|
||||
provider_api_key_id=provider_api_key_id,
|
||||
response_headers=response_headers,
|
||||
db=db,
|
||||
)
|
||||
return finalized
|
||||
|
||||
@classmethod
|
||||
def update_settled_billing(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
*,
|
||||
total_cost_usd: float,
|
||||
request_cost_usd: float | None = None,
|
||||
status: str = "completed",
|
||||
status_code: int = 200,
|
||||
error_message: str | None = None,
|
||||
response_time_ms: int | None = None,
|
||||
billing_snapshot: dict[str, Any] | None = None,
|
||||
extra_metadata: dict[str, Any] | None = None,
|
||||
finalized_at: datetime | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
写入异步任务最终账单（轮询完成后调用）。
|
||||
|
||||
语义：
|
||||
- 仅允许 pending -> settled / void（首次最终结算）
|
||||
- settled / void 一旦进入即不可再修改
|
||||
|
||||
约定：
|
||||
- 不在本方法内 commit，由调用方决定事务提交时机
|
||||
"""
|
||||
now = finalized_at or datetime.now(timezone.utc)
|
||||
cost = to_money_decimal(total_cost_usd)
|
||||
request_cost = to_money_decimal(request_cost_usd) if request_cost_usd is not None else cost
|
||||
|
||||
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
|
||||
if not usage or usage.billing_status != "pending":
|
||||
return False
|
||||
|
||||
usage.total_cost_usd = cost
|
||||
usage.request_cost_usd = request_cost
|
||||
usage.status = status
|
||||
usage.status_code = status_code
|
||||
if error_message is not None:
|
||||
usage.error_message = error_message
|
||||
if response_time_ms is not None:
|
||||
usage.response_time_ms = response_time_ms
|
||||
usage.finalized_at = now
|
||||
if cost > 0:
|
||||
usage.billing_status = "settled"
|
||||
WalletService.apply_usage_charge(db, usage=usage, amount_usd=cost)
|
||||
else:
|
||||
usage.billing_status = "void" if status in {"failed", "cancelled"} else "settled"
|
||||
|
||||
# 写入审计快照
|
||||
metadata = usage.request_metadata or {}
|
||||
if billing_snapshot is not None:
|
||||
metadata["billing_snapshot"] = billing_snapshot
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
metadata["billing_updated_at"] = now.isoformat()
|
||||
usage.request_metadata = cls._sanitize_request_metadata(metadata)
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def void_settled(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
*,
|
||||
reason: str | None = None,
|
||||
status_code: int = 499,
|
||||
) -> bool:
|
||||
"""
|
||||
已废弃：settled 为账务终态，不允许再回滚为 void。
|
||||
"""
|
||||
logger.warning(
|
||||
"void_settled is deprecated and ignored: request_id={}, reason={}, status_code={}",
|
||||
request_id,
|
||||
reason,
|
||||
status_code,
|
||||
)
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def update_usage_status(
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
status: str,
|
||||
error_message: str | None = None,
|
||||
provider: str | None = None,
|
||||
target_model: str | None = None,
|
||||
first_byte_time_ms: int | None = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
api_format: str | None = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool | None = None,
|
||||
status_code: int | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
provider_request_body: Any | None = None,
|
||||
) -> Usage | None:
|
||||
"""
|
||||
快速更新使用记录状态
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
request_id: 请求ID
|
||||
status: 新状态 (pending, streaming, completed, failed)
|
||||
error_message: 错误消息（仅在 failed 状态时使用）
|
||||
provider: 提供商名称（可选，streaming 状态时更新）
|
||||
target_model: 映射后的目标模型名（可选）
|
||||
first_byte_time_ms: 首字时间/TTFB（可选，streaming 状态时更新）
|
||||
provider_id: Provider ID（可选，streaming 状态时更新）
|
||||
provider_endpoint_id: Endpoint ID（可选，streaming 状态时更新）
|
||||
provider_api_key_id: Provider API Key ID（可选，streaming 状态时更新）
|
||||
api_format: API 格式（可选，用于获取按格式配置的倍率）
|
||||
endpoint_api_format: 端点原生 API 格式（可选）
|
||||
has_format_conversion: 是否发生了格式转换（可选）
|
||||
status_code: HTTP 状态码（可选）
|
||||
request_headers: 客户端请求头（可选，用于补写 pending/streaming 记录）
|
||||
request_body: 客户端请求体（可选，用于补写 pending/streaming 记录）
|
||||
provider_request_headers: 提供商请求头（可选，streaming 时可写入）
|
||||
provider_request_body: 提供商请求体（可选，streaming 时可写入）
|
||||
|
||||
Returns:
|
||||
更新后的 Usage 记录，如果未找到则返回 None
|
||||
"""
|
||||
usage = db.query(Usage).filter(Usage.request_id == request_id).first()
|
||||
if not usage:
|
||||
logger.warning("未找到 request_id={} 的使用记录，无法更新状态", request_id)
|
||||
return None
|
||||
|
||||
if cls._is_billing_terminal(usage):
|
||||
logger.debug(
|
||||
"跳过已终态 Usage 状态更新: request_id={}, status={}, billing_status={}",
|
||||
request_id,
|
||||
status,
|
||||
getattr(usage, "billing_status", None),
|
||||
)
|
||||
return usage
|
||||
|
||||
# 避免状态回退：streaming 只能从 pending/streaming 进入
|
||||
if status == "streaming" and usage.status not in ("pending", "streaming"):
|
||||
logger.debug(
|
||||
f"跳过 streaming 状态更新（避免回退）: request_id={request_id}, "
|
||||
f"{usage.status} -> {status}"
|
||||
)
|
||||
return usage
|
||||
|
||||
old_status = usage.status
|
||||
usage.status = status
|
||||
if error_message:
|
||||
usage.error_message = error_message
|
||||
if provider:
|
||||
usage.provider_name = provider
|
||||
elif status == "streaming" and usage.provider_name == "pending":
|
||||
# 状态变为 streaming 但 provider_name 仍为 pending，记录警告
|
||||
logger.warning(
|
||||
f"状态更新为 streaming 但 provider_name 为空: request_id={request_id}, "
|
||||
f"当前 provider_name={usage.provider_name}"
|
||||
)
|
||||
if target_model:
|
||||
usage.target_model = target_model
|
||||
if first_byte_time_ms is not None:
|
||||
usage.first_byte_time_ms = first_byte_time_ms
|
||||
if provider_id is not None:
|
||||
usage.provider_id = provider_id
|
||||
if provider_endpoint_id is not None:
|
||||
usage.provider_endpoint_id = provider_endpoint_id
|
||||
if provider_api_key_id is not None:
|
||||
usage.provider_api_key_id = provider_api_key_id
|
||||
# 当设置 provider_api_key_id 时，同步获取并更新 rate_multiplier
|
||||
# 这样前端在 streaming 状态就能显示倍率
|
||||
rate_multiplier = cls._get_rate_multiplier_sync(
|
||||
db, provider_api_key_id, api_format or usage.api_format
|
||||
)
|
||||
if rate_multiplier is not None:
|
||||
usage.rate_multiplier = rate_multiplier
|
||||
if endpoint_api_format is not None:
|
||||
usage.endpoint_api_format = endpoint_api_format
|
||||
if has_format_conversion is not None:
|
||||
usage.has_format_conversion = has_format_conversion
|
||||
if status_code is not None:
|
||||
usage.status_code = status_code
|
||||
|
||||
should_log_headers = SystemConfigService.should_log_headers(db)
|
||||
should_log_body = SystemConfigService.should_log_body(db)
|
||||
|
||||
if should_log_headers:
|
||||
if isinstance(request_headers, dict):
|
||||
usage.request_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, request_headers
|
||||
)
|
||||
if isinstance(provider_request_headers, dict):
|
||||
usage.provider_request_headers = SystemConfigService.mask_sensitive_headers(
|
||||
db, provider_request_headers
|
||||
)
|
||||
if should_log_body:
|
||||
if request_body is not None:
|
||||
usage.request_body = SystemConfigService.truncate_body(
|
||||
db, request_body, is_request=True
|
||||
)
|
||||
if provider_request_body is not None:
|
||||
usage.provider_request_body = SystemConfigService.truncate_body(
|
||||
db, provider_request_body, is_request=True
|
||||
)
|
||||
|
||||
# 仅在“明确不会收费”的终态下直接关闭账单。
|
||||
# completed 的费用通常要由后续 record_usage / update_settled_billing 写入，
|
||||
# 这里不能提前把 billing_status 置为 settled，否则会阻断真正扣费。
|
||||
if (
|
||||
status in ("failed", "cancelled")
|
||||
and getattr(usage, "billing_status", None) == "pending"
|
||||
):
|
||||
usage.billing_status = "void"
|
||||
if getattr(usage, "finalized_at", None) is None:
|
||||
usage.finalized_at = datetime.now(timezone.utc)
|
||||
usage.total_cost_usd = to_money_decimal(0)
|
||||
usage.request_cost_usd = to_money_decimal(0)
|
||||
usage.actual_total_cost_usd = to_money_decimal(0)
|
||||
usage.actual_request_cost_usd = to_money_decimal(0)
|
||||
|
||||
db.commit()
|
||||
|
||||
logger.debug("更新使用记录状态: request_id={}, {} -> {}", request_id, old_status, status)
|
||||
|
||||
return usage
|
||||
176
_deprecated_py_src/services/usage/pricing.py
Normal file
176
_deprecated_py_src/services/usage/pricing.py
Normal file
@@ -0,0 +1,176 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.model.cost import ModelCostService
|
||||
|
||||
|
||||
class UsagePricingMixin:
|
||||
"""定价相关方法"""
|
||||
|
||||
@classmethod
|
||||
async def get_model_price_async(
|
||||
cls, db: Session, provider: str, model: str
|
||||
) -> tuple[float, float]:
|
||||
"""异步获取模型价格（输入价格，输出价格）每1M tokens
|
||||
|
||||
查找逻辑：
|
||||
1. 直接通过 GlobalModel.name 匹配
|
||||
2. 查找该 Provider 的 Model 实现并获取价格
|
||||
3. 如果找不到则使用系统默认价格
|
||||
"""
|
||||
|
||||
service = ModelCostService(db)
|
||||
return await service.get_model_price_async(provider, model)
|
||||
|
||||
@classmethod
|
||||
def get_model_price(cls, db: Session, provider: str, model: str) -> tuple[float, float]:
|
||||
"""获取模型价格（输入价格，输出价格）每1M tokens
|
||||
|
||||
查找逻辑：
|
||||
1. 直接通过 GlobalModel.name 匹配
|
||||
2. 查找该 Provider 的 Model 实现并获取价格
|
||||
3. 如果找不到则使用系统默认价格
|
||||
"""
|
||||
|
||||
service = ModelCostService(db)
|
||||
return service.get_model_price(provider, model)
|
||||
|
||||
@classmethod
|
||||
async def get_cache_prices_async(
|
||||
cls, db: Session, provider: str, model: str, input_price: float
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""异步获取模型缓存价格（缓存创建价格，缓存读取价格）每1M tokens"""
|
||||
service = ModelCostService(db)
|
||||
return await service.get_cache_prices_async(provider, model, input_price)
|
||||
|
||||
@classmethod
|
||||
def get_cache_prices(
|
||||
cls, db: Session, provider: str, model: str, input_price: float
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""获取模型缓存价格（缓存创建价格，缓存读取价格）每1M tokens"""
|
||||
service = ModelCostService(db)
|
||||
return service.get_cache_prices(provider, model, input_price)
|
||||
|
||||
@classmethod
|
||||
async def get_request_price_async(cls, db: Session, provider: str, model: str) -> float | None:
|
||||
"""异步获取模型按次计费价格"""
|
||||
service = ModelCostService(db)
|
||||
return await service.get_request_price_async(provider, model)
|
||||
|
||||
@classmethod
|
||||
def get_request_price(cls, db: Session, provider: str, model: str) -> float | None:
|
||||
"""获取模型按次计费价格"""
|
||||
service = ModelCostService(db)
|
||||
return service.get_request_price(provider, model)
|
||||
|
||||
@staticmethod
|
||||
def calculate_cost(
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
input_price_per_1m: float,
|
||||
output_price_per_1m: float,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
cache_creation_price_per_1m: float | None = None,
|
||||
cache_read_price_per_1m: float | None = None,
|
||||
price_per_request: float | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float]:
|
||||
"""计算成本（价格是每百万tokens）- 固定价格模式
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost, cache_creation_cost,
|
||||
cache_read_cost, cache_cost, request_cost, total_cost)
|
||||
"""
|
||||
return ModelCostService.compute_cost(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
input_price_per_1m=input_price_per_1m,
|
||||
output_price_per_1m=output_price_per_1m,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
cache_creation_price_per_1m=cache_creation_price_per_1m,
|
||||
cache_read_price_per_1m=cache_read_price_per_1m,
|
||||
price_per_request=price_per_request,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def calculate_cost_with_strategy_async(
|
||||
cls,
|
||||
db: Session,
|
||||
provider: str,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: str | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, int | None]:
|
||||
"""使用策略模式计算成本（支持阶梯计费）
|
||||
|
||||
根据 api_format 选择对应的计费策略，支持阶梯计费和 TTL 差异化。
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost, cache_creation_cost,
|
||||
cache_read_cost, cache_cost, request_cost, total_cost, tier_index)
|
||||
"""
|
||||
service = ModelCostService(db)
|
||||
return await service.compute_cost_with_strategy_async(
|
||||
provider=provider,
|
||||
model=model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
api_format=api_format,
|
||||
cache_ttl_minutes=cache_ttl_minutes,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _get_rate_multiplier_and_free_tier(
|
||||
cls,
|
||||
db: Session,
|
||||
provider_api_key_id: str | None,
|
||||
provider_id: str | None,
|
||||
api_format: str | None = None,
|
||||
) -> tuple[float, bool]:
|
||||
"""获取费率倍数和是否免费套餐（使用缓存）"""
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
|
||||
return await ProviderCacheService.get_rate_multiplier_and_free_tier(
|
||||
db, provider_api_key_id, provider_id, api_format
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_rate_multiplier_sync(
|
||||
db: Session,
|
||||
provider_api_key_id: str,
|
||||
api_format: str | None = None,
|
||||
) -> float | None:
|
||||
"""
|
||||
同步获取 ProviderAPIKey 的 rate_multiplier
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
provider_api_key_id: ProviderAPIKey ID
|
||||
api_format: API 格式（可选），如 "CLAUDE"、"OPENAI"
|
||||
|
||||
Returns:
|
||||
rate_multiplier 或 None
|
||||
"""
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
|
||||
provider_key = (
|
||||
db.query(ProviderAPIKey.rate_multipliers)
|
||||
.filter(ProviderAPIKey.id == provider_api_key_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not provider_key:
|
||||
return None
|
||||
|
||||
return ProviderCacheService.compute_rate_multiplier(
|
||||
provider_key.rate_multipliers, api_format
|
||||
)
|
||||
578
_deprecated_py_src/services/usage/query.py
Normal file
578
_deprecated_py_src/services/usage/query.py
Normal file
@@ -0,0 +1,578 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy import case, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.sql.elements import ColumnElement
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Usage, User
|
||||
|
||||
|
||||
def input_context_expr() -> ColumnElement[int]:
|
||||
"""计算缓存命中率口径下的总输入上下文 token 数。
|
||||
|
||||
为了与 usage 表中“输入 tokens + 缓存读取 tokens”的展示口径保持一致，
|
||||
聚合统计统一使用 `input_tokens + cache_read_input_tokens` 作为分母。
|
||||
"""
|
||||
return Usage.input_tokens + Usage.cache_read_input_tokens
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RequestBalanceCheckResult:
|
||||
allowed: bool
|
||||
message: str
|
||||
remaining: float | None
|
||||
|
||||
|
||||
class UsageQueryMixin:
|
||||
"""查询/统计相关方法"""
|
||||
|
||||
# 热力图缓存键前缀（依赖 TTL 自动过期，用户角色变更时主动清除）
|
||||
HEATMAP_CACHE_KEY_PREFIX = "activity_heatmap"
|
||||
|
||||
@classmethod
|
||||
def _get_heatmap_cache_key(cls, user_id: str | None, include_actual_cost: bool) -> str:
|
||||
"""生成热力图缓存键"""
|
||||
cost_suffix = "with_cost" if include_actual_cost else "no_cost"
|
||||
if user_id:
|
||||
return f"{cls.HEATMAP_CACHE_KEY_PREFIX}:user:{user_id}:{cost_suffix}"
|
||||
else:
|
||||
return f"{cls.HEATMAP_CACHE_KEY_PREFIX}:admin:all:{cost_suffix}"
|
||||
|
||||
@classmethod
|
||||
async def clear_user_heatmap_cache(cls, user_id: str) -> None:
|
||||
"""
|
||||
清除用户的热力图缓存（用户角色变更时调用）
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
"""
|
||||
from src.clients.redis_client import get_redis_client
|
||||
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if not redis_client:
|
||||
return
|
||||
|
||||
# 清除该用户的所有热力图缓存（with_cost 和 no_cost）
|
||||
keys_to_delete = [
|
||||
cls._get_heatmap_cache_key(user_id, include_actual_cost=True),
|
||||
cls._get_heatmap_cache_key(user_id, include_actual_cost=False),
|
||||
]
|
||||
|
||||
for key in keys_to_delete:
|
||||
try:
|
||||
await redis_client.delete(key)
|
||||
logger.debug("已清除热力图缓存: {}", key)
|
||||
except Exception as e:
|
||||
logger.warning("清除热力图缓存失败: {}, error={}", key, e)
|
||||
|
||||
@classmethod
|
||||
async def get_cached_heatmap(
|
||||
cls,
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
include_actual_cost: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取带缓存的热力图数据
|
||||
|
||||
缓存策略：
|
||||
- TTL: 10分钟（CacheTTL.ACTIVITY_HEATMAP = 600）
|
||||
- 仅依赖 TTL 自动过期，新使用记录最多延迟 10 分钟出现
|
||||
- 用户角色变更时通过 clear_user_heatmap_cache() 主动清除
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
user_id: 用户ID，None 表示获取全局热力图（管理员）
|
||||
include_actual_cost: 是否包含实际成本
|
||||
|
||||
Returns:
|
||||
热力图数据字典
|
||||
"""
|
||||
import json
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.config.constants import CacheTTL
|
||||
|
||||
cache_key = cls._get_heatmap_cache_key(user_id, include_actual_cost)
|
||||
|
||||
cache_ttl = CacheTTL.ACTIVITY_HEATMAP
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
# 尝试从缓存获取
|
||||
if redis_client:
|
||||
try:
|
||||
cached = await redis_client.get(cache_key)
|
||||
if cached:
|
||||
try:
|
||||
return json.loads(cached) # type: ignore[no-any-return]
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(
|
||||
"热力图缓存解析失败，删除损坏缓存: {}, error={}", cache_key, e
|
||||
)
|
||||
try:
|
||||
await redis_client.delete(cache_key)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error("读取热力图缓存出错: {}, error={}", cache_key, e)
|
||||
|
||||
# 从数据库查询
|
||||
result = cls.get_daily_activity(
|
||||
db=db,
|
||||
user_id=user_id,
|
||||
window_days=365,
|
||||
include_actual_cost=include_actual_cost,
|
||||
)
|
||||
|
||||
# 保存到缓存（失败不影响返回结果）
|
||||
if redis_client:
|
||||
try:
|
||||
await redis_client.setex(
|
||||
cache_key,
|
||||
cache_ttl,
|
||||
json.dumps(result, ensure_ascii=False, default=str),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("保存热力图缓存失败: {}, error={}", cache_key, e)
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def check_request_balance_details(
|
||||
db: Session,
|
||||
user: User,
|
||||
estimated_tokens: int = 0,
|
||||
estimated_cost: float = 0,
|
||||
api_key: ApiKey | None = None,
|
||||
) -> RequestBalanceCheckResult:
|
||||
"""Return a structured balance-check result."""
|
||||
from src.services.wallet import WalletService
|
||||
|
||||
wallet_access = WalletService.check_request_allowed(
|
||||
db,
|
||||
user=None if (api_key and api_key.is_standalone) else user,
|
||||
api_key=api_key,
|
||||
)
|
||||
snapshot = wallet_access.balance_snapshot
|
||||
if snapshot is None:
|
||||
snapshot = wallet_access.remaining
|
||||
remaining = float(snapshot) if snapshot is not None else None
|
||||
if wallet_access.allowed:
|
||||
return RequestBalanceCheckResult(True, "OK", remaining)
|
||||
|
||||
if wallet_access.message in {"钱包欠费，请先充值", "账户欠费，请先充值"}:
|
||||
if api_key and api_key.is_standalone:
|
||||
return RequestBalanceCheckResult(False, "Key欠费，请先调账或充值", remaining)
|
||||
return RequestBalanceCheckResult(False, "账户欠费，请先充值", remaining)
|
||||
|
||||
if wallet_access.message == "钱包不可用":
|
||||
if api_key and api_key.is_standalone:
|
||||
return RequestBalanceCheckResult(False, "Key钱包不可用", remaining)
|
||||
return RequestBalanceCheckResult(False, "钱包不可用", remaining)
|
||||
|
||||
if api_key and api_key.is_standalone:
|
||||
if remaining is None:
|
||||
return RequestBalanceCheckResult(False, "Key余额不足", remaining)
|
||||
return RequestBalanceCheckResult(
|
||||
False, f"Key余额不足（剩余: ${remaining:.2f}）", remaining
|
||||
)
|
||||
|
||||
# Admin users are already allowed in WalletService.check_request_allowed.
|
||||
if remaining is None:
|
||||
return RequestBalanceCheckResult(False, wallet_access.message or "余额不足", remaining)
|
||||
return RequestBalanceCheckResult(False, f"余额不足（剩余: ${remaining:.2f}）", remaining)
|
||||
|
||||
@staticmethod
|
||||
def check_request_balance(
|
||||
db: Session,
|
||||
user: User,
|
||||
estimated_tokens: int = 0,
|
||||
estimated_cost: float = 0,
|
||||
api_key: ApiKey | None = None,
|
||||
) -> tuple[bool, str]:
|
||||
"""Check whether the request passes balance rules."""
|
||||
result = UsageQueryMixin.check_request_balance_details(
|
||||
db,
|
||||
user,
|
||||
estimated_tokens=estimated_tokens,
|
||||
estimated_cost=estimated_cost,
|
||||
api_key=api_key,
|
||||
)
|
||||
return result.allowed, result.message
|
||||
|
||||
@staticmethod
|
||||
def get_usage_summary(
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
group_by: str | None = "day", # day, week, month, None(不按时间分桶)
|
||||
) -> list[dict[str, Any]]:
|
||||
"""获取使用汇总"""
|
||||
|
||||
query = db.query(Usage)
|
||||
# 过滤掉 pending/streaming 状态的请求（尚未完成的请求不应计入统计）
|
||||
query = query.filter(Usage.status.notin_(["pending", "streaming"]))
|
||||
|
||||
if user_id:
|
||||
query = query.filter(Usage.user_id == user_id)
|
||||
if api_key_id:
|
||||
query = query.filter(Usage.api_key_id == api_key_id)
|
||||
if start_date:
|
||||
query = query.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(Usage.created_at < end_date)
|
||||
|
||||
select_columns = [Usage.provider_name, Usage.model]
|
||||
group_columns = [Usage.provider_name, Usage.model]
|
||||
|
||||
if group_by is not None:
|
||||
from src.utils.database_helpers import date_trunc_portable
|
||||
|
||||
bind = db.bind
|
||||
dialect = bind.dialect.name if bind is not None else "sqlite"
|
||||
|
||||
if group_by == "day":
|
||||
date_func = date_trunc_portable(dialect, "day", Usage.created_at)
|
||||
elif group_by == "week":
|
||||
date_func = date_trunc_portable(dialect, "week", Usage.created_at)
|
||||
elif group_by == "month":
|
||||
date_func = date_trunc_portable(dialect, "month", Usage.created_at)
|
||||
else:
|
||||
date_func = date_trunc_portable(dialect, "day", Usage.created_at)
|
||||
select_columns.insert(0, date_func.label("period"))
|
||||
group_columns.insert(0, date_func)
|
||||
|
||||
summary = db.query(
|
||||
*select_columns,
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.total_tokens).label("total_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(input_context_expr()).label("total_input_context"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost_usd"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost_usd"),
|
||||
func.sum(case((Usage.status_code == 200, 1), else_=0)).label("success_count"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
func.sum(
|
||||
case(
|
||||
(
|
||||
(Usage.status_code == 200) & Usage.response_time_ms.isnot(None),
|
||||
Usage.response_time_ms,
|
||||
),
|
||||
else_=0,
|
||||
)
|
||||
).label("success_response_time_sum"),
|
||||
func.sum(
|
||||
case(
|
||||
((Usage.status_code == 200) & Usage.response_time_ms.isnot(None), 1),
|
||||
else_=0,
|
||||
)
|
||||
).label("success_response_time_count"),
|
||||
)
|
||||
|
||||
# 过滤掉 pending/streaming 状态的请求（与上方明细查询一致）
|
||||
summary = summary.filter(Usage.status.notin_(["pending", "streaming"]))
|
||||
|
||||
if user_id:
|
||||
summary = summary.filter(Usage.user_id == user_id)
|
||||
if api_key_id:
|
||||
summary = summary.filter(Usage.api_key_id == api_key_id)
|
||||
if start_date:
|
||||
summary = summary.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
summary = summary.filter(Usage.created_at < end_date)
|
||||
|
||||
summary = summary.group_by(*group_columns).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"period": getattr(row, "period", None),
|
||||
"provider": row.provider_name,
|
||||
"model": row.model,
|
||||
"requests": row.requests,
|
||||
"input_tokens": row.input_tokens,
|
||||
"output_tokens": row.output_tokens,
|
||||
"total_tokens": row.total_tokens,
|
||||
"cache_read_tokens": int(row.cache_read_tokens or 0),
|
||||
"cache_creation_tokens": int(row.cache_creation_tokens or 0),
|
||||
"total_input_context": int(row.total_input_context or 0),
|
||||
"total_cost_usd": float(row.total_cost_usd or 0.0),
|
||||
"actual_total_cost_usd": float(row.actual_total_cost_usd or 0.0),
|
||||
"success_count": int(row.success_count or 0),
|
||||
"avg_response_time_ms": (
|
||||
float(row.avg_response_time) if row.avg_response_time else 0
|
||||
),
|
||||
"success_response_time_sum_ms": float(row.success_response_time_sum or 0.0),
|
||||
"success_response_time_count": int(row.success_response_time_count or 0),
|
||||
}
|
||||
for row in summary
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def get_daily_activity(
|
||||
db: Session,
|
||||
user_id: str | None = None,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
window_days: int = 365,
|
||||
include_actual_cost: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""按天统计请求活跃度，用于渲染热力图。
|
||||
|
||||
优化策略：
|
||||
- 历史数据从预计算的 StatsDaily/StatsUserDaily 表读取
|
||||
- 只有"今天"的数据才实时查询 Usage 表
|
||||
"""
|
||||
|
||||
def ensure_timezone(value: datetime) -> datetime:
|
||||
if value.tzinfo is None:
|
||||
return value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
# 如果调用方未指定时间范围，则默认统计最近 window_days 天
|
||||
now = datetime.now(timezone.utc)
|
||||
end_dt = ensure_timezone(end_date) if end_date else now
|
||||
start_dt = (
|
||||
ensure_timezone(start_date) if start_date else end_dt - timedelta(days=window_days - 1)
|
||||
)
|
||||
|
||||
# 对齐到自然日的开始/结束
|
||||
start_dt = datetime.combine(start_dt.date(), datetime.min.time(), tzinfo=timezone.utc)
|
||||
end_dt = datetime.combine(end_dt.date(), datetime.max.time(), tzinfo=timezone.utc)
|
||||
|
||||
today = now.date()
|
||||
today_start_dt = datetime.combine(today, datetime.min.time(), tzinfo=timezone.utc)
|
||||
aggregated: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# 1. 从预计算表读取历史数据（不包括今天）
|
||||
if user_id:
|
||||
from src.models.database import StatsUserDaily
|
||||
|
||||
hist_query = db.query(StatsUserDaily).filter(
|
||||
StatsUserDaily.user_id == user_id,
|
||||
StatsUserDaily.date >= start_dt,
|
||||
StatsUserDaily.date < today_start_dt,
|
||||
)
|
||||
for row in hist_query.all():
|
||||
key = (
|
||||
row.date.date().isoformat()
|
||||
if isinstance(row.date, datetime)
|
||||
else str(row.date)[:10]
|
||||
)
|
||||
aggregated[key] = {
|
||||
"requests": row.total_requests or 0,
|
||||
"total_tokens": (
|
||||
(row.input_tokens or 0)
|
||||
+ (row.output_tokens or 0)
|
||||
+ (row.cache_creation_tokens or 0)
|
||||
+ (row.cache_read_tokens or 0)
|
||||
),
|
||||
"total_cost_usd": float(row.total_cost or 0.0),
|
||||
}
|
||||
# StatsUserDaily 没有 actual_total_cost 字段，用户视图不需要倍率成本
|
||||
else:
|
||||
from src.models.database import StatsDaily
|
||||
|
||||
hist_query = db.query(StatsDaily).filter(
|
||||
StatsDaily.date >= start_dt,
|
||||
StatsDaily.date < today_start_dt,
|
||||
)
|
||||
for row in hist_query.all():
|
||||
key = (
|
||||
row.date.date().isoformat()
|
||||
if isinstance(row.date, datetime)
|
||||
else str(row.date)[:10]
|
||||
)
|
||||
aggregated[key] = {
|
||||
"requests": row.total_requests or 0,
|
||||
"total_tokens": (
|
||||
(row.input_tokens or 0)
|
||||
+ (row.output_tokens or 0)
|
||||
+ (row.cache_creation_tokens or 0)
|
||||
+ (row.cache_read_tokens or 0)
|
||||
),
|
||||
"total_cost_usd": float(row.total_cost or 0.0),
|
||||
}
|
||||
if include_actual_cost:
|
||||
aggregated[key]["actual_total_cost_usd"] = float(
|
||||
row.actual_total_cost or 0.0 # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
# 2. 实时查询今天的数据（如果在查询范围内）
|
||||
if today >= start_dt.date() and today <= end_dt.date():
|
||||
today_start = datetime.combine(today, datetime.min.time(), tzinfo=timezone.utc)
|
||||
today_end = datetime.combine(today, datetime.max.time(), tzinfo=timezone.utc)
|
||||
|
||||
if include_actual_cost:
|
||||
today_query = db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("total_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost_usd"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost_usd"),
|
||||
).filter(
|
||||
Usage.created_at >= today_start,
|
||||
Usage.created_at <= today_end,
|
||||
)
|
||||
else:
|
||||
today_query = db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("total_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost_usd"),
|
||||
).filter(
|
||||
Usage.created_at >= today_start,
|
||||
Usage.created_at <= today_end,
|
||||
)
|
||||
|
||||
if user_id:
|
||||
today_query = today_query.filter(Usage.user_id == user_id)
|
||||
|
||||
today_row = today_query.first()
|
||||
if today_row and today_row.requests:
|
||||
aggregated[today.isoformat()] = {
|
||||
"requests": int(today_row.requests or 0),
|
||||
"total_tokens": int(today_row.total_tokens or 0),
|
||||
"total_cost_usd": float(today_row.total_cost_usd or 0.0),
|
||||
}
|
||||
if include_actual_cost:
|
||||
aggregated[today.isoformat()]["actual_total_cost_usd"] = float(
|
||||
today_row.actual_total_cost_usd or 0.0
|
||||
)
|
||||
|
||||
# 3. 构建返回结果
|
||||
days: list[dict[str, Any]] = []
|
||||
cursor = start_dt.date()
|
||||
end_date_only = end_dt.date()
|
||||
max_requests = 0
|
||||
|
||||
while cursor <= end_date_only:
|
||||
iso_date = cursor.isoformat()
|
||||
stats = aggregated.get(iso_date, {})
|
||||
requests = stats.get("requests", 0)
|
||||
total_tokens = stats.get("total_tokens", 0)
|
||||
total_cost = stats.get("total_cost_usd", 0.0)
|
||||
|
||||
entry: dict[str, Any] = {
|
||||
"date": iso_date,
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": total_cost,
|
||||
}
|
||||
|
||||
if include_actual_cost:
|
||||
entry["actual_total_cost"] = stats.get("actual_total_cost_usd", 0.0)
|
||||
|
||||
days.append(entry)
|
||||
max_requests = max(max_requests, requests)
|
||||
cursor += timedelta(days=1)
|
||||
|
||||
return {
|
||||
"start_date": start_dt.date().isoformat(),
|
||||
"end_date": end_dt.date().isoformat(),
|
||||
"total_days": len(days),
|
||||
"max_requests": max_requests,
|
||||
"days": days,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_top_users(
|
||||
db: Session,
|
||||
limit: int = 10,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
order_by: str = "cost", # cost, tokens, requests
|
||||
) -> list[dict[str, Any]]:
|
||||
"""获取使用量最高的用户"""
|
||||
|
||||
query = (
|
||||
db.query(
|
||||
User.id,
|
||||
User.email,
|
||||
User.username,
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_tokens).label("tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("cost_usd"),
|
||||
)
|
||||
.join(Usage, User.id == Usage.user_id)
|
||||
.filter(Usage.user_id.isnot(None))
|
||||
)
|
||||
|
||||
if start_date:
|
||||
query = query.filter(Usage.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(Usage.created_at <= end_date)
|
||||
|
||||
query = query.group_by(User.id, User.email, User.username)
|
||||
|
||||
# 排序
|
||||
if order_by == "cost":
|
||||
query = query.order_by(func.sum(Usage.total_cost_usd).desc())
|
||||
elif order_by == "tokens":
|
||||
query = query.order_by(func.sum(Usage.total_tokens).desc())
|
||||
else:
|
||||
query = query.order_by(func.count(Usage.id).desc())
|
||||
|
||||
results = query.limit(limit).all()
|
||||
|
||||
return [
|
||||
{
|
||||
"user_id": row.id,
|
||||
"email": row.email,
|
||||
"username": row.username,
|
||||
"requests": row.requests,
|
||||
"tokens": row.tokens,
|
||||
"cost_usd": float(row.cost_usd),
|
||||
}
|
||||
for row in results
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def cleanup_old_usage_records(
|
||||
db: Session, days_to_keep: int = 90, batch_size: int = 1000
|
||||
) -> int:
|
||||
"""清理旧的使用记录（分批删除避免长事务锁定）
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
days_to_keep: 保留天数，默认 90 天
|
||||
batch_size: 每批删除数量，默认 1000 条
|
||||
|
||||
Returns:
|
||||
删除的总记录数
|
||||
"""
|
||||
cutoff_date = datetime.now(timezone.utc) - timedelta(days=days_to_keep)
|
||||
total_deleted = 0
|
||||
|
||||
while True:
|
||||
# 查询待删除的 ID（使用新索引 idx_usage_user_created）
|
||||
batch_ids = (
|
||||
db.query(Usage.id).filter(Usage.created_at < cutoff_date).limit(batch_size).all()
|
||||
)
|
||||
|
||||
if not batch_ids:
|
||||
break
|
||||
|
||||
# 批量删除
|
||||
deleted_count = (
|
||||
db.query(Usage)
|
||||
.filter(Usage.id.in_([row.id for row in batch_ids]))
|
||||
.delete(synchronize_session=False)
|
||||
)
|
||||
db.commit()
|
||||
total_deleted += deleted_count
|
||||
|
||||
logger.debug("清理使用记录: 本批删除 {} 条", deleted_count)
|
||||
|
||||
logger.info("清理使用记录: 共删除 {} 条超过 {} 天的记录", total_deleted, days_to_keep)
|
||||
|
||||
return total_deleted
|
||||
162
_deprecated_py_src/services/usage/quota_scheduler.py
Normal file
162
_deprecated_py_src/services/usage/quota_scheduler.py
Normal file
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
额度周期重置定时任务
|
||||
|
||||
支持按天数周期重置额度：
|
||||
- quota_reset_day: 重置周期（天数），例如7=每周，30=每月
|
||||
- quota_last_reset_at: 上次重置时间，用于计算下次重置
|
||||
|
||||
使用统一的 TaskScheduler 进行调度。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from src.core.enums import ProviderBillingType
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import Provider
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
|
||||
|
||||
class QuotaScheduler:
|
||||
"""额度周期重置调度器"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.running = False
|
||||
|
||||
async def start(self) -> Any:
|
||||
"""启动调度器"""
|
||||
if self.running:
|
||||
logger.warning("Quota scheduler already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
logger.info("Quota scheduler started")
|
||||
|
||||
scheduler = get_scheduler()
|
||||
|
||||
# 每小时检查一次额度重置
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_quota_check,
|
||||
hours=1,
|
||||
job_id="quota_reset_check",
|
||||
name="额度周期重置检查",
|
||||
)
|
||||
|
||||
# 启动时立即执行一次检查
|
||||
await self._check_and_reset_quotas()
|
||||
|
||||
async def stop(self) -> Any:
|
||||
"""停止调度器"""
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
self.running = False
|
||||
scheduler = get_scheduler()
|
||||
scheduler.remove_job("quota_reset_check")
|
||||
logger.info("Quota scheduler stopped")
|
||||
|
||||
async def _scheduled_quota_check(self) -> None:
|
||||
"""额度检查任务（定时调用）"""
|
||||
if not self.running:
|
||||
return
|
||||
await self._check_and_reset_quotas()
|
||||
|
||||
async def _check_and_reset_quotas(self) -> None:
|
||||
"""检查并重置周期额度"""
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
# 获取所有定额类型的提供商
|
||||
providers = (
|
||||
db.query(Provider)
|
||||
.filter(
|
||||
Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA,
|
||||
Provider.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not providers:
|
||||
logger.debug("No quota providers to check")
|
||||
return
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
reset_count = 0
|
||||
|
||||
for provider in providers:
|
||||
try:
|
||||
# 如果没有上次重置时间，初始化为当前时间
|
||||
if provider.quota_last_reset_at is None:
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Initialized quota_last_reset_at for provider {provider.name}")
|
||||
continue
|
||||
|
||||
# 计算距离上次重置的天数
|
||||
days_since_reset = (now - provider.quota_last_reset_at).days
|
||||
|
||||
# 如果达到或超过重置周期，执行重置
|
||||
if days_since_reset >= provider.quota_reset_day:
|
||||
logger.info(f"Resetting quota for provider {provider.name}")
|
||||
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.quota_last_reset_at = now
|
||||
reset_count += 1
|
||||
|
||||
# 检查是否过期
|
||||
if provider.quota_expires_at and provider.quota_expires_at < now:
|
||||
logger.warning(f"Provider {provider.name} quota expired")
|
||||
# 可以选择禁用过期的提供商
|
||||
# provider.is_active = False
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Error processing provider {provider.name}: {e}")
|
||||
|
||||
if reset_count > 0:
|
||||
db.commit()
|
||||
logger.info(f"Reset quotas for {reset_count} providers")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def force_reset(self, provider_id: str | None = None) -> Any:
|
||||
"""手动强制重置额度"""
|
||||
db = create_session()
|
||||
try:
|
||||
now = datetime.now(timezone.utc)
|
||||
if provider_id:
|
||||
# 重置指定提供商
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if provider and provider.billing_type == ProviderBillingType.MONTHLY_QUOTA:
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Force reset quota for provider {provider.name}")
|
||||
else:
|
||||
# 重置所有定额提供商
|
||||
providers = (
|
||||
db.query(Provider)
|
||||
.filter(Provider.billing_type == ProviderBillingType.MONTHLY_QUOTA)
|
||||
.all()
|
||||
)
|
||||
for provider in providers:
|
||||
provider.monthly_used_usd = 0.0
|
||||
provider.quota_last_reset_at = now
|
||||
db.commit()
|
||||
logger.info(f"Force reset quotas for {len(providers)} providers")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
# 全局单例
|
||||
_quota_scheduler = None
|
||||
|
||||
|
||||
def get_quota_scheduler() -> QuotaScheduler:
|
||||
"""获取全局调度器实例"""
|
||||
global _quota_scheduler
|
||||
if _quota_scheduler is None:
|
||||
_quota_scheduler = QuotaScheduler()
|
||||
return _quota_scheduler
|
||||
261
_deprecated_py_src/services/usage/recorder.py
Normal file
261
_deprecated_py_src/services/usage/recorder.py
Normal file
@@ -0,0 +1,261 @@
|
||||
"""
|
||||
统一的 Usage 记录器
|
||||
|
||||
设计原则：
|
||||
1. 单一入口：所有 Usage 记录都通过 UsageRecorder
|
||||
2. 自动处理：根据 RequestResult 自动判断成功/失败
|
||||
3. 完整记录：确保所有必要字段都被记录
|
||||
4. 异步友好：支持后台异步记录，不阻塞主流程
|
||||
|
||||
使用方式：
|
||||
```python
|
||||
recorder = UsageRecorder(db, user, api_key)
|
||||
|
||||
# 记录成功请求
|
||||
await recorder.record_success(result)
|
||||
|
||||
# 记录失败请求
|
||||
await recorder.record_failure(result)
|
||||
|
||||
# 或者自动判断
|
||||
await recorder.record(result)
|
||||
```
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.api_format import filter_response_headers as filter_proxy_response_headers
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, User
|
||||
from src.services.request.result import RequestResult
|
||||
from src.services.system.audit import audit_service
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
|
||||
class UsageRecorder:
|
||||
"""
|
||||
统一的 Usage 记录器
|
||||
|
||||
职责：
|
||||
1. 记录成功请求的 Usage（包含 token 使用量和费用）
|
||||
2. 记录失败请求的 Usage（token=0，记录错误信息）
|
||||
3. 记录审计日志
|
||||
4. 更新用户配额
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: Session,
|
||||
user: User,
|
||||
api_key: ApiKey,
|
||||
client_ip: str = "unknown",
|
||||
request_id: str | None = None,
|
||||
):
|
||||
self.db = db
|
||||
self.user = user
|
||||
self.api_key = api_key
|
||||
self.client_ip = client_ip
|
||||
self.request_id = request_id
|
||||
|
||||
async def record(self, result: RequestResult) -> None:
|
||||
"""
|
||||
根据 RequestResult 自动判断并记录 Usage
|
||||
|
||||
Args:
|
||||
result: 请求结果
|
||||
"""
|
||||
if result.is_success:
|
||||
await self.record_success(result)
|
||||
else:
|
||||
await self.record_failure(result)
|
||||
|
||||
async def record_success(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录成功请求的 Usage
|
||||
|
||||
Args:
|
||||
result: 成功的请求结果
|
||||
request_headers: 原始请求头（可选，用于调试）
|
||||
request_body: 原始请求体（可选，用于调试）
|
||||
"""
|
||||
metadata = result.metadata
|
||||
usage = result.usage
|
||||
|
||||
# 确定 target_model：当存在 original_model 且与 model 不同时，说明发生了映射
|
||||
target_model = None
|
||||
if metadata.original_model and metadata.original_model != metadata.model:
|
||||
target_model = metadata.model
|
||||
|
||||
# 非流式成功时，返回给客户端的是提供商响应头（透传）+ content-type
|
||||
client_response_headers = filter_proxy_response_headers(metadata.provider_response_headers)
|
||||
client_response_headers["content-type"] = "application/json"
|
||||
|
||||
await UsageService.record_usage(
|
||||
db=self.db,
|
||||
user=self.user,
|
||||
api_key=self.api_key,
|
||||
provider=metadata.provider,
|
||||
model=metadata.original_model or metadata.model,
|
||||
target_model=target_model,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cache_creation_input_tokens=usage.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=usage.cache_read_input_tokens,
|
||||
request_type="chat",
|
||||
api_format=metadata.api_format,
|
||||
api_family=metadata.api_family,
|
||||
endpoint_kind=metadata.endpoint_kind,
|
||||
is_stream=result.is_stream,
|
||||
response_time_ms=result.response_time_ms,
|
||||
status_code=200,
|
||||
error_message=None,
|
||||
metadata=metadata.response_metadata if metadata.response_metadata else None,
|
||||
request_headers=request_headers or result.request_headers,
|
||||
request_body=request_body or result.request_body,
|
||||
provider_request_headers=metadata.provider_request_headers,
|
||||
response_headers=metadata.provider_response_headers,
|
||||
client_response_headers=client_response_headers,
|
||||
response_body=result.response_data if isinstance(result.response_data, dict) else {},
|
||||
request_id=self.request_id,
|
||||
provider_id=metadata.provider_id,
|
||||
provider_endpoint_id=metadata.provider_endpoint_id,
|
||||
provider_api_key_id=metadata.provider_api_key_id,
|
||||
status="completed", # 成功请求
|
||||
)
|
||||
|
||||
# 记录审计日志
|
||||
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 or "",
|
||||
model=metadata.original_model or metadata.model,
|
||||
provider=metadata.provider,
|
||||
success=True,
|
||||
ip_address=self.client_ip,
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[UsageRecorder] 成功记录: provider={metadata.provider}, "
|
||||
f"model={metadata.model}, api_format={metadata.api_format}, "
|
||||
f"tokens={usage.input_tokens}+{usage.output_tokens}"
|
||||
)
|
||||
|
||||
async def record_failure(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录失败请求的 Usage
|
||||
|
||||
Args:
|
||||
result: 失败的请求结果
|
||||
request_headers: 原始请求头
|
||||
request_body: 原始请求体
|
||||
"""
|
||||
metadata = result.metadata
|
||||
|
||||
# 确定 target_model：当存在 original_model 且与 model 不同时，说明发生了映射
|
||||
target_model = None
|
||||
if metadata.original_model and metadata.original_model != metadata.model:
|
||||
target_model = metadata.model
|
||||
|
||||
await UsageService.record_usage(
|
||||
db=self.db,
|
||||
user=self.user,
|
||||
api_key=self.api_key,
|
||||
provider=metadata.provider,
|
||||
model=metadata.original_model or metadata.model,
|
||||
target_model=target_model,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
request_type="chat",
|
||||
api_format=metadata.api_format,
|
||||
api_family=metadata.api_family,
|
||||
endpoint_kind=metadata.endpoint_kind,
|
||||
is_stream=result.is_stream,
|
||||
response_time_ms=result.response_time_ms,
|
||||
status_code=result.status_code,
|
||||
error_message=result.error_message,
|
||||
metadata=metadata.response_metadata if metadata.response_metadata else None,
|
||||
request_headers=request_headers or result.request_headers,
|
||||
request_body=request_body or result.request_body,
|
||||
provider_request_headers=metadata.provider_request_headers,
|
||||
response_headers={},
|
||||
# 失败请求返回给客户端的是 JSON 错误响应
|
||||
client_response_headers={"content-type": "application/json"},
|
||||
response_body={"error": result.error_message} if result.error_message else {},
|
||||
request_id=self.request_id,
|
||||
provider_id=metadata.provider_id,
|
||||
provider_endpoint_id=metadata.provider_endpoint_id,
|
||||
provider_api_key_id=metadata.provider_api_key_id,
|
||||
status="failed", # 失败请求
|
||||
)
|
||||
|
||||
# 记录审计日志
|
||||
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 or "",
|
||||
model=metadata.original_model or metadata.model,
|
||||
provider=metadata.provider,
|
||||
success=False,
|
||||
ip_address=self.client_ip,
|
||||
status_code=result.status_code,
|
||||
error_message=result.error_message,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[UsageRecorder] 失败记录: provider={metadata.provider}, "
|
||||
f"model={metadata.model}, api_format={metadata.api_format}, "
|
||||
f"status={result.status_code}, error={result.error_message[:100] if result.error_message else 'N/A'}"
|
||||
)
|
||||
|
||||
async def record_from_exception(
|
||||
self,
|
||||
exception: Exception,
|
||||
api_format: str,
|
||||
model: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
从异常创建 RequestResult 并记录失败
|
||||
|
||||
这是一个便捷方法，用于在异常处理中快速记录失败请求。
|
||||
|
||||
Args:
|
||||
exception: 捕获的异常
|
||||
api_format: API 格式（必须提供，确保始终有值）
|
||||
model: 模型名称
|
||||
response_time_ms: 响应时间
|
||||
is_stream: 是否流式请求
|
||||
request_headers: 原始请求头
|
||||
request_body: 原始请求体
|
||||
"""
|
||||
result = RequestResult.from_exception(
|
||||
exception=exception,
|
||||
api_format=api_format,
|
||||
model=model,
|
||||
response_time_ms=response_time_ms,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
result.request_headers = request_headers or {}
|
||||
result.request_body = request_body or {}
|
||||
|
||||
await self.record_failure(result, request_headers, request_body)
|
||||
1299
_deprecated_py_src/services/usage/recording.py
Normal file
1299
_deprecated_py_src/services/usage/recording.py
Normal file
File diff suppressed because it is too large
Load Diff
27
_deprecated_py_src/services/usage/service.py
Normal file
27
_deprecated_py_src/services/usage/service.py
Normal file
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
用量统计和配额管理服务
|
||||
"""
|
||||
|
||||
from src.services.usage._types import UsageRecordParams
|
||||
from src.services.usage.active_requests import UsageActiveRequestsMixin
|
||||
from src.services.usage.cache_analysis import UsageCacheAnalysisMixin
|
||||
from src.services.usage.lifecycle import UsageLifecycleMixin
|
||||
from src.services.usage.pricing import UsagePricingMixin
|
||||
from src.services.usage.query import UsageQueryMixin
|
||||
from src.services.usage.recording import UsageRecordingMixin
|
||||
|
||||
|
||||
class UsageService(
|
||||
UsagePricingMixin,
|
||||
UsageRecordingMixin,
|
||||
UsageLifecycleMixin,
|
||||
UsageQueryMixin,
|
||||
UsageActiveRequestsMixin,
|
||||
UsageCacheAnalysisMixin,
|
||||
):
|
||||
"""用量统计服务"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
__all__ = ["UsageService", "UsageRecordParams"]
|
||||
1250
_deprecated_py_src/services/usage/stream.py
Normal file
1250
_deprecated_py_src/services/usage/stream.py
Normal file
File diff suppressed because it is too large
Load Diff
371
_deprecated_py_src/services/usage/telemetry.py
Normal file
371
_deprecated_py_src/services/usage/telemetry.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""
|
||||
消息遥测记录器。
|
||||
|
||||
从 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,
|
||||
)
|
||||
309
_deprecated_py_src/services/usage/telemetry_writer.py
Normal file
309
_deprecated_py_src/services/usage/telemetry_writer.py
Normal file
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
Telemetry writer abstraction for stream usage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from src.clients.redis_client import get_usage_queue_redis_client as get_redis_client
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.services.usage.events import UsageEventType, build_usage_event
|
||||
from src.services.usage.telemetry import MessageTelemetry
|
||||
|
||||
|
||||
class TelemetryWriter(ABC):
|
||||
def supports_background_submission(self) -> bool:
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
async def record_success(self, **kwargs: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def record_failure(self, **kwargs: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def record_cancelled(self, **kwargs: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class DbTelemetryWriter(TelemetryWriter):
|
||||
"""通过 MessageTelemetry 写入数据库的 Writer"""
|
||||
|
||||
# MessageTelemetry 不支持的参数，需要过滤掉
|
||||
# - request_type: MessageTelemetry 内部固定为 "chat"，无需外部传入
|
||||
# - metadata: 由本 writer 映射到 request_metadata（用于落库追踪信息）
|
||||
_IGNORED_KWARGS = frozenset({"request_type"})
|
||||
|
||||
def __init__(self, telemetry: MessageTelemetry) -> None:
|
||||
self._telemetry = telemetry
|
||||
|
||||
def _filter_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""过滤掉 MessageTelemetry 不支持的参数"""
|
||||
out = {k: v for k, v in kwargs.items() if k not in self._IGNORED_KWARGS}
|
||||
# 兼容 stream 侧传入的 metadata 字段：映射到 MessageTelemetry 的 request_metadata
|
||||
if "metadata" in out and "request_metadata" not in out:
|
||||
out["request_metadata"] = out.get("metadata")
|
||||
out.pop("metadata", None)
|
||||
return out
|
||||
|
||||
async def record_success(self, **kwargs: Any) -> None:
|
||||
await self._telemetry.record_success(**self._filter_kwargs(kwargs))
|
||||
|
||||
async def record_failure(self, **kwargs: Any) -> None:
|
||||
await self._telemetry.record_failure(**self._filter_kwargs(kwargs))
|
||||
|
||||
async def record_cancelled(self, **kwargs: Any) -> None:
|
||||
await self._telemetry.record_cancelled(**self._filter_kwargs(kwargs))
|
||||
|
||||
|
||||
class QueueTelemetryWriter(TelemetryWriter):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
request_id: str,
|
||||
user_id: str,
|
||||
api_key_id: str,
|
||||
log_level: str = "basic",
|
||||
sensitive_headers: list[str] | None = None,
|
||||
max_request_body_size: int = 0,
|
||||
max_response_body_size: int = 0,
|
||||
) -> None:
|
||||
self.request_id = request_id
|
||||
self.user_id = user_id
|
||||
self.api_key_id = api_key_id
|
||||
self.log_level = (log_level or "basic").strip().lower()
|
||||
self._sensitive_headers = sensitive_headers or [
|
||||
"authorization",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"cookie",
|
||||
"set-cookie",
|
||||
]
|
||||
self._max_request_body_size = int(max_request_body_size or 0)
|
||||
self._max_response_body_size = int(max_response_body_size or 0)
|
||||
|
||||
@property
|
||||
def include_headers(self) -> bool:
|
||||
return self.log_level in {"headers", "full"}
|
||||
|
||||
@property
|
||||
def include_bodies(self) -> bool:
|
||||
return self.log_level == "full"
|
||||
|
||||
def supports_background_submission(self) -> bool:
|
||||
return True
|
||||
|
||||
async def record_success(self, **kwargs: Any) -> None:
|
||||
await self._publish_event(UsageEventType.COMPLETED, **kwargs)
|
||||
|
||||
async def record_failure(self, **kwargs: Any) -> None:
|
||||
await self._publish_event(UsageEventType.FAILED, **kwargs)
|
||||
|
||||
async def record_cancelled(self, **kwargs: Any) -> None:
|
||||
await self._publish_event(UsageEventType.CANCELLED, **kwargs)
|
||||
|
||||
async def _publish_event(self, event_type: UsageEventType, **kwargs: Any) -> None:
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if not redis_client:
|
||||
raise RuntimeError("Redis unavailable for usage queue")
|
||||
|
||||
data = self._build_event_data(**kwargs)
|
||||
event = build_usage_event(
|
||||
event_type=event_type,
|
||||
request_id=self.request_id,
|
||||
data=data,
|
||||
)
|
||||
maxlen = config.usage_queue_stream_maxlen
|
||||
try:
|
||||
if maxlen > 0:
|
||||
await redis_client.xadd(
|
||||
config.usage_queue_stream_key,
|
||||
event.to_stream_fields(),
|
||||
maxlen=maxlen,
|
||||
approximate=True,
|
||||
)
|
||||
else:
|
||||
await redis_client.xadd(config.usage_queue_stream_key, event.to_stream_fields())
|
||||
except Exception as exc:
|
||||
logger.error("[usage-queue] XADD failed: {}", exc)
|
||||
raise
|
||||
|
||||
def _mask_headers(self, headers: Any) -> Any:
|
||||
"""Mask sensitive headers before putting them into Redis."""
|
||||
if not isinstance(headers, dict) or not headers:
|
||||
return headers
|
||||
sensitive = {h.lower() for h in self._sensitive_headers if isinstance(h, str) and h}
|
||||
if not sensitive:
|
||||
return headers
|
||||
out: dict[str, Any] = {}
|
||||
for k, v in headers.items():
|
||||
key = str(k)
|
||||
if key.lower() in sensitive:
|
||||
s = str(v)
|
||||
if len(s) > 8:
|
||||
out[key] = s[:4] + "****" + s[-4:]
|
||||
else:
|
||||
out[key] = "****"
|
||||
else:
|
||||
out[key] = v
|
||||
return out
|
||||
|
||||
def _truncate_body(self, value: Any, *, max_size: int, is_request: bool) -> Any:
|
||||
"""Best-effort truncate body based on SystemConfigService max_*_body_size."""
|
||||
if value is None:
|
||||
return None
|
||||
limit = int(max_size or 0)
|
||||
if limit <= 0:
|
||||
return value
|
||||
|
||||
body_str = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
|
||||
if len(body_str) <= limit:
|
||||
return value
|
||||
|
||||
# Match SystemConfigService.truncate_body contract.
|
||||
if isinstance(value, (dict, list)):
|
||||
return {
|
||||
"_truncated": True,
|
||||
"_original_size": len(body_str),
|
||||
"_content": body_str[:limit],
|
||||
}
|
||||
kind = "request" if is_request else "response"
|
||||
return (
|
||||
body_str[:limit]
|
||||
+ f"\n... (truncated {kind} body, original size: {len(body_str)} bytes)"
|
||||
)
|
||||
|
||||
def _build_event_data(self, **kwargs: Any) -> dict[str, Any]:
|
||||
# 必需字段
|
||||
data: dict[str, Any] = {
|
||||
"request_id": self.request_id,
|
||||
"user_id": self.user_id,
|
||||
"api_key_id": self.api_key_id,
|
||||
}
|
||||
|
||||
# 可选字段 - 只添加非 None/非默认值，减少 payload 大小
|
||||
# 注意：消费者端需要处理缺失字段的默认值
|
||||
if kwargs.get("provider"):
|
||||
data["provider"] = kwargs["provider"]
|
||||
if kwargs.get("model"):
|
||||
data["model"] = kwargs["model"]
|
||||
if kwargs.get("target_model"):
|
||||
data["target_model"] = kwargs["target_model"]
|
||||
|
||||
# Token 计数 - 0 是常见值，但仍需传递
|
||||
input_tokens = kwargs.get("input_tokens", 0)
|
||||
output_tokens = kwargs.get("output_tokens", 0)
|
||||
if input_tokens:
|
||||
data["input_tokens"] = input_tokens
|
||||
if output_tokens:
|
||||
data["output_tokens"] = output_tokens
|
||||
|
||||
# 缓存 token（cache_creation_tokens -> cache_creation_input_tokens 映射）
|
||||
cache_creation = kwargs.get("cache_creation_tokens", 0)
|
||||
cache_read = kwargs.get("cache_read_tokens", 0)
|
||||
if cache_creation:
|
||||
data["cache_creation_input_tokens"] = cache_creation
|
||||
if cache_read:
|
||||
data["cache_read_input_tokens"] = cache_read
|
||||
|
||||
# 缓存 5m/1h 细分
|
||||
cache_creation_5m = kwargs.get("cache_creation_tokens_5m", 0)
|
||||
cache_creation_1h = kwargs.get("cache_creation_tokens_1h", 0)
|
||||
if cache_creation_5m:
|
||||
data["cache_creation_input_tokens_5m"] = cache_creation_5m
|
||||
if cache_creation_1h:
|
||||
data["cache_creation_input_tokens_1h"] = cache_creation_1h
|
||||
|
||||
# 时间指标
|
||||
if kwargs.get("response_time_ms") is not None:
|
||||
data["response_time_ms"] = kwargs["response_time_ms"]
|
||||
if kwargs.get("first_byte_time_ms") is not None:
|
||||
data["first_byte_time_ms"] = kwargs["first_byte_time_ms"]
|
||||
|
||||
# 状态信息
|
||||
status_code = kwargs.get("status_code", 200)
|
||||
if status_code != 200:
|
||||
data["status_code"] = status_code
|
||||
if kwargs.get("error_message"):
|
||||
data["error_message"] = kwargs["error_message"]
|
||||
|
||||
# 格式信息
|
||||
request_type = kwargs.get("request_type", "chat")
|
||||
if request_type != "chat":
|
||||
data["request_type"] = request_type
|
||||
if kwargs.get("api_format"):
|
||||
data["api_format"] = kwargs["api_format"]
|
||||
if kwargs.get("api_family"):
|
||||
data["api_family"] = kwargs["api_family"]
|
||||
if kwargs.get("endpoint_kind"):
|
||||
data["endpoint_kind"] = kwargs["endpoint_kind"]
|
||||
if kwargs.get("endpoint_api_format"):
|
||||
data["endpoint_api_format"] = kwargs["endpoint_api_format"]
|
||||
if kwargs.get("has_format_conversion"):
|
||||
data["has_format_conversion"] = True
|
||||
|
||||
# 流式标记 - 默认 True，只记录 False
|
||||
if not kwargs.get("is_stream", True):
|
||||
data["is_stream"] = False
|
||||
|
||||
# Provider 追踪
|
||||
if kwargs.get("provider_id"):
|
||||
data["provider_id"] = kwargs["provider_id"]
|
||||
if kwargs.get("provider_endpoint_id"):
|
||||
data["provider_endpoint_id"] = kwargs["provider_endpoint_id"]
|
||||
if kwargs.get("provider_api_key_id"):
|
||||
data["provider_api_key_id"] = kwargs["provider_api_key_id"]
|
||||
|
||||
# 元数据
|
||||
if kwargs.get("metadata"):
|
||||
data["metadata"] = kwargs["metadata"]
|
||||
|
||||
# Optional: Headers (masked)
|
||||
if self.include_headers:
|
||||
for _hdr_key in (
|
||||
"request_headers",
|
||||
"provider_request_headers",
|
||||
"response_headers",
|
||||
"client_response_headers",
|
||||
):
|
||||
if kwargs.get(_hdr_key) is not None:
|
||||
data[_hdr_key] = self._mask_headers(kwargs[_hdr_key])
|
||||
|
||||
# Optional: Bodies (truncated)
|
||||
if self.include_bodies:
|
||||
request_body = self._truncate_body(
|
||||
kwargs.get("request_body"),
|
||||
max_size=self._max_request_body_size,
|
||||
is_request=True,
|
||||
)
|
||||
provider_request_body = self._truncate_body(
|
||||
kwargs.get("provider_request_body"),
|
||||
max_size=self._max_request_body_size,
|
||||
is_request=True,
|
||||
)
|
||||
response_body = self._truncate_body(
|
||||
kwargs.get("response_body"),
|
||||
max_size=self._max_response_body_size,
|
||||
is_request=False,
|
||||
)
|
||||
client_response_body = self._truncate_body(
|
||||
kwargs.get("client_response_body"),
|
||||
max_size=self._max_response_body_size,
|
||||
is_request=False,
|
||||
)
|
||||
if request_body is not None:
|
||||
data["request_body"] = request_body
|
||||
if provider_request_body is not None:
|
||||
data["provider_request_body"] = provider_request_body
|
||||
if response_body is not None:
|
||||
data["response_body"] = response_body
|
||||
if client_response_body is not None:
|
||||
data["client_response_body"] = client_response_body
|
||||
|
||||
return data
|
||||
Reference in New Issue
Block a user