Spaces:
Paused
Paused
File size: 5,593 Bytes
5b063c2 536e671 5b063c2 8501103 5b063c2 536e671 5b063c2 8501103 5b063c2 8501103 5b063c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """The rupee budget for AI usage: one shared pool for the test group, an equal allocation per tester.
Every model call is recorded (api.logstore `usage`). A call is billable when a provider charges for it;
OpenRouter-exempt calls (free models, or anything OpenRouter reports at zero cost) are recorded with the
reason and never deducted. GPU/CPU jobs are not part of this pool: they run on Modal's free credit and have
their own dollar limit (api.jobs, MONTHLY_CREDIT_LIMIT_USD).
Before an AI task runs, its estimate is checked against the tester's remaining allocation and the pool; after
it runs, the actual cost (as OpenRouter reports it, else computed from list prices) is what gets deducted.
"""
from __future__ import annotations
from dataclasses import dataclass
from .routing import ModelSpec, call_cost_usd
CATEGORIES = ("coding", "simulation_ai", "debugging", "testing")
class BudgetError(Exception):
def __init__(self, status: int, message: str, detail: dict | None = None):
super().__init__(message)
self.status = status
self.detail = detail or {}
@dataclass
class CallCost:
cost_usd: float
billable: bool
exempt_reason: str | None
estimated: bool
def cost_of_call(model: ModelSpec, usage: dict) -> CallCost:
"""What one model call cost, from the provider's own figure when it reports one."""
tin = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
tout = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
reported = usage.get("cost")
if model.free:
return CallCost(0.0, False, "opencode_zen_free", False)
if model.id.endswith(":free"):
return CallCost(0.0, False, "free model (no OpenRouter charge)", False)
if reported is not None:
cost = float(reported)
if cost <= 0:
return CallCost(0.0, False, "OpenRouter reported no charge", False)
return CallCost(cost, True, None, False)
return CallCost(call_cost_usd(model, tin, tout), True, None, True)
def record_call(logs, settings, *, tester, generation_id, category, model: ModelSpec, usage: dict, estimate_inr=None):
c = cost_of_call(model, usage)
if tester == "owner":
c.billable = False
c.exempt_reason = "owner key usage"
logs.record_usage(
tester=tester,
generation_id=generation_id,
category=category,
model=model.id,
tier=model.tier,
provider="opencode-zen"
if model.free
else ("openrouter" if "openrouter" in settings.llm_base_url else settings.llm_base_url),
input_tokens=int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0),
output_tokens=int(usage.get("completion_tokens") or usage.get("output_tokens") or 0),
cost_usd=round(c.cost_usd, 6),
cost_inr=round(c.cost_usd * settings.usd_inr, 4),
billable=int(c.billable),
exempt_reason=c.exempt_reason,
estimated=int(c.estimated),
estimate_inr=estimate_inr,
)
def summary(logs, settings, tester: str | None = None) -> dict:
"""Used / remaining for one tester (their allocation) and for the whole pool, split by category."""
rows = logs.usage_rows()
mine = [r for r in rows if tester is None or r["tester"] == tester]
billable = [r for r in mine if r["billable"]]
by_cat = {c: round(sum(r["cost_inr"] for r in billable if r["category"] == c), 2) for c in CATEGORIES}
used = round(sum(r["cost_inr"] for r in billable), 2)
pool_used = round(sum(r["cost_inr"] for r in rows if r["billable"]), 2)
return {
"currency": "INR",
"allocation": settings.budget_per_user_inr,
"used": used,
"remaining": round(max(settings.budget_per_user_inr - used, 0.0), 2),
"by_category": by_cat,
"exempt_calls": sum(1 for r in mine if not r["billable"]),
"calls": len(mine),
"pool": {
"total": settings.budget_total_inr,
"used": pool_used,
"remaining": round(max(settings.budget_total_inr - pool_used, 0.0), 2),
},
}
def team(logs, settings, testers: list[str]) -> list[dict]:
rows = logs.usage_rows()
out = []
for name in sorted(set(testers) | {r["tester"] for r in rows}):
if name == "owner":
continue
used = round(sum(r["cost_inr"] for r in rows if r["tester"] == name and r["billable"]), 2)
out.append(
{
"tester": name,
"allocation": settings.budget_per_user_inr,
"used": used,
"remaining": round(max(settings.budget_per_user_inr - used, 0.0), 2),
}
)
return out
def check(logs, settings, tester: str, estimate: tuple[float, float], confirmed: bool) -> None:
"""Refuse when the tester's allocation or the pool cannot cover the likely cost; ask above the threshold."""
if tester == "owner":
return
s = summary(logs, settings, tester)
likely, worst = estimate
if s["pool"]["remaining"] < likely:
raise BudgetError(402, "The team's AI budget is used up. Please tell the Orcanion team.", {"budget": s})
if s["remaining"] < likely:
raise BudgetError(402, "Your AI budget is used up. Please tell the Orcanion team.", {"budget": s})
if worst >= settings.confirm_above_inr and not confirmed:
raise BudgetError(
409,
"This task may cost more than usual; confirm to run it.",
{"needs_confirmation": True, "estimate_inr": list(estimate), "budget": s},
)
|