nex_docus/backend/app/services/llm_provider_service.py

"""
LLM 提供方测试服务
"""
import asyncio
import json
import socket
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Dict, List, Optional


PROVIDER_CATALOG: List[Dict[str, str]] = [
    {
        "value": "openai",
        "label": "OpenAI",
        "default_endpoint_url": "https://api.openai.com/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "deepseek",
        "label": "DeepSeek",
        "default_endpoint_url": "https://api.deepseek.com/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "anthropic",
        "label": "Anthropic",
        "default_endpoint_url": "https://api.anthropic.com",
        "protocol": "anthropic",
    },
    {
        "value": "gemini",
        "label": "Google Gemini",
        "default_endpoint_url": "https://generativelanguage.googleapis.com/v1beta",
        "protocol": "gemini",
    },
    {
        "value": "dashscope",
        "label": "阿里百炼",
        "default_endpoint_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "zhipu",
        "label": "智谱 AI",
        "default_endpoint_url": "https://open.bigmodel.cn/api/paas/v4",
        "protocol": "openai_compatible",
    },
    {
        "value": "moonshot",
        "label": "Moonshot AI",
        "default_endpoint_url": "https://api.moonshot.cn/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "groq",
        "label": "Groq",
        "default_endpoint_url": "https://api.groq.com/openai/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "openrouter",
        "label": "OpenRouter",
        "default_endpoint_url": "https://openrouter.ai/api/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "siliconflow",
        "label": "SiliconFlow",
        "default_endpoint_url": "https://api.siliconflow.cn/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "ollama",
        "label": "Ollama",
        "default_endpoint_url": "http://localhost:11434/v1",
        "protocol": "openai_compatible",
    },
    {
        "value": "ark",
        "label": "火山方舟",
        "default_endpoint_url": "https://ark.cn-beijing.volces.com/api/v3",
        "protocol": "openai_compatible",
    },
    {
        "value": "custom",
        "label": "自定义兼容接口",
        "default_endpoint_url": "",
        "protocol": "openai_compatible",
    },
]

PROVIDER_MAP = {item["value"]: item for item in PROVIDER_CATALOG}


class LLMProviderService:
    """LLM 提供方测试服务"""

    @staticmethod
    def get_provider_catalog() -> List[Dict[str, str]]:
        return PROVIDER_CATALOG

    @staticmethod
    def get_provider_label(provider: Optional[str]) -> str:
        if not provider:
            return "自定义模型"
        return PROVIDER_MAP.get(provider, {}).get("label", provider)

    @staticmethod
    def get_default_endpoint_url(provider: Optional[str]) -> str:
        if not provider:
            return ""
        return PROVIDER_MAP.get(provider, {}).get("default_endpoint_url", "")

    @classmethod
    def build_model_name(cls, provider: Optional[str], llm_model_name: str) -> str:
        model_name = (llm_model_name or "").strip()
        label = cls.get_provider_label(provider)
        if not model_name:
            return label
        return f"{label} {model_name}"

    @staticmethod
    def build_model_code(provider: Optional[str], llm_model_name: str) -> str:
        provider_part = (provider or "custom").strip().lower()
        model_part = (llm_model_name or "").strip().lower()
        sanitized = []
        previous_is_separator = False
        for char in model_part:
            if char.isalnum():
                sanitized.append(char)
                previous_is_separator = False
            else:
                if not previous_is_separator:
                    sanitized.append("_")
                    previous_is_separator = True

        model_slug = "".join(sanitized).strip("_") or "model"
        return f"llm_{provider_part}_{model_slug}"

    @classmethod
    async def test_model_connection(cls, payload: Dict[str, Any]) -> Dict[str, Any]:
        provider = payload.get("provider")
        provider_meta = PROVIDER_MAP.get(provider, PROVIDER_MAP["custom"])
        endpoint_url = (payload.get("endpoint_url") or provider_meta.get("default_endpoint_url") or "").strip()
        llm_model_name = (payload.get("llm_model_name") or "").strip()
        api_key = (payload.get("api_key") or "").strip()

        if not endpoint_url:
            raise ValueError("缺少接口地址，请先选择提供方或手动填写 base_url")
        if not llm_model_name:
            raise ValueError("缺少模型名称，请填写 llm_model_name")
        if provider not in {"ollama"} and not api_key:
            raise ValueError("缺少 API Key，请填写后再测试")

        timeout = int(payload.get("llm_timeout") or 120)
        temperature = float(payload.get("llm_temperature") or 0.7)
        top_p = float(payload.get("llm_top_p") or 0.9)
        max_tokens = int(payload.get("llm_max_tokens") or 2048)
        system_prompt = (payload.get("llm_system_prompt") or "").strip()

        started_at = time.perf_counter()
        preview = await asyncio.to_thread(
            cls._send_test_request,
            provider_meta.get("protocol", "openai_compatible"),
            provider,
            endpoint_url,
            api_key,
            llm_model_name,
            timeout,
            temperature,
            top_p,
            max_tokens,
            system_prompt,
        )
        latency_ms = int((time.perf_counter() - started_at) * 1000)

        return {
            "provider": provider,
            "endpoint_url": endpoint_url,
            "llm_model_name": llm_model_name,
            "latency_ms": latency_ms,
            "preview": preview[:200],
        }

    @classmethod
    def _send_test_request(
        cls,
        protocol: str,
        provider: Optional[str],
        endpoint_url: str,
        api_key: str,
        llm_model_name: str,
        timeout: int,
        temperature: float,
        top_p: float,
        max_tokens: int,
        system_prompt: str,
    ) -> str:
        if protocol == "anthropic":
            return cls._test_anthropic(
                endpoint_url,
                api_key,
                llm_model_name,
                timeout,
                temperature,
                top_p,
                max_tokens,
                system_prompt,
            )

        if protocol == "gemini":
            return cls._test_gemini(
                endpoint_url,
                api_key,
                llm_model_name,
                timeout,
                temperature,
                top_p,
                max_tokens,
                system_prompt,
            )

        return cls._test_openai_compatible(
            provider,
            endpoint_url,
            api_key,
            llm_model_name,
            timeout,
            temperature,
            top_p,
            max_tokens,
            system_prompt,
        )

    @classmethod
    def _test_openai_compatible(
        cls,
        provider: Optional[str],
        endpoint_url: str,
        api_key: str,
        llm_model_name: str,
        timeout: int,
        temperature: float,
        top_p: float,
        max_tokens: int,
        system_prompt: str,
    ) -> str:
        url = cls._join_endpoint(endpoint_url, "/chat/completions")
        headers = {
            "Content-Type": "application/json",
        }
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        payload = {
            "model": llm_model_name,
            "messages": cls._build_openai_messages(system_prompt),
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": min(max_tokens, 256),
            "stream": False,
        }
        response = cls._request_json(url, headers, payload, timeout)
        choices = response.get("choices") or []
        if not choices:
            raise ValueError("测试请求已发送，但未收到模型返回内容")

        content = choices[0].get("message", {}).get("content", "")
        if isinstance(content, list):
            content = "".join(
                item.get("text", "") if isinstance(item, dict) else str(item)
                for item in content
            )
        content = str(content).strip()
        if not content:
            raise ValueError("模型返回成功，但内容为空")
        return content

    @classmethod
    def _test_anthropic(
        cls,
        endpoint_url: str,
        api_key: str,
        llm_model_name: str,
        timeout: int,
        temperature: float,
        top_p: float,
        max_tokens: int,
        system_prompt: str,
    ) -> str:
        url = cls._join_endpoint(endpoint_url, "/v1/messages")
        headers = {
            "Content-Type": "application/json",
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
        }
        payload = {
            "model": llm_model_name,
            "max_tokens": min(max_tokens, 256),
            "temperature": temperature,
            "top_p": top_p,
            "messages": [
                {
                    "role": "user",
                    "content": "请只回复“连接测试成功”。",
                }
            ],
        }
        if system_prompt:
            payload["system"] = system_prompt

        response = cls._request_json(url, headers, payload, timeout)
        content = response.get("content") or []
        texts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                texts.append(item.get("text", ""))
        preview = "".join(texts).strip()
        if not preview:
            raise ValueError("模型返回成功，但内容为空")
        return preview

    @classmethod
    def _test_gemini(
        cls,
        endpoint_url: str,
        api_key: str,
        llm_model_name: str,
        timeout: int,
        temperature: float,
        top_p: float,
        max_tokens: int,
        system_prompt: str,
    ) -> str:
        model_path = llm_model_name if llm_model_name.startswith("models/") else f"models/{llm_model_name}"
        encoded_model_path = "/".join(urllib.parse.quote(part) for part in model_path.split("/"))
        url = f"{endpoint_url.rstrip('/')}/{encoded_model_path}:generateContent?key={urllib.parse.quote(api_key)}"
        headers = {
            "Content-Type": "application/json",
        }
        payload = {
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "请只回复“连接测试成功”。"}],
                }
            ],
            "generationConfig": {
                "temperature": temperature,
                "topP": top_p,
                "maxOutputTokens": min(max_tokens, 256),
            },
        }
        if system_prompt:
            payload["systemInstruction"] = {
                "parts": [{"text": system_prompt}],
            }

        response = cls._request_json(url, headers, payload, timeout)
        candidates = response.get("candidates") or []
        if not candidates:
            raise ValueError("测试请求已发送，但未收到模型返回内容")

        parts = candidates[0].get("content", {}).get("parts", [])
        preview = "".join(
            part.get("text", "") for part in parts if isinstance(part, dict)
        ).strip()
        if not preview:
            raise ValueError("模型返回成功，但内容为空")
        return preview

    @staticmethod
    def _build_openai_messages(system_prompt: str) -> List[Dict[str, str]]:
        messages: List[Dict[str, str]] = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": "请只回复“连接测试成功”。"})
        return messages

    @staticmethod
    def _join_endpoint(base_url: str, suffix: str) -> str:
        normalized = base_url.rstrip("/")
        if normalized.endswith(suffix):
            return normalized
        return f"{normalized}{suffix}"

    @classmethod
    def _request_json(
        cls,
        url: str,
        headers: Dict[str, str],
        payload: Dict[str, Any],
        timeout: int,
    ) -> Dict[str, Any]:
        request = urllib.request.Request(
            url,
            data=json.dumps(payload).encode("utf-8"),
            headers=headers,
            method="POST",
        )

        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                body = response.read().decode("utf-8")
                if not body:
                    return {}
                return json.loads(body)
        except urllib.error.HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="ignore")
            message = cls._extract_error_message(error_body) or error_body[:300] or str(exc)
            raise ValueError(f"模型测试失败（HTTP {exc.code}）：{message}") from exc
        except urllib.error.URLError as exc:
            reason = exc.reason
            if isinstance(reason, socket.timeout):
                raise ValueError("模型测试超时，请检查网络或调大超时时间") from exc
            raise ValueError(f"模型测试失败：{reason}") from exc
        except socket.timeout as exc:
            raise ValueError("模型测试超时，请检查网络或调大超时时间") from exc
        except json.JSONDecodeError as exc:
            raise ValueError("模型服务返回了无法解析的响应，请检查接口地址是否正确") from exc

    @staticmethod
    def _extract_error_message(error_body: str) -> Optional[str]:
        if not error_body:
            return None

        try:
            payload = json.loads(error_body)
        except json.JSONDecodeError:
            return error_body.strip()

        if isinstance(payload, dict):
            if isinstance(payload.get("error"), dict):
                return payload["error"].get("message") or payload["error"].get("type")
            return payload.get("message") or payload.get("detail")

        return None