tzt
aa7cb0704c
- billing.BillingMixin 新增 budget_mode(student, est, ts)：整数基点判定 normal(<80%) / optimize(>=80%) / cheap(>=95%) / block(>100%)； cap<=0 不限额恒 normal；学生缺失 normal（扣费与拒绝权威仍在 try_hold） - routes._run_chat：optimize/cheap 档自动降档池条目（_downgrade_entry， TIERS 排名制降档、杜绝反向升档、跳过停用与 mock），重算 model/est； 非正常档位经 X-Budget-Mode 响应头如实上报（缓存命中路径同样携带） - 黄金用例锁死边界：7999/8000/9499/9500/10000/10001bp 九组参数化断言， 全整数运算无浮点漂移；跨日重置与不限额口径与 try_hold 一致 pytest 443 passed（基线 425 + 18）
197 lines
8.9 KiB
Python
197 lines
8.9 KiB
Python
"""计费两阶段（D-P11）：预扣-结算-退款，作为 Ledger 的 mixin。
|
||
|
||
拆分为独立模块的工程原因：ledger.py 承载 DDL/CRUD/查询；金额敏感的
|
||
两阶段语义集中在此，便于评审与测试聚焦。
|
||
|
||
并发正确性：全局锁内 check-then-act（单写者模型）+ 余额原子防线
|
||
（语句条件 balance_milli >= ?，受影响行数为 0 即不足）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import threading
|
||
import time
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from gateway.proxy.ledgerutil import _today
|
||
|
||
# 预扣：余额原子防线；日上限判定在锁内 Python 侧（单写者模型下等价安全）
|
||
# 占位顺序：est, spent_total, today, sid, est
|
||
_SQL_HOLD = """UPDATE students SET
|
||
balance_milli = balance_milli - ?,
|
||
spent_today_milli = ?, spent_date = ?
|
||
WHERE id = ? AND balance_milli >= ?"""
|
||
|
||
# 预算四档阈值（T-X1，基点 bp：1% = 100bp；整数比较，杜绝浮点边界漂移）
|
||
BUDGET_OPTIMIZE_BP = 8000 # >= 80% 建议降一档
|
||
BUDGET_CHEAP_BP = 9500 # >= 95% 强制最低档
|
||
|
||
_SQL_HOLD_MARK = """INSERT INTO usage_ledger
|
||
(request_id, ts, key_id, model, bucket, charged_milli, status)
|
||
VALUES (?, ?, ?, ?, ?, ?, 'holding')"""
|
||
|
||
_SQL_INSUFFICIENT = """INSERT OR REPLACE INTO usage_ledger
|
||
(request_id, ts, key_id, model, bucket, charged_milli, status)
|
||
VALUES (?, ?, ?, ?, ?, 0, 'insufficient')"""
|
||
|
||
_SQL_SETTLE_REFUND = """UPDATE students SET
|
||
balance_milli = balance_milli + ?,
|
||
spent_today_milli = MAX(0, spent_today_milli - ?)
|
||
WHERE id = ?"""
|
||
|
||
_SQL_SETTLE_USAGE = """UPDATE usage_ledger SET
|
||
in_miss_tok = ?, in_hit_tok = ?, out_tok = ?, gateway_cached = ?,
|
||
upstream_cost_milli = ?, charged_milli = ?, margin_milli = ? - ?,
|
||
ttfb_ms = ?, total_ms = ?, status = ?
|
||
WHERE request_id = ?"""
|
||
|
||
|
||
class BillingMixin:
|
||
"""try_hold / settle / void（由 Ledger 继承；依赖 _lock/_connect）。"""
|
||
|
||
_lock: threading.Lock
|
||
|
||
def _connect(self): # 由 Ledger 提供
|
||
raise NotImplementedError
|
||
|
||
def try_hold(self, request_id: str, key_id: int, student_id: int,
|
||
model: str, bucket: str, est_milli: int, ts: float) -> bool:
|
||
"""预扣：锁内读现状 -> 日上限判定 -> 余额原子扣减 -> 写 holding 流水。
|
||
|
||
受影响行数为 0 即不足（余额不够）-> False（402）；
|
||
request_id 幂等：重复 hold 直接返回 True（不重复扣）。
|
||
"""
|
||
today = _today(ts)
|
||
with self._lock, self._connect() as conn:
|
||
dup = conn.execute(
|
||
"SELECT 1 FROM usage_ledger WHERE request_id = ?",
|
||
(request_id,)).fetchone()
|
||
if dup:
|
||
return True
|
||
stu = conn.execute(
|
||
"SELECT balance_milli, daily_cap_milli, spent_today_milli, spent_date"
|
||
" FROM students WHERE id = ?", (student_id,)).fetchone()
|
||
if stu is None:
|
||
return False
|
||
spent_base = stu["spent_today_milli"] if stu["spent_date"] == today else 0
|
||
if stu["daily_cap_milli"] > 0 and spent_base + est_milli > stu["daily_cap_milli"]:
|
||
conn.execute(_SQL_INSUFFICIENT,
|
||
(request_id, int(ts), key_id, model, bucket))
|
||
return False
|
||
cur = conn.execute(_SQL_HOLD, (
|
||
est_milli, spent_base + est_milli, today, student_id, est_milli))
|
||
if cur.rowcount == 0:
|
||
conn.execute(_SQL_INSUFFICIENT,
|
||
(request_id, int(ts), key_id, model, bucket))
|
||
return False
|
||
conn.execute(_SQL_HOLD_MARK,
|
||
(request_id, int(ts), key_id, model, bucket, est_milli))
|
||
return True
|
||
|
||
def settle(self, request_id: str, actual_milli: int, *,
|
||
in_miss_tok: int = 0, in_hit_tok: int = 0, out_tok: int = 0,
|
||
gateway_cached: int = 0, upstream_cost_milli: int = 0,
|
||
ttfb_ms: Optional[int] = None, total_ms: Optional[int] = None,
|
||
status: str = "ok") -> bool:
|
||
"""结算：按真实值更新流水并回补（预扣额 − 实际额）差额。"""
|
||
with self._lock, self._connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT key_id, charged_milli FROM usage_ledger WHERE request_id = ?",
|
||
(request_id,)).fetchone()
|
||
if row is None:
|
||
return False
|
||
key_id = row["key_id"]
|
||
est = row["charged_milli"]
|
||
student = conn.execute(
|
||
"SELECT student_id FROM proxy_keys WHERE id = ?", (key_id,)).fetchone()
|
||
if student is None:
|
||
return False
|
||
refund = est - actual_milli
|
||
if refund != 0:
|
||
conn.execute(_SQL_SETTLE_REFUND,
|
||
(refund, refund, student["student_id"]))
|
||
conn.execute(_SQL_SETTLE_USAGE, (
|
||
in_miss_tok, in_hit_tok, out_tok, gateway_cached,
|
||
upstream_cost_milli, actual_milli,
|
||
actual_milli, upstream_cost_milli, ttfb_ms, total_ms,
|
||
status, request_id))
|
||
return True
|
||
|
||
def void(self, request_id: str, status: str = "aborted") -> bool:
|
||
"""全额退款（上游失败）；流水保留审计。"""
|
||
with self._lock, self._connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT key_id, charged_milli, status FROM usage_ledger"
|
||
" WHERE request_id = ?", (request_id,)).fetchone()
|
||
if row is None or row["status"] not in ("holding", "ok"):
|
||
return False
|
||
est = row["charged_milli"]
|
||
student = conn.execute(
|
||
"SELECT student_id FROM proxy_keys WHERE id = ?",
|
||
(row["key_id"],)).fetchone()
|
||
if student is not None and est > 0:
|
||
conn.execute(_SQL_SETTLE_REFUND,
|
||
(est, est, student["student_id"]))
|
||
conn.execute(
|
||
"UPDATE usage_ledger SET charged_milli = 0, status = ?"
|
||
" WHERE request_id = ?", (status, request_id))
|
||
return True
|
||
|
||
# ---------- 预算四档（T-X1） ----------
|
||
def budget_mode(self, student_id: int, est_milli: int, ts: float) -> str:
|
||
"""预算档位判定：normal / optimize / cheap / block。
|
||
|
||
以「当日已 spent + 本次预估」占日上限的比例判定（基点 bp，全整数运算，
|
||
无浮点边界漂移；黄金用例锁死）：
|
||
< 8000bp（80%） -> normal（现行行为）
|
||
>= 8000bp（80%） -> optimize（调用方应降一档模型）
|
||
>= 9500bp（95%） -> cheap（调用方应强制最低档）
|
||
> 10000bp（100%） -> block（try_hold 的硬拒绝语义兜底）
|
||
daily_cap_milli <= 0 视为不限额 -> 恒 normal；学生不存在 -> normal
|
||
（try_hold 才是扣费与拒绝的唯一权威，本方法只做档位建议）。
|
||
"""
|
||
today = _today(ts)
|
||
with self._lock, self._connect() as conn:
|
||
stu = conn.execute(
|
||
"SELECT daily_cap_milli, spent_today_milli, spent_date"
|
||
" FROM students WHERE id = ?", (student_id,)).fetchone()
|
||
if stu is None:
|
||
return "normal"
|
||
cap = int(stu["daily_cap_milli"] or 0)
|
||
if cap <= 0:
|
||
return "normal"
|
||
spent = int(stu["spent_today_milli"] or 0) \
|
||
if stu["spent_date"] == today else 0
|
||
projected = spent + max(0, int(est_milli))
|
||
projected_bp = projected * 10000 // cap
|
||
if projected_bp > 10000:
|
||
return "block"
|
||
if projected_bp >= BUDGET_CHEAP_BP:
|
||
return "cheap"
|
||
if projected_bp >= BUDGET_OPTIMIZE_BP:
|
||
return "optimize"
|
||
return "normal"
|
||
|
||
# ---------- 流水查询 ----------
|
||
def get_usage(self, request_id: str) -> Optional[Dict[str, Any]]:
|
||
with self._lock, self._connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM usage_ledger WHERE request_id = ?",
|
||
(request_id,)).fetchone()
|
||
return dict(row) if row else None
|
||
|
||
def list_usage(self, student_id: Optional[int] = None,
|
||
limit: int = 50, offset: int = 0) -> List[Dict[str, Any]]:
|
||
"""流水分页（可按学生过滤，经其名下 key）；按学生过滤走联表常量语句。"""
|
||
if student_id is None:
|
||
with self._lock, self._connect() as conn:
|
||
rows = conn.execute(
|
||
"SELECT * FROM usage_ledger ORDER BY ts DESC LIMIT ? OFFSET ?",
|
||
(limit, offset)).fetchall()
|
||
return [dict(r) for r in rows]
|
||
with self._lock, self._connect() as conn:
|
||
rows = conn.execute(
|
||
"SELECT u.* FROM usage_ledger u JOIN proxy_keys k ON k.id = u.key_id"
|
||
" WHERE k.student_id = ? ORDER BY u.ts DESC LIMIT ? OFFSET ?",
|
||
(student_id, limit, offset)).fetchall()
|
||
return [dict(r) for r in rows]
|