mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 22:15:10 +00:00
feat(quotas): admin quota API and GET /api/user/quota
Admins read and set the instance default, team allowances and user overrides under /api/admin/quotas. A user's endpoint also returns the limits those layers resolve to, the layer each came from and the usage against them. The overview lists catalog models used this period that have no price, since a cost limit cannot see them. Every write is audited. GET /api/user/quota gives a user their own limited buckets, usage and reset time without naming the policies behind them; any valid token may call it.
This commit is contained in:
1 parent
5406550bca
commit
21f43b2966
6 files changed
+616
No files matched your search
@@ -1,3 +1,4 @@
|
||||
from .routes import admin_ns
|
||||
from . import quotas # noqa: F401 (registers the quota resources on admin_ns)
|
||||
|
||||
__all__ = ["admin_ns"]
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Admin endpoints for usage quotas (RBAC ``admin`` role required).
|
||||
|
||||
Policies are set at three layers: the instance default, a team's per-member
|
||||
allowance and a single user's override. Every write is audited to
|
||||
``auth_events`` with the acting admin recorded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any, Optional
|
||||
|
||||
from flask import jsonify, make_response, request
|
||||
from flask_restx import Resource
|
||||
|
||||
from docsgpt.api.admin.routes import _actor, admin_ns
|
||||
from docsgpt.api.user.authz import admin_required
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.pricing import is_priced
|
||||
from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService
|
||||
from docsgpt.quotas.windows import window_bounds
|
||||
from docsgpt.storage.db.base_repository import looks_like_uuid
|
||||
from docsgpt.storage.db.repositories.auth_events import AuthEventsRepository
|
||||
from docsgpt.storage.db.repositories.quota_policies import BUCKETS, QuotaPoliciesRepository
|
||||
from docsgpt.storage.db.repositories.teams import TeamsRepository
|
||||
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from docsgpt.storage.db.repositories.users import UsersRepository
|
||||
from docsgpt.storage.db.session import db_readonly, db_session
|
||||
|
||||
_MAX_TOKEN_LIMIT = 2**62
|
||||
_MAX_COST_LIMIT = 99_999_999.0
|
||||
_MAX_NOTE_LENGTH = 500
|
||||
|
||||
|
||||
class _BadPolicy(ValueError):
|
||||
"""The request body does not describe a valid policy."""
|
||||
|
||||
|
||||
def _policy_json(row: dict) -> dict:
|
||||
cost = row.get("cost_limit_usd")
|
||||
return {
|
||||
"scope": row["scope"],
|
||||
"subject_id": row.get("subject_id"),
|
||||
"bucket": row["bucket"],
|
||||
"token_limit": row.get("token_limit"),
|
||||
"token_unlimited": bool(row.get("token_unlimited")),
|
||||
"cost_limit_usd": float(cost) if cost is not None else None,
|
||||
"cost_unlimited": bool(row.get("cost_unlimited")),
|
||||
"enabled": bool(row.get("enabled", True)),
|
||||
"note": row.get("note"),
|
||||
"updated_by": row.get("updated_by"),
|
||||
"updated_at": row.get("updated_at"),
|
||||
}
|
||||
|
||||
|
||||
def _error(message: str, status: int):
|
||||
return make_response(jsonify({"success": False, "message": message}), status)
|
||||
|
||||
|
||||
def _bucket(value: Any) -> str:
|
||||
if value is None:
|
||||
return "all"
|
||||
if value not in BUCKETS:
|
||||
raise _BadPolicy(f"bucket must be one of: {', '.join(BUCKETS)}")
|
||||
return value
|
||||
|
||||
|
||||
def _flag(data: dict, key: str, default: bool) -> bool:
|
||||
value = data.get(key, default)
|
||||
if not isinstance(value, bool):
|
||||
raise _BadPolicy(f"{key} must be a boolean")
|
||||
return value
|
||||
|
||||
|
||||
def _parse_policy(data: Any) -> dict:
|
||||
"""Validate a policy body into ``QuotaPoliciesRepository.upsert`` kwargs."""
|
||||
if not isinstance(data, dict):
|
||||
raise _BadPolicy("Body must be a JSON object")
|
||||
token_limit = data.get("token_limit")
|
||||
if token_limit is not None:
|
||||
if isinstance(token_limit, bool) or not isinstance(token_limit, int):
|
||||
raise _BadPolicy("token_limit must be a whole number or null")
|
||||
if not 0 <= token_limit <= _MAX_TOKEN_LIMIT:
|
||||
raise _BadPolicy("token_limit is out of range")
|
||||
cost_limit = data.get("cost_limit_usd")
|
||||
if cost_limit is not None:
|
||||
if isinstance(cost_limit, bool) or not isinstance(cost_limit, (int, float)):
|
||||
raise _BadPolicy("cost_limit_usd must be a number or null")
|
||||
if not math.isfinite(cost_limit) or not 0 <= cost_limit <= _MAX_COST_LIMIT:
|
||||
raise _BadPolicy("cost_limit_usd is out of range")
|
||||
cost_limit = round(float(cost_limit), 4)
|
||||
token_unlimited = _flag(data, "token_unlimited", False)
|
||||
cost_unlimited = _flag(data, "cost_unlimited", False)
|
||||
if token_unlimited and token_limit is not None:
|
||||
raise _BadPolicy("Set token_limit or token_unlimited, not both")
|
||||
if cost_unlimited and cost_limit is not None:
|
||||
raise _BadPolicy("Set cost_limit_usd or cost_unlimited, not both")
|
||||
if token_limit is None and cost_limit is None and not token_unlimited and not cost_unlimited:
|
||||
raise _BadPolicy("Set a limit or mark a budget unlimited; delete the policy to remove it")
|
||||
note = data.get("note")
|
||||
if note is not None:
|
||||
if not isinstance(note, str):
|
||||
raise _BadPolicy("note must be a string")
|
||||
note = note.strip()[:_MAX_NOTE_LENGTH] or None
|
||||
return {
|
||||
"bucket": _bucket(data.get("bucket")),
|
||||
"token_limit": token_limit,
|
||||
"token_unlimited": token_unlimited,
|
||||
"cost_limit_usd": cost_limit,
|
||||
"cost_unlimited": cost_unlimited,
|
||||
"enabled": _flag(data, "enabled", True),
|
||||
"note": note,
|
||||
}
|
||||
|
||||
|
||||
def _audit(conn, event: str, scope: str, subject_id: Optional[str], detail: dict) -> None:
|
||||
actor = _actor()
|
||||
AuthEventsRepository(conn).insert(
|
||||
# A user policy is filed under that user; the rest under the acting admin.
|
||||
subject_id if scope == "user" else (actor or "unknown"),
|
||||
event,
|
||||
ip=request.remote_addr,
|
||||
user_agent=request.headers.get("User-Agent"),
|
||||
metadata={"by": actor, "via": "admin_api", "scope": scope, "subject_id": subject_id, **detail},
|
||||
)
|
||||
|
||||
|
||||
def _put_policy(scope: str, subject_id: Optional[str]):
|
||||
try:
|
||||
fields = _parse_policy(request.get_json(silent=True))
|
||||
except _BadPolicy as exc:
|
||||
return _error(str(exc), 400)
|
||||
with db_session() as conn:
|
||||
row = QuotaPoliciesRepository(conn).upsert(
|
||||
scope=scope, subject_id=subject_id, actor=_actor(), **fields
|
||||
)
|
||||
_audit(conn, "quota_policy_set", scope, subject_id, fields)
|
||||
return make_response(jsonify({"success": True, "policy": _policy_json(row)}), 200)
|
||||
|
||||
|
||||
def _delete_policy(scope: str, subject_id: Optional[str]):
|
||||
try:
|
||||
bucket = _bucket(request.args.get("bucket")) if "bucket" in request.args else None
|
||||
except _BadPolicy as exc:
|
||||
return _error(str(exc), 400)
|
||||
with db_session() as conn:
|
||||
deleted = QuotaPoliciesRepository(conn).delete(scope, subject_id, bucket)
|
||||
if deleted:
|
||||
_audit(conn, "quota_policy_deleted", scope, subject_id, {"bucket": bucket or "*"})
|
||||
return make_response(jsonify({"success": True, "deleted": deleted}), 200)
|
||||
|
||||
|
||||
def _unpriced_models(conn) -> list[dict]:
|
||||
"""Catalog models used this period that no cost limit can see."""
|
||||
start, _ = window_bounds(settings.QUOTA_PERIOD)
|
||||
return [
|
||||
row
|
||||
for row in TokenUsageRepository(conn).tokens_by_model(start=start)
|
||||
# BYOM ids are UUIDs; those calls are $0 by design, not by omission.
|
||||
if not looks_like_uuid(row["model_id"]) and not is_priced(row["model_id"])
|
||||
]
|
||||
|
||||
|
||||
@admin_ns.route("/admin/quotas")
|
||||
class AdminQuotasResource(Resource):
|
||||
@admin_required
|
||||
def get(self):
|
||||
"""Every stored policy, grouped by layer, plus the models cost limits cannot see."""
|
||||
start, resets_at = window_bounds(settings.QUOTA_PERIOD)
|
||||
with db_readonly() as conn:
|
||||
repo = QuotaPoliciesRepository(conn)
|
||||
teams = {str(t["id"]): t for t in TeamsRepository(conn).list_all()}
|
||||
team_policies = []
|
||||
for row in repo.list_by_scope("team"):
|
||||
team = teams.get(str(row["subject_id"]), {})
|
||||
team_policies.append(
|
||||
{
|
||||
**_policy_json(row),
|
||||
"team_name": team.get("name"),
|
||||
"team_slug": team.get("slug"),
|
||||
"member_count": team.get("member_count"),
|
||||
}
|
||||
)
|
||||
body = {
|
||||
"success": True,
|
||||
"period": settings.QUOTA_PERIOD,
|
||||
"period_start": start.isoformat(),
|
||||
"resets_at": resets_at.isoformat(),
|
||||
"instance": [_policy_json(r) for r in repo.list_by_scope("instance")],
|
||||
"teams": team_policies,
|
||||
"users": [_policy_json(r) for r in repo.list_by_scope("user")],
|
||||
"unpriced_models": _unpriced_models(conn),
|
||||
}
|
||||
return make_response(jsonify(body), 200)
|
||||
|
||||
|
||||
@admin_ns.route("/admin/quotas/instance")
|
||||
class AdminInstanceQuotaResource(Resource):
|
||||
@admin_required
|
||||
def put(self):
|
||||
"""Set the instance default for one bucket."""
|
||||
return _put_policy("instance", None)
|
||||
|
||||
@admin_required
|
||||
def delete(self):
|
||||
"""Remove the instance default for ``?bucket=``, or for every bucket."""
|
||||
return _delete_policy("instance", None)
|
||||
|
||||
|
||||
@admin_ns.route("/admin/quotas/teams/<string:team_id>")
|
||||
class AdminTeamQuotaResource(Resource):
|
||||
@admin_required
|
||||
def get(self, team_id):
|
||||
"""The per-member allowance of one team."""
|
||||
if not looks_like_uuid(team_id):
|
||||
return _error("Team not found", 404)
|
||||
with db_readonly() as conn:
|
||||
if TeamsRepository(conn).get(team_id) is None:
|
||||
return _error("Team not found", 404)
|
||||
rows = QuotaPoliciesRepository(conn).list_for_subject("team", team_id)
|
||||
return make_response(
|
||||
jsonify({"success": True, "policies": [_policy_json(r) for r in rows]}), 200
|
||||
)
|
||||
|
||||
@admin_required
|
||||
def put(self, team_id):
|
||||
"""Set the allowance each member of the team gets."""
|
||||
if not looks_like_uuid(team_id):
|
||||
return _error("Team not found", 404)
|
||||
with db_readonly() as conn:
|
||||
if TeamsRepository(conn).get(team_id) is None:
|
||||
return _error("Team not found", 404)
|
||||
return _put_policy("team", team_id)
|
||||
|
||||
@admin_required
|
||||
def delete(self, team_id):
|
||||
"""Remove the team's allowance for ``?bucket=``, or for every bucket."""
|
||||
if not looks_like_uuid(team_id):
|
||||
return _error("Team not found", 404)
|
||||
return _delete_policy("team", team_id)
|
||||
|
||||
|
||||
@admin_ns.route("/admin/quotas/users/<string:user_id>")
|
||||
class AdminUserQuotaResource(Resource):
|
||||
@admin_required
|
||||
def get(self, user_id):
|
||||
"""A user's overrides and the limits and usage they resolve to."""
|
||||
with db_readonly() as conn:
|
||||
if UsersRepository(conn).get(user_id) is None:
|
||||
return _error("User not found", 404)
|
||||
rows = QuotaPoliciesRepository(conn).list_for_subject("user", user_id)
|
||||
statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS))
|
||||
return make_response(
|
||||
jsonify(
|
||||
{
|
||||
"success": True,
|
||||
"period": settings.QUOTA_PERIOD,
|
||||
"policies": [_policy_json(r) for r in rows],
|
||||
"effective": [s.to_dict() for s in statuses],
|
||||
}
|
||||
),
|
||||
200,
|
||||
)
|
||||
|
||||
@admin_required
|
||||
def put(self, user_id):
|
||||
"""Set one user's override, which beats team allowances and the default."""
|
||||
with db_readonly() as conn:
|
||||
if UsersRepository(conn).get(user_id) is None:
|
||||
return _error("User not found", 404)
|
||||
return _put_policy("user", user_id)
|
||||
|
||||
@admin_required
|
||||
def delete(self, user_id):
|
||||
"""Remove the user's override for ``?bucket=``, or for every bucket."""
|
||||
return _delete_policy("user", user_id)
|
||||
@@ -161,6 +161,7 @@ _CHAT = dict(
|
||||
RULES: dict[tuple[str, str], Rule] = {
|
||||
# Identity and public metadata: any valid token.
|
||||
("/api/user/me", "GET"): _rule(open=True),
|
||||
("/api/user/quota", "GET"): _rule(open=True),
|
||||
("/api/health", "GET"): _rule(open=True),
|
||||
("/api/config", "GET"): _rule(open=True),
|
||||
# Agents
|
||||
|
||||
@@ -5,6 +5,9 @@ only from ``request.decoded_token`` (already populated and role-resolved by the
|
||||
auth chokepoint in ``app.py``). Auth-mode-agnostic. ``email``/``name``/
|
||||
``picture`` are OIDC-only and optional — they are echoed from the token and are
|
||||
never present for ``simple_jwt``/``session_jwt``/no-auth modes.
|
||||
|
||||
``GET /api/user/quota`` returns the caller's usage against the limits an admin
|
||||
set for them, without naming the policies behind those limits.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -13,6 +16,8 @@ from flask import jsonify, make_response, request
|
||||
from flask_restx import Namespace, Resource
|
||||
|
||||
from docsgpt.api.pat.tokens import is_pat
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService
|
||||
|
||||
me_ns = Namespace("me", description="Current user identity and roles", path="/api")
|
||||
|
||||
@@ -43,3 +48,34 @@ class MeResource(Resource):
|
||||
"resource_filter": decoded_token.get("resource_filter") or {},
|
||||
}
|
||||
return make_response(jsonify(body), 200)
|
||||
|
||||
|
||||
def _own_budget(budget: dict) -> dict:
|
||||
return {"limit": budget["limit"], "used": budget["used"]}
|
||||
|
||||
|
||||
@me_ns.route("/user/quota")
|
||||
class MyQuotaResource(Resource):
|
||||
def get(self):
|
||||
"""Return the caller's limited buckets: ``{bucket, tokens, cost, resets_at}`` each."""
|
||||
decoded_token = getattr(request, "decoded_token", None)
|
||||
user_id = decoded_token.get("sub") if decoded_token else None
|
||||
if not user_id:
|
||||
return make_response(jsonify({"success": False}), 401)
|
||||
statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS))
|
||||
buckets = []
|
||||
for status in statuses:
|
||||
if status.limits.unlimited:
|
||||
continue
|
||||
data = status.to_dict()
|
||||
buckets.append(
|
||||
{
|
||||
"bucket": data["bucket"],
|
||||
"tokens": _own_budget(data["tokens"]),
|
||||
"cost": _own_budget(data["cost"]),
|
||||
"resets_at": data["resets_at"],
|
||||
}
|
||||
)
|
||||
return make_response(
|
||||
jsonify({"success": True, "period": settings.QUOTA_PERIOD, "buckets": buckets}), 200
|
||||
)
|
||||
@@ -165,6 +165,28 @@ class TokenUsageRepository:
|
||||
).one()
|
||||
return int(row[0]), float(row[1])
|
||||
|
||||
def tokens_by_model(self, *, start: datetime) -> list[dict]:
|
||||
"""Return ``{model_id, tokens, cost}`` per model since ``start``, busiest first."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT model_id,
|
||||
COALESCE(SUM(prompt_tokens + generated_tokens), 0) AS tokens,
|
||||
COALESCE(SUM(cost), 0) AS cost
|
||||
FROM token_usage
|
||||
WHERE timestamp >= :start AND model_id IS NOT NULL
|
||||
AND source <> ALL(:rollup_sources)
|
||||
GROUP BY model_id
|
||||
ORDER BY tokens DESC, model_id
|
||||
"""
|
||||
),
|
||||
{"start": start, "rollup_sources": list(self.ROLLUP_SOURCES)},
|
||||
)
|
||||
return [
|
||||
{"model_id": row[0], "tokens": int(row[1]), "cost": float(row[2])}
|
||||
for row in result.fetchall()
|
||||
]
|
||||
|
||||
# Token usage written outside a user-initiated request (conversation
|
||||
# title generation, history compression, RAG question condensing,
|
||||
# provider fallback). Mirrors the exclusion list in ``count_in_range``.
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Endpoint tests for the admin quota API and ``GET /api/user/quota``.
|
||||
|
||||
Driven through the real app.py chokepoint against an ephemeral Postgres; only
|
||||
``handle_auth`` / ``resolve_roles`` are patched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository
|
||||
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from docsgpt.storage.db.repositories.users import UsersRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
from docsgpt.app import app as flask_app
|
||||
|
||||
flask_app.config["TESTING"] = True
|
||||
return flask_app.test_client()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db(pg_conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with ExitStack() as stack:
|
||||
for target in (
|
||||
"docsgpt.api.admin.quotas.db_readonly",
|
||||
"docsgpt.api.admin.quotas.db_session",
|
||||
"docsgpt.quotas.service.db_readonly",
|
||||
):
|
||||
stack.enter_context(patch(target, _yield))
|
||||
yield pg_conn
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _as(sub, *roles):
|
||||
with patch("docsgpt.app.handle_auth", return_value={"sub": sub}), patch(
|
||||
"docsgpt.app.resolve_roles", return_value=list(roles) or ["user"]
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _admin():
|
||||
return _as("admin1", "admin", "user")
|
||||
|
||||
|
||||
def _body(resp):
|
||||
return json.loads(resp.data)
|
||||
|
||||
|
||||
def _team(conn, slug="q-team", member=None):
|
||||
team_id = str(
|
||||
conn.execute(
|
||||
text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'o') RETURNING id"),
|
||||
{"n": slug, "s": slug},
|
||||
).scalar()
|
||||
)
|
||||
if member:
|
||||
conn.execute(
|
||||
text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), :u, 'team_member')"),
|
||||
{"t": team_id, "u": member},
|
||||
)
|
||||
return team_id
|
||||
|
||||
|
||||
ADMIN_ROUTES = [
|
||||
("get", "/api/admin/quotas"),
|
||||
("put", "/api/admin/quotas/instance"),
|
||||
("delete", "/api/admin/quotas/instance"),
|
||||
("get", "/api/admin/quotas/users/u1"),
|
||||
("put", "/api/admin/quotas/users/u1"),
|
||||
("delete", "/api/admin/quotas/users/u1"),
|
||||
("get", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"),
|
||||
("put", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"),
|
||||
("delete", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"),
|
||||
]
|
||||
|
||||
|
||||
class TestGuard:
|
||||
@pytest.mark.parametrize("method, path", ADMIN_ROUTES)
|
||||
def test_non_admin_forbidden(self, client, db, method, path):
|
||||
with _as("u1"):
|
||||
assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 403
|
||||
|
||||
@pytest.mark.parametrize("method, path", ADMIN_ROUTES)
|
||||
def test_unauthenticated(self, client, method, path):
|
||||
with patch("docsgpt.app.handle_auth", return_value=None):
|
||||
assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 401
|
||||
|
||||
def test_team_admin_cannot_set_their_teams_allowance(self, client, db):
|
||||
team_id = _team(db)
|
||||
db.execute(
|
||||
text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), 'lead', 'team_admin')"),
|
||||
{"t": team_id},
|
||||
)
|
||||
with _as("lead"):
|
||||
resp = client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_unlimited": True})
|
||||
assert resp.status_code == 403
|
||||
assert QuotaPoliciesRepository(db).get("team", team_id) is None
|
||||
|
||||
|
||||
class TestInstancePolicy:
|
||||
def test_set_read_delete(self, client, db):
|
||||
with _admin():
|
||||
put = client.put("/api/admin/quotas/instance", json={"token_limit": 1000, "note": " default "})
|
||||
assert put.status_code == 200
|
||||
policy = _body(put)["policy"]
|
||||
assert (policy["token_limit"], policy["note"], policy["updated_by"]) == (1000, "default", "admin1")
|
||||
|
||||
client.put("/api/admin/quotas/instance", json={"bucket": "agent", "cost_limit_usd": 2.5})
|
||||
overview = _body(client.get("/api/admin/quotas"))
|
||||
assert [(p["bucket"], p["token_limit"], p["cost_limit_usd"]) for p in overview["instance"]] == [
|
||||
("all", 1000, None),
|
||||
("agent", None, 2.5),
|
||||
]
|
||||
assert overview["period"] == "month"
|
||||
|
||||
assert _body(client.delete("/api/admin/quotas/instance?bucket=agent"))["deleted"] == 1
|
||||
assert _body(client.delete("/api/admin/quotas/instance"))["deleted"] == 1
|
||||
assert _body(client.get("/api/admin/quotas"))["instance"] == []
|
||||
|
||||
def test_writes_are_audited(self, client, db):
|
||||
with _admin():
|
||||
client.put("/api/admin/quotas/instance", json={"token_limit": 5})
|
||||
client.delete("/api/admin/quotas/instance")
|
||||
client.delete("/api/admin/quotas/instance")
|
||||
events = db.execute(
|
||||
text("SELECT user_id, event, metadata FROM auth_events WHERE event LIKE 'quota_policy_%'")
|
||||
).fetchall()
|
||||
by_event = {e[1]: e for e in events}
|
||||
# The second delete removed nothing, so it left no event.
|
||||
assert sorted((e[0], e[1]) for e in events) == [
|
||||
("admin1", "quota_policy_deleted"),
|
||||
("admin1", "quota_policy_set"),
|
||||
]
|
||||
metadata = by_event["quota_policy_set"][2]
|
||||
assert metadata["token_limit"] == 5 and metadata["by"] == "admin1"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
None,
|
||||
[],
|
||||
{},
|
||||
{"note": "only a note"},
|
||||
{"token_limit": -1},
|
||||
{"token_limit": 1.5},
|
||||
{"token_limit": True},
|
||||
{"token_limit": "10"},
|
||||
{"token_limit": 2**63},
|
||||
{"cost_limit_usd": -0.01},
|
||||
{"cost_limit_usd": "5"},
|
||||
{"cost_limit_usd": float("inf")},
|
||||
{"cost_limit_usd": 1e12},
|
||||
{"token_limit": 1, "token_unlimited": True},
|
||||
{"cost_limit_usd": 1, "cost_unlimited": True},
|
||||
{"token_unlimited": "yes"},
|
||||
{"token_limit": 1, "enabled": "no"},
|
||||
{"token_limit": 1, "bucket": "everything"},
|
||||
{"token_limit": 1, "note": 7},
|
||||
],
|
||||
)
|
||||
def test_invalid_bodies_rejected(self, client, db, body):
|
||||
with _admin():
|
||||
resp = client.put("/api/admin/quotas/instance", json=body)
|
||||
assert resp.status_code == 400
|
||||
assert QuotaPoliciesRepository(db).list_by_scope("instance") == []
|
||||
|
||||
def test_unknown_bucket_on_delete(self, client, db):
|
||||
with _admin():
|
||||
assert client.delete("/api/admin/quotas/instance?bucket=nope").status_code == 400
|
||||
|
||||
def test_zero_is_accepted_as_a_block(self, client, db):
|
||||
with _admin():
|
||||
resp = client.put("/api/admin/quotas/instance", json={"token_limit": 0, "cost_limit_usd": 0})
|
||||
assert resp.status_code == 200
|
||||
assert _body(resp)["policy"]["token_limit"] == 0
|
||||
|
||||
|
||||
class TestTeamPolicy:
|
||||
def test_set_and_list_with_team_details(self, client, db):
|
||||
team_id = _team(db, "q-eng", member="u1")
|
||||
with _admin():
|
||||
assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 500}).status_code == 200
|
||||
(row,) = _body(client.get("/api/admin/quotas"))["teams"]
|
||||
assert (row["team_slug"], row["token_limit"], row["member_count"]) == ("q-eng", 500, 1)
|
||||
assert _body(client.get(f"/api/admin/quotas/teams/{team_id}"))["policies"][0]["token_limit"] == 500
|
||||
|
||||
@pytest.mark.parametrize("team_id", ["not-a-uuid", "00000000-0000-0000-0000-000000000000"])
|
||||
def test_unknown_team(self, client, db, team_id):
|
||||
with _admin():
|
||||
assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 1}).status_code == 404
|
||||
assert client.get(f"/api/admin/quotas/teams/{team_id}").status_code == 404
|
||||
|
||||
|
||||
class TestUserPolicy:
|
||||
def test_unknown_user(self, client, db):
|
||||
with _admin():
|
||||
assert client.put("/api/admin/quotas/users/ghost", json={"token_limit": 1}).status_code == 404
|
||||
assert client.get("/api/admin/quotas/users/ghost").status_code == 404
|
||||
|
||||
def test_effective_limits_name_their_source(self, client, db):
|
||||
UsersRepository(db).upsert("u1")
|
||||
small, big = _team(db, "q-small", member="u1"), _team(db, "q-big", member="u1")
|
||||
repo = QuotaPoliciesRepository(db)
|
||||
repo.upsert(scope="instance", subject_id=None, token_limit=10, cost_limit_usd=1)
|
||||
repo.upsert(scope="team", subject_id=small, token_limit=100)
|
||||
repo.upsert(scope="team", subject_id=big, token_limit=900)
|
||||
TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=40, cost=0.25)
|
||||
|
||||
with _admin():
|
||||
body = _body(client.get("/api/admin/quotas/users/u1"))
|
||||
overall = body["effective"][0]
|
||||
assert overall["bucket"] == "all"
|
||||
assert overall["tokens"] == {"limit": 900.0, "used": 40, "source": "team", "source_id": big}
|
||||
assert overall["cost"] == {"limit": 1.0, "used": 0.25, "source": "instance", "source_id": None}
|
||||
assert body["policies"] == []
|
||||
|
||||
with _admin():
|
||||
client.put("/api/admin/quotas/users/u1", json={"token_limit": 50})
|
||||
body = _body(client.get("/api/admin/quotas/users/u1"))
|
||||
assert body["effective"][0]["tokens"]["source"] == "user"
|
||||
assert body["policies"][0]["token_limit"] == 50
|
||||
|
||||
def test_user_policy_audit_is_filed_under_the_user(self, client, db):
|
||||
UsersRepository(db).upsert("u1")
|
||||
with _admin():
|
||||
client.put("/api/admin/quotas/users/u1", json={"cost_unlimited": True})
|
||||
row = db.execute(
|
||||
text("SELECT user_id, metadata FROM auth_events WHERE event = 'quota_policy_set'")
|
||||
).one()
|
||||
assert row[0] == "u1" and row[1]["by"] == "admin1" and row[1]["scope"] == "user"
|
||||
|
||||
|
||||
class TestUnpricedModels:
|
||||
def test_lists_used_catalog_models_without_a_price(self, client, db):
|
||||
usage = TokenUsageRepository(db)
|
||||
usage.insert(user_id="u1", prompt_tokens=10, model_id="local-llama")
|
||||
usage.insert(user_id="u1", prompt_tokens=5, model_id="claude-haiku-4-5", cost=0.1)
|
||||
usage.insert(user_id="u1", prompt_tokens=7, model_id="7d0c1a52-2f5e-4c53-9a0e-111111111111")
|
||||
with _admin(), patch("docsgpt.api.admin.quotas.is_priced", lambda m: m == "claude-haiku-4-5"):
|
||||
unpriced = _body(client.get("/api/admin/quotas"))["unpriced_models"]
|
||||
assert unpriced == [{"model_id": "local-llama", "tokens": 10, "cost": 0.0}]
|
||||
|
||||
|
||||
class TestMyQuota:
|
||||
def test_unlimited_user_sees_no_buckets(self, client, db):
|
||||
with _as("u1"):
|
||||
body = _body(client.get("/api/user/quota"))
|
||||
assert body == {"success": True, "period": "month", "buckets": []}
|
||||
|
||||
def test_limited_user_sees_usage_without_policy_internals(self, client, db):
|
||||
QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100, note="secret note")
|
||||
TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=30)
|
||||
with _as("u1"):
|
||||
body = _body(client.get("/api/user/quota"))
|
||||
(bucket,) = body["buckets"]
|
||||
assert bucket["bucket"] == "all"
|
||||
assert bucket["tokens"] == {"limit": 100.0, "used": 30}
|
||||
assert bucket["cost"] == {"limit": None, "used": 0.0}
|
||||
assert "source" not in json.dumps(body) and "secret" not in json.dumps(body)
|
||||
|
||||
def test_a_user_only_sees_their_own_quota(self, client, db):
|
||||
QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100)
|
||||
with _as("u2"):
|
||||
assert _body(client.get("/api/user/quota"))["buckets"] == []
|
||||
|
||||
def test_unauthenticated(self, client):
|
||||
with patch("docsgpt.app.handle_auth", return_value=None):
|
||||
assert client.get("/api/user/quota").status_code == 401
|
||||
Reference in new issue
Block a user