app/llm_client.py

from __future__ import annotations

import json
import itertools
import threading
from typing import Any

from openai import OpenAI
from openai import AuthenticationError


def mask_key(key: str) -> str:
    if not key:
        return ""
    if len(key) <= 10:
        return "*" * len(key)
    return f"{key[:6]}...{key[-4:]}"


class ApiKeyPool:
    def __init__(self, api_keys: list[str], base_url: str):
        if not api_keys:
            raise ValueError("ApiKeyPool 初始化失败：api_keys 为空。")
        self._base_url = base_url
        self._lock = threading.Lock()
        self._cycle = itertools.cycle(api_keys)

    def get_client(self) -> OpenAI:
        with self._lock:
            key = next(self._cycle)
        return OpenAI(api_key=key, base_url=self._base_url)


def build_client(api_key: str, base_url: str) -> OpenAI:
    return OpenAI(api_key=api_key, base_url=base_url)


def safe_json_loads(text: str) -> dict[str, Any]:
    text = (text or "").strip()

    if text.startswith("```"):
        lines = text.splitlines()
        if len(lines) >= 3:
            text = "\n".join(lines[1:-1]).strip()

    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        text = text[start:end + 1]

    return json.loads(text)


def chat_completion_json(
    client: OpenAI,
    model: str,
    messages: list[dict[str, Any]],
    temperature: float = 0.1,
    timeout: int = 300,
) -> str:
    resp = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        timeout=timeout,
    )
    return resp.choices[0].message.content or ""


def validate_api_key(api_key: str, base_url: str, model: str) -> tuple[bool, str]:
    try:
        client = build_client(api_key=api_key, base_url=base_url)
        _ = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "你是一个AI助手"},
                {"role": "user", "content": "Reply only with OK."},
            ],
            temperature=0,
            timeout=60,
        )
        return True, "OK"
    except AuthenticationError as e:
        return False, f"AuthenticationError: {e}"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"