Spaces:
Sleeping
Sleeping
File size: 2,323 Bytes
6167a25 302aeaa 6167a25 b6d0232 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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}" |