LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 08:33:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feat: Agent Token or Request Limiting (#2041)

Browse Source 
* Update routes.py, added token and request limits to create/update agent operations

* added usage limit check to api endpoints

cannot create agents with usage limit right now that will be implemented

* implemented api limiting as either token limiting or request limiting modes

* minor typo & bug fix
This commit is contained in:
Ali Arda Fincan
2025-10-13 21:32:46 +03:00
committed by GitHub
parent 72bc24a490
commit ce32dd2907
9 changed files with 352 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
application/api/user/agents/routes.py
View File

@@ -10,6 +10,7 @@ from flask import current_app, jsonify, make_response, request
from flask_restx import fields, Namespace, Resource
from application.api import api
from application.core.settings import settings
from application.api.user.base import (
agents_collection,
db,
@@ -74,6 +75,10 @@ class GetAgent(Resource):
"agent_type": agent.get("agent_type", ""),
"status": agent.get("status", ""),
"json_schema": agent.get("json_schema"),
"limited_token_mode": agent.get("limited_token_mode", False),
"token_limit": agent.get("token_limit", settings.DEFAULT_AGENT_LIMITS["token_limit"]),
"limited_request_mode": agent.get("limited_request_mode", False),
"request_limit": agent.get("request_limit", settings.DEFAULT_AGENT_LIMITS["request_limit"]),
"created_at": agent.get("createdAt", ""),
"updated_at": agent.get("updatedAt", ""),
"last_used_at": agent.get("lastUsedAt", ""),
@@ -143,6 +148,10 @@ class GetAgents(Resource):
"agent_type": agent.get("agent_type", ""),
"status": agent.get("status", ""),
"json_schema": agent.get("json_schema"),
"limited_token_mode": agent.get("limited_token_mode", False),
"token_limit": agent.get("token_limit", settings.DEFAULT_AGENT_LIMITS["token_limit"]),
"limited_request_mode": agent.get("limited_request_mode", False),
"request_limit": agent.get("request_limit", settings.DEFAULT_AGENT_LIMITS["request_limit"]),
"created_at": agent.get("createdAt", ""),
"updated_at": agent.get("updatedAt", ""),
"last_used_at": agent.get("lastUsedAt", ""),
@@ -199,6 +208,22 @@ class CreateAgent(Resource):
required=False,
description="JSON schema for enforcing structured output format",
),
"limited_token_mode": fields.Boolean(
required=False,
description="Whether the agent is in limited token mode"
),
"token_limit": fields.Integer(
required=False,
description="Token limit for the agent in limited mode"
),
"limited_request_mode": fields.Boolean(
required=False,
description="Whether the agent is in limited request mode"
),
"request_limit": fields.Integer(
required=False,
description="Request limit for the agent in limited mode"
)
},
)
@@ -344,6 +369,10 @@ class CreateAgent(Resource):
"agent_type": data.get("agent_type", ""),
"status": data.get("status"),
"json_schema": data.get("json_schema"),
"limited_token_mode": data.get("limited_token_mode", False),
"token_limit": data.get("token_limit", settings.DEFAULT_AGENT_LIMITS["token_limit"]),
"limited_request_mode": data.get("limited_request_mode", False),
"request_limit": data.get("request_limit", settings.DEFAULT_AGENT_LIMITS["request_limit"]),
"createdAt": datetime.datetime.now(datetime.timezone.utc),
"updatedAt": datetime.datetime.now(datetime.timezone.utc),
"lastUsedAt": None,
@@ -399,6 +428,22 @@ class UpdateAgent(Resource):
required=False,
description="JSON schema for enforcing structured output format",
),
"limited_token_mode": fields.Boolean(
required=False,
description="Whether the agent is in limited token mode"
),
"token_limit": fields.Integer(
required=False,
description="Token limit for the agent in limited mode"
),
"limited_request_mode": fields.Boolean(
require=False,
description="Whether the agent is in limited request mode"
),
"request_limit": fields.Integer(
required=False,
description="Request limit for the agent in limited mode"
)
},
)
@@ -486,6 +531,10 @@ class UpdateAgent(Resource):
"agent_type",
"status",
"json_schema",
"limited_token_mode",
"token_limit",
"limited_request_mode",
"request_limit"
]
for field in allowed_fields:
@@ -602,6 +651,54 @@ class UpdateAgent(Resource):
update_fields[field] = json_schema
else:
update_fields[field] = None
elif field == "limited_token_mode":
is_mode_enabled = data.get("limited_token_mode", False)
if is_mode_enabled and data.get("token_limit") is None:
return make_response(
jsonify(
{
"success": False,
"message": "Token limit must be provided when limited token mode is enabled",
}
),
400,
)
elif field == "limited_request_mode":
is_mode_enabled = data.get("limited_request_mode", False)
if is_mode_enabled and data.get("request_limit") is None:
return make_response(
jsonify(
{
"success": False,
"message": "Request limit must be provided when limited request mode is enabled",
}
),
400,
)
elif field == "token_limit":
token_limit = data.get("token_limit")
if token_limit is not None and not data.get("limited_token_mode"):
return make_response(
jsonify(
{
"success": False,
"message": "Token limit cannot be set when limited token mode is disabled",
}
),
400,
)
elif field == "request_limit":
request_limit = data.get("request_limit")
if request_limit is not None and not data.get("limited_request_mode"):
return make_response(
jsonify(
{
"success": False,
"message": "Request limit cannot be set when limited request mode is disabled",
}
),
400,
)
else:
value = data[field]
if field in ["name", "description", "prompt_id", "agent_type"]:

9
application/api/user/agents/sharing.py
View File

@@ -9,6 +9,7 @@ from flask import current_app, jsonify, make_response, request
from flask_restx import fields, Namespace, Resource
from application.api import api
from application.core.settings import settings
from application.api.user.base import (
agents_collection,
db,
@@ -75,6 +76,10 @@ class SharedAgent(Resource):
"agent_type": shared_agent.get("agent_type", ""),
"status": shared_agent.get("status", ""),
"json_schema": shared_agent.get("json_schema"),
"limited_token_mode": shared_agent.get("limited_token_mode", False),
"token_limit": shared_agent.get("token_limit", settings.DEFAULT_AGENT_LIMITS["token_limit"]),
"limited_request_mode": shared_agent.get("limited_request_mode", False),
"request_limit": shared_agent.get("request_limit", settings.DEFAULT_AGENT_LIMITS["request_limit"]),
"created_at": shared_agent.get("createdAt", ""),
"updated_at": shared_agent.get("updatedAt", ""),
"shared": shared_agent.get("shared_publicly", False),
@@ -149,6 +154,10 @@ class SharedAgents(Resource):
"agent_type": agent.get("agent_type", ""),
"status": agent.get("status", ""),
"json_schema": agent.get("json_schema"),
"limited_token_mode": agent.get("limited_token_mode", False),
"token_limit": agent.get("token_limit", settings.DEFAULT_AGENT_LIMITS["token_limit"]),
"limited_request_mode": agent.get("limited_request_mode", False),
"request_limit": agent.get("request_limit", settings.DEFAULT_AGENT_LIMITS["request_limit"]),
"created_at": agent.get("createdAt", ""),
"updated_at": agent.get("updatedAt", ""),
"pinned": str(agent["_id"]) in pinned_ids,