AIPortPilot/backend/app/services/llm_client.py

"""LLM 客户端 — 千问 DashScope OpenAI 兼容模式。

提供统一的 LLM 调用接口，全部使用流式输出（SSE）。
"""

import json
import logging
from collections.abc import AsyncGenerator
from typing import Any

import httpx

from app.core.config import settings

logger = logging.getLogger(__name__)


class LLMClient:
    """千问 LLM 客户端（OpenAI 兼容接口，流式输出）。"""

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        model: str | None = None,
        timeout: int | None = None,
    ):
        self.api_key = api_key or settings.llm_api_key
        self.base_url = base_url or settings.llm_base_url
        self.model = model or settings.llm_model
        self.timeout = timeout or settings.llm_timeout_seconds

    def _build_headers(self) -> dict[str, str]:
        """构建请求头。"""
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    def _build_payload(
        self,
        messages: list[dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> dict[str, Any]:
        """构建请求体。"""
        return {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": True,
        }

    async def chat_stream(
        self,
        messages: list[dict[str, str]],
        temperature: float = 0.3,
        max_tokens: int = 2000,
    ) -> AsyncGenerator[str, None]:
        """流式 chat completion，逐 token yield。

        Args:
            messages: OpenAI 格式的消息列表
            temperature: 温度参数
            max_tokens: 最大 token 数

        Yields:
            每个 token 的文本片段

        Raises:
            RuntimeError: API 调用失败
        """
        if not self.api_key:
            raise RuntimeError("LLM_API_KEY 未配置")

        headers = self._build_headers()
        payload = self._build_payload(messages, temperature, max_tokens)

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                async with client.stream(
                    "POST",
                    f"{self.base_url}/chat/completions",
                    headers=headers,
                    json=payload,
                ) as response:
                    response.raise_for_status()
                    async for line in response.aiter_lines():
                        if not line.startswith("data: "):
                            continue
                        data_str = line[6:]
                        if data_str.strip() == "[DONE]":
                            break
                        try:
                            chunk = json.loads(data_str)
                            delta = chunk.get("choices", [{}])[0].get("delta", {})
                            content = delta.get("content", "")
                            if content:
                                yield content
                        except json.JSONDecodeError:
                            continue
        except httpx.TimeoutException:
            logger.error("LLM 流式请求超时")
            raise RuntimeError("LLM 请求超时")
        except httpx.HTTPStatusError as e:
            logger.error("LLM API 错误: %s", e.response.status_code)
            raise RuntimeError(f"LLM API 错误: {e.response.status_code}")
        except RuntimeError:
            raise
        except Exception as e:
            logger.error("LLM 流式调用异常: %s", e)
            raise RuntimeError(f"LLM 调用失败: {e}")

    async def chat(
        self,
        messages: list[dict[str, str]],
        temperature: float = 0.3,
        max_tokens: int = 2000,
    ) -> str:
        """流式调用但收集为完整字符串（兼容非流式调用方）。"""
        parts: list[str] = []
        async for token in self.chat_stream(messages, temperature, max_tokens):
            parts.append(token)
        return "".join(parts)

    async def chat_json_stream(
        self,
        messages: list[dict[str, str]],
        temperature: float = 0.3,
        max_tokens: int = 2000,
    ) -> AsyncGenerator[str, None]:
        """流式 JSON 输出，逐 token yield 原始文本片段。

        调用方自行收集并解析 JSON。
        """
        # 确保 system prompt 要求 JSON 输出
        messages = list(messages)
        if messages and messages[0]["role"] == "system":
            if "json" not in messages[0]["content"].lower():
                messages[0]["content"] += "\n\n请严格以 JSON 格式输出，不要包含 markdown 代码块标记。"
        else:
            messages.insert(0, {
                "role": "system",
                "content": "你是一个专业的投后管理分析助手。请严格以 JSON 格式输出，不要包含 markdown 代码块标记。",
            })

        async for token in self.chat_stream(messages, temperature, max_tokens):
            yield token

    async def chat_json(
        self,
        messages: list[dict[str, str]],
        temperature: float = 0.3,
        max_tokens: int = 2000,
    ) -> dict[str, Any]:
        """流式调用但收集为完整 JSON 对象（兼容非流式调用方）。"""
        parts: list[str] = []
        async for token in self.chat_json_stream(messages, temperature, max_tokens):
            parts.append(token)

        text = "".join(parts)

        # 清理可能的 markdown 代码块
        text = text.strip()
        if text.startswith("```"):
            text = text.split("\n", 1)[1] if "\n" in text else text[3:]
        if text.endswith("```"):
            text = text[:-3]
        text = text.strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError as e:
            logger.error("LLM JSON 解析失败: %s, 原始文本: %s", e, text[:500])
            raise RuntimeError(f"LLM 输出 JSON 解析失败: {e}")


llm_client = LLMClient()