adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #899 from wanxunyang/feature/add-quota-check-decorator

Browse Source 
Feature/add quota check decorator
This commit is contained in:
Mark
2026-04-16 13:48:40 +08:00
committed by GitHub
parent 5a2fe738dc b59e2b5bcd
commit 5ce0bdb0f5
21 changed files with 745 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
api/app/config/default_free_plan.py Normal file
View File

@@ -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,
},
}

4
api/app/controllers/__init__.py
View File

@@ -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"]

2
api/app/controllers/app_controller.py
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),

2
api/app/controllers/file_controller.py
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,

2
api/app/controllers/knowledge_controller.py
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),

2
api/app/controllers/memory_storage_controller.py
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),

6
api/app/controllers/model_controller.py
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,

3
api/app/controllers/ontology_controller.py
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),

2
api/app/controllers/service/end_user_api_controller.py
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,

2
api/app/controllers/service/memory_api_controller.py
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,

2
api/app/controllers/skill_controller.py
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),

82
api/app/controllers/tenant_subscription_controller.py Normal file
View File

@@ -0,0 +1,82 @@
"""
租户套餐查询接口（普通用户可访问）
"""
import datetime
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 模块，从配置文件读取免费套餐
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="获取套餐信息失败"))

13
api/app/controllers/user_controller.py
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"]

2
api/app/controllers/workspace_controller.py
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"),

32
api/app/core/api_key_auth.py
View File

@@ -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:

473
api/app/core/quota_manager.py Normal file
View File

@@ -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": "次/秒"},
}

36
api/app/core/quota_stub.py Normal file
View File

@@ -0,0 +1,36 @@
"""
配额检查 stub - 社区版和 SaaS 版统一使用 core.quota_manager 实现
所有配额检查逻辑统一在 core 层实现，两个版本共用：
- 社区版：从 default_free_plan.py 读取配额限制
- SaaS 版：优先从 tenant_subscriptions 表读取，降级到配置文件
"""
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,
)
__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",
]

7
api/app/models/tenant_model.py
View File

@@ -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")

25
api/app/repositories/model_repository.py
View File

@@ -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

29
api/app/services/api_key_service.py
View File

@@ -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限制

13
api/app/services/model_service.py
View File

@@ -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