mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 02:14:23 +00:00
Alex
6826313b60
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.
100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
"""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"),
|
|
)
|