UnknownPixel/askPESU

1"""LLM Quota Management using a State Machine."""2 3import datetime4import logging5from dataclasses import dataclass6 7import pytz8 9IST = pytz.timezone("Asia/Kolkata")10 11 12@dataclass13class QuotaState:14    """Manages LLM quota state with cooldown logic."""15 16    name: str17    enabled: bool = True18    disabled_until: datetime.datetime | None = None19    cooldown_hours: int = 2420 21    def refresh(self) -> None:22        """Re-enable if cooldown has expired."""23        now = datetime.datetime.now(IST)24        if not self.enabled and self.disabled_until and now >= self.disabled_until:25            self.enabled, self.disabled_until = True, None26            logging.info(f"{self.name} cooldown expired, re-enabled for use.")27 28    def disable(self) -> None:29        """Disable for cooldown period."""30        now = datetime.datetime.now(IST)31        self.enabled = False32        self.disabled_until = now + datetime.timedelta(hours=self.cooldown_hours)33        logging.warning(f"Quota exceeded on llm:{self.name}. Disabled until {self.disabled_until}")34 35    def status(self) -> dict:36        """Get current status."""37        return {38            "available": self.enabled,39            "next_available": self.disabled_until.isoformat() if not self.enabled and self.disabled_until else None,40        }41