finbot/limits.py

"""Per-user throttling for model-backed work.

Only relevant in open-access mode. Reading a receipt or a statement spends the
Space owner's inference credits, so without a cap a single stranger with a
folder of PDFs can run up an unbounded bill on someone else's token.

Deliberately in-memory: it resets when the Space restarts or sleeps, which is
fine for a budget guard and avoids putting other people's usage records in the
owner's dataset. Plain text costs nothing (the deterministic parser handles it)
and is never throttled.
"""

from __future__ import annotations

import threading
import time
from collections import defaultdict

WINDOW_SECONDS = 24 * 60 * 60


class RateLimiter:
    def __init__(self, per_day: int):
        #: Zero or negative disables the limiter entirely.
        self.per_day = per_day
        self._hits: dict[int, list[float]] = defaultdict(list)
        self._lock = threading.Lock()

    @property
    def enabled(self) -> bool:
        return self.per_day > 0

    def _prune(self, user_id: int, now: float) -> list[float]:
        recent = [t for t in self._hits[user_id] if now - t < WINDOW_SECONDS]
        self._hits[user_id] = recent
        return recent

    def check(self, user_id: int, now: float | None = None) -> bool:
        """Record an attempt. False means the user is over their daily cap."""
        if not self.enabled:
            return True
        now = now if now is not None else time.time()
        with self._lock:
            recent = self._prune(user_id, now)
            if len(recent) >= self.per_day:
                return False
            recent.append(now)
            return True

    def remaining(self, user_id: int, now: float | None = None) -> int:
        if not self.enabled:
            return -1  # unlimited
        now = now if now is not None else time.time()
        with self._lock:
            return max(0, self.per_day - len(self._prune(user_id, now)))