From 72be9f75f92ae6c2f59306bf9a4a9629e94c637a Mon Sep 17 00:00:00 2001 From: wxy Date: Mon, 13 Apr 2026 11:58:14 +0800 Subject: [PATCH 1/4] 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 --- api/app/controllers/app_controller.py | 2 + api/app/controllers/file_controller.py | 2 + api/app/controllers/knowledge_controller.py | 2 + .../controllers/memory_storage_controller.py | 2 + api/app/controllers/model_controller.py | 6 +++ api/app/controllers/ontology_controller.py | 3 ++ .../service/end_user_api_controller.py | 2 + .../service/memory_api_controller.py | 2 + api/app/controllers/skill_controller.py | 2 + api/app/controllers/user_controller.py | 13 +++--- api/app/controllers/workspace_controller.py | 2 + api/app/core/api_key_auth.py | 32 ++++++++++++++ api/app/core/quota_stub.py | 44 +++++++++++++++++++ api/app/models/tenant_model.py | 7 +-- api/app/services/api_key_service.py | 29 ++++++++++++ 15 files changed, 140 insertions(+), 10 deletions(-) create mode 100644 api/app/core/quota_stub.py diff --git a/api/app/controllers/app_controller.py b/api/app/controllers/app_controller.py index db3c7536..3b7240de 100644 --- a/api/app/controllers/app_controller.py +++ b/api/app/controllers/app_controller.py @@ -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), diff --git a/api/app/controllers/file_controller.py b/api/app/controllers/file_controller.py index f7bd0e7a..6f8b1b97 100644 --- a/api/app/controllers/file_controller.py +++ b/api/app/controllers/file_controller.py @@ -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, diff --git a/api/app/controllers/knowledge_controller.py b/api/app/controllers/knowledge_controller.py index afda7cce..5cd87647 100644 --- a/api/app/controllers/knowledge_controller.py +++ b/api/app/controllers/knowledge_controller.py @@ -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), diff --git a/api/app/controllers/memory_storage_controller.py b/api/app/controllers/memory_storage_controller.py index 76eed50f..545f8302 100644 --- a/api/app/controllers/memory_storage_controller.py +++ b/api/app/controllers/memory_storage_controller.py @@ -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), diff --git a/api/app/controllers/model_controller.py b/api/app/controllers/model_controller.py index 71fd41ad..6105c3d8 100644 --- a/api/app/controllers/model_controller.py +++ b/api/app/controllers/model_controller.py @@ -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, diff --git a/api/app/controllers/ontology_controller.py b/api/app/controllers/ontology_controller.py index fe6b3598..83f75888 100644 --- a/api/app/controllers/ontology_controller.py +++ b/api/app/controllers/ontology_controller.py @@ -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), diff --git a/api/app/controllers/service/end_user_api_controller.py b/api/app/controllers/service/end_user_api_controller.py index df9996c2..92a9d7c8 100644 --- a/api/app/controllers/service/end_user_api_controller.py +++ b/api/app/controllers/service/end_user_api_controller.py @@ -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, diff --git a/api/app/controllers/service/memory_api_controller.py b/api/app/controllers/service/memory_api_controller.py index dc5e0408..16f1e223 100644 --- a/api/app/controllers/service/memory_api_controller.py +++ b/api/app/controllers/service/memory_api_controller.py @@ -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, diff --git a/api/app/controllers/skill_controller.py b/api/app/controllers/skill_controller.py index 6e673679..4ee07c7d 100644 --- a/api/app/controllers/skill_controller.py +++ b/api/app/controllers/skill_controller.py @@ -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), diff --git a/api/app/controllers/user_controller.py b/api/app/controllers/user_controller.py index cc16a6b4..5a329165 100644 --- a/api/app/controllers/user_controller.py +++ b/api/app/controllers/user_controller.py @@ -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"] diff --git a/api/app/controllers/workspace_controller.py b/api/app/controllers/workspace_controller.py index 6f4a4fa8..47068288 100644 --- a/api/app/controllers/workspace_controller.py +++ b/api/app/controllers/workspace_controller.py @@ -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"), diff --git a/api/app/core/api_key_auth.py b/api/app/core/api_key_auth.py index 342405b8..91d6bd8a 100644 --- a/api/app/core/api_key_auth.py +++ b/api/app/core/api_key_auth.py @@ -96,6 +96,38 @@ def require_api_key( resource_id=api_key_obj.resource_id, ) + # ── Tenant 级别限速（来自套餐配额 api_ops_rate_limit）────────── + try: + from app.models.workspace_model import Workspace + from premium.platform_admin.package_plan_service import TenantSubscriptionService + + workspace = db.query(Workspace).filter( + Workspace.id == api_key_obj.workspace_id + ).first() + if workspace: + quota = TenantSubscriptionService(db).get_effective_quota(workspace.tenant_id) + tenant_qps_limit = quota.get("api_ops_rate_limit") if quota else None + if tenant_qps_limit: + rate_limiter = RateLimiterService() + tenant_ok, tenant_info = await rate_limiter.check_tenant_rate_limit( + workspace.tenant_id, tenant_qps_limit + ) + if not tenant_ok: + raise RateLimitException( + "租户 API 调用速率超限", + BizCode.API_KEY_QPS_LIMIT_EXCEEDED, + rate_headers={ + "X-RateLimit-Tenant-Limit": str(tenant_info["limit"]), + "X-RateLimit-Tenant-Remaining": str(tenant_info["remaining"]), + "X-RateLimit-Tenant-Reset": str(tenant_info["reset"]), + } + ) + except RateLimitException: + raise + except Exception as e: + logger.warning(f"Tenant 限速检查异常，跳过: {e}") + # ───────────────────────────────────────────────────────────── + rate_limiter = RateLimiterService() is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj) if not is_allowed: diff --git a/api/app/core/quota_stub.py b/api/app/core/quota_stub.py new file mode 100644 index 00000000..b8f82e75 --- /dev/null +++ b/api/app/core/quota_stub.py @@ -0,0 +1,44 @@ +""" +配额检查 stub - 社区版使用，所有检查直接放行。 +企业版通过 premium.platform_admin.quota_decorator 提供真实实现。 +""" +from functools import wraps +from typing import Callable + + +def _noop_decorator(func: Callable) -> Callable: + """空装饰器，直接放行""" + return func + + +def _noop_check(*args, **kwargs): + """空检查函数，直接放行""" + pass + + +try: + from premium.platform_admin.quota_decorator import ( + check_workspace_quota, + check_skill_quota, + check_app_quota, + check_knowledge_capacity_quota, + check_memory_engine_quota, + check_end_user_quota, + check_ontology_project_quota, + check_model_quota, + check_model_activation_quota, + get_quota_usage, + _check_quota, + ) +except ModuleNotFoundError: + check_workspace_quota = _noop_decorator + check_skill_quota = _noop_decorator + check_app_quota = _noop_decorator + check_knowledge_capacity_quota = _noop_decorator + check_memory_engine_quota = _noop_decorator + check_end_user_quota = _noop_decorator + check_ontology_project_quota = _noop_decorator + check_model_quota = _noop_decorator + check_model_activation_quota = _noop_decorator + get_quota_usage = lambda db, tenant_id: {} + _check_quota = _noop_check diff --git a/api/app/models/tenant_model.py b/api/app/models/tenant_model.py index a92b5629..c3fd82df 100644 --- a/api/app/models/tenant_model.py +++ b/api/app/models/tenant_model.py @@ -29,11 +29,8 @@ class Tenants(Base): contact_email = Column(String(255), nullable=True) # 联系人邮箱 contact_phone = Column(String(50), nullable=True) # 联系人电话 - # 租户套餐信息 - plan = Column(String(50), nullable=True) # 套餐类型 - plan_expired_at = Column(DateTime, nullable=True) # 套餐到期时间 - api_ops_rate_limit = Column(String(100), nullable=True) # API 调用频率限制 - status = Column(String(50), nullable=True, default='active') # 租户状态 + # 租户套餐信息（只读，从 tenant_subscriptions 动态获取） + status = Column(String(50), nullable=True, default='active', server_default='active') # 租户状态 # Relationship to users - one tenant has many users users = relationship("User", back_populates="tenant") diff --git a/api/app/services/api_key_service.py b/api/app/services/api_key_service.py index a49e8fe0..07d55198 100644 --- a/api/app/services/api_key_service.py +++ b/api/app/services/api_key_service.py @@ -248,6 +248,35 @@ class RateLimiterService: def __init__(self): self.redis = aio_redis + async def check_tenant_rate_limit(self, tenant_id: uuid.UUID, limit: int) -> Tuple[bool, dict]: + """ + 按 tenant_id 做 1 秒滑动窗口限速，限制值来自套餐配额 api_ops_rate_limit + """ + now = time.time() + window_start = now - 1 # 1 秒窗口 + key = f"rate_limit:tenant_qps:{tenant_id}" + + async with self.redis.pipeline() as pipe: + # 清理 1 秒前的旧记录 + pipe.zremrangebyscore(key, 0, window_start) + # 加入当前请求（score=时间戳，member=时间戳+随机数保证唯一） + pipe.zadd(key, {f"{now}:{uuid.uuid4().hex}": now}) + # 统计窗口内请求数 + pipe.zcard(key) + # 设置 key 过期（2 秒后自动清理） + pipe.expire(key, 2) + results = await pipe.execute() + + current = results[2] + remaining = max(0, limit - current) + reset_time = int(now) + 1 + + return current <= limit, { + "limit": limit, + "remaining": remaining, + "reset": reset_time, + } + async def check_qps(self, api_key_id: uuid.UUID, limit: int) -> Tuple[bool, dict]: """ 检查QPS限制 From 18be1a9f89dd0f90499dfed3dfeeb4eb26faa07f Mon Sep 17 00:00:00 2001 From: wxy Date: Tue, 14 Apr 2026 18:14:45 +0800 Subject: [PATCH 2/4] feat(tenant): add tenant package query endpoint Add tenant package query functionality. Regular users can access this endpoint to retrieve their tenant's package information. --- api/app/controllers/__init__.py | 4 +- .../tenant_subscription_controller.py | 53 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 api/app/controllers/tenant_subscription_controller.py diff --git a/api/app/controllers/__init__.py b/api/app/controllers/__init__.py index 50e9e0b0..377205c4 100644 --- a/api/app/controllers/__init__.py +++ b/api/app/controllers/__init__.py @@ -47,7 +47,8 @@ from . import ( user_memory_controllers, workspace_controller, ontology_controller, - skill_controller + skill_controller, + tenant_subscription_controller, ) # 创建管理端 API 路由器 @@ -98,5 +99,6 @@ manager_router.include_router(file_storage_controller.router) manager_router.include_router(ontology_controller.router) manager_router.include_router(skill_controller.router) manager_router.include_router(i18n_controller.router) +manager_router.include_router(tenant_subscription_controller.router) __all__ = ["manager_router"] diff --git a/api/app/controllers/tenant_subscription_controller.py b/api/app/controllers/tenant_subscription_controller.py new file mode 100644 index 00000000..2629f7f1 --- /dev/null +++ b/api/app/controllers/tenant_subscription_controller.py @@ -0,0 +1,53 @@ +""" +租户套餐查询接口（普通用户可访问） +""" +from typing import Callable + +from fastapi import APIRouter, Depends +from fastapi.responses import JSONResponse +from sqlalchemy.orm import Session + +from app.core.logging_config import get_api_logger +from app.core.response_utils import success, fail +from app.db import get_db +from app.dependencies import get_current_user +from app.i18n.dependencies import get_translator +from app.models.user_model import User +from app.schemas.response_schema import ApiResponse + +logger = get_api_logger() + +router = APIRouter(prefix="/tenant", tags=["Tenant"]) + + +@router.get("/subscription", response_model=ApiResponse, summary="获取当前用户所属租户的套餐信息") +async def get_my_tenant_subscription( + current_user: User = Depends(get_current_user), + db: Session = Depends(get_db), + t: Callable = Depends(get_translator), +): + """ + 获取当前登录用户所属租户的有效套餐订阅信息。 + 包含套餐名称、版本、配额、到期时间等。 + """ + try: + from premium.platform_admin.package_plan_service import TenantSubscriptionService + + if not current_user.tenant: + return JSONResponse(status_code=404, content=fail(code=404, msg="用户未关联租户")) + + tenant_id = current_user.tenant.id + svc = TenantSubscriptionService(db) + sub = svc.get_subscription(tenant_id) + + if not sub: + return success(data=None, msg="暂无有效套餐") + + return success(data=svc.build_response(sub)) + + except ModuleNotFoundError: + # 社区版无 premium 模块，返回空 + return success(data=None, msg="套餐功能未启用") + except Exception as e: + logger.error(f"获取租户套餐信息失败: {e}", exc_info=True) + return JSONResponse(status_code=500, content=fail(code=500, msg="获取套餐信息失败")) From 1faa258e233eb5712526ac1abe9086ac5a165035 Mon Sep 17 00:00:00 2001 From: wwq Date: Wed, 15 Apr 2026 18:48:09 +0800 Subject: [PATCH 3/4] feat(quota): implement unified quota management system and add community free plan - Add `default_free_plan.py` to define the configuration for the Community Free Plan. - Refactor `quota_stub.py` as a unified entry point, delegating checks to `core/quota_manager`. - Implement core logic in `quota_manager.py` to support retrieving quotas from the premium module or configuration files. - Update `tenant_subscription_controller` to return Community Free Plan information. --- api/app/config/default_free_plan.py | 30 ++ .../tenant_subscription_controller.py | 33 +- api/app/core/quota_manager.py | 473 ++++++++++++++++++ api/app/core/quota_stub.py | 74 ++- 4 files changed, 567 insertions(+), 43 deletions(-) create mode 100644 api/app/config/default_free_plan.py create mode 100644 api/app/core/quota_manager.py diff --git a/api/app/config/default_free_plan.py b/api/app/config/default_free_plan.py new file mode 100644 index 00000000..23a3a10e --- /dev/null +++ b/api/app/config/default_free_plan.py @@ -0,0 +1,30 @@ +""" +社区版默认免费套餐配置 +当无法从 SaaS 版获取 premium 模块时，使用此配置作为兜底 +""" + +DEFAULT_FREE_PLAN = { + "name": "记忆体验版", + "category": "saas_personal", + "tier_level": 0, + "version": "1.0", + "status": True, + "price": 0, + "billing_cycle": "permanent_free", + "core_value": "感受永久记忆", + "tech_support": "社群交流", + "sla_compliance": "无", + "page_customization": "无", + "theme_color": "#64748B", + "quotas": { + "workspace_quota": 1, + "skill_quota": 5, + "app_quota": 2, + "knowledge_capacity_quota": 0.3, + "memory_engine_quota": 1, + "end_user_quota": 1, + "ontology_project_quota": 3, + "model_quota": 1, + "api_ops_rate_limit": 50, + }, +} diff --git a/api/app/controllers/tenant_subscription_controller.py b/api/app/controllers/tenant_subscription_controller.py index 2629f7f1..c3fde572 100644 --- a/api/app/controllers/tenant_subscription_controller.py +++ b/api/app/controllers/tenant_subscription_controller.py @@ -1,6 +1,7 @@ """ 租户套餐查询接口（普通用户可访问） """ +import datetime from typing import Callable from fastapi import APIRouter, Depends @@ -46,8 +47,36 @@ async def get_my_tenant_subscription( return success(data=svc.build_response(sub)) except ModuleNotFoundError: - # 社区版无 premium 模块，返回空 - return success(data=None, msg="套餐功能未启用") + # 社区版无 premium 模块，从配置文件读取免费套餐 + if not current_user.tenant: + return JSONResponse(status_code=404, content=fail(code=404, msg="用户未关联租户")) + + from app.config.default_free_plan import DEFAULT_FREE_PLAN + + plan = DEFAULT_FREE_PLAN + response_data = { + "subscription_id": None, + "tenant_id": str(current_user.tenant.id), + "package_plan_id": None, + "package_version": plan["version"], + "package_plan": { + "id": None, + "name": plan["name"], + "version": plan["version"], + "category": plan["category"], + "tier_level": plan["tier_level"], + "price": float(plan["price"]), + "billing_cycle": plan["billing_cycle"], + }, + "started_at": None, + "expired_at": None, + "status": "active", + "quota": plan["quotas"], + "created_at": int(datetime.datetime.utcnow().timestamp() * 1000), + "updated_at": int(datetime.datetime.utcnow().timestamp() * 1000), + } + return success(data=response_data, msg="社区版免费套餐") + except Exception as e: logger.error(f"获取租户套餐信息失败: {e}", exc_info=True) return JSONResponse(status_code=500, content=fail(code=500, msg="获取套餐信息失败")) diff --git a/api/app/core/quota_manager.py b/api/app/core/quota_manager.py new file mode 100644 index 00000000..6c02ac7a --- /dev/null +++ b/api/app/core/quota_manager.py @@ -0,0 +1,473 @@ +""" +统一配额管理器 - 社区版和 SaaS 版共用 + +配额来源策略： +1. 优先从 premium 模块的 tenant_subscriptions 表读取（SaaS 版） +2. 降级到 default_free_plan.py 配置文件（社区版兜底） +""" +import asyncio +import time +from functools import wraps +from typing import Optional, Callable, Dict, Any +from uuid import UUID + +from sqlalchemy import func +from sqlalchemy.orm import Session + +from app.core.logging_config import get_auth_logger +from app.i18n.exceptions import QuotaExceededError + +logger = get_auth_logger() + + +def _get_user_from_kwargs(kwargs: dict): + """从 kwargs 中获取 user 对象""" + for key in ["user", "current_user"]: + if key in kwargs: + return kwargs[key] + return None + + +def _get_tenant_id_from_kwargs(db: Session, kwargs: dict): + """从 kwargs 中获取 tenant_id""" + user = _get_user_from_kwargs(kwargs) + if user and hasattr(user, 'tenant_id'): + return user.tenant_id + + workspace_id = kwargs.get("workspace_id") + if workspace_id: + from app.models.workspace_model import Workspace + workspace = db.query(Workspace).filter(Workspace.id == workspace_id).first() + if workspace: + return workspace.tenant_id + + api_key_auth = kwargs.get("api_key_auth") + if api_key_auth and hasattr(api_key_auth, 'workspace_id'): + from app.models.workspace_model import Workspace + workspace = db.query(Workspace).filter(Workspace.id == api_key_auth.workspace_id).first() + if workspace: + return workspace.tenant_id + + data = kwargs.get("data") or kwargs.get("body") or kwargs.get("payload") + if data and hasattr(data, "workspace_id"): + from app.models.workspace_model import Workspace + workspace = db.query(Workspace).filter(Workspace.id == data.workspace_id).first() + if workspace: + return workspace.tenant_id + + return None + + +def _get_quota_config(db: Session, tenant_id: UUID) -> Optional[Dict[str, Any]]: + """ + 获取租户的配额配置 + + 优先级： + 1. premium 模块的 tenant_subscriptions（SaaS 版） + 2. default_free_plan.py 配置文件（社区版兜底） + """ + # 尝试从 premium 模块获取 + try: + from premium.platform_admin.package_plan_service import TenantSubscriptionService + quota_config = TenantSubscriptionService(db).get_effective_quota(tenant_id) + if quota_config: + logger.debug(f"从 premium 模块获取租户 {tenant_id} 配额配置") + return quota_config + except (ModuleNotFoundError, ImportError, Exception) as e: + logger.debug(f"无法从 premium 模块获取配额配置: {e}") + + # 降级到配置文件 + try: + from app.config.default_free_plan import DEFAULT_FREE_PLAN + logger.info(f"使用配置文件中的免费套餐配额: tenant={tenant_id}") + return DEFAULT_FREE_PLAN.get("quotas") + except Exception as e: + logger.error(f"无法从配置文件获取配额: {e}") + return None + + +class QuotaUsageRepository: + """配额使用量数据访问层""" + + def __init__(self, db: Session): + self.db = db + + def count_workspaces(self, tenant_id: UUID) -> int: + from app.models.workspace_model import Workspace + return self.db.query(Workspace).filter( + Workspace.tenant_id == tenant_id, + Workspace.is_active.is_(True) + ).count() + + def count_apps(self, tenant_id: UUID) -> int: + from app.models.app_model import App + from app.models.workspace_model import Workspace + return self.db.query(App).join( + Workspace, App.workspace_id == Workspace.id + ).filter( + Workspace.tenant_id == tenant_id, + App.is_active.is_(True) + ).count() + + def count_skills(self, tenant_id: UUID) -> int: + from app.models.skill_model import Skill + return self.db.query(Skill).filter( + Skill.tenant_id == tenant_id, + Skill.is_active.is_(True) + ).count() + + def sum_knowledge_capacity_gb(self, tenant_id: UUID) -> float: + from app.models.document_model import Document + from app.models.knowledge_model import Knowledge + from app.models.workspace_model import Workspace + result = self.db.query(func.coalesce(func.sum(Document.file_size), 0)).join( + Knowledge, Document.kb_id == Knowledge.id + ).join( + Workspace, Knowledge.workspace_id == Workspace.id + ).filter( + Workspace.tenant_id == tenant_id, + Document.status == 1, + ).scalar() + return float(result) / (1024 ** 3) if result else 0.0 + + def count_memory_engines(self, tenant_id: UUID) -> int: + from app.models.memory_config_model import MemoryConfig + from app.models.workspace_model import Workspace + return self.db.query(MemoryConfig).join( + Workspace, MemoryConfig.workspace_id == Workspace.id + ).filter( + Workspace.tenant_id == tenant_id + ).count() + + def count_end_users(self, tenant_id: UUID) -> int: + from app.models.end_user_model import EndUser + from app.models.workspace_model import Workspace + return self.db.query(EndUser).join( + Workspace, EndUser.workspace_id == Workspace.id + ).filter( + Workspace.tenant_id == tenant_id + ).count() + + def count_models(self, tenant_id: UUID) -> int: + from app.models.models_model import ModelConfig + return self.db.query(ModelConfig).filter( + ModelConfig.tenant_id == tenant_id, + ModelConfig.is_active == True + ).count() + + def count_ontology_projects(self, tenant_id: UUID) -> int: + from app.models.ontology_scene import OntologyScene + from app.models.workspace_model import Workspace + return self.db.query(OntologyScene).join( + Workspace, OntologyScene.workspace_id == Workspace.id + ).filter( + Workspace.tenant_id == tenant_id + ).count() + + def get_usage_by_quota_type(self, tenant_id: UUID, quota_type: str): + """按配额类型分发，返回当前使用量""" + dispatch = { + "workspace_quota": self.count_workspaces, + "app_quota": self.count_apps, + "skill_quota": self.count_skills, + "knowledge_capacity_quota": self.sum_knowledge_capacity_gb, + "memory_engine_quota": self.count_memory_engines, + "end_user_quota": self.count_end_users, + "model_quota": self.count_models, + "ontology_project_quota": self.count_ontology_projects, + } + fn = dispatch.get(quota_type) + return fn(tenant_id) if fn else 0 + + +def _check_quota( + db: Session, + tenant_id: UUID, + quota_type: str, + resource_name: str, + usage_func: Optional[Callable] = None, +) -> None: + """核心配额检查逻辑：对比使用量和配额限制""" + try: + quota_config = _get_quota_config(db, tenant_id) + if not quota_config: + logger.warning(f"租户 {tenant_id} 无有效配额配置，跳过配额检查") + return + + quota_limit = quota_config.get(quota_type) + if quota_limit is None: + logger.warning(f"配额配置未包含 {quota_type}，跳过配额检查") + return + + if usage_func: + current_usage = usage_func(db, tenant_id) + else: + current_usage = QuotaUsageRepository(db).get_usage_by_quota_type(tenant_id, quota_type) + + if current_usage >= quota_limit: + logger.warning( + f"配额不足: tenant={tenant_id}, type={quota_type}, " + f"usage={current_usage}, limit={quota_limit}" + ) + raise QuotaExceededError( + resource=resource_name, + current_usage=current_usage, + quota_limit=quota_limit, + ) + + logger.debug( + f"配额检查通过: tenant={tenant_id}, type={quota_type}, " + f"usage={current_usage}, limit={quota_limit}" + ) + + except QuotaExceededError: + raise + except Exception as e: + logger.error( + f"配额检查异常: tenant={tenant_id}, type={quota_type}, " + f"error_type={type(e).__name__}, error={str(e)}", + exc_info=True, + ) + raise + + +# ─── 具名装饰器 ──────────────────────────────────────────────────────────── + +def check_workspace_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "workspace_quota", "workspace") + return func(*args, **kwargs) + return wrapper + + +def check_skill_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "skill_quota", "skill") + return func(*args, **kwargs) + return wrapper + + +def check_app_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "app_quota", "app") + return func(*args, **kwargs) + return wrapper + + +def check_knowledge_capacity_quota(func: Callable) -> Callable: + @wraps(func) + async def async_wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + if not db: + logger.warning("配额检查失败：缺少 db 参数") + return await func(*args, **kwargs) + tenant_id = _get_tenant_id_from_kwargs(db, kwargs) + if not tenant_id: + logger.warning("配额检查失败：无法获取 tenant_id") + return await func(*args, **kwargs) + _check_quota(db, tenant_id, "knowledge_capacity_quota", "knowledge_capacity") + return await func(*args, **kwargs) + + @wraps(func) + def sync_wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "knowledge_capacity_quota", "knowledge_capacity") + return func(*args, **kwargs) + + return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper + + +def check_memory_engine_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "memory_engine_quota", "memory_engine") + return func(*args, **kwargs) + return wrapper + + +def check_end_user_quota(func: Callable) -> Callable: + @wraps(func) + async def async_wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + if not db: + logger.warning("配额检查失败：缺少 db 参数") + return await func(*args, **kwargs) + tenant_id = _get_tenant_id_from_kwargs(db, kwargs) + if not tenant_id: + logger.warning("配额检查失败：无法获取 tenant_id") + return await func(*args, **kwargs) + _check_quota(db, tenant_id, "end_user_quota", "end_user") + return await func(*args, **kwargs) + + @wraps(func) + def sync_wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + if not db: + logger.warning("配额检查失败：缺少 db 参数") + return func(*args, **kwargs) + tenant_id = _get_tenant_id_from_kwargs(db, kwargs) + if not tenant_id: + logger.warning("配额检查失败：无法获取 tenant_id") + return func(*args, **kwargs) + _check_quota(db, tenant_id, "end_user_quota", "end_user") + return func(*args, **kwargs) + + return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper + + +def check_ontology_project_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "ontology_project_quota", "ontology_project") + return func(*args, **kwargs) + return wrapper + + +def check_model_quota(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, "model_quota", "model") + return func(*args, **kwargs) + return wrapper + + +def check_model_activation_quota(func: Callable) -> Callable: + """模型激活时的配额检查装饰器""" + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + + model_id = kwargs.get("model_id") or (args[1] if len(args) > 1 else None) + model_data = kwargs.get("model_data") + + if not model_id or not model_data: + logger.warning("模型激活配额检查失败：缺少 model_id 或 model_data 参数") + return func(*args, **kwargs) + + if model_data.is_active is True: + try: + from app.models.models_model import ModelConfig + from app.services.model_service import ModelConfigService + + existing_model = ModelConfigService.get_model_by_id( + db=db, + model_id=model_id, + tenant_id=user.tenant_id + ) + + if not existing_model.is_active: + logger.info(f"模型激活操作，检查配额: model_id={model_id}, tenant_id={user.tenant_id}") + _check_quota(db, user.tenant_id, "model_quota", "model") + except Exception as e: + logger.error(f"模型激活配额检查异常: model_id={model_id}, error={str(e)}") + raise + + return func(*args, **kwargs) + return wrapper + + +def check_quota(quota_type: str, resource_name: str, usage_func: Optional[Callable] = None): + """通用配额检查装饰器，支持自定义使用量获取函数""" + def decorator(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + db: Session = kwargs.get("db") + user = _get_user_from_kwargs(kwargs) + if not db or not user: + logger.warning("配额检查失败：缺少 db 或 user 参数") + return func(*args, **kwargs) + _check_quota(db, user.tenant_id, quota_type, resource_name, usage_func) + return func(*args, **kwargs) + return wrapper + return decorator + + +# ─── 配额使用统计 ──────────────────────────────────────────────────────────── + +def get_quota_usage(db: Session, tenant_id: UUID) -> dict: + """获取租户所有配额的使用情况""" + quota_config = _get_quota_config(db, tenant_id) + if not quota_config: + return {} + + repo = QuotaUsageRepository(db) + + def pct(used, limit): + return round(used / limit * 100, 1) if limit else None + + workspace_count = repo.count_workspaces(tenant_id) + skill_count = repo.count_skills(tenant_id) + app_count = repo.count_apps(tenant_id) + knowledge_gb = repo.sum_knowledge_capacity_gb(tenant_id) + memory_count = repo.count_memory_engines(tenant_id) + end_user_count = repo.count_end_users(tenant_id) + model_count = repo.count_models(tenant_id) + ontology_count = repo.count_ontology_projects(tenant_id) + + api_ops_current = 0 + try: + from app.core.config import settings + import redis + _now = time.time() + _rk = f"rate_limit:tenant_qps:{tenant_id}" + _r = redis.StrictRedis( + host=settings.REDIS_HOST, port=settings.REDIS_PORT, + db=settings.REDIS_DB, password=settings.REDIS_PASSWORD, + decode_responses=True + ) + api_ops_current = int(_r.zcount(_rk, _now - 1, "+inf")) + except Exception: + pass + + return { + "workspace": {"used": workspace_count, "limit": quota_config.get("workspace_quota"), "percentage": pct(workspace_count, quota_config.get("workspace_quota"))}, + "skill": {"used": skill_count, "limit": quota_config.get("skill_quota"), "percentage": pct(skill_count, quota_config.get("skill_quota"))}, + "app": {"used": app_count, "limit": quota_config.get("app_quota"), "percentage": pct(app_count, quota_config.get("app_quota"))}, + "knowledge_capacity": {"used": round(knowledge_gb, 2), "limit": quota_config.get("knowledge_capacity_quota"), "percentage": pct(knowledge_gb, quota_config.get("knowledge_capacity_quota")), "unit": "GB"}, + "memory_engine": {"used": memory_count, "limit": quota_config.get("memory_engine_quota"), "percentage": pct(memory_count, quota_config.get("memory_engine_quota"))}, + "end_user": {"used": end_user_count, "limit": quota_config.get("end_user_quota"), "percentage": pct(end_user_count, quota_config.get("end_user_quota"))}, + "ontology_project": {"used": ontology_count, "limit": quota_config.get("ontology_project_quota"), "percentage": pct(ontology_count, quota_config.get("ontology_project_quota"))}, + "model": {"used": model_count, "limit": quota_config.get("model_quota"), "percentage": pct(model_count, quota_config.get("model_quota"))}, + "api_ops_rate_limit": {"current": api_ops_current, "limit": quota_config.get("api_ops_rate_limit"), "percentage": None, "unit": "次/秒"}, + } diff --git a/api/app/core/quota_stub.py b/api/app/core/quota_stub.py index b8f82e75..577dfadb 100644 --- a/api/app/core/quota_stub.py +++ b/api/app/core/quota_stub.py @@ -1,44 +1,36 @@ """ -配额检查 stub - 社区版使用，所有检查直接放行。 -企业版通过 premium.platform_admin.quota_decorator 提供真实实现。 +配额检查 stub - 社区版和 SaaS 版统一使用 core.quota_manager 实现 + +所有配额检查逻辑统一在 core 层实现，两个版本共用： +- 社区版：从 default_free_plan.py 读取配额限制 +- SaaS 版：优先从 tenant_subscriptions 表读取，降级到配置文件 """ -from functools import wraps -from typing import Callable +from app.core.quota_manager import ( + check_workspace_quota, + check_skill_quota, + check_app_quota, + check_knowledge_capacity_quota, + check_memory_engine_quota, + check_end_user_quota, + check_ontology_project_quota, + check_model_quota, + check_model_activation_quota, + get_quota_usage, + _check_quota, + QuotaUsageRepository, +) - -def _noop_decorator(func: Callable) -> Callable: - """空装饰器，直接放行""" - return func - - -def _noop_check(*args, **kwargs): - """空检查函数，直接放行""" - pass - - -try: - from premium.platform_admin.quota_decorator import ( - check_workspace_quota, - check_skill_quota, - check_app_quota, - check_knowledge_capacity_quota, - check_memory_engine_quota, - check_end_user_quota, - check_ontology_project_quota, - check_model_quota, - check_model_activation_quota, - get_quota_usage, - _check_quota, - ) -except ModuleNotFoundError: - check_workspace_quota = _noop_decorator - check_skill_quota = _noop_decorator - check_app_quota = _noop_decorator - check_knowledge_capacity_quota = _noop_decorator - check_memory_engine_quota = _noop_decorator - check_end_user_quota = _noop_decorator - check_ontology_project_quota = _noop_decorator - check_model_quota = _noop_decorator - check_model_activation_quota = _noop_decorator - get_quota_usage = lambda db, tenant_id: {} - _check_quota = _noop_check +__all__ = [ + "check_workspace_quota", + "check_skill_quota", + "check_app_quota", + "check_knowledge_capacity_quota", + "check_memory_engine_quota", + "check_end_user_quota", + "check_ontology_project_quota", + "check_model_quota", + "check_model_activation_quota", + "get_quota_usage", + "_check_quota", + "QuotaUsageRepository", +] From b59e2b5bcd22457cdd6b03c6bfa75222c438bd53 Mon Sep 17 00:00:00 2001 From: wwq Date: Thu, 16 Apr 2026 13:35:35 +0800 Subject: [PATCH 4/4] fix(model): fix issue where associated model config status was not updated when deleting API Key When deleting an API Key, check if the associated model configuration has other active keys; if not, automatically set it to inactive. Also optimize the model configuration query method to support multi-type queries and add sorting conditions. --- api/app/repositories/model_repository.py | 25 ++++++++++++------------ api/app/services/model_service.py | 13 +++++++++++- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/api/app/repositories/model_repository.py b/api/app/repositories/model_repository.py index 8c477d39..03870b4d 100644 --- a/api/app/repositories/model_repository.py +++ b/api/app/repositories/model_repository.py @@ -263,16 +263,15 @@ class ModelConfigRepository: raise @staticmethod - def get_by_type(db: Session, model_type: ModelType, tenant_id: uuid.UUID | None = None, is_active: bool = True) -> List[ModelConfig]: - """根据类型获取模型配置""" - db_logger.debug(f"根据类型查询模型配置: type={model_type}, tenant_id={tenant_id}, is_active={is_active}") - + def get_by_type(db: Session, model_types: List[ModelType], tenant_id: uuid.UUID | None = None, is_active: bool = True) -> List[ModelConfig]: + """根据类型获取模型配置，支持多类型查询""" + db_logger.debug(f"根据类型查询模型配置: types={[t.value for t in model_types]}, tenant_id={tenant_id}, is_active={is_active}") + try: query = db.query(ModelConfig).options( joinedload(ModelConfig.api_keys) - ).filter(ModelConfig.type == model_type) - - # 添加租户过滤 + ).filter(ModelConfig.type.in_([t.value for t in model_types])) + if tenant_id: query = query.filter( or_( @@ -280,16 +279,18 @@ class ModelConfigRepository: ModelConfig.is_public ) ) - + if is_active: query = query.filter(ModelConfig.is_active) - - models = query.order_by(ModelConfig.name).all() + + query = query.filter(ModelConfig.is_composite == False) + + models = query.order_by(ModelConfig.created_at.desc()).all() db_logger.debug(f"根据类型查询模型配置成功: 数量={len(models)}") return models - + except Exception as e: - db_logger.error(f"根据类型查询模型配置失败: type={model_type} - {str(e)}") + db_logger.error(f"根据类型查询模型配置失败: types={model_types} - {str(e)}") raise @staticmethod diff --git a/api/app/services/model_service.py b/api/app/services/model_service.py index 4cbb3509..d202b83a 100644 --- a/api/app/services/model_service.py +++ b/api/app/services/model_service.py @@ -729,10 +729,21 @@ class ModelApiKeyService: @staticmethod def delete_api_key(db: Session, api_key_id: uuid.UUID) -> bool: """删除API Key""" - if not ModelApiKeyRepository.get_by_id(db, api_key_id): + api_key = ModelApiKeyRepository.get_by_id(db, api_key_id) + if not api_key: raise BusinessException("API Key不存在", BizCode.NOT_FOUND) + model_config_ids = [mc.id for mc in api_key.model_configs] + success = ModelApiKeyRepository.delete(db, api_key_id) + + for model_config_id in model_config_ids: + model_config = ModelConfigRepository.get_by_id(db, model_config_id) + if model_config: + has_active_key = any(key.is_active for key in model_config.api_keys) + if not has_active_key and model_config.is_active: + model_config.is_active = False + db.commit() return success