mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 12:13:38 +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
+24
-1
@@ -2,6 +2,7 @@ import logging
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
|
||||
from docsgpt.pricing import compute_cost_usd
|
||||
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from docsgpt.storage.db.session import db_session
|
||||
from docsgpt.utils import num_tokens_from_object_or_list, num_tokens_from_string
|
||||
@@ -119,6 +120,12 @@ def _persist_call_usage(llm, call_usage):
|
||||
},
|
||||
)
|
||||
return
|
||||
model_id = getattr(llm, "_canonical_model_id", None)
|
||||
# Bring-your-own models run on the user's own provider key: recorded, never priced.
|
||||
if getattr(llm, "_is_byom", False):
|
||||
cost = 0.0
|
||||
else:
|
||||
cost = _call_cost_usd(model_id, call_usage)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
# ``timestamp`` is omitted so Postgres ``server_default
|
||||
@@ -136,16 +143,32 @@ def _persist_call_usage(llm, call_usage):
|
||||
# "0% cache hits".
|
||||
cached_tokens=call_usage.get("cached_tokens"),
|
||||
cache_write_tokens=call_usage.get("cache_write_tokens"),
|
||||
cost=cost,
|
||||
source=(
|
||||
getattr(llm, "_token_usage_source", None) or "agent_stream"
|
||||
),
|
||||
request_id=getattr(llm, "_request_id", None),
|
||||
model_id=getattr(llm, "_canonical_model_id", None),
|
||||
model_id=model_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("token_usage persist failed")
|
||||
|
||||
|
||||
def _call_cost_usd(model_id, call_usage) -> float:
|
||||
"""Price one call; a pricing failure records $0 rather than dropping the row."""
|
||||
try:
|
||||
return compute_cost_usd(
|
||||
model_id,
|
||||
call_usage["prompt_tokens"],
|
||||
call_usage["generated_tokens"],
|
||||
cached_tokens=call_usage.get("cached_tokens"),
|
||||
cache_write_tokens=call_usage.get("cache_write_tokens"),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("token_usage cost computation failed")
|
||||
return 0.0
|
||||
|
||||
|
||||
def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, int]:
|
||||
"""Replace estimates with upstream counts when a provider reported them.
|
||||
|
||||
|
||||
Reference in new issue
Block a user