mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-17 06:14:25 +00:00
feat(quotas): limit resolution and the quota service
docsgpt/quotas resolves a user's effective token and cost limits from the policy rows that apply to them: their own override, then the most generous allowance among their teams, then the instance default, then any defaults a registered provider supplies. Each budget resolves on its own, and a team counts once however many memberships the user holds in it. QuotaService compares those limits with the user's token_usage totals over the current QUOTA_PERIOD window (calendar-aligned, UTC, computed at read time). It fails open, and skips the usage query for unlimited users.
This commit is contained in:
1 parent
6c42139224
commit
6826313b60
12 files changed
+1175
No files matched your search
@@ -0,0 +1,99 @@
|
||||
"""Resolve the policy rows that apply to a user into effective limits.
|
||||
|
||||
Each budget (tokens, cost) resolves on its own: the user's row wins, then the
|
||||
most generous of the user's team rows, then the instance row, then the
|
||||
registered defaults. A row with neither a limit nor the unlimited flag for a
|
||||
budget has no opinion on it and is skipped.
|
||||
|
||||
Teams resolve to the most generous allowance because team membership is not
|
||||
controlled by the instance admin: under "most restrictive", any team admin
|
||||
could throttle a user by adding them to a low-allowance team.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Mapping, Optional
|
||||
|
||||
LAYERS = ("user", "team", "instance", "default")
|
||||
|
||||
_FIELDS = {"tokens": ("token_limit", "token_unlimited"), "cost": ("cost_limit_usd", "cost_unlimited")}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedLimit:
|
||||
"""One budget's effective limit and the layer it came from.
|
||||
|
||||
``limit`` is ``None`` when the budget is unlimited. ``source`` is ``None``
|
||||
when no layer had an opinion, otherwise one of ``LAYERS``; ``source_id`` is
|
||||
the team id for a team-sourced limit.
|
||||
"""
|
||||
|
||||
limit: Optional[float] = None
|
||||
source: Optional[str] = None
|
||||
source_id: Optional[str] = None
|
||||
|
||||
@property
|
||||
def unlimited(self) -> bool:
|
||||
return self.limit is None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedLimits:
|
||||
"""A user's effective token and cost limits for one bucket."""
|
||||
|
||||
tokens: ResolvedLimit
|
||||
cost: ResolvedLimit
|
||||
|
||||
@property
|
||||
def unlimited(self) -> bool:
|
||||
return self.tokens.unlimited and self.cost.unlimited
|
||||
|
||||
|
||||
def _opinion(row: Mapping, budget: str) -> Optional[tuple[bool, Optional[float]]]:
|
||||
"""Return ``(unlimited, limit)`` for a row's budget, or ``None`` if it defers."""
|
||||
limit_field, unlimited_field = _FIELDS[budget]
|
||||
if row.get(unlimited_field):
|
||||
return True, None
|
||||
value = row.get(limit_field)
|
||||
if value is None:
|
||||
return None
|
||||
return False, float(value)
|
||||
|
||||
|
||||
def _resolve_budget(rows: Iterable[Mapping], budget: str) -> ResolvedLimit:
|
||||
by_layer: dict[str, list[tuple[Mapping, tuple[bool, Optional[float]]]]] = {}
|
||||
for row in rows:
|
||||
opinion = _opinion(row, budget)
|
||||
if opinion is not None:
|
||||
by_layer.setdefault(row["scope"], []).append((row, opinion))
|
||||
for layer in LAYERS:
|
||||
candidates = by_layer.get(layer)
|
||||
if not candidates:
|
||||
continue
|
||||
# Most generous first: unlimited, then the larger limit. Only the team
|
||||
# layer can hold more than one candidate. Ties break on subject id so
|
||||
# the reported source is stable.
|
||||
row, (unlimited, limit) = min(
|
||||
candidates,
|
||||
key=lambda c: (not c[1][0], -(c[1][1] or 0.0), str(c[0].get("subject_id") or "")),
|
||||
)
|
||||
source_id = str(row["subject_id"]) if layer == "team" else None
|
||||
return ResolvedLimit(limit=None if unlimited else limit, source=layer, source_id=source_id)
|
||||
return ResolvedLimit()
|
||||
|
||||
|
||||
def resolve_limits(rows: Iterable[Mapping], bucket: str = "all") -> ResolvedLimits:
|
||||
"""Return the effective limits for ``bucket`` from a user's applicable rows.
|
||||
|
||||
Args:
|
||||
rows: Policy rows that apply to the user (their own, their teams', the
|
||||
instance's and any provider defaults). Disabled rows and rows for
|
||||
other buckets are ignored.
|
||||
bucket: The policy bucket to resolve.
|
||||
"""
|
||||
applicable = [r for r in rows if r.get("bucket", "all") == bucket and r.get("enabled", True)]
|
||||
return ResolvedLimits(
|
||||
tokens=_resolve_budget(applicable, "tokens"),
|
||||
cost=_resolve_budget(applicable, "cost"),
|
||||
)
|
||||
Reference in new issue
Block a user