mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat: 视频计费增强与影子计费系统
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user