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

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

View File

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

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( is_allowed, _ = check_model_allowed_with_mappings(
model_name=model_id, model_name=model_id,
allowed_models=allowed_models, allowed_models=allowed_models,
resolved_model_name=model.provider_model_name,
model_mappings=model_mappings, model_mappings=model_mappings,
) )
if is_allowed: if is_allowed:

View File

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

View File

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

View File

@@ -6,7 +6,6 @@ class TestCheckModelAllowedWithMappings:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="gpt-4o", model_name="gpt-4o",
allowed_models=["gpt-4o"], allowed_models=["gpt-4o"],
resolved_model_name="gpt-4o",
model_mappings=[r"gpt-4o-.*"], model_mappings=[r"gpt-4o-.*"],
) )
assert is_allowed is True assert is_allowed is True
@@ -16,7 +15,6 @@ class TestCheckModelAllowedWithMappings:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="target", model_name="target",
allowed_models=["b", "a"], allowed_models=["b", "a"],
resolved_model_name="target",
model_mappings=[r".*"], model_mappings=[r".*"],
) )
assert is_allowed is True assert is_allowed is True
@@ -26,7 +24,6 @@ class TestCheckModelAllowedWithMappings:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="target", model_name="target",
allowed_models=["other-1", "allowed-1"], allowed_models=["other-1", "allowed-1"],
resolved_model_name="target",
model_mappings=[r".*-1"], model_mappings=[r".*-1"],
candidate_models={"allowed-1"}, candidate_models={"allowed-1"},
) )
@@ -37,7 +34,6 @@ class TestCheckModelAllowedWithMappings:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="target", model_name="target",
allowed_models=["allowed-1"], allowed_models=["allowed-1"],
resolved_model_name="target",
model_mappings=[r".*-1"], model_mappings=[r".*-1"],
candidate_models={"not-present"}, candidate_models={"not-present"},
) )

View File

@@ -447,7 +447,6 @@ class TestModelMappingsIntegration:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="claude-haiku-4-5-20251001", # GlobalModel.name model_name="claude-haiku-4-5-20251001", # GlobalModel.name
allowed_models=["claude-3-5-haiku-20251001"], # Key.allowed_models allowed_models=["claude-3-5-haiku-20251001"], # Key.allowed_models
resolved_model_name="claude-haiku-4-5-20251001", # Model.provider_model_name
model_mappings=["claude-3-5-haiku-.*"], # GlobalModel.config.model_mappings model_mappings=["claude-3-5-haiku-.*"], # GlobalModel.config.model_mappings
) )
@@ -461,7 +460,6 @@ class TestModelMappingsIntegration:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="claude-haiku-4-5-20251001", model_name="claude-haiku-4-5-20251001",
allowed_models=["claude-haiku-4-5-20251001"], # 精确匹配 allowed_models=["claude-haiku-4-5-20251001"], # 精确匹配
resolved_model_name="claude-haiku-4-5-20251001",
model_mappings=["claude-3-5-haiku-.*"], model_mappings=["claude-3-5-haiku-.*"],
) )
@@ -475,7 +473,6 @@ class TestModelMappingsIntegration:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="claude-haiku-4-5-20251001", model_name="claude-haiku-4-5-20251001",
allowed_models=["gpt-4o"], # 与映射模式不匹配 allowed_models=["gpt-4o"], # 与映射模式不匹配
resolved_model_name="claude-haiku-4-5-20251001",
model_mappings=["claude-3-5-haiku-.*"], model_mappings=["claude-3-5-haiku-.*"],
) )
@@ -489,7 +486,6 @@ class TestModelMappingsIntegration:
is_allowed, matched = check_model_allowed_with_mappings( is_allowed, matched = check_model_allowed_with_mappings(
model_name="target-model", model_name="target-model",
allowed_models=["b-model-1", "a-model-1"], allowed_models=["b-model-1", "a-model-1"],
resolved_model_name="target-model",
model_mappings=[".*-model-1"], # 匹配两个 model_mappings=[".*-model-1"], # 匹配两个
) )