refactor: 拆分大型模块为 mixin/子模块结构

- cli_handler_base.py 拆分为 7 个 mixin (event/monitor/prefetch/request/sse_helpers/stream/sync)
- usage/service.py 拆分为 6 个子模块 (types/active_requests/cache_analysis/lifecycle/pricing/query/recording)
- models/database 拆分为独立模型文件 (auth/misc/model/provider/stats/usage/user)
- DUMMY_THOUGHT_SIGNATURE 常量提升到 core/api_format/conversion/constants 统一管理
- task/service.py 内联导入提升为顶层导入
- 流处理函数签名移除冗余的 http_client 参数
This commit is contained in:
fawney19
2026-02-14 12:55:58 +08:00
parent f464f32e48
commit 26ede849e2
32 changed files with 10173 additions and 7218 deletions

View File

@@ -62,6 +62,7 @@ def update_user_agent_version(version: str) -> None:
# Backward compat: keep module-level constant in sync.
HTTP_USER_AGENT = f"antigravity/{_ua_version} {_PLATFORM_TAG}"
def parse_version_string(text: str) -> str | None:
"""从任意文本中提取 X.Y.Z 格式的版本号。"""
m = _VERSION_RE.search(text)
@@ -72,7 +73,9 @@ def parse_version_string(text: str) -> str | None:
URL_UNAVAILABLE_TTL_SECONDS = 300 # 5 分钟
# ============== Thinking Signature ==============
DUMMY_THOUGHT_SIGNATURE = "skip_thought_signature_validator"
# 统一从 core 层导入,避免多处定义
from src.core.api_format.conversion.constants import DUMMY_THOUGHT_SIGNATURE # noqa: E402
MIN_SIGNATURE_LENGTH = 50 # 与 Antigravity-Manager 对齐
# ============== Thinking Budget ==============

View File

@@ -3,22 +3,52 @@ from __future__ import annotations
import re
from collections.abc import Callable
from typing import Any
from uuid import uuid4
import httpx
from sqlalchemy.orm import Session
from src.config.settings import config
from src.core.error_utils import extract_error_message
from src.core.exceptions import (
ConcurrencyLimitError,
EmbeddedErrorException,
ProviderNotAvailableException,
ProxyNodeUnavailableError,
ThinkingSignatureException,
UpstreamClientException,
)
from src.core.logger import logger
from src.core.provider_types import ProviderType
from src.models.database import ApiKey
from src.models.database import (
ApiKey,
Provider,
ProviderAPIKey,
ProviderEndpoint,
RequestCandidate,
Usage,
User,
VideoTask,
)
from src.services.cache.aware_scheduler import (
CacheAwareScheduler,
get_cache_aware_scheduler,
)
from src.services.candidate.failover import FailoverEngine
from src.services.candidate.policy import RetryPolicy, SkipPolicy
from src.services.candidate.recorder import CandidateRecorder
from src.services.candidate.resolver import CandidateResolver
from src.services.orchestration.error_classifier import ErrorClassifier
from src.services.orchestration.request_dispatcher import RequestDispatcher
from src.services.provider.format import normalize_endpoint_signature
from src.services.request.candidate import RequestCandidateService
from src.services.request.result import RequestMetadata
from src.services.system.config import SystemConfigService
from src.services.task.context import TaskMode
from src.services.task.exceptions import TaskNotFoundError
from src.services.task.protocol import AttemptKind, AttemptResult
from src.services.task.schema import ExecutionResult, TaskStatusResult
from src.services.usage.service import UsageService
_SENSITIVE_PATTERN = re.compile(
r"(api[_-]?key|token|bearer|authorization)[=:\s]+\S+",
@@ -173,21 +203,9 @@ class TaskService:
- RequestDispatcher execution
- Error classification/rectify logic ported from the previous SYNC implementation
"""
from uuid import uuid4
from src.models.database import User
from src.services.cache.aware_scheduler import (
CacheAwareScheduler,
get_cache_aware_scheduler,
)
from src.services.candidate.resolver import CandidateResolver
from src.services.orchestration.error_classifier import ErrorClassifier
from src.services.orchestration.request_dispatcher import RequestDispatcher
from src.services.rate_limit.adaptive_rpm import get_adaptive_rpm_manager
from src.services.rate_limit.concurrency_manager import get_concurrency_manager
from src.services.request.executor import RequestExecutor
from src.services.system.config import SystemConfigService
from src.services.usage.service import UsageService
if not request_id:
request_id = str(uuid4())
@@ -432,8 +450,6 @@ class TaskService:
if not error or not candidate:
return
from src.services.request.result import RequestMetadata
existing_metadata = getattr(error, "request_metadata", None)
if existing_metadata and getattr(existing_metadata, "api_format", None):
return
@@ -469,10 +485,6 @@ class TaskService:
last_error: Exception | None = None,
) -> None:
"""Raise a unified 'all candidates failed' exception."""
import httpx
from src.core.exceptions import ProviderNotAvailableException
logger.error(" [{}] 所有 {} 个组合均失败", request_id, max_attempts)
request_metadata = None
@@ -530,8 +542,6 @@ class TaskService:
extra_data: dict[str, Any],
) -> None:
"""Mark ThinkingSignatureException as failed for the candidate."""
from src.core.exceptions import ThinkingSignatureException
if not isinstance(error, ThinkingSignatureException):
return
@@ -560,7 +570,6 @@ class TaskService:
request_body_ref: dict[str, Any] | None,
) -> str:
"""Try to rectify thinking signature errors and request a retry."""
from src.core.exceptions import ThinkingSignatureException
from src.services.message.thinking_rectifier import ThinkingRectifier
if not isinstance(converted_error, ThinkingSignatureException):
@@ -697,17 +706,7 @@ class TaskService:
- "break": move to next candidate
- "raise": raise the underlying exception
"""
import httpx
from src.core.api_format.conversion.exceptions import FormatConversionError
from src.core.error_utils import extract_error_message
from src.core.exceptions import (
ConcurrencyLimitError,
EmbeddedErrorException,
ProxyNodeUnavailableError,
ThinkingSignatureException,
UpstreamClientException,
)
from src.services.proxy_node.resolver import resolve_effective_proxy, resolve_proxy_info
from src.services.request.executor import ExecutionError
@@ -980,20 +979,14 @@ class TaskService:
"""
from datetime import datetime, timezone
import httpx
from sqlalchemy import update
from src.models.database import RequestCandidate
from src.services.billing.rule_service import BillingRuleLookupResult, BillingRuleService
from src.services.cache.aware_scheduler import ProviderCandidate, get_cache_aware_scheduler
from src.services.candidate.resolver import CandidateResolver
from src.services.candidate.submit import (
AllCandidatesFailedError,
SubmitOutcome,
UpstreamClientRequestError,
)
from src.services.orchestration.error_classifier import ErrorClassifier
from src.services.system.config import SystemConfigService
def _sanitize(message: str, max_length: int = 200) -> str:
if not message:
@@ -1131,8 +1124,6 @@ class TaskService:
continue
# 2. global switch (from database config)
from src.services.system.config import SystemConfigService
if not SystemConfigService.is_format_conversion_enabled(self.db):
skip_reason = "format_conversion_disabled"
candidate_info.update(
@@ -1398,8 +1389,6 @@ class TaskService:
- internal UUID (VideoTask.id)
- external operation id (VideoTask.short_id)
"""
from src.models.database import VideoTask
task = (
self.db.query(VideoTask)
.filter(VideoTask.id == task_id, VideoTask.user_id == user_id)
@@ -1498,9 +1487,7 @@ class TaskService:
)
from src.core.api_format.conversion.internal_video import VideoStatus
from src.core.crypto import crypto_service
from src.models.database import ProviderAPIKey, ProviderEndpoint
from src.services.provider.transport import build_provider_url
from src.services.usage.service import UsageService
try:
task = self._get_video_task_for_user(task_id, user_id=user_id)
@@ -1617,9 +1604,6 @@ class TaskService:
This keeps behavior compatible with the old Phase2 finalize logic.
"""
from src.models.database import ApiKey, Provider, User
from src.services.usage.service import UsageService
user_obj = self.db.query(User).filter(User.id == task.user_id).first()
api_key_obj = (
self.db.query(ApiKey).filter(ApiKey.id == task.api_key_id).first()
@@ -1707,11 +1691,9 @@ class TaskService:
from datetime import datetime, timezone
from src.core.api_format.conversion.internal_video import VideoStatus
from src.models.database import Usage
from src.services.billing.dimension_collector_service import DimensionCollectorService
from src.services.billing.formula_engine import BillingIncompleteError, FormulaEngine
from src.services.billing.rule_service import BillingRuleService
from src.services.usage.service import UsageService
request_id = getattr(task, "request_id", None) or getattr(task, "id", None)
if not request_id:
@@ -1916,8 +1898,6 @@ class TaskService:
async def finalize(self, task_id: str) -> bool:
"""Finalize a task by internal id (best-effort)."""
from src.models.database import VideoTask
task = self.db.query(VideoTask).filter(VideoTask.id == task_id).first()
if not task:
return False

View File

@@ -0,0 +1,81 @@
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
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
response_headers: dict[str, Any] | None
client_response_headers: dict[str, Any] | None
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
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}")

View File

@@ -0,0 +1,331 @@
from __future__ import annotations
from datetime import datetime, timedelta, timezone
from typing import Any
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:
"""活跃请求管理方法"""
@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,
) -> int:
"""
清理超时的 pending/streaming 请求
将超过指定时间仍处于 pending 或 streaming 状态的请求标记为 failed。
这些请求可能是由于网络问题、服务重启或其他异常导致未能正常完成。
Args:
db: 数据库会话
timeout_minutes: 超时时间(分钟),默认 10 分钟
Returns:
清理的记录数
"""
cutoff_time = datetime.now(timezone.utc) - timedelta(minutes=timeout_minutes)
# 查找超时的请求
stale_requests = (
db.query(Usage)
.filter(
Usage.status.in_(["pending", "streaming"]),
Usage.created_at < cutoff_time,
)
.all()
)
count = 0
for usage in stale_requests:
old_status = usage.status
usage.status = "failed"
usage.error_message = f"请求超时: 状态 '{old_status}' 超过 {timeout_minutes} 分钟未完成"
usage.status_code = 504 # Gateway Timeout
count += 1
if count > 0:
db.commit()
logger.info(
f"清理超时请求: 将 {count} 条超过 {timeout_minutes} 分钟的 pending/streaming 请求标记为 failed"
)
return count
@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 (
db.query(Usage)
.filter(
Usage.status.in_(["pending", "streaming"]),
Usage.created_at < cutoff_time,
)
.count()
)
@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,
) -> list[dict[str, Any]]:
"""
获取活跃请求状态(用于前端轮询),并自动清理超时的 pending/streaming 请求
与 get_active_requests 不同,此方法:
1. 返回轻量级的状态字典而非完整 Usage 对象
2. 自动检测并清理超时的 pending/streaming 请求
3. 支持按 ID 列表查询特定请求
Args:
db: 数据库会话
ids: 指定要查询的请求 ID 列表(可选)
user_id: 限制只查询该用户的请求(可选,用于普通用户接口)
default_timeout_seconds: 默认超时时间(秒),当端点未配置时使用
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()
# 检查超时的 pending/streaming 请求
# 收集可能超时的 usage_id 列表
timeout_candidates: list[str] = []
for r in records:
if r.status in ("pending", "streaming") and r.created_at:
# 使用全局配置的超时时间
timeout_seconds = default_timeout_seconds
# 处理时区:如果 created_at 没有时区信息,假定为 UTC
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:
# 需要获取 request_id 以便检查 RequestCandidate 表
# r.id 是 usage_id需要查询 request_id
timeout_candidates.append(r.id)
# 批量更新超时的请求(排除已有成功完成记录的请求)
timeout_ids = []
if timeout_candidates:
# 检查 RequestCandidate 表是否有成功完成的记录
# 如果流已经成功完成stream_completed: true不应该标记为超时
# 先获取这些 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
# 包括两种情况:
# 1. status='success' 且 stream_completed=True正常完成
# 2. status='streaming' 且 status_code=200流传输中但 Provider 已返回 200可能是服务重启导致回调丢失
completed_usage_ids = set()
if request_ids:
from sqlalchemy import or_
candidates = (
db.query(
RequestCandidate.request_id,
RequestCandidate.status,
RequestCandidate.status_code,
RequestCandidate.extra_data,
)
.filter(
RequestCandidate.request_id.in_(request_ids),
or_(
RequestCandidate.status == "success",
# streaming 状态且 status_code=200说明 Provider 响应成功
# 但流传输可能因服务重启而中断
(RequestCandidate.status == "streaming")
& (RequestCandidate.status_code == 200),
),
)
.all()
)
for candidate in candidates:
extra_data = candidate.extra_data or {}
# 情况1status='success' 且 stream_completed=True
if candidate.status == "success" and extra_data.get("stream_completed", False):
usage_id = request_id_to_usage_id.get(candidate.request_id)
if usage_id:
completed_usage_ids.add(usage_id)
# 情况2status='streaming' 且 status_code=200
# 这表示 Provider 返回了 200但流传输可能因服务重启而未正常结束
# 此时应该恢复为 completed 而不是标记为 failed
elif candidate.status == "streaming" and candidate.status_code == 200:
usage_id = request_id_to_usage_id.get(candidate.request_id)
if usage_id:
completed_usage_ids.add(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,
)
db.commit()
logger.info(
f"[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

View 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

View File

@@ -0,0 +1,524 @@
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.system.config import SystemConfigService
class UsageLifecycleMixin:
"""使用记录生命周期管理方法"""
@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:
processed_request_headers = SystemConfigService.mask_sensitive_headers(
db, request_headers
)
# 处理请求体
processed_request_body = None
if should_log_body and request_body:
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,
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,
) -> bool:
"""
并发安全的幂等 finalizesettled
约定:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
cost = float(total_cost_usd)
request_cost = float(request_cost_usd) if request_cost_usd is not None else cost
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(
billing_status="settled",
finalized_at=now,
total_cost_usd=cost,
request_cost_usd=request_cost,
status=status,
status_code=status_code,
error_message=error_message,
response_time_ms=response_time_ms,
)
)
if result.rowcount != 1:
return False
# 写入审计快照(只在本次 finalize 生效时执行)
usage = db.query(Usage).filter(Usage.request_id == request_id).first()
if usage:
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,
) -> bool:
"""
并发安全的幂等 finalizevoid不收费
约定:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(
billing_status="void",
finalized_at=now,
total_cost_usd=0.0,
request_cost_usd=0.0,
status="cancelled",
status_code=status_code,
error_message=reason,
response_time_ms=None,
)
)
return result.rowcount == 1
@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 使用记录标记为 settled费用暂时为 0。
后续轮询完成后通过 update_settled_billing 更新实际费用。
约定:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
# 处理响应头和响应体
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:
processed_provider_headers = SystemConfigService.mask_sensitive_headers(
db, provider_request_headers
)
processed_response_headers = None
if should_log_headers and response_headers:
processed_response_headers = dict(response_headers)
processed_response_body = None
if should_log_body and response_body:
processed_response_body = SystemConfigService.truncate_body(
db, response_body, is_request=False
)
values: dict[str, Any] = {
"billing_status": "settled",
"finalized_at": now,
"total_cost_usd": 0.0,
"request_cost_usd": 0.0,
"status": "completed",
"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
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(**values)
)
return result.rowcount == 1
@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,
) -> bool:
"""
更新已结算记录的计费信息(用于异步任务轮询完成后)。
与 finalize_settled 不同:
- finalize_settled: pending -> settled首次结算
- update_settled_billing: settled -> settled更新费用
约定:
- 仅当 billing_status='settled' 时才会生效
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
cost = float(total_cost_usd)
request_cost = float(request_cost_usd) if request_cost_usd is not None else cost
values: dict[str, Any] = {
"total_cost_usd": cost,
"request_cost_usd": request_cost,
"status": status,
"status_code": status_code,
}
if error_message is not None:
values["error_message"] = error_message
if response_time_ms is not None:
values["response_time_ms"] = response_time_ms
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "settled",
)
.values(**values)
)
if result.rowcount != 1:
return False
# 写入审计快照
usage = db.query(Usage).filter(Usage.request_id == request_id).first()
if usage:
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:
"""
将已结算的记录作废(用于异步任务取消)。
与 finalize_void 不同:
- finalize_void: pending -> void未结算时作废
- void_settled: settled -> void已结算后取消费用归零
约定:
- 仅当 billing_status='settled' 时才会生效
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "settled",
)
.values(
billing_status="void",
finalized_at=now,
total_cost_usd=0.0,
request_cost_usd=0.0,
status="cancelled",
status_code=status_code,
error_message=reason,
)
)
return result.rowcount == 1
@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,
) -> 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 状态码(可选)
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
# 避免状态回退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
# 结算状态:当请求进入终态时,将 billing_status 标记为 settled
# 注意:取消是否应 VOID/部分结算由更高层策略决定;这里默认终态均视为已结算。
if status in ("completed", "failed", "cancelled"):
if getattr(usage, "billing_status", None) == "pending":
usage.billing_status = "settled"
if getattr(usage, "finalized_at", None) is None:
usage.finalized_at = datetime.now(timezone.utc)
db.commit()
logger.debug("更新使用记录状态: request_id={}, {} -> {}", request_id, old_status, status)
return usage

View 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
)

531
src/services/usage/query.py Normal file
View File

@@ -0,0 +1,531 @@
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.logger import logger
from src.models.database import ApiKey, Usage, User, UserRole
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: 用户IDNone 表示获取全局热力图(管理员)
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_user_quota(
db: Session,
user: User,
estimated_tokens: int = 0,
estimated_cost: float = 0,
api_key: ApiKey | None = None,
) -> tuple[bool, str]:
"""检查用户配额或独立Key余额
Args:
db: 数据库会话
user: 用户对象
estimated_tokens: 预估token数
estimated_cost: 预估费用
api_key: API Key对象用于检查独立余额Key
Returns:
(是否通过, 消息)
"""
# 如果是独立余额Key检查Key的余额而不是用户配额
if api_key and api_key.is_standalone:
# 导入 ApiKeyService 以使用统一的余额计算方法
from src.services.user.apikey import ApiKeyService
# NULL 表示无限制
if api_key.current_balance_usd is None:
return True, "OK"
# 使用统一的余额计算方法
remaining_balance = ApiKeyService.get_remaining_balance(api_key)
if remaining_balance is None:
return True, "OK"
# 检查余额是否充足
if remaining_balance < estimated_cost:
return (
False,
f"Key余额不足剩余: ${remaining_balance:.2f},需要: ${estimated_cost:.2f}",
)
return True, "OK"
# 普通Key检查用户配额
# 管理员无限制
if user.role == UserRole.ADMIN:
return True, "OK"
# NULL 表示无限制
if user.quota_usd is None:
return True, "OK"
# 有配额限制,检查是否超额
used_usd = float(user.used_usd or 0)
quota_usd = float(user.quota_usd)
if used_usd + estimated_cost > quota_usd:
remaining = quota_usd - used_usd
return False, f"配额不足(剩余: ${remaining:.2f}"
return True, "OK"
@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 = "day", # day, week, month
) -> 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)
# 使用跨数据库兼容的日期函数
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)
# 汇总查询
summary = db.query(
date_func.label("period"),
Usage.provider_name,
Usage.model,
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.total_cost_usd).label("total_cost_usd"),
func.avg(Usage.response_time_ms).label("avg_response_time"),
)
# 过滤掉 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(date_func, Usage.provider_name, Usage.model).all()
return [
{
"period": row.period,
"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,
"total_cost_usd": float(row.total_cost_usd),
"avg_response_time_ms": (
float(row.avg_response_time) if row.avg_response_time else 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

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff