mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 18:15:58 +00:00
feat(quotas): quota_policies table and a per-call cost on token_usage
Migration 0033 adds quota_policies (instance default, team per-member allowance, user override; a token budget and a USD budget per row) and token_usage.cost. The column is added IF NOT EXISTS so a database that already carries it upgrades cleanly. Every usage row now records the call's USD cost from the model catalog; bring-your-own models are recorded at $0.
This commit is contained in:
1 parent
b77561288d
commit
43fad2a865
9 files changed
+415
-3
No files matched your search
@@ -28,6 +28,7 @@ from sqlalchemy import (
|
||||
Index,
|
||||
Integer,
|
||||
MetaData,
|
||||
Numeric,
|
||||
PrimaryKeyConstraint,
|
||||
UniqueConstraint,
|
||||
Table,
|
||||
@@ -241,6 +242,9 @@ token_usage_table = Table(
|
||||
# cache activity, so hit-rate queries stay honest across providers.
|
||||
Column("cached_tokens", Integer),
|
||||
Column("cache_write_tokens", Integer),
|
||||
# Added in ``0033_quotas``. USD cost of the call at write time; 0 for
|
||||
# unpriced and bring-your-own models.
|
||||
Column("cost", Numeric(12, 8), nullable=False, server_default="0"),
|
||||
)
|
||||
|
||||
user_logs_table = Table(
|
||||
@@ -1121,3 +1125,42 @@ Index(
|
||||
personal_access_tokens_table.c.user_id,
|
||||
personal_access_tokens_table.c.created_at.desc(),
|
||||
)
|
||||
|
||||
# --- Usage quotas (migration 0033) ------------------------------------------
|
||||
# Admin-set limits at three layers: instance (``subject_id`` NULL), team
|
||||
# per-member allowance (``teams.id``) and user override (auth ``sub``). Per
|
||||
# budget a row sets a limit, marks it unlimited, or defers to the next layer.
|
||||
|
||||
quota_policies_table = Table(
|
||||
"quota_policies",
|
||||
metadata,
|
||||
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
|
||||
Column("scope", Text, nullable=False),
|
||||
Column("subject_id", Text),
|
||||
Column("bucket", Text, nullable=False, server_default="all"),
|
||||
Column("token_limit", BigInteger),
|
||||
Column("token_unlimited", Boolean, nullable=False, server_default="false"),
|
||||
Column("cost_limit_usd", Numeric(12, 4)),
|
||||
Column("cost_unlimited", Boolean, nullable=False, server_default="false"),
|
||||
Column("enabled", Boolean, nullable=False, server_default="true"),
|
||||
Column("note", Text),
|
||||
Column("created_by", Text),
|
||||
Column("updated_by", Text),
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
CheckConstraint("scope IN ('instance', 'team', 'user')", name="quota_policies_scope_check"),
|
||||
CheckConstraint("bucket IN ('all', 'direct', 'agent')", name="quota_policies_bucket_check"),
|
||||
CheckConstraint("token_limit >= 0", name="quota_policies_token_limit_check"),
|
||||
CheckConstraint("cost_limit_usd >= 0", name="quota_policies_cost_limit_check"),
|
||||
CheckConstraint("(scope = 'instance') = (subject_id IS NULL)", name="quota_policies_subject_chk"),
|
||||
CheckConstraint("NOT (token_unlimited AND token_limit IS NOT NULL)", name="quota_policies_token_chk"),
|
||||
CheckConstraint("NOT (cost_unlimited AND cost_limit_usd IS NOT NULL)", name="quota_policies_cost_chk"),
|
||||
)
|
||||
|
||||
Index(
|
||||
"quota_policies_subject_uidx",
|
||||
quota_policies_table.c.scope,
|
||||
func.coalesce(quota_policies_table.c.subject_id, ""),
|
||||
quota_policies_table.c.bucket,
|
||||
unique=True,
|
||||
)
|
||||
@@ -37,6 +37,7 @@ class TokenUsageRepository:
|
||||
timestamp: Optional[datetime] = None,
|
||||
cached_tokens: Optional[int] = None,
|
||||
cache_write_tokens: Optional[int] = None,
|
||||
cost: float = 0.0,
|
||||
) -> None:
|
||||
# Attribution guard: the ``token_usage_attribution_chk`` CHECK
|
||||
# constraint requires at least one of ``user_id`` / ``api_key``
|
||||
@@ -62,14 +63,14 @@ class TokenUsageRepository:
|
||||
INSERT INTO token_usage (
|
||||
user_id, api_key, agent_id,
|
||||
prompt_tokens, generated_tokens,
|
||||
cached_tokens, cache_write_tokens,
|
||||
cached_tokens, cache_write_tokens, cost,
|
||||
source, request_id, model_id, timestamp
|
||||
)
|
||||
VALUES (
|
||||
:user_id, :api_key,
|
||||
CAST(:agent_id AS uuid),
|
||||
:prompt_tokens, :generated_tokens,
|
||||
:cached_tokens, :cache_write_tokens,
|
||||
:cached_tokens, :cache_write_tokens, :cost,
|
||||
:source, :request_id, :model_id, COALESCE(:timestamp, now())
|
||||
)
|
||||
"""
|
||||
@@ -82,6 +83,7 @@ class TokenUsageRepository:
|
||||
"generated_tokens": generated_tokens,
|
||||
"cached_tokens": cached_tokens,
|
||||
"cache_write_tokens": cache_write_tokens,
|
||||
"cost": cost,
|
||||
"source": source,
|
||||
"request_id": request_id,
|
||||
"model_id": model_id,
|
||||
|
||||
Reference in new issue
Block a user