baodan/api/insurance/ppt/llm_client.py

"""统一 LLM 客户端 — 多供应商自动切换 + 限流保护。

支持: DeepSeek → MiniMax → Gemini
特点:
  - 限流保护：多用户并发时自动排队
  - 失败切换：一个供应商失败自动切换下一个
  - 统一接口：所有 LLM 调用走这里
"""
import os
import re
import json
import time
import asyncio
import logging
from dataclasses import dataclass, field
from typing import Optional

import httpx

logger = logging.getLogger(__name__)


# ─── 配置 ───────────────────────────────────────────────

@dataclass
class LLMProviderConfig:
    name: str
    base_url: str
    model: str
    max_retries: int = 2
    rate_limit: int = 0  # 每分钟请求数，0=无限制


PROVIDERS = {
    "deepseek": LLMProviderConfig(
        name="deepseek",
        base_url="https://api.deepseek.com/v1",
        model="deepseek-v4-pro",
        max_retries=2,
        rate_limit=0,
    ),
    "minimax": LLMProviderConfig(
        name="minimax",
        base_url="https://api.minimax.chat/v1",
        model="MiniMax-2.7-Flash",
        max_retries=2,
        rate_limit=30,
    ),
    "gemini": LLMProviderConfig(
        name="gemini",
        base_url="https://generativelanguage.googleapis.com/v1/models",
        model="gemini-2.5-flash",
        max_retries=1,
        rate_limit=60,
    ),
}


# ─── 限流器 ─────────────────────────────────────────────

class RateLimiter:
    """Token Bucket 限流器。"""

    def __init__(self, max_tokens: int, refill_ms: int = 60_000):
        self.max_tokens = max_tokens
        self.refill_ms = refill_ms / 1000  # 转换为秒
        self.tokens = [0.0] * max_tokens
        self.last_refill = time.monotonic()

    async def acquire(self, timeout_ms: int = 30_000) -> bool:
        """获取一个令牌，超时返回 False。"""
        if self.max_tokens == 0:
            return True

        start = time.monotonic()
        timeout_s = timeout_ms / 1000

        while True:
            self._refill_if_needed()
            now = time.monotonic()

            # 找最早可用的令牌槽位
            for i, t in enumerate(self.tokens):
                if now - t >= self.refill_ms:
                    self.tokens[i] = now
                    return True

            # 所有槽位都在使用中，等待最老的释放
            oldest = min(self.tokens)
            wait = min(self.refill_ms - (now - oldest), timeout_s)
            if wait <= 0 or (time.monotonic() - start) >= timeout_s:
                return False
            await asyncio.sleep(min(wait, 2.0))

    def _refill_if_needed(self):
        now = time.monotonic()
        if now - self.last_refill >= self.refill_ms:
            self.tokens = [0.0] * self.max_tokens
            self.last_refill = now


# ─── 响应类型 ────────────────────────────────────────────

@dataclass
class LLMResponse:
    content: str
    provider: str
    tokens: Optional[dict] = None  # {"input": int, "output": int}
    latency_ms: float = 0


# ─── 单供应商调用 ─────────────────────────────────────────

async def _call_provider(
    config: LLMProviderConfig,
    api_key: str,
    messages: list[dict],
    timeout_ms: int = 60_000,
) -> LLMResponse:
    """调用单个 LLM 供应商。"""
    start = time.monotonic()
    timeout_s = timeout_ms / 1000

    headers = {"Content-Type": "application/json"}

    if config.name in ("deepseek", "minimax"):
        # OpenAI 兼容格式
        headers["Authorization"] = f"Bearer {api_key}"
        if config.name == "minimax":
            url = f"{config.base_url}/text/chatcompletion_v2"
        else:
            url = f"{config.base_url}/chat/completions"
        body = {
            "model": config.model,
            "messages": messages,
            "temperature": 0.3,
            "max_tokens": 4096,
        }
    elif config.name == "gemini":
        # Gemini 格式
        model_part = f"{config.model}:generateContent" if ":" not in config.model else config.model
        url = f"{config.base_url}/{model_part}?key={api_key}"
        contents = []
        for m in messages:
            role = "model" if m["role"] == "assistant" else "user"
            contents.append({"role": role, "parts": [{"text": m["content"]}]})
        body = {
            "contents": contents,
            "generationConfig": {"temperature": 0.3, "maxOutputTokens": 4096},
        }
    else:
        # 自定义供应商：OpenAI 兼容格式
        headers["Authorization"] = f"Bearer {api_key}"
        base = config.base_url.rstrip("/")
        url = f"{base}/chat/completions"
        body = {
            "model": config.model,
            "messages": messages,
            "temperature": 0.3,
            "max_tokens": 4096,
        }

    async with httpx.AsyncClient(timeout=timeout_s) as client:
        resp = await client.post(url, json=body, headers=headers)
        resp.raise_for_status()
        data = resp.json()

    latency_ms = (time.monotonic() - start) * 1000

    # 解析响应
    if config.name == "gemini":
        content = ""
        candidates = data.get("candidates", [])
        if candidates:
            parts = candidates[0].get("content", {}).get("parts", [])
            if parts:
                content = parts[0].get("text", "")
        usage_meta = data.get("usageMetadata")
        tokens = None
        if usage_meta:
            tokens = {
                "input": usage_meta.get("promptTokenCount", 0),
                "output": usage_meta.get("candidatesTokenCount", 0),
            }
    else:
        # OpenAI 兼容格式（deepseek / minimax / 自定义供应商）
        content = ""
        choices = data.get("choices", [])
        if choices:
            content = choices[0].get("message", {}).get("content", "")
        usage = data.get("usage")
        tokens = None
        if usage:
            tokens = {
                "input": usage.get("prompt_tokens", 0),
                "output": usage.get("completion_tokens", 0),
            }

    return LLMResponse(
        content=content,
        provider=config.name,
        tokens=tokens,
        latency_ms=latency_ms,
    )


async def _call_dify(
    model: str,
    messages: list[dict],
    timeout_ms: int = 60_000,
) -> LLMResponse:
    """通过 Dify Chat API 调用模型（复用 Dify 已配置的 API Key）。"""
    import flask
    dify_base_url = flask.current_app.config.get("DIFY_BASE_URL", os.getenv("DIFY_BASE_URL", "http://localhost:5001"))
    dify_api_key = flask.current_app.config.get("BAODAN_CHAT_API_KEY", os.getenv("DIFY_CHAT_APP_API_KEY", ""))

    if not dify_api_key:
        raise RuntimeError("未配置 DIFY_CHAT_APP_API_KEY，无法使用 Dify 模式")

    # 合并 messages 为 query（Dify Chat API 不支持多轮 messages 格式）
    system_parts = []
    user_parts = []
    for m in messages:
        if m["role"] == "system":
            system_parts.append(m["content"])
        else:
            user_parts.append(m["content"])
    query = "\n\n".join(user_parts)
    if system_parts:
        query = "\n\n".join(system_parts) + "\n\n" + query

    # 如果有指定模型，在 query 前加上模型提示
    if model:
        query = f"[请使用模型 {model} 回答]\n\n{query}"

    start = time.monotonic()
    timeout_s = timeout_ms / 1000
    url = f"{dify_base_url.rstrip('/')}/v1/chat-messages"
    headers = {
        "Authorization": f"Bearer {dify_api_key}",
        "Content-Type": "application/json",
    }
    body = {
        "inputs": {},
        "query": query,
        "response_mode": "blocking",
        "user": "insurance-system",
    }

    async with httpx.AsyncClient(timeout=timeout_s) as client:
        resp = await client.post(url, json=body, headers=headers)
        resp.raise_for_status()
        data = resp.json()

    latency_ms = (time.monotonic() - start) * 1000
    content = data.get("answer", "")
    tokens = None
    usage = data.get("metadata", {}).get("usage", {})
    if usage:
        tokens = {
            "input": usage.get("prompt_tokens", 0),
            "output": usage.get("completion_tokens", 0),
        }

    return LLMResponse(content=content, provider="dify", tokens=tokens, latency_ms=latency_ms)


# ─── 统一客户端 ───────────────────────────────────────────

class LLMClient:
    """多供应商 LLM 客户端，支持自动切换和速率限制。"""

    def __init__(self):
        self._configs: list[tuple[LLMProviderConfig, str]] = []
        self._limiters: dict[str, RateLimiter] = {}
        self._active_idx = 0
        self._db_config_time: float = 0  # 上次从数据库加载配置的时间戳
        self._db_config_ttl: float = 60  # 配置缓存有效期（秒）

        self._load_env_config()

    def _load_env_config(self):
        """从环境变量加载默认配置。"""
        deepseek_key = os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("OPENAI_API_KEY", "")
        minimax_key = os.environ.get("MINIMAX_API_KEY", "")
        gemini_key = os.environ.get("GEMINI_API_KEY", "")

        if deepseek_key:
            self._configs.append((PROVIDERS["deepseek"], deepseek_key))
            self._limiters["deepseek"] = RateLimiter(PROVIDERS["deepseek"].rate_limit)
        if minimax_key:
            self._configs.append((PROVIDERS["minimax"], minimax_key))
            self._limiters["minimax"] = RateLimiter(PROVIDERS["minimax"].rate_limit)
        if gemini_key:
            self._configs.append((PROVIDERS["gemini"], gemini_key))
            self._limiters["gemini"] = RateLimiter(PROVIDERS["gemini"].rate_limit)

        if not self._configs:
            logger.warning("[LLMClient] 未配置任何 API Key，使用 mock 模式")

    def _try_load_db_config(self):
        """尝试从数据库加载模型配置（带缓存，不阻塞）。"""
        now = time.monotonic()
        if now - self._db_config_time < self._db_config_ttl:
            return
        self._db_config_time = now
        try:
            from insurance.models.system_setting import SystemSetting
            settings = {s.key: s.value for s in SystemSetting.query.filter(
                SystemSetting.key.in_([
                    "ppt_llm_provider", "ppt_llm_model", "ppt_llm_api_key", "ppt_llm_base_url",
                ])
            ).all()}
            provider = settings.get("ppt_llm_provider", "").strip()
            if not provider:
                return

            model = settings.get("ppt_llm_model", "").strip()

            # Dify 模式：通过 Dify Chat API 调用，无需独立 API Key
            if provider == "dify":
                cfg = LLMProviderConfig(
                    name="dify", base_url="", model=model or "",
                )
                self._configs = [(cfg, "")]
                self._limiters = {"dify": RateLimiter(0)}
                self._active_idx = 0
                logger.info(f"[LLMClient] 使用 Dify 模式: {model}")
                return

            api_key = settings.get("ppt_llm_api_key", "").strip()
            if not api_key:
                return

            base_url = settings.get("ppt_llm_base_url", "").strip()

            # 内置供应商：替换对应配置
            if provider in PROVIDERS and not base_url:
                cfg = PROVIDERS[provider]
                if model:
                    cfg = LLMProviderConfig(
                        name=cfg.name, base_url=cfg.base_url, model=model,
                        max_retries=cfg.max_retries, rate_limit=cfg.rate_limit,
                    )
                self._configs = [(cfg, api_key)]
                self._limiters = {provider: RateLimiter(cfg.rate_limit)}
                self._active_idx = 0
                logger.info(f"[LLMClient] 使用数据库配置: {provider}/{cfg.model}")
                return

            # 自定义供应商
            if not base_url:
                return
            cfg = LLMProviderConfig(
                name=provider, base_url=base_url, model=model or "gpt-4o-mini",
            )
            self._configs = [(cfg, api_key)]
            self._limiters = {provider: RateLimiter(0)}
            self._active_idx = 0
            logger.info(f"[LLMClient] 使用自定义模型: {provider}/{cfg.model}")
        except Exception:
            pass  # 无 Flask 上下文或数据库不可用，使用环境变量配置

    async def chat(self, prompt: str, system_prompt: str = "") -> LLMResponse:
        """简单聊天。"""
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})
        return await self._call(messages)

    async def structured_output(
        self,
        prompt: str,
        system_prompt: str = "",
        schema: Optional[dict] = None,
    ) -> tuple[dict, LLMResponse]:
        """结构化输出（返回 JSON）。返回 (parsed_data, response)。"""
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})

        full_prompt = prompt
        if schema:
            full_prompt += f"\n\n请以JSON格式输出，格式如下：\n{json.dumps(schema, ensure_ascii=False, indent=2)}"
            full_prompt += "\n重要：只输出JSON，不要任何额外文字。"
        messages.append({"role": "user", "content": full_prompt})

        response = await self._call(messages)

        # 解析 JSON
        json_str = response.content.strip()
        # 尝试提取 markdown 代码块或裸 JSON
        match = re.search(r"```(?:json)?\s*([\s\S]*?)```|(\{[\s\S]*\}|\[[\s\S]*\])$", json_str)
        if match:
            json_str = match.group(1) or match.group(2)

        try:
            data = json.loads(json_str)
            return data, response
        except json.JSONDecodeError:
            # 清理尾逗号
            json_str = re.sub(r",\s*([\]}])", r"\1", json_str)
            try:
                data = json.loads(json_str)
                return data, response
            except json.JSONDecodeError:
                raise ValueError(f"[LLMClient] JSON 解析失败: {json_str[:200]}")

    async def _call(self, messages: list[dict], attempt: int = 0) -> LLMResponse:
        """多供应商自动切换调用。"""
        self._try_load_db_config()
        if not self._configs:
            return LLMResponse(content="{}", provider="mock", latency_ms=0)

        start_idx = self._active_idx
        tried = set()

        for i in range(len(self._configs)):
            idx = (start_idx + i) % len(self._configs)
            config, api_key = self._configs[idx]

            if config.name in tried:
                continue
            tried.add(config.name)

            # Dify 模式：通过 Dify Chat API 调用
            if config.name == "dify":
                try:
                    return await _call_dify(config.model, messages, timeout_ms=60_000)
                except Exception as e:
                    logger.warning(f"[LLMClient] Dify 调用失败: {e}")
                    continue

            # 限流
            limiter = self._limiters.get(config.name)
            if limiter:
                acquired = await limiter.acquire(timeout_ms=30_000)
                if not acquired:
                    logger.warning(f"[LLMClient] {config.name} 限流超时")
                    continue

            # 调用
            try:
                response = await _call_provider(config, api_key, messages, timeout_ms=60_000)
                self._active_idx = idx
                return response
            except Exception as e:
                logger.warning(f"[LLMClient] {config.name} 失败: {e}")
                if attempt < 3 and i < len(self._configs) - 1:
                    self._active_idx = (idx + 1) % len(self._configs)

        raise RuntimeError("所有 LLM 供应商均失败")

    def get_status(self) -> dict:
        """获取当前供应商信息。"""
        available = [c[0].name for c in self._configs]
        active = self._configs[self._active_idx][0].name if self._configs else "none"
        return {"available": available, "active": active}


# ─── 单例 ────────────────────────────────────────────────
llm_client = LLMClient()