Aether/src/api/base/models_service.py

"""
公共模型查询服务

为 Claude/OpenAI/Gemini 的 /models 端点提供统一的查询逻辑

查询逻辑:
1. 找到指定 api_format 的活跃端点
2. 端点下有活跃的 Key
3. Provider 关联了该模型（Model 表）
4. Key 的 allowed_models 允许该模型（null = 允许所有）
"""

from dataclasses import asdict, dataclass
from typing import Any, Optional

from sqlalchemy.orm import Session, joinedload

from src.config.constants import CacheTTL
from src.core.cache_service import CacheService
from src.core.logger import logger
from src.models.database import (
    ApiKey,
    GlobalModel,
    Model,
    Provider,
    ProviderAPIKey,
    ProviderEndpoint,
    User,
)

# 缓存 key 前缀
_CACHE_KEY_PREFIX = "models:list"
_CACHE_TTL = CacheTTL.MODEL  # 300 秒


def _get_cache_key(api_formats: list[str]) -> str:
    """生成缓存 key"""
    formats_str = ",".join(sorted(api_formats))
    return f"{_CACHE_KEY_PREFIX}:{formats_str}"


async def _get_cached_models(api_formats: list[str]) -> Optional[list["ModelInfo"]]:
    """从缓存获取模型列表"""
    cache_key = _get_cache_key(api_formats)
    try:
        cached = await CacheService.get(cache_key)
        if cached:
            logger.debug(f"[ModelsService] 缓存命中: {cache_key}, {len(cached)} 个模型")
            return [ModelInfo(**item) for item in cached]
    except Exception as e:
        logger.warning(f"[ModelsService] 缓存读取失败: {e}")
    return None


async def _set_cached_models(api_formats: list[str], models: list["ModelInfo"]) -> None:
    """将模型列表写入缓存"""
    cache_key = _get_cache_key(api_formats)
    try:
        data = [asdict(m) for m in models]
        await CacheService.set(cache_key, data, ttl_seconds=_CACHE_TTL)
        logger.debug(f"[ModelsService] 已缓存: {cache_key}, {len(models)} 个模型, TTL={_CACHE_TTL}s")
    except Exception as e:
        logger.warning(f"[ModelsService] 缓存写入失败: {e}")


async def invalidate_models_list_cache() -> None:
    """
    清除所有 /v1/models 列表缓存

    在模型创建、更新、删除时调用，确保模型列表实时更新
    """
    # 清除所有格式的缓存
    all_formats = ["CLAUDE", "OPENAI", "GEMINI"]
    for fmt in all_formats:
        cache_key = f"{_CACHE_KEY_PREFIX}:{fmt}"
        try:
            await CacheService.delete(cache_key)
            logger.debug(f"[ModelsService] 已清除缓存: {cache_key}")
        except Exception as e:
            logger.warning(f"[ModelsService] 清除缓存失败 {cache_key}: {e}")


@dataclass
class ModelInfo:
    """统一的模型信息结构"""

    id: str  # 模型 ID (GlobalModel.name 或 provider_model_name)
    display_name: str
    description: Optional[str]
    created_at: Optional[str]  # ISO 格式
    created_timestamp: int  # Unix 时间戳
    provider_name: str
    provider_id: str = ""  # Provider ID，用于权限过滤
    # 能力配置
    streaming: bool = True
    vision: bool = False
    function_calling: bool = False
    extended_thinking: bool = False
    image_generation: bool = False
    structured_output: bool = False
    # 规格参数
    context_limit: Optional[int] = None
    output_limit: Optional[int] = None
    # 元信息
    family: Optional[str] = None
    knowledge_cutoff: Optional[str] = None
    input_modalities: Optional[list[str]] = None
    output_modalities: Optional[list[str]] = None


@dataclass
class AccessRestrictions:
    """API Key 或 User 的访问限制"""

    allowed_providers: Optional[list[str]] = None  # 允许的 Provider ID 列表
    allowed_models: Optional[list[str]] = None  # 允许的模型名称列表
    allowed_api_formats: Optional[list[str]] = None  # 允许的 API 格式列表

    @classmethod
    def from_api_key_and_user(
        cls, api_key: Optional[ApiKey], user: Optional[User]
    ) -> "AccessRestrictions":
        """
        从 API Key 和 User 合并访问限制

        限制逻辑:
        - API Key 的限制优先于 User 的限制
        - 如果 API Key 有限制，使用 API Key 的限制
        - 如果 API Key 无限制但 User 有限制，使用 User 的限制
        - 两者都无限制则返回空限制
        """
        allowed_providers: Optional[list[str]] = None
        allowed_models: Optional[list[str]] = None
        allowed_api_formats: Optional[list[str]] = None

        # 优先使用 API Key 的限制
        if api_key:
            if api_key.allowed_providers is not None:
                allowed_providers = api_key.allowed_providers
            if api_key.allowed_models is not None:
                allowed_models = api_key.allowed_models
            if api_key.allowed_api_formats is not None:
                allowed_api_formats = api_key.allowed_api_formats

        # 如果 API Key 没有限制，检查 User 的限制
        # 注意: User 没有 allowed_api_formats 字段
        if user:
            if allowed_providers is None and user.allowed_providers is not None:
                allowed_providers = user.allowed_providers
            if allowed_models is None and user.allowed_models is not None:
                allowed_models = user.allowed_models

        return cls(
            allowed_providers=allowed_providers,
            allowed_models=allowed_models,
            allowed_api_formats=allowed_api_formats,
        )

    def is_api_format_allowed(self, api_format: str) -> bool:
        """
        检查 API 格式是否被允许

        Args:
            api_format: API 格式 (如 "OPENAI", "CLAUDE", "GEMINI")

        Returns:
            True 如果格式被允许，False 否则
        """
        if self.allowed_api_formats is None:
            return True
        return api_format in self.allowed_api_formats

    def is_model_allowed(self, model_id: str, provider_id: str) -> bool:
        """
        检查模型是否被允许访问

        Args:
            model_id: 模型 ID
            provider_id: Provider ID

        Returns:
            True 如果模型被允许，False 否则
        """
        # 检查 Provider 限制
        if self.allowed_providers is not None:
            if provider_id not in self.allowed_providers:
                return False

        # 检查模型限制
        if self.allowed_models is not None:
            if model_id not in self.allowed_models:
                return False

        return True


def get_available_provider_ids(db: Session, api_formats: list[str]) -> set[str]:
    """
    返回有可用端点的 Provider IDs

    条件:
    - 端点 api_format 匹配
    - 端点是活跃的
    - 端点下有活跃的 Key
    """
    rows = (
        db.query(ProviderEndpoint.provider_id)
        .join(ProviderAPIKey, ProviderAPIKey.endpoint_id == ProviderEndpoint.id)
        .filter(
            ProviderEndpoint.api_format.in_(api_formats),
            ProviderEndpoint.is_active.is_(True),
            ProviderAPIKey.is_active.is_(True),
        )
        .distinct()
        .all()
    )
    return {row[0] for row in rows}


def _get_available_model_ids_for_format(db: Session, api_formats: list[str]) -> set[str]:
    """
    获取指定格式下真正可用的模型 ID 集合

    一个模型可用需满足:
    1. 端点 api_format 匹配且活跃
    2. 端点下有活跃的 Key
    3. **该端点的 Provider 关联了该模型**
    4. Key 的 allowed_models 允许该模型（null = 允许该 Provider 关联的所有模型）
    """
    # 查询所有匹配格式的活跃端点及其活跃 Key，同时获取 endpoint_id
    endpoint_keys = (
        db.query(
            ProviderEndpoint.id.label("endpoint_id"),
            ProviderEndpoint.provider_id,
            ProviderAPIKey.allowed_models,
        )
        .join(ProviderAPIKey, ProviderAPIKey.endpoint_id == ProviderEndpoint.id)
        .filter(
            ProviderEndpoint.api_format.in_(api_formats),
            ProviderEndpoint.is_active.is_(True),
            ProviderAPIKey.is_active.is_(True),
        )
        .all()
    )

    if not endpoint_keys:
        return set()

    # 收集每个 (provider_id, endpoint_id) 对应的 allowed_models
    # 使用 provider_id 作为 key，因为模型是关联到 Provider 的
    provider_allowed_models: dict[str, list[Optional[list[str]]]] = {}
    provider_ids_with_format: set[str] = set()

    for endpoint_id, provider_id, allowed_models in endpoint_keys:
        provider_ids_with_format.add(provider_id)
        if provider_id not in provider_allowed_models:
            provider_allowed_models[provider_id] = []
        provider_allowed_models[provider_id].append(allowed_models)

    # 只查询那些有匹配格式端点的 Provider 下的模型
    models = (
        db.query(Model)
        .options(joinedload(Model.global_model))
        .join(Provider)
        .filter(
            Model.provider_id.in_(provider_ids_with_format),
            Model.is_active.is_(True),
            Provider.is_active.is_(True),
        )
        .all()
    )

    available_model_ids: set[str] = set()

    for model in models:
        model_provider_id = model.provider_id
        global_model = model.global_model
        model_id = global_model.name if global_model else model.provider_model_name  # type: ignore

        if not model_provider_id or not model_id:
            continue

        # 该模型的 Provider 必须有匹配格式的端点
        if model_provider_id not in provider_ids_with_format:
            continue

        # 检查该 provider 下是否有 Key 允许这个模型
        allowed_lists = provider_allowed_models.get(model_provider_id, [])
        for allowed_models in allowed_lists:
            if allowed_models is None:
                # null = 允许该 Provider 关联的所有模型（已通过上面的查询限制）
                available_model_ids.add(model_id)
                break
            elif model_id in allowed_models:
                # 明确在允许列表中
                available_model_ids.add(model_id)
                break
            elif global_model and model.provider_model_name in allowed_models:
                # 也检查 provider_model_name
                available_model_ids.add(model_id)
                break

    return available_model_ids


def _extract_model_info(model: Any) -> ModelInfo:
    """从 Model 对象提取 ModelInfo"""
    global_model = model.global_model
    model_id: str = global_model.name if global_model else model.provider_model_name
    display_name: str = global_model.display_name if global_model else model.provider_model_name
    created_at: Optional[str] = (
        model.created_at.strftime("%Y-%m-%dT%H:%M:%SZ") if model.created_at else None
    )
    created_timestamp: int = int(model.created_at.timestamp()) if model.created_at else 0
    provider_name: str = model.provider.name if model.provider else "unknown"
    provider_id: str = model.provider_id or ""

    # 从 GlobalModel.config 提取配置信息
    config: dict = {}
    description: Optional[str] = None
    if global_model:
        config = global_model.config or {}
        description = config.get("description")

    return ModelInfo(
        id=model_id,
        display_name=display_name,
        description=description,
        created_at=created_at,
        created_timestamp=created_timestamp,
        provider_name=provider_name,
        provider_id=provider_id,
        # 能力配置
        streaming=config.get("streaming", True),
        vision=config.get("vision", False),
        function_calling=config.get("function_calling", False),
        extended_thinking=config.get("extended_thinking", False),
        image_generation=config.get("image_generation", False),
        structured_output=config.get("structured_output", False),
        # 规格参数
        context_limit=config.get("context_limit"),
        output_limit=config.get("output_limit"),
        # 元信息
        family=config.get("family"),
        knowledge_cutoff=config.get("knowledge_cutoff"),
        input_modalities=config.get("input_modalities"),
        output_modalities=config.get("output_modalities"),
    )


async def list_available_models(
    db: Session,
    available_provider_ids: set[str],
    api_formats: Optional[list[str]] = None,
    restrictions: Optional[AccessRestrictions] = None,
) -> list[ModelInfo]:
    """
    获取可用模型列表（已去重，带缓存）

    Args:
        db: 数据库会话
        available_provider_ids: 有可用端点的 Provider ID 集合
        api_formats: API 格式列表，用于检查 Key 的 allowed_models
        restrictions: API Key/User 的访问限制

    Returns:
        去重后的 ModelInfo 列表，按创建时间倒序
    """
    if not available_provider_ids:
        return []

    # 缓存策略：只有完全无访问限制时才使用缓存
    # - restrictions is None: 未传入限制对象
    # - restrictions 的两个字段都为 None: 传入了限制对象但无实际限制
    # 以上两种情况返回的结果相同，可以共享全局缓存
    use_cache = restrictions is None or (
        restrictions.allowed_providers is None and restrictions.allowed_models is None
    )

    # 尝试从缓存获取
    if api_formats and use_cache:
        cached = await _get_cached_models(api_formats)
        if cached is not None:
            return cached

    # 如果提供了 api_formats，获取真正可用的模型 ID
    available_model_ids: Optional[set[str]] = None
    if api_formats:
        available_model_ids = _get_available_model_ids_for_format(db, api_formats)
        if not available_model_ids:
            return []

    query = (
        db.query(Model)
        .options(joinedload(Model.global_model), joinedload(Model.provider))
        .join(Provider)
        .filter(
            Model.is_active.is_(True),
            Provider.is_active.is_(True),
            Model.provider_id.in_(available_provider_ids),
        )
        .order_by(Model.created_at.desc())
    )
    all_models = query.all()

    result: list[ModelInfo] = []
    seen_model_ids: set[str] = set()

    for model in all_models:
        info = _extract_model_info(model)

        # 如果有 available_model_ids 限制，检查是否在其中
        if available_model_ids is not None and info.id not in available_model_ids:
            continue

        # 检查 API Key/User 访问限制
        if restrictions is not None:
            if not restrictions.is_model_allowed(info.id, info.provider_id):
                continue

        if info.id in seen_model_ids:
            continue
        seen_model_ids.add(info.id)

        result.append(info)

    # 只有无限制的情况才写入缓存
    if api_formats and use_cache:
        await _set_cached_models(api_formats, result)

    return result


def find_model_by_id(
    db: Session,
    model_id: str,
    available_provider_ids: set[str],
    api_formats: Optional[list[str]] = None,
    restrictions: Optional[AccessRestrictions] = None,
) -> Optional[ModelInfo]:
    """
    按 ID 查找模型

    查找顺序：
    1. 先按 GlobalModel.name 查找
    2. 如果没找到任何候选，再按 provider_model_name 查找
    3. 如果有候选但都不可用，返回 None（不回退）

    Args:
        db: 数据库会话
        model_id: 模型 ID
        available_provider_ids: 有可用端点的 Provider ID 集合
        api_formats: API 格式列表，用于检查 Key 的 allowed_models
        restrictions: API Key/User 的访问限制

    Returns:
        ModelInfo 或 None
    """
    if not available_provider_ids:
        return None

    # 如果提供了 api_formats，获取真正可用的模型 ID
    available_model_ids: Optional[set[str]] = None
    if api_formats:
        available_model_ids = _get_available_model_ids_for_format(db, api_formats)
        # 快速检查：如果目标模型不在可用列表中，直接返回 None
        if available_model_ids is not None and model_id not in available_model_ids:
            return None

    # 快速检查：如果 restrictions 明确限制了模型列表且目标模型不在其中，直接返回 None
    if restrictions is not None and restrictions.allowed_models is not None:
        if model_id not in restrictions.allowed_models:
            return None

    # 先按 GlobalModel.name 查找
    models_by_global = (
        db.query(Model)
        .options(joinedload(Model.global_model), joinedload(Model.provider))
        .join(Provider)
        .join(GlobalModel, Model.global_model_id == GlobalModel.id)
        .filter(
            GlobalModel.name == model_id,
            Model.is_active.is_(True),
            Provider.is_active.is_(True),
        )
        .order_by(Model.created_at.desc())
        .all()
    )

    def is_model_accessible(m: Model) -> bool:
        """检查模型是否可访问"""
        if m.provider_id not in available_provider_ids:
            return False
        # 检查 API Key/User 访问限制
        if restrictions is not None:
            provider_id = m.provider_id or ""
            if not restrictions.is_model_allowed(model_id, provider_id):
                return False
        return True

    model = next(
        (m for m in models_by_global if is_model_accessible(m)),
        None,
    )

    # 如果有候选但都不可用，直接返回 None（不回退 provider_model_name）
    if not model and models_by_global:
        return None

    # 如果找不到任何候选，按 provider_model_name 查找
    if not model:
        models_by_provider_name = (
            db.query(Model)
            .options(joinedload(Model.global_model), joinedload(Model.provider))
            .join(Provider)
            .filter(
                Model.provider_model_name == model_id,
                Model.is_active.is_(True),
                Provider.is_active.is_(True),
            )
            .order_by(Model.created_at.desc())
            .all()
        )

        model = next(
            (m for m in models_by_provider_name if is_model_accessible(m)),
            None,
        )

    if not model:
        return None

    return _extract_model_info(model)