Spaces:
Paused
Paused
| """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 {} | |
| 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}, | |
| ) | |