226 lines
9.0 KiB
Python
226 lines
9.0 KiB
Python
"""按版本化模型价格重算历史 usage_events，默认只输出 dry-run 汇总。
|
||
|
||
示例：
|
||
.venv/Scripts/python.exe scripts/reprice_usage.py --from 2026-08-16T16:00:00Z
|
||
.venv/Scripts/python.exe scripts/reprice_usage.py --from 2026-08-16T16:00:00Z \
|
||
--apply --confirm APPLY_DEEPSEEK_REPRICE
|
||
|
||
脚本不加载 .env；调用方必须显式提供 ZCBOT_DB_URL。生产写入前先核对 dry-run 输出和
|
||
打印的脱敏目标。缺少 cache_hit_tokens 的旧辅助调用按全部未命中估算并留标记。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
from collections import defaultdict
|
||
from datetime import datetime, timezone
|
||
from decimal import Decimal
|
||
from pathlib import Path
|
||
import sys
|
||
from typing import Any
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.engine import make_url
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
if str(ROOT) not in sys.path:
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
from core.capabilities import ModelCapabilities # noqa: E402
|
||
from core.pricing import calculate_chat_cost, pricing_snapshot, resolve_chat_price # noqa: E402
|
||
from core.storage import session_scope # noqa: E402
|
||
from core.storage.models import UsageEvent # noqa: E402
|
||
|
||
|
||
DEFAULT_PROFILES = ("deepseek_v4.flash", "deepseek_v4.pro")
|
||
COST_KINDS = {"chat", "prompt_optimize", "context_fold", "task_title", "kb_ingest"}
|
||
CONFIRM_TEXT = "APPLY_DEEPSEEK_REPRICE"
|
||
PRICE_SNAPSHOT_KEYS = {
|
||
"pricing_revision",
|
||
"pricing_source",
|
||
"pricing_source_url",
|
||
"pricing_currency",
|
||
"pricing_at",
|
||
"price_tier",
|
||
"fx_to_cny",
|
||
"input_price_per_mtoken",
|
||
"cache_hit_price_per_mtoken",
|
||
"output_price_per_mtoken",
|
||
"input_cost_cny",
|
||
"cache_hit_cost_cny",
|
||
"output_cost_cny",
|
||
}
|
||
|
||
|
||
def _parse_datetime(value: str) -> datetime:
|
||
dt = datetime.fromisoformat(value.strip().replace("Z", "+00:00"))
|
||
if dt.tzinfo is None:
|
||
dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt.astimezone(timezone.utc)
|
||
|
||
|
||
def _args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(description="重算历史模型成本（默认 dry-run）")
|
||
p.add_argument("--from", dest="from_at", required=True, type=_parse_datetime)
|
||
p.add_argument("--profile", action="append", dest="profiles")
|
||
p.add_argument("--apply", action="store_true", help="事务内写回；默认只预览")
|
||
p.add_argument("--rollback", action="store_true", help="恢复本 revision 保存的旧成本")
|
||
p.add_argument("--confirm", default="", help=f"写入确认串：{CONFIRM_TEXT}")
|
||
return p.parse_args()
|
||
|
||
|
||
def _target() -> str:
|
||
raw = os.environ.get("ZCBOT_DB_URL", "").strip()
|
||
if not raw:
|
||
raise RuntimeError("ZCBOT_DB_URL is not set")
|
||
return make_url(raw).render_as_string(hide_password=True)
|
||
|
||
|
||
def main() -> int:
|
||
args = _args()
|
||
if args.apply and args.rollback:
|
||
raise SystemExit("[ERR] --apply 与 --rollback 不能同时使用")
|
||
if (args.apply or args.rollback) and args.confirm != CONFIRM_TEXT:
|
||
raise SystemExit(f"[ERR] 写入必须传 --confirm {CONFIRM_TEXT}")
|
||
profiles = tuple(args.profiles or DEFAULT_PROFILES)
|
||
caps_by_profile = {
|
||
profile: ModelCapabilities.load(profile, ROOT / "config" / "models")
|
||
for profile in profiles
|
||
}
|
||
revisions = {
|
||
period.get("revision")
|
||
for caps in caps_by_profile.values()
|
||
for period in caps.pricing.get("periods", [])
|
||
if period.get("revision")
|
||
}
|
||
target = _target()
|
||
mode = "ROLLBACK" if args.rollback else ("APPLY" if args.apply else "DRY-RUN")
|
||
print(f"[INFO] mode={mode} target={target} from={args.from_at.isoformat()}")
|
||
print(f"[INFO] profiles={','.join(profiles)} revisions={','.join(sorted(revisions))}")
|
||
|
||
summary: dict[tuple[str, str], dict[str, Any]] = defaultdict(
|
||
lambda: {
|
||
"rows": 0,
|
||
"estimated": 0,
|
||
"old": Decimal("0"),
|
||
"new": Decimal("0"),
|
||
"estimated_old": Decimal("0"),
|
||
"estimated_new": Decimal("0"),
|
||
}
|
||
)
|
||
estimated_by_kind: dict[tuple[str, str], dict[str, Any]] = defaultdict(
|
||
lambda: {"rows": 0, "old": Decimal("0"), "new": Decimal("0")}
|
||
)
|
||
changed = 0
|
||
with session_scope() as s:
|
||
events = s.execute(
|
||
select(UsageEvent).where(
|
||
UsageEvent.model_profile.in_(profiles),
|
||
UsageEvent.kind.in_(COST_KINDS),
|
||
UsageEvent.created_at >= args.from_at,
|
||
).order_by(UsageEvent.created_at, UsageEvent.event_id)
|
||
).scalars().all()
|
||
for event in events:
|
||
units = dict(event.units or {})
|
||
caps = caps_by_profile[event.model_profile]
|
||
if args.rollback:
|
||
if units.get("repricing_revision") not in revisions:
|
||
continue
|
||
previous = units.get("previous_cost_cny")
|
||
if previous is None:
|
||
continue
|
||
old = Decimal(str(event.cost_cny or 0))
|
||
new = Decimal(str(previous))
|
||
previous_pricing = units.get("previous_pricing_units") or {}
|
||
for price_key in PRICE_SNAPSHOT_KEYS:
|
||
units.pop(price_key, None)
|
||
if isinstance(previous_pricing, dict):
|
||
units.update(previous_pricing)
|
||
for key in list(units):
|
||
if (
|
||
key.startswith("repricing_")
|
||
or key in {"previous_cost_cny", "previous_pricing_units"}
|
||
):
|
||
units.pop(key, None)
|
||
key = (event.model_profile, "rollback")
|
||
else:
|
||
quote = resolve_chat_price(caps.pricing, occurred_at=event.created_at)
|
||
if quote is None or units.get("repricing_revision") == quote.revision:
|
||
continue
|
||
tokens_in = int(units.get("tokens_in") or 0)
|
||
tokens_out = int(units.get("tokens_out") or 0)
|
||
# 旧主循环一直提取 DeepSeek cache usage，但 record_chat_usage 的
|
||
# extra_units 会省略值为 0 的字段，所以 chat 缺 key 可判定为零命中；
|
||
# 旧辅助调用没有提取缓存明细，缺 key 才是真正的不确定。
|
||
has_cache = "cache_hit_tokens" in units or event.kind == "chat"
|
||
cache_hit = int(units.get("cache_hit_tokens") or 0)
|
||
breakdown = calculate_chat_cost(
|
||
quote,
|
||
prompt_tokens=tokens_in,
|
||
completion_tokens=tokens_out,
|
||
cache_hit_tokens=cache_hit,
|
||
)
|
||
old = Decimal(str(event.cost_cny or 0))
|
||
new = breakdown.total_cny
|
||
previous_pricing = {
|
||
key: units[key] for key in PRICE_SNAPSHOT_KEYS if key in units
|
||
}
|
||
units.update(pricing_snapshot(quote, breakdown))
|
||
units.update({
|
||
"repricing_revision": quote.revision,
|
||
"previous_cost_cny": float(old),
|
||
"repriced_at": datetime.now(timezone.utc).isoformat(),
|
||
"repricing_time_basis": "usage_event.created_at",
|
||
"repricing_estimated": not has_cache,
|
||
"previous_pricing_units": previous_pricing,
|
||
})
|
||
key = (event.model_profile, quote.tier)
|
||
if not has_cache:
|
||
summary[key]["estimated"] += 1
|
||
summary[key]["estimated_old"] += old
|
||
summary[key]["estimated_new"] += new
|
||
uncertain = estimated_by_kind[(event.model_profile, event.kind)]
|
||
uncertain["rows"] += 1
|
||
uncertain["old"] += old
|
||
uncertain["new"] += new
|
||
|
||
bucket = summary[key]
|
||
bucket["rows"] += 1
|
||
bucket["old"] += old
|
||
bucket["new"] += new
|
||
changed += 1
|
||
if args.apply or args.rollback:
|
||
event.cost_cny = new.quantize(Decimal("0.000001"))
|
||
event.units = units
|
||
|
||
if not (args.apply or args.rollback):
|
||
s.rollback()
|
||
|
||
total_old = Decimal("0")
|
||
total_new = Decimal("0")
|
||
for (profile, tier), item in sorted(summary.items()):
|
||
total_old += item["old"]
|
||
total_new += item["new"]
|
||
print(
|
||
f"[INFO] {profile} tier={tier} rows={item['rows']} "
|
||
f"estimated={item['estimated']} old={item['old']:.6f} "
|
||
f"new={item['new']:.6f} delta={(item['new'] - item['old']):.6f} "
|
||
f"estimated_delta={(item['estimated_new'] - item['estimated_old']):.6f}"
|
||
)
|
||
for (profile, kind), item in sorted(estimated_by_kind.items()):
|
||
print(
|
||
f"[WARN] estimated profile={profile} kind={kind} rows={item['rows']} "
|
||
f"old={item['old']:.6f} new={item['new']:.6f} "
|
||
f"delta={(item['new'] - item['old']):.6f}"
|
||
)
|
||
print(
|
||
f"[OK] rows={changed} old={total_old:.6f} new={total_new:.6f} "
|
||
f"delta={(total_new - total_old):.6f} mode={mode}"
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|