Oleksandr Kozachuk
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""Daily usage metering + hard budget (SPEC-003, SAF-04..07).
|
|
|
|
The ledger estimates spend from configured prices and answers the one
|
|
question that matters fail-closed: may the bot still call the API
|
|
today? Counters live in the store's usage table when a store exists
|
|
(SAF-05), else in memory (degraded but safe).
|
|
"""
|
|
|
|
import time
|
|
from typing import Callable, Dict, Optional, Tuple
|
|
|
|
from .persistence import PersistentStore
|
|
|
|
DEFAULT_PRICE_INPUT_PER_M = 1.0
|
|
DEFAULT_PRICE_OUTPUT_PER_M = 6.0
|
|
DEFAULT_PRICE_PER_IMAGE = 0.05
|
|
|
|
|
|
class QuotaLedger:
|
|
def __init__(self, store: Optional[PersistentStore], config_getter: Callable[[], Dict]) -> None:
|
|
self.store = store
|
|
self._config = config_getter
|
|
self._memory: Dict[Tuple[str, str], float] = {}
|
|
|
|
@staticmethod
|
|
def _day() -> str:
|
|
return time.strftime("%Y-%m-%d")
|
|
|
|
def _add(self, key: str, amount: float) -> None:
|
|
if self.store is not None:
|
|
self.store.usage_add(self._day(), key, amount)
|
|
else:
|
|
slot = (self._day(), key)
|
|
self._memory[slot] = self._memory.get(slot, 0.0) + amount
|
|
|
|
def _get(self, key: str) -> float:
|
|
if self.store is not None:
|
|
return self.store.usage_get(self._day(), key)
|
|
return self._memory.get((self._day(), key), 0.0)
|
|
|
|
def add_tokens(self, prompt_tokens: int, completion_tokens: int) -> None:
|
|
self._add("tokens-in", prompt_tokens)
|
|
self._add("tokens-out", completion_tokens)
|
|
|
|
def add_images(self, count: int = 1) -> None:
|
|
self._add("images", count)
|
|
|
|
def tokens_today(self) -> Tuple[int, int]:
|
|
return int(self._get("tokens-in")), int(self._get("tokens-out"))
|
|
|
|
def images_today(self) -> int:
|
|
return int(self._get("images"))
|
|
|
|
def count_user_message(self, user: str) -> int:
|
|
self._add(f"msg-user:{user}", 1)
|
|
return int(self._get(f"msg-user:{user}"))
|
|
|
|
def count_user_image(self, user: str) -> int:
|
|
self._add(f"img-user:{user}", 1)
|
|
return int(self._get(f"img-user:{user}"))
|
|
|
|
def user_images(self, user: str) -> int:
|
|
return int(self._get(f"img-user:{user}"))
|
|
|
|
def spent_usd(self) -> float:
|
|
config = self._config()
|
|
tokens_in, tokens_out = self.tokens_today()
|
|
price_in = float(config.get("price-input-per-m", DEFAULT_PRICE_INPUT_PER_M))
|
|
price_out = float(config.get("price-output-per-m", DEFAULT_PRICE_OUTPUT_PER_M))
|
|
price_image = float(config.get("price-per-image", DEFAULT_PRICE_PER_IMAGE))
|
|
return (tokens_in * price_in + tokens_out * price_out) / 1_000_000.0 + self.images_today() * price_image
|
|
|
|
def budget_ok(self) -> bool:
|
|
config = self._config()
|
|
if "daily-budget-usd" not in config:
|
|
return True
|
|
return self.spent_usd() < float(config["daily-budget-usd"])
|