mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 01:11:59 +00:00
核查花费账：engine/spend.py 照结论句裁判的每日计数做一本核查用的每日账（设置表 verify_spend 一行，按本地日期滚动）。调用前按估算预留，加上这笔会超过每日上限就不记、返回 None，上限为 None 表示不限；调用后按实际用量找补，跨了天不找补；价格目录里没有的模型只计次数。和裁判分开记：核查一次要查好几轮库、花得多，记在同一本账上会把当天的裁判挤到判不了、让报告出具无故降档。读改写按事件循环各持一把锁，并发预留不会一起越过上限
This commit is contained in:
ThearchyHelios
@@ -0,0 +1,98 @@
|
||||
"""核查花费账：双盲复算、反方审稿、审批质询里的模型调用，按本地日期记一本每日账。
|
||||
|
||||
照结论句裁判的每日计数做（judge.py 的 SPEND_KEY、_charge、_settle）：调用前按目录价估算、先预留，加上这笔会超过
|
||||
每日上限就不记；调用后按实际用量找补差额。账落在设置表里，按本地日期滚动。
|
||||
|
||||
**为什么和裁判分开记。** 核查是一个 agent 查库、来回好几轮，一次花的钱比裁判判一份报告多得多。记在同一本账上，
|
||||
核查一多，当天的裁判就全因为「每日上限」判不了——报告的出具会因此降档，而降档的原因和这份报告毫无关系。两本账
|
||||
各管各的上限，互不挤占。
|
||||
|
||||
目录里没有价格的模型估不出金额：只计次数（unpriced_calls），金额上限对它不起作用，由调用方照实说明。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import weakref
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Setting
|
||||
|
||||
#: 每日计数存在设置表的这一行：{date: 本地日期, usd, calls, unpriced_calls}。不从设置接口进出（api/settings.py）
|
||||
SPEND_KEY = "verify_spend"
|
||||
|
||||
_locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]" = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
def _lock() -> asyncio.Lock:
|
||||
"""读改写这一行时持的锁，按事件循环各一把（asyncio.Lock 不能跨循环用）。
|
||||
|
||||
几个核查并行时各自预留：不串行的话两边读到同一个余额，会一起越过上限。
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
lock = _locks.get(loop)
|
||||
if lock is None:
|
||||
lock = _locks[loop] = asyncio.Lock()
|
||||
return lock
|
||||
|
||||
|
||||
def _today() -> str:
|
||||
"""本地日期：每日上限按它滚动。"""
|
||||
return datetime.now().astimezone().date().isoformat()
|
||||
|
||||
|
||||
def _spend_of(row: Any) -> dict[str, Any]:
|
||||
value = row.value if row is not None and isinstance(row.value, dict) else {}
|
||||
today = _today()
|
||||
if value.get("date") != today:
|
||||
return {"date": today, "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
return {"date": today, "usd": float(value.get("usd") or 0.0), "calls": int(value.get("calls") or 0),
|
||||
"unpriced_calls": int(value.get("unpriced_calls") or 0)}
|
||||
|
||||
|
||||
async def daily_spend() -> dict[str, Any]:
|
||||
"""今天（本地日期）核查已经花了多少：{date, usd, calls, unpriced_calls}。估不出金额的调用只计次数。"""
|
||||
async with SessionLocal() as session:
|
||||
return _spend_of(await session.get(Setting, SPEND_KEY))
|
||||
|
||||
|
||||
async def charge(amount: float, *, limit: float | None, priced: bool) -> str | None:
|
||||
"""记一次调用、先按估算预留金额。加上这笔会超过每日上限就不记，返回 None；记下了返回记在哪天（找补时要用）。
|
||||
|
||||
limit 为 None 表示不限。估不出金额的调用（priced=False）照记次数，金额记 0，调用方应传 limit=None：
|
||||
拿 0 去比上限，等于说它不花钱。
|
||||
"""
|
||||
async with _lock():
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, SPEND_KEY)
|
||||
spend = _spend_of(row)
|
||||
if limit is not None and spend["usd"] + amount > limit + 1e-12:
|
||||
return None
|
||||
spend["usd"] = round(spend["usd"] + amount, 8)
|
||||
spend["calls"] += 1
|
||||
spend["unpriced_calls"] += 0 if priced else 1
|
||||
if row is None:
|
||||
session.add(Setting(key=SPEND_KEY, value=spend))
|
||||
else:
|
||||
row.value = spend
|
||||
await session.commit()
|
||||
return spend["date"]
|
||||
|
||||
|
||||
async def settle(date: str, delta: float) -> None:
|
||||
"""调用结束后按实际用量找补预留的差额。跨了天就不找补：预留记在前一天，今天从零算。"""
|
||||
if not delta:
|
||||
return
|
||||
async with _lock():
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, SPEND_KEY)
|
||||
spend = _spend_of(row)
|
||||
if row is None or spend["date"] != date:
|
||||
return
|
||||
spend["usd"] = max(0.0, round(spend["usd"] + delta, 8))
|
||||
row.value = spend
|
||||
await session.commit()
|
||||
|
||||
|
||||
__all__ = ["SPEND_KEY", "charge", "daily_spend", "settle"]
|
||||
@@ -0,0 +1,107 @@
|
||||
"""核查花费账（engine/spend.py）：双盲复算、反方审稿、审批质询的模型调用按本地日期记一本每日账。
|
||||
|
||||
照结论句裁判的每日花费做（judge.py 的 SPEND_KEY、_charge、_settle），但是独立的一行设置：
|
||||
|
||||
- 调用前按估算预留，加上这笔会超过每日上限就不记、返回 None；上限是 None 表示不限
|
||||
- 调用后按实际用量找补差额；跨了天的不找补（预留记在前一天，今天从零算）
|
||||
- 价格目录里没有的模型估不出金额：只计次数（unpriced_calls），不计金额
|
||||
- 和裁判的账互不挤占：核查花掉的钱不算进裁判的每日上限，反过来也一样
|
||||
- 并发预留不会一起越过上限
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import delete
|
||||
|
||||
from app.db.base import SessionLocal
|
||||
from app.db.models import Setting
|
||||
from app.engine import spend
|
||||
from app.engine import judge as judge_mod
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def clean(monkeypatch):
|
||||
monkeypatch.setattr(spend, "_today", lambda: "2026-10-02")
|
||||
async with SessionLocal() as session:
|
||||
await session.execute(delete(Setting).where(Setting.key.in_([spend.SPEND_KEY, judge_mod.SPEND_KEY])))
|
||||
await session.commit()
|
||||
yield
|
||||
|
||||
|
||||
async def stored() -> dict | None:
|
||||
async with SessionLocal() as session:
|
||||
row = await session.get(Setting, spend.SPEND_KEY)
|
||||
return row.value if row is not None else None
|
||||
|
||||
|
||||
async def test_a_call_within_the_cap_is_recorded():
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
assert await spend.charge(0.4, limit=1.0, priced=True) == "2026-10-02"
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.4, "calls": 1, "unpriced_calls": 0}
|
||||
|
||||
|
||||
async def test_a_call_that_would_cross_the_cap_is_refused_and_nothing_is_recorded():
|
||||
assert await spend.charge(0.9, limit=1.0, priced=True)
|
||||
assert await spend.charge(0.2, limit=1.0, priced=True) is None
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.9, "calls": 1, "unpriced_calls": 0}
|
||||
# 正好到上限的还能记：上限是「不超过」
|
||||
assert await spend.charge(0.1, limit=1.0, priced=True) == "2026-10-02"
|
||||
|
||||
|
||||
async def test_no_cap_means_unlimited():
|
||||
assert await spend.charge(50.0, limit=None, priced=True) == "2026-10-02"
|
||||
assert (await spend.daily_spend())["usd"] == 50.0
|
||||
|
||||
|
||||
async def test_unpriced_calls_are_counted_without_money():
|
||||
assert await spend.charge(0.0, limit=0.5, priced=False) == "2026-10-02"
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 1, "unpriced_calls": 1}
|
||||
|
||||
|
||||
async def test_yesterdays_spend_does_not_count_today():
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=spend.SPEND_KEY, value={"date": "2026-10-01", "usd": 1.9, "calls": 7,
|
||||
"unpriced_calls": 2}))
|
||||
await session.commit()
|
||||
assert await spend.daily_spend() == {"date": "2026-10-02", "usd": 0.0, "calls": 0, "unpriced_calls": 0}
|
||||
assert await spend.charge(1.5, limit=2.0, priced=True) == "2026-10-02"
|
||||
assert await stored() == {"date": "2026-10-02", "usd": 1.5, "calls": 1, "unpriced_calls": 0}
|
||||
|
||||
|
||||
async def test_settle_corrects_the_reservation_and_never_goes_below_zero():
|
||||
day = await spend.charge(0.5, limit=None, priced=True)
|
||||
await spend.settle(day, -0.2)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.3)
|
||||
await spend.settle(day, 0.15)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.45)
|
||||
await spend.settle(day, -5)
|
||||
assert (await spend.daily_spend())["usd"] == 0.0
|
||||
assert (await spend.daily_spend())["calls"] == 1, "找补只改金额，不改次数"
|
||||
|
||||
|
||||
async def test_settling_after_midnight_does_not_touch_the_new_day(monkeypatch):
|
||||
day = await spend.charge(0.5, limit=None, priced=True)
|
||||
monkeypatch.setattr(spend, "_today", lambda: "2026-10-03")
|
||||
assert await spend.charge(0.2, limit=None, priced=True) == "2026-10-03"
|
||||
await spend.settle(day, -0.4)
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.2), "预留记在前一天，今天从零算"
|
||||
|
||||
|
||||
async def test_the_verify_ledger_is_separate_from_the_judges(monkeypatch):
|
||||
monkeypatch.setattr(judge_mod, "_today", lambda: "2026-10-02")
|
||||
assert spend.SPEND_KEY != judge_mod.SPEND_KEY
|
||||
await spend.charge(1.2, limit=2.0, priced=True)
|
||||
assert (await judge_mod.daily_spend())["usd"] == 0.0, "核查花的钱不进裁判的账"
|
||||
async with SessionLocal() as session:
|
||||
session.add(Setting(key=judge_mod.SPEND_KEY, value={"date": "2026-10-02", "usd": 1.9, "calls": 4}))
|
||||
await session.commit()
|
||||
assert await spend.charge(0.7, limit=2.0, priced=True) == "2026-10-02", "裁判花的钱也不挤占核查的上限"
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(1.9)
|
||||
|
||||
|
||||
async def test_concurrent_reservations_never_cross_the_cap_together():
|
||||
results = await asyncio.gather(*(spend.charge(0.3, limit=1.0, priced=True) for _ in range(10)))
|
||||
assert sum(1 for r in results if r) == 3
|
||||
assert (await spend.daily_spend())["usd"] == pytest.approx(0.9)
|
||||
Reference in New Issue
Block a user