# -*- coding: utf-8 -*- """ 统一模型网关（V3.5）：多供应商 OpenAI 兼容协议调用 + 流式 + 精细计量 + 计价 ===================================================================== 核心变化： 1. 所有模型输出按 token 流式接收（SSE），不再一次性等完整响应； 2. 超时语义改为「单 token 返回超时 / 首字延迟超时 / 整体兜底」， 三个值均可在 设置 页面配置（settings 表，动态生效）： - token_timeout 相邻两个 token 数据块的最大间隔（默认 60s） - first_token_timeout 请求发出后首块数据的最长等待（默认 120s） - request_timeout 整体兜底上限（默认 600s） 3. 大模型接口库（llm_endpoints）：base_url / api_key / 模型列表 / 逐模型定价， 计费方式支持「按 token 数」与「按调用次数」两种，Worker 直接选用接口库； 4. 精细计量：每次调用记录 prompt / completion / 缓存命中 cached / 调用次数 / 首字延迟 / 总耗时，全部入库（cost_records / agent_steps / chat_messages）。 """ import json import time import requests import config import db class LLMError(Exception): """模型调用异常。partial_text：超时中断前已收到的部分输出。""" def __init__(self, msg, partial_text=''): super().__init__(msg) self.partial_text = partial_text or '' def get_provider_cfg(provider): cfg = config.PROVIDERS.get(provider) if not cfg: raise LLMError(f'未知供应商: {provider}') return cfg # --------------------------------------------------------------------------- # 大模型接口库 # --------------------------------------------------------------------------- def get_endpoint(endpoint_id): if not endpoint_id: return None try: return db.q('SELECT * FROM llm_endpoints WHERE id=? AND status="enabled"', (int(endpoint_id),), one=True) except Exception: return None def _endpoint_models(ep): try: return json.loads(ep.get('models') or '[]') or [] except Exception: return [] def _endpoint_pricing_map(ep): try: return json.loads(ep.get('pricing') or '{}') or {} except Exception: return {} # --------------------------------------------------------------------------- # 模型能力标签（V3.5.2）：chat/thinking/vision/audio_in/audio_out/image_gen/video_gen/embedding/rerank # --------------------------------------------------------------------------- CAP_LABELS = { 'chat': '💬 对话', 'thinking': '🧠 思考', 'vision': '👁️ 视觉输入', 'audio_in': '🎤 语音输入', 'audio_out': '🔊 语音输出', 'image_gen': '🎨 图片生成', 'video_gen': '🎬 视频生成', 'embedding': '🔢 Embedding', 'rerank': '🔀 Rerank', } ALL_CAPS = list(CAP_LABELS.keys()) def endpoint_model_caps(endpoint, model): """某接口下某模型的 能力标签 列表（缺省按模型名启发式推断）""" if endpoint: try: caps_map = json.loads(endpoint.get('capabilities') or '{}') or {} except Exception: caps_map = {} caps = list(caps_map.get(model) or []) if caps: return caps # 兜底：模型名启发式 return guess_model_caps(model) return guess_model_caps(model) def guess_model_caps(model): """按模型名启发式推断能力（未显式配置时）""" m = (model or '').lower() caps = {'chat'} if any(k in m for k in ('think', 'reason', 'r1', 'deepseek-reasoner', 'o1', 'o3', 'glm-4.6')): caps.add('thinking') if any(k in m for k in ('vision', 'vl', 'omni', 'qwen2.5-vl', 'glm-4v', 'llava', 'internvl', 'mini-omni')): caps.add('vision') if any(k in m for k in ('audio', 'voice', 'tts', 'speech', 'asr', 'whisper', 'transcri', 's2t', 'funasr', 'sensevoice')): caps.add('audio_in') caps.add('audio_out') if any(k in m for k in ('image', 'dall', 'stable', 'flux', 'sd3', 'sora', 'wan', 'cogview', 'draw')): caps.add('image_gen') if any(k in m for k in ('video', 'sora', 'wan', 'kling', 'runway')): caps.add('video_gen') if any(k in m for k in ('embedding', 'text-embedding', 'bge')): caps.add('embedding') if any(k in m for k in ('rerank', 'bge-rerank')): caps.add('rerank') return sorted(caps) def model_has_cap(endpoint, model, cap): return cap in endpoint_model_caps(endpoint, model) def worker_llm_cfg(worker): """解析 Worker 的大模型接口配置（V3.5）： 优先取绑定的接口库 endpoint_id（统一鉴权/计价），worker 自带 base_url/api_key 可覆盖。 返回 dict(provider, model, base_url, api_key, endpoint, in_price, out_price, price_per_call, billing, pricing_map)""" ep = get_endpoint(worker.get('endpoint_id')) if worker else None if ep: models = _endpoint_models(ep) return { 'provider': ep.get('provider') or 'custom', 'model': worker.get('model') or (models[0] if models else ''), 'base_url': worker.get('base_url') or ep.get('base_url') or '', 'api_key': worker.get('api_key') or ep.get('api_key') or '', 'endpoint': ep, 'in_price': float(ep.get('input_price') or 0), 'out_price': float(ep.get('output_price') or 0), 'price_per_call': float(ep.get('price_per_call') or 0), 'billing': ep.get('billing') or 'token', 'pricing_map': _endpoint_pricing_map(ep), } return { 'provider': (worker or {}).get('provider', ''), 'model': (worker or {}).get('model', ''), 'base_url': (worker or {}).get('base_url') or '', 'api_key': (worker or {}).get('api_key') or '', 'endpoint': None, 'in_price': None, 'out_price': None, 'price_per_call': None, 'billing': 'token', 'pricing_map': {}, } # --------------------------------------------------------------------------- # 计价：优先接口库（逐模型定价 / 按调用次数），否则 config.MODEL_PRICING # --------------------------------------------------------------------------- def model_price(model): p = config.MODEL_PRICING.get(model, config.DEFAULT_PRICE) return p['input'], p['output'] def calc_cost(model, prompt_tokens, completion_tokens): pin, pout = model_price(model) return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6) def calc_cost_ex(model, prompt_tokens, completion_tokens, calls=1, cached_tokens=0, cfg=None): """按 Worker/接口库配置计价。cfg = worker_llm_cfg() 结果。 - 按调用次数计费：price_per_call 单位为 元/千次，cost = price_per_call/1000 * calls - 按 token 计费：逐模型 {input, input_cache, output}，缓存命中 token 按 input_cache 计价""" if cfg and cfg.get('endpoint'): ep = cfg['endpoint'] if ep.get('billing') == 'call': return round(float(ep.get('price_per_call') or 0) / 1000.0 * max(1, calls), 6) p = (cfg.get('pricing_map') or {}).get(model) if p: pin = float(p.get('input') or 0) pcache = float(p.get('input_cache') if p.get('input_cache') is not None else p.get('cache') or 0) pout = float(p.get('output') or 0) if pin or pout: cached = max(0, min(cached_tokens or 0, prompt_tokens)) return round((prompt_tokens - cached) / 1e6 * pin + cached / 1e6 * pcache + completion_tokens / 1e6 * pout, 6) pin, pout = cfg.get('in_price') or 0, cfg.get('out_price') or 0 if pin or pout: return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6) pin, pout = model_price(model) return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6) def worker_unit_price(worker): """自动路由用：估算 Worker 单次调用成本（按 token 计费 = 输入价+0.5*输出价；按次计费 = 单价/千次）""" cfg = worker_llm_cfg(worker) if cfg.get('endpoint') and cfg.get('billing') == 'call': return float(cfg.get('price_per_call') or 0) / 1000.0 p = (cfg.get('pricing_map') or {}).get(cfg['model']) if p: return float(p.get('input') or 0) + float(p.get('output') or 0) * 0.5 if cfg.get('in_price') is not None: return float(cfg['in_price']) + float(cfg['out_price']) * 0.5 pin, pout = model_price(cfg['model']) return pin + pout * 0.5 # --------------------------------------------------------------------------- # Token 估算（流式响应未带 usage 时的兜底） # --------------------------------------------------------------------------- def _estimate_prompt_tokens(messages): n = 0 for m in messages or []: c = m.get('content') if isinstance(m, dict) else '' if isinstance(c, str): n += max(1, int(len(c) * 0.6)) elif isinstance(c, list): for part in c: if not isinstance(part, dict): continue t = part.get('text') or '' n += max(1, int(len(t) * 0.6)) if part.get('image_url') or part.get('image'): n += 1000 # 图片按约 1000 token 估算 return max(1, n) def _estimate_completion_tokens(text): return max(1, int(len(text or '') * 0.6)) # --------------------------------------------------------------------------- # SSE 流式核心（单次尝试，无重试；超时语义见模块说明） # --------------------------------------------------------------------------- def _sock_of(resp): """尽力获取底层 socket 以调整读超时（兼容不同 requests/urllib3 版本）""" try: raw = resp.raw fp = getattr(raw, '_fp', None) or getattr(raw, 'fp', None) fpp = getattr(fp, 'fp', None) for obj in (fpp, fp): if obj is None: continue s = getattr(obj, 'raw', None) or getattr(obj, '_sock', None) if s is not None: return s except Exception: pass return None def _chat_stream_raw(provider, model, messages, temperature=0.7, max_tokens=None, base_url=None, api_key=None, token_timeout=60, first_token_timeout=120, extra_payload=None): """发起一次流式请求，产出事件： yield ('delta', piece) | ('usage', usage_dict) | ('finish', finish_reason) 结束前若收到 usage 则正常给出；异常抛 LLMError（含 partial_text）。 超时：首块数据等待 first_token_timeout；相邻数据块间隔 token_timeout；整体 request_timeout 兜底。""" cfg = get_provider_cfg(provider) url = (base_url or cfg['base_url']).rstrip('/') + '/chat/completions' key = api_key or cfg['api_key'] if not key: raise LLMError(f'供应商 {provider} 未配置 API Key') headers = { 'Authorization': f'Bearer {key}', 'Content-Type': 'application/json', 'Accept': 'text/event-stream', } payload = { 'model': model, 'messages': messages, 'temperature': temperature, 'stream': True, 'stream_options': {'include_usage': True}, } if max_tokens: payload['max_tokens'] = max_tokens if extra_payload: payload.update(extra_payload) t0 = time.time() first_token_at = None resp = None try: resp = requests.post(url, json=payload, headers=headers, stream=True, timeout=(min(30, first_token_timeout), first_token_timeout)) if resp.status_code != 200: body = resp.text[:300] resp.close() if resp.status_code == 429: raise LLMError(f'模型限流(429): {body}') if resp.status_code >= 500: raise LLMError(f'服务端错误({resp.status_code}): {body}') raise LLMError(f'调用失败({resp.status_code}): {body}') # 首块之后，读超时降为「单 token 返回超时」（首次 read 保持 first_token_timeout） sock = _sock_of(resp) first_line_seen = False for raw_line in resp.iter_lines(decode_unicode=True): line = (raw_line or '').strip() if not first_line_seen: first_line_seen = True if sock is not None: try: sock.settimeout(token_timeout) except Exception: pass if not line or not line.startswith('data:'): continue data = line[5:].strip() if data == '[DONE]': break try: evt = json.loads(data) except Exception: continue if evt.get('usage'): yield ('usage', evt['usage']) continue choices = evt.get('choices') or [] if not choices: continue ch = choices[0] delta = ch.get('delta') or {} piece = delta.get('content') or '' if piece: if first_token_at is None: first_token_at = time.time() yield ('delta', piece) else: rp = delta.get('reasoning_content') or '' if rp: if first_token_at is None: first_token_at = time.time() yield ('reasoning', rp) if ch.get('finish_reason'): yield ('finish', ch.get('finish_reason')) return except requests.exceptions.ReadTimeout: elapsed = time.time() - t0 if first_token_at is None: raise LLMError(f'首字延迟超时（>{first_token_timeout}s 无输出）') raise LLMError(f'Token 返回超时（>{token_timeout}s 无新数据，已输出 {elapsed:.0f}s）') except requests.exceptions.Timeout: raise LLMError(f'请求超时（{(time.time() - t0):.0f}s）') except requests.exceptions.ConnectionError as e: raise LLMError(f'连接失败: {e}') finally: if resp is not None: try: resp.close() except Exception: pass def _usage_fields(usage): """从 usage 中解析精细计量字段""" usage = usage or {} pt = int(usage.get('prompt_tokens') or 0) ct = int(usage.get('completion_tokens') or 0) det = usage.get('prompt_tokens_details') or {} cached = int(det.get('cached_tokens') or 0) if not cached: cached = int(usage.get('prompt_cache_hit_tokens') or 0) return pt, ct, cached # --------------------------------------------------------------------------- # 流式调用（带重试，聚合结果） # --------------------------------------------------------------------------- def chat_stream(provider, model, messages, temperature=0.7, max_tokens=None, base_url=None, api_key=None, timeout=None, retries=None, token_timeout=None, first_token_timeout=None, on_chunk=None, extra_payload=None): """流式接收全部输出，返回聚合结果 dict： {text, model, prompt_tokens, completion_tokens, total_tokens, cached_tokens, cost, calls, first_token_ms, elapsed_ms, usage_estimated, finish_reason} - 超时按「单 token 返回 / 首字延迟」语义（设置中可配） - 已有部分输出时不再重试（避免重复内容），否则按 retries 重试 - on_chunk(delta) 逐块回调（用于对话流式转发）""" if token_timeout is None or first_token_timeout is None: tk, fk, rk = db.get_llm_timeouts() token_timeout = token_timeout or tk first_token_timeout = first_token_timeout or fk request_timeout = timeout or max(first_token_timeout + 5, 60) retries = config.MAX_RETRY if retries is None else retries last_err = None for attempt in range(retries + 1): parts, usage = [], None thinking = [] finish_reason = None first_token_at = None t0 = time.time() try: for evt, val in _chat_stream_raw( provider, model, messages, temperature=temperature, max_tokens=max_tokens, base_url=base_url, api_key=api_key, token_timeout=token_timeout, first_token_timeout=first_token_timeout, extra_payload=extra_payload): if evt == 'delta': if first_token_at is None: first_token_at = time.time() parts.append(val) if on_chunk: try: on_chunk(val) except Exception: pass elif evt == 'reasoning': if first_token_at is None: first_token_at = time.time() thinking.append(val) elif evt == 'usage': usage = val elif evt == 'finish': finish_reason = val text = ''.join(parts) if not text and not usage: last_err = LLMError('模型返回空内容，重试中…') continue pt, ct, cached = _usage_fields(usage) usage_estimated = usage is None if usage is None: pt, ct = _estimate_prompt_tokens(messages), _estimate_completion_tokens(text) elapsed_ms = int((time.time() - t0) * 1000) first_ms = int((first_token_at - t0) * 1000) if first_token_at else None return { 'text': text, 'thinking': ''.join(thinking), 'model': model, 'prompt_tokens': pt, 'completion_tokens': ct, 'total_tokens': pt + ct, 'cached_tokens': cached, 'cost': calc_cost(model, pt, ct), 'calls': 1, 'first_token_ms': first_ms, 'elapsed_ms': elapsed_ms, 'usage_estimated': usage_estimated, 'finish_reason': finish_reason, } except LLMError as e: last_err = e if e.partial_text: raise continue raise last_err or LLMError('未知错误') # --------------------------------------------------------------------------- # 兼容接口（非 Worker 场景） # --------------------------------------------------------------------------- def chat(provider, model, messages, temperature=0.7, max_tokens=None, base_url=None, api_key=None, timeout=None, retries=None, token_timeout=None, first_token_timeout=None, on_chunk=None, cfg=None): """兼容旧接口：流式接收全部输出后返回聚合结果。 cfg = worker_llm_cfg() 结果时按接口库计价/鉴权。""" if cfg: provider = cfg['provider'] base_url = cfg['base_url'] or None api_key = cfg['api_key'] or None model = cfg['model'] r = chat_stream(provider, model, messages, temperature=temperature, max_tokens=max_tokens, base_url=base_url, api_key=api_key, timeout=timeout, retries=retries, token_timeout=token_timeout, first_token_timeout=first_token_timeout, on_chunk=on_chunk) if cfg: r['cost'] = calc_cost_ex(model, r['prompt_tokens'], r['completion_tokens'], 1, cfg) return r def chat_worker(worker, messages, temperature=None, max_tokens=None, on_chunk=None): """按 Worker 配置（含接口库）调用模型，返回聚合结果（含接口库计价与精细计量）""" cfg = worker_llm_cfg(worker) if not cfg['model']: raise LLMError(f'Worker「{worker.get("name", "")}」未配置模型') r = chat_stream(cfg['provider'], cfg['model'], messages, temperature=temperature if temperature is not None else worker.get('temperature', 0.7), max_tokens=max_tokens or worker.get('max_tokens') or 2000, base_url=cfg['base_url'] or None, api_key=cfg['api_key'] or None, on_chunk=on_chunk) r['cost'] = calc_cost_ex(cfg['model'], r['prompt_tokens'], r['completion_tokens'], 1, r.get('cached_tokens', 0), cfg) r['worker_id'] = worker['id'] r['provider'] = cfg['provider'] r['model'] = cfg['model'] return r def stream_chat(cfg, messages, temperature=0.7, max_tokens=None, extra_payload=None): """对话专用的 实时流式生成器（V3.5.2）——真正边收边吐，思考与回答分开。 cfg: {'provider','model','base_url','api_key','endpoint'}（worker 或 接口 目标解析结果） yield ('reasoning', piece) | ('delta', piece) | ('done', result) | ('error', msg) result 含 thinking/text/usage/cost(含缓存命中价与按次计费)/first_token_ms/elapsed_ms。""" provider = cfg['provider'] model = cfg['model'] if not model: yield ('error', '未选择模型') return tk, fk, rk = db.get_llm_timeouts() parts, thinking, usage = [], [], None finish_reason = None first_token_at = None t0 = time.time() tried = False while True: try: for evt, val in _chat_stream_raw( provider, model, messages, temperature=temperature, max_tokens=max_tokens, base_url=cfg.get('base_url') or None, api_key=cfg.get('api_key') or None, token_timeout=tk, first_token_timeout=fk, extra_payload=extra_payload): if evt == 'reasoning': if first_token_at is None: first_token_at = time.time() thinking.append(val) yield ('reasoning', val) elif evt == 'delta': if first_token_at is None: first_token_at = time.time() parts.append(val) yield ('delta', val) elif evt == 'usage': usage = val elif evt == 'finish': finish_reason = val break except LLMError as e: if parts or tried: yield ('error', str(e)) return tried = True # 无任何输出时允许重试一次 continue text = ''.join(parts) pt, ct, cached = _usage_fields(usage) usage_estimated = usage is None if usage is None: pt, ct = _estimate_prompt_tokens(messages), _estimate_completion_tokens(text) elapsed_ms = int((time.time() - t0) * 1000) first_ms = int((first_token_at - t0) * 1000) if first_token_at else None cost = calc_cost_ex(model, pt, ct, 1, cached, cfg) if cfg.get('endpoint') else calc_cost(model, pt, ct) yield ('done', { 'text': text, 'thinking': ''.join(thinking), 'model': model, 'prompt_tokens': pt, 'completion_tokens': ct, 'total_tokens': pt + ct, 'cached_tokens': cached, 'cost': cost, 'calls': 1, 'first_token_ms': first_ms, 'elapsed_ms': elapsed_ms, 'usage_estimated': usage_estimated, 'finish_reason': finish_reason, }) def chat_vision(provider, model, text, image_url=None, image_path=None, temperature=0.4, max_tokens=2000, base_url=None, api_key=None, retries=3): """多模态视觉调用：文本 + 图片（URL 或本地路径/base64）。 返回与 chat() 相同结构。 容错：聚合 API 偶发路由到纯文本后端（不认识 image_url），自动重试。""" import base64 as _b64 import time as _time content = [{'type': 'text', 'text': text}] img_url = image_url if image_path: with open(image_path, 'rb') as f: raw = f.read() mime = 'image/png' if image_path.lower().endswith(('.jpg', '.jpeg')): mime = 'image/jpeg' elif image_path.lower().endswith('.gif'): mime = 'image/gif' elif image_path.lower().endswith('.webp'): mime = 'image/webp' img_url = f'data:{mime};base64,{_b64.b64encode(raw).decode()}' if img_url: content.append({'type': 'image_url', 'image_url': {'url': img_url}}) messages = [{'role': 'user', 'content': content}] last_err = None for attempt in range(max(1, retries)): try: return chat(provider, model, messages, temperature=temperature, max_tokens=max_tokens, base_url=base_url, api_key=api_key) except LLMError as e: last_err = e msg = str(e) if any(k in msg for k in ('image_url', 'InvalidParameter', 'invalid_parameter', 'does not appear to be valid', 'image')): _time.sleep(2 * (attempt + 1)) continue raise raise last_err or LLMError('视觉调用失败') def test_connection(provider, model, base_url=None, api_key=None, endpoint=None): """连通性测试：发一条最小请求（流式），返回延迟/回复/成本""" t0 = time.time() cfg = None if endpoint: cfg = {'provider': endpoint.get('provider') or 'custom', 'model': model, 'base_url': base_url or endpoint.get('base_url') or '', 'api_key': api_key or endpoint.get('api_key') or '', 'endpoint': endpoint} r = chat(provider, model, [{'role': 'user', 'content': '请回复"OK"两个字'}], temperature=0, max_tokens=16, base_url=base_url, api_key=api_key, timeout=30, retries=0, token_timeout=30, first_token_timeout=30, cfg=cfg) return {'ok': True, 'latency_ms': int((time.time() - t0) * 1000), 'first_token_ms': r.get('first_token_ms'), 'reply': r['text'][:50], 'cost': r['cost'], 'tokens': r['total_tokens']}