i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 简化模型权限检查，移除 resolved_model_name 参数

Browse Source 
- 从 check_model_allowed 和 check_model_allowed_with_mappings 移除 resolved_model_name 参数
- 优化匹配逻辑：先精确匹配 model_name，再检查 candidate_models 交集，最后进行映射匹配
- 更新 CacheAwareScheduler 中的相关调用，传入 provider_model_names 作为 candidate_models
- 同步更新测试用例
This commit is contained in:
fawney19
2026-01-23 00:57:47 +08:00
parent 7ef0488419
commit 08a3148ab3
6 changed files with 40 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/api/admin/models/routing.py
View File

@@ -300,7 +300,6 @@ class AdminGetModelRoutingPreviewAdapter(AdminApiAdapter):
is_allowed, _ = check_model_allowed_with_mappings(
model_name=global_model.name,
allowed_models=allowed_models_list,
resolved_model_name=global_model.name,
model_mappings=global_model_mappings,
)
return is_allowed

1
src/api/base/models_service.py
View File

@@ -270,7 +270,6 @@ def _get_available_model_ids_for_format(db: Session, api_formats: list[str]) ->
is_allowed, _ = check_model_allowed_with_mappings(
model_name=model_id,
allowed_models=allowed_models,
resolved_model_name=model.provider_model_name,
model_mappings=model_mappings,
)
if is_allowed:

56
src/core/model_permissions.py
View File

@@ -49,7 +49,6 @@ def normalize_allowed_models(allowed_models: AllowedModels) -> Optional[set[str]
def check_model_allowed(
model_name: str,
allowed_models: AllowedModels,
resolved_model_name: Optional[str] = None,
) -> bool:
"""
检查模型是否被允许
@@ -57,7 +56,6 @@ def check_model_allowed(
Args:
model_name: 请求的模型名称
allowed_models: 允许的模型配置
resolved_model_name: 解析后的 GlobalModel.name（可选）
Returns:
True: 允许使用该模型
@@ -73,14 +71,8 @@ def check_model_allowed(
# 空集合 = 拒绝所有
return False
# 检查请求的模型名或解析后的名称是否在白名单中
if model_name in allowed_set:
return True
if resolved_model_name and resolved_model_name in allowed_set:
return True
return False
# 检查请求的模型名是否在白名单中
return model_name in allowed_set
def merge_allowed_models(
@@ -350,7 +342,6 @@ def match_model_with_pattern(pattern: str, model_name: str) -> bool:
def check_model_allowed_with_mappings(
model_name: str,
allowed_models: AllowedModels,
resolved_model_name: Optional[str] = None,
model_mappings: Optional[List[str]] = None,
candidate_models: Optional[set[str]] = None,
) -> tuple[bool, Optional[str]]:
@@ -358,8 +349,8 @@ def check_model_allowed_with_mappings(
检查模型是否被允许（支持映射通配符匹配）
匹配优先级：
1. 精确匹配 model_name（用户请求的模型名）
2. 精确匹配 resolved_model_name（GlobalModel.name）
1. 精确匹配 model_name（用户请求的模型名，即 GlobalModel.name）
2. 精确匹配 candidate_models 中的任一模型名（Provider 的 provider_model_mappings）
3. 遍历 model_mappings，检查每个映射是否匹配 allowed_models 中的任一项
映射匹配顺序说明：
@@ -369,25 +360,22 @@ def check_model_allowed_with_mappings(
- 如需确定性行为，请确保 model_mappings 中的规则从最具体到最通用排序
Args:
model_name: 请求的模型名称
model_name: 请求的模型名称（GlobalModel.name）
allowed_models: 允许的模型配置（来自 Provider Key）
resolved_model_name: 解析后的 GlobalModel.name
model_mappings: GlobalModel 的映射列表（来自 config.model_mappings）
candidate_models: 可选的候选模型集合（用于限制映射匹配只能落到这些模型名上）
candidate_models: 可选的候选模型集合（Provider 的 provider_model_names，包含 provider_model_name 和 provider_model_mappings）
Returns:
(is_allowed, matched_model_name):
- is_allowed: 是否允许使用该模型
- matched_model_name: 通过映射匹配到的模型名（仅映射匹配时有值，精确匹配时为 None）
- matched_model_name: 匹配到的模型名（用于实际请求时的模型名替换）
- model_name 精确匹配时为 None（无需替换）
- candidate_models 或 model_mappings 匹配时返回匹配到的模型名
"""
# 先尝试精确匹配（使用原有逻辑）
if check_model_allowed(model_name, allowed_models, resolved_model_name):
# 先尝试精确匹配 model_name
if check_model_allowed(model_name, allowed_models):
return True, None
# 如果精确匹配失败且有映射配置，尝试映射匹配
if not model_mappings:
return False, None
# 获取 allowed_models 的集合
allowed_set = normalize_allowed_models(allowed_models)
@@ -399,19 +387,31 @@ def check_model_allowed_with_mappings(
# 空集合 = 拒绝所有
return False, None
# 如果提供了候选集合，只允许在候选集合中进行映射匹配
# 检查 candidate_models 与 allowed_models 的交集
# candidate_models = Provider 实际支持的模型名（provider_model_name + provider_model_mappings）
# 如果有交集，说明 Key 的 allowed_models 中有 Provider 支持的模型名，可以直接使用
if candidate_models:
intersection = allowed_set & candidate_models
if intersection:
# 返回第一个匹配的模型名（排序确保确定性），用于实际请求时替换 model_name
return True, sorted(intersection)[0]
# 如果精确匹配失败且有映射配置，尝试映射匹配
if not model_mappings:
return False, None
# 映射匹配的搜索空间：allowed_models ∩ candidate_models
# 只在 Provider 实际支持的模型名中进行正则匹配，避免匹配到 Provider 不支持的模型
if candidate_models is not None:
allowed_set = allowed_set & candidate_models
if len(allowed_set) == 0:
return False, None
# 遍历 allowed_models 中的每个模型名，检查是否有映射能匹配
# 注意：为了避免 set 迭代顺序带来的非确定性，这里对 allowed_set 做排序
# 返回第一个匹配的模型名，匹配顺序由 allowed_models 排序结果和 model_mappings 数组顺序共同决定
# 遍历 allowed_set，检查是否有模型名能匹配 model_mappings 中的任一正则
# 排序确保确定性行为
for allowed_model in sorted(allowed_set):
for mapping_pattern in model_mappings:
if match_model_with_pattern(mapping_pattern, allowed_model):
# 返回匹配到的模型名，用于实际请求
return True, allowed_model
return False, None

40
src/services/cache/aware_scheduler.py vendored
View File

@@ -597,8 +597,6 @@ class CacheAwareScheduler:
# 使用 GlobalModel.id 作为缓存亲和性的模型标识，确保映射名和规范名都能命中同一个缓存
global_model_id: str = str(global_model.id)
requested_model_name = model_name
resolved_model_name = str(global_model.name)
# 提取模型映射（用于 Provider Key 的 allowed_models 匹配）
model_mappings: List[str] = (global_model.config or {}).get("model_mappings", [])
@@ -628,17 +626,11 @@ class CacheAwareScheduler:
from src.core.model_permissions import check_model_allowed, get_allowed_models_preview
if not check_model_allowed(
model_name=requested_model_name,
model_name=model_name,
allowed_models=allowed_models,
resolved_model_name=resolved_model_name,
):
resolved_note = (
f" (解析为 {resolved_model_name})"
if resolved_model_name != requested_model_name
else ""
)
logger.debug(
f"用户/API Key 不允许使用模型 {requested_model_name}{resolved_note}, "
f"用户/API Key 不允许使用模型 {model_name}, "
f"允许的模型: {get_allowed_models_preview(allowed_models)}"
)
return [], global_model_id
@@ -674,8 +666,7 @@ class CacheAwareScheduler:
db=db,
providers=providers,
client_format=target_format,
model_name=requested_model_name,
resolved_model_name=resolved_model_name,
model_name=model_name,
model_mappings=model_mappings,
affinity_key=affinity_key,
max_candidates=max_candidates,
@@ -905,7 +896,6 @@ class CacheAwareScheduler:
api_format: Optional[str],
model_name: str,
capability_requirements: Optional[Dict[str, bool]] = None,
resolved_model_name: Optional[str] = None,
model_mappings: Optional[List[str]] = None,
candidate_models: Optional[set[str]] = None,
) -> Tuple[bool, Optional[str], Optional[str]]:
@@ -917,9 +907,8 @@ class CacheAwareScheduler:
Args:
key: API Key 对象
model_name: 模型名称
model_name: 模型名称（GlobalModel.name）
capability_requirements: 能力需求（可选）
resolved_model_name: 解析后的 GlobalModel.name（可选）
model_mappings: GlobalModel 的映射列表（用于通配符匹配）
candidate_models: Provider 侧可用的模型名称集合（用于限制映射匹配范围）
@@ -939,24 +928,20 @@ class CacheAwareScheduler:
# 模型权限检查：使用 allowed_models 白名单
# None = 允许所有模型，[] = 拒绝所有模型，["a","b"] = 只允许指定模型
# 支持通配符映射匹配（通过 model_mappings）
from src.core.model_permissions import (
check_model_allowed_with_mappings,
get_allowed_models_preview,
)
from src.core.model_permissions import check_model_allowed_with_mappings
try:
is_allowed, mapping_matched_model = check_model_allowed_with_mappings(
model_name=model_name,
allowed_models=key.allowed_models,
resolved_model_name=resolved_model_name,
model_mappings=model_mappings,
candidate_models=candidate_models,
)
if mapping_matched_model:
logger.debug(
f"[Scheduler] Key {key.id[:8]}... 映射匹配成功: "
f"[Scheduler] Key {key.id[:8]}... 模型名匹配: "
f"model={model_name} -> {mapping_matched_model}, "
f"allowed_models={key.allowed_models}, model_mappings={model_mappings}"
f"allowed_models={key.allowed_models}"
)
except TimeoutError:
# 正则匹配超时（可能是 ReDoS 攻击或复杂模式）
@@ -977,7 +962,7 @@ class CacheAwareScheduler:
if not is_allowed:
return (
False,
f"模型权限不匹配(允许: {get_allowed_models_preview(key.allowed_models)})",
f"Key 不支持 {model_name}",
None,
)
@@ -1001,7 +986,6 @@ class CacheAwareScheduler:
client_format: APIFormat,
model_name: str,
affinity_key: Optional[str],
resolved_model_name: Optional[str] = None,
model_mappings: Optional[List[str]] = None,
max_candidates: Optional[int] = None,
is_stream: bool = False,
@@ -1017,9 +1001,8 @@ class CacheAwareScheduler:
db: 数据库会话
providers: Provider 列表
client_format: 客户端请求的 API 格式
model_name: 模型名称（用户请求的名称，可能是映射名）
model_name: 模型名称（GlobalModel.name）
affinity_key: 亲和性标识符（通常为API Key ID）
resolved_model_name: 解析后的 GlobalModel.name（用于 Key.allowed_models 校验）
model_mappings: GlobalModel 的映射列表（用于 Key.allowed_models 通配符匹配）
max_candidates: 最大候选数
is_stream: 是否是流式请求，如果为 True 则过滤不支持流式的 Provider
@@ -1108,15 +1091,16 @@ class CacheAwareScheduler:
for key in keys:
# Key 级别检查（健康度/熔断按 provider_format bucket）
# 注意：不传入 candidate_models，保持原有映射匹配行为
# 传入 provider_model_names 作为 candidate_models，
# 用于检查 Key 的 allowed_models 是否支持 Provider 定义的模型名称
is_available, key_skip_reason, mapping_matched_model = (
self._check_key_availability(
key,
endpoint_format_str,
model_name,
capability_requirements,
resolved_model_name=resolved_model_name,
model_mappings=model_mappings,
candidate_models=provider_model_names,
)
)