101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import fnmatch
|
|
from typing import Any
|
|
|
|
|
|
class PermissionRule:
|
|
"""模型能力访问控制的单条权限规则。
|
|
|
|
规则使用 ``fnmatch`` 风格的通配符模式匹配模型名称。
|
|
|
|
参数:
|
|
model_pattern: ``fnmatch`` 模式，用于匹配模型名（如 ``"deepseek-*"``）。
|
|
capability: 此规则管控的能力（如 ``"vision"``）。
|
|
allowed: 能力是否允许（``True`` 允许，``False`` 拒绝）。
|
|
|
|
示例::
|
|
|
|
# 拒绝 "basic-*" 模型上的 vision 能力
|
|
PermissionRule(model_pattern="basic-*", capability="vision", allowed=False)
|
|
|
|
# 允许所有模型的 thinking 能力
|
|
PermissionRule(model_pattern="*", capability="thinking", allowed=True)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_pattern: str,
|
|
capability: str,
|
|
allowed: bool = True,
|
|
) -> None:
|
|
self.model_pattern = model_pattern
|
|
self.capability = capability
|
|
self.allowed = allowed
|
|
|
|
def match(self, model_name: str) -> bool:
|
|
"""检查此规则是否适用于给定的模型名。"""
|
|
return fnmatch.fnmatch(model_name, self.model_pattern)
|
|
|
|
def __repr__(self) -> str:
|
|
verb = "ALLOW" if self.allowed else "DENY"
|
|
return f"PermissionRule({verb} {self.capability} on {self.model_pattern!r})"
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
if not isinstance(other, PermissionRule):
|
|
return NotImplemented
|
|
return (
|
|
self.model_pattern == other.model_pattern
|
|
and self.capability == other.capability
|
|
and self.allowed == other.allowed
|
|
)
|
|
|
|
def __hash__(self) -> int:
|
|
return hash((self.model_pattern, self.capability, self.allowed))
|
|
|
|
|
|
class CapabilityPermissionController:
|
|
"""控制哪些模型允许哪些能力。
|
|
|
|
按顺序评估 :class:`PermissionRule` 列表。
|
|
**最后匹配的规则** 胜出（最具体的规则应放在较宽泛的规则之后）。
|
|
如果没有规则匹配某个模型/能力对，该能力**默认允许**。
|
|
|
|
用法::
|
|
|
|
controller = CapabilityPermissionController([
|
|
PermissionRule(model_pattern="*", capability="vision", allowed=False),
|
|
PermissionRule(model_pattern="vision-*", capability="vision", allowed=True),
|
|
])
|
|
|
|
controller.check("fast-model", "vision") # False
|
|
controller.check("vision-model", "vision") # True
|
|
"""
|
|
|
|
def __init__(self, rules: list[PermissionRule] | None = None) -> None:
|
|
self._rules = list(rules) if rules else []
|
|
|
|
def check(self, model_name: str, capability: str) -> bool:
|
|
"""检查 *capability* 是否在 *model_name* 上被允许。
|
|
|
|
返回 ``True`` 表示允许，``False`` 表示拒绝。
|
|
"""
|
|
result: bool | None = None
|
|
for rule in self._rules:
|
|
if rule.capability == capability and rule.match(model_name):
|
|
result = rule.allowed
|
|
# Default: allow if no rule matched
|
|
return result if result is not None else True
|
|
|
|
def is_allowed(self, model_name: str, capability: str) -> bool:
|
|
""":meth:`check` 的别名。"""
|
|
return self.check(model_name, capability)
|
|
|
|
def filter_capabilities(
|
|
self,
|
|
model_name: str,
|
|
capabilities: set[str],
|
|
) -> set[str]:
|
|
"""返回 *capabilities* 中在 *model_name* 上允许的子集。"""
|
|
return {c for c in capabilities if self.check(model_name, c)}
|