hz4th_coder
07f597b57e
1. 修复假流式：对话接口边收边吐(stream_chat 生成器直驱 _chat_stream_raw)；思考模型先流式思考内容(折叠)再流式回答；回答块显示 tok/s 2. 历史会话列表：重命名/置顶/删除(pinned/use_kb 列)；📝Markdown 一键开关，默认 Markdown 渲染 3. 模型能力标签(chat/thinking/vision/audio_in/audio_out/image_gen/video_gen/embedding/rerank)：接口库每模型勾选； 对话按能力适配：视觉→传图、语音入→🎤录音(默认ASR转写)、语音出→🔊朗读(默认TTS)、思考→思考折叠；不匹配提前提示 4. 模型库标签：能力矩阵 + 系统默认模型(语音识别/合成/生图/生视频/embedding/rerank) 5. 计费修正：按次=元/千次；逐模型定价=模型名 输入价 缓存输入价 输出价(缓存价计缓存命中token) 6. 弹窗未保存提醒(点外部/X 时 confirm，全弹窗通用) 7. 知识库导航：#/kb 全局文档 CRUD + txt/md/pdf/docx 上传解析 + jieba BM25 检索；对话可一键注入知识库
599 lines
26 KiB
Python
599 lines
26 KiB
Python
# -*- 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']}
|