318 lines
12 KiB
Python
318 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""大模型提供商适配器：OpenAI 兼容 / Anthropic / Google Gemini
|
||
|
||
统一接口: call_stream(cfg, prompt, gen, log, should_stop) -> metrics(dict)
|
||
度量指标:
|
||
prompt_tokens 上文/提示词 token 数（来自 provider usage）
|
||
output_tokens 输出 token 数
|
||
cached_tokens 命中缓存的 token 数（provider 返回时才有）
|
||
ttft_ms 首字延迟（time to first token）
|
||
decode_ms 解码阶段耗时（首字 -> 结束）
|
||
total_ms 总耗时（连接开始 -> 结束）
|
||
prefill_speed 预填充速度 = prompt_tokens / ttft（tok/s）
|
||
decode_speed 解码速度 = output_tokens / decode_time（tok/s）
|
||
"""
|
||
import json
|
||
import time
|
||
|
||
import requests
|
||
|
||
import config
|
||
|
||
|
||
class ProviderError(Exception):
|
||
"""API 调用失败"""
|
||
|
||
|
||
class StopRequested(Exception):
|
||
"""用户请求停止"""
|
||
|
||
|
||
DEFAULT_URLS = {
|
||
"openai": "https://api.openai.com/v1",
|
||
"anthropic": "https://api.anthropic.com",
|
||
"google": "https://generativelanguage.googleapis.com",
|
||
}
|
||
|
||
PROVIDER_LABELS = {
|
||
"openai": "OpenAI 兼容",
|
||
"anthropic": "Anthropic",
|
||
"google": "Google Gemini",
|
||
}
|
||
|
||
|
||
def _parse_sse_line(line):
|
||
line = (line or "").strip()
|
||
if not line.startswith("data:"):
|
||
return None
|
||
data = line[5:].strip()
|
||
if not data or data == "[DONE]":
|
||
return None
|
||
try:
|
||
return json.loads(data)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _metrics(start, first_token_at, end, prompt_tokens, output_tokens,
|
||
cached_tokens, output_chars, prompt_chars):
|
||
if first_token_at is None:
|
||
first_token_at = end # 未收到正文但请求完成（如纯 usage 响应）
|
||
ttft_ms = (first_token_at - start) * 1000
|
||
decode_ms = (end - first_token_at) * 1000
|
||
total_ms = (end - start) * 1000
|
||
prefill = (prompt_tokens / (ttft_ms / 1000)) if prompt_tokens and ttft_ms > 0 else None
|
||
decode = (output_tokens / (decode_ms / 1000)) if output_tokens and decode_ms > 0 else None
|
||
return {
|
||
"prompt_tokens": int(prompt_tokens or 0),
|
||
"output_tokens": int(output_tokens or 0),
|
||
"cached_tokens": int(cached_tokens or 0),
|
||
"prompt_chars": int(prompt_chars or 0),
|
||
"output_chars": int(output_chars or 0),
|
||
"ttft_ms": round(ttft_ms, 1),
|
||
"decode_ms": round(decode_ms, 1),
|
||
"total_ms": round(total_ms, 1),
|
||
"prefill_speed": round(prefill, 1) if prefill else None,
|
||
"decode_speed": round(decode, 1) if decode else None,
|
||
}
|
||
|
||
|
||
# ───────────────────────── OpenAI 兼容 ─────────────────────────
|
||
|
||
def stream_openai(cfg, prompt, gen, log, should_stop=None):
|
||
base = (cfg.get("base_url") or DEFAULT_URLS["openai"]).rstrip("/")
|
||
url = base + "/chat/completions"
|
||
headers = {
|
||
"Authorization": "Bearer " + (cfg.get("api_key") or ""),
|
||
"Content-Type": "application/json",
|
||
}
|
||
payload = {
|
||
"model": cfg["model"],
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"stream": True,
|
||
"max_tokens": int(gen.get("max_tokens", 256)),
|
||
"temperature": float(cfg.get("temperature", 0.7)),
|
||
}
|
||
use_usage = True # stream_options include_usage（部分网关不支持则自动去掉重试）
|
||
|
||
def build_payload():
|
||
p = dict(payload)
|
||
if use_usage:
|
||
p["stream_options"] = {"include_usage": True}
|
||
return p
|
||
|
||
start = time.time()
|
||
first_token_at = None
|
||
output_chars = 0
|
||
prompt_tokens = output_tokens = cached_tokens = 0
|
||
event_count = 0
|
||
resp = None
|
||
try:
|
||
while True:
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
resp = requests.post(url, json=build_payload(), headers=headers, stream=True,
|
||
timeout=(config.CONNECT_TIMEOUT, config.STREAM_READ_TIMEOUT))
|
||
if resp.status_code == 200:
|
||
break
|
||
err = resp.text[:400]
|
||
code = resp.status_code
|
||
resp.close()
|
||
resp = None
|
||
if use_usage and _bad_stream_options(err):
|
||
use_usage = False
|
||
if log:
|
||
log("WARN", "提供商不支持 stream_options=include_usage，已去掉参数重试")
|
||
continue
|
||
raise ProviderError("HTTP %s: %s" % (code, err))
|
||
|
||
for obj in _iter_json(resp):
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
event_count += 1
|
||
if obj.get("choices"):
|
||
delta = obj["choices"][0].get("delta") or {}
|
||
# 兼容推理型模型：Qwen3/DeepSeek 思维链在 reasoning_content
|
||
piece = delta.get("content") or delta.get("reasoning_content") or ""
|
||
if piece:
|
||
if first_token_at is None:
|
||
first_token_at = time.time()
|
||
output_chars += len(piece)
|
||
usage = obj.get("usage")
|
||
if usage:
|
||
prompt_tokens = usage.get("prompt_tokens") or 0
|
||
output_tokens = usage.get("completion_tokens") or 0
|
||
details = usage.get("prompt_tokens_details") or {}
|
||
if isinstance(details, dict):
|
||
cached_tokens = details.get("cached_tokens") or 0
|
||
except StopRequested:
|
||
raise
|
||
except Exception as e:
|
||
raise ProviderError("流式请求异常: %s" % e)
|
||
finally:
|
||
if resp is not None:
|
||
resp.close()
|
||
|
||
if event_count == 0:
|
||
raise ProviderError("未收到任何输出内容（HTTP 200 但响应流为空）")
|
||
end = time.time()
|
||
return _metrics(start, first_token_at, end, prompt_tokens, output_tokens,
|
||
cached_tokens, output_chars, len(prompt))
|
||
|
||
|
||
def _bad_stream_options(err: str):
|
||
err = (err or "").lower()
|
||
return ("stream_options" in err or "unknown parameter" in err or "unknown field" in err
|
||
or "additional properties" in err)
|
||
|
||
|
||
def _iter_json(resp):
|
||
"""解析 SSE data: 行，逐个返回 JSON 对象"""
|
||
for raw in resp.iter_lines(decode_unicode=True):
|
||
obj = _parse_sse_line(raw)
|
||
if obj is not None:
|
||
yield obj
|
||
|
||
|
||
# ───────────────────────── Anthropic ─────────────────────────
|
||
|
||
def stream_anthropic(cfg, prompt, gen, log, should_stop=None):
|
||
base = (cfg.get("base_url") or DEFAULT_URLS["anthropic"]).rstrip("/")
|
||
url = base + "/v1/messages"
|
||
headers = {
|
||
"x-api-key": cfg.get("api_key") or "",
|
||
"anthropic-version": "2023-06-01",
|
||
"Content-Type": "application/json",
|
||
}
|
||
payload = {
|
||
"model": cfg["model"],
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"max_tokens": int(gen.get("max_tokens", 256)),
|
||
"temperature": float(cfg.get("temperature", 0.7)),
|
||
"stream": True,
|
||
}
|
||
start = time.time()
|
||
first_token_at = None
|
||
output_chars = 0
|
||
prompt_tokens = output_tokens = 0
|
||
event_count = 0
|
||
resp = None
|
||
try:
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
resp = requests.post(url, json=payload, headers=headers, stream=True,
|
||
timeout=(config.CONNECT_TIMEOUT, config.STREAM_READ_TIMEOUT))
|
||
if resp.status_code != 200:
|
||
err = resp.text[:400]
|
||
resp.close()
|
||
resp = None
|
||
raise ProviderError("HTTP %s: %s" % (resp.status_code, err))
|
||
for obj in _iter_json(resp):
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
event_count += 1
|
||
etype = obj.get("type")
|
||
if etype == "message_start":
|
||
usage = (obj.get("message") or {}).get("usage") or {}
|
||
prompt_tokens = usage.get("input_tokens") or 0
|
||
elif etype == "content_block_delta":
|
||
delta = obj.get("delta") or {}
|
||
# 兼容 extended thinking：thinking 文本也算输出
|
||
text = delta.get("text") or delta.get("thinking") or ""
|
||
if text:
|
||
if first_token_at is None:
|
||
first_token_at = time.time()
|
||
output_chars += len(text)
|
||
elif etype == "message_delta":
|
||
usage = obj.get("usage") or {}
|
||
output_tokens = usage.get("output_tokens") or output_tokens
|
||
except StopRequested:
|
||
raise
|
||
except Exception as e:
|
||
raise ProviderError("流式请求异常: %s" % e)
|
||
finally:
|
||
if resp is not None:
|
||
resp.close()
|
||
|
||
if event_count == 0:
|
||
raise ProviderError("未收到任何输出内容（HTTP 200 但响应流为空）")
|
||
end = time.time()
|
||
return _metrics(start, first_token_at, end, prompt_tokens, output_tokens,
|
||
0, output_chars, len(prompt))
|
||
|
||
|
||
# ───────────────────────── Google Gemini ─────────────────────────
|
||
|
||
def stream_google(cfg, prompt, gen, log, should_stop=None):
|
||
base = (cfg.get("base_url") or DEFAULT_URLS["google"]).rstrip("/")
|
||
model = cfg["model"]
|
||
url = "%s/v1beta/models/%s:streamGenerateContent" % (base, model)
|
||
params = {"alt": "sse", "key": cfg.get("api_key") or ""}
|
||
headers = {"Content-Type": "application/json"}
|
||
payload = {
|
||
"contents": [{"parts": [{"text": prompt}]}],
|
||
"generationConfig": {
|
||
"temperature": float(cfg.get("temperature", 0.7)),
|
||
"maxOutputTokens": int(gen.get("max_tokens", 256)),
|
||
"candidateCount": 1,
|
||
},
|
||
}
|
||
start = time.time()
|
||
first_token_at = None
|
||
output_chars = 0
|
||
prompt_tokens = output_tokens = cached_tokens = 0
|
||
event_count = 0
|
||
resp = None
|
||
try:
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
resp = requests.post(url, params=params, json=payload, headers=headers, stream=True,
|
||
timeout=(config.CONNECT_TIMEOUT, config.STREAM_READ_TIMEOUT))
|
||
if resp.status_code != 200:
|
||
err = resp.text[:400]
|
||
resp.close()
|
||
resp = None
|
||
raise ProviderError("HTTP %s: %s" % (resp.status_code, err))
|
||
for obj in _iter_json(resp):
|
||
if should_stop and should_stop():
|
||
raise StopRequested()
|
||
event_count += 1
|
||
cands = obj.get("candidates") or []
|
||
if cands:
|
||
parts = (cands[0].get("content") or {}).get("parts") or []
|
||
for part in parts:
|
||
# 兼容 thinking 模型：thought 文本也算输出
|
||
text = part.get("text") or part.get("thought") or ""
|
||
if text:
|
||
if first_token_at is None:
|
||
first_token_at = time.time()
|
||
output_chars += len(text)
|
||
um = obj.get("usageMetadata") or {}
|
||
if um:
|
||
prompt_tokens = um.get("promptTokenCount") or 0
|
||
output_tokens = um.get("candidatesTokenCount") or 0
|
||
cached_tokens = um.get("cachedContentTokenCount") or 0
|
||
except StopRequested:
|
||
raise
|
||
except Exception as e:
|
||
raise ProviderError("流式请求异常: %s" % e)
|
||
finally:
|
||
if resp is not None:
|
||
resp.close()
|
||
|
||
if event_count == 0:
|
||
raise ProviderError("未收到任何输出内容（HTTP 200 但响应流为空）")
|
||
end = time.time()
|
||
return _metrics(start, first_token_at, end, prompt_tokens, output_tokens,
|
||
cached_tokens, output_chars, len(prompt))
|
||
|
||
|
||
# ───────────────────────── 统一入口 ─────────────────────────
|
||
|
||
def call_stream(cfg, prompt, gen, log=None, should_stop=None):
|
||
provider = cfg.get("provider", "openai")
|
||
fn = {"openai": stream_openai, "anthropic": stream_anthropic, "google": stream_google}.get(provider)
|
||
if fn is None:
|
||
raise ProviderError("不支持的提供商类型: %s" % provider)
|
||
return fn(cfg, prompt, gen, log, should_stop)
|