Aether/src/core/key_capabilities.py

"""
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"],  # 上下文超限错误
)