59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
"""Cost ledger for cloud model usage and a budget guard.
|
|
|
|
.env: BUDGET_LIMIT_USD (default 45), BUDGET_CYCLE_START (YYYY-MM-DD, default 2026-10-02).
|
|
Prices are DeepSeek V4.1 Flash list prices without cache discount (upper bound).
|
|
"""
|
|
import json
|
|
import os
|
|
import time
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
LEDGER = os.path.join(ROOT, "runs", "ledger.jsonl")
|
|
PRICE_IN, PRICE_OUT = 0.30 / 1e6, 1.20 / 1e6
|
|
|
|
|
|
class BudgetExceeded(Exception):
|
|
pass
|
|
|
|
|
|
def add_usage(model, usage, kind, ref=""):
|
|
if ":cloud" not in model:
|
|
return
|
|
pin, pout = usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0)
|
|
os.makedirs(os.path.dirname(LEDGER), exist_ok=True)
|
|
with open(LEDGER, "a") as f:
|
|
f.write(json.dumps({"t": time.time(), "day": time.strftime("%Y-%m-%d"), "model": model, "kind": kind,
|
|
"ref": ref, "in": pin, "out": pout,
|
|
"usd": round(pin * PRICE_IN + pout * PRICE_OUT, 5)}) + "\n")
|
|
|
|
|
|
def spent(since=None):
|
|
since = since or os.environ.get("BUDGET_CYCLE_START", "2026-10-02")
|
|
if not os.path.exists(LEDGER):
|
|
return 0.0
|
|
return round(sum(json.loads(l)["usd"] for l in open(LEDGER) if json.loads(l)["day"] >= since), 4)
|
|
|
|
|
|
def _env_budget():
|
|
"""BUDGET_LIMIT_USD and BUDGET_RESERVE_USD from .env (read at each check: a change applies at once)."""
|
|
vals = {"BUDGET_LIMIT_USD": os.environ.get("BUDGET_LIMIT_USD", "45"),
|
|
"BUDGET_RESERVE_USD": os.environ.get("BUDGET_RESERVE_USD", "0")}
|
|
try:
|
|
for line in open(os.path.join(ROOT, ".env")):
|
|
k, _, v = line.strip().partition("=")
|
|
if k in vals and v:
|
|
vals[k] = v
|
|
except OSError:
|
|
pass
|
|
return float(vals["BUDGET_LIMIT_USD"]), float(vals["BUDGET_RESERVE_USD"])
|
|
|
|
|
|
def check_budget():
|
|
# BUDGET_RESERVE_USD: ledger USD kept back (second teacher test later); the guard stops at limit - reserve
|
|
limit_raw, reserve = _env_budget()
|
|
limit = limit_raw - reserve
|
|
s = spent()
|
|
if s >= limit:
|
|
raise BudgetExceeded(f"cycle spend {s} USD >= limit {limit} USD")
|
|
return s
|