# retoor from __future__ import annotations import logging import re import uuid from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Optional from pydantic import BaseModel, Field, field_validator, model_validator from devplacepy.database import bump_cache_version, db, get_table, sync_local_cache logger = logging.getLogger(__name__) RULES_TABLE = "gateway_quota_rules" RESETS_TABLE = "gateway_quota_resets" CACHE_NAME = "gateway_quota" APP_REFERENCE_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]{1,30}$") OWNER_KINDS = ("internal", "key", "user", "admin", "anonymous") FIELD_DEFAULT_USER = "gateway_default_user_daily_usd" FIELD_DEFAULT_ADMIN = "gateway_default_admin_daily_usd" FIELD_DEFAULT_GUEST = "gateway_default_guest_daily_usd" FIELD_DEFAULT_INTERNAL = "gateway_default_internal_daily_usd" FIELD_DEFAULT_KEY = "gateway_default_key_daily_usd" _DEFAULT_FIELD_BY_KIND = { "user": FIELD_DEFAULT_USER, "admin": FIELD_DEFAULT_ADMIN, "anonymous": FIELD_DEFAULT_GUEST, "internal": FIELD_DEFAULT_INTERNAL, "key": FIELD_DEFAULT_KEY, } _QUOTA_CACHE: dict = {} def _now() -> str: return datetime.now(timezone.utc).isoformat() def ensure_tables() -> None: db.query( "CREATE TABLE IF NOT EXISTS " + RULES_TABLE + " (id INTEGER PRIMARY KEY, uid TEXT, owner_kind TEXT, owner_id TEXT, " "app_reference TEXT, limit_usd REAL DEFAULT 0, is_active INTEGER DEFAULT 1, " "label TEXT, created_by TEXT, created_at TEXT, updated_at TEXT)" ) try: db.query( "CREATE UNIQUE INDEX IF NOT EXISTS idx_gateway_quota_rules_uid ON " + RULES_TABLE + " (uid)" ) db.query( "CREATE INDEX IF NOT EXISTS idx_gateway_quota_rules_lookup ON " + RULES_TABLE + " (owner_kind, owner_id, app_reference)" ) except Exception as exc: logger.warning("gateway quota rule index creation failed: %s", exc) db.query( "CREATE TABLE IF NOT EXISTS " + RESETS_TABLE + " (id INTEGER PRIMARY KEY, uid TEXT, owner_kind TEXT, owner_id TEXT, " "app_reference TEXT, reset_at TEXT, created_by TEXT)" ) try: db.query( "CREATE UNIQUE INDEX IF NOT EXISTS idx_gateway_quota_resets_uid ON " + RESETS_TABLE + " (uid)" ) db.query( "CREATE INDEX IF NOT EXISTS idx_gateway_quota_resets_lookup ON " + RESETS_TABLE + " (owner_kind, owner_id, app_reference)" ) except Exception as exc: logger.warning("gateway quota reset index creation failed: %s", exc) class QuotaScopeIn(BaseModel): owner_kind: Optional[str] = None owner_id: Optional[str] = Field(default=None, max_length=64) app_reference: Optional[str] = Field(default=None, max_length=30) @field_validator("owner_kind") @classmethod def _clean_owner_kind(cls, value: Optional[str]) -> Optional[str]: value = (value or "").strip().lower() if not value: return None if value not in OWNER_KINDS: raise ValueError(f"owner_kind must be one of {', '.join(OWNER_KINDS)}") return value @field_validator("owner_id") @classmethod def _clean_owner_id(cls, value: Optional[str]) -> Optional[str]: return (value or "").strip() or None @field_validator("app_reference") @classmethod def _clean_app_reference(cls, value: Optional[str]) -> Optional[str]: value = (value or "").strip() if not value: return None if not APP_REFERENCE_PATTERN.match(value): raise ValueError("app_reference must match ^[a-zA-Z0-9_.-]{1,30}$") return value class QuotaRuleIn(QuotaScopeIn): limit_usd: float = Field(default=0.0, ge=0) is_active: bool = True label: str = Field(default="", max_length=200) @field_validator("label") @classmethod def _clean_label(cls, value: str) -> str: return (value or "").strip() @model_validator(mode="after") def _require_a_dimension(self) -> "QuotaRuleIn": if self.owner_kind is None and self.owner_id is None and self.app_reference is None: raise ValueError( "At least one of owner_kind, owner_id, or app_reference is required - " "an unscoped cap belongs in the global default fields, not a rule" ) return self class QuotaResetIn(QuotaScopeIn): pass @dataclass(frozen=True) class QuotaRule: uid: str owner_kind: Optional[str] owner_id: Optional[str] app_reference: Optional[str] limit_usd: float is_active: bool label: str created_by: str created_at: str updated_at: str @property def specificity(self) -> int: return sum( 1 for v in (self.owner_kind, self.owner_id, self.app_reference) if v is not None ) def matches(self, owner_kind: str, owner_id: str, app_reference: str) -> bool: if not self.is_active: return False if self.owner_kind is not None and self.owner_kind != owner_kind: return False if self.owner_id is not None and self.owner_id != owner_id: return False if self.app_reference is not None and self.app_reference != app_reference: return False return True def as_dict(self) -> dict: return { "uid": self.uid, "owner_kind": self.owner_kind, "owner_id": self.owner_id, "app_reference": self.app_reference, "limit_usd": self.limit_usd, "is_active": self.is_active, "label": self.label, "created_by": self.created_by, "created_at": self.created_at, "updated_at": self.updated_at, "specificity": self.specificity, } def _row_to_rule(row: dict) -> QuotaRule: return QuotaRule( uid=str(row.get("uid") or ""), owner_kind=row.get("owner_kind") or None, owner_id=row.get("owner_id") or None, app_reference=row.get("app_reference") or None, limit_usd=float(row.get("limit_usd") or 0.0), is_active=bool(row.get("is_active", 1)), label=str(row.get("label") or ""), created_by=str(row.get("created_by") or ""), created_at=str(row.get("created_at") or ""), updated_at=str(row.get("updated_at") or ""), ) def _load() -> list[QuotaRule]: sync_local_cache(CACHE_NAME, _QUOTA_CACHE) if "rules" not in _QUOTA_CACHE: rules: list[QuotaRule] = [] try: if RULES_TABLE in db.tables: for row in get_table(RULES_TABLE).all(): if row.get("uid"): rules.append(_row_to_rule(row)) except Exception as exc: logger.warning("gateway quota rule load failed: %s", exc) _QUOTA_CACHE["rules"] = rules return _QUOTA_CACHE["rules"] class QuotaRuleStore: def list(self) -> list[dict]: rules = sorted(_load(), key=lambda r: (-r.specificity, r.created_at)) return [r.as_dict() for r in rules] def get(self, uid: str) -> Optional[QuotaRule]: if not uid: return None for rule in _load(): if rule.uid == uid: return rule return None def count(self) -> int: return len(_load()) def set( self, payload: QuotaRuleIn, *, uid: Optional[str] = None, created_by: str = "" ) -> dict: ensure_tables() table = get_table(RULES_TABLE) existing = table.find_one(uid=uid) if uid else None record: dict = { "owner_kind": payload.owner_kind, "owner_id": payload.owner_id, "app_reference": payload.app_reference, "limit_usd": payload.limit_usd, "is_active": 1 if payload.is_active else 0, "label": payload.label, "updated_at": _now(), } if existing: record["uid"] = existing["uid"] record["created_by"] = existing.get("created_by") or created_by record["created_at"] = existing.get("created_at") or _now() table.update({**record, "id": existing["id"]}, ["id"]) else: record["uid"] = uid or uuid.uuid4().hex record["created_by"] = created_by record["created_at"] = _now() table.insert(record) bump_cache_version(CACHE_NAME) _QUOTA_CACHE.clear() saved = self.get(record["uid"]) return saved.as_dict() if saved else record def remove(self, uid: str) -> bool: uid = (uid or "").strip() if not uid or RULES_TABLE not in db.tables: return False removed = int(get_table(RULES_TABLE).delete(uid=uid)) if removed: bump_cache_version(CACHE_NAME) _QUOTA_CACHE.clear() return bool(removed) quota_rule_store = QuotaRuleStore() def _load_resets() -> list[dict]: sync_local_cache(CACHE_NAME, _QUOTA_CACHE) if "resets" not in _QUOTA_CACHE: resets: list[dict] = [] try: if RESETS_TABLE in db.tables: for row in get_table(RESETS_TABLE).all(): if row.get("reset_at"): resets.append( { "owner_kind": row.get("owner_kind") or None, "owner_id": row.get("owner_id") or None, "app_reference": row.get("app_reference") or None, "reset_at": str(row["reset_at"]), } ) except Exception as exc: logger.warning("gateway quota reset load failed: %s", exc) _QUOTA_CACHE["resets"] = resets return _QUOTA_CACHE["resets"] def _reset_applies(reset: dict, scope: dict) -> bool: for field in ("owner_kind", "owner_id", "app_reference"): wanted = reset.get(field) if wanted is None: continue if scope.get(field) is None or scope[field] != wanted: return False return True def reset_watermark( owner_kind: Optional[str], owner_id: Optional[str], app_reference: Optional[str] ) -> str: scope = { "owner_kind": owner_kind, "owner_id": owner_id, "app_reference": app_reference, } stamps = [r["reset_at"] for r in _load_resets() if _reset_applies(r, scope)] return max(stamps) if stamps else "" def reset(payload: Optional[QuotaResetIn] = None, *, created_by: str = "") -> dict: ensure_tables() payload = payload or QuotaResetIn() table = get_table(RESETS_TABLE) scope = { "owner_kind": payload.owner_kind, "owner_id": payload.owner_id, "app_reference": payload.app_reference, } stamp = _now() existing = table.find_one(**scope) if existing: table.update({"id": existing["id"], "reset_at": stamp, "created_by": created_by}, ["id"]) uid = existing.get("uid") or uuid.uuid4().hex else: uid = uuid.uuid4().hex table.insert({**scope, "uid": uid, "reset_at": stamp, "created_by": created_by}) bump_cache_version(CACHE_NAME) _QUOTA_CACHE.clear() return {**scope, "uid": uid, "reset_at": stamp} def scope_label(scope: dict, fallback: str = "") -> str: parts = [] if scope.get("owner_kind"): parts.append(f"role={scope['owner_kind']}") if scope.get("owner_id"): parts.append(f"user={scope['owner_id']}") if scope.get("app_reference"): parts.append(f"app={scope['app_reference']}") return ", ".join(parts) or fallback def default_limit(owner_kind: str, cfg: dict) -> float: field = _DEFAULT_FIELD_BY_KIND.get(owner_kind, FIELD_DEFAULT_USER) return float(cfg.get(field, 0.0) or 0.0) def resolve( owner_kind: str, owner_id: str, app_reference: str, cfg: dict ) -> tuple[float, tuple[Optional[str], Optional[str], Optional[str]], Optional[QuotaRule]]: matches = [r for r in _load() if r.matches(owner_kind, owner_id, app_reference)] if matches: def _sort_key(rule: QuotaRule): tie = float("inf") if rule.limit_usd == 0 else rule.limit_usd return (-rule.specificity, tie) best = sorted(matches, key=_sort_key)[0] return best.limit_usd, (best.owner_kind, best.owner_id, best.app_reference), best return default_limit(owner_kind, cfg), (owner_kind, owner_id, None), None def resolve_for_owner( owner_kind: str, owner_id: str, cfg: dict ) -> tuple[float, tuple[Optional[str], Optional[str], None], Optional[QuotaRule]]: matches = [ r for r in _load() if r.app_reference is None and r.matches(owner_kind, owner_id, "") ] if matches: def _sort_key(rule: QuotaRule): tie = float("inf") if rule.limit_usd == 0 else rule.limit_usd return (-rule.specificity, tie) best = sorted(matches, key=_sort_key)[0] return best.limit_usd, (best.owner_kind, best.owner_id, None), best return default_limit(owner_kind, cfg), (owner_kind, owner_id, None), None def spent_24h( owner_kind: Optional[str], owner_id: Optional[str], app_reference: Optional[str] ) -> float: from devplacepy.services.openai_gateway.usage import GATEWAY_LEDGER if GATEWAY_LEDGER not in db.tables: return 0.0 cutoff = (datetime.now(timezone.utc) - timedelta(hours=24)).isoformat() watermark = reset_watermark(owner_kind, owner_id, app_reference) clauses = ["created_at >= :cutoff"] params: dict = {"cutoff": max(cutoff, watermark) if watermark else cutoff} if owner_kind is not None: clauses.append("owner_kind = :owner_kind") params["owner_kind"] = owner_kind if owner_id is not None: clauses.append("owner_id = :owner_id") params["owner_id"] = owner_id if app_reference is not None: clauses.append("app_reference = :app_reference") params["app_reference"] = app_reference where = " AND ".join(clauses) rows = list( db.query( f"SELECT COALESCE(SUM(cost_usd), 0) AS spent FROM {GATEWAY_LEDGER} WHERE {where}", **params, ) ) return float(rows[0].get("spent") or 0.0) if rows else 0.0