Codex
153 lines
3.9 KiB
Python
153 lines
3.9 KiB
Python
"""
|
|
提供商注册中心
|
|
"""
|
|
import logging
|
|
from typing import Dict, List, Optional, Type
|
|
|
|
from .base import BaseLLMProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ProviderRegistry:
|
|
"""
|
|
提供商注册中心
|
|
|
|
管理所有 LLM 提供商适配器的注册和获取
|
|
"""
|
|
|
|
_providers: Dict[str, Type[BaseLLMProvider]] = {}
|
|
|
|
@classmethod
|
|
def register(cls, provider_class: Type[BaseLLMProvider]) -> Type[BaseLLMProvider]:
|
|
"""
|
|
注册提供商（可作为装饰器使用）
|
|
|
|
Args:
|
|
provider_class: 提供商类
|
|
|
|
Returns:
|
|
提供商类
|
|
"""
|
|
provider_type = provider_class.provider_type
|
|
if not provider_type:
|
|
raise ValueError(f'Provider {provider_class.__name__} must have a provider_type')
|
|
|
|
cls._providers[provider_type] = provider_class
|
|
logger.info(f'Registered LLM provider: {provider_type}')
|
|
return provider_class
|
|
|
|
@classmethod
|
|
def get(cls, provider_type: str) -> Optional[Type[BaseLLMProvider]]:
|
|
"""
|
|
获取提供商类
|
|
|
|
Args:
|
|
provider_type: 提供商类型
|
|
|
|
Returns:
|
|
提供商类或 None
|
|
"""
|
|
return cls._providers.get(provider_type)
|
|
|
|
@classmethod
|
|
def create_instance(
|
|
cls,
|
|
provider_type: str,
|
|
api_key: str = '',
|
|
api_base: str = '',
|
|
**kwargs
|
|
) -> Optional[BaseLLMProvider]:
|
|
"""
|
|
创建提供商实例
|
|
|
|
Args:
|
|
provider_type: 提供商类型
|
|
api_key: API Key
|
|
api_base: API 地址
|
|
**kwargs: 其他配置
|
|
|
|
Returns:
|
|
提供商实例或 None
|
|
"""
|
|
provider_class = cls.get(provider_type)
|
|
if not provider_class:
|
|
logger.warning(f'Unknown provider type: {provider_type}')
|
|
return None
|
|
|
|
return provider_class(api_key=api_key, api_base=api_base, **kwargs)
|
|
|
|
@classmethod
|
|
def get_all_types(cls) -> List[Dict[str, str]]:
|
|
"""
|
|
获取所有已注册的提供商类型
|
|
|
|
Returns:
|
|
提供商类型列表
|
|
"""
|
|
return [
|
|
{
|
|
'type': provider_type,
|
|
'name': provider_class.provider_name,
|
|
}
|
|
for provider_type, provider_class in cls._providers.items()
|
|
]
|
|
|
|
@classmethod
|
|
def get_default_models(cls, provider_type: str) -> List[Dict]:
|
|
"""
|
|
获取提供商的默认模型列表
|
|
|
|
Args:
|
|
provider_type: 提供商类型
|
|
|
|
Returns:
|
|
默认模型列表
|
|
"""
|
|
provider_class = cls.get(provider_type)
|
|
if not provider_class:
|
|
return []
|
|
return provider_class.get_default_models()
|
|
|
|
@classmethod
|
|
async def fetch_models_from_api(
|
|
cls,
|
|
provider_type: str,
|
|
api_key: str = '',
|
|
api_base: str = '',
|
|
**kwargs
|
|
) -> List[Dict]:
|
|
"""
|
|
通过 API 在线拉取提供商的最新模型列表
|
|
|
|
Args:
|
|
provider_type: 提供商类型
|
|
api_key: API Key
|
|
api_base: API 地址
|
|
|
|
Returns:
|
|
模型列表，失败时返回空列表
|
|
"""
|
|
instance = cls.create_instance(provider_type, api_key=api_key, api_base=api_base, **kwargs)
|
|
if not instance:
|
|
return []
|
|
return await instance.fetch_models_from_api()
|
|
|
|
|
|
# 自动加载所有提供商
|
|
def _load_providers():
|
|
"""加载所有提供商适配器"""
|
|
from . import openai_provider
|
|
from . import claude_provider
|
|
from . import qwen_provider
|
|
from . import ollama_provider
|
|
from . import deepseek_provider
|
|
from . import siliconflow_provider
|
|
|
|
|
|
# 延迟加载
|
|
try:
|
|
_load_providers()
|
|
except ImportError as e:
|
|
logger.warning(f'Failed to load some providers: {e}')
|