| """Query/spend accountant for Tier B. All paid pressure passes through here.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| class Budget: | |
| max_queries: int | |
| usd_cap: float | |
| queries_used: int = 0 | |
| usd_spent: float = 0.0 | |
| per_model_queries: dict[str, int] = field(default_factory=dict) | |
| def can_spend(self, n: int = 1) -> bool: | |
| return self.queries_used + n <= self.max_queries and self.usd_spent < self.usd_cap | |
| def record(self, model: str, usd: float, n: int = 1) -> None: | |
| self.queries_used += n | |
| self.usd_spent += usd | |
| self.per_model_queries[model] = self.per_model_queries.get(model, 0) + n | |