xiaoxia-saas/packages/shared/ai_client.py

"""豆包大模型 API 客户端（共享层）.

API 和 Worker 两边共用。基于火山引擎方舟平台的 OpenAI 兼容接口。

使用方式：
    from packages.shared.ai_client import get_doubao_client

    client = get_doubao_client()
    if client.is_available:
        result = client.chat_completion(messages=[...])
"""

from __future__ import annotations

import logging
import time
from typing import Any, Optional

import httpx

from packages.shared.config import get_shared_settings

logger = logging.getLogger(__name__)


class DoubaoClient:
    """豆包大模型 API 客户端.

    封装 OpenAI 兼容的 Chat Completion 接口，支持自动重试。
    未配置 API Key 时 is_available 为 False，调用方应降级处理。
    """

    def __init__(self) -> None:
        settings = get_shared_settings()
        self.api_key: str = settings.doubao_api_key
        self.model: str = settings.doubao_model
        self.base_url: str = settings.doubao_base_url.rstrip("/")
        self.timeout: int = settings.doubao_timeout
        self.max_retries: int = settings.doubao_max_retries

    @property
    def is_available(self) -> bool:
        """是否可用（配置了 API Key）."""
        return bool(self.api_key)

    def chat_completion(
        self,
        messages: list[dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 1024,
    ) -> Optional[str]:
        """调用 Chat Completion 接口.

        Args:
            messages: 对话消息列表，[{"role": "user"/"system"/"assistant", "content": "..."}]
            temperature: 采样温度，0-2，默认0.7
            max_tokens: 最大生成token数，默认1024

        Returns:
            模型返回的文本内容，失败返回 None
        """
        if not self.is_available:
            return None

        url = f"{self.base_url}/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }

        last_error: Optional[Exception] = None
        for attempt in range(self.max_retries + 1):
            try:
                response = httpx.post(
                    url,
                    headers=headers,
                    json=payload,
                    timeout=self.timeout,
                )
                response.raise_for_status()
                data = response.json()
                content = data["choices"][0]["message"]["content"]
                return content.strip()
            except Exception as e:
                last_error = e
                if attempt < self.max_retries:
                    wait = 0.5 * (2**attempt)
                    logger.warning(
                        "豆包API调用失败，%.1fs后重试 (第%d/%d次): %s",
                        wait,
                        attempt + 1,
                        self.max_retries + 1,
                        e,
                    )
                    time.sleep(wait)

        logger.error("豆包API调用最终失败: %s", last_error)
        return None


# ── 单例 ─────────────────────────────────────────────────────────────────────


_client: Optional[DoubaoClient] = None


def get_doubao_client() -> DoubaoClient:
    """获取豆包客户端单例."""
    global _client
    if _client is None:
        _client = DoubaoClient()
    return _client