Initial commit of LLM speed test app
This commit is contained in:
@@ -0,0 +1,213 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user