mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 12:13:38 +00:00
feat(quotas): enforce user quotas on chat, agent, scheduled and webhook runs
check_usage now checks the billable user's quota on every request, before the per-agent 24h limits, which keep applying to traffic through an agent. Until now a request without an agent key skipped every limit. A refusal is a 429 with Retry-After and a body naming the budget, usage, limit, the layer the limit came from and when it resets. Headless runs check the agent owner's quota before starting. A refused scheduled run is recorded as budget_exceeded; a refused webhook run returns a quota_exceeded result instead of raising, so Celery does not retry it.
This commit is contained in:
1 parent
6826313b60
commit
5406550bca
13 files changed
+238
-9
No files matched your search
@@ -103,7 +103,9 @@ class AnswerResource(Resource, BaseAnswerResource):
|
||||
)
|
||||
if not processor.decoded_token:
|
||||
return make_response({"error": "Unauthorized"}, 401)
|
||||
if error := self.check_usage(processor.agent_config):
|
||||
if error := self.check_usage(
|
||||
processor.agent_config, processor.decoded_token
|
||||
):
|
||||
return error
|
||||
stream = self.complete_stream(
|
||||
question="",
|
||||
@@ -129,7 +131,9 @@ class AnswerResource(Resource, BaseAnswerResource):
|
||||
if not processor.decoded_token:
|
||||
return make_response({"error": "Unauthorized"}, 401)
|
||||
|
||||
if error := self.check_usage(processor.agent_config):
|
||||
if error := self.check_usage(
|
||||
processor.agent_config, processor.decoded_token
|
||||
):
|
||||
return error
|
||||
|
||||
should_persist, visibility = resolve_persistence(
|
||||
|
||||
@@ -23,6 +23,8 @@ from docsgpt.core.model_utils import (
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.error import sanitize_api_error
|
||||
from docsgpt.llm.llm_creator import LLMCreator
|
||||
from docsgpt.quotas.http import quota_exceeded_response
|
||||
from docsgpt.quotas.service import QuotaService
|
||||
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
||||
from docsgpt.storage.db.repositories.conversations import (
|
||||
HeartbeatState,
|
||||
@@ -112,17 +114,29 @@ class BaseAnswerResource:
|
||||
prepared.append(item)
|
||||
return prepared
|
||||
|
||||
def check_usage(self, agent_config: Dict) -> Optional[Response]:
|
||||
"""Check if there is a usage limit and if it is exceeded
|
||||
def check_usage(
|
||||
self, agent_config: Dict, decoded_token: Optional[Dict] = None
|
||||
) -> Optional[Response]:
|
||||
"""Refuse the request when a usage limit is exhausted.
|
||||
|
||||
The billable user's quota is checked first, for every request; the
|
||||
agent's own 24h token and request limits then apply to traffic that
|
||||
runs through an agent.
|
||||
|
||||
Args:
|
||||
agent_config: The config dict of agent instance
|
||||
decoded_token: The request's resolved identity; its ``sub`` is the
|
||||
billable user.
|
||||
|
||||
Returns:
|
||||
None or Response if either of limits exceeded.
|
||||
|
||||
"""
|
||||
api_key = agent_config.get("user_api_key")
|
||||
user_id = (decoded_token or {}).get("sub") or agent_config.get("user_id")
|
||||
exceeded = QuotaService.check(user_id, "agent" if api_key else "direct")
|
||||
if exceeded is not None:
|
||||
return quota_exceeded_response(exceeded)
|
||||
if not api_key:
|
||||
return None
|
||||
with db_readonly() as conn:
|
||||
|
||||
@@ -115,7 +115,9 @@ class StreamResource(Resource, BaseAnswerResource):
|
||||
status=401,
|
||||
mimetype="text/event-stream",
|
||||
)
|
||||
if error := self.check_usage(processor.agent_config):
|
||||
if error := self.check_usage(
|
||||
processor.agent_config, processor.decoded_token
|
||||
):
|
||||
return error
|
||||
return Response(
|
||||
with_sse_keepalive(
|
||||
@@ -151,7 +153,9 @@ class StreamResource(Resource, BaseAnswerResource):
|
||||
mimetype="text/event-stream",
|
||||
)
|
||||
|
||||
if error := self.check_usage(processor.agent_config):
|
||||
if error := self.check_usage(
|
||||
processor.agent_config, processor.decoded_token
|
||||
):
|
||||
return error
|
||||
should_persist, visibility = resolve_persistence(
|
||||
visibility_flag=data.get("visibility"),
|
||||
|
||||
Reference in new issue
Block a user