mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: 视频计费增强与影子计费系统
This commit is contained in:
@@ -23,6 +23,7 @@ from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.database import get_db
|
||||
from src.models.database import BillingRule, DimensionCollector
|
||||
from src.services.billing.formula_engine import SafeExpressionEvaluator, UnsafeExpressionError
|
||||
from src.services.billing.presets import BillingPresetService, PresetApplyMode, list_preset_packs
|
||||
|
||||
router = APIRouter(prefix="/api/admin/billing", tags=["Admin - Billing"])
|
||||
pipeline = ApiRequestPipeline()
|
||||
@@ -127,6 +128,30 @@ class DimensionCollectorResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class BillingPresetInfoResponse(BaseModel):
|
||||
name: str
|
||||
version: str
|
||||
description: str
|
||||
collector_count: int
|
||||
|
||||
|
||||
class ApplyBillingPresetRequest(BaseModel):
|
||||
preset: str = Field(..., min_length=1, max_length=100)
|
||||
mode: PresetApplyMode = "merge"
|
||||
|
||||
|
||||
@router.get("/presets")
|
||||
async def list_billing_presets(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||
adapter = BillingPresetListAdapter()
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.post("/presets/apply")
|
||||
async def apply_billing_preset(request: Request, db: Session = Depends(get_db)) -> Any:
|
||||
adapter = BillingPresetApplyAdapter()
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/rules")
|
||||
async def list_billing_rules(
|
||||
request: Request,
|
||||
@@ -525,3 +550,38 @@ def _validate_dimension_collector_request(
|
||||
raise InvalidRequestException(
|
||||
"default_value already exists for this (api_format, task_type, dimension_name)"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BillingPresetListAdapter(AdminApiAdapter):
|
||||
async def handle(self, context: ApiRequestContext) -> dict[str, Any]:
|
||||
items = []
|
||||
for p in list_preset_packs():
|
||||
items.append(
|
||||
BillingPresetInfoResponse(
|
||||
name=p.name,
|
||||
version=p.version,
|
||||
description=p.description,
|
||||
collector_count=len(p.collectors or []),
|
||||
).model_dump()
|
||||
)
|
||||
return {"items": items}
|
||||
|
||||
|
||||
class BillingPresetApplyAdapter(AdminApiAdapter):
|
||||
async def handle(self, context: ApiRequestContext) -> dict[str, Any]:
|
||||
payload = context.ensure_json_body()
|
||||
try:
|
||||
req = ApplyBillingPresetRequest.model_validate(payload)
|
||||
except Exception as exc:
|
||||
raise InvalidRequestException(f"Invalid request body: {exc}")
|
||||
|
||||
result = BillingPresetService.apply_preset(
|
||||
context.db,
|
||||
preset_name=req.preset,
|
||||
mode=req.mode,
|
||||
)
|
||||
if result.errors:
|
||||
# still return counts; caller can display partial results
|
||||
return {"ok": False, **result.to_dict()}
|
||||
return {"ok": True, **result.to_dict()}
|
||||
|
||||
@@ -726,9 +726,9 @@ class AdminGetApiFormatsAdapter(AdminApiAdapter):
|
||||
def _label_for(sig: str) -> str:
|
||||
fam, kind = (sig.split(":", 1) + [""])[:2]
|
||||
fam_title = {"claude": "Claude", "openai": "OpenAI", "gemini": "Gemini"}.get(fam, fam)
|
||||
if kind == "chat":
|
||||
return fam_title
|
||||
kind_title = {"cli": "CLI", "video": "Video", "image": "Image"}.get(kind, kind)
|
||||
kind_title = {"chat": "Chat", "cli": "CLI", "video": "Video", "image": "Image"}.get(
|
||||
kind, kind
|
||||
)
|
||||
return f"{fam_title} {kind_title}".strip()
|
||||
|
||||
endpoint_defs = list_endpoint_definitions()
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
@@ -14,6 +14,8 @@ from sqlalchemy.orm import Session
|
||||
from src.api.base.admin_adapter import AdminApiAdapter
|
||||
from src.api.base.context import ApiRequestContext
|
||||
from src.api.base.pipeline import ApiRequestPipeline
|
||||
from src.config.constants import CacheTTL
|
||||
from src.config.settings import config
|
||||
from src.database import get_db
|
||||
from src.models.database import (
|
||||
ApiKey,
|
||||
@@ -25,11 +27,31 @@ from src.models.database import (
|
||||
User,
|
||||
)
|
||||
from src.services.usage.service import UsageService
|
||||
from src.utils.cache_decorator import cache_result
|
||||
|
||||
router = APIRouter(prefix="/api/admin/usage", tags=["Admin - Usage"])
|
||||
pipeline = ApiRequestPipeline()
|
||||
|
||||
|
||||
def _apply_admin_default_range(
|
||||
start_date: datetime | None, end_date: datetime | None
|
||||
) -> tuple[datetime | None, datetime | None]:
|
||||
"""
|
||||
Apply a default time range for admin usage endpoints to protect DB from unbounded scans.
|
||||
|
||||
Enabled by setting ADMIN_USAGE_DEFAULT_DAYS>0.
|
||||
"""
|
||||
if start_date is not None or end_date is not None:
|
||||
return start_date, end_date
|
||||
|
||||
days = int(getattr(config, "admin_usage_default_days", 0) or 0)
|
||||
if days <= 0:
|
||||
return start_date, end_date
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
return now - timedelta(days=days), now
|
||||
|
||||
|
||||
# ==================== RESTful Routes ====================
|
||||
|
||||
|
||||
@@ -267,9 +289,14 @@ async def get_usage_detail(
|
||||
|
||||
class AdminUsageStatsAdapter(AdminApiAdapter):
|
||||
def __init__(self, start_date: datetime | None, end_date: datetime | None):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:stats",
|
||||
ttl=CacheTTL.ADMIN_USAGE_AGGREGATION,
|
||||
user_specific=False,
|
||||
vary_by=["start_date", "end_date"],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
# Perf: use a single aggregate query (avoid 3 full scans).
|
||||
from sqlalchemy import case
|
||||
@@ -347,10 +374,15 @@ class AdminActivityHeatmapAdapter(AdminApiAdapter):
|
||||
|
||||
class AdminUsageByModelAdapter(AdminApiAdapter):
|
||||
def __init__(self, start_date: datetime | None, end_date: datetime | None, limit: int):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
self.limit = limit
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:agg:model",
|
||||
ttl=CacheTTL.ADMIN_USAGE_AGGREGATION,
|
||||
user_specific=False,
|
||||
vary_by=["start_date", "end_date", "limit"],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
query = db.query(
|
||||
@@ -394,10 +426,15 @@ class AdminUsageByModelAdapter(AdminApiAdapter):
|
||||
|
||||
class AdminUsageByUserAdapter(AdminApiAdapter):
|
||||
def __init__(self, start_date: datetime | None, end_date: datetime | None, limit: int):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
self.limit = limit
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:agg:user",
|
||||
ttl=CacheTTL.ADMIN_USAGE_AGGREGATION,
|
||||
user_specific=False,
|
||||
vary_by=["start_date", "end_date", "limit"],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
query = (
|
||||
@@ -444,10 +481,15 @@ class AdminUsageByUserAdapter(AdminApiAdapter):
|
||||
|
||||
class AdminUsageByProviderAdapter(AdminApiAdapter):
|
||||
def __init__(self, start_date: datetime | None, end_date: datetime | None, limit: int):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
self.limit = limit
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:agg:provider",
|
||||
ttl=CacheTTL.ADMIN_USAGE_AGGREGATION,
|
||||
user_specific=False,
|
||||
vary_by=["start_date", "end_date", "limit"],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
|
||||
@@ -558,10 +600,15 @@ class AdminUsageByProviderAdapter(AdminApiAdapter):
|
||||
|
||||
class AdminUsageByApiFormatAdapter(AdminApiAdapter):
|
||||
def __init__(self, start_date: datetime | None, end_date: datetime | None, limit: int):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
self.limit = limit
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:agg:api_format",
|
||||
ttl=CacheTTL.ADMIN_USAGE_AGGREGATION,
|
||||
user_specific=False,
|
||||
vary_by=["start_date", "end_date", "limit"],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
query = db.query(
|
||||
@@ -624,8 +671,7 @@ class AdminUsageRecordsAdapter(AdminApiAdapter):
|
||||
limit: int,
|
||||
offset: int,
|
||||
):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.start_date, self.end_date = _apply_admin_default_range(start_date, end_date)
|
||||
self.search = search
|
||||
self.user_id = user_id
|
||||
self.username = username
|
||||
@@ -635,6 +681,23 @@ class AdminUsageRecordsAdapter(AdminApiAdapter):
|
||||
self.limit = limit
|
||||
self.offset = offset
|
||||
|
||||
@cache_result(
|
||||
key_prefix="admin:usage:records",
|
||||
ttl=CacheTTL.ADMIN_USAGE_RECORDS,
|
||||
user_specific=False,
|
||||
vary_by=[
|
||||
"start_date",
|
||||
"end_date",
|
||||
"search",
|
||||
"user_id",
|
||||
"username",
|
||||
"model",
|
||||
"provider",
|
||||
"status",
|
||||
"limit",
|
||||
"offset",
|
||||
],
|
||||
)
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
from sqlalchemy import or_
|
||||
from sqlalchemy.orm import load_only
|
||||
@@ -955,7 +1018,11 @@ class AdminUsageDetailAdapter(AdminApiAdapter):
|
||||
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
# 先通过主键 id 查找,如果找不到再尝试通过 request_id 查找
|
||||
usage_record = db.query(Usage).filter(Usage.id == self.usage_id).first()
|
||||
if not usage_record:
|
||||
# 兼容通过 request_id 查找(用于异步任务等场景)
|
||||
usage_record = db.query(Usage).filter(Usage.request_id == self.usage_id).first()
|
||||
if not usage_record:
|
||||
raise HTTPException(status_code=404, detail="Usage record not found")
|
||||
|
||||
@@ -970,6 +1037,9 @@ class AdminUsageDetailAdapter(AdminApiAdapter):
|
||||
usage_id=self.usage_id,
|
||||
)
|
||||
|
||||
# 提取视频/图像/音频计费信息
|
||||
video_billing_info = self._extract_video_billing_info(usage_record)
|
||||
|
||||
return {
|
||||
"id": usage_record.id,
|
||||
"request_id": usage_record.request_id,
|
||||
@@ -1022,6 +1092,7 @@ class AdminUsageDetailAdapter(AdminApiAdapter):
|
||||
"response_body": usage_record.get_response_body(),
|
||||
"metadata": usage_record.request_metadata,
|
||||
"tiered_pricing": tiered_pricing_info,
|
||||
"video_billing": video_billing_info,
|
||||
}
|
||||
|
||||
async def _get_tiered_pricing_info(self, db: Session, usage_record: Any) -> dict | None:
|
||||
@@ -1077,6 +1148,75 @@ class AdminUsageDetailAdapter(AdminApiAdapter):
|
||||
"source": pricing_source, # 定价来源: 'provider' 或 'global'
|
||||
}
|
||||
|
||||
def _extract_video_billing_info(self, usage_record: Any) -> dict | None:
|
||||
"""
|
||||
从 request_metadata.billing_snapshot 和 dimensions 中提取视频/图像/音频计费信息。
|
||||
|
||||
返回结构:
|
||||
{
|
||||
"task_type": "video" | "image" | "audio",
|
||||
"duration_seconds": 10.5, # 视频时长(秒)
|
||||
"resolution": "1080p", # 分辨率
|
||||
"video_price_per_second": 0.1, # 每秒单价
|
||||
"video_cost": 1.05, # 视频费用
|
||||
"rule_name": "...", # 计费规则名称
|
||||
"expression": "...", # 计费公式
|
||||
"status": "complete", # 计费状态
|
||||
}
|
||||
"""
|
||||
request_type = getattr(usage_record, "request_type", None)
|
||||
if request_type not in {"video", "image", "audio"}:
|
||||
return None
|
||||
|
||||
metadata = getattr(usage_record, "request_metadata", None)
|
||||
if not metadata:
|
||||
return None
|
||||
|
||||
billing_snapshot = metadata.get("billing_snapshot") if isinstance(metadata, dict) else None
|
||||
dimensions = metadata.get("dimensions") if isinstance(metadata, dict) else None
|
||||
|
||||
result: dict = {
|
||||
"task_type": request_type,
|
||||
}
|
||||
|
||||
# 从 billing_snapshot 中提取计费规则信息
|
||||
if billing_snapshot and isinstance(billing_snapshot, dict):
|
||||
result["rule_name"] = billing_snapshot.get("rule_name")
|
||||
result["expression"] = billing_snapshot.get("expression")
|
||||
result["status"] = billing_snapshot.get("status")
|
||||
result["cost"] = billing_snapshot.get("cost")
|
||||
|
||||
# 从 dimensions_used 中提取维度
|
||||
dims_used = billing_snapshot.get("dimensions_used")
|
||||
if dims_used and isinstance(dims_used, dict):
|
||||
if "duration_seconds" in dims_used:
|
||||
result["duration_seconds"] = dims_used["duration_seconds"]
|
||||
if "video_resolution_key" in dims_used:
|
||||
result["resolution"] = dims_used["video_resolution_key"]
|
||||
if "video_price_per_second" in dims_used:
|
||||
result["video_price_per_second"] = dims_used["video_price_per_second"]
|
||||
if "video_cost" in dims_used:
|
||||
result["video_cost"] = dims_used["video_cost"]
|
||||
|
||||
# 补充从 dimensions 中提取(备用)
|
||||
if dimensions and isinstance(dimensions, dict):
|
||||
if "duration_seconds" not in result and "duration_seconds" in dimensions:
|
||||
result["duration_seconds"] = dimensions["duration_seconds"]
|
||||
if "resolution" not in result and "video_resolution_key" in dimensions:
|
||||
result["resolution"] = dimensions["video_resolution_key"]
|
||||
|
||||
# 如果没有有意义的视频计费信息,返回 None
|
||||
has_video_info = (
|
||||
result.get("duration_seconds")
|
||||
or result.get("resolution")
|
||||
or result.get("video_cost")
|
||||
or result.get("cost")
|
||||
)
|
||||
if not has_video_info:
|
||||
return None
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ==================== 缓存亲和性分析 ====================
|
||||
|
||||
|
||||
@@ -7,18 +7,22 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.api.base.context import ApiRequestContext
|
||||
from src.api.base.pipeline import ApiRequestPipeline
|
||||
from src.api.dashboard.routes import DashboardAdapter
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.enums import UserRole
|
||||
from src.core.logger import logger
|
||||
from src.database import get_db
|
||||
from src.models.database import Provider, ProviderEndpoint, User, VideoTask
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint, User, VideoTask
|
||||
|
||||
router = APIRouter(prefix="/api/admin/video-tasks", tags=["Admin - Video Tasks"])
|
||||
pipeline = ApiRequestPipeline()
|
||||
@@ -119,6 +123,117 @@ async def cancel_video_task(
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/{task_id}/video")
|
||||
async def proxy_video_stream(
|
||||
task_id: str,
|
||||
request: Request,
|
||||
token: str | None = Query(None, description="JWT access token"),
|
||||
db: Session = Depends(get_db),
|
||||
) -> StreamingResponse:
|
||||
"""
|
||||
代理视频流(用于需要认证的视频链接)
|
||||
|
||||
**路径参数**:
|
||||
- `task_id`: 任务 ID
|
||||
|
||||
**查询参数**:
|
||||
- `token`: JWT access token(用于 video 标签请求)
|
||||
|
||||
**返回**:
|
||||
- 视频流
|
||||
"""
|
||||
from src.services.auth.service import AuthService
|
||||
|
||||
# 尝试从多个来源获取 token:query param > cookie > header
|
||||
auth_token = token
|
||||
if not auth_token:
|
||||
auth_token = request.cookies.get("access_token")
|
||||
if not auth_token:
|
||||
auth_header = request.headers.get("Authorization")
|
||||
if auth_header and auth_header.startswith("Bearer "):
|
||||
auth_token = auth_header[7:]
|
||||
|
||||
if not auth_token:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
|
||||
try:
|
||||
# 验证 token 并获取 payload
|
||||
payload = await AuthService.verify_token(auth_token, token_type="access")
|
||||
user_id = payload.get("user_id") or payload.get("sub")
|
||||
if not user_id:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
# 查询用户
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
|
||||
# 查询任务
|
||||
query = db.query(VideoTask).filter(VideoTask.id == task_id)
|
||||
if user.role != UserRole.ADMIN:
|
||||
query = query.filter(VideoTask.user_id == user.id)
|
||||
|
||||
task = query.first()
|
||||
if not task:
|
||||
raise HTTPException(status_code=404, detail="Video task not found")
|
||||
|
||||
if not task.video_url:
|
||||
raise HTTPException(status_code=404, detail="Video not available")
|
||||
|
||||
# 检查是否需要代理(Google API 链接需要认证)
|
||||
video_url = task.video_url
|
||||
needs_proxy = "generativelanguage.googleapis.com" in video_url
|
||||
|
||||
if not needs_proxy:
|
||||
# 不需要代理,重定向到原始 URL
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
return RedirectResponse(url=video_url)
|
||||
|
||||
# 需要代理:获取 provider key 进行认证
|
||||
if not task.key_id:
|
||||
raise HTTPException(status_code=500, detail="Missing provider key")
|
||||
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == task.key_id).first()
|
||||
if not key or not key.api_key:
|
||||
raise HTTPException(status_code=500, detail="Provider key not found")
|
||||
|
||||
try:
|
||||
api_key = crypto_service.decrypt(key.api_key)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to decrypt provider key")
|
||||
|
||||
# 构建认证头
|
||||
headers = {"x-goog-api-key": api_key}
|
||||
|
||||
async def stream_video() -> AsyncIterator[bytes]:
|
||||
"""流式下载并返回视频"""
|
||||
try:
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
async with client.stream("GET", video_url, headers=headers) as response:
|
||||
if response.status_code >= 400:
|
||||
logger.warning(
|
||||
"Video proxy failed: task={} status={}", task_id, response.status_code
|
||||
)
|
||||
return
|
||||
async for chunk in response.aiter_bytes(chunk_size=65536):
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.exception("Video proxy error: task={} error={}", task_id, str(e))
|
||||
|
||||
return StreamingResponse(
|
||||
stream_video(),
|
||||
media_type="video/mp4",
|
||||
headers={
|
||||
"Content-Disposition": f'inline; filename="video_{task_id}.mp4"',
|
||||
"Cache-Control": "private, max-age=3600",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ==================== Adapters ====================
|
||||
|
||||
|
||||
@@ -359,6 +474,7 @@ class VideoTaskDetailAdapter(DashboardAdapter):
|
||||
"video_urls": task.video_urls,
|
||||
"thumbnail_url": task.thumbnail_url,
|
||||
"video_size_bytes": task.video_size_bytes,
|
||||
"video_duration_seconds": task.video_duration_seconds,
|
||||
"video_expires_at": (
|
||||
task.video_expires_at.isoformat() if task.video_expires_at else None
|
||||
),
|
||||
|
||||
@@ -211,8 +211,10 @@ class AdminDashboardStatsAdapter(AdminApiAdapter):
|
||||
month_start = month_start_local.astimezone(timezone.utc)
|
||||
|
||||
# ==================== 使用预聚合数据 ====================
|
||||
# 今日实时数据只查询一次,避免重复扫描 Usage 表
|
||||
today_stats = StatsAggregatorService.get_today_realtime_stats(db)
|
||||
# 从 stats_summary + 今日实时数据获取全局统计
|
||||
combined_stats = StatsAggregatorService.get_combined_stats(db)
|
||||
combined_stats = StatsAggregatorService.get_combined_stats(db, today_stats=today_stats)
|
||||
|
||||
all_time_requests = combined_stats["total_requests"]
|
||||
all_time_success_requests = combined_stats["success_requests"]
|
||||
@@ -237,7 +239,6 @@ class AdminDashboardStatsAdapter(AdminApiAdapter):
|
||||
)
|
||||
|
||||
# ==================== 今日实时统计 ====================
|
||||
today_stats = StatsAggregatorService.get_today_realtime_stats(db)
|
||||
requests_today = today_stats["total_requests"]
|
||||
cost_today = today_stats["total_cost"]
|
||||
actual_cost_today = today_stats["actual_total_cost"]
|
||||
@@ -951,26 +952,9 @@ class DashboardDailyStatsAdapter(DashboardAdapter):
|
||||
today_stats = StatsAggregatorService.get_today_realtime_stats(db)
|
||||
today_str = today_local.date().isoformat()
|
||||
if today_stats["total_requests"] > 0:
|
||||
# 今日平均响应时间需要单独查询
|
||||
today_avg_rt = (
|
||||
db.query(func.avg(Usage.response_time_ms))
|
||||
.filter(Usage.created_at >= today, Usage.response_time_ms.isnot(None))
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
# 今日 unique_models 和 unique_providers
|
||||
today_unique_models = (
|
||||
db.query(func.count(func.distinct(Usage.model)))
|
||||
.filter(Usage.created_at >= today)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
today_unique_providers = (
|
||||
db.query(func.count(func.distinct(Usage.provider_name)))
|
||||
.filter(Usage.created_at >= today)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
today_avg_rt_ms = float(today_stats.get("avg_response_time_ms") or 0.0)
|
||||
today_unique_models = int(today_stats.get("unique_models") or 0)
|
||||
today_unique_providers = int(today_stats.get("unique_providers") or 0)
|
||||
# 今日 fallback_count
|
||||
today_fallback_count = (
|
||||
db.query(func.count())
|
||||
@@ -996,7 +980,7 @@ class DashboardDailyStatsAdapter(DashboardAdapter):
|
||||
+ today_stats["cache_read_tokens"]
|
||||
),
|
||||
"cost": today_stats["total_cost"],
|
||||
"avg_response_time": float(today_avg_rt) / 1000.0 if today_avg_rt else 0,
|
||||
"avg_response_time": today_avg_rt_ms / 1000.0 if today_avg_rt_ms else 0,
|
||||
"unique_models": today_unique_models,
|
||||
"unique_providers": today_unique_providers,
|
||||
"fallback_count": today_fallback_count,
|
||||
|
||||
@@ -100,10 +100,7 @@ class StreamTelemetryRecorder:
|
||||
return
|
||||
actual_request_body = ctx.provider_request_body or original_request_body
|
||||
response_body = None
|
||||
if (
|
||||
not isinstance(writer, QueueTelemetryWriter)
|
||||
or config.usage_queue_include_bodies
|
||||
):
|
||||
if not isinstance(writer, QueueTelemetryWriter) or writer.include_bodies:
|
||||
response_body = ctx.build_response_body(response_time_ms)
|
||||
|
||||
try:
|
||||
@@ -403,10 +400,25 @@ class StreamTelemetryRecorder:
|
||||
self, bg_db: Session, ctx: StreamContext, response_time_ms: int
|
||||
) -> TelemetryWriter | None:
|
||||
if config.usage_queue_enabled and self.user_id and self.api_key_id:
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
# Queue payload detail follows system config request_record_level.
|
||||
log_level = SystemConfigService.get_request_record_level(bg_db).value
|
||||
sensitive_headers = SystemConfigService.get_sensitive_headers(bg_db) or []
|
||||
max_request_body_size = int(
|
||||
SystemConfigService.get_config(bg_db, "max_request_body_size", 5242880) or 0
|
||||
)
|
||||
max_response_body_size = int(
|
||||
SystemConfigService.get_config(bg_db, "max_response_body_size", 5242880) or 0
|
||||
)
|
||||
return QueueTelemetryWriter(
|
||||
request_id=self.request_id,
|
||||
user_id=self.user_id,
|
||||
api_key_id=self.api_key_id,
|
||||
log_level=log_level,
|
||||
sensitive_headers=sensitive_headers,
|
||||
max_request_body_size=max_request_body_size,
|
||||
max_response_body_size=max_response_body_size,
|
||||
)
|
||||
db_writer = self._build_db_writer(bg_db)
|
||||
if db_writer is None:
|
||||
|
||||
@@ -75,6 +75,17 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
)
|
||||
self._normalizer = GeminiNormalizer()
|
||||
|
||||
@staticmethod
|
||||
def _get_request_base_url(http_request: Request) -> str:
|
||||
"""从 HTTP 请求中获取基础 URL(协议 + 主机)"""
|
||||
# 优先使用 X-Forwarded-Proto 和 X-Forwarded-Host(代理场景)
|
||||
proto = http_request.headers.get("x-forwarded-proto") or http_request.url.scheme
|
||||
host = http_request.headers.get("x-forwarded-host") or http_request.headers.get("host")
|
||||
if host:
|
||||
return f"{proto}://{host}"
|
||||
# 回退到 request.url
|
||||
return f"{http_request.url.scheme}://{http_request.url.netloc}"
|
||||
|
||||
async def handle_create_task(
|
||||
self,
|
||||
*,
|
||||
@@ -259,7 +270,8 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
created_at=task.created_at,
|
||||
original_request=internal_request,
|
||||
)
|
||||
response_body = self._normalizer.video_task_from_internal(internal_task)
|
||||
base_url = self._get_request_base_url(http_request)
|
||||
response_body = self._normalizer.video_task_from_internal(internal_task, base_url=base_url)
|
||||
|
||||
# 提交成功后立即结算 Usage(费用暂时为 0,轮询完成后更新)
|
||||
response_time_ms = int((time.time() - self.start_time) * 1000)
|
||||
@@ -313,7 +325,8 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
|
||||
# 直接从数据库返回任务状态(后台轮询服务会持续更新状态)
|
||||
internal_task = self._task_to_internal(task)
|
||||
response_body = self._normalizer.video_task_from_internal(internal_task)
|
||||
base_url = self._get_request_base_url(http_request)
|
||||
response_body = self._normalizer.video_task_from_internal(internal_task, base_url=base_url)
|
||||
return JSONResponse(response_body)
|
||||
|
||||
async def handle_list_tasks(
|
||||
@@ -331,8 +344,10 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
.limit(100)
|
||||
.all()
|
||||
)
|
||||
base_url = self._get_request_base_url(http_request)
|
||||
items = [
|
||||
self._normalizer.video_task_from_internal(self._task_to_internal(t)) for t in tasks
|
||||
self._normalizer.video_task_from_internal(self._task_to_internal(t), base_url=base_url)
|
||||
for t in tasks
|
||||
]
|
||||
return JSONResponse({"operations": items})
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, AsyncIterator
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
@@ -650,8 +651,10 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
try:
|
||||
# 使用 httpx 的 stream 方法并正确管理上下文
|
||||
# 视频下载可能较大,设置 5 分钟超时
|
||||
request = client.build_request("GET", upstream_url, headers=headers)
|
||||
response = await client.send(request, stream=True, timeout=300.0)
|
||||
request = client.build_request(
|
||||
"GET", upstream_url, headers=headers, timeout=httpx.Timeout(300.0)
|
||||
)
|
||||
response = await client.send(request, stream=True)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[VideoDownload] Upstream connection failed task={} url={}: {}",
|
||||
@@ -723,8 +726,8 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
"""代理直接的视频 URL(如 CDN URL),保持与官方 API 一致的流式返回行为"""
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
try:
|
||||
request = client.build_request("GET", url)
|
||||
response = await client.send(request, stream=True, timeout=300.0)
|
||||
request = client.build_request("GET", url, timeout=httpx.Timeout(300.0))
|
||||
response = await client.send(request, stream=True)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[VideoDownload] Direct URL connection failed task={} url={}: {}",
|
||||
|
||||
@@ -73,11 +73,11 @@ def _get_formats_for_api(api_format: str) -> list[str]:
|
||||
return _OPENAI_FORMATS
|
||||
|
||||
|
||||
def _is_format_conversion_enabled() -> bool:
|
||||
"""检查全局格式转换开关(从环境变量读取,默认开启)"""
|
||||
from src.config.settings import config
|
||||
def _is_format_conversion_enabled(db: Session) -> bool:
|
||||
"""检查全局格式转换开关(从数据库配置读取,默认开启)"""
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
return config.format_conversion_enabled
|
||||
return SystemConfigService.is_format_conversion_enabled(db)
|
||||
|
||||
|
||||
def _get_convertible_formats(client_format: str, global_conversion_enabled: bool) -> list[str]:
|
||||
@@ -500,7 +500,7 @@ async def list_models(
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(key_record, user)
|
||||
|
||||
# 获取可用格式(包括可转换的格式)
|
||||
global_conversion_enabled = _is_format_conversion_enabled()
|
||||
global_conversion_enabled = _is_format_conversion_enabled(db)
|
||||
candidate_formats = _get_convertible_formats(api_format, global_conversion_enabled)
|
||||
candidate_formats, empty_response = _filter_formats_by_restrictions(
|
||||
candidate_formats, restrictions, api_format
|
||||
@@ -604,7 +604,7 @@ async def retrieve_model(
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(key_record, user)
|
||||
|
||||
# 获取可用格式(包括可转换的格式)
|
||||
global_conversion_enabled = _is_format_conversion_enabled()
|
||||
global_conversion_enabled = _is_format_conversion_enabled(db)
|
||||
candidate_formats = _get_convertible_formats(api_format, global_conversion_enabled)
|
||||
candidate_formats, _ = _filter_formats_by_restrictions(
|
||||
candidate_formats, restrictions, api_format
|
||||
@@ -687,7 +687,7 @@ async def list_models_gemini(
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(key_record, user)
|
||||
|
||||
# 获取可用格式(包括可转换的格式)
|
||||
global_conversion_enabled = _is_format_conversion_enabled()
|
||||
global_conversion_enabled = _is_format_conversion_enabled(db)
|
||||
candidate_formats = _get_convertible_formats(api_format, global_conversion_enabled)
|
||||
candidate_formats, empty_response = _filter_formats_by_restrictions(
|
||||
candidate_formats, restrictions, api_format
|
||||
@@ -766,7 +766,7 @@ async def get_model_gemini(
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(key_record, user)
|
||||
|
||||
# 获取可用格式(包括可转换的格式)
|
||||
global_conversion_enabled = _is_format_conversion_enabled()
|
||||
global_conversion_enabled = _is_format_conversion_enabled(db)
|
||||
candidate_formats = _get_convertible_formats(api_format, global_conversion_enabled)
|
||||
candidate_formats, _ = _filter_formats_by_restrictions(
|
||||
candidate_formats, restrictions, api_format
|
||||
|
||||
@@ -1110,7 +1110,7 @@ class ListAvailableModelsAdapter(AuthenticatedApiAdapter):
|
||||
from sqlalchemy import or_
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
from src.config.settings import config as app_config
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
db = context.db
|
||||
user = context.user
|
||||
@@ -1118,8 +1118,8 @@ class ListAvailableModelsAdapter(AuthenticatedApiAdapter):
|
||||
# 使用 AccessRestrictions 类来处理限制(与 /v1/models 逻辑一致)
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(api_key=None, user=user)
|
||||
|
||||
# 检查全局格式转换开关
|
||||
global_conversion_enabled = app_config.format_conversion_enabled
|
||||
# 检查全局格式转换开关(从数据库配置读取)
|
||||
global_conversion_enabled = SystemConfigService.is_format_conversion_enabled(db)
|
||||
|
||||
# 获取所有可用的 Provider ID(考虑格式转换)
|
||||
available_provider_ids = self._get_all_available_provider_ids(db, global_conversion_enabled)
|
||||
|
||||
@@ -28,6 +28,10 @@ class CacheTTL:
|
||||
DASHBOARD_STATS = 120 # 2分钟(管理员)
|
||||
DASHBOARD_DAILY = 600 # 10分钟(每日统计)
|
||||
|
||||
# Admin usage pages (heavy DB aggregations / list queries)
|
||||
ADMIN_USAGE_AGGREGATION = 30 # 30秒
|
||||
ADMIN_USAGE_RECORDS = 10 # 10秒(列表页短缓存,避免轮询/重复刷新打爆 DB)
|
||||
|
||||
# 并发锁 TTL - 防止死锁
|
||||
CONCURRENCY_LOCK = 600 # 10分钟
|
||||
|
||||
|
||||
@@ -148,21 +148,6 @@ class Config:
|
||||
# HTTP_REQUEST_TIMEOUT: 非流式请求整体超时(秒),默认 300 秒
|
||||
self.http_request_timeout = float(os.getenv("HTTP_REQUEST_TIMEOUT", "300.0"))
|
||||
|
||||
# 格式转换配置
|
||||
# FORMAT_CONVERSION_ENABLED: 全局格式转换总开关,默认开启
|
||||
# 注意:即使开启,也需要端点配置 format_acceptance_config.enabled=true 才能生效
|
||||
self.format_conversion_enabled = (
|
||||
os.getenv("FORMAT_CONVERSION_ENABLED", "true").lower() == "true"
|
||||
)
|
||||
|
||||
# KEEP_PRIORITY_ON_CONVERSION: 格式转换时是否保持提供商原优先级,默认关闭
|
||||
# - false(默认): 需要格式转换的候选整体降级到不需要转换的候选之后
|
||||
# - true: 所有提供商保持原优先级,不因格式转换降级
|
||||
# 注意:即使全局关闭,单个提供商也可以通过 keep_priority_on_conversion 字段保持自己的优先级
|
||||
self.keep_priority_on_conversion = (
|
||||
os.getenv("KEEP_PRIORITY_ON_CONVERSION", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# HTTP 连接池配置
|
||||
# HTTP_MAX_CONNECTIONS: 最大连接数,影响并发能力
|
||||
# - 每个连接占用一个 socket,过多会耗尽系统资源
|
||||
@@ -192,15 +177,8 @@ class Config:
|
||||
# Usage 队列配置(Redis Streams)
|
||||
# 默认启用队列模式,通过 Redis Streams 异步写入 DB,提升响应性能
|
||||
self.usage_queue_enabled = os.getenv("USAGE_QUEUE_ENABLED", "true").lower() == "true"
|
||||
# 默认传输 headers/bodies,由系统设置(request_log_level)决定最终存储内容
|
||||
self.usage_queue_include_headers = (
|
||||
os.getenv("USAGE_QUEUE_INCLUDE_HEADERS", "true").lower() == "true"
|
||||
)
|
||||
self.usage_queue_include_bodies = (
|
||||
os.getenv("USAGE_QUEUE_INCLUDE_BODIES", "true").lower() == "true"
|
||||
)
|
||||
# 0 表示不截断,由系统设置(max_request/response_body_size)统一控制
|
||||
self.usage_queue_body_max_bytes = int(os.getenv("USAGE_QUEUE_BODY_MAX_BYTES", "0"))
|
||||
# 队列事件是否包含 headers/bodies 由系统配置(request_record_level)决定;
|
||||
# 最终写入 DB 前仍会按 SystemConfigService 做脱敏与截断。
|
||||
self.usage_queue_stream_key = os.getenv("USAGE_QUEUE_STREAM_KEY", "usage:events")
|
||||
self.usage_queue_stream_group = os.getenv("USAGE_QUEUE_STREAM_GROUP", "usage_consumers")
|
||||
self.usage_queue_stream_maxlen = int(os.getenv("USAGE_QUEUE_STREAM_MAXLEN", "200000"))
|
||||
@@ -217,6 +195,17 @@ class Config:
|
||||
os.getenv("USAGE_QUEUE_METRICS_INTERVAL_SECONDS", "30")
|
||||
)
|
||||
|
||||
# Admin analytics query defaults (protect DB from unbounded scans)
|
||||
# ADMIN_USAGE_DEFAULT_DAYS:
|
||||
# - 0: keep current behavior (no implicit time filter)
|
||||
# - >0: when admin usage endpoints omit start_date/end_date, default to "last N days"
|
||||
default_admin_usage_default_days = (
|
||||
"0" if self.environment in {"development", "test", "testing"} else "30"
|
||||
)
|
||||
self.admin_usage_default_days = int(
|
||||
os.getenv("ADMIN_USAGE_DEFAULT_DAYS", default_admin_usage_default_days)
|
||||
)
|
||||
|
||||
# Thinking 整流器配置
|
||||
# THINKING_RECTIFIER_ENABLED: 是否启用 Thinking 整流器
|
||||
# 当遇到跨 Provider 的 thinking 签名错误时,自动整流请求体后重试
|
||||
@@ -252,6 +241,37 @@ class Config:
|
||||
self.billing_require_rule = os.getenv("BILLING_REQUIRE_RULE", "false").lower() == "true"
|
||||
self.billing_strict_mode = os.getenv("BILLING_STRICT_MODE", "false").lower() == "true"
|
||||
|
||||
# 计费迁移运行时开关(用于灰度/影子计费/快速止血)
|
||||
# BILLING_ENGINE:
|
||||
# - legacy: 仅旧系统(当前默认)
|
||||
# - shadow: 旧系统为真值 + 新系统影子计算(对账期)
|
||||
# - new_with_fallback: 新系统为真值,差异过大时回退旧系统
|
||||
# - new: 仅新系统
|
||||
# Default to "new" per unified billing architecture.
|
||||
self.billing_engine = os.getenv("BILLING_ENGINE", "new").strip().lower()
|
||||
# 按 provider/model 粒度覆盖(JSON 字符串)
|
||||
# 示例: {"anthropic/*": "shadow", "openai/gpt-4*": "new"}
|
||||
self.billing_engine_overrides = os.getenv("BILLING_ENGINE_OVERRIDES", "{}")
|
||||
# 影子计费差异阈值(美元)
|
||||
self.billing_diff_threshold_usd = float(os.getenv("BILLING_DIFF_THRESHOLD_USD", "0.0001"))
|
||||
# 差异日志级别(DEBUG/INFO/WARNING/ERROR)
|
||||
self.billing_shadow_log_level = os.getenv("BILLING_SHADOW_LOG_LEVEL", "INFO").strip()
|
||||
# 是否启用差异告警(预留扩展)
|
||||
self.billing_diff_alert_enabled = (
|
||||
os.getenv("BILLING_DIFF_ALERT_ENABLED", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# Usage.request_metadata 体积控制(用于降低 DB/CPU/内存压力)
|
||||
# USAGE_METADATA_MAX_BYTES:
|
||||
# - 0: unlimited (backward compatible)
|
||||
# - >0: best-effort prune large keys when metadata JSON exceeds this size
|
||||
default_usage_metadata_max_bytes = (
|
||||
"0" if self.environment in {"development", "test", "testing"} else "65536"
|
||||
)
|
||||
self.usage_metadata_max_bytes = int(
|
||||
os.getenv("USAGE_METADATA_MAX_BYTES", default_usage_metadata_max_bytes)
|
||||
)
|
||||
|
||||
# 视频任务轮询配置
|
||||
# VIDEO_POLL_INTERVAL_SECONDS: 轮询间隔(秒),默认 10 秒
|
||||
# VIDEO_MAX_POLL_COUNT: 最大轮询次数,默认 360 次(约 1 小时)
|
||||
|
||||
@@ -81,7 +81,7 @@ def is_format_compatible(
|
||||
# 2. 格式不同 -> 需要检查格式转换开关
|
||||
# 如果总开关为 False,直接拒绝(禁用任何跨格式转换)
|
||||
if not effective_conversion_enabled:
|
||||
return False, False, "格式转换已禁用(FORMAT_CONVERSION_ENABLED=false)"
|
||||
return False, False, "格式转换已禁用(enable_format_conversion=false)"
|
||||
|
||||
# 3. 如果全局或提供商开关为 ON,跳过端点配置检查
|
||||
if not skip_endpoint_check:
|
||||
|
||||
@@ -74,6 +74,7 @@ class InternalVideoPollResult:
|
||||
error_code: str | None = None
|
||||
error_message: str | None = None
|
||||
raw_response: dict[str, Any] | None = None
|
||||
video_duration_seconds: float | None = None # 实际视频时长
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
@@ -102,8 +102,15 @@ class FormatNormalizer(ABC):
|
||||
"""将视频任务响应转换为内部表示"""
|
||||
raise NotImplementedError(f"{self.__class__.__name__} does not support video conversion")
|
||||
|
||||
def video_task_from_internal(self, internal: InternalVideoTask) -> dict[str, Any]:
|
||||
"""将内部视频任务转换为格式特定响应"""
|
||||
def video_task_from_internal(
|
||||
self, internal: InternalVideoTask, *, base_url: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""将内部视频任务转换为格式特定响应
|
||||
|
||||
Args:
|
||||
internal: 内部视频任务表示
|
||||
base_url: 可选的基础 URL,用于构建完整的下载链接
|
||||
"""
|
||||
raise NotImplementedError(f"{self.__class__.__name__} does not support video conversion")
|
||||
|
||||
def video_poll_to_internal(self, response: dict[str, Any]) -> InternalVideoPollResult:
|
||||
|
||||
@@ -872,7 +872,9 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
extra={"metadata": metadata},
|
||||
)
|
||||
|
||||
def video_task_from_internal(self, internal: InternalVideoTask) -> dict[str, Any]:
|
||||
def video_task_from_internal(
|
||||
self, internal: InternalVideoTask, *, base_url: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
# 从 external_id 中提取 model 名称,用于构建 operation name
|
||||
# external_id 格式: models/{model}/operations/{gemini_id}
|
||||
model_name = "unknown"
|
||||
@@ -889,7 +891,9 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
# 使用我们的内部 task_id 构建下载 URL,不暴露真实的 Gemini file_id
|
||||
# 使用 aev_ 前缀标识这是视频任务的下载链接
|
||||
# 格式:/v1beta/files/aev_{task_id}:download?alt=media
|
||||
proxy_download_url = f"/v1beta/files/aev_{internal.id}:download?alt=media"
|
||||
download_path = f"/v1beta/files/aev_{internal.id}:download?alt=media"
|
||||
# 如果提供了 base_url,返回完整 URL;否则返回相对路径
|
||||
proxy_download_url = f"{base_url}{download_path}" if base_url else download_path
|
||||
return {
|
||||
"name": operation_name,
|
||||
"done": True,
|
||||
@@ -938,12 +942,15 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
if isinstance(s, dict) and s.get("video", {}).get("uri")
|
||||
]
|
||||
video_url = video_urls[0] if video_urls else None
|
||||
# 提取实际视频时长
|
||||
video_duration = self._extract_gemini_video_duration(response, samples)
|
||||
return InternalVideoPollResult(
|
||||
status=VideoStatus.COMPLETED,
|
||||
progress_percent=100,
|
||||
video_url=video_url,
|
||||
video_urls=video_urls,
|
||||
raw_response=response,
|
||||
video_duration_seconds=video_duration,
|
||||
)
|
||||
|
||||
return InternalVideoPollResult(
|
||||
@@ -952,6 +959,41 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
raw_response=response,
|
||||
)
|
||||
|
||||
def _extract_gemini_video_duration(
|
||||
self, response: dict[str, Any], samples: list[dict[str, Any]]
|
||||
) -> float | None:
|
||||
"""从 Gemini 响应中提取实际视频时长"""
|
||||
# 尝试从 samples 中获取时长
|
||||
for sample in samples:
|
||||
if not isinstance(sample, dict):
|
||||
continue
|
||||
video = sample.get("video", {})
|
||||
if isinstance(video, dict):
|
||||
# 尝试多种字段名
|
||||
for field in ["durationSeconds", "duration_seconds", "duration"]:
|
||||
val = video.get(field)
|
||||
if val is not None:
|
||||
try:
|
||||
# duration 可能是 "5s" 格式
|
||||
if isinstance(val, str) and val.endswith("s"):
|
||||
return float(val[:-1])
|
||||
return float(val)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
# 尝试从 response.metadata 获取
|
||||
metadata = response.get("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
for field in ["durationSeconds", "duration_seconds", "duration"]:
|
||||
val = metadata.get(field)
|
||||
if val is not None:
|
||||
try:
|
||||
if isinstance(val, str) and val.endswith("s"):
|
||||
return float(val[:-1])
|
||||
return float(val)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
return None
|
||||
|
||||
# =========================
|
||||
# Helpers
|
||||
# =========================
|
||||
|
||||
@@ -396,7 +396,9 @@ class OpenAINormalizer(FormatNormalizer):
|
||||
# OpenAI streaming may send a final "usage-only" chunk when
|
||||
# stream_options.include_usage=true, where `choices` is empty but `usage` exists.
|
||||
usage_info = self._openai_usage_to_internal(chunk.get("usage"))
|
||||
if usage_info is not None and (usage_info.total_tokens or usage_info.input_tokens or usage_info.output_tokens):
|
||||
if usage_info is not None and (
|
||||
usage_info.total_tokens or usage_info.input_tokens or usage_info.output_tokens
|
||||
):
|
||||
# For cross-format targets (e.g. Gemini), emitting usage as a late MessageStopEvent
|
||||
# allows the target normalizer to surface usage metadata even if the stop chunk
|
||||
# didn't carry it.
|
||||
@@ -730,7 +732,9 @@ class OpenAINormalizer(FormatNormalizer):
|
||||
},
|
||||
)
|
||||
|
||||
def video_task_from_internal(self, internal: InternalVideoTask) -> dict[str, Any]:
|
||||
def video_task_from_internal(
|
||||
self, internal: InternalVideoTask, *, base_url: str | None = None
|
||||
) -> dict[str, Any]:
|
||||
status_map = {
|
||||
VideoStatus.PENDING: "queued",
|
||||
VideoStatus.SUBMITTED: "queued",
|
||||
@@ -797,6 +801,8 @@ class OpenAINormalizer(FormatNormalizer):
|
||||
error_message="Upstream response missing video url",
|
||||
raw_response=response,
|
||||
)
|
||||
# 提取实际视频时长(尝试多种字段名)
|
||||
video_duration = self._extract_video_duration(response)
|
||||
return InternalVideoPollResult(
|
||||
status=VideoStatus.COMPLETED,
|
||||
progress_percent=100,
|
||||
@@ -805,6 +811,7 @@ class OpenAINormalizer(FormatNormalizer):
|
||||
datetime.fromtimestamp(expires_at, tz=timezone.utc) if expires_at else None
|
||||
),
|
||||
raw_response=response,
|
||||
video_duration_seconds=video_duration,
|
||||
)
|
||||
if status == "failed":
|
||||
error = response.get("error") or {}
|
||||
@@ -821,6 +828,36 @@ class OpenAINormalizer(FormatNormalizer):
|
||||
raw_response=response,
|
||||
)
|
||||
|
||||
def _extract_video_duration(self, response: dict[str, Any]) -> float | None:
|
||||
"""从响应中提取实际视频时长"""
|
||||
# 尝试多种可能的字段名
|
||||
duration_fields = [
|
||||
"duration_seconds",
|
||||
"duration",
|
||||
"video_duration",
|
||||
"video_duration_seconds",
|
||||
"length",
|
||||
"length_seconds",
|
||||
]
|
||||
for field in duration_fields:
|
||||
val = response.get(field)
|
||||
if val is not None:
|
||||
try:
|
||||
return float(val)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
# 尝试从嵌套的 metadata 中获取
|
||||
metadata = response.get("metadata") or response.get("video_metadata") or {}
|
||||
if isinstance(metadata, dict):
|
||||
for field in duration_fields:
|
||||
val = metadata.get(field)
|
||||
if val is not None:
|
||||
try:
|
||||
return float(val)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
return None
|
||||
|
||||
# =========================
|
||||
# Helpers
|
||||
# =========================
|
||||
|
||||
@@ -80,3 +80,28 @@ format_conversion_duration_seconds = Histogram(
|
||||
["direction", "source_format", "target_format"],
|
||||
buckets=[0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0],
|
||||
)
|
||||
|
||||
# ==================== Billing migration / shadow billing ====================
|
||||
|
||||
billing_requests_total = Counter(
|
||||
"billing_requests_total",
|
||||
"Total number of billing calculations",
|
||||
["engine_mode", "truth_engine"], # low-cardinality labels
|
||||
)
|
||||
|
||||
billing_fallback_total = Counter(
|
||||
"billing_fallback_total",
|
||||
"Total number of billing fallbacks to legacy engine",
|
||||
)
|
||||
|
||||
billing_diff_exceeds_threshold_total = Counter(
|
||||
"billing_diff_exceeds_threshold_total",
|
||||
"Total number of shadow billing diffs exceeding threshold",
|
||||
["engine_mode"],
|
||||
)
|
||||
|
||||
billing_invariant_violation_total = Counter(
|
||||
"billing_invariant_violation_total",
|
||||
"Total number of billing invariant violations (sum(breakdown)!=total)",
|
||||
["engine_mode", "truth_engine"],
|
||||
)
|
||||
|
||||
@@ -616,6 +616,9 @@ class ModelResponse(BaseModel):
|
||||
global_model_name: str | None = None
|
||||
global_model_display_name: str | None = None
|
||||
|
||||
# 有效配置(合并 Model 和 GlobalModel 的 config)
|
||||
effective_config: dict | None = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
|
||||
@@ -666,7 +666,7 @@ class Provider(Base):
|
||||
# 格式转换时是否保持优先级(默认 False)
|
||||
# - False: 需要格式转换时,该提供商的候选会被降级到不需要转换的候选之后
|
||||
# - True: 即使需要格式转换,也保持原优先级排名
|
||||
# 注意:如果全局配置 KEEP_PRIORITY_ON_CONVERSION=true,此字段被忽略(所有提供商都保持优先级)
|
||||
# 注意:如果系统配置 keep_priority_on_conversion=true,此字段被忽略(所有提供商都保持优先级)
|
||||
keep_priority_on_conversion = Column(Boolean, default=False, nullable=False)
|
||||
|
||||
# 是否允许格式转换(默认 True)
|
||||
@@ -1038,6 +1038,32 @@ class Model(Base):
|
||||
def get_effective_supports_image_generation(self) -> bool:
|
||||
return self._get_effective_capability("supports_image_generation", False)
|
||||
|
||||
def get_effective_config(self) -> dict | None:
|
||||
"""获取有效的 config(合并 Model 和 GlobalModel 的 config)
|
||||
|
||||
合并策略:
|
||||
- GlobalModel.config 作为基础
|
||||
- Model.config 覆盖 GlobalModel.config
|
||||
- 深度合并 billing 子字段
|
||||
"""
|
||||
global_config = {}
|
||||
if self.global_model and self.global_model.config:
|
||||
global_config = dict(self.global_model.config)
|
||||
|
||||
if not self.config:
|
||||
return global_config if global_config else None
|
||||
|
||||
# 深度合并 config
|
||||
result = dict(global_config)
|
||||
for key, value in self.config.items():
|
||||
if key == "billing" and isinstance(value, dict) and isinstance(result.get(key), dict):
|
||||
# 深度合并 billing
|
||||
result[key] = {**result[key], **value}
|
||||
else:
|
||||
result[key] = value
|
||||
|
||||
return result if result else None
|
||||
|
||||
def select_provider_model_name(
|
||||
self, affinity_key: str | None = None, api_format: str | None = None
|
||||
) -> str:
|
||||
@@ -1429,6 +1455,7 @@ class VideoTask(Base):
|
||||
video_urls = Column(JSON)
|
||||
thumbnail_url = Column(String(2000))
|
||||
video_size_bytes = Column(BigInteger)
|
||||
video_duration_seconds = Column(Float) # 实际视频时长(秒)
|
||||
video_expires_at = Column(DateTime(timezone=True))
|
||||
|
||||
# 存储 (可选)
|
||||
|
||||
@@ -31,6 +31,7 @@ from src.services.billing.models import (
|
||||
)
|
||||
from src.services.billing.schema import BillingSnapshot, CostResult
|
||||
from src.services.billing.service import BillingService
|
||||
from src.services.billing.shadow import ShadowBillingService
|
||||
from src.services.billing.templates import BILLING_TEMPLATE_REGISTRY, BillingTemplates
|
||||
from src.services.billing.usage_mapper import UsageMapper, map_usage, map_usage_from_response
|
||||
|
||||
@@ -50,6 +51,7 @@ __all__ = [
|
||||
"BillingService",
|
||||
"BillingSnapshot",
|
||||
"CostResult",
|
||||
"ShadowBillingService",
|
||||
# 映射器
|
||||
"UsageMapper",
|
||||
"map_usage",
|
||||
|
||||
130
src/services/billing/cache.py
Normal file
130
src/services/billing/cache.py
Normal file
@@ -0,0 +1,130 @@
|
||||
"""
|
||||
Billing in-process cache.
|
||||
|
||||
This module provides a small TTL cache for billing rule lookups and other
|
||||
high-read, low-churn billing configuration objects.
|
||||
|
||||
Important:
|
||||
- Keep cached values *session-agnostic*. Avoid caching SQLAlchemy ORM objects
|
||||
bound to a specific Session; prefer plain dataclasses / dicts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BillingCache:
|
||||
"""
|
||||
Simple TTL + LRU cache.
|
||||
|
||||
- TTL: 300s (5 minutes)
|
||||
- Max entries per cache: 2048 (evict oldest on overflow)
|
||||
"""
|
||||
|
||||
TTL_SECONDS = 300
|
||||
MAX_ENTRIES = 2048
|
||||
|
||||
_rule_cache: dict[str, tuple[Any, float]] = {}
|
||||
_collector_cache: dict[str, tuple[Any, float]] = {}
|
||||
_default_rule_cache: dict[str, tuple[Any, float]] = {}
|
||||
|
||||
# ----------------------------
|
||||
# Rule cache
|
||||
# ----------------------------
|
||||
@classmethod
|
||||
def get_rule(cls, cache_key: str) -> Any | None:
|
||||
return cls._get(cls._rule_cache, cache_key)
|
||||
|
||||
@classmethod
|
||||
def set_rule(cls, cache_key: str, value: Any) -> None:
|
||||
cls._set(cls._rule_cache, cache_key, value)
|
||||
|
||||
# ----------------------------
|
||||
# Default-rule cache
|
||||
# ----------------------------
|
||||
@classmethod
|
||||
def get_default_rule(cls, cache_key: str) -> Any | None:
|
||||
return cls._get(cls._default_rule_cache, cache_key)
|
||||
|
||||
@classmethod
|
||||
def set_default_rule(cls, cache_key: str, value: Any) -> None:
|
||||
cls._set(cls._default_rule_cache, cache_key, value)
|
||||
|
||||
# ----------------------------
|
||||
# Collector cache (reserved)
|
||||
# ----------------------------
|
||||
@classmethod
|
||||
def get_collectors(cls, cache_key: str) -> Any | None:
|
||||
return cls._get(cls._collector_cache, cache_key)
|
||||
|
||||
@classmethod
|
||||
def set_collectors(cls, cache_key: str, value: Any) -> None:
|
||||
cls._set(cls._collector_cache, cache_key, value)
|
||||
|
||||
# ----------------------------
|
||||
# Invalidation
|
||||
# ----------------------------
|
||||
@classmethod
|
||||
def invalidate_all(cls) -> None:
|
||||
cls._rule_cache.clear()
|
||||
cls._collector_cache.clear()
|
||||
cls._default_rule_cache.clear()
|
||||
|
||||
@classmethod
|
||||
def invalidate_model(cls, model_name: str) -> None:
|
||||
"""
|
||||
Invalidate cache entries referencing a model name.
|
||||
|
||||
Note:
|
||||
- This is best-effort string matching (cache key format must include model_name).
|
||||
"""
|
||||
cls._invalidate_by_substring(cls._rule_cache, model_name)
|
||||
cls._invalidate_by_substring(cls._default_rule_cache, model_name)
|
||||
|
||||
# ----------------------------
|
||||
# Internal helpers
|
||||
# ----------------------------
|
||||
@classmethod
|
||||
def _get(cls, cache: dict[str, tuple[Any, float]], key: str) -> Any | None:
|
||||
item = cache.get(key)
|
||||
if item is None:
|
||||
return None
|
||||
value, ts = item
|
||||
if time.time() - ts < cls.TTL_SECONDS:
|
||||
return value
|
||||
# expired
|
||||
try:
|
||||
del cache[key]
|
||||
except KeyError:
|
||||
pass
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _set(cls, cache: dict[str, tuple[Any, float]], key: str, value: Any) -> None:
|
||||
"""Set with LRU eviction when cache exceeds MAX_ENTRIES."""
|
||||
now = time.time()
|
||||
cache[key] = (value, now)
|
||||
|
||||
# Evict oldest entries if over limit
|
||||
if len(cache) > cls.MAX_ENTRIES:
|
||||
cls._evict_oldest(cache, cls.MAX_ENTRIES // 4)
|
||||
|
||||
@classmethod
|
||||
def _evict_oldest(cls, cache: dict[str, tuple[Any, float]], count: int) -> None:
|
||||
"""Evict the oldest `count` entries from cache."""
|
||||
if not cache or count <= 0:
|
||||
return
|
||||
# Sort by timestamp (oldest first) and remove
|
||||
sorted_keys = sorted(cache.keys(), key=lambda k: cache[k][1])
|
||||
for k in sorted_keys[:count]:
|
||||
cache.pop(k, None)
|
||||
|
||||
@staticmethod
|
||||
def _invalidate_by_substring(cache: dict[str, tuple[Any, float]], needle: str) -> None:
|
||||
if not needle:
|
||||
return
|
||||
keys = [k for k in cache.keys() if needle in k]
|
||||
for k in keys:
|
||||
cache.pop(k, None)
|
||||
31
src/services/billing/collector_defs/__init__.py
Normal file
31
src/services/billing/collector_defs/__init__.py
Normal file
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
Collector definitions (config-file mode).
|
||||
|
||||
Goal:
|
||||
- Developers define dimension collectors in code, grouped by api_format.
|
||||
- Adding support for a new api_format should only require adding a new file here.
|
||||
- No DB seeding required.
|
||||
|
||||
Each module should export:
|
||||
- COLLECTORS: list[dict[str, Any]]
|
||||
|
||||
Each dict supports keys (aligned with DimensionCollector):
|
||||
- api_format: "openai:chat" (canonical family:kind)
|
||||
- task_type: "chat" | "cli" | "video" | "image" | "audio"
|
||||
- dimension_name: string
|
||||
- source_type: "request" | "response" | "metadata" | "computed"
|
||||
- source_path: string | None
|
||||
- value_type: "float" | "int" | "string"
|
||||
- transform_expression: string | None
|
||||
- default_value: string | None
|
||||
- priority: int
|
||||
- is_enabled: bool
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# This package is discovered dynamically by `src.services.billing.presets`.
|
||||
|
||||
COLLECTORS: list[dict[str, Any]] = []
|
||||
27
src/services/billing/collector_defs/claude_chat.py
Normal file
27
src/services/billing/collector_defs/claude_chat.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# Anthropic / Claude messages
|
||||
COLLECTORS: list[dict[str, Any]] = [
|
||||
{
|
||||
"api_format": "claude:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "input_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usage.input_tokens",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "claude:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "output_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usage.output_tokens",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
]
|
||||
27
src/services/billing/collector_defs/gemini_chat.py
Normal file
27
src/services/billing/collector_defs/gemini_chat.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# Gemini generateContent
|
||||
COLLECTORS: list[dict[str, Any]] = [
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "input_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usageMetadata.promptTokenCount",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "output_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usageMetadata.candidatesTokenCount",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
]
|
||||
27
src/services/billing/collector_defs/openai_chat.py
Normal file
27
src/services/billing/collector_defs/openai_chat.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# OpenAI chat completions
|
||||
COLLECTORS: list[dict[str, Any]] = [
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "input_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usage.prompt_tokens",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "output_tokens",
|
||||
"source_type": "response",
|
||||
"source_path": "usage.completion_tokens",
|
||||
"value_type": "int",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
]
|
||||
116
src/services/billing/collector_defs/video_common.py
Normal file
116
src/services/billing/collector_defs/video_common.py
Normal file
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# Async video finalize flow: extra dims from metadata (base_dimensions already provided by caller).
|
||||
#
|
||||
# Note:
|
||||
# - DimensionCollectorService has a "video -> base api_format fallback" that may query
|
||||
# base api_format collectors when api_format is "openai:video"/"gemini:video" etc.
|
||||
COLLECTORS: list[dict[str, Any]] = [
|
||||
# Prefer size as "resolution key" (e.g. 1024x1792), fallback to resolution label (e.g. 720p/4k).
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_resolution_key",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.size",
|
||||
"value_type": "string",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_resolution_key",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.resolution",
|
||||
"value_type": "string",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_size_bytes",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.video_size_bytes",
|
||||
"value_type": "int",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
# 实际视频时长(秒),优先使用从 provider 响应中提取的实际时长
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_duration_seconds",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.video_duration_seconds",
|
||||
"value_type": "float",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
# 回退到请求的 duration_seconds(如果没有实际时长)
|
||||
{
|
||||
"api_format": "openai:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_duration_seconds",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.duration_seconds",
|
||||
"value_type": "int",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_resolution_key",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.size",
|
||||
"value_type": "string",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_resolution_key",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.resolution",
|
||||
"value_type": "string",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_size_bytes",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.video_size_bytes",
|
||||
"value_type": "int",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
# 实际视频时长(秒),优先使用从 provider 响应中提取的实际时长
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_duration_seconds",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.video_duration_seconds",
|
||||
"value_type": "float",
|
||||
"priority": 10,
|
||||
"is_enabled": True,
|
||||
},
|
||||
# 回退到请求的 duration_seconds(如果没有实际时长)
|
||||
{
|
||||
"api_format": "gemini:chat",
|
||||
"task_type": "video",
|
||||
"dimension_name": "video_duration_seconds",
|
||||
"source_type": "metadata",
|
||||
"source_path": "task.duration_seconds",
|
||||
"value_type": "int",
|
||||
"priority": 0,
|
||||
"is_enabled": True,
|
||||
},
|
||||
]
|
||||
259
src/services/billing/default_rules.py
Normal file
259
src/services/billing/default_rules.py
Normal file
@@ -0,0 +1,259 @@
|
||||
"""
|
||||
Default billing rules (runtime-generated).
|
||||
|
||||
Goal:
|
||||
- Keep backward compatibility with existing GlobalModel/Model pricing config
|
||||
(tiered_pricing + price_per_request)
|
||||
- Provide a virtual BillingRule when no explicit BillingRule is configured in DB.
|
||||
|
||||
This module MUST NOT write to DB.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from src.models.database import GlobalModel, Model
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VirtualBillingRule:
|
||||
"""A rule object compatible with BillingRule fields, generated at runtime."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
task_type: str
|
||||
expression: str
|
||||
variables: dict[str, Any]
|
||||
dimension_mappings: dict[str, Any]
|
||||
is_virtual: bool = True
|
||||
|
||||
|
||||
def _as_float(value: Any, *, default: float = 0.0) -> float:
|
||||
try:
|
||||
if value is None:
|
||||
return default
|
||||
# avoid bool being treated as int
|
||||
if isinstance(value, bool):
|
||||
return default
|
||||
return float(value)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def _get_tiers(tiered_pricing: dict | None) -> list[dict[str, Any]]:
|
||||
if not isinstance(tiered_pricing, dict):
|
||||
return []
|
||||
tiers = tiered_pricing.get("tiers")
|
||||
if not isinstance(tiers, list):
|
||||
return []
|
||||
return [t for t in tiers if isinstance(t, dict)]
|
||||
|
||||
|
||||
class DefaultBillingRuleGenerator:
|
||||
"""
|
||||
Build a virtual BillingRule from GlobalModel/Model pricing fields.
|
||||
|
||||
Pricing sources:
|
||||
- Tiered pricing: Model.tiered_pricing overrides GlobalModel.default_tiered_pricing
|
||||
- Per-request price: Model.price_per_request overrides GlobalModel.default_price_per_request
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def generate_for_model(
|
||||
*,
|
||||
global_model: GlobalModel,
|
||||
model: Model | None = None,
|
||||
task_type: str = "chat",
|
||||
) -> VirtualBillingRule:
|
||||
tiered_pricing = (
|
||||
model.get_effective_tiered_pricing()
|
||||
if model is not None
|
||||
else global_model.default_tiered_pricing
|
||||
)
|
||||
tiers = _get_tiers(tiered_pricing)
|
||||
|
||||
# Base prices (used as defaults if tier_key missing)
|
||||
first_tier = tiers[0] if tiers else {}
|
||||
base_input_price = _as_float(first_tier.get("input_price_per_1m"), default=0.0)
|
||||
base_output_price = _as_float(first_tier.get("output_price_per_1m"), default=0.0)
|
||||
|
||||
# Cache prices: keep legacy behavior (derive from input price when missing)
|
||||
base_cache_creation_price = _as_float(
|
||||
first_tier.get("cache_creation_price_per_1m"),
|
||||
default=base_input_price * 1.25,
|
||||
)
|
||||
base_cache_read_price = _as_float(
|
||||
first_tier.get("cache_read_price_per_1m"),
|
||||
default=base_input_price * 0.1,
|
||||
)
|
||||
|
||||
# Per-request price
|
||||
if model is not None:
|
||||
request_price = model.get_effective_price_per_request()
|
||||
else:
|
||||
request_price = global_model.default_price_per_request
|
||||
base_request_price = _as_float(request_price, default=0.0)
|
||||
|
||||
# Expression uses per-1M prices and token counts.
|
||||
# v2-friendly expression: total cost is sum of component costs.
|
||||
expression = (
|
||||
"input_cost + output_cost + cache_creation_cost + cache_read_cost + request_cost"
|
||||
)
|
||||
|
||||
variables: dict[str, Any] = {
|
||||
"input_price_per_1m": base_input_price,
|
||||
"output_price_per_1m": base_output_price,
|
||||
"cache_creation_price_per_1m": base_cache_creation_price,
|
||||
"cache_read_price_per_1m": base_cache_read_price,
|
||||
"price_per_request": base_request_price,
|
||||
}
|
||||
|
||||
dimension_mappings: dict[str, Any] = {
|
||||
# Raw dimensions
|
||||
"input_tokens": {
|
||||
"source": "dimension",
|
||||
"key": "input_tokens",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 0,
|
||||
},
|
||||
"output_tokens": {
|
||||
"source": "dimension",
|
||||
"key": "output_tokens",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 0,
|
||||
},
|
||||
"cache_creation_tokens": {
|
||||
"source": "dimension",
|
||||
"key": "cache_creation_tokens",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 0,
|
||||
},
|
||||
"cache_read_tokens": {
|
||||
"source": "dimension",
|
||||
"key": "cache_read_tokens",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 0,
|
||||
},
|
||||
"request_count": {
|
||||
"source": "dimension",
|
||||
"key": "request_count",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 1,
|
||||
},
|
||||
# Component costs (computed)
|
||||
"input_cost": {
|
||||
"source": "computed",
|
||||
"expression": "input_tokens * input_price_per_1m / 1000000",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
},
|
||||
"output_cost": {
|
||||
"source": "computed",
|
||||
"expression": "output_tokens * output_price_per_1m / 1000000",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
},
|
||||
"cache_creation_cost": {
|
||||
"source": "computed",
|
||||
"expression": "cache_creation_tokens * cache_creation_price_per_1m / 1000000",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
},
|
||||
"cache_read_cost": {
|
||||
"source": "computed",
|
||||
"expression": "cache_read_tokens * cache_read_price_per_1m / 1000000",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
},
|
||||
"request_cost": {
|
||||
"source": "computed",
|
||||
"expression": "request_count * price_per_request",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
},
|
||||
}
|
||||
|
||||
# Tiered pricing: resolve effective prices based on total_input_context
|
||||
# (legacy definition: input_tokens + cache_read_tokens)
|
||||
if tiers:
|
||||
# Build tier lists with legacy cache fallbacks per tier.
|
||||
def _tier_value(
|
||||
t: dict[str, Any], key: str, *, default_multiplier: float | None = None
|
||||
) -> float:
|
||||
if key in t and t.get(key) is not None:
|
||||
return _as_float(t.get(key), default=0.0)
|
||||
if default_multiplier is not None:
|
||||
input_price = _as_float(t.get("input_price_per_1m"), default=0.0)
|
||||
return input_price * default_multiplier
|
||||
return 0.0
|
||||
|
||||
def _tiers_for(
|
||||
key: str,
|
||||
*,
|
||||
default_multiplier: float | None = None,
|
||||
include_cache_ttl_pricing: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
out: list[dict[str, Any]] = []
|
||||
for t in tiers:
|
||||
item: dict[str, Any] = {
|
||||
"up_to": t.get("up_to"),
|
||||
"value": _tier_value(t, key, default_multiplier=default_multiplier),
|
||||
}
|
||||
if include_cache_ttl_pricing and isinstance(t.get("cache_ttl_pricing"), list):
|
||||
# Preserve raw ttl pricing list for FormulaEngine tiered resolver.
|
||||
item["cache_ttl_pricing"] = t.get("cache_ttl_pricing")
|
||||
out.append(item)
|
||||
return out
|
||||
|
||||
tier_key = "total_input_context"
|
||||
dimension_mappings["input_price_per_1m"] = {
|
||||
"source": "tiered",
|
||||
"tier_key": tier_key,
|
||||
"allow_zero": True,
|
||||
"tiers": _tiers_for("input_price_per_1m"),
|
||||
"default": base_input_price,
|
||||
}
|
||||
dimension_mappings["output_price_per_1m"] = {
|
||||
"source": "tiered",
|
||||
"tier_key": tier_key,
|
||||
"allow_zero": True,
|
||||
"tiers": _tiers_for("output_price_per_1m"),
|
||||
"default": base_output_price,
|
||||
}
|
||||
dimension_mappings["cache_creation_price_per_1m"] = {
|
||||
"source": "tiered",
|
||||
"tier_key": tier_key,
|
||||
"allow_zero": True,
|
||||
"tiers": _tiers_for("cache_creation_price_per_1m", default_multiplier=1.25),
|
||||
"default": base_cache_creation_price,
|
||||
}
|
||||
dimension_mappings["cache_read_price_per_1m"] = {
|
||||
"source": "tiered",
|
||||
"tier_key": tier_key,
|
||||
"allow_zero": True,
|
||||
# TTL override supported when dims include cache_ttl_minutes
|
||||
"ttl_key": "cache_ttl_minutes",
|
||||
"ttl_value_key": "cache_read_price_per_1m",
|
||||
"tiers": _tiers_for(
|
||||
"cache_read_price_per_1m",
|
||||
default_multiplier=0.1,
|
||||
include_cache_ttl_pricing=True,
|
||||
),
|
||||
"default": base_cache_read_price,
|
||||
}
|
||||
|
||||
return VirtualBillingRule(
|
||||
id="__default__",
|
||||
name=f"Default rule for {getattr(global_model, 'name', 'unknown')}",
|
||||
task_type=task_type,
|
||||
expression=expression,
|
||||
variables=variables,
|
||||
dimension_mappings=dimension_mappings,
|
||||
)
|
||||
@@ -10,24 +10,40 @@ DimensionCollector 运行时维度采集
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Literal, Protocol
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import DimensionCollector
|
||||
from src.services.billing.cache import BillingCache
|
||||
from src.services.billing.formula_engine import (
|
||||
ExpressionEvaluationError,
|
||||
SafeExpressionEvaluator,
|
||||
UnsafeExpressionError,
|
||||
extract_variable_names,
|
||||
)
|
||||
from src.services.billing.presets import CORE_PRESET_PACK
|
||||
|
||||
ValueType = Literal["float", "int", "string"]
|
||||
|
||||
|
||||
class CollectorLike(Protocol):
|
||||
api_format: str
|
||||
task_type: str
|
||||
dimension_name: str
|
||||
source_type: str
|
||||
source_path: str | None
|
||||
value_type: str
|
||||
transform_expression: str | None
|
||||
default_value: str | None
|
||||
priority: int
|
||||
is_enabled: bool
|
||||
|
||||
|
||||
def _normalize_api_format(api_format: str | None) -> str:
|
||||
if not api_format:
|
||||
return ""
|
||||
@@ -88,6 +104,23 @@ def _type_default(value_type: ValueType) -> Any:
|
||||
return "" if value_type == "string" else (0 if value_type == "int" else 0.0)
|
||||
|
||||
|
||||
_WXH_PATTERN = re.compile(r"^(\d+)x(\d+)$")
|
||||
|
||||
|
||||
def _normalize_resolution_key(raw: str) -> str:
|
||||
"""
|
||||
Normalize resolution key:
|
||||
- lowercase, remove spaces, × → x
|
||||
- For WxH format, sort dimensions so smaller comes first (1080x720 → 720x1080)
|
||||
"""
|
||||
k = (raw or "").strip().lower().replace(" ", "").replace("×", "x")
|
||||
match = _WXH_PATTERN.match(k)
|
||||
if match:
|
||||
a, b = int(match.group(1)), int(match.group(2))
|
||||
k = f"{a}x{b}" if a <= b else f"{b}x{a}"
|
||||
return k
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DimensionCollectInput:
|
||||
request: dict[str, Any] | None = None
|
||||
@@ -105,13 +138,13 @@ class DimensionCollectorRuntime:
|
||||
def collect(
|
||||
self,
|
||||
*,
|
||||
collectors: list[DimensionCollector],
|
||||
collectors: list[CollectorLike],
|
||||
inp: DimensionCollectInput,
|
||||
) -> dict[str, Any]:
|
||||
dims: dict[str, Any] = dict(inp.base_dimensions or {})
|
||||
|
||||
# dimension_name -> collectors (priority desc)
|
||||
grouped: dict[str, list[DimensionCollector]] = {}
|
||||
grouped: dict[str, list[CollectorLike]] = {}
|
||||
for c in collectors:
|
||||
grouped.setdefault(c.dimension_name, []).append(c)
|
||||
for name in grouped:
|
||||
@@ -144,7 +177,7 @@ class DimensionCollectorRuntime:
|
||||
def _resolve_dimension(
|
||||
self,
|
||||
dim_name: str,
|
||||
collectors: list[DimensionCollector],
|
||||
collectors: list[CollectorLike],
|
||||
dims: dict[str, Any],
|
||||
inp: DimensionCollectInput,
|
||||
) -> Any:
|
||||
@@ -207,7 +240,7 @@ class DimensionCollectorRuntime:
|
||||
def _resolve_computed_dimension(
|
||||
self,
|
||||
dim_name: str,
|
||||
collectors: list[DimensionCollector],
|
||||
collectors: list[CollectorLike],
|
||||
dims: dict[str, Any],
|
||||
) -> Any:
|
||||
fallback_default: str | None = None
|
||||
@@ -245,7 +278,7 @@ class DimensionCollectorRuntime:
|
||||
|
||||
def _toposort_computed(
|
||||
self,
|
||||
grouped: dict[str, list[DimensionCollector]],
|
||||
grouped: dict[str, list[CollectorLike]],
|
||||
computed_only: set[str],
|
||||
) -> list[str]:
|
||||
# 建图:dependency -> dim
|
||||
@@ -300,7 +333,7 @@ class DimensionCollectorRuntime:
|
||||
|
||||
|
||||
class DimensionCollectorService:
|
||||
"""DB + runtime 的封装:读取 collectors 并执行采集。"""
|
||||
"""运行时读取 collectors 并执行采集(code-only)。"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
@@ -311,14 +344,67 @@ class DimensionCollectorService:
|
||||
*,
|
||||
api_format: str | None,
|
||||
task_type: str | None,
|
||||
) -> list[DimensionCollector]:
|
||||
) -> list[CollectorLike]:
|
||||
api = _normalize_api_format(api_format)
|
||||
task = _normalize_task_type(task_type)
|
||||
if not api or not task:
|
||||
return []
|
||||
|
||||
# Code-defined collectors cache.
|
||||
cache_key = f"code:{api}:{task}"
|
||||
cached = BillingCache.get_collectors(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
built = self._list_builtin_collectors(api_format=api_format, task_type=task_type)
|
||||
BillingCache.set_collectors(cache_key, built)
|
||||
return built
|
||||
|
||||
def _list_builtin_collectors(
|
||||
self,
|
||||
*,
|
||||
api_format: str | None,
|
||||
task_type: str | None,
|
||||
) -> list[DimensionCollector]:
|
||||
"""
|
||||
Built-in (code) collectors.
|
||||
|
||||
Developers ship a curated set of collectors in code (config-file mode).
|
||||
"""
|
||||
api = _normalize_api_format(api_format)
|
||||
task = _normalize_task_type(task_type)
|
||||
if not api or not task:
|
||||
return []
|
||||
|
||||
def _preset_query(api_keys: list[str], task_t: str) -> list[DimensionCollector]:
|
||||
out: list[DimensionCollector] = []
|
||||
allowed = {k for k in api_keys if k}
|
||||
for p in CORE_PRESET_PACK.collectors:
|
||||
if not p.is_enabled:
|
||||
continue
|
||||
if _normalize_api_format(p.api_format) not in allowed:
|
||||
continue
|
||||
if _normalize_task_type(p.task_type) != task_t:
|
||||
continue
|
||||
out.append(
|
||||
DimensionCollector(
|
||||
api_format=_normalize_api_format(p.api_format),
|
||||
task_type=_normalize_task_type(p.task_type),
|
||||
dimension_name=p.dimension_name,
|
||||
source_type=p.source_type,
|
||||
source_path=p.source_path,
|
||||
value_type=p.value_type,
|
||||
transform_expression=p.transform_expression,
|
||||
default_value=p.default_value,
|
||||
priority=int(p.priority or 0),
|
||||
is_enabled=True,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
api_variants = list({api, api.lower()})
|
||||
|
||||
if task == "video":
|
||||
# VIDEO → base 回退:优先使用 family:video 专用 collector;
|
||||
# 缺失的维度再回退到 family:chat。
|
||||
from src.core.api_format.signature import parse_signature_key
|
||||
|
||||
base_api = api
|
||||
@@ -330,24 +416,8 @@ class DimensionCollectorService:
|
||||
base_api = api
|
||||
base_variants = list({base_api, base_api.lower()})
|
||||
|
||||
video_collectors = (
|
||||
self.db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format.in_(api_variants),
|
||||
DimensionCollector.task_type == "video",
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
base_collectors = (
|
||||
self.db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format.in_(base_variants),
|
||||
DimensionCollector.task_type == "video",
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
video_collectors = _preset_query(api_variants, "video")
|
||||
base_collectors = _preset_query(base_variants, "video")
|
||||
video_dims: set[str] = {c.dimension_name for c in video_collectors}
|
||||
result: list[DimensionCollector] = list(video_collectors)
|
||||
for c in base_collectors:
|
||||
@@ -356,25 +426,8 @@ class DimensionCollectorService:
|
||||
return result
|
||||
|
||||
if task == "cli":
|
||||
# CLI → chat:按维度回退(维度存在 cli collector 则用 cli,否则用 chat)
|
||||
cli_collectors = (
|
||||
self.db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format.in_(api_variants),
|
||||
DimensionCollector.task_type == "cli",
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
chat_collectors = (
|
||||
self.db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format.in_(api_variants),
|
||||
DimensionCollector.task_type == "chat",
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
cli_collectors = _preset_query(api_variants, "cli")
|
||||
chat_collectors = _preset_query(api_variants, "chat")
|
||||
cli_dims: set[str] = {c.dimension_name for c in cli_collectors}
|
||||
result: list[DimensionCollector] = list(cli_collectors)
|
||||
for c in chat_collectors:
|
||||
@@ -382,15 +435,7 @@ class DimensionCollectorService:
|
||||
result.append(c)
|
||||
return result
|
||||
|
||||
return (
|
||||
self.db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format.in_(api_variants),
|
||||
DimensionCollector.task_type == task,
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
return _preset_query(api_variants, task)
|
||||
|
||||
def collect_dimensions(
|
||||
self,
|
||||
@@ -403,7 +448,7 @@ class DimensionCollectorService:
|
||||
base_dimensions: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
collectors = self.list_enabled_collectors(api_format=api_format, task_type=task_type)
|
||||
return self._runtime.collect(
|
||||
dims = self._runtime.collect(
|
||||
collectors=collectors,
|
||||
inp=DimensionCollectInput(
|
||||
request=request,
|
||||
@@ -412,3 +457,9 @@ class DimensionCollectorService:
|
||||
base_dimensions=base_dimensions,
|
||||
),
|
||||
)
|
||||
# Post-process: normalize video_resolution_key (e.g., 1080x720 → 720x1080)
|
||||
if "video_resolution_key" in dims:
|
||||
raw = dims["video_resolution_key"]
|
||||
if isinstance(raw, str) and raw:
|
||||
dims["video_resolution_key"] = _normalize_resolution_key(raw)
|
||||
return dims
|
||||
|
||||
@@ -12,9 +12,13 @@ FormulaEngine - 配置驱动的安全计费表达式引擎
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from decimal import Decimal
|
||||
from functools import lru_cache
|
||||
from typing import Any, Iterable, Literal
|
||||
|
||||
from src.services.billing.precision import DECIMAL_CONTEXT_PRECISION, to_decimal
|
||||
|
||||
|
||||
class UnsafeExpressionError(ValueError):
|
||||
"""表达式包含不安全/不支持的 AST 结构。"""
|
||||
@@ -35,9 +39,13 @@ class BillingIncompleteError(RuntimeError):
|
||||
@dataclass(frozen=True)
|
||||
class FormulaEvaluationResult:
|
||||
status: Literal["complete", "incomplete"]
|
||||
cost: float
|
||||
resolved_values: dict[str, Any]
|
||||
missing_required: list[str]
|
||||
cost: Decimal
|
||||
resolved_dimensions: dict[str, Any]
|
||||
resolved_variables: dict[str, Any]
|
||||
cost_breakdown: dict[str, Decimal] = field(default_factory=dict)
|
||||
tier_index: int | None = None
|
||||
tier_info: dict[str, Any] | None = None
|
||||
missing_required: list[str] = field(default_factory=list)
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@@ -53,6 +61,9 @@ _ALLOWED_BINOPS = (
|
||||
_ALLOWED_UNARYOPS = (ast.UAdd, ast.USub)
|
||||
_ALLOWED_OP_NODES = _ALLOWED_BINOPS + _ALLOWED_UNARYOPS
|
||||
|
||||
# Allowed function names used in expressions.
|
||||
_ALLOWED_FUNC_NAMES = frozenset(("min", "max", "abs", "round", "int", "float"))
|
||||
|
||||
|
||||
def _iter_ast_nodes(node: ast.AST) -> Iterable[ast.AST]:
|
||||
yield node
|
||||
@@ -60,13 +71,92 @@ def _iter_ast_nodes(node: ast.AST) -> Iterable[ast.AST]:
|
||||
yield from _iter_ast_nodes(child)
|
||||
|
||||
|
||||
def extract_variable_names(expression: str) -> set[str]:
|
||||
"""提取表达式中出现的变量名(不含函数名)。"""
|
||||
@lru_cache(maxsize=2048)
|
||||
def _validate_expression_cached(expression: str) -> ast.Expression:
|
||||
"""
|
||||
Parse + validate an expression and cache the resulting AST.
|
||||
|
||||
This is a hot path (called for every billing evaluation and many collector transforms),
|
||||
so we cache validated ASTs to avoid repeated ast.parse + whitelist scans.
|
||||
"""
|
||||
try:
|
||||
tree = ast.parse(expression, mode="eval")
|
||||
except SyntaxError as exc:
|
||||
raise UnsafeExpressionError(f"Invalid expression syntax: {exc}") from exc
|
||||
|
||||
for node in _iter_ast_nodes(tree):
|
||||
if isinstance(node, ast.Expression):
|
||||
continue
|
||||
# 运算符节点本身也会出现在 iter_child_nodes 中
|
||||
if isinstance(node, _ALLOWED_OP_NODES):
|
||||
continue
|
||||
if isinstance(node, ast.Constant):
|
||||
# 仅允许数字常量(bool 是 int 子类,需要显式排除)
|
||||
if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
|
||||
raise UnsafeExpressionError("Only int/float constants are allowed")
|
||||
continue
|
||||
if isinstance(node, ast.BinOp):
|
||||
if not isinstance(node.op, _ALLOWED_BINOPS):
|
||||
raise UnsafeExpressionError(f"Operator not allowed: {type(node.op).__name__}")
|
||||
continue
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
if not isinstance(node.op, _ALLOWED_UNARYOPS):
|
||||
raise UnsafeExpressionError(f"Unary operator not allowed: {type(node.op).__name__}")
|
||||
continue
|
||||
if isinstance(node, ast.Name):
|
||||
# 防御:拒绝双下划线变量名
|
||||
if node.id.startswith("__"):
|
||||
raise UnsafeExpressionError("Dunder names are not allowed")
|
||||
continue
|
||||
if isinstance(node, ast.Load):
|
||||
continue
|
||||
if isinstance(node, ast.keyword):
|
||||
continue
|
||||
if isinstance(node, ast.Call):
|
||||
if not isinstance(node.func, ast.Name):
|
||||
raise UnsafeExpressionError("Only direct function calls are allowed")
|
||||
func_name = node.func.id
|
||||
if func_name not in _ALLOWED_FUNC_NAMES:
|
||||
raise UnsafeExpressionError(f"Function not allowed: {func_name}")
|
||||
if any(k.arg is None for k in node.keywords):
|
||||
raise UnsafeExpressionError("**kwargs is not allowed")
|
||||
continue
|
||||
|
||||
# 明确禁止的/不需要的节点类型(属性访问、下标、推导式、比较等)
|
||||
if isinstance(
|
||||
node,
|
||||
(
|
||||
ast.Attribute,
|
||||
ast.Subscript,
|
||||
ast.Compare,
|
||||
ast.BoolOp,
|
||||
ast.IfExp,
|
||||
ast.Lambda,
|
||||
ast.Dict,
|
||||
ast.List,
|
||||
ast.Tuple,
|
||||
ast.Set,
|
||||
ast.ListComp,
|
||||
ast.SetComp,
|
||||
ast.DictComp,
|
||||
ast.GeneratorExp,
|
||||
ast.Await,
|
||||
ast.Yield,
|
||||
ast.YieldFrom,
|
||||
),
|
||||
):
|
||||
raise UnsafeExpressionError(f"AST node not allowed: {type(node).__name__}")
|
||||
|
||||
raise UnsafeExpressionError(f"AST node not allowed: {type(node).__name__}")
|
||||
|
||||
assert isinstance(tree, ast.Expression)
|
||||
return tree
|
||||
|
||||
|
||||
def extract_variable_names(expression: str) -> set[str]:
|
||||
"""提取表达式中出现的变量名(不含函数名)。"""
|
||||
tree = _validate_expression_cached(expression)
|
||||
|
||||
names: set[str] = set()
|
||||
for node in _iter_ast_nodes(tree):
|
||||
if isinstance(node, ast.Name):
|
||||
@@ -81,110 +171,154 @@ def extract_variable_names(expression: str) -> set[str]:
|
||||
class SafeExpressionEvaluator:
|
||||
"""AST 白名单 + 无 builtins 的安全求值器。"""
|
||||
|
||||
ALLOWED_FUNCS: dict[str, Any] = {
|
||||
"min": min,
|
||||
"max": max,
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
"int": int,
|
||||
"float": float,
|
||||
}
|
||||
def __init__(self) -> None:
|
||||
# Decimal-friendly allowed functions (return Decimal)
|
||||
self.ALLOWED_FUNCS: dict[str, Any] = {
|
||||
"min": self._min,
|
||||
"max": self._max,
|
||||
"abs": self._abs,
|
||||
"round": self._round,
|
||||
"int": self._int,
|
||||
"float": self._float,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _min(*args: Any) -> Decimal:
|
||||
return min(to_decimal(a) for a in args)
|
||||
|
||||
@staticmethod
|
||||
def _max(*args: Any) -> Decimal:
|
||||
return max(to_decimal(a) for a in args)
|
||||
|
||||
@staticmethod
|
||||
def _abs(x: Any) -> Decimal:
|
||||
return abs(to_decimal(x))
|
||||
|
||||
@staticmethod
|
||||
def _round(x: Any, ndigits: Any = 0) -> Decimal:
|
||||
# Round Decimal returns Decimal; coerce ndigits to int safely.
|
||||
try:
|
||||
n = int(ndigits)
|
||||
except Exception:
|
||||
n = 0
|
||||
return round(to_decimal(x), n)
|
||||
|
||||
@staticmethod
|
||||
def _int(x: Any) -> Decimal:
|
||||
return to_decimal(int(to_decimal(x)))
|
||||
|
||||
@staticmethod
|
||||
def _float(x: Any) -> Decimal:
|
||||
# Keep numeric chain in Decimal even if caller used float()
|
||||
return to_decimal(float(to_decimal(x)))
|
||||
|
||||
def validate(self, expression: str) -> ast.Expression:
|
||||
try:
|
||||
tree = ast.parse(expression, mode="eval")
|
||||
except SyntaxError as exc:
|
||||
raise UnsafeExpressionError(f"Invalid expression syntax: {exc}") from exc
|
||||
return _validate_expression_cached(expression)
|
||||
|
||||
for node in _iter_ast_nodes(tree):
|
||||
if isinstance(node, ast.Expression):
|
||||
continue
|
||||
# 运算符节点本身也会出现在 iter_child_nodes 中
|
||||
if isinstance(node, _ALLOWED_OP_NODES):
|
||||
continue
|
||||
if isinstance(node, ast.Constant):
|
||||
# 仅允许数字常量(bool 是 int 子类,需要显式排除)
|
||||
if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
|
||||
raise UnsafeExpressionError("Only int/float constants are allowed")
|
||||
continue
|
||||
if isinstance(node, ast.BinOp):
|
||||
if not isinstance(node.op, _ALLOWED_BINOPS):
|
||||
raise UnsafeExpressionError(f"Operator not allowed: {type(node.op).__name__}")
|
||||
continue
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
if not isinstance(node.op, _ALLOWED_UNARYOPS):
|
||||
raise UnsafeExpressionError(
|
||||
f"Unary operator not allowed: {type(node.op).__name__}"
|
||||
)
|
||||
continue
|
||||
if isinstance(node, ast.Name):
|
||||
# 防御:拒绝双下划线变量名
|
||||
if node.id.startswith("__"):
|
||||
raise UnsafeExpressionError("Dunder names are not allowed")
|
||||
continue
|
||||
if isinstance(node, ast.Load):
|
||||
continue
|
||||
if isinstance(node, ast.keyword):
|
||||
continue
|
||||
if isinstance(node, ast.Call):
|
||||
if not isinstance(node.func, ast.Name):
|
||||
raise UnsafeExpressionError("Only direct function calls are allowed")
|
||||
func_name = node.func.id
|
||||
if func_name not in self.ALLOWED_FUNCS:
|
||||
raise UnsafeExpressionError(f"Function not allowed: {func_name}")
|
||||
if any(k.arg is None for k in node.keywords):
|
||||
raise UnsafeExpressionError("**kwargs is not allowed")
|
||||
continue
|
||||
def eval_decimal(self, expression: str, variables: dict[str, Any]) -> Decimal:
|
||||
"""
|
||||
Evaluate expression into Decimal.
|
||||
|
||||
# 明确禁止的/不需要的节点类型(属性访问、下标、推导式、比较等)
|
||||
if isinstance(
|
||||
node,
|
||||
(
|
||||
ast.Attribute,
|
||||
ast.Subscript,
|
||||
ast.Compare,
|
||||
ast.BoolOp,
|
||||
ast.IfExp,
|
||||
ast.Lambda,
|
||||
ast.Dict,
|
||||
ast.List,
|
||||
ast.Tuple,
|
||||
ast.Set,
|
||||
ast.ListComp,
|
||||
ast.SetComp,
|
||||
ast.DictComp,
|
||||
ast.GeneratorExp,
|
||||
ast.Await,
|
||||
ast.Yield,
|
||||
ast.YieldFrom,
|
||||
),
|
||||
):
|
||||
raise UnsafeExpressionError(f"AST node not allowed: {type(node).__name__}")
|
||||
|
||||
raise UnsafeExpressionError(f"AST node not allowed: {type(node).__name__}")
|
||||
|
||||
assert isinstance(tree, ast.Expression)
|
||||
return tree
|
||||
|
||||
def eval_number(self, expression: str, variables: dict[str, Any]) -> float:
|
||||
We avoid Python eval() here to ensure:
|
||||
- float literals don't leak binary float arithmetic
|
||||
- all arithmetic stays within Decimal
|
||||
"""
|
||||
tree = self.validate(expression)
|
||||
|
||||
safe_globals = {"__builtins__": {}}
|
||||
safe_locals = dict(self.ALLOWED_FUNCS)
|
||||
safe_locals.update(variables or {})
|
||||
|
||||
try:
|
||||
compiled = compile(tree, "<billing_expr>", "eval")
|
||||
value = eval(compiled, safe_globals, safe_locals) # noqa: S307 - validated AST
|
||||
with _decimal_context(DECIMAL_CONTEXT_PRECISION):
|
||||
return _eval_decimal(tree.body, variables or {}, self.ALLOWED_FUNCS)
|
||||
except NameError:
|
||||
raise
|
||||
except ExpressionEvaluationError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ExpressionEvaluationError(str(exc)) from exc
|
||||
|
||||
def eval_number(self, expression: str, variables: dict[str, Any]) -> float:
|
||||
"""Backward-compatible float evaluation (used by DimensionCollector transforms)."""
|
||||
value = self.eval_decimal(expression, variables)
|
||||
try:
|
||||
return float(value)
|
||||
except Exception as exc:
|
||||
raise ExpressionEvaluationError(f"Expression result is not numeric: {value!r}") from exc
|
||||
|
||||
|
||||
class _decimal_context:
|
||||
def __init__(self, prec: int):
|
||||
self.prec = prec
|
||||
|
||||
def __enter__(self) -> None:
|
||||
from decimal import getcontext
|
||||
|
||||
self._ctx = getcontext().copy()
|
||||
getcontext().prec = self.prec
|
||||
|
||||
def __exit__(self, exc_type: type | None, exc: BaseException | None, tb: Any) -> None:
|
||||
from decimal import setcontext
|
||||
|
||||
# Restore full context to avoid leaking settings.
|
||||
setcontext(self._ctx)
|
||||
|
||||
|
||||
def _eval_decimal(node: ast.AST, variables: dict[str, Any], funcs: dict[str, Any]) -> Decimal:
|
||||
if isinstance(node, ast.Constant):
|
||||
return to_decimal(node.value)
|
||||
if isinstance(node, ast.Name):
|
||||
if node.id not in variables:
|
||||
raise NameError(node.id)
|
||||
return to_decimal(variables[node.id])
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
v = _eval_decimal(node.operand, variables, funcs)
|
||||
if isinstance(node.op, ast.UAdd):
|
||||
return v
|
||||
if isinstance(node.op, ast.USub):
|
||||
return -v
|
||||
raise ExpressionEvaluationError(f"Unary operator not allowed: {type(node.op).__name__}")
|
||||
if isinstance(node, ast.BinOp):
|
||||
left = _eval_decimal(node.left, variables, funcs)
|
||||
right = _eval_decimal(node.right, variables, funcs)
|
||||
if isinstance(node.op, ast.Add):
|
||||
return left + right
|
||||
if isinstance(node.op, ast.Sub):
|
||||
return left - right
|
||||
if isinstance(node.op, ast.Mult):
|
||||
return left * right
|
||||
if isinstance(node.op, ast.Div):
|
||||
return left / right
|
||||
if isinstance(node.op, ast.FloorDiv):
|
||||
return left // right
|
||||
if isinstance(node.op, ast.Mod):
|
||||
return left % right
|
||||
if isinstance(node.op, ast.Pow):
|
||||
# Decimal power is only well-defined for integer exponents here.
|
||||
try:
|
||||
exp_int = int(right)
|
||||
if to_decimal(exp_int) != right:
|
||||
raise ValueError("non-integer exponent")
|
||||
except Exception as exc:
|
||||
raise ExpressionEvaluationError("Pow only supports integer exponents") from exc
|
||||
return left**exp_int
|
||||
raise ExpressionEvaluationError(f"Operator not allowed: {type(node.op).__name__}")
|
||||
if isinstance(node, ast.Call):
|
||||
if not isinstance(node.func, ast.Name):
|
||||
raise ExpressionEvaluationError("Only direct function calls are allowed")
|
||||
func_name = node.func.id
|
||||
func = funcs.get(func_name)
|
||||
if func is None:
|
||||
raise ExpressionEvaluationError(f"Function not allowed: {func_name}")
|
||||
args = [_eval_decimal(a, variables, funcs) for a in node.args]
|
||||
kwargs = {
|
||||
kw.arg: _eval_decimal(kw.value, variables, funcs) for kw in node.keywords if kw.arg
|
||||
}
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
except Exception as exc:
|
||||
raise ExpressionEvaluationError(str(exc)) from exc
|
||||
return to_decimal(result)
|
||||
|
||||
raise ExpressionEvaluationError(f"AST node not allowed: {type(node).__name__}")
|
||||
|
||||
|
||||
class FormulaEngine:
|
||||
"""计费表达式引擎:解析 dimension_mappings 并进行安全求值。"""
|
||||
|
||||
@@ -205,19 +339,59 @@ class FormulaEngine:
|
||||
resolved: dict[str, Any] = dict(variables or {})
|
||||
|
||||
missing_required: list[str] = []
|
||||
tier_index: int | None = None
|
||||
tier_info: dict[str, Any] | None = None
|
||||
|
||||
# 先解析 dimension_mappings,产出 expression 变量表
|
||||
computed: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# 1) Resolve non-computed mappings first
|
||||
for var_name, mapping in mappings.items():
|
||||
source = (mapping.get("source") or "constant").lower()
|
||||
# 显式 constant 映射属于“兜底行为”:如果 variables 已经提供该变量,则不覆盖。
|
||||
if source == "computed":
|
||||
computed[var_name] = mapping
|
||||
continue
|
||||
# Explicit constant mapping is fallback-only when variable already exists.
|
||||
if source == "constant" and var_name in resolved:
|
||||
continue
|
||||
value, is_missing = self._resolve_mapping(var_name, mapping, dims)
|
||||
value, is_missing, tier_meta = self._resolve_mapping(var_name, mapping, dims)
|
||||
if tier_meta and tier_index is None:
|
||||
tier_index = tier_meta.get("tier_index")
|
||||
tier_info = tier_meta.get("tier_info")
|
||||
if is_missing:
|
||||
missing_required.append(var_name)
|
||||
continue
|
||||
resolved[var_name] = value
|
||||
|
||||
# 2) Resolve computed mappings (iterative dependency resolution)
|
||||
if computed:
|
||||
unresolved = dict(computed)
|
||||
for _ in range(max(4, len(unresolved) + 1)):
|
||||
progressed = False
|
||||
for var_name, mapping in list(unresolved.items()):
|
||||
if var_name in resolved:
|
||||
unresolved.pop(var_name, None)
|
||||
continue
|
||||
value, status = self._try_resolve_computed(var_name, mapping, dims, resolved)
|
||||
if status == "pending":
|
||||
continue
|
||||
unresolved.pop(var_name, None)
|
||||
if status == "missing_required":
|
||||
missing_required.append(var_name)
|
||||
continue
|
||||
resolved[var_name] = value
|
||||
progressed = True
|
||||
if not progressed:
|
||||
break
|
||||
|
||||
# any remaining unresolved computed vars
|
||||
for var_name, mapping in unresolved.items():
|
||||
required = bool(mapping.get("required", False))
|
||||
default = mapping.get("default", 0)
|
||||
if required:
|
||||
missing_required.append(var_name)
|
||||
else:
|
||||
resolved[var_name] = default
|
||||
|
||||
# required 维度缺失:直接标记 incomplete(并由 strict_mode 决定是否抛错)
|
||||
if missing_required:
|
||||
if strict_mode:
|
||||
@@ -227,45 +401,120 @@ class FormulaEngine:
|
||||
)
|
||||
return FormulaEvaluationResult(
|
||||
status="incomplete",
|
||||
cost=0.0,
|
||||
resolved_values=resolved,
|
||||
cost=Decimal("0"),
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables=resolved,
|
||||
missing_required=missing_required,
|
||||
tier_index=tier_index,
|
||||
tier_info=tier_info,
|
||||
)
|
||||
|
||||
# 3) Evaluate total cost
|
||||
try:
|
||||
cost = self._evaluator.eval_number(expression, resolved)
|
||||
cost = self._evaluator.eval_decimal(expression, resolved)
|
||||
if cost < 0:
|
||||
# 防御:不允许负数成本(通常表示配置错误)
|
||||
return FormulaEvaluationResult(
|
||||
status="incomplete",
|
||||
cost=0.0,
|
||||
resolved_values=resolved,
|
||||
cost=Decimal("0"),
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables=resolved,
|
||||
missing_required=[],
|
||||
tier_index=tier_index,
|
||||
tier_info=tier_info,
|
||||
error="negative_cost",
|
||||
)
|
||||
|
||||
breakdown = self._extract_cost_breakdown(resolved)
|
||||
return FormulaEvaluationResult(
|
||||
status="complete",
|
||||
cost=cost,
|
||||
resolved_values=resolved,
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables=resolved,
|
||||
cost_breakdown=breakdown,
|
||||
tier_index=tier_index,
|
||||
tier_info=tier_info,
|
||||
missing_required=[],
|
||||
)
|
||||
except NameError as exc:
|
||||
# expression references missing vars
|
||||
if strict_mode:
|
||||
raise ExpressionEvaluationError(f"Missing variable: {exc}") from exc
|
||||
return FormulaEvaluationResult(
|
||||
status="incomplete",
|
||||
cost=Decimal("0"),
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables=resolved,
|
||||
missing_required=[],
|
||||
tier_index=tier_index,
|
||||
tier_info=tier_info,
|
||||
error=f"missing_variable:{exc}",
|
||||
)
|
||||
except (UnsafeExpressionError, ExpressionEvaluationError) as exc:
|
||||
if strict_mode:
|
||||
raise
|
||||
return FormulaEvaluationResult(
|
||||
status="incomplete",
|
||||
cost=0.0,
|
||||
resolved_values=resolved,
|
||||
cost=Decimal("0"),
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables=resolved,
|
||||
missing_required=[],
|
||||
tier_index=tier_index,
|
||||
tier_info=tier_info,
|
||||
error=str(exc),
|
||||
)
|
||||
|
||||
def _extract_cost_breakdown(self, resolved: dict[str, Any]) -> dict[str, Decimal]:
|
||||
breakdown: dict[str, Decimal] = {}
|
||||
for k, v in resolved.items():
|
||||
if not k.endswith("_cost"):
|
||||
continue
|
||||
try:
|
||||
breakdown[k] = to_decimal(v)
|
||||
except Exception:
|
||||
continue
|
||||
return breakdown
|
||||
|
||||
def _try_resolve_computed(
|
||||
self,
|
||||
var_name: str,
|
||||
mapping: dict[str, Any],
|
||||
dims: dict[str, Any],
|
||||
resolved: dict[str, Any],
|
||||
) -> tuple[Any, Literal["ok", "pending", "missing_required"]]:
|
||||
"""
|
||||
Try resolve a computed mapping.
|
||||
|
||||
Returns:
|
||||
(value, status)
|
||||
- ok: value computed
|
||||
- pending: missing dependencies, retry later
|
||||
- missing_required: required=true and cannot resolve
|
||||
"""
|
||||
required = bool(mapping.get("required", False))
|
||||
default = mapping.get("default", 0)
|
||||
expr = mapping.get("expression") or mapping.get("transform_expression")
|
||||
if not expr:
|
||||
return (None, "missing_required") if required else (default, "ok")
|
||||
# Computed vars can reference both resolved variables and raw dims.
|
||||
env: dict[str, Any] = {}
|
||||
env.update(dims)
|
||||
env.update(resolved)
|
||||
try:
|
||||
value = self._evaluator.eval_decimal(str(expr), env)
|
||||
return value, "ok"
|
||||
except NameError:
|
||||
# dependency not ready yet
|
||||
return (None, "pending") if required else (default, "pending")
|
||||
except Exception:
|
||||
# treat as config error: fallback to default unless required
|
||||
return (None, "missing_required") if required else (default, "ok")
|
||||
|
||||
def _resolve_mapping(
|
||||
self,
|
||||
var_name: str,
|
||||
mapping: dict[str, Any],
|
||||
dims: dict[str, Any],
|
||||
) -> tuple[Any, bool]:
|
||||
) -> tuple[Any, bool, dict[str, Any] | None]:
|
||||
"""
|
||||
Returns:
|
||||
(value, is_missing_required)
|
||||
@@ -287,82 +536,185 @@ class FormulaEngine:
|
||||
|
||||
if source == "constant":
|
||||
# constant 默认行为:由 variables 提供;dimension_mappings 显式 constant 时仅做兜底
|
||||
return default, False
|
||||
return default, False, None
|
||||
|
||||
if source == "dimension":
|
||||
key = mapping.get("key") or var_name
|
||||
raw = dims.get(key)
|
||||
if raw is None:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
if isinstance(raw, str):
|
||||
if raw == "":
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
# 尝试将字符串解析为数字,否则按字符串返回(供上层自行决定)
|
||||
try:
|
||||
num = float(raw)
|
||||
num = to_decimal(raw)
|
||||
if num == 0 and not allow_zero:
|
||||
return _missing()
|
||||
return num, False
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
return num, False, None
|
||||
except Exception:
|
||||
return raw, False
|
||||
if isinstance(raw, (int, float)):
|
||||
if float(raw) == 0 and not allow_zero:
|
||||
return _missing()
|
||||
return raw, False
|
||||
return raw, False, None
|
||||
if isinstance(raw, (int, float, Decimal)):
|
||||
num = to_decimal(raw)
|
||||
if num == 0 and not allow_zero:
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
return num, False, None
|
||||
# 其他类型:尽量转为 float,否则视为缺失
|
||||
try:
|
||||
num = float(raw)
|
||||
num = to_decimal(raw)
|
||||
if num == 0 and not allow_zero:
|
||||
return _missing()
|
||||
return num, False
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
return num, False, None
|
||||
except Exception:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
|
||||
if source == "matrix":
|
||||
key = mapping.get("key") or var_name
|
||||
raw = dims.get(key)
|
||||
if raw is None or raw == "":
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
raw_key = str(raw)
|
||||
matrix = mapping.get("map") or {}
|
||||
if raw_key in matrix:
|
||||
return matrix[raw_key], False
|
||||
try:
|
||||
return to_decimal(matrix[raw_key]), False, None
|
||||
except Exception:
|
||||
return matrix[raw_key], False, None
|
||||
# matrix 未命中:若 required=true 则仍视为缺失;否则使用 default
|
||||
if required:
|
||||
return None, True
|
||||
return default, False
|
||||
return None, True, None
|
||||
return default, False, None
|
||||
|
||||
if source == "tiered":
|
||||
tier_key = mapping.get("tier_key")
|
||||
if not tier_key:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
raw_tier_value = dims.get(tier_key)
|
||||
if raw_tier_value is None:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
try:
|
||||
tier_value = float(raw_tier_value)
|
||||
tier_value = to_decimal(raw_tier_value)
|
||||
except Exception:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
|
||||
if tier_value == 0 and not allow_zero:
|
||||
return _missing()
|
||||
v, m = _missing()
|
||||
return v, m, None
|
||||
|
||||
# Optional TTL override (legacy: Claude cache pricing)
|
||||
ttl_key = mapping.get("ttl_key")
|
||||
ttl_value_key = mapping.get("ttl_value_key")
|
||||
ttl_minutes: Decimal | None = None
|
||||
if ttl_key and ttl_value_key and dims.get(ttl_key) is not None:
|
||||
try:
|
||||
ttl_minutes = to_decimal(dims.get(ttl_key))
|
||||
except Exception:
|
||||
ttl_minutes = None
|
||||
|
||||
tiers = mapping.get("tiers") or []
|
||||
# tiers: [{up_to: 128000, value: 2.5}, {up_to: null, value: 1.25}]
|
||||
for tier in tiers:
|
||||
for idx, tier in enumerate(tiers):
|
||||
up_to = tier.get("up_to")
|
||||
if up_to is None:
|
||||
return tier.get("value", default), False
|
||||
value = to_decimal(tier.get("value", default))
|
||||
if (
|
||||
ttl_minutes is not None
|
||||
and ttl_value_key
|
||||
and isinstance(tier.get("cache_ttl_pricing"), list)
|
||||
):
|
||||
value = self._resolve_ttl_pricing(
|
||||
tier.get("cache_ttl_pricing") or [],
|
||||
ttl_minutes,
|
||||
str(ttl_value_key),
|
||||
fallback=value,
|
||||
)
|
||||
return value, False, {"tier_index": idx, "tier_info": dict(tier)}
|
||||
try:
|
||||
if tier_value <= float(up_to):
|
||||
return tier.get("value", default), False
|
||||
if tier_value <= to_decimal(up_to):
|
||||
value = to_decimal(tier.get("value", default))
|
||||
if (
|
||||
ttl_minutes is not None
|
||||
and ttl_value_key
|
||||
and isinstance(tier.get("cache_ttl_pricing"), list)
|
||||
):
|
||||
value = self._resolve_ttl_pricing(
|
||||
tier.get("cache_ttl_pricing") or [],
|
||||
ttl_minutes,
|
||||
str(ttl_value_key),
|
||||
fallback=value,
|
||||
)
|
||||
return value, False, {"tier_index": idx, "tier_info": dict(tier)}
|
||||
except Exception:
|
||||
# up_to 配置异常:忽略并继续
|
||||
continue
|
||||
# 无匹配:使用最后一个或 default
|
||||
if tiers:
|
||||
return tiers[-1].get("value", default), False
|
||||
return default, False
|
||||
last = tiers[-1]
|
||||
value = to_decimal(last.get("value", default))
|
||||
if (
|
||||
ttl_minutes is not None
|
||||
and ttl_value_key
|
||||
and isinstance(last.get("cache_ttl_pricing"), list)
|
||||
):
|
||||
value = self._resolve_ttl_pricing(
|
||||
last.get("cache_ttl_pricing") or [],
|
||||
ttl_minutes,
|
||||
str(ttl_value_key),
|
||||
fallback=value,
|
||||
)
|
||||
return value, False, {"tier_index": len(tiers) - 1, "tier_info": dict(last)}
|
||||
return default, False, None
|
||||
|
||||
# 未知 source:视为配置错误,但不直接中断计费(返回 default)
|
||||
return default, False
|
||||
return default, False, None
|
||||
|
||||
def _resolve_ttl_pricing(
|
||||
self,
|
||||
ttl_pricing: list[Any],
|
||||
ttl_minutes: Decimal,
|
||||
ttl_value_key: str,
|
||||
*,
|
||||
fallback: Decimal,
|
||||
) -> Decimal:
|
||||
"""
|
||||
Resolve TTL-dependent pricing (legacy: cache_ttl_pricing).
|
||||
|
||||
Rules:
|
||||
- pick the first entry whose ttl_minutes >= requested ttl
|
||||
- otherwise pick the last entry
|
||||
- if missing/invalid, fallback to base value
|
||||
"""
|
||||
try:
|
||||
entries = [e for e in ttl_pricing if isinstance(e, dict)]
|
||||
if not entries:
|
||||
return fallback
|
||||
|
||||
def _ttl_key(e: dict[str, Any]) -> Decimal:
|
||||
return to_decimal(e.get("ttl_minutes") or 0)
|
||||
|
||||
entries_sorted = sorted(entries, key=_ttl_key)
|
||||
chosen: dict[str, Any] = entries_sorted[-1]
|
||||
for e in entries_sorted:
|
||||
try:
|
||||
if ttl_minutes <= to_decimal(e.get("ttl_minutes") or 0):
|
||||
chosen = e
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
v = chosen.get(ttl_value_key)
|
||||
if v is None:
|
||||
return fallback
|
||||
return to_decimal(v)
|
||||
except Exception:
|
||||
return fallback
|
||||
|
||||
49
src/services/billing/precision.py
Normal file
49
src/services/billing/precision.py
Normal file
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Precision helpers for billing calculations.
|
||||
|
||||
We standardize money arithmetic with `Decimal` and quantize to stable precisions.
|
||||
|
||||
Notes:
|
||||
- `Decimal` context precision (`DECIMAL_CONTEXT_PRECISION`) is **significant digits**,
|
||||
not "decimal places".
|
||||
- We keep these as constants (not runtime-configurable) to avoid drift between
|
||||
environments during billing reconciliation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from decimal import ROUND_HALF_UP, Decimal, localcontext
|
||||
|
||||
# Decimal context precision (significant digits)
|
||||
DECIMAL_CONTEXT_PRECISION = 28
|
||||
|
||||
# Money precisions
|
||||
BILLING_STORAGE_PRECISION = 8 # persisted to DB / metadata
|
||||
BILLING_DISPLAY_PRECISION = 6 # UI display
|
||||
|
||||
|
||||
def to_decimal(value: float | int | str | Decimal | None) -> Decimal:
|
||||
"""Convert values to Decimal safely (float via str to avoid binary artifacts)."""
|
||||
if value is None:
|
||||
return Decimal("0")
|
||||
if isinstance(value, Decimal):
|
||||
return value
|
||||
return Decimal(str(value))
|
||||
|
||||
|
||||
def quantize_decimal(value: Decimal, *, precision: int) -> Decimal:
|
||||
"""Quantize a Decimal to the given number of decimal places (ROUND_HALF_UP)."""
|
||||
quantizer = Decimal(10) ** -precision
|
||||
with localcontext() as ctx:
|
||||
ctx.prec = DECIMAL_CONTEXT_PRECISION
|
||||
return value.quantize(quantizer, rounding=ROUND_HALF_UP)
|
||||
|
||||
|
||||
def quantize_cost(value: Decimal) -> Decimal:
|
||||
"""Quantize to storage precision."""
|
||||
return quantize_decimal(value, precision=BILLING_STORAGE_PRECISION)
|
||||
|
||||
|
||||
def quantize_display(value: Decimal) -> Decimal:
|
||||
"""Quantize to display precision."""
|
||||
return quantize_decimal(value, precision=BILLING_DISPLAY_PRECISION)
|
||||
254
src/services/billing/presets.py
Normal file
254
src/services/billing/presets.py
Normal file
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Billing presets (developer-provided defaults).
|
||||
|
||||
Why:
|
||||
- Asking end-users to configure DimensionCollectors / BillingRules from scratch is too complex.
|
||||
- We ship a curated set of "known-good" collector presets per api_format/task_type
|
||||
and provide an Admin API to apply them into DB (merge or overwrite).
|
||||
|
||||
Notes:
|
||||
- BillingRule presets are intentionally NOT materialized here, because the unified
|
||||
billing architecture already provides a runtime default rule generator that stays
|
||||
in-sync with Model/GlobalModel pricing. Persisting those prices into BillingRule
|
||||
rows would become stale when model pricing changes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import pkgutil
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
from src.models.database import DimensionCollector
|
||||
|
||||
PresetApplyMode = Literal["merge", "overwrite"]
|
||||
|
||||
|
||||
def _norm_api(api_format: str) -> str:
|
||||
return normalize_signature_key(api_format or "")
|
||||
|
||||
|
||||
def _norm_task(task_type: str) -> str:
|
||||
return (task_type or "").strip().lower()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CollectorPreset:
|
||||
api_format: str
|
||||
task_type: str
|
||||
dimension_name: str
|
||||
|
||||
source_type: str
|
||||
source_path: str | None = None
|
||||
|
||||
value_type: str = "float" # float/int/string
|
||||
transform_expression: str | None = None
|
||||
default_value: str | None = None
|
||||
|
||||
priority: int = 0
|
||||
is_enabled: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PresetPack:
|
||||
name: str
|
||||
version: str
|
||||
description: str
|
||||
collectors: list[CollectorPreset]
|
||||
|
||||
|
||||
def _discover_collectors() -> list[CollectorPreset]:
|
||||
"""
|
||||
Config-file mode: discover collectors from `src.services.billing.collector_defs`.
|
||||
|
||||
Developers add a new file under that package; no central registry edits required.
|
||||
"""
|
||||
out: list[CollectorPreset] = []
|
||||
|
||||
try:
|
||||
pkg = importlib.import_module("src.services.billing.collector_defs")
|
||||
pkg_path = getattr(pkg, "__path__", None)
|
||||
if not pkg_path:
|
||||
return out
|
||||
except Exception:
|
||||
return out
|
||||
|
||||
for mod in pkgutil.iter_modules(pkg_path):
|
||||
if mod.ispkg:
|
||||
continue
|
||||
mod_name = f"src.services.billing.collector_defs.{mod.name}"
|
||||
try:
|
||||
m = importlib.import_module(mod_name)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
items = getattr(m, "COLLECTORS", None)
|
||||
if not isinstance(items, list):
|
||||
continue
|
||||
|
||||
for raw in items:
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
try:
|
||||
out.append(
|
||||
CollectorPreset(
|
||||
api_format=str(raw.get("api_format") or "").strip(),
|
||||
task_type=str(raw.get("task_type") or "").strip().lower(),
|
||||
dimension_name=str(raw.get("dimension_name") or "").strip(),
|
||||
source_type=str(raw.get("source_type") or "").strip().lower(),
|
||||
source_path=raw.get("source_path"),
|
||||
value_type=str(raw.get("value_type") or "float").strip().lower(),
|
||||
transform_expression=raw.get("transform_expression"),
|
||||
default_value=raw.get("default_value"),
|
||||
priority=int(raw.get("priority") or 0),
|
||||
is_enabled=bool(raw.get("is_enabled", True)),
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return out
|
||||
|
||||
|
||||
CORE_PRESET_PACK = PresetPack(
|
||||
name="aether-core",
|
||||
version="1.0",
|
||||
description="Aether built-in dimension collectors for common api_formats/task_types.",
|
||||
collectors=_discover_collectors(),
|
||||
)
|
||||
|
||||
|
||||
def list_preset_packs() -> list[PresetPack]:
|
||||
return [CORE_PRESET_PACK]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PresetApplyResult:
|
||||
preset: str
|
||||
mode: PresetApplyMode
|
||||
created: int
|
||||
updated: int
|
||||
skipped: int
|
||||
errors: list[str]
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"preset": self.preset,
|
||||
"mode": self.mode,
|
||||
"created": self.created,
|
||||
"updated": self.updated,
|
||||
"skipped": self.skipped,
|
||||
"errors": list(self.errors),
|
||||
}
|
||||
|
||||
|
||||
class BillingPresetService:
|
||||
@staticmethod
|
||||
def apply_preset(
|
||||
db: Session,
|
||||
*,
|
||||
preset_name: str,
|
||||
mode: PresetApplyMode = "merge",
|
||||
) -> PresetApplyResult:
|
||||
preset_name = (preset_name or "").strip()
|
||||
packs = {p.name: p for p in list_preset_packs()}
|
||||
pack = packs.get(preset_name)
|
||||
if pack is None:
|
||||
available = ", ".join(sorted(packs.keys()))
|
||||
return PresetApplyResult(
|
||||
preset=preset_name,
|
||||
mode=mode,
|
||||
created=0,
|
||||
updated=0,
|
||||
skipped=0,
|
||||
errors=[f"Unknown preset: {preset_name!r}. Available: {available}"],
|
||||
)
|
||||
|
||||
created = 0
|
||||
updated = 0
|
||||
skipped = 0
|
||||
errors: list[str] = []
|
||||
|
||||
for item in pack.collectors:
|
||||
api_format = _norm_api(item.api_format)
|
||||
task_type = _norm_task(item.task_type)
|
||||
dim = (item.dimension_name or "").strip()
|
||||
|
||||
if not api_format or not task_type or not dim:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
try:
|
||||
existing = (
|
||||
db.query(DimensionCollector)
|
||||
.filter(
|
||||
DimensionCollector.api_format == api_format,
|
||||
DimensionCollector.task_type == task_type,
|
||||
DimensionCollector.dimension_name == dim,
|
||||
DimensionCollector.priority == int(item.priority or 0),
|
||||
DimensionCollector.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.first()
|
||||
)
|
||||
except Exception as exc:
|
||||
errors.append(
|
||||
f"Failed to query collector: api_format={api_format} task_type={task_type} dim={dim}: {exc}"
|
||||
)
|
||||
continue
|
||||
|
||||
if existing is not None:
|
||||
if mode == "overwrite":
|
||||
try:
|
||||
existing.source_type = (item.source_type or "").strip().lower()
|
||||
existing.source_path = item.source_path
|
||||
existing.value_type = (item.value_type or "float").strip().lower()
|
||||
existing.transform_expression = item.transform_expression
|
||||
existing.default_value = item.default_value
|
||||
existing.is_enabled = bool(item.is_enabled)
|
||||
updated += 1
|
||||
except Exception as exc:
|
||||
errors.append(
|
||||
f"Failed to update collector {getattr(existing, 'id', None)}: {exc}"
|
||||
)
|
||||
else:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
try:
|
||||
c = DimensionCollector(
|
||||
api_format=api_format,
|
||||
task_type=task_type,
|
||||
dimension_name=dim,
|
||||
source_type=(item.source_type or "").strip().lower(),
|
||||
source_path=item.source_path,
|
||||
value_type=(item.value_type or "float").strip().lower(),
|
||||
transform_expression=item.transform_expression,
|
||||
default_value=item.default_value,
|
||||
priority=int(item.priority or 0),
|
||||
is_enabled=bool(item.is_enabled),
|
||||
)
|
||||
db.add(c)
|
||||
created += 1
|
||||
except Exception as exc:
|
||||
errors.append(
|
||||
f"Failed to create collector: api_format={api_format} task_type={task_type} dim={dim}: {exc}"
|
||||
)
|
||||
|
||||
try:
|
||||
db.commit()
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
errors.append(f"DB commit failed: {exc}")
|
||||
|
||||
return PresetApplyResult(
|
||||
preset=pack.name,
|
||||
mode=mode,
|
||||
created=created,
|
||||
updated=updated,
|
||||
skipped=skipped,
|
||||
errors=errors,
|
||||
)
|
||||
10
src/services/billing/rule_defs/__init__.py
Normal file
10
src/services/billing/rule_defs/__init__.py
Normal file
@@ -0,0 +1,10 @@
|
||||
"""
|
||||
Billing rule definitions (config-file mode).
|
||||
|
||||
Each module should export:
|
||||
- TEMPLATES: list[CodeBillingRuleTemplate]
|
||||
|
||||
Design goal:
|
||||
- Add a new billing mode by adding a new file here.
|
||||
- No central registry edits required.
|
||||
"""
|
||||
195
src/services/billing/rule_defs/universal.py
Normal file
195
src/services/billing/rule_defs/universal.py
Normal file
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
Universal billing template.
|
||||
|
||||
This is the single unified billing template for all task types.
|
||||
Formula: total = (input_cost + output_cost + cache_creation_cost + cache_read_cost) + request_cost + video_cost
|
||||
|
||||
Each component can be 0 if not applicable for the specific task type.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from src.services.billing.default_rules import DefaultBillingRuleGenerator, VirtualBillingRule
|
||||
from src.services.billing.rule_templates import CodeBillingRuleTemplate, RuleTemplateContext
|
||||
|
||||
|
||||
def _get_nested(obj: object | None, path: str) -> object | None:
|
||||
if not isinstance(obj, dict):
|
||||
return None
|
||||
cur: object = obj
|
||||
for part in (path or "").split("."):
|
||||
if not part:
|
||||
continue
|
||||
if not isinstance(cur, dict):
|
||||
return None
|
||||
cur = cur.get(part) # type: ignore[assignment]
|
||||
return cur
|
||||
|
||||
|
||||
def _as_float(v: object | None) -> float | None:
|
||||
try:
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, bool):
|
||||
return None
|
||||
return float(v)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
_WXH_PATTERN = re.compile(r"^(\d+)x(\d+)$")
|
||||
|
||||
|
||||
def _normalize_resolution_key(raw: str) -> str:
|
||||
"""
|
||||
Normalize resolution key:
|
||||
- lowercase, remove spaces, × → x
|
||||
- For WxH format, sort dimensions so smaller comes first (1080x720 → 720x1080)
|
||||
"""
|
||||
k = (raw or "").strip().lower().replace(" ", "").replace("×", "x")
|
||||
match = _WXH_PATTERN.match(k)
|
||||
if match:
|
||||
a, b = int(match.group(1)), int(match.group(2))
|
||||
k = f"{a}x{b}" if a <= b else f"{b}x{a}"
|
||||
return k
|
||||
|
||||
|
||||
def _effective_unit_price(ctx: RuleTemplateContext) -> float:
|
||||
"""Get video price per second from config."""
|
||||
if ctx.model is not None:
|
||||
v = _as_float(
|
||||
_get_nested(getattr(ctx.model, "config", None), "billing.video.price_per_second")
|
||||
)
|
||||
if v is not None:
|
||||
return v
|
||||
v = _as_float(
|
||||
_get_nested(getattr(ctx.global_model, "config", None), "billing.video.price_per_second")
|
||||
)
|
||||
if v is not None:
|
||||
return v
|
||||
return 0.0
|
||||
|
||||
|
||||
def _effective_resolution_price_per_second(ctx: RuleTemplateContext) -> dict[str, float]:
|
||||
"""
|
||||
Resolution (or size) -> price_per_second.
|
||||
"""
|
||||
for conf in (
|
||||
getattr(ctx.model, "config", None) if ctx.model is not None else None,
|
||||
getattr(ctx.global_model, "config", None),
|
||||
):
|
||||
raw = _get_nested(conf, "billing.video.price_per_second_by_resolution")
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
out: dict[str, float] = {}
|
||||
for k, v in raw.items():
|
||||
fk = _normalize_resolution_key(str(k))
|
||||
fv = _as_float(v)
|
||||
if not fk:
|
||||
continue
|
||||
if fv is None:
|
||||
continue
|
||||
out[fk] = fv
|
||||
if out:
|
||||
return out
|
||||
|
||||
# Backward-compat: resolution multipliers
|
||||
base = _effective_unit_price(ctx)
|
||||
if base and base > 0:
|
||||
for conf in (
|
||||
getattr(ctx.model, "config", None) if ctx.model is not None else None,
|
||||
getattr(ctx.global_model, "config", None),
|
||||
):
|
||||
raw = _get_nested(conf, "billing.video.resolution_multipliers")
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
out2: dict[str, float] = {}
|
||||
for k, v in raw.items():
|
||||
fk = _normalize_resolution_key(str(k))
|
||||
mv = _as_float(v)
|
||||
if not fk:
|
||||
continue
|
||||
if mv is None:
|
||||
continue
|
||||
out2[fk] = float(base) * float(mv)
|
||||
if out2:
|
||||
return out2
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def build_universal(ctx: RuleTemplateContext) -> VirtualBillingRule:
|
||||
"""
|
||||
Build the universal billing rule.
|
||||
|
||||
Formula:
|
||||
total = (input_cost + output_cost + cache_creation_cost + cache_read_cost) + request_cost + video_cost
|
||||
|
||||
Each component defaults to 0 if not configured or not applicable.
|
||||
"""
|
||||
# Base rule: token + per-request
|
||||
base = DefaultBillingRuleGenerator.generate_for_model(
|
||||
global_model=ctx.global_model,
|
||||
model=ctx.model,
|
||||
task_type=ctx.task_type,
|
||||
)
|
||||
|
||||
unit_price = _effective_unit_price(ctx)
|
||||
resolution_price_map = _effective_resolution_price_per_second(ctx)
|
||||
|
||||
variables = dict(base.variables or {})
|
||||
|
||||
dimension_mappings = dict(base.dimension_mappings or {})
|
||||
|
||||
# Video duration dimension
|
||||
dimension_mappings["duration_seconds"] = {
|
||||
"source": "dimension",
|
||||
"key": "duration_seconds",
|
||||
"required": False,
|
||||
"allow_zero": True,
|
||||
"default": 0,
|
||||
}
|
||||
|
||||
# Video price per second (resolved from resolution map or fallback to unit price)
|
||||
dimension_mappings["video_price_per_second"] = {
|
||||
"source": "matrix",
|
||||
"key": "video_resolution_key",
|
||||
"required": False,
|
||||
"default": unit_price,
|
||||
"map": resolution_price_map,
|
||||
}
|
||||
|
||||
# Video cost component
|
||||
dimension_mappings["video_cost"] = {
|
||||
"source": "computed",
|
||||
"required": False,
|
||||
"default": 0,
|
||||
"expression": "duration_seconds * video_price_per_second",
|
||||
}
|
||||
|
||||
# Universal formula: token costs + request cost + video cost
|
||||
# base.expression = "input_cost + output_cost + cache_creation_cost + cache_read_cost + request_cost"
|
||||
expression = f"({base.expression}) + video_cost"
|
||||
|
||||
return VirtualBillingRule(
|
||||
id="__default__",
|
||||
name="Universal Billing Rule",
|
||||
task_type=ctx.task_type,
|
||||
expression=expression,
|
||||
variables=variables,
|
||||
dimension_mappings=dimension_mappings,
|
||||
is_virtual=True,
|
||||
)
|
||||
|
||||
|
||||
TEMPLATES = [
|
||||
CodeBillingRuleTemplate(
|
||||
name="universal",
|
||||
description="Universal billing: (input + output + cache) + request + video. All components default to 0 if not applicable.",
|
||||
task_types={"chat", "cli", "video", "image", "audio"},
|
||||
priority=100, # Highest priority - used for all task types
|
||||
build=build_universal,
|
||||
)
|
||||
]
|
||||
@@ -2,7 +2,7 @@
|
||||
BillingRule 查找逻辑
|
||||
|
||||
查找顺序(与 .plans/humming-seeking-marble.md 一致):
|
||||
1) Model(Provider 级)→ 2) GlobalModel(默认)
|
||||
1) 读取 GlobalModel/Model 价格配置 → 2) 使用代码内置计费模板生成规则(config-file mode)
|
||||
|
||||
注意:
|
||||
- CLI 在计费域等同于 chat:billing_rules.task_type 不含 "cli"
|
||||
@@ -11,13 +11,26 @@ BillingRule 查找逻辑
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
from typing import Any, Literal, Protocol
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.database import BillingRule, GlobalModel, Model
|
||||
from src.config.settings import config
|
||||
from src.models.database import GlobalModel, Model
|
||||
from src.services.billing.cache import BillingCache
|
||||
from src.services.billing.default_rules import DefaultBillingRuleGenerator, VirtualBillingRule
|
||||
from src.services.billing.rule_templates import CodeBillingRuleTemplateService
|
||||
|
||||
TaskType = Literal["chat", "cli", "video", "image", "audio"]
|
||||
BillingRuleScope = Literal["model", "global", "default"]
|
||||
|
||||
|
||||
class BillingRuleLike(Protocol):
|
||||
id: str
|
||||
name: str
|
||||
expression: str
|
||||
variables: dict[str, Any]
|
||||
dimension_mappings: dict[str, Any]
|
||||
|
||||
|
||||
def effective_rule_task_type(task_type: str) -> str:
|
||||
@@ -28,8 +41,8 @@ def effective_rule_task_type(task_type: str) -> str:
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BillingRuleLookupResult:
|
||||
rule: BillingRule
|
||||
scope: Literal["model", "global"]
|
||||
rule: BillingRuleLike
|
||||
scope: BillingRuleScope
|
||||
effective_task_type: str
|
||||
|
||||
|
||||
@@ -44,6 +57,16 @@ class BillingRuleService:
|
||||
) -> BillingRuleLookupResult | None:
|
||||
effective_task = effective_rule_task_type(task_type)
|
||||
|
||||
# Normalize provider_id for cache key to avoid duplicate entries (None vs "").
|
||||
pid = provider_id or ""
|
||||
# Cache must include runtime knobs that affect fallback behavior.
|
||||
cache_key = (
|
||||
f"{pid}:{model_name}:{effective_task}:require={int(config.billing_require_rule)}"
|
||||
)
|
||||
cached = BillingCache.get_rule(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
global_model = (
|
||||
db.query(GlobalModel)
|
||||
.filter(
|
||||
@@ -55,7 +78,9 @@ class BillingRuleService:
|
||||
if not global_model:
|
||||
return None
|
||||
|
||||
# 1) Provider Model 覆盖
|
||||
model_obj: Model | None = None
|
||||
|
||||
# Provider Model(用于覆盖价格配置)
|
||||
if provider_id:
|
||||
model_obj = (
|
||||
db.query(Model)
|
||||
@@ -66,36 +91,43 @@ class BillingRuleService:
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if model_obj:
|
||||
rule = (
|
||||
db.query(BillingRule)
|
||||
.filter(
|
||||
BillingRule.model_id == model_obj.id,
|
||||
BillingRule.task_type == effective_task,
|
||||
BillingRule.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if rule:
|
||||
return BillingRuleLookupResult(
|
||||
rule=rule,
|
||||
scope="model",
|
||||
effective_task_type=effective_task,
|
||||
)
|
||||
|
||||
# 2) GlobalModel 默认规则
|
||||
rule = (
|
||||
db.query(BillingRule)
|
||||
.filter(
|
||||
BillingRule.global_model_id == global_model.id,
|
||||
BillingRule.task_type == effective_task,
|
||||
BillingRule.is_enabled == True, # noqa: E712
|
||||
)
|
||||
.first()
|
||||
# Code templates (config-file mode)
|
||||
code_rule = CodeBillingRuleTemplateService.resolve_rule(
|
||||
global_model=global_model,
|
||||
model=model_obj,
|
||||
provider_id=provider_id,
|
||||
model_name=model_name,
|
||||
task_type=effective_task,
|
||||
)
|
||||
if rule:
|
||||
return BillingRuleLookupResult(
|
||||
rule=rule, scope="global", effective_task_type=effective_task
|
||||
if code_rule is not None:
|
||||
result = BillingRuleLookupResult(
|
||||
rule=code_rule,
|
||||
scope="default",
|
||||
effective_task_type=effective_task,
|
||||
)
|
||||
BillingCache.set_rule(cache_key, result)
|
||||
return result
|
||||
|
||||
# Runtime default rule (backward compatible)
|
||||
#
|
||||
# - Always applies to chat-domain billing (cli is normalized to chat).
|
||||
# - For video/image/audio:
|
||||
# - When BILLING_REQUIRE_RULE=true, caller expects an explicit BillingRule (missing -> no_rule/error).
|
||||
# - When BILLING_REQUIRE_RULE=false, fallback to default rule to preserve legacy pricing semantics
|
||||
# (avoid silent $0 billing due to missing rule).
|
||||
if effective_task == "chat" or not config.billing_require_rule:
|
||||
default_rule = DefaultBillingRuleGenerator.generate_for_model(
|
||||
global_model=global_model,
|
||||
model=model_obj,
|
||||
task_type=effective_task,
|
||||
)
|
||||
result = BillingRuleLookupResult(
|
||||
rule=default_rule,
|
||||
scope="default",
|
||||
effective_task_type=effective_task,
|
||||
)
|
||||
BillingCache.set_rule(cache_key, result)
|
||||
return result
|
||||
|
||||
return None
|
||||
|
||||
129
src/services/billing/rule_templates.py
Normal file
129
src/services/billing/rule_templates.py
Normal file
@@ -0,0 +1,129 @@
|
||||
"""
|
||||
Billing rule templates (config-file mode).
|
||||
|
||||
Goal:
|
||||
- Developers define billing rules in code as templates ("模式一/二/三/四/五 ...").
|
||||
- Adding a new billing template should only require adding a new `*.py` file under
|
||||
`src.services.billing.rule_defs` (no DB / no UI).
|
||||
|
||||
How it works:
|
||||
- Each module under `rule_defs/` exports `TEMPLATES: list[CodeBillingRuleTemplate]`.
|
||||
- We dynamically discover all templates at runtime and pick the best one by:
|
||||
- task_type match
|
||||
- optional match(ctx) predicate
|
||||
- highest priority wins
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import pkgutil
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterable
|
||||
|
||||
from src.models.database import GlobalModel, Model
|
||||
from src.services.billing.default_rules import VirtualBillingRule
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RuleTemplateContext:
|
||||
global_model: GlobalModel
|
||||
model: Model | None
|
||||
provider_id: str | None
|
||||
model_name: str
|
||||
task_type: str
|
||||
|
||||
|
||||
MatchFn = Callable[[RuleTemplateContext], bool]
|
||||
BuildFn = Callable[[RuleTemplateContext], VirtualBillingRule]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodeBillingRuleTemplate:
|
||||
"""
|
||||
A code-defined billing template.
|
||||
|
||||
Notes:
|
||||
- `task_types` are billing-domain task types ("cli" is normalized to "chat" by rule_service).
|
||||
- `build()` must return a VirtualBillingRule-like object (VirtualBillingRule is used here).
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
task_types: set[str]
|
||||
priority: int = 0
|
||||
match: MatchFn | None = None
|
||||
build: BuildFn | None = None
|
||||
|
||||
def supports(self, task_type: str) -> bool:
|
||||
return (task_type or "").lower() in {t.lower() for t in (self.task_types or set())}
|
||||
|
||||
|
||||
def _iter_modules() -> Iterable[str]:
|
||||
try:
|
||||
pkg = importlib.import_module("src.services.billing.rule_defs")
|
||||
pkg_path = getattr(pkg, "__path__", None)
|
||||
if not pkg_path:
|
||||
return []
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
out: list[str] = []
|
||||
for mod in pkgutil.iter_modules(pkg_path):
|
||||
if mod.ispkg:
|
||||
continue
|
||||
out.append(f"src.services.billing.rule_defs.{mod.name}")
|
||||
return out
|
||||
|
||||
|
||||
def discover_rule_templates() -> list[CodeBillingRuleTemplate]:
|
||||
templates: list[CodeBillingRuleTemplate] = []
|
||||
for mod_name in _iter_modules():
|
||||
try:
|
||||
m = importlib.import_module(mod_name)
|
||||
except Exception:
|
||||
continue
|
||||
items = getattr(m, "TEMPLATES", None)
|
||||
if not isinstance(items, list):
|
||||
continue
|
||||
for t in items:
|
||||
if isinstance(t, CodeBillingRuleTemplate):
|
||||
templates.append(t)
|
||||
# higher priority first, stable within same module import order
|
||||
templates.sort(key=lambda x: int(getattr(x, "priority", 0) or 0), reverse=True)
|
||||
return templates
|
||||
|
||||
|
||||
class CodeBillingRuleTemplateService:
|
||||
@staticmethod
|
||||
def resolve_rule(
|
||||
*,
|
||||
global_model: GlobalModel,
|
||||
model: Model | None,
|
||||
provider_id: str | None,
|
||||
model_name: str,
|
||||
task_type: str,
|
||||
) -> VirtualBillingRule | None:
|
||||
ctx = RuleTemplateContext(
|
||||
global_model=global_model,
|
||||
model=model,
|
||||
provider_id=provider_id,
|
||||
model_name=model_name,
|
||||
task_type=(task_type or "").lower(),
|
||||
)
|
||||
for t in discover_rule_templates():
|
||||
if not t.supports(ctx.task_type):
|
||||
continue
|
||||
if t.match is not None:
|
||||
try:
|
||||
if not bool(t.match(ctx)):
|
||||
continue
|
||||
except Exception:
|
||||
continue
|
||||
if t.build is None:
|
||||
continue
|
||||
try:
|
||||
return t.build(ctx)
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
@@ -10,14 +10,26 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal
|
||||
|
||||
BILLING_SNAPSHOT_SCHEMA_VERSION = "1.0"
|
||||
BILLING_SNAPSHOT_SCHEMA_VERSION = "2.0"
|
||||
|
||||
BillingSnapshotStatus = Literal["complete", "incomplete", "no_rule", "legacy"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BillingSnapshot:
|
||||
"""Stable billing snapshot for audit."""
|
||||
"""
|
||||
Stable billing snapshot for audit.
|
||||
|
||||
v2.0 semantics:
|
||||
- resolved_dimensions: final dimension values used (tokens, request_count, etc.)
|
||||
- resolved_variables: final variables used (prices, tier-resolved values, etc.)
|
||||
- cost_breakdown: itemized costs (quantized)
|
||||
- total_cost: quantized total cost (equals sum(cost_breakdown) when breakdown present)
|
||||
|
||||
Backward compatibility:
|
||||
- dimensions_used aliases resolved_dimensions
|
||||
- cost aliases total_cost
|
||||
"""
|
||||
|
||||
schema_version: str = BILLING_SNAPSHOT_SCHEMA_VERSION
|
||||
|
||||
@@ -29,29 +41,65 @@ class BillingSnapshot:
|
||||
# Rule expression (internal, do not expose to clients)
|
||||
expression: str | None = None
|
||||
|
||||
# Dimensions
|
||||
dimensions_used: dict[str, Any] = field(default_factory=dict)
|
||||
# v2: resolved inputs
|
||||
resolved_dimensions: dict[str, Any] = field(default_factory=dict)
|
||||
resolved_variables: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
# v2: breakdown and totals
|
||||
cost_breakdown: dict[str, float] = field(default_factory=dict)
|
||||
total_cost: float = 0.0
|
||||
|
||||
# Tier info (optional)
|
||||
tier_index: int | None = None
|
||||
tier_info: dict[str, Any] | None = None
|
||||
|
||||
# Missing dims
|
||||
missing_required: list[str] = field(default_factory=list)
|
||||
|
||||
# Result
|
||||
cost: float = 0.0
|
||||
# Result status
|
||||
status: BillingSnapshotStatus = "no_rule"
|
||||
|
||||
# Audit
|
||||
calculated_at: str = "" # ISO 8601
|
||||
engine_version: str = "2.0"
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Backward-compatible aliases (v1 fields)
|
||||
# ---------------------------------------------------------------------
|
||||
@property
|
||||
def dimensions_used(self) -> dict[str, Any]:
|
||||
return self.resolved_dimensions
|
||||
|
||||
@property
|
||||
def cost(self) -> float:
|
||||
return self.total_cost
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Serialize snapshot.
|
||||
|
||||
Includes both v2 keys and v1-compatible keys for safer rollouts.
|
||||
"""
|
||||
return {
|
||||
"schema_version": self.schema_version,
|
||||
"rule_id": self.rule_id,
|
||||
"rule_name": self.rule_name,
|
||||
"scope": self.scope,
|
||||
"expression": self.expression,
|
||||
"dimensions_used": self.dimensions_used,
|
||||
# v2
|
||||
"resolved_dimensions": self.resolved_dimensions,
|
||||
"resolved_variables": self.resolved_variables,
|
||||
"cost_breakdown": self.cost_breakdown,
|
||||
"total_cost": self.total_cost,
|
||||
"tier_index": self.tier_index,
|
||||
"tier_info": self.tier_info,
|
||||
"missing_required": self.missing_required,
|
||||
"cost": self.cost,
|
||||
"status": self.status,
|
||||
"calculated_at": self.calculated_at,
|
||||
"engine_version": self.engine_version,
|
||||
# v1 compat
|
||||
"dimensions_used": self.resolved_dimensions,
|
||||
"cost": self.total_cost,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -9,8 +10,8 @@ from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.services.billing.dimension_collector_service import DimensionCollectorService
|
||||
from src.services.billing.formula_engine import BillingIncompleteError, FormulaEngine
|
||||
from src.services.billing.precision import quantize_cost, to_decimal
|
||||
from src.services.billing.rule_service import BillingRuleService
|
||||
from src.services.model.cost import ModelCostService
|
||||
|
||||
from .schema import BILLING_SNAPSHOT_SCHEMA_VERSION, BillingSnapshot, CostResult
|
||||
|
||||
@@ -24,10 +25,25 @@ class BillingService:
|
||||
- It may read billing rules & collectors from DB.
|
||||
"""
|
||||
|
||||
# FormulaEngine is stateless and safe to share within a process.
|
||||
_shared_formula_engine: FormulaEngine | None = None
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
self._formula_engine = FormulaEngine()
|
||||
self._dimension_collector = DimensionCollectorService(db)
|
||||
self._formula_engine = self._get_formula_engine()
|
||||
# Lazy-init: most call sites already provide dimensions (hot path).
|
||||
self._dimension_collector: DimensionCollectorService | None = None
|
||||
|
||||
@classmethod
|
||||
def _get_formula_engine(cls) -> FormulaEngine:
|
||||
if cls._shared_formula_engine is None:
|
||||
cls._shared_formula_engine = FormulaEngine()
|
||||
return cls._shared_formula_engine
|
||||
|
||||
def _get_dimension_collector(self) -> DimensionCollectorService:
|
||||
if self._dimension_collector is None:
|
||||
self._dimension_collector = DimensionCollectorService(self.db)
|
||||
return self._dimension_collector
|
||||
|
||||
def collect_dimensions(
|
||||
self,
|
||||
@@ -39,7 +55,7 @@ class BillingService:
|
||||
metadata: dict[str, Any] | None = None,
|
||||
base_dimensions: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return self._dimension_collector.collect_dimensions(
|
||||
return self._get_dimension_collector().collect_dimensions(
|
||||
api_format=api_format,
|
||||
task_type=task_type,
|
||||
request=request,
|
||||
@@ -68,6 +84,42 @@ class BillingService:
|
||||
"""
|
||||
strict = config.billing_strict_mode if strict_mode is None else bool(strict_mode)
|
||||
|
||||
# Normalize & enrich dimensions (do not mutate caller dict)
|
||||
dims: dict[str, Any] = dict(dimensions or {})
|
||||
|
||||
# Compatibility aliases (legacy fields in some call sites)
|
||||
if "cache_creation_tokens" not in dims and "cache_creation_input_tokens" in dims:
|
||||
dims["cache_creation_tokens"] = dims.get("cache_creation_input_tokens")
|
||||
if "cache_read_tokens" not in dims and "cache_read_input_tokens" in dims:
|
||||
dims["cache_read_tokens"] = dims.get("cache_read_input_tokens")
|
||||
|
||||
# Default request_count=1 for per-request billing
|
||||
if "request_count" not in dims:
|
||||
dims["request_count"] = 1
|
||||
|
||||
# total_input_context is the tier-key for legacy tiered pricing:
|
||||
# default: input_tokens + cache_creation_tokens + cache_read_tokens
|
||||
#
|
||||
# NOTE:
|
||||
# Some adapters (e.g. Claude) include cache_creation tokens in the tier context.
|
||||
# Making this the default avoids per-callsite inconsistency.
|
||||
if "total_input_context" not in dims:
|
||||
try:
|
||||
input_tokens_i = int(float(dims.get("input_tokens") or 0))
|
||||
except Exception:
|
||||
input_tokens_i = 0
|
||||
try:
|
||||
cache_creation_tokens_i = int(float(dims.get("cache_creation_tokens") or 0))
|
||||
except Exception:
|
||||
cache_creation_tokens_i = 0
|
||||
try:
|
||||
cache_read_tokens_i = int(float(dims.get("cache_read_tokens") or 0))
|
||||
except Exception:
|
||||
cache_read_tokens_i = 0
|
||||
dims["total_input_context"] = (
|
||||
input_tokens_i + cache_creation_tokens_i + cache_read_tokens_i
|
||||
)
|
||||
|
||||
lookup = BillingRuleService.find_rule(
|
||||
self.db,
|
||||
provider_id=provider_id,
|
||||
@@ -80,49 +132,60 @@ class BillingService:
|
||||
result = self._formula_engine.evaluate(
|
||||
expression=rule.expression,
|
||||
variables=rule.variables or {},
|
||||
dimensions=dimensions,
|
||||
dimensions=dims,
|
||||
dimension_mappings=rule.dimension_mappings or {},
|
||||
strict_mode=strict,
|
||||
)
|
||||
cost = float(result.cost) if result.status == "complete" else 0.0
|
||||
# ------------------------------------------------------------
|
||||
# Quantize: component costs first, then total = sum(components)
|
||||
# ------------------------------------------------------------
|
||||
breakdown_dec: dict[str, Decimal] = {
|
||||
k: to_decimal(v) for k, v in (result.cost_breakdown or {}).items()
|
||||
}
|
||||
|
||||
breakdown_quantized: dict[str, Decimal] = {
|
||||
k: quantize_cost(v) for k, v in breakdown_dec.items()
|
||||
}
|
||||
total_dec = (
|
||||
quantize_cost(sum(breakdown_quantized.values(), Decimal("0")))
|
||||
if breakdown_quantized
|
||||
else quantize_cost(to_decimal(result.cost))
|
||||
)
|
||||
|
||||
cost_breakdown = {k: float(v) for k, v in breakdown_quantized.items()}
|
||||
total_cost = float(total_dec) if result.status == "complete" else 0.0
|
||||
|
||||
# Filter resolved_variables for JSON safety + semantics clarity:
|
||||
# - remove dims (they live in resolved_dimensions)
|
||||
# - remove *_cost (they live in cost_breakdown)
|
||||
resolved_vars: dict[str, Any] = {}
|
||||
for k, v in (result.resolved_variables or {}).items():
|
||||
if k in (result.resolved_dimensions or {}):
|
||||
continue
|
||||
if k.endswith("_cost"):
|
||||
continue
|
||||
if isinstance(v, Decimal):
|
||||
resolved_vars[k] = str(v)
|
||||
else:
|
||||
resolved_vars[k] = v
|
||||
|
||||
snapshot = BillingSnapshot(
|
||||
schema_version=BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
rule_id=str(rule.id),
|
||||
rule_name=str(rule.name),
|
||||
scope=str(getattr(lookup, "scope", None) or ""),
|
||||
expression=str(rule.expression),
|
||||
dimensions_used=dimensions,
|
||||
resolved_dimensions=result.resolved_dimensions or dims,
|
||||
resolved_variables=resolved_vars,
|
||||
cost_breakdown=cost_breakdown,
|
||||
total_cost=total_cost,
|
||||
tier_index=result.tier_index,
|
||||
tier_info=result.tier_info,
|
||||
missing_required=result.missing_required,
|
||||
cost=cost,
|
||||
status=result.status,
|
||||
calculated_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
return CostResult(cost=cost, status=result.status, snapshot=snapshot)
|
||||
|
||||
# No rule fallback
|
||||
if task_type in ("chat", "cli"):
|
||||
input_tokens = int(dimensions.get("input_tokens") or 0)
|
||||
output_tokens = int(dimensions.get("output_tokens") or 0)
|
||||
cost = float(
|
||||
ModelCostService.calculate_cost(
|
||||
model=model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
)
|
||||
snapshot = BillingSnapshot(
|
||||
schema_version=BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
rule_id=None,
|
||||
rule_name=None,
|
||||
scope=None,
|
||||
expression=None,
|
||||
dimensions_used=dimensions,
|
||||
missing_required=[],
|
||||
cost=cost,
|
||||
status="legacy",
|
||||
calculated_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
return CostResult(cost=cost, status="legacy", snapshot=snapshot)
|
||||
return CostResult(cost=total_cost, status=result.status, snapshot=snapshot)
|
||||
|
||||
logger.warning(
|
||||
"No billing rule for task (task_type={}, model={}, provider_id={})",
|
||||
@@ -136,10 +199,46 @@ class BillingService:
|
||||
rule_name=None,
|
||||
scope=None,
|
||||
expression=None,
|
||||
dimensions_used=dimensions,
|
||||
resolved_dimensions=dims,
|
||||
resolved_variables={},
|
||||
cost_breakdown={},
|
||||
total_cost=0.0,
|
||||
missing_required=[],
|
||||
cost=0.0,
|
||||
status="no_rule",
|
||||
calculated_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
return CostResult(cost=0.0, status="no_rule", snapshot=snapshot)
|
||||
|
||||
def calculate_from_response(
|
||||
self,
|
||||
*,
|
||||
task_type: str,
|
||||
model: str,
|
||||
provider_id: str,
|
||||
api_format: str | None,
|
||||
request: dict[str, Any] | None = None,
|
||||
response: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
base_dimensions: dict[str, Any] | None = None,
|
||||
strict_mode: bool | None = None,
|
||||
) -> CostResult:
|
||||
"""
|
||||
Convenience wrapper:
|
||||
- collect dimensions from request/response/metadata
|
||||
- run billing calculation
|
||||
"""
|
||||
dimensions = self.collect_dimensions(
|
||||
api_format=api_format,
|
||||
task_type=task_type,
|
||||
request=request,
|
||||
response=response,
|
||||
metadata=metadata,
|
||||
base_dimensions=base_dimensions,
|
||||
)
|
||||
return self.calculate(
|
||||
task_type=task_type,
|
||||
model=model,
|
||||
provider_id=provider_id,
|
||||
dimensions=dimensions,
|
||||
strict_mode=strict_mode,
|
||||
)
|
||||
|
||||
323
src/services/billing/shadow.py
Normal file
323
src/services/billing/shadow.py
Normal file
@@ -0,0 +1,323 @@
|
||||
"""
|
||||
Shadow billing (reconciliation period).
|
||||
|
||||
This module runs the new billing engine alongside the legacy billing outcome.
|
||||
Truth vs Shadow is kept strictly separated:
|
||||
- truth_breakdown: the values written into Usage rows (the "billable truth")
|
||||
- shadow_snapshot: new engine snapshot stored only in request_metadata.billing_shadow
|
||||
|
||||
Runtime switch:
|
||||
- config.billing_engine: legacy | shadow | new_with_fallback | new
|
||||
- config.billing_engine_overrides: JSON mapping of "provider/model" patterns -> mode
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Any, Literal
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.core.metrics import (
|
||||
billing_diff_exceeds_threshold_total,
|
||||
billing_fallback_total,
|
||||
billing_invariant_violation_total,
|
||||
billing_requests_total,
|
||||
)
|
||||
from src.services.billing.schema import BillingSnapshot
|
||||
from src.services.billing.service import BillingService
|
||||
|
||||
EngineMode = Literal["legacy", "shadow", "new_with_fallback", "new"]
|
||||
TruthEngine = Literal["legacy", "new"]
|
||||
|
||||
|
||||
@lru_cache(maxsize=32)
|
||||
def _compile_engine_overrides(overrides_raw: str) -> tuple[dict[str, str], list[tuple[str, str]]]:
|
||||
"""
|
||||
Parse and normalize engine overrides.
|
||||
|
||||
Cached to avoid json.loads + dict walk on every request.
|
||||
"""
|
||||
try:
|
||||
overrides = json.loads(overrides_raw or "{}")
|
||||
except Exception:
|
||||
overrides = {}
|
||||
|
||||
exact: dict[str, str] = {}
|
||||
patterns: list[tuple[str, str]] = []
|
||||
|
||||
if isinstance(overrides, dict):
|
||||
for pattern, mode in overrides.items():
|
||||
p = str(pattern)
|
||||
m = str(mode).strip().lower()
|
||||
# fnmatch supports *, ?, and [] character classes.
|
||||
if any(ch in p for ch in ("*", "?", "[")):
|
||||
patterns.append((p, m))
|
||||
else:
|
||||
exact[p] = m
|
||||
|
||||
return exact, patterns
|
||||
|
||||
|
||||
@lru_cache(maxsize=4096)
|
||||
def _resolve_engine_mode_cached(key: str, base_mode: str, overrides_raw: str) -> str:
|
||||
exact, patterns = _compile_engine_overrides(overrides_raw)
|
||||
if key in exact:
|
||||
return exact[key]
|
||||
for pattern, mode in patterns:
|
||||
try:
|
||||
if fnmatch.fnmatch(key, pattern):
|
||||
return mode
|
||||
except Exception:
|
||||
continue
|
||||
return base_mode
|
||||
|
||||
|
||||
def resolve_engine_mode(provider: str, model: str) -> EngineMode:
|
||||
"""Resolve engine mode with overrides (pure function, no DB)."""
|
||||
base_mode = (config.billing_engine or "legacy").strip().lower()
|
||||
overrides_raw = getattr(config, "billing_engine_overrides", "{}") or "{}"
|
||||
|
||||
key = f"{provider}/{model}"
|
||||
return _resolve_engine_mode_cached(key, base_mode, overrides_raw) # type: ignore[return-value]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CostBreakdown:
|
||||
"""Cost breakdown written into Usage rows (truth)."""
|
||||
|
||||
input_cost: float
|
||||
output_cost: float
|
||||
cache_creation_cost: float
|
||||
cache_read_cost: float
|
||||
request_cost: float
|
||||
total_cost: float
|
||||
|
||||
@property
|
||||
def cache_cost(self) -> float:
|
||||
return float(self.cache_creation_cost) + float(self.cache_read_cost)
|
||||
|
||||
def validate(self) -> bool:
|
||||
"""
|
||||
Invariant: total_cost == sum(components) (within tiny tolerance).
|
||||
|
||||
For new engine we quantize and sum components deterministically, so this should be exact.
|
||||
For legacy floats, we allow a tiny epsilon.
|
||||
"""
|
||||
computed_total = (
|
||||
float(self.input_cost)
|
||||
+ float(self.output_cost)
|
||||
+ float(self.cache_creation_cost)
|
||||
+ float(self.cache_read_cost)
|
||||
+ float(self.request_cost)
|
||||
)
|
||||
return abs(computed_total - float(self.total_cost)) < 1e-8
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShadowBillingResult:
|
||||
# billable truth (written to Usage table)
|
||||
truth_breakdown: CostBreakdown
|
||||
# shadow snapshot (written to request_metadata.billing_shadow only)
|
||||
shadow_snapshot: BillingSnapshot | None
|
||||
# reconciliation information (diffs etc.)
|
||||
comparison: dict[str, Any]
|
||||
|
||||
# policy vs actual
|
||||
engine_mode: EngineMode = "legacy"
|
||||
truth_engine: TruthEngine = "legacy"
|
||||
was_fallback: bool = False
|
||||
|
||||
|
||||
class ShadowBillingService:
|
||||
"""
|
||||
Shadow billing orchestrator.
|
||||
|
||||
This service does NOT write DB rows. Callers decide how to persist truth and shadow data.
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session) -> None:
|
||||
self.db = db
|
||||
# Lazy init: many call sites only need resolve_engine_mode(), and legacy mode
|
||||
# should not pay the cost of constructing BillingService.
|
||||
self._new_billing: BillingService | None = None
|
||||
|
||||
def _get_new_billing(self) -> BillingService:
|
||||
if self._new_billing is None:
|
||||
self._new_billing = BillingService(self.db)
|
||||
return self._new_billing
|
||||
|
||||
def get_engine_mode(self, provider: str, model: str) -> EngineMode:
|
||||
return resolve_engine_mode(provider, model)
|
||||
|
||||
def calculate_with_shadow(
|
||||
self,
|
||||
*,
|
||||
provider: str,
|
||||
provider_id: str | None,
|
||||
model: str,
|
||||
task_type: str,
|
||||
api_format: str | None,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
legacy_truth: CostBreakdown,
|
||||
is_failed_request: bool,
|
||||
) -> ShadowBillingResult:
|
||||
"""
|
||||
Compute shadow billing outcome given the legacy truth.
|
||||
|
||||
Notes:
|
||||
- When engine_mode is legacy, we skip new engine calculation.
|
||||
- When engine_mode is shadow, we compute new engine snapshot and compare, but keep truth legacy.
|
||||
- new/new_with_fallback are supported for later phases; callers can choose to honor truth_engine.
|
||||
"""
|
||||
engine_mode = resolve_engine_mode(provider, model)
|
||||
|
||||
# Default response (legacy only)
|
||||
if engine_mode == "legacy":
|
||||
billing_requests_total.labels(engine_mode=engine_mode, truth_engine="legacy").inc()
|
||||
return ShadowBillingResult(
|
||||
truth_breakdown=legacy_truth,
|
||||
shadow_snapshot=None,
|
||||
comparison={"engine_mode": engine_mode},
|
||||
engine_mode=engine_mode,
|
||||
truth_engine="legacy",
|
||||
was_fallback=False,
|
||||
)
|
||||
|
||||
# Build dimensions for new engine
|
||||
request_count = 0 if is_failed_request else 1
|
||||
dimensions: dict[str, Any] = {
|
||||
"input_tokens": int(input_tokens or 0),
|
||||
"output_tokens": int(output_tokens or 0),
|
||||
"cache_creation_input_tokens": int(cache_creation_input_tokens or 0),
|
||||
"cache_read_input_tokens": int(cache_read_input_tokens or 0),
|
||||
"request_count": int(request_count),
|
||||
}
|
||||
if cache_ttl_minutes is not None:
|
||||
dimensions["cache_ttl_minutes"] = int(cache_ttl_minutes)
|
||||
|
||||
# Normalize task_type
|
||||
tt = (task_type or "").lower()
|
||||
if tt not in {"chat", "cli", "video", "image", "audio"}:
|
||||
tt = "chat"
|
||||
|
||||
new_result = self._get_new_billing().calculate(
|
||||
task_type=tt,
|
||||
model=model,
|
||||
provider_id=provider_id or "",
|
||||
dimensions=dimensions,
|
||||
strict_mode=None,
|
||||
)
|
||||
shadow_snapshot = new_result.snapshot
|
||||
|
||||
new_breakdown = CostBreakdown(
|
||||
input_cost=float(shadow_snapshot.cost_breakdown.get("input_cost", 0.0)),
|
||||
output_cost=float(shadow_snapshot.cost_breakdown.get("output_cost", 0.0)),
|
||||
cache_creation_cost=float(
|
||||
shadow_snapshot.cost_breakdown.get("cache_creation_cost", 0.0)
|
||||
),
|
||||
cache_read_cost=float(shadow_snapshot.cost_breakdown.get("cache_read_cost", 0.0)),
|
||||
request_cost=float(shadow_snapshot.cost_breakdown.get("request_cost", 0.0)),
|
||||
total_cost=float(shadow_snapshot.total_cost),
|
||||
)
|
||||
|
||||
diff = abs(float(new_breakdown.total_cost) - float(legacy_truth.total_cost))
|
||||
diff_pct = (
|
||||
(diff / float(legacy_truth.total_cost) * 100.0) if legacy_truth.total_cost > 0 else 0.0
|
||||
)
|
||||
|
||||
comparison = {
|
||||
"engine_mode": engine_mode,
|
||||
"old_total": legacy_truth.total_cost,
|
||||
"new_total": new_breakdown.total_cost,
|
||||
"diff_usd": diff,
|
||||
"diff_pct": diff_pct,
|
||||
"breakdown_diff": {
|
||||
"input_cost": new_breakdown.input_cost - legacy_truth.input_cost,
|
||||
"output_cost": new_breakdown.output_cost - legacy_truth.output_cost,
|
||||
"cache_creation_cost": new_breakdown.cache_creation_cost
|
||||
- legacy_truth.cache_creation_cost,
|
||||
"cache_read_cost": new_breakdown.cache_read_cost - legacy_truth.cache_read_cost,
|
||||
"request_cost": new_breakdown.request_cost - legacy_truth.request_cost,
|
||||
},
|
||||
}
|
||||
|
||||
# Diff logging / metrics
|
||||
threshold = float(getattr(config, "billing_diff_threshold_usd", 0.0001) or 0.0001)
|
||||
if diff > threshold:
|
||||
billing_diff_exceeds_threshold_total.labels(engine_mode=engine_mode).inc()
|
||||
log_level = (
|
||||
(getattr(config, "billing_shadow_log_level", "INFO") or "INFO").strip().lower()
|
||||
)
|
||||
log_fn = getattr(logger, log_level, logger.info)
|
||||
log_fn(
|
||||
"Billing diff detected: provider={}, model={}, old={:.8f}, new={:.8f}, diff={:.8f} ({:.4f}%), mode={}",
|
||||
provider,
|
||||
model,
|
||||
legacy_truth.total_cost,
|
||||
new_breakdown.total_cost,
|
||||
diff,
|
||||
diff_pct,
|
||||
engine_mode,
|
||||
)
|
||||
|
||||
# Invariant monitoring (should be 0)
|
||||
truth_engine: TruthEngine = "legacy"
|
||||
was_fallback = False
|
||||
|
||||
if engine_mode == "shadow":
|
||||
truth_engine = "legacy"
|
||||
truth = legacy_truth
|
||||
elif engine_mode == "new":
|
||||
truth_engine = "new"
|
||||
truth = new_breakdown
|
||||
elif engine_mode == "new_with_fallback":
|
||||
# new is truth unless diff is too large
|
||||
fallback_threshold = threshold * 10.0
|
||||
if diff > fallback_threshold:
|
||||
truth_engine = "legacy"
|
||||
truth = legacy_truth
|
||||
was_fallback = True
|
||||
billing_fallback_total.inc()
|
||||
else:
|
||||
truth_engine = "new"
|
||||
truth = new_breakdown
|
||||
else:
|
||||
# Unknown value -> behave like legacy
|
||||
truth_engine = "legacy"
|
||||
truth = legacy_truth
|
||||
|
||||
billing_requests_total.labels(engine_mode=engine_mode, truth_engine=truth_engine).inc()
|
||||
|
||||
if not truth.validate():
|
||||
billing_invariant_violation_total.labels(
|
||||
engine_mode=engine_mode, truth_engine=truth_engine
|
||||
).inc()
|
||||
logger.warning(
|
||||
"Billing invariant violation: provider={}, model={}, engine_mode={}, truth_engine={}, truth_total={}",
|
||||
provider,
|
||||
model,
|
||||
engine_mode,
|
||||
truth_engine,
|
||||
truth.total_cost,
|
||||
)
|
||||
|
||||
return ShadowBillingResult(
|
||||
truth_breakdown=truth,
|
||||
shadow_snapshot=(
|
||||
shadow_snapshot if engine_mode in {"shadow", "new_with_fallback", "new"} else None
|
||||
),
|
||||
comparison=comparison,
|
||||
engine_mode=engine_mode,
|
||||
truth_engine=truth_engine,
|
||||
was_fallback=was_fallback,
|
||||
)
|
||||
41
src/services/cache/aware_scheduler.py
vendored
41
src/services/cache/aware_scheduler.py
vendored
@@ -721,18 +721,11 @@ class CacheAwareScheduler:
|
||||
return [], global_model_id
|
||||
|
||||
# 2. 构建候选列表(传入 is_stream 和 capability_requirements 用于过滤)
|
||||
from src.config.settings import config
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
# 格式转换总开关(环境变量):关闭时禁止任何跨格式候选进入队列
|
||||
master_conversion_enabled = bool(config.format_conversion_enabled)
|
||||
|
||||
# 全局覆盖开关(数据库):开启时强制允许所有提供商的格式转换(跳过端点格式接受策略)
|
||||
# 格式转换总开关(数据库配置):关闭时禁止任何跨格式候选进入队列
|
||||
global_conversion_enabled = SystemConfigService.is_format_conversion_enabled(db)
|
||||
|
||||
# 如果环境变量明确禁用,则全局覆盖也视为关闭(并最终禁止跨格式转换)
|
||||
if not master_conversion_enabled:
|
||||
global_conversion_enabled = False
|
||||
candidates = await self._build_candidates(
|
||||
db=db,
|
||||
providers=providers,
|
||||
@@ -744,11 +737,10 @@ class CacheAwareScheduler:
|
||||
is_stream=is_stream,
|
||||
capability_requirements=capability_requirements,
|
||||
global_conversion_enabled=global_conversion_enabled,
|
||||
master_conversion_enabled=master_conversion_enabled,
|
||||
)
|
||||
|
||||
# 3. 应用优先级模式排序
|
||||
candidates = self._apply_priority_mode_sort(candidates, affinity_key, target_format)
|
||||
candidates = self._apply_priority_mode_sort(candidates, db, affinity_key, target_format)
|
||||
|
||||
# 更新指标
|
||||
self._metrics["total_candidates"] += len(candidates)
|
||||
@@ -765,6 +757,7 @@ class CacheAwareScheduler:
|
||||
if affinity_key and candidates:
|
||||
candidates = await self._apply_cache_affinity(
|
||||
candidates=candidates,
|
||||
db=db,
|
||||
affinity_key=affinity_key,
|
||||
api_format=target_format,
|
||||
global_model_id=global_model_id,
|
||||
@@ -1068,8 +1061,7 @@ class CacheAwareScheduler:
|
||||
max_candidates: int | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
global_conversion_enabled: bool = False,
|
||||
master_conversion_enabled: bool = True,
|
||||
global_conversion_enabled: bool = True,
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
构建候选列表
|
||||
@@ -1086,8 +1078,7 @@ class CacheAwareScheduler:
|
||||
max_candidates: 最大候选数
|
||||
is_stream: 是否是流式请求,如果为 True 则过滤不支持流式的 Provider
|
||||
capability_requirements: 能力需求(可选)
|
||||
global_conversion_enabled: 全局覆盖开关(DB),开启时跳过端点格式接受策略检查
|
||||
master_conversion_enabled: 总开关(ENV),关闭时禁止任何跨格式转换
|
||||
global_conversion_enabled: 格式转换总开关(数据库配置),关闭时禁止任何跨格式转换
|
||||
|
||||
Returns:
|
||||
候选列表
|
||||
@@ -1179,12 +1170,11 @@ class CacheAwareScheduler:
|
||||
|
||||
# 计算格式转换开关状态(三层优先级)
|
||||
#
|
||||
# 1) 总开关(ENV)关闭 -> 禁止任何跨格式转换
|
||||
# 2) 全局覆盖(DB)开启 -> 强制允许(跳过端点检查)
|
||||
# 1) 全局开关(数据库配置)关闭 -> 禁止任何跨格式转换
|
||||
# 2) 全局开关开启 -> 允许跨格式转换
|
||||
# 3) 提供商覆盖(Provider.enable_format_conversion)开启 -> 强制允许(跳过端点检查)
|
||||
# 4) 否则 -> 由端点配置 format_acceptance_config 决定是否允许
|
||||
provider_allows_conversion = getattr(provider, "enable_format_conversion", True)
|
||||
effective_conversion_enabled = bool(master_conversion_enabled)
|
||||
skip_endpoint_check = global_conversion_enabled or provider_allows_conversion
|
||||
|
||||
is_compatible, needs_conversion, _compat_reason = is_format_compatible(
|
||||
@@ -1192,16 +1182,15 @@ class CacheAwareScheduler:
|
||||
endpoint_format_str,
|
||||
getattr(endpoint, "format_acceptance_config", None),
|
||||
is_stream,
|
||||
effective_conversion_enabled,
|
||||
global_conversion_enabled,
|
||||
skip_endpoint_check=skip_endpoint_check,
|
||||
)
|
||||
logger.debug(
|
||||
"[Scheduler] Format compatibility: client={}, endpoint={}, compatible={}, "
|
||||
"master={}, global={}, provider={}, skip_endpoint={}, reason={}",
|
||||
"global={}, provider={}, skip_endpoint={}, reason={}",
|
||||
client_format_str,
|
||||
endpoint_format_str,
|
||||
is_compatible,
|
||||
master_conversion_enabled,
|
||||
global_conversion_enabled,
|
||||
provider_allows_conversion,
|
||||
skip_endpoint_check,
|
||||
@@ -1302,6 +1291,7 @@ class CacheAwareScheduler:
|
||||
async def _apply_cache_affinity(
|
||||
self,
|
||||
candidates: list[ProviderCandidate],
|
||||
db: Session,
|
||||
affinity_key: str,
|
||||
api_format: str,
|
||||
global_model_id: str,
|
||||
@@ -1334,9 +1324,9 @@ class CacheAwareScheduler:
|
||||
return candidates
|
||||
|
||||
# 判断候选是否应该被降级(用于分组)
|
||||
from src.config.settings import config
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
global_keep_priority = config.keep_priority_on_conversion
|
||||
global_keep_priority = SystemConfigService.is_keep_priority_on_conversion(db)
|
||||
|
||||
def should_demote(c: ProviderCandidate) -> bool:
|
||||
"""判断候选是否应该被降级"""
|
||||
@@ -1467,13 +1457,14 @@ class CacheAwareScheduler:
|
||||
def _apply_priority_mode_sort(
|
||||
self,
|
||||
candidates: list[ProviderCandidate],
|
||||
db: Session,
|
||||
affinity_key: str | None = None,
|
||||
api_format: str | None = None,
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
根据优先级模式对候选列表排序(数字越小越优先)
|
||||
|
||||
排序规则(受 KEEP_PRIORITY_ON_CONVERSION 配置影响):
|
||||
排序规则(受 keep_priority_on_conversion 配置影响):
|
||||
1. 如果全局配置 keep_priority_on_conversion=True,所有候选保持原优先级
|
||||
2. 否则,按 needs_conversion 和 provider.keep_priority_on_conversion 分组:
|
||||
- 保持优先级的候选(exact 或 provider.keep_priority_on_conversion=True)按原优先级排序
|
||||
@@ -1485,10 +1476,10 @@ class CacheAwareScheduler:
|
||||
if not candidates:
|
||||
return candidates
|
||||
|
||||
from src.config.settings import config
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
# 全局配置:如果开启,所有候选保持原优先级
|
||||
global_keep_priority = config.keep_priority_on_conversion
|
||||
global_keep_priority = SystemConfigService.is_keep_priority_on_conversion(db)
|
||||
|
||||
if global_keep_priority:
|
||||
# 全局开启:不分组,直接按优先级模式排序
|
||||
|
||||
@@ -479,13 +479,11 @@ class EndpointHealthService:
|
||||
fam, kind = normalized.split(":", 1)
|
||||
fam_label = {"claude": "Claude", "openai": "OpenAI", "gemini": "Gemini"}.get(fam, fam)
|
||||
kind_label = {
|
||||
"chat": "",
|
||||
"chat": "Chat",
|
||||
"cli": "CLI",
|
||||
"video": "Video",
|
||||
"image": "Image",
|
||||
}.get(kind, kind)
|
||||
if not kind_label:
|
||||
return fam_label
|
||||
return f"{fam_label} {kind_label}"
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -472,4 +472,6 @@ class ModelService:
|
||||
global_model_display_name=(
|
||||
model.global_model.display_name if model.global_model else None
|
||||
),
|
||||
# 有效配置(合并 Model 和 GlobalModel 的 config)
|
||||
effective_config=model.get_effective_config(),
|
||||
)
|
||||
|
||||
@@ -14,9 +14,12 @@ from sqlalchemy.orm import Session
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, SystemConfig
|
||||
|
||||
REQUEST_RECORD_LEVEL_KEY = "request_record_level"
|
||||
_LEGACY_REQUEST_LOG_LEVEL_KEY = "request_log_level"
|
||||
|
||||
class LogLevel(str, Enum):
|
||||
"""日志记录级别"""
|
||||
|
||||
class RequestRecordLevel(str, Enum):
|
||||
"""请求记录级别(控制请求/响应详情入库)"""
|
||||
|
||||
BASIC = "basic" # 仅记录基本信息(tokens、成本等)
|
||||
HEADERS = "headers" # 记录基本信息+请求/响应头(敏感信息会脱敏)
|
||||
@@ -71,9 +74,9 @@ class SystemConfigService:
|
||||
|
||||
# 默认配置
|
||||
DEFAULT_CONFIGS = {
|
||||
"request_log_level": {
|
||||
"value": LogLevel.BASIC.value,
|
||||
"description": "请求记录级别:basic(基本信息), headers(含请求头), full(完整请求响应)",
|
||||
REQUEST_RECORD_LEVEL_KEY: {
|
||||
"value": RequestRecordLevel.BASIC.value,
|
||||
"description": "请求记录级别:basic(基本信息), headers(含请求/响应头), full(完整请求/响应)",
|
||||
},
|
||||
"max_request_body_size": {
|
||||
"value": 5242880, # 5MB
|
||||
@@ -136,10 +139,14 @@ class SystemConfigService:
|
||||
"value": [],
|
||||
"description": "邮箱后缀列表,配合 email_suffix_mode 使用",
|
||||
},
|
||||
# 格式转换开关
|
||||
# 格式转换配置
|
||||
"enable_format_conversion": {
|
||||
"value": True,
|
||||
"description": "格式转换总开关:开启时允许跨格式转换;关闭时禁止任何跨格式转换",
|
||||
},
|
||||
"keep_priority_on_conversion": {
|
||||
"value": False,
|
||||
"description": "全局格式转换开关:开启时强制允许所有提供商的格式转换;关闭时由各提供商自行决定",
|
||||
"description": "格式转换时保持优先级:开启时需要转换的候选保持原优先级;关闭时降级到不需要转换的候选之后",
|
||||
},
|
||||
"audit_log_retention_days": {
|
||||
"value": 30,
|
||||
@@ -183,6 +190,17 @@ class SystemConfigService:
|
||||
@classmethod
|
||||
def get_config(cls, db: Session, key: str, default: Any | None = None) -> Any | None:
|
||||
"""获取系统配置值(带进程内缓存)"""
|
||||
# Backward-compatible alias: request_log_level -> request_record_level
|
||||
if key in {REQUEST_RECORD_LEVEL_KEY, _LEGACY_REQUEST_LOG_LEVEL_KEY}:
|
||||
value = cls._get_request_record_level_raw(db)
|
||||
if value is not None:
|
||||
return value
|
||||
if REQUEST_RECORD_LEVEL_KEY in cls.DEFAULT_CONFIGS:
|
||||
value = cls.DEFAULT_CONFIGS[REQUEST_RECORD_LEVEL_KEY]["value"]
|
||||
_set_cached_config(REQUEST_RECORD_LEVEL_KEY, value)
|
||||
return value
|
||||
return default
|
||||
|
||||
# 1. 检查进程内缓存
|
||||
hit, cached_value = _get_cached_config(key)
|
||||
if hit:
|
||||
@@ -236,6 +254,44 @@ class SystemConfigService:
|
||||
db: Session, key: str, value: Any, description: str | None = None
|
||||
) -> SystemConfig:
|
||||
"""设置系统配置值"""
|
||||
# Backward-compatible alias: request_log_level -> request_record_level
|
||||
if key in {REQUEST_RECORD_LEVEL_KEY, _LEGACY_REQUEST_LOG_LEVEL_KEY}:
|
||||
config = (
|
||||
db.query(SystemConfig).filter(SystemConfig.key == REQUEST_RECORD_LEVEL_KEY).first()
|
||||
)
|
||||
legacy = (
|
||||
db.query(SystemConfig)
|
||||
.filter(SystemConfig.key == _LEGACY_REQUEST_LOG_LEVEL_KEY)
|
||||
.first()
|
||||
)
|
||||
|
||||
if config:
|
||||
config.value = value
|
||||
if description:
|
||||
config.description = description
|
||||
# 如果同时存在旧 key,删除它避免混乱
|
||||
if legacy:
|
||||
db.delete(legacy)
|
||||
elif legacy:
|
||||
# 原地迁移旧 key -> 新 key
|
||||
legacy.key = REQUEST_RECORD_LEVEL_KEY
|
||||
legacy.value = value
|
||||
if description:
|
||||
legacy.description = description
|
||||
config = legacy
|
||||
else:
|
||||
config = SystemConfig(
|
||||
key=REQUEST_RECORD_LEVEL_KEY, value=value, description=description
|
||||
)
|
||||
db.add(config)
|
||||
|
||||
db.commit()
|
||||
db.refresh(config)
|
||||
|
||||
invalidate_config_cache(REQUEST_RECORD_LEVEL_KEY)
|
||||
invalidate_config_cache(_LEGACY_REQUEST_LOG_LEVEL_KEY)
|
||||
return config
|
||||
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == key).first()
|
||||
|
||||
if config:
|
||||
@@ -289,10 +345,20 @@ class SystemConfigService:
|
||||
def get_all_configs(cls, db: Session) -> list:
|
||||
"""获取所有系统配置"""
|
||||
configs = db.query(SystemConfig).all()
|
||||
by_key = {c.key: c for c in configs}
|
||||
result = []
|
||||
for config in configs:
|
||||
# Hide legacy key in list; present as canonical key instead.
|
||||
if config.key == _LEGACY_REQUEST_LOG_LEVEL_KEY:
|
||||
if REQUEST_RECORD_LEVEL_KEY in by_key:
|
||||
continue
|
||||
# Expose as canonical key name
|
||||
config_key = REQUEST_RECORD_LEVEL_KEY
|
||||
else:
|
||||
config_key = config.key
|
||||
|
||||
item = {
|
||||
"key": config.key,
|
||||
"key": config_key,
|
||||
"description": config.description,
|
||||
"updated_at": config.updated_at.isoformat(),
|
||||
}
|
||||
@@ -308,6 +374,24 @@ class SystemConfigService:
|
||||
@classmethod
|
||||
def delete_config(cls, db: Session, key: str) -> bool:
|
||||
"""删除系统配置"""
|
||||
# Backward-compatible alias: request_log_level -> request_record_level
|
||||
if key in {REQUEST_RECORD_LEVEL_KEY, _LEGACY_REQUEST_LOG_LEVEL_KEY}:
|
||||
configs = (
|
||||
db.query(SystemConfig)
|
||||
.filter(
|
||||
SystemConfig.key.in_([REQUEST_RECORD_LEVEL_KEY, _LEGACY_REQUEST_LOG_LEVEL_KEY])
|
||||
)
|
||||
.all()
|
||||
)
|
||||
if not configs:
|
||||
return False
|
||||
for c in configs:
|
||||
db.delete(c)
|
||||
db.commit()
|
||||
invalidate_config_cache(REQUEST_RECORD_LEVEL_KEY)
|
||||
invalidate_config_cache(_LEGACY_REQUEST_LOG_LEVEL_KEY)
|
||||
return True
|
||||
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == key).first()
|
||||
if config:
|
||||
db.delete(config)
|
||||
@@ -333,24 +417,56 @@ class SystemConfigService:
|
||||
logger.info("初始化默认系统配置完成")
|
||||
|
||||
@classmethod
|
||||
def get_log_level(cls, db: Session) -> LogLevel:
|
||||
"""获取日志记录级别"""
|
||||
level = cls.get_config(db, "request_log_level", LogLevel.BASIC.value)
|
||||
def _get_request_record_level_raw(cls, db: Session) -> Any | None:
|
||||
"""Raw value from DB/cache for request record level (supports legacy key)."""
|
||||
hit, cached_value = _get_cached_config(REQUEST_RECORD_LEVEL_KEY)
|
||||
if hit:
|
||||
return cached_value
|
||||
|
||||
config = db.query(SystemConfig).filter(SystemConfig.key == REQUEST_RECORD_LEVEL_KEY).first()
|
||||
if config:
|
||||
_set_cached_config(REQUEST_RECORD_LEVEL_KEY, config.value)
|
||||
return config.value
|
||||
|
||||
hit, cached_value = _get_cached_config(_LEGACY_REQUEST_LOG_LEVEL_KEY)
|
||||
if hit:
|
||||
_set_cached_config(REQUEST_RECORD_LEVEL_KEY, cached_value)
|
||||
return cached_value
|
||||
|
||||
legacy = (
|
||||
db.query(SystemConfig).filter(SystemConfig.key == _LEGACY_REQUEST_LOG_LEVEL_KEY).first()
|
||||
)
|
||||
if legacy:
|
||||
_set_cached_config(_LEGACY_REQUEST_LOG_LEVEL_KEY, legacy.value)
|
||||
_set_cached_config(REQUEST_RECORD_LEVEL_KEY, legacy.value)
|
||||
return legacy.value
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_request_record_level(cls, db: Session) -> RequestRecordLevel:
|
||||
"""获取请求记录级别(控制请求/响应详情入库)"""
|
||||
level = cls.get_config(db, REQUEST_RECORD_LEVEL_KEY, RequestRecordLevel.BASIC.value)
|
||||
if isinstance(level, str):
|
||||
return LogLevel(level)
|
||||
return RequestRecordLevel(level)
|
||||
return level
|
||||
|
||||
@classmethod
|
||||
def get_log_level(cls, db: Session) -> RequestRecordLevel:
|
||||
"""Deprecated: use get_request_record_level."""
|
||||
return cls.get_request_record_level(db)
|
||||
|
||||
@classmethod
|
||||
def should_log_headers(cls, db: Session) -> bool:
|
||||
"""是否应该记录请求头"""
|
||||
log_level = cls.get_log_level(db)
|
||||
return log_level in [LogLevel.HEADERS, LogLevel.FULL]
|
||||
level = cls.get_request_record_level(db)
|
||||
return level in [RequestRecordLevel.HEADERS, RequestRecordLevel.FULL]
|
||||
|
||||
@classmethod
|
||||
def should_log_body(cls, db: Session) -> bool:
|
||||
"""是否应该记录请求体和响应体"""
|
||||
log_level = cls.get_log_level(db)
|
||||
return log_level == LogLevel.FULL
|
||||
level = cls.get_request_record_level(db)
|
||||
return level == RequestRecordLevel.FULL
|
||||
|
||||
@classmethod
|
||||
def should_mask_sensitive_data(cls, db: Session) -> bool:
|
||||
@@ -368,6 +484,11 @@ class SystemConfigService:
|
||||
"""检查全局格式转换是否启用"""
|
||||
return bool(cls.get_config(db, "enable_format_conversion", True))
|
||||
|
||||
@classmethod
|
||||
def is_keep_priority_on_conversion(cls, db: Session) -> bool:
|
||||
"""检查格式转换时是否保持优先级"""
|
||||
return bool(cls.get_config(db, "keep_priority_on_conversion", False))
|
||||
|
||||
@classmethod
|
||||
def mask_sensitive_headers(cls, db: Session, headers: dict[str, Any]) -> dict[str, Any]:
|
||||
"""脱敏敏感请求头"""
|
||||
|
||||
@@ -9,7 +9,7 @@ import os
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy import and_, case, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
@@ -61,12 +61,30 @@ class StatsAggregatorService:
|
||||
"""计算指定业务日期的统计数据(不写入数据库)"""
|
||||
day_start, day_end = _get_business_day_range(date)
|
||||
|
||||
base_query = db.query(Usage).filter(
|
||||
and_(Usage.created_at >= day_start, Usage.created_at < day_end)
|
||||
error_cond = (Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("total_requests"),
|
||||
func.sum(case((error_cond, 1), else_=0)).label("error_requests"),
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
func.sum(Usage.input_cost_usd).label("input_cost"),
|
||||
func.sum(Usage.output_cost_usd).label("output_cost"),
|
||||
func.sum(Usage.cache_creation_cost_usd).label("cache_creation_cost"),
|
||||
func.sum(Usage.cache_read_cost_usd).label("cache_read_cost"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
func.count(func.distinct(Usage.model)).label("unique_models"),
|
||||
func.count(func.distinct(Usage.provider_name)).label("unique_providers"),
|
||||
)
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.first()
|
||||
)
|
||||
|
||||
total_requests = base_query.count()
|
||||
|
||||
total_requests = int(getattr(aggregated, "total_requests", 0) or 0)
|
||||
if total_requests == 0:
|
||||
return {
|
||||
"day_start": day_start,
|
||||
@@ -89,27 +107,7 @@ class StatsAggregatorService:
|
||||
"unique_providers": 0,
|
||||
}
|
||||
|
||||
error_requests = base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
func.sum(Usage.input_cost_usd).label("input_cost"),
|
||||
func.sum(Usage.output_cost_usd).label("output_cost"),
|
||||
func.sum(Usage.cache_creation_cost_usd).label("cache_creation_cost"),
|
||||
func.sum(Usage.cache_read_cost_usd).label("cache_read_cost"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
)
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.first()
|
||||
)
|
||||
error_requests = int(getattr(aggregated, "error_requests", 0) or 0)
|
||||
|
||||
# Fallback 统计 (执行候选数 > 1 的请求数)
|
||||
fallback_subquery = (
|
||||
@@ -135,42 +133,25 @@ class StatsAggregatorService:
|
||||
or 0
|
||||
)
|
||||
|
||||
unique_models = (
|
||||
db.query(func.count(func.distinct(Usage.model)))
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
unique_providers = (
|
||||
db.query(func.count(func.distinct(Usage.provider_name)))
|
||||
.filter(and_(Usage.created_at >= day_start, Usage.created_at < day_end))
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
return {
|
||||
"day_start": day_start,
|
||||
"total_requests": total_requests,
|
||||
"success_requests": total_requests - error_requests,
|
||||
"error_requests": error_requests,
|
||||
"input_tokens": int(aggregated.input_tokens or 0) if aggregated else 0,
|
||||
"output_tokens": int(aggregated.output_tokens or 0) if aggregated else 0,
|
||||
"cache_creation_tokens": (
|
||||
int(aggregated.cache_creation_tokens or 0) if aggregated else 0
|
||||
),
|
||||
"cache_read_tokens": int(aggregated.cache_read_tokens or 0) if aggregated else 0,
|
||||
"total_cost": float(aggregated.total_cost or 0) if aggregated else 0.0,
|
||||
"actual_total_cost": float(aggregated.actual_total_cost or 0) if aggregated else 0.0,
|
||||
"input_cost": float(aggregated.input_cost or 0) if aggregated else 0.0,
|
||||
"output_cost": float(aggregated.output_cost or 0) if aggregated else 0.0,
|
||||
"cache_creation_cost": (
|
||||
float(aggregated.cache_creation_cost or 0) if aggregated else 0.0
|
||||
),
|
||||
"cache_read_cost": float(aggregated.cache_read_cost or 0) if aggregated else 0.0,
|
||||
"avg_response_time_ms": float(aggregated.avg_response_time or 0) if aggregated else 0.0,
|
||||
"input_tokens": int(getattr(aggregated, "input_tokens", 0) or 0),
|
||||
"output_tokens": int(getattr(aggregated, "output_tokens", 0) or 0),
|
||||
"cache_creation_tokens": (int(getattr(aggregated, "cache_creation_tokens", 0) or 0)),
|
||||
"cache_read_tokens": int(getattr(aggregated, "cache_read_tokens", 0) or 0),
|
||||
"total_cost": float(getattr(aggregated, "total_cost", 0) or 0.0),
|
||||
"actual_total_cost": float(getattr(aggregated, "actual_total_cost", 0) or 0.0),
|
||||
"input_cost": float(getattr(aggregated, "input_cost", 0) or 0.0),
|
||||
"output_cost": float(getattr(aggregated, "output_cost", 0) or 0.0),
|
||||
"cache_creation_cost": (float(getattr(aggregated, "cache_creation_cost", 0) or 0.0)),
|
||||
"cache_read_cost": float(getattr(aggregated, "cache_read_cost", 0) or 0.0),
|
||||
"avg_response_time_ms": float(getattr(aggregated, "avg_response_time", 0) or 0.0),
|
||||
"fallback_count": fallback_count,
|
||||
"unique_models": unique_models,
|
||||
"unique_providers": unique_providers,
|
||||
"unique_models": int(getattr(aggregated, "unique_models", 0) or 0),
|
||||
"unique_providers": int(getattr(aggregated, "unique_providers", 0) or 0),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@@ -426,38 +407,11 @@ class StatsAggregatorService:
|
||||
else:
|
||||
stats = StatsUserDaily(id=str(uuid.uuid4()), user_id=user_id, date=day_start)
|
||||
|
||||
# 用户请求统计
|
||||
base_query = db.query(Usage).filter(
|
||||
and_(
|
||||
Usage.user_id == user_id,
|
||||
Usage.created_at >= day_start,
|
||||
Usage.created_at < day_end,
|
||||
)
|
||||
)
|
||||
|
||||
total_requests = base_query.count()
|
||||
|
||||
if total_requests == 0:
|
||||
stats.total_requests = 0
|
||||
stats.success_requests = 0
|
||||
stats.error_requests = 0
|
||||
stats.input_tokens = 0
|
||||
stats.output_tokens = 0
|
||||
stats.cache_creation_tokens = 0
|
||||
stats.cache_read_tokens = 0
|
||||
stats.total_cost = 0.0
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
return stats
|
||||
|
||||
error_requests = base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
|
||||
error_cond = (Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("total_requests"),
|
||||
func.sum(case((error_cond, 1), else_=0)).label("error_requests"),
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
@@ -474,14 +428,32 @@ class StatsAggregatorService:
|
||||
.first()
|
||||
)
|
||||
|
||||
total_requests = int(getattr(aggregated, "total_requests", 0) or 0)
|
||||
if total_requests == 0:
|
||||
stats.total_requests = 0
|
||||
stats.success_requests = 0
|
||||
stats.error_requests = 0
|
||||
stats.input_tokens = 0
|
||||
stats.output_tokens = 0
|
||||
stats.cache_creation_tokens = 0
|
||||
stats.cache_read_tokens = 0
|
||||
stats.total_cost = 0.0
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
db.commit()
|
||||
return stats
|
||||
|
||||
error_requests = int(getattr(aggregated, "error_requests", 0) or 0)
|
||||
|
||||
stats.total_requests = total_requests
|
||||
stats.success_requests = total_requests - error_requests
|
||||
stats.error_requests = error_requests
|
||||
stats.input_tokens = int(aggregated.input_tokens or 0)
|
||||
stats.output_tokens = int(aggregated.output_tokens or 0)
|
||||
stats.cache_creation_tokens = int(aggregated.cache_creation_tokens or 0)
|
||||
stats.cache_read_tokens = int(aggregated.cache_read_tokens or 0)
|
||||
stats.total_cost = float(aggregated.total_cost or 0)
|
||||
stats.input_tokens = int(getattr(aggregated, "input_tokens", 0) or 0)
|
||||
stats.output_tokens = int(getattr(aggregated, "output_tokens", 0) or 0)
|
||||
stats.cache_creation_tokens = int(getattr(aggregated, "cache_creation_tokens", 0) or 0)
|
||||
stats.cache_read_tokens = int(getattr(aggregated, "cache_read_tokens", 0) or 0)
|
||||
stats.total_cost = float(getattr(aggregated, "total_cost", 0) or 0.0)
|
||||
|
||||
if not existing:
|
||||
db.add(stats)
|
||||
@@ -571,10 +543,26 @@ class StatsAggregatorService:
|
||||
# 转换为 UTC 用于查询
|
||||
today_utc = today_local.astimezone(timezone.utc)
|
||||
|
||||
base_query = db.query(Usage).filter(Usage.created_at >= today_utc)
|
||||
|
||||
total_requests = base_query.count()
|
||||
error_cond = (Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("total_requests"),
|
||||
func.sum(case((error_cond, 1), else_=0)).label("error_requests"),
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
func.avg(Usage.response_time_ms).label("avg_response_time"),
|
||||
func.count(func.distinct(Usage.model)).label("unique_models"),
|
||||
func.count(func.distinct(Usage.provider_name)).label("unique_providers"),
|
||||
)
|
||||
.filter(Usage.created_at >= today_utc)
|
||||
.first()
|
||||
)
|
||||
|
||||
total_requests = int(getattr(aggregated, "total_requests", 0) or 0)
|
||||
if total_requests == 0:
|
||||
return {
|
||||
"total_requests": 0,
|
||||
@@ -586,42 +574,33 @@ class StatsAggregatorService:
|
||||
"cache_read_tokens": 0,
|
||||
"total_cost": 0.0,
|
||||
"actual_total_cost": 0.0,
|
||||
"avg_response_time_ms": 0.0,
|
||||
"unique_models": 0,
|
||||
"unique_providers": 0,
|
||||
}
|
||||
|
||||
error_requests = base_query.filter(
|
||||
(Usage.status_code >= 400) | (Usage.error_message.isnot(None))
|
||||
).count()
|
||||
|
||||
aggregated = (
|
||||
db.query(
|
||||
func.sum(Usage.input_tokens).label("input_tokens"),
|
||||
func.sum(Usage.output_tokens).label("output_tokens"),
|
||||
func.sum(Usage.cache_creation_input_tokens).label("cache_creation_tokens"),
|
||||
func.sum(Usage.cache_read_input_tokens).label("cache_read_tokens"),
|
||||
func.sum(Usage.total_cost_usd).label("total_cost"),
|
||||
func.sum(Usage.actual_total_cost_usd).label("actual_total_cost"),
|
||||
)
|
||||
.filter(Usage.created_at >= today_utc)
|
||||
.first()
|
||||
)
|
||||
error_requests = int(getattr(aggregated, "error_requests", 0) or 0)
|
||||
|
||||
return {
|
||||
"total_requests": total_requests,
|
||||
"success_requests": total_requests - error_requests,
|
||||
"error_requests": error_requests,
|
||||
"input_tokens": int(aggregated.input_tokens or 0),
|
||||
"output_tokens": int(aggregated.output_tokens or 0),
|
||||
"cache_creation_tokens": int(aggregated.cache_creation_tokens or 0),
|
||||
"cache_read_tokens": int(aggregated.cache_read_tokens or 0),
|
||||
"total_cost": float(aggregated.total_cost or 0),
|
||||
"actual_total_cost": float(aggregated.actual_total_cost or 0),
|
||||
"input_tokens": int(getattr(aggregated, "input_tokens", 0) or 0),
|
||||
"output_tokens": int(getattr(aggregated, "output_tokens", 0) or 0),
|
||||
"cache_creation_tokens": int(getattr(aggregated, "cache_creation_tokens", 0) or 0),
|
||||
"cache_read_tokens": int(getattr(aggregated, "cache_read_tokens", 0) or 0),
|
||||
"total_cost": float(getattr(aggregated, "total_cost", 0) or 0.0),
|
||||
"actual_total_cost": float(getattr(aggregated, "actual_total_cost", 0) or 0.0),
|
||||
"avg_response_time_ms": float(getattr(aggregated, "avg_response_time", 0) or 0.0),
|
||||
"unique_models": int(getattr(aggregated, "unique_models", 0) or 0),
|
||||
"unique_providers": int(getattr(aggregated, "unique_providers", 0) or 0),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_combined_stats(db: Session) -> dict:
|
||||
def get_combined_stats(db: Session, today_stats: dict | None = None) -> dict:
|
||||
"""获取合并后的统计数据(预聚合 + 今日实时)"""
|
||||
summary = db.query(StatsSummary).first()
|
||||
today_stats = StatsAggregatorService.get_today_realtime_stats(db)
|
||||
today_stats = today_stats or StatsAggregatorService.get_today_realtime_stats(db)
|
||||
|
||||
if not summary:
|
||||
# 如果没有预聚合数据,返回今日数据
|
||||
|
||||
@@ -5,6 +5,8 @@ API密钥统计同步服务
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -39,7 +41,7 @@ class SyncStatsService:
|
||||
else:
|
||||
# 分页处理,避免一次加载所有数据
|
||||
offset = 0
|
||||
api_keys = []
|
||||
api_keys: list[ApiKey] = []
|
||||
while True:
|
||||
batch = db.query(ApiKey).offset(offset).limit(SyncStatsService.BATCH_SIZE).all()
|
||||
if not batch:
|
||||
@@ -47,28 +49,53 @@ class SyncStatsService:
|
||||
api_keys.extend(batch)
|
||||
offset += SyncStatsService.BATCH_SIZE
|
||||
|
||||
# Pre-aggregate Usage stats in ONE query to avoid per-key N+1 scans.
|
||||
# This is critical for large datasets (DB CPU killer otherwise).
|
||||
usage_stats_map: dict[str, dict[str, Any]] = {}
|
||||
if not api_key_id:
|
||||
rows = (
|
||||
db.query(
|
||||
Usage.api_key_id,
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_cost_usd).label("cost"),
|
||||
func.max(Usage.created_at).label("last_used"),
|
||||
)
|
||||
.filter(Usage.api_key_id.isnot(None))
|
||||
.group_by(Usage.api_key_id)
|
||||
.all()
|
||||
)
|
||||
usage_stats_map = {
|
||||
str(r.api_key_id): {
|
||||
"requests": int(r.requests or 0),
|
||||
"cost": float(r.cost or 0),
|
||||
"last_used": r.last_used,
|
||||
}
|
||||
for r in rows
|
||||
if r.api_key_id is not None
|
||||
}
|
||||
|
||||
for api_key in api_keys:
|
||||
try:
|
||||
# 计算实际的使用统计
|
||||
stats = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_cost_usd).label("cost"),
|
||||
if api_key_id:
|
||||
# 单 key 路径:直接查(数据量小)
|
||||
stats = (
|
||||
db.query(
|
||||
func.count(Usage.id).label("requests"),
|
||||
func.sum(Usage.total_cost_usd).label("cost"),
|
||||
func.max(Usage.created_at).label("last_used"),
|
||||
)
|
||||
.filter(Usage.api_key_id == api_key.id)
|
||||
.first()
|
||||
)
|
||||
.filter(Usage.api_key_id == api_key.id)
|
||||
.first()
|
||||
)
|
||||
|
||||
actual_requests = stats.requests or 0
|
||||
actual_cost = float(stats.cost or 0)
|
||||
|
||||
# 获取最后使用时间
|
||||
last_usage = (
|
||||
db.query(Usage.created_at)
|
||||
.filter(Usage.api_key_id == api_key.id)
|
||||
.order_by(Usage.created_at.desc())
|
||||
.first()
|
||||
)
|
||||
actual_requests = int(stats.requests or 0) if stats else 0
|
||||
actual_cost = float(stats.cost or 0) if stats else 0.0
|
||||
last_used_at = stats.last_used if stats else None
|
||||
else:
|
||||
# 批量路径:使用预聚合结果
|
||||
s = usage_stats_map.get(str(api_key.id)) or {}
|
||||
actual_requests = int(s.get("requests") or 0)
|
||||
actual_cost = float(s.get("cost") or 0.0)
|
||||
last_used_at = s.get("last_used")
|
||||
|
||||
# 检查是否需要更新
|
||||
needs_update = False
|
||||
@@ -86,8 +113,8 @@ class SyncStatsService:
|
||||
api_key.total_cost_usd = actual_cost
|
||||
needs_update = True
|
||||
|
||||
if last_usage and api_key.last_used_at != last_usage[0]:
|
||||
api_key.last_used_at = last_usage[0]
|
||||
if last_used_at and api_key.last_used_at != last_used_at:
|
||||
api_key.last_used_at = last_used_at
|
||||
needs_update = True
|
||||
|
||||
result["synced"] += 1
|
||||
|
||||
@@ -255,6 +255,8 @@ class VideoTaskPollerAdapter:
|
||||
task.progress_percent = 100
|
||||
if result.video_urls:
|
||||
task.video_urls = result.video_urls
|
||||
if result.video_duration_seconds is not None:
|
||||
task.video_duration_seconds = result.video_duration_seconds
|
||||
self._attach_poll_raw_response(task, result)
|
||||
elif result.status == VideoStatus.FAILED:
|
||||
task.status = VideoStatus.FAILED.value
|
||||
@@ -376,6 +378,8 @@ class VideoTaskPollerAdapter:
|
||||
task.progress_percent = 100
|
||||
if result.video_urls:
|
||||
task.video_urls = result.video_urls
|
||||
if result.video_duration_seconds is not None:
|
||||
task.video_duration_seconds = result.video_duration_seconds
|
||||
self._attach_poll_raw_response(task, result)
|
||||
elif result.status == VideoStatus.FAILED:
|
||||
task.status = VideoStatus.FAILED.value
|
||||
|
||||
@@ -1041,8 +1041,10 @@ class TaskService:
|
||||
)
|
||||
continue
|
||||
|
||||
# 2. master switch
|
||||
if not config.format_conversion_enabled:
|
||||
# 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(
|
||||
{
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
@@ -756,6 +757,67 @@ class UsageService:
|
||||
cache_ttl_minutes=cache_ttl_minutes,
|
||||
)
|
||||
|
||||
# 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_shadow",
|
||||
"billing_updated_at",
|
||||
"_metadata_truncated",
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _sanitize_request_metadata(cls, 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 cls._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 cls._METADATA_KEEP_KEYS if k in metadata}
|
||||
return reduced
|
||||
|
||||
@classmethod
|
||||
async def _prepare_usage_record(
|
||||
cls,
|
||||
@@ -779,35 +841,169 @@ class UsageService:
|
||||
params.db, params.provider_api_key_id, params.provider_id, params.api_format
|
||||
)
|
||||
|
||||
# 计算成本
|
||||
metadata = dict(params.metadata or {})
|
||||
is_failed_request = params.status_code >= 400 or params.error_message is not None
|
||||
(
|
||||
input_price,
|
||||
output_price,
|
||||
cache_creation_price,
|
||||
cache_read_price,
|
||||
request_price,
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_creation_cost,
|
||||
cache_read_cost,
|
||||
cache_cost,
|
||||
request_cost,
|
||||
total_cost,
|
||||
_tier_index,
|
||||
) = await cls._calculate_costs(
|
||||
db=params.db,
|
||||
provider=params.provider,
|
||||
model=params.model,
|
||||
input_tokens=params.input_tokens,
|
||||
output_tokens=params.output_tokens,
|
||||
cache_creation_input_tokens=params.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=params.cache_read_input_tokens,
|
||||
api_format=params.api_format,
|
||||
cache_ttl_minutes=params.cache_ttl_minutes,
|
||||
use_tiered_pricing=params.use_tiered_pricing,
|
||||
is_failed_request=is_failed_request,
|
||||
)
|
||||
|
||||
# Resolve engine mode early to avoid unnecessary legacy computations.
|
||||
from src.services.billing.shadow import resolve_engine_mode
|
||||
|
||||
engine_mode = resolve_engine_mode(params.provider, params.model)
|
||||
|
||||
# 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"
|
||||
|
||||
# Defaults (filled by either legacy or new path)
|
||||
input_price: float = 0.0
|
||||
output_price: float = 0.0
|
||||
cache_creation_price: float | None = None
|
||||
cache_read_price: float | None = None
|
||||
request_price: float | None = None
|
||||
|
||||
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
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# NEW: new engine as truth (no reconciliation)
|
||||
# ------------------------------------------------------------------
|
||||
if engine_mode == "new":
|
||||
from src.services.billing.service import BillingService
|
||||
|
||||
request_count = 0 if is_failed_request else 1
|
||||
dims: dict[str, Any] = {
|
||||
"input_tokens": params.input_tokens,
|
||||
"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 params.cache_ttl_minutes is not None:
|
||||
dims["cache_ttl_minutes"] = params.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 for new engine (pruned later by _sanitize_request_metadata)
|
||||
metadata["billing_snapshot"] = snap.to_dict()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# LEGACY truth (legacy or shadow or new_with_fallback)
|
||||
# ------------------------------------------------------------------
|
||||
else:
|
||||
(
|
||||
input_price,
|
||||
output_price,
|
||||
cache_creation_price,
|
||||
cache_read_price,
|
||||
request_price,
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_creation_cost,
|
||||
cache_read_cost,
|
||||
cache_cost,
|
||||
request_cost,
|
||||
total_cost,
|
||||
_tier_index,
|
||||
) = await cls._calculate_costs(
|
||||
db=params.db,
|
||||
provider=params.provider,
|
||||
model=params.model,
|
||||
input_tokens=params.input_tokens,
|
||||
output_tokens=params.output_tokens,
|
||||
cache_creation_input_tokens=params.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=params.cache_read_input_tokens,
|
||||
api_format=params.api_format,
|
||||
cache_ttl_minutes=params.cache_ttl_minutes,
|
||||
use_tiered_pricing=params.use_tiered_pricing,
|
||||
is_failed_request=is_failed_request,
|
||||
)
|
||||
|
||||
# Shadow mode: compute new snapshot and store in metadata.billing_shadow only.
|
||||
if engine_mode == "shadow":
|
||||
try:
|
||||
from src.services.billing.shadow import CostBreakdown as ShadowCostBreakdown
|
||||
from src.services.billing.shadow import (
|
||||
ShadowBillingService,
|
||||
)
|
||||
|
||||
legacy_truth = ShadowCostBreakdown(
|
||||
input_cost=input_cost,
|
||||
output_cost=output_cost,
|
||||
cache_creation_cost=cache_creation_cost,
|
||||
cache_read_cost=cache_read_cost,
|
||||
request_cost=request_cost,
|
||||
total_cost=total_cost,
|
||||
)
|
||||
|
||||
shadow = ShadowBillingService(params.db)
|
||||
shadow_result = shadow.calculate_with_shadow(
|
||||
provider=params.provider,
|
||||
provider_id=params.provider_id,
|
||||
model=params.model,
|
||||
task_type=billing_task_type,
|
||||
api_format=params.api_format,
|
||||
input_tokens=params.input_tokens,
|
||||
output_tokens=params.output_tokens,
|
||||
cache_creation_input_tokens=params.cache_creation_input_tokens,
|
||||
cache_read_input_tokens=params.cache_read_input_tokens,
|
||||
cache_ttl_minutes=params.cache_ttl_minutes,
|
||||
legacy_truth=legacy_truth,
|
||||
is_failed_request=is_failed_request,
|
||||
)
|
||||
if shadow_result.shadow_snapshot is not None:
|
||||
metadata["billing_shadow"] = {
|
||||
"engine_mode": shadow_result.engine_mode,
|
||||
"truth_engine": shadow_result.truth_engine,
|
||||
"was_fallback": shadow_result.was_fallback,
|
||||
"comparison": shadow_result.comparison,
|
||||
"snapshot": shadow_result.shadow_snapshot.to_dict(),
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.debug("Shadow billing skipped/failed: {}", str(exc))
|
||||
|
||||
# Best-effort prune metadata to reduce DB/memory pressure.
|
||||
metadata = cls._sanitize_request_metadata(metadata)
|
||||
|
||||
# 构建 Usage 参数
|
||||
usage_params = cls._build_usage_params(
|
||||
@@ -829,7 +1025,7 @@ class UsageService:
|
||||
first_byte_time_ms=params.first_byte_time_ms,
|
||||
status_code=params.status_code,
|
||||
error_message=params.error_message,
|
||||
metadata=params.metadata,
|
||||
metadata=metadata,
|
||||
request_headers=params.request_headers,
|
||||
request_body=params.request_body,
|
||||
provider_request_headers=params.provider_request_headers,
|
||||
@@ -2367,7 +2563,7 @@ class UsageService:
|
||||
metadata["billing_snapshot"] = billing_snapshot
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
usage.request_metadata = metadata
|
||||
usage.request_metadata = cls._sanitize_request_metadata(metadata)
|
||||
|
||||
return True
|
||||
|
||||
@@ -2555,7 +2751,7 @@ class UsageService:
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
metadata["billing_updated_at"] = now.isoformat()
|
||||
usage.request_metadata = metadata
|
||||
usage.request_metadata = cls._sanitize_request_metadata(metadata)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@@ -61,10 +61,32 @@ class QueueTelemetryWriter(TelemetryWriter):
|
||||
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"
|
||||
|
||||
async def record_success(self, **kwargs: Any) -> None:
|
||||
await self._publish_event(UsageEventType.COMPLETED, **kwargs)
|
||||
@@ -101,20 +123,50 @@ class QueueTelemetryWriter(TelemetryWriter):
|
||||
logger.error(f"[usage-queue] XADD failed: {exc}")
|
||||
raise
|
||||
|
||||
def _truncate_body(self, value: Any) -> str | None:
|
||||
"""将 body 序列化为字符串,超长时截断并添加标记"""
|
||||
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
|
||||
try:
|
||||
raw = json.dumps(value, ensure_ascii=False)
|
||||
except TypeError:
|
||||
raw = str(value)
|
||||
max_bytes = config.usage_queue_body_max_bytes
|
||||
if max_bytes > 0 and len(raw) > max_bytes:
|
||||
# 截断并添加标记,预留 15 字符给标记
|
||||
truncate_at = max(0, max_bytes - 15)
|
||||
raw = raw[:truncate_at] + "...[truncated]"
|
||||
return raw
|
||||
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]:
|
||||
# 必需字段
|
||||
@@ -189,24 +241,36 @@ class QueueTelemetryWriter(TelemetryWriter):
|
||||
if kwargs.get("metadata"):
|
||||
data["metadata"] = kwargs["metadata"]
|
||||
|
||||
# 可选:Headers
|
||||
if config.usage_queue_include_headers:
|
||||
# Optional: Headers (masked)
|
||||
if self.include_headers:
|
||||
if kwargs.get("request_headers"):
|
||||
data["request_headers"] = kwargs["request_headers"]
|
||||
data["request_headers"] = self._mask_headers(kwargs["request_headers"])
|
||||
if kwargs.get("provider_request_headers"):
|
||||
data["provider_request_headers"] = kwargs["provider_request_headers"]
|
||||
data["provider_request_headers"] = self._mask_headers(
|
||||
kwargs["provider_request_headers"]
|
||||
)
|
||||
if kwargs.get("response_headers"):
|
||||
data["response_headers"] = kwargs["response_headers"]
|
||||
data["response_headers"] = self._mask_headers(kwargs["response_headers"])
|
||||
if kwargs.get("client_response_headers"):
|
||||
data["client_response_headers"] = kwargs["client_response_headers"]
|
||||
data["client_response_headers"] = self._mask_headers(
|
||||
kwargs["client_response_headers"]
|
||||
)
|
||||
|
||||
# 可选:Bodies
|
||||
if config.usage_queue_include_bodies:
|
||||
request_body = self._truncate_body(kwargs.get("request_body"))
|
||||
response_body = self._truncate_body(kwargs.get("response_body"))
|
||||
if request_body:
|
||||
# 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,
|
||||
)
|
||||
response_body = self._truncate_body(
|
||||
kwargs.get("response_body"),
|
||||
max_size=self._max_response_body_size,
|
||||
is_request=False,
|
||||
)
|
||||
if request_body is not None:
|
||||
data["request_body"] = request_body
|
||||
if response_body:
|
||||
if response_body is not None:
|
||||
data["response_body"] = response_body
|
||||
|
||||
return data
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""缓存装饰器工具"""
|
||||
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
@@ -24,7 +25,22 @@ def _is_api_context(obj: Any) -> bool:
|
||||
return hasattr(obj, "user") and hasattr(obj, "db")
|
||||
|
||||
|
||||
def cache_result(key_prefix: str, ttl: int = 60, user_specific: bool = True) -> Callable:
|
||||
def _hash_vary(vary: dict[str, Any]) -> str:
|
||||
"""Build a short stable hash for cache key variations."""
|
||||
try:
|
||||
raw = json.dumps(vary, sort_keys=True, ensure_ascii=False, default=str)
|
||||
except Exception:
|
||||
raw = str(vary)
|
||||
return hashlib.sha1(raw.encode("utf-8")).hexdigest()[:16]
|
||||
|
||||
|
||||
def cache_result(
|
||||
key_prefix: str,
|
||||
ttl: int = 60,
|
||||
user_specific: bool = True,
|
||||
*,
|
||||
vary_by: list[str] | None = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
缓存函数结果的装饰器
|
||||
|
||||
@@ -68,13 +84,22 @@ def cache_result(key_prefix: str, ttl: int = 60, user_specific: bool = True) ->
|
||||
else:
|
||||
cache_key = f"{key_prefix}:global"
|
||||
|
||||
# 如果有额外的参数(如 days),添加到键中
|
||||
# 从 adapter_self 获取(dataclass 属性)
|
||||
# If there are extra parameters, include them into the key.
|
||||
# - When vary_by is provided: hash the selected attributes to keep key short.
|
||||
# - Otherwise keep backward-compatible "days/limit" suffix behavior.
|
||||
if adapter_self and hasattr(adapter_self, "__dict__"):
|
||||
for attr_name in ["days", "limit"]:
|
||||
if hasattr(adapter_self, attr_name):
|
||||
attr_value = getattr(adapter_self, attr_name)
|
||||
cache_key += f":{attr_name}:{attr_value}"
|
||||
if vary_by:
|
||||
vary: dict[str, Any] = {}
|
||||
for attr_name in vary_by:
|
||||
if hasattr(adapter_self, attr_name):
|
||||
vary[attr_name] = getattr(adapter_self, attr_name)
|
||||
if vary:
|
||||
cache_key += f":v:{_hash_vary(vary)}"
|
||||
else:
|
||||
for attr_name in ["days", "limit"]:
|
||||
if hasattr(adapter_self, attr_name):
|
||||
attr_value = getattr(adapter_self, attr_name)
|
||||
cache_key += f":{attr_name}:{attr_value}"
|
||||
|
||||
# 尝试从缓存获取
|
||||
cached = await redis_client.get(cache_key)
|
||||
|
||||
Reference in New Issue
Block a user