214 lines
8.0 KiB
Python
214 lines
8.0 KiB
Python
"""OpenAI 兼容 API 客户端
|
||
|
||
对接任意 OpenAI /v1 兼容服务（OpenAI、Together、硅基流动、llama.cpp server 等），
|
||
与 OllamaClient 保持相同的 generate/chat 接口，测试用例无需区分后端。
|
||
|
||
指标来源：
|
||
- TTFT: 流式请求中首个 content 片段到达时间（真实测量）
|
||
- Prefill: 优先取服务端 timings.prompt_ms（llama.cpp server 暴露），否则回退 TTFT
|
||
- Decode Speed: 优先取服务端 timings.predicted_per_second，否则按 completion_tokens/e2e 估算
|
||
- E2E: 客户端请求总耗时
|
||
"""
|
||
import json
|
||
import logging
|
||
import time
|
||
|
||
import requests
|
||
|
||
from . import config
|
||
from .client_base import InferenceResult, ClientError
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class OpenAIError(ClientError):
|
||
"""OpenAI 兼容 API 调用错误"""
|
||
|
||
|
||
class OpenAIClient:
|
||
"""封装对 OpenAI 兼容 /v1 服务的调用"""
|
||
|
||
def __init__(self, base_url: str = None, model: str = None, api_key: str = None,
|
||
timeout: int = 300):
|
||
self.base_url = (base_url or config.OPENAI_BASE_URL).rstrip("/")
|
||
self.model = model or config.OPENAI_DEFAULT_MODEL
|
||
self.api_key = api_key if api_key not in (None, "") else config.OPENAI_API_KEY
|
||
self.timeout = timeout
|
||
self._session = requests.Session()
|
||
|
||
# ---------- 请求头 ----------
|
||
|
||
@property
|
||
def _headers(self) -> dict:
|
||
headers = {"Content-Type": "application/json"}
|
||
if self.api_key:
|
||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||
return headers
|
||
|
||
# ---------- 服务探测 ----------
|
||
|
||
def ping(self) -> bool:
|
||
"""探测服务是否可用"""
|
||
try:
|
||
return self._session.get(
|
||
f"{self.base_url}/models", headers=self._headers, timeout=5
|
||
).ok
|
||
except Exception:
|
||
return False
|
||
|
||
def list_models(self) -> list:
|
||
"""列出可用模型"""
|
||
try:
|
||
resp = self._session.get(
|
||
f"{self.base_url}/models", headers=self._headers, timeout=10
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json().get("data", [])
|
||
models = []
|
||
for m in data:
|
||
mid = m.get("id") or m.get("name")
|
||
if mid:
|
||
models.append(mid)
|
||
return models
|
||
except requests.exceptions.ConnectionError as e:
|
||
raise OpenAIError(
|
||
f"无法连接服务（{self.base_url}），请检查地址与网络"
|
||
) from e
|
||
except requests.exceptions.HTTPError as e:
|
||
if resp.status_code in (401, 403):
|
||
raise OpenAIError("API Key 无效或被拒绝（401/403）") from e
|
||
raise OpenAIError(f"获取模型列表失败: {resp.status_code} {resp.text[:200]}") from e
|
||
except Exception as e:
|
||
raise OpenAIError(f"获取模型列表失败: {e}") from e
|
||
|
||
# ---------- 文本生成 ----------
|
||
|
||
def generate(self, prompt: str, system: str = None, options: dict = None) -> InferenceResult:
|
||
"""单轮文本生成（流式），等价于单条 user 消息的 chat"""
|
||
messages = []
|
||
if system:
|
||
messages.append({"role": "system", "content": system})
|
||
messages.append({"role": "user", "content": prompt})
|
||
return self.chat(messages, options)
|
||
|
||
# ---------- 多轮对话 ----------
|
||
|
||
def chat(self, messages: list, options: dict = None) -> InferenceResult:
|
||
"""多轮对话，messages 为 [{"role": ..., "content": ...}]"""
|
||
options = options or {}
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"stream": True,
|
||
"temperature": options.get("temperature", 0.3),
|
||
"max_tokens": options.get("num_predict", 256),
|
||
}
|
||
|
||
start = time.perf_counter()
|
||
ttft_ms = 0.0
|
||
chunks = []
|
||
final = {}
|
||
timings = {}
|
||
|
||
try:
|
||
resp = self._session.post(
|
||
f"{self.base_url}/chat/completions",
|
||
headers=self._headers,
|
||
json=payload,
|
||
stream=True,
|
||
timeout=self.timeout,
|
||
)
|
||
except requests.exceptions.ConnectionError as e:
|
||
raise OpenAIError(
|
||
f"无法连接服务（{self.base_url}），请检查地址与网络"
|
||
) from e
|
||
except requests.exceptions.Timeout as e:
|
||
raise OpenAIError(f"请求超时（{self.timeout}s）: {e}") from e
|
||
|
||
if resp.status_code != 200:
|
||
err_text = resp.text[:300]
|
||
raise OpenAIError(f"API 错误 {resp.status_code}: {err_text}")
|
||
|
||
# 流式解析 SSE
|
||
try:
|
||
for line in resp.iter_lines(decode_unicode=True):
|
||
if not line:
|
||
continue
|
||
line = line.strip()
|
||
if line.startswith("data:"):
|
||
line = line[5:].strip()
|
||
if line == "[DONE]":
|
||
break
|
||
if not line:
|
||
continue
|
||
try:
|
||
data = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
choices = data.get("choices") or []
|
||
if not choices:
|
||
continue
|
||
delta = choices[0].get("delta") or {}
|
||
# 兼容推理模型：token 可能放在 reasoning_content（思考过程）或 content
|
||
piece = delta.get("content") or delta.get("reasoning_content") or ""
|
||
|
||
if piece:
|
||
# 首个内容片段到达时间 = TTFT
|
||
if ttft_ms == 0.0:
|
||
ttft_ms = (time.perf_counter() - start) * 1000
|
||
chunks.append(piece)
|
||
|
||
# 服务端最终块可能携带 usage / timings
|
||
if choices[0].get("finish_reason"):
|
||
final = data
|
||
if data.get("usage"):
|
||
final = data
|
||
except (requests.exceptions.ConnectionError, requests.exceptions.ChunkedEncodingError) as e:
|
||
raise OpenAIError(f"流式读取中断: {e}") from e
|
||
|
||
e2e_ms = (time.perf_counter() - start) * 1000
|
||
result = InferenceResult()
|
||
result.response = "".join(chunks)
|
||
result.e2e_ms = e2e_ms
|
||
result.ttft_ms = ttft_ms
|
||
|
||
# usage: token 计数（部分服务端流式响应不含 usage，回退用 timings）
|
||
usage = final.get("usage", {}) or {}
|
||
prompt_tok = int(usage.get("prompt_tokens", 0) or 0)
|
||
completion_tok = int(usage.get("completion_tokens", 0) or 0)
|
||
|
||
# timings: llama.cpp / llama-server 暴露的服务端计时（单位毫秒）
|
||
timings = final.get("timings", {}) or {}
|
||
prompt_n = int(timings.get("prompt_n", 0) or 0)
|
||
predicted_n = int(timings.get("predicted_n", 0) or 0)
|
||
prompt_ms = float(timings.get("prompt_ms", 0) or 0)
|
||
predicted_ms = float(timings.get("predicted_ms", 0) or 0)
|
||
predicted_pps = float(timings.get("predicted_per_second", 0) or 0)
|
||
|
||
result.prompt_tokens = prompt_tok or prompt_n
|
||
result.completion_tokens = completion_tok or predicted_n
|
||
result.total_tokens = result.prompt_tokens + result.completion_tokens
|
||
|
||
result.prefill_ms = prompt_ms
|
||
|
||
# 解码速度：优先服务端计时
|
||
if predicted_pps > 0:
|
||
result.decode_speed_tok_s = predicted_pps
|
||
elif result.completion_tokens > 0 and result.e2e_ms > 0:
|
||
result.decode_speed_tok_s = result.completion_tokens / (result.e2e_ms / 1000)
|
||
|
||
# TTFT 回退：流式被缓冲或未测到时，用服务端 prompt_ms
|
||
if result.ttft_ms <= 0 or result.ttft_ms > e2e_ms * 0.95:
|
||
if prompt_ms > 0:
|
||
result.ttft_ms = prompt_ms
|
||
else:
|
||
result.ttft_ms = e2e_ms if e2e_ms > 0 else 0.0
|
||
|
||
logger.debug(
|
||
"推理完成(openai): e2e=%.1fms ttft=%.1fms decode=%.1f tok/s tokens=%d/%d",
|
||
result.e2e_ms, result.ttft_ms, result.decode_speed_tok_s,
|
||
result.prompt_tokens, result.completion_tokens,
|
||
)
|
||
return result
|