mirror of
https://github.com/fawney19/Aether.git
synced 2026-01-09 11:12:28 +08:00
248 lines
7.8 KiB
Python
248 lines
7.8 KiB
Python
"""
|
||
Key 能力系统
|
||
|
||
能力类型：
|
||
1. 互斥能力 (EXCLUSIVE): 需要时选有的，不需要时选没有的（如 cache_1h）
|
||
2. 兼容能力 (COMPATIBLE): 需要时选有的，不需要时都可选（如 context_1m）
|
||
|
||
配置模式：
|
||
1. user_configurable: 用户可配置（模型级 + Key级强制）
|
||
2. auto_detect: 自动检测（请求失败后升级）
|
||
3. request_param: 从请求参数检测
|
||
"""
|
||
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
|
||
class CapabilityMatchMode(Enum):
|
||
"""能力匹配模式"""
|
||
|
||
EXCLUSIVE = "exclusive" # 互斥：需要时选有的，不需要时选没有的
|
||
COMPATIBLE = "compatible" # 兼容：需要时选有的，不需要时都可选
|
||
|
||
|
||
class CapabilityConfigMode(Enum):
|
||
"""能力配置模式"""
|
||
|
||
USER_CONFIGURABLE = "user_configurable" # 用户可配置（模型级 + Key级强制）
|
||
AUTO_DETECT = "auto_detect" # 自动检测（请求失败后升级）
|
||
REQUEST_PARAM = "request_param" # 从请求参数检测
|
||
|
||
|
||
@dataclass
|
||
class CapabilityDefinition:
|
||
"""能力定义"""
|
||
|
||
name: str
|
||
display_name: str
|
||
description: str
|
||
match_mode: CapabilityMatchMode
|
||
config_mode: CapabilityConfigMode
|
||
short_name: str = "" # 简短展示名称（用于列表等紧凑场景）
|
||
error_patterns: List[str] = field(default_factory=list) # 错误检测关键词组
|
||
|
||
|
||
# ============ 能力注册表 ============
|
||
|
||
_capabilities: Dict[str, CapabilityDefinition] = {}
|
||
|
||
|
||
def register_capability(
|
||
name: str,
|
||
display_name: str,
|
||
description: str,
|
||
match_mode: CapabilityMatchMode,
|
||
config_mode: CapabilityConfigMode,
|
||
short_name: str = "",
|
||
error_patterns: Optional[List[str]] = None,
|
||
) -> CapabilityDefinition:
|
||
"""注册能力"""
|
||
cap = CapabilityDefinition(
|
||
name=name,
|
||
display_name=display_name,
|
||
description=description,
|
||
match_mode=match_mode,
|
||
config_mode=config_mode,
|
||
short_name=short_name or display_name, # 默认使用 display_name
|
||
error_patterns=error_patterns or [],
|
||
)
|
||
_capabilities[name] = cap
|
||
return cap
|
||
|
||
|
||
def get_capability(name: str) -> Optional[CapabilityDefinition]:
|
||
"""获取能力定义"""
|
||
return _capabilities.get(name)
|
||
|
||
|
||
def get_all_capabilities() -> List[CapabilityDefinition]:
|
||
"""获取所有能力定义"""
|
||
return list(_capabilities.values())
|
||
|
||
|
||
def get_user_configurable_capabilities() -> List[CapabilityDefinition]:
|
||
"""获取用户可配置的能力列表"""
|
||
return [c for c in _capabilities.values() if c.config_mode == CapabilityConfigMode.USER_CONFIGURABLE]
|
||
|
||
|
||
# ============ 能力匹配检查 ============
|
||
|
||
|
||
def check_capability_match(
|
||
key_capabilities: Optional[Dict[str, bool]],
|
||
requirements: Optional[Dict[str, bool]],
|
||
) -> Tuple[bool, Optional[str]]:
|
||
"""
|
||
检查 Key 能力是否满足需求
|
||
|
||
匹配逻辑：
|
||
1. EXCLUSIVE（互斥）能力：
|
||
- 请求需要且 Key 有 → 通过
|
||
- 请求需要但 Key 没有 → 拒绝
|
||
- 请求不需要但 Key 有 → 拒绝（避免浪费高价资源）
|
||
- 请求不需要且 Key 没有 → 通过
|
||
- 请求未声明但 Key 有 → 拒绝（关键：未声明等同于不需要）
|
||
|
||
2. COMPATIBLE（兼容）能力：
|
||
- 请求需要且 Key 有 → 通过
|
||
- 请求需要但 Key 没有 → 拒绝
|
||
- 请求不需要/未声明且 Key 有 → 通过（无额外成本，不浪费）
|
||
- 请求不需要/未声明且 Key 没有 → 通过
|
||
|
||
Args:
|
||
key_capabilities: Key 拥有的能力 {"cache_1h": True, ...}
|
||
requirements: 请求需要的能力 {"cache_1h": True, "context_1m": False}
|
||
|
||
Returns:
|
||
(is_match, skip_reason) - 是否匹配及跳过原因
|
||
"""
|
||
key_caps = key_capabilities or {}
|
||
reqs = requirements or {}
|
||
|
||
# 第一步：检查请求声明的需求
|
||
for cap_name, is_required in reqs.items():
|
||
cap_def = _capabilities.get(cap_name)
|
||
if not cap_def:
|
||
continue
|
||
|
||
key_has_cap = key_caps.get(cap_name, False)
|
||
|
||
if cap_def.match_mode == CapabilityMatchMode.EXCLUSIVE:
|
||
if is_required and not key_has_cap:
|
||
return False, f"需要{cap_def.display_name}但 Key 不支持"
|
||
if not is_required and key_has_cap:
|
||
return False, f"不需要{cap_def.display_name}(避免浪费高价资源)"
|
||
|
||
elif cap_def.match_mode == CapabilityMatchMode.COMPATIBLE:
|
||
if is_required and not key_has_cap:
|
||
return False, f"需要{cap_def.display_name}但 Key 不支持"
|
||
|
||
# 第二步：检查 Key 拥有的 EXCLUSIVE 能力是否被请求需要
|
||
# 如果 Key 有某个 EXCLUSIVE 能力，但请求没有声明需要，应该跳过这个 Key
|
||
for cap_name, key_has_cap in key_caps.items():
|
||
if not key_has_cap:
|
||
continue
|
||
|
||
cap_def = _capabilities.get(cap_name)
|
||
if not cap_def:
|
||
continue
|
||
|
||
if cap_def.match_mode == CapabilityMatchMode.EXCLUSIVE:
|
||
# 如果请求没有声明需要这个 EXCLUSIVE 能力，视为不需要
|
||
if cap_name not in reqs:
|
||
return False, f"不需要{cap_def.display_name}(避免浪费高价资源)"
|
||
|
||
return True, None
|
||
|
||
|
||
def _match_error_patterns(error_msg: str, patterns: List[str]) -> bool:
|
||
"""检查错误信息是否匹配模式（所有关键词都要出现）"""
|
||
if not patterns:
|
||
return False
|
||
msg_lower = error_msg.lower()
|
||
return all(p.lower() in msg_lower for p in patterns)
|
||
|
||
|
||
def detect_capability_upgrade_from_error(
|
||
error_msg: str,
|
||
current_requirements: Optional[Dict[str, bool]] = None,
|
||
) -> Optional[str]:
|
||
"""
|
||
从错误信息检测是否需要升级某能力
|
||
|
||
Args:
|
||
error_msg: 错误信息
|
||
current_requirements: 当前已有的能力需求
|
||
|
||
Returns:
|
||
需要升级的能力名称，如果不需要升级则返回 None
|
||
"""
|
||
current_reqs = current_requirements or {}
|
||
|
||
for cap in _capabilities.values():
|
||
if not current_reqs.get(cap.name) and cap.error_patterns:
|
||
if _match_error_patterns(error_msg, cap.error_patterns):
|
||
return cap.name
|
||
|
||
return None
|
||
|
||
|
||
# ============ 兼容性别名 ============
|
||
|
||
# 保留旧 API 兼容
|
||
get_capability_definition = get_capability
|
||
|
||
|
||
class _CapabilityDefinitionsProxy:
|
||
"""CAPABILITY_DEFINITIONS 代理，提供字典式访问（兼容旧代码）"""
|
||
|
||
def get(self, name: str) -> Optional[CapabilityDefinition]:
|
||
return _capabilities.get(name)
|
||
|
||
def __getitem__(self, name: str) -> CapabilityDefinition:
|
||
result = _capabilities.get(name)
|
||
if result is None:
|
||
raise KeyError(name)
|
||
return result
|
||
|
||
def __contains__(self, name: str) -> bool:
|
||
return name in _capabilities
|
||
|
||
def values(self) -> List[CapabilityDefinition]:
|
||
return list(_capabilities.values())
|
||
|
||
def items(self) -> List[Tuple[str, CapabilityDefinition]]:
|
||
return list(_capabilities.items())
|
||
|
||
|
||
CAPABILITY_DEFINITIONS = _CapabilityDefinitionsProxy()
|
||
|
||
|
||
# ============ 兼容旧的插件基类（逐步废弃） ============
|
||
|
||
CapabilityPlugin = CapabilityDefinition # 类型别名，兼容旧代码
|
||
|
||
|
||
# ============ 注册内置能力 ============
|
||
|
||
register_capability(
|
||
name="cache_1h",
|
||
display_name="1 小时缓存",
|
||
description="使用 1 小时缓存 TTL（价格更高，适合长对话）",
|
||
match_mode=CapabilityMatchMode.EXCLUSIVE,
|
||
config_mode=CapabilityConfigMode.USER_CONFIGURABLE,
|
||
short_name="1h缓存",
|
||
)
|
||
|
||
register_capability(
|
||
name="context_1m",
|
||
display_name="CLI 1M 上下文",
|
||
description="支持 1M tokens 上下文窗口",
|
||
match_mode=CapabilityMatchMode.COMPATIBLE,
|
||
config_mode=CapabilityConfigMode.REQUEST_PARAM,
|
||
short_name="CLI 1M",
|
||
error_patterns=["context", "token", "length", "exceed"], # 上下文超限错误
|
||
)
|