210 lines
6.9 KiB
Python
210 lines
6.9 KiB
Python
"""版本化 chat 价格解析与成本计算。
|
||
|
||
模型 API 通常只返回 token usage，不返回实际金额。这里以模型 YAML 中经过核对的
|
||
provider 价格为事实源，并把命中的价格版本、时段和汇率快照交给 usage_events 留痕。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, time, timezone
|
||
from decimal import Decimal
|
||
from typing import Any, Mapping, Optional
|
||
|
||
|
||
_MTOKEN = Decimal("1000000")
|
||
_WEEKDAYS = {
|
||
"mon": 0,
|
||
"tue": 1,
|
||
"wed": 2,
|
||
"thu": 3,
|
||
"fri": 4,
|
||
"sat": 5,
|
||
"sun": 6,
|
||
}
|
||
|
||
|
||
def _decimal(value: Any, default: str = "0") -> Decimal:
|
||
if value is None or value == "":
|
||
return Decimal(default)
|
||
return Decimal(str(value))
|
||
|
||
|
||
def _parse_datetime(value: Any) -> Optional[datetime]:
|
||
if not value:
|
||
return None
|
||
if isinstance(value, datetime):
|
||
dt = value
|
||
else:
|
||
raw = str(value).strip().replace("Z", "+00:00")
|
||
dt = datetime.fromisoformat(raw)
|
||
if dt.tzinfo is None:
|
||
dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt.astimezone(timezone.utc)
|
||
|
||
|
||
def _parse_time(value: str) -> time:
|
||
hour, minute = (int(part) for part in value.split(":", 1))
|
||
return time(hour=hour, minute=minute)
|
||
|
||
|
||
def _in_window(now: time, start: time, end: time) -> bool:
|
||
if start <= end:
|
||
return start <= now < end
|
||
return now >= start or now < end
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PriceQuote:
|
||
revision: str
|
||
source_url: str
|
||
currency: str
|
||
fx_to_cny: Decimal
|
||
tier: str
|
||
input_per_mtoken: Decimal
|
||
output_per_mtoken: Decimal
|
||
cache_hit_per_mtoken: Decimal
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CostBreakdown:
|
||
total_cny: Decimal
|
||
input_cny: Decimal
|
||
cache_hit_cny: Decimal
|
||
output_cny: Decimal
|
||
cache_hit_tokens: int
|
||
cache_miss_tokens: int
|
||
|
||
|
||
def _periods(pricing: Mapping[str, Any]) -> list[Mapping[str, Any]]:
|
||
periods = pricing.get("periods")
|
||
if isinstance(periods, list):
|
||
return [p for p in periods if isinstance(p, Mapping)]
|
||
# 兼容一个价格版本直接写在 pricing 顶层的简写。
|
||
return [pricing]
|
||
|
||
|
||
def _select_period(
|
||
pricing: Mapping[str, Any], occurred_at: datetime
|
||
) -> Optional[Mapping[str, Any]]:
|
||
candidates: list[tuple[datetime, Mapping[str, Any]]] = []
|
||
for period in _periods(pricing):
|
||
start = _parse_datetime(period.get("effective_from")) or datetime.min.replace(
|
||
tzinfo=timezone.utc
|
||
)
|
||
end = _parse_datetime(period.get("effective_to"))
|
||
if start <= occurred_at and (end is None or occurred_at < end):
|
||
candidates.append((start, period))
|
||
if not candidates:
|
||
return None
|
||
return max(candidates, key=lambda item: item[0])[1]
|
||
|
||
|
||
def _tier_matches(tier: Mapping[str, Any], occurred_at: datetime) -> bool:
|
||
tz_name = str(tier.get("timezone") or "UTC").upper()
|
||
if tz_name != "UTC":
|
||
raise ValueError(f"pricing time tier 暂只支持 UTC，收到 {tz_name!r}")
|
||
now = occurred_at.astimezone(timezone.utc)
|
||
weekdays = tier.get("weekdays") or []
|
||
if weekdays:
|
||
allowed = {_WEEKDAYS[str(day).lower()] for day in weekdays}
|
||
if now.weekday() not in allowed:
|
||
return False
|
||
windows = tier.get("windows") or []
|
||
if not windows:
|
||
return True
|
||
current = now.time().replace(tzinfo=None)
|
||
return any(
|
||
_in_window(current, _parse_time(window[0]), _parse_time(window[1]))
|
||
for window in windows
|
||
if isinstance(window, (list, tuple)) and len(window) == 2
|
||
)
|
||
|
||
|
||
def resolve_chat_price(
|
||
pricing: Mapping[str, Any] | None,
|
||
*,
|
||
occurred_at: Optional[datetime] = None,
|
||
) -> Optional[PriceQuote]:
|
||
"""解析调用时刻适用的价格；无配置或尚未生效时返回 ``None``。"""
|
||
if not pricing:
|
||
return None
|
||
when = occurred_at or datetime.now(timezone.utc)
|
||
if when.tzinfo is None:
|
||
when = when.replace(tzinfo=timezone.utc)
|
||
when = when.astimezone(timezone.utc)
|
||
period = _select_period(pricing, when)
|
||
if period is None:
|
||
return None
|
||
|
||
rates = period.get("default") or {}
|
||
tier_name = "default"
|
||
for tier in period.get("time_tiers") or []:
|
||
if isinstance(tier, Mapping) and _tier_matches(tier, when):
|
||
rates = tier
|
||
tier_name = str(tier.get("name") or "time_tier")
|
||
break
|
||
|
||
currency = str(period.get("currency") or "CNY").upper()
|
||
if currency not in {"CNY", "USD"}:
|
||
raise ValueError(f"不支持的 pricing currency: {currency!r}")
|
||
fx = _decimal(period.get("fx_to_cny"), "1") if currency == "USD" else Decimal("1")
|
||
return PriceQuote(
|
||
revision=str(period.get("revision") or "unversioned"),
|
||
source_url=str(period.get("source_url") or ""),
|
||
currency=currency,
|
||
fx_to_cny=fx,
|
||
tier=tier_name,
|
||
input_per_mtoken=_decimal(rates.get("input_per_mtoken")),
|
||
output_per_mtoken=_decimal(rates.get("output_per_mtoken")),
|
||
cache_hit_per_mtoken=_decimal(
|
||
rates.get("cache_hit_per_mtoken"),
|
||
str(rates.get("input_per_mtoken") or 0),
|
||
),
|
||
)
|
||
|
||
|
||
def calculate_chat_cost(
|
||
quote: PriceQuote,
|
||
*,
|
||
prompt_tokens: int,
|
||
completion_tokens: int,
|
||
cache_hit_tokens: int = 0,
|
||
) -> CostBreakdown:
|
||
"""按价格快照拆分缓存命中、未命中和输出成本，返回人民币金额。"""
|
||
tokens_in = max(0, int(prompt_tokens))
|
||
hit = max(0, min(int(cache_hit_tokens), tokens_in))
|
||
miss = tokens_in - hit
|
||
tokens_out = max(0, int(completion_tokens))
|
||
input_cost = Decimal(miss) * quote.input_per_mtoken * quote.fx_to_cny / _MTOKEN
|
||
hit_cost = Decimal(hit) * quote.cache_hit_per_mtoken * quote.fx_to_cny / _MTOKEN
|
||
output_cost = Decimal(tokens_out) * quote.output_per_mtoken * quote.fx_to_cny / _MTOKEN
|
||
quant = Decimal("0.000001")
|
||
return CostBreakdown(
|
||
total_cny=(input_cost + hit_cost + output_cost).quantize(quant),
|
||
input_cny=input_cost.quantize(quant),
|
||
cache_hit_cny=hit_cost.quantize(quant),
|
||
output_cny=output_cost.quantize(quant),
|
||
cache_hit_tokens=hit,
|
||
cache_miss_tokens=miss,
|
||
)
|
||
|
||
|
||
def pricing_snapshot(quote: PriceQuote, breakdown: CostBreakdown) -> dict[str, Any]:
|
||
"""生成可直接合并进 usage_events.units 的价格快照。"""
|
||
return {
|
||
"pricing_revision": quote.revision,
|
||
"pricing_source": "local_catalog",
|
||
"pricing_source_url": quote.source_url,
|
||
"pricing_currency": quote.currency,
|
||
"price_tier": quote.tier,
|
||
"fx_to_cny": float(quote.fx_to_cny),
|
||
"input_price_per_mtoken": float(quote.input_per_mtoken),
|
||
"cache_hit_price_per_mtoken": float(quote.cache_hit_per_mtoken),
|
||
"output_price_per_mtoken": float(quote.output_per_mtoken),
|
||
"cache_hit_tokens": breakdown.cache_hit_tokens,
|
||
"cache_miss_tokens": breakdown.cache_miss_tokens,
|
||
"input_cost_cny": float(breakdown.input_cny),
|
||
"cache_hit_cost_cny": float(breakdown.cache_hit_cny),
|
||
"output_cost_cny": float(breakdown.output_cny),
|
||
}
|