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()