Merge pull request #961 from SuanmoSuanyangTechnology/fix/wxy_031
fix(api): fix API Key rate limiting and terminal user quota checks
This commit is contained in:
@@ -167,6 +167,8 @@ def update_api_key(
|
||||
|
||||
return success(data=api_key_schema.ApiKey.model_validate(api_key), msg="API Key 更新成功")
|
||||
|
||||
except BusinessException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"未知错误: {str(e)}", extra={
|
||||
"api_key_id": str(api_key_id),
|
||||
|
||||
@@ -219,9 +219,20 @@ def list_conversations(
|
||||
end_user_repo = EndUserRepository(db)
|
||||
app_service = AppService(db)
|
||||
app = app_service._get_app_or_404(share.app_id)
|
||||
workspace_id = app.workspace_id
|
||||
|
||||
# 仅在新建终端用户时检查配额
|
||||
existing_end_user = end_user_repo.get_end_user_by_other_id(workspace_id=workspace_id, other_id=other_id)
|
||||
if existing_end_user is None:
|
||||
from app.core.quota_manager import _check_quota
|
||||
from app.models.workspace_model import Workspace
|
||||
ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
|
||||
if ws:
|
||||
_check_quota(db, ws.tenant_id, "end_user_quota", "end_user")
|
||||
|
||||
new_end_user = end_user_repo.get_or_create_end_user(
|
||||
app_id=share.app_id,
|
||||
workspace_id=app.workspace_id,
|
||||
workspace_id=workspace_id,
|
||||
other_id=other_id
|
||||
)
|
||||
logger.debug(new_end_user.id)
|
||||
@@ -309,7 +320,6 @@ def get_conversation(
|
||||
"/chat",
|
||||
summary="发送消息(支持流式和非流式)"
|
||||
)
|
||||
@check_end_user_quota
|
||||
async def chat(
|
||||
payload: conversation_schema.ChatRequest,
|
||||
share_data: ShareTokenData = Depends(get_share_user_id),
|
||||
@@ -350,6 +360,18 @@ async def chat(
|
||||
app_service = AppService(db)
|
||||
app = app_service._get_app_or_404(share.app_id)
|
||||
workspace_id = app.workspace_id
|
||||
|
||||
# 仅在新建终端用户时检查配额,已有用户复用不受限制
|
||||
existing_end_user = end_user_repo.get_end_user_by_other_id(workspace_id=workspace_id, other_id=other_id)
|
||||
logger.info(f"终端用户配额检查: workspace_id={workspace_id}, other_id={other_id}, existing={existing_end_user is not None}")
|
||||
if existing_end_user is None:
|
||||
from app.core.quota_manager import _check_quota
|
||||
from app.models.workspace_model import Workspace
|
||||
ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
|
||||
if ws:
|
||||
logger.info(f"新终端用户,执行配额检查: tenant_id={ws.tenant_id}")
|
||||
_check_quota(db, ws.tenant_id, "end_user_quota", "end_user")
|
||||
|
||||
new_end_user = end_user_repo.get_or_create_end_user(
|
||||
app_id=share.app_id,
|
||||
workspace_id=workspace_id,
|
||||
|
||||
@@ -66,6 +66,17 @@ class EndUserRepository:
|
||||
db_logger.error(f"查询宿主 {end_user_id} 时出错: {str(e)}")
|
||||
raise
|
||||
|
||||
def get_end_user_by_other_id(self, workspace_id: uuid.UUID, other_id: str) -> Optional["EndUser"]:
|
||||
"""按 workspace_id + other_id 查找终端用户,不存在返回 None"""
|
||||
return (
|
||||
self.db.query(EndUser)
|
||||
.filter(
|
||||
EndUser.workspace_id == workspace_id,
|
||||
EndUser.other_id == other_id
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
def get_or_create_end_user(
|
||||
self,
|
||||
app_id: uuid.UUID,
|
||||
|
||||
@@ -51,7 +51,7 @@ class ApiKeyService:
|
||||
if existing:
|
||||
raise BusinessException(f"API Key 名称 {data.name} 已存在", BizCode.API_KEY_DUPLICATE_NAME)
|
||||
|
||||
# 若 rate_limit 超过租户套餐的 api_ops_rate_limit,自动截断到套餐上限
|
||||
# 若 rate_limit 超过租户套餐的 api_ops_rate_limit,直接报错
|
||||
from app.models.workspace_model import Workspace
|
||||
from app.core.quota_manager import get_api_ops_rate_limit
|
||||
|
||||
@@ -59,7 +59,10 @@ class ApiKeyService:
|
||||
if workspace:
|
||||
tenant_api_ops_limit = get_api_ops_rate_limit(db, workspace.tenant_id)
|
||||
if tenant_api_ops_limit and data.rate_limit > tenant_api_ops_limit:
|
||||
data.rate_limit = tenant_api_ops_limit
|
||||
raise BusinessException(
|
||||
f"API Key QPS 不能超过套餐上限 {tenant_api_ops_limit}",
|
||||
BizCode.BAD_REQUEST
|
||||
)
|
||||
|
||||
# 生成 API Key
|
||||
api_key = generate_api_key(data.type)
|
||||
@@ -162,7 +165,7 @@ class ApiKeyService:
|
||||
if existing:
|
||||
raise BusinessException(f"API Key 名称 {data.name} 已存在", BizCode.API_KEY_DUPLICATE_NAME)
|
||||
|
||||
# 若 rate_limit 超过租户套餐的 api_ops_rate_limit,自动截断到套餐上限
|
||||
# 若 rate_limit 超过租户套餐的 api_ops_rate_limit,直接报错
|
||||
if data.rate_limit is not None:
|
||||
from app.models.workspace_model import Workspace
|
||||
from app.core.quota_manager import get_api_ops_rate_limit
|
||||
@@ -171,7 +174,10 @@ class ApiKeyService:
|
||||
if workspace:
|
||||
tenant_api_ops_limit = get_api_ops_rate_limit(db, workspace.tenant_id)
|
||||
if tenant_api_ops_limit and data.rate_limit > tenant_api_ops_limit:
|
||||
data.rate_limit = tenant_api_ops_limit
|
||||
raise BusinessException(
|
||||
f"API Key QPS 不能超过套餐上限 {tenant_api_ops_limit}",
|
||||
BizCode.BAD_REQUEST
|
||||
)
|
||||
|
||||
update_data = data.model_dump(exclude_unset=True)
|
||||
ApiKeyRepository.update(db, api_key_id, update_data)
|
||||
|
||||
Reference in New Issue
Block a user