hunjing/backend/app/ai_provider/core/registry.py

from __future__ import annotations

import random

from app.ai_provider.config.models import ModelConfig, ProviderConfig
from app.ai_provider.core.billing import PricingFactory, PricingRegistry
from app.ai_provider.types.errors import AIProviderError


class ModelRegistry:
    """内存模型索引。支持逻辑名和别名查找。"""

    def __init__(
        self,
        models: list[ModelConfig],
        aliases: dict[str, str] | None = None,
    ) -> None:
        self._models: dict[str, ModelConfig] = {}
        self._aliases: dict[str, str] = dict(aliases or {})
        self._default: str | None = None
        self._fallback: str | None = None
        self._capability_index: dict[str, list[ModelConfig]] = {}

        for m in models:
            self._models[m.name] = m
            if "default" in m.tags and self._default is None:
                self._default = m.name
            if "fallback" in m.tags and self._fallback is None:
                self._fallback = m.name
            caps = m.resolved_capabilities()
            for cap, supported in caps.items():
                if supported:
                    self._capability_index.setdefault(cap, []).append(m)
            # 索引 tags（如 tts / asr / voice_clone）——这些不在
            # resolved_capabilities() 中，但 AudioService 等模块依赖它们查找模型。
            for tag in m.tags:
                self._capability_index.setdefault(tag, []).append(m)

        if self._default is None and models:
            self._default = models[0].name

    def get(self, name: str) -> ModelConfig:
        model = self._models.get(name)
        if model is None:
            raise AIProviderError(
                message=f"Model '{name}' not found in registry",
                code="MODEL_NOT_FOUND",
                status_code=404,
            )
        return model

    def resolve(self, name_or_alias: str | None = None) -> ModelConfig:
        """解析逻辑名或别名为 ModelConfig。

        无参数时返回默认模型。
        """
        if name_or_alias is None:
            return self._models[self.get_default_name()]
        if name_or_alias in self._models:
            return self._models[name_or_alias]
        if name_or_alias in self._aliases:
            logical = self._aliases[name_or_alias]
            return self._models[logical]
        raise AIProviderError(
            message=f"Model or alias '{name_or_alias}' not found",
            code="MODEL_NOT_FOUND",
            status_code=404,
        )

    def resolve_name(self, name_or_alias: str) -> str:
        """解析为规范逻辑名（不返回完整配置）。"""
        if name_or_alias in self._models:
            return name_or_alias
        if name_or_alias in self._aliases:
            return self._aliases[name_or_alias]
        raise AIProviderError(
            message=f"Model or alias '{name_or_alias}' not found",
            code="MODEL_NOT_FOUND",
            status_code=404,
        )

    def get_default_name(self) -> str:
        if self._default is None:
            raise AIProviderError(
                message="No default model configured",
                code="NO_DEFAULT_MODEL",
                status_code=500,
            )
        return self._default

    def get_fallback_name(self) -> str | None:
        return self._fallback

    def contains(self, name: str) -> bool:
        return name in self._models

    def list_all(self) -> list[ModelConfig]:
        return list(self._models.values())

    def list_by_capability(self, capability: str) -> list[ModelConfig]:
        return list(self._capability_index.get(capability, []))

    def resolve_by_capability(self, capability: str) -> str:
        candidates = self._capability_index.get(capability)
        if not candidates:
            raise AIProviderError(
                message=f"No active models with capability '{capability}'",
                code="CAPABILITY_NOT_FOUND",
                status_code=404,
            )
        weights = [max(c.weight, 1) for c in candidates]
        chosen = random.choices(candidates, weights=weights, k=1)[0]
        return chosen.name


class ProviderRegistry:
    """内存提供商配置索引。"""

    def __init__(self, providers: dict[str, ProviderConfig]) -> None:
        self._providers = dict(providers)

    def get(self, name: str) -> ProviderConfig:
        prov = self._providers.get(name)
        if prov is None:
            raise AIProviderError(
                message=f"Provider '{name}' not configured",
                code="PROVIDER_NOT_FOUND",
                status_code=404,
            )
        return prov

    def list_all(self) -> dict[str, ProviderConfig]:
        return dict(self._providers)


class RegistryFactory:
    """工厂：从配置列表创建注册表。"""

    @staticmethod
    def create(
        models: list[ModelConfig],
        providers: dict[str, ProviderConfig],
        aliases: dict[str, str] | None = None,
        pricing_data: dict[str, dict] | None = None,
    ) -> RegistryHolder:
        return RegistryHolder(
            model_registry=ModelRegistry(models, aliases=aliases),
            provider_registry=ProviderRegistry(providers),
            pricing_registry=PricingFactory.from_dict(pricing_data or {}),
        )


class RegistryHolder:
    """同时持有所有注册表。"""

    def __init__(
        self,
        model_registry: ModelRegistry,
        provider_registry: ProviderRegistry,
        pricing_registry: PricingRegistry | None = None,
    ) -> None:
        self.model_registry = model_registry
        self.provider_registry = provider_registry
        self.pricing_registry = pricing_registry or PricingRegistry()