Files
company-haness/.claude/hooks/bench_cascade/budget.py
T

50 lines
1.7 KiB
Python

"""전 유료 모델 호출(arm-run·calibrate·judge·retry·LLM sanitize)에 대한 run-level 예산 receipt.
receipt 없이 실행 거부 — 우발적 대량 API 소비 방지(Blocker 4)."""
import json
import os
def approve(plan_id, max_tokens, max_cost, out_path):
rec = {"plan-id": plan_id, "max-tokens": int(max_tokens), "max-cost": float(max_cost),
"spent-tokens": 0, "spent-cost": 0.0}
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "w", encoding="utf-8") as f:
json.dump(rec, f)
return rec
def load(path):
if not os.path.exists(path):
return None
with open(path, encoding="utf-8") as f:
return json.load(f)
def remaining(path):
r = load(path)
if r is None:
return {"tokens": 0, "cost": 0.0}
return {"tokens": r["max-tokens"] - r["spent-tokens"], "cost": r["max-cost"] - r["spent-cost"]}
def charge(path, tokens, cost):
r = load(path)
if r is None:
raise RuntimeError("예산 receipt 없음 — approve-budget 먼저")
if r["spent-tokens"] + tokens > r["max-tokens"] or r["spent-cost"] + cost > r["max-cost"]:
raise ValueError(f"예산 초과: 요구 {tokens}tok/{cost}$ > 잔여 {remaining(path)}")
r["spent-tokens"] += int(tokens)
r["spent-cost"] += float(cost)
with open(path, "w", encoding="utf-8") as f:
json.dump(r, f)
return r
def require(path):
r = load(path)
if r is None:
raise RuntimeError("예산 receipt 없음 — 유료 실행 거부(approve-budget 필요)")
if r["max-tokens"] - r["spent-tokens"] <= 0 or r["max-cost"] - r["spent-cost"] <= 0:
raise RuntimeError("예산 소진 — 유료 실행 거부")
return r