# -*- 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)