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}"