dataClean/app/ai_client.py

"""LLM API 客户端，兼容 OpenAI API 格式"""
import json
import logging
import re
from typing import Optional

from openai import OpenAI, APIError

from config import settings

logger = logging.getLogger(__name__)

# 匹配 reasoning 模型（MiniMax-M3 / DeepSeek-R1 / GLM-Z1 等）的 <think>...</think> 推理块
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)


def _parse_llm_json(content: str) -> dict:
    """从 LLM 输出中提取 JSON。

    兼容 reasoning 模型在 json_object 模式下仍输出 <think>...</think>
    推理块、以及 JSON 前后有多余文本的情况。
    """
    if not content or not content.strip():
        raise ValueError("LLM 返回空内容，无法解析 JSON")

    text = content.strip()
    # 1) 去掉闭合的 <think>...</think> 块
    text = _THINK_RE.sub("", text).strip()
    # 2) 处理只有 <think> 开头但未闭合（content 被截断）的情况
    if text.startswith("<think>"):
        text = text.split("</think>", 1)[-1].strip()

    # 3) 尝试直接解析
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 4) 提取首个 { 到最后 } 之间的子串
    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end > start:
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            pass

    # 5) 兜底：尝试数组
    start = text.find("[")
    end = text.rfind("]")
    if start != -1 and end > start:
        return json.loads(text[start : end + 1])

    logger.error("无法从 LLM 输出提取 JSON: %s", content[:500])
    raise ValueError("LLM 输出无法解析为 JSON")


class AIClient:
    """封装 LLM 调用，支持重试和 JSON 输出。

    配置以 property 形式运行时从 settings 读取，避免模块 import 时
    固化旧值（settings 在 FastAPI lifespan 启动后才会被数据库配置覆盖）。
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        timeout: Optional[int] = None,
        max_retries: Optional[int] = None,
    ):
        # 仅保存显式传入的覆盖值；为 None 时运行时回退到 settings
        self._api_key = api_key
        self._base_url = base_url
        self._model = model
        self._timeout = timeout
        self._max_retries = max_retries

    @property
    def api_key(self) -> str:
        return self._api_key or settings.OPENAI_API_KEY

    @property
    def base_url(self) -> str:
        return self._base_url or settings.OPENAI_BASE_URL

    @property
    def model(self) -> str:
        return self._model or settings.OPENAI_MODEL

    @property
    def timeout(self) -> int:
        return self._timeout or settings.OPENAI_TIMEOUT

    @property
    def max_retries(self) -> int:
        return self._max_retries or settings.OPENAI_MAX_RETRIES

    @property
    def client(self) -> OpenAI:
        # 每次按最新配置创建，确保用到启动后覆盖的真实配置
        return OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            timeout=self.timeout,
            max_retries=self.max_retries,
        )

    def chat_completion(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.3,
        json_mode: bool = False,
    ) -> str:
        """调用 LLM 返回文本"""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

        kwargs = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}

        try:
            resp = self.client.chat.completions.create(**kwargs)
            content = resp.choices[0].message.content or ""
            return content.strip()
        except APIError as exc:
            logger.error("LLM API 调用失败: %s", exc)
            raise

    def chat_completion_json(
        self,
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.3,
    ) -> dict:
        """调用 LLM 并解析返回的 JSON（兼容 reasoning 模型的 <think> 块）"""
        content = self.chat_completion(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            temperature=temperature,
            json_mode=True,
        )
        return _parse_llm_json(content)


ai_client = AIClient()