feat: Add quota check decorator and implement tenant-level API rate limiting

- Add quota check decorator module quota_stub.py, providing community edition stub implementation
- Add quota check decorators to multiple controllers
- Implement tenant-level API call rate limiting
- Remove redundant plan fields from tenant_model.py
- Optimize user permission check logic with added error handling
This commit is contained in:
wxy
2026-04-13 11:58:14 +08:00
parent 807dee8460
commit 72be9f75f9
15 changed files with 140 additions and 10 deletions

View File

@@ -28,6 +28,7 @@ from app.services.app_statistics_service import AppStatisticsService
from app.services.workflow_import_service import WorkflowImportService
from app.services.workflow_service import WorkflowService, get_workflow_service
from app.services.app_dsl_service import AppDslService
from app.core.quota_stub import check_app_quota
router = APIRouter(prefix="/apps", tags=["Apps"])
logger = get_business_logger()
@@ -35,6 +36,7 @@ logger = get_business_logger()
@router.post("", summary="创建应用(可选创建 Agent 配置)")
@cur_workspace_access_guard()
@check_app_quota
def create_app(
payload: app_schema.AppCreate,
db: Session = Depends(get_db),

View File

@@ -19,6 +19,7 @@ from app.models.user_model import User
from app.schemas import file_schema, document_schema
from app.schemas.response_schema import ApiResponse
from app.services import file_service, document_service
from app.core.quota_stub import check_knowledge_capacity_quota
# Obtain a dedicated API logger
@@ -131,6 +132,7 @@ async def create_folder(
@router.post("/file", response_model=ApiResponse)
@check_knowledge_capacity_quota
async def upload_file(
kb_id: uuid.UUID,
parent_id: uuid.UUID,

View File

@@ -27,6 +27,7 @@ from app.schemas import knowledge_schema
from app.schemas.response_schema import ApiResponse
from app.services import knowledge_service, document_service
from app.services.model_service import ModelConfigService
from app.core.quota_stub import check_knowledge_capacity_quota
# Obtain a dedicated API logger
api_logger = get_api_logger()
@@ -179,6 +180,7 @@ async def get_knowledges(
@router.post("/knowledge", response_model=ApiResponse)
@check_knowledge_capacity_quota
async def create_knowledge(
create_data: knowledge_schema.KnowledgeCreate,
db: Session = Depends(get_db),

View File

@@ -34,6 +34,7 @@ from app.services.memory_storage_service import (
search_entity,
search_statement,
)
from app.core.quota_stub import check_memory_engine_quota
from fastapi import APIRouter, Depends, Header
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
@@ -76,6 +77,7 @@ async def get_storage_info(
@router.post("/create_config", response_model=ApiResponse) # 创建配置文件,其他参数默认
@check_memory_engine_quota
def create_config(
payload: ConfigParamsCreate,
current_user: User = Depends(get_current_user),

View File

@@ -15,6 +15,7 @@ from app.core.response_utils import success
from app.schemas.response_schema import ApiResponse, PageData
from app.services.model_service import ModelConfigService, ModelApiKeyService, ModelBaseService
from app.core.logging_config import get_api_logger
from app.core.quota_stub import check_model_quota, check_model_activation_quota
# 获取API专用日志器
api_logger = get_api_logger()
@@ -236,6 +237,7 @@ def delete_model_base(
@router.post("/model_plaza/{model_base_id}/add", response_model=ApiResponse)
@check_model_quota
def add_model_from_plaza(
model_base_id: uuid.UUID,
db: Session = Depends(get_db),
@@ -273,6 +275,7 @@ def get_model_by_id(
@router.post("", response_model=ApiResponse)
@check_model_quota
async def create_model(
model_data: model_schema.ModelConfigCreate,
db: Session = Depends(get_db),
@@ -303,6 +306,7 @@ async def create_model(
@router.post("/composite", response_model=ApiResponse)
@check_model_quota
async def create_composite_model(
model_data: model_schema.CompositeModelCreate,
db: Session = Depends(get_db),
@@ -329,6 +333,7 @@ async def create_composite_model(
@router.put("/composite/{model_id}", response_model=ApiResponse)
@check_model_activation_quota
async def update_composite_model(
model_id: uuid.UUID,
model_data: model_schema.CompositeModelCreate,
@@ -370,6 +375,7 @@ def delete_composite_model(
@router.put("/{model_id}", response_model=ApiResponse)
@check_model_activation_quota
def update_model(
model_id: uuid.UUID,
model_data: model_schema.ModelConfigUpdate,

View File

@@ -28,6 +28,8 @@ from fastapi import APIRouter, Depends, HTTPException, File, UploadFile, Form, H
from fastapi.responses import StreamingResponse, JSONResponse
from sqlalchemy.orm import Session
from app.core.quota_stub import check_ontology_project_quota
from app.core.config import settings
from app.core.error_codes import BizCode
from app.core.language_utils import get_language_from_header
@@ -287,6 +289,7 @@ async def extract_ontology(
# ==================== 本体场景管理接口 ====================
@router.post("/scene", response_model=ApiResponse)
@check_ontology_project_quota
async def create_scene(
request: SceneCreateRequest,
db: Session = Depends(get_db),

View File

@@ -9,6 +9,7 @@ from app.core.api_key_auth import require_api_key
from app.core.error_codes import BizCode
from app.core.exceptions import BusinessException
from app.core.logging_config import get_business_logger
from app.core.quota_stub import check_end_user_quota
from app.core.response_utils import success
from app.db import get_db
from app.repositories.end_user_repository import EndUserRepository
@@ -22,6 +23,7 @@ logger = get_business_logger()
@router.post("/create")
@require_api_key(scopes=["memory"])
@check_end_user_quota
async def create_end_user(
request: Request,
api_key_auth: ApiKeyAuth = None,

View File

@@ -2,6 +2,7 @@
from app.core.api_key_auth import require_api_key
from app.core.logging_config import get_business_logger
from app.core.quota_stub import check_end_user_quota
from app.core.response_utils import success
from app.db import get_db
from app.schemas.api_key_schema import ApiKeyAuth
@@ -119,6 +120,7 @@ async def list_memory_configs(
@router.post("/end_users")
@require_api_key(scopes=["memory"])
@check_end_user_quota
async def create_end_user(
request: Request,
api_key_auth: ApiKeyAuth = None,

View File

@@ -11,11 +11,13 @@ from app.schemas import skill_schema
from app.schemas.response_schema import PageData, PageMeta
from app.services.skill_service import SkillService
from app.core.response_utils import success
from app.core.quota_stub import check_skill_quota
router = APIRouter(prefix="/skills", tags=["Skills"])
@router.post("", summary="创建技能")
@check_skill_quota
def create_skill(
data: skill_schema.SkillCreate,
db: Session = Depends(get_db),

View File

@@ -114,11 +114,14 @@ def get_current_user_info(
# 设置权限:如果用户来自 SSO Source则使用该 Source 的 permissions否则返回 "all" 表示拥有所有权限
if current_user.external_source:
from premium.sso.models import SSOSource
source = db.query(SSOSource).filter(SSOSource.source_code == current_user.external_source).first()
if source and source.permissions:
result_schema.permissions = source.permissions
else:
try:
from premium.sso.models import SSOSource
source = db.query(SSOSource).filter(SSOSource.source_code == current_user.external_source).first()
if source and source.permissions:
result_schema.permissions = source.permissions
else:
result_schema.permissions = []
except ModuleNotFoundError:
result_schema.permissions = []
else:
result_schema.permissions = ["all"]

View File

@@ -35,6 +35,7 @@ from app.schemas.workspace_schema import (
WorkspaceUpdate,
)
from app.services import workspace_service
from app.core.quota_stub import check_workspace_quota
# 获取API专用日志器
api_logger = get_api_logger()
@@ -106,6 +107,7 @@ def get_workspaces(
@router.post("", response_model=ApiResponse)
@check_workspace_quota
def create_workspace(
workspace: WorkspaceCreate,
language_type: str = Header(default="zh", alias="X-Language-Type"),