# retoor <retoor@molodetz.nl>
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
Reference in New Issue View Git Blame Copy Permalink