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