SillyTavern_replica/backend/services/llm_model_service.py

"""
LLM 模型管理服务

提供获取不同 LLM 提供商可用模型列表的功能
"""
from typing import List, Dict, Any, Optional
import requests


class LLMModelService:
    """LLM 模型管理服务"""
    
    @staticmethod
    def get_openai_models(api_key: str, base_url: Optional[str] = None) -> List[str]:
        """
        获取 OpenAI 兼容 API 的模型列表
        
        Args:
            api_key: API Key
            base_url: API 基础 URL，默认为 OpenAI 官方 API
            
        Returns:
            模型名称列表
        """
        try:
            # 默认使用 OpenAI 官方 API
            if not base_url:
                base_url = "https://api.openai.com/v1"
            
            # 规范化 base_url：确保有协议前缀
            base_url = base_url.strip()
            if not base_url.startswith(('http://', 'https://')):
                base_url = 'https://' + base_url
            
            # 移除末尾的斜杠和常见 endpoint 路径
            base_url = base_url.rstrip('/')
            # 移除可能已经存在的 endpoint 路径
            for endpoint in ['/chat/completions', '/completions', '/embeddings', '/models']:
                if base_url.endswith(endpoint):
                    base_url = base_url[:-len(endpoint)]
                    break
            
            # 调用 models API
            try:
                response = requests.get(
                    f"{base_url}/models",
                    headers={
                        "Authorization": f"Bearer {api_key}",
                        "Content-Type": "application/json"
                    },
                    timeout=10
                )
            except requests.exceptions.InvalidSchema as e:
                raise Exception(f"URL 格式错误: {base_url}/models - 请确保 URL 以 http:// 或 https:// 开头")
            except requests.exceptions.ConnectionError as e:
                raise Exception(f"无法连接到 API: {base_url}/models - 请检查网络连接和 API 地址")
            except requests.exceptions.Timeout as e:
                raise Exception(f"请求超时: {base_url}/models - 请检查网络连接")
            
            if response.status_code != 200:
                raise Exception(f"HTTP {response.status_code}: {response.text}")
            
            data = response.json()
            models = [model['id'] for model in data.get('data', [])]
            
            # 返回所有模型，不做过滤
            return models
            
        except Exception as e:
            raise Exception(f"获取 OpenAI 模型列表失败: {str(e)}")
    
    @staticmethod
    def get_anthropic_models(api_key: str) -> List[str]:
        """
        获取 Anthropic Claude 模型列表
        
        Args:
            api_key: API Key
            
        Returns:
            模型名称列表
        """
        try:
            # Anthropic 没有公开的模型列表 API，返回已知模型
            return [
                "claude-3-5-sonnet-20241022",
                "claude-3-5-haiku-20241022",
                "claude-3-opus-20240229",
                "claude-3-sonnet-20240229",
                "claude-3-haiku-20240307",
                "claude-2.1",
                "claude-2.0",
                "claude-instant-1.2"
            ]
            
        except Exception as e:
            raise Exception(f"获取 Anthropic 模型列表失败: {str(e)}")
    
    @staticmethod
    def get_ollama_models(base_url: str = "http://localhost:11434") -> List[str]:
        """
        获取 Ollama 本地模型列表
        
        Args:
            base_url: Ollama API 地址
            
        Returns:
            模型名称列表
        """
        try:
            response = requests.get(
                f"{base_url}/api/tags",
                timeout=10
            )
            
            if response.status_code != 200:
                raise Exception(f"HTTP {response.status_code}: {response.text}")
            
            data = response.json()
            models = [model['name'] for model in data.get('models', [])]
            
            return models
            
        except Exception as e:
            raise Exception(f"获取 Ollama 模型列表失败: {str(e)}")
    
    @staticmethod
    def detect_provider(api_url: str) -> str:
        """
        根据 API URL 检测提供商类型
        
        Args:
            api_url: API 地址
            
        Returns:
            提供商类型: 'openai', 'anthropic', 'ollama', 'unknown'
        """
        api_url_lower = api_url.lower()
        
        if 'openai' in api_url_lower or 'api.openai.com' in api_url_lower:
            return 'openai'
        elif 'anthropic' in api_url_lower or 'api.anthropic.com' in api_url_lower:
            return 'anthropic'
        elif 'ollama' in api_url_lower or 'localhost:11434' in api_url_lower or '127.0.0.1:11434' in api_url_lower:
            return 'ollama'
        elif 'bigmodel' in api_url_lower or 'glm' in api_url_lower:
            # 智谱AI GLM - 兼容 OpenAI API
            return 'openai'
        elif 'siliconflow' in api_url_lower or 'silicon.cloud' in api_url_lower:
            # SiliconFlow 等兼容 OpenAI API 的服务
            return 'openai'
        elif 'deepseek' in api_url_lower:
            # DeepSeek 等兼容 OpenAI API 的服务
            return 'openai'
        else:
            # 默认尝试 OpenAI 兼容 API
            return 'openai'
    
    @staticmethod
    def get_models_by_provider(
        provider: str,
        api_key: str,
        api_url: Optional[str] = None
    ) -> List[str]:
        """
        根据提供商类型获取模型列表
        
        Args:
            provider: 提供商类型 ('openai', 'anthropic', 'ollama')
            api_key: API Key
            api_url: API 地址（可选）
            
        Returns:
            模型名称列表
        """
        if provider == 'openai':
            return LLMModelService.get_openai_models(api_key, api_url)
        elif provider == 'anthropic':
            return LLMModelService.get_anthropic_models(api_key)
        elif provider == 'ollama':
            base_url = api_url or "http://localhost:11434"
            # 移除 /v1 后缀（如果有）
            base_url = base_url.replace('/v1', '').replace('/v1/', '')
            return LLMModelService.get_ollama_models(base_url)
        else:
            raise Exception(f"不支持的提供商: {provider}")