mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-20 02:16:39 +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
@@ -15,6 +15,7 @@ from docsgpt.api.answer.services.prompt_renderer import (
|
||||
)
|
||||
from docsgpt.api.answer.services.stream_processor import get_prompt
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.quotas.service import QuotaExceededError, QuotaService
|
||||
from docsgpt.retriever.retriever_creator import RetrieverCreator
|
||||
from docsgpt.storage.db.repositories.sources import SourcesRepository
|
||||
from docsgpt.storage.db.session import db_readonly
|
||||
@@ -69,7 +70,11 @@ def run_agent_headless(
|
||||
chat_history: Optional[List[Dict[str, Any]]] = None,
|
||||
conversation_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run an agent with no live client; returns a structured outcome dict."""
|
||||
"""Run an agent with no live client; returns a structured outcome dict.
|
||||
|
||||
Raises:
|
||||
QuotaExceededError: If the agent owner's usage quota is exhausted.
|
||||
"""
|
||||
from docsgpt.core.model_utils import (
|
||||
get_api_key_for_provider,
|
||||
get_default_model_id,
|
||||
@@ -82,6 +87,9 @@ def run_agent_headless(
|
||||
if not owner:
|
||||
raise ValueError("Agent config is missing user_id; cannot run headless.")
|
||||
decoded_token = {"sub": owner}
|
||||
exceeded = QuotaService.check(owner, "agent" if agent_config.get("key") else "direct")
|
||||
if exceeded is not None:
|
||||
raise QuotaExceededError(exceeded)
|
||||
|
||||
retriever_kind = agent_config.get("retriever", "classic")
|
||||
source_id = agent_config.get("source_id") or agent_config.get("source")
|
||||
|
||||
Reference in new issue
Block a user