safety layer: hard daily budget, user quotas, spend report, forgetme + privacy

fjerkroa_bot/quota.py

@@ -0,0 +1,77 @@
+"""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"])