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__)
DEFAULT_TIMEOUT_MS = 180_000
MAX_OUTPUT_TOKENS = 8192


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

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


PROVIDERS = {
    "openai": LLMProviderConfig(
        name="openai",
        base_url="https://api.openai.com/v1",
        model="gpt-4o",
        max_retries=2,
        rate_limit=0,
    ),
    "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,
    ),
}


def _format_exception(exc: Exception) -> str:
    message = str(exc) or repr(exc)
    return f"{exc.__class__.__name__}: {message}"


def _is_retryable_llm_error(exc: Exception) -> bool:
    if isinstance(exc, httpx.TransportError):
        return True
    if isinstance(exc, httpx.HTTPStatusError):
        return exc.response.status_code in {408, 429, 500, 502, 503, 504}
    return False


def _parse_timeout_ms(value: Optional[str], default: int = DEFAULT_TIMEOUT_MS) -> int:
    try:
        timeout_ms = int(str(value or "").strip())
    except (TypeError, ValueError):
        return default
    return timeout_ms if timeout_ms > 0 else default


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

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


def _parse_json_content(content: str):
    """从模型输出中提取 JSON，兼容代码块、前后说明和尾逗号。"""
    text = (content or "").strip().lstrip("\ufeff")
    if not text:
        raise ValueError("模型返回了空内容")

    candidates = re.findall(
        r"```(?:json)?\s*([\s\S]*?)```",
        text,
        flags=re.IGNORECASE,
    )
    candidates.append(text)

    decoder = json.JSONDecoder(strict=False)
    last_error = None
    for candidate in candidates:
        cleaned = re.sub(r",\s*([\]}])", r"\1", candidate.strip())
        start_positions = [0]
        start_positions.extend(
            index for index, char in enumerate(cleaned) if char in "[{" and index != 0
        )
        for start in start_positions:
            try:
                data, _end = decoder.raw_decode(cleaned[start:])
                return data
            except json.JSONDecodeError as exc:
                last_error = exc

    detail = str(last_error) if last_error else "未找到 JSON 对象"
    raise ValueError(f"模型输出不是有效 JSON（{detail}）")


def _validate_required_fields(data, schema: Optional[dict], path: str = "$") -> None:
    """校验 JSON Schema 中声明的必填容器和字段，避免把空对象当成成功。"""
    if not schema:
        return
    schema_type = schema.get("type")
    if schema_type == "object":
        if not isinstance(data, dict):
            raise ValueError(f"{path} 应为对象")
        missing = [key for key in schema.get("required", []) if key not in data]
        if missing:
            raise ValueError(f"{path} 缺少必填字段: {', '.join(missing)}")
        for key, child_schema in schema.get("properties", {}).items():
            if key in data and data[key] is not None:
                _validate_required_fields(data[key], child_schema, f"{path}.{key}")
    elif schema_type == "array" and not isinstance(data, list):
        raise ValueError(f"{path} 应为数组")


def _normalize_schema_root(data, schema: Optional[dict]):
    """兼容模型把“单一必填数组”直接输出为数组或换了包装键的情况。"""
    if not schema or schema.get("type") != "object":
        return data

    required = schema.get("required", [])
    properties = schema.get("properties", {})
    if len(required) != 1 or properties.get(required[0], {}).get("type") != "array":
        return data

    required_key = required[0]
    if isinstance(data, list):
        return {required_key: data}
    if not isinstance(data, dict) or required_key in data or len(data) != 1:
        return data

    only_key, only_value = next(iter(data.items()))
    compact_key = required_key.replace("_", "")
    safe_aliases = {
        "data", "rows", "items", "result", compact_key,
        "benefits", "withdrawals", "benefitillustration", "withdrawalillustration",
    }
    if str(only_key).replace("_", "").lower() in safe_aliases and isinstance(only_value, list):
        return {required_key: only_value}
    return data


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

async def _call_provider(
    config: LLMProviderConfig,
    api_key: str,
    messages: list[dict],
    timeout_ms: int = 60_000,
    json_mode: bool = False,
    temperature: float = 0.3,
) -> 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": temperature,
            "max_tokens": MAX_OUTPUT_TOKENS,
        }
        if json_mode and config.name == "deepseek":
            body["response_format"] = {"type": "json_object"}
            # DeepSeek V4 默认开启 Thinking Mode。结构化提取若不显式关闭，
            # 可能耗尽输出额度后只返回 reasoning_content，content 为空。
            body["thinking"] = {"type": "disabled"}
    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": temperature, "maxOutputTokens": MAX_OUTPUT_TOKENS},
        }
        if json_mode:
            body["generationConfig"]["responseMimeType"] = "application/json"
    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": temperature,
            "max_tokens": MAX_OUTPUT_TOKENS,
        }
        if json_mode:
            body["response_format"] = {"type": "json_object"}

    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 客户端，支持自动切换和速率限制。

    config_prefix: 数据库配置键前缀，如 "ppt" 读取 ppt_llm_*，"poster" 读取 poster_llm_*。
    """

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

        self._load_env_config()

    def _load_env_config(self):
        """从环境变量加载默认配置。"""
        prefix = self._config_prefix.upper()
        deepseek_key = os.environ.get(f"{prefix}_LLM_API_KEY") or 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.debug("[LLMClient] 未从环境变量加载 API Key，将在首次调用时从数据库读取配置")

    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
        prefix = self._config_prefix
        try:
            from insurance.models.system_setting import SystemSetting
            settings = {s.key: s.value for s in SystemSetting.query.filter(
                SystemSetting.key.in_([
                    f"{prefix}_llm_provider", f"{prefix}_llm_model",
                    f"{prefix}_llm_api_key", f"{prefix}_llm_base_url",
                    f"{prefix}_llm_timeout_ms",
                ])
            ).all()}
            self._timeout_ms = _parse_timeout_ms(
                settings.get(f"{prefix}_llm_timeout_ms"),
                _parse_timeout_ms(os.getenv(f"{prefix.upper()}_LLM_TIMEOUT_MS")),
            )
            provider = settings.get(f"{prefix}_llm_provider", "").strip()
            if not provider:
                return

            model = settings.get(f"{prefix}_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:{prefix}] 使用 Dify 模式: {model}")
                return

            api_key = settings.get(f"{prefix}_llm_api_key", "").strip()
            if not api_key:
                return

            base_url = settings.get(f"{prefix}_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:{prefix}] 使用数据库配置: {provider}/{cfg.model}")
                return

            # 自定义供应商
            if not base_url:
                return

            # 安全校验：拒绝私网地址（SEC-P1-02）
            from insurance.utils.security import is_safe_base_url
            is_safe, err_msg = is_safe_base_url(base_url)
            if not is_safe:
                logger.warning(f"[LLMClient:{prefix}] 不安全的 Base URL: {err_msg}")
                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:{prefix}] 使用自定义模型: {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,
        temperature: float = 0.3,
    ) -> tuple[dict, LLMResponse]:
        """结构化输出（返回 JSON）。返回 (parsed_data, response)。

        Args:
            temperature: 输出温度。结构化提取建议 0，分析任务可用 0.3。
        """
        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_mode=True, temperature=temperature)
        validation_error = ""
        try:
            parsed = _normalize_schema_root(_parse_json_content(response.content), schema)
            _validate_required_fields(parsed, schema)
            return parsed, response
        except ValueError as first_error:
            validation_error = str(first_error)
            logger.warning(
                "[LLMClient] 首次结构化输出无效，正在请求模型纠正: %s",
                first_error,
            )

        repair_messages = messages + [
            {"role": "assistant", "content": (response.content or "")[:12000]},
            {
                "role": "user",
                "content": (
                    f"上一次回答未通过校验：{validation_error}。请严格按照以下 JSON Schema 重新输出：\n"
                    f"{json.dumps(schema or {}, ensure_ascii=False, indent=2)}\n"
                    "不要使用 Markdown 代码块，不要解释，不要省略必填字段。"
                ),
            },
        ]
        repaired_response = await self._call(repair_messages, json_mode=True, temperature=temperature)
        try:
            parsed = _normalize_schema_root(_parse_json_content(repaired_response.content), schema)
            _validate_required_fields(parsed, schema)
            return parsed, repaired_response
        except ValueError as repair_error:
            raise ValueError(f"[LLMClient] JSON 解析失败，纠正重试仍无效: {repair_error}")

    async def _call(
        self,
        messages: list[dict],
        attempt: int = 0,
        json_mode: bool = False,
        temperature: float = 0.3,
    ) -> LLMResponse:
        """多供应商自动切换调用。"""
        self._try_load_db_config()
        if not self._configs:
            raise RuntimeError(
                "未配置任何 LLM API Key，请在管理后台 > 系统配置中设置 "
                "ppt_llm_provider / ppt_llm_api_key，或设置环境变量 DEEPSEEK_API_KEY / OPENAI_API_KEY"
            )

        start_idx = self._active_idx
        tried = set()

        provider_errors = []
        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)

            # 限流
            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

            max_attempts = max(1, int(config.max_retries) + 1)
            for retry_index in range(max_attempts):
                try:
                    if config.name == "dify":
                        response = await _call_dify(
                            config.model, messages, timeout_ms=self._timeout_ms
                        )
                    else:
                        response = await _call_provider(
                            config,
                            api_key,
                            messages,
                            timeout_ms=self._timeout_ms,
                            json_mode=json_mode,
                            temperature=temperature,
                        )
                    self._active_idx = idx
                    return response
                except Exception as e:
                    error_text = _format_exception(e)
                    provider_errors.append(f"{config.name}: {error_text}")
                    should_retry = (
                        _is_retryable_llm_error(e)
                        and retry_index + 1 < max_attempts
                    )
                    if should_retry:
                        delay = min(2 ** retry_index, 4)
                        logger.warning(
                            f"[LLMClient] {config.name} 传输失败，"
                            f"{delay}s 后重试 {retry_index + 2}/{max_attempts}: {error_text}"
                        )
                        await asyncio.sleep(delay)
                        continue
                    logger.warning(
                        f"[LLMClient] {config.name} 失败 "
                        f"attempt={retry_index + 1}/{max_attempts}: {error_text}"
                    )
                    break

            if attempt < 3 and i < len(self._configs) - 1:
                self._active_idx = (idx + 1) % len(self._configs)

        provider_count = len(self._configs)
        last_error = provider_errors[-1] if provider_errors else "未知错误"
        raise RuntimeError(
            f"已配置的 {provider_count} 个 LLM 供应商均调用失败"
            f"（当前超时 {self._timeout_ms}ms，最后错误: {last_error}）"
        )

    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(config_prefix="ppt")
poster_llm_client = LLMClient(config_prefix="poster")