132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""LLM 客户端（OpenAI 兼容接口）
|
|
|
|
支持按模型配置调用（base_url / api_key / model / temperature / timeout），
|
|
支持多模态消息（content 为 [{type:text},{type:image_url}] 列表，用于视觉模型分析截图）。
|
|
"""
|
|
import json
|
|
import re
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
from config import LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, LLM_TEMPERATURE, LLM_TIMEOUT
|
|
|
|
|
|
class LLMError(Exception):
|
|
pass
|
|
|
|
|
|
def _cfg(cfg):
|
|
"""补齐默认值，返回完整配置 dict"""
|
|
base = {
|
|
'base_url': LLM_BASE_URL,
|
|
'api_key': LLM_API_KEY,
|
|
'model': LLM_MODEL,
|
|
'temperature': LLM_TEMPERATURE,
|
|
'timeout': LLM_TIMEOUT,
|
|
'vision': False,
|
|
}
|
|
if cfg:
|
|
base.update({k: v for k, v in cfg.items() if v is not None})
|
|
return base
|
|
|
|
|
|
def _extract_json(text):
|
|
"""从 LLM 输出中提取 JSON（容忍 markdown 代码块等包裹）"""
|
|
if not text:
|
|
return None
|
|
text = text.strip()
|
|
# 去掉 markdown 代码块
|
|
fence = re.search(r'```(?:json)?\s*(.*?)```', text, re.S)
|
|
if fence:
|
|
text = fence.group(1).strip()
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
# 尝试截取第一个 { 到最后一个 }
|
|
s, e = text.find('{'), text.rfind('}')
|
|
if s != -1 and e > s:
|
|
try:
|
|
return json.loads(text[s:e + 1])
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return None
|
|
|
|
|
|
def chat(messages, cfg=None, temperature=None, max_tokens=None, timeout=None):
|
|
"""调用 OpenAI 兼容 chat/completions，返回 content 字符串
|
|
|
|
messages 中的 content 可以是字符串，也可以是多模态列表：
|
|
[{"type":"text","text":"..."},
|
|
{"type":"image_url","image_url":{"url":"data:image/png;base64,..."}}]
|
|
"""
|
|
c = _cfg(cfg)
|
|
url = f"{c['base_url'].rstrip('/')}/chat/completions"
|
|
body = {
|
|
'model': c['model'],
|
|
'messages': messages,
|
|
'temperature': temperature if temperature is not None else c['temperature'],
|
|
}
|
|
if max_tokens:
|
|
body['max_tokens'] = max_tokens
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(body).encode('utf-8'),
|
|
headers={
|
|
'Content-Type': 'application/json',
|
|
'Authorization': f'Bearer {c["api_key"]}',
|
|
},
|
|
method='POST',
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=timeout or c['timeout']) as resp:
|
|
data = json.loads(resp.read().decode('utf-8'))
|
|
except urllib.error.HTTPError as e:
|
|
detail = e.read().decode('utf-8', 'ignore')[:300]
|
|
raise LLMError(f'LLM HTTP {e.code}: {detail}')
|
|
except Exception as e:
|
|
raise LLMError(f'LLM 调用失败: {e}')
|
|
|
|
try:
|
|
msg = data['choices'][0]['message']
|
|
content = msg.get('content') or ''
|
|
if not content.strip():
|
|
# thinking 模型（如 qwen3/deepseek-r1 系）：最终答案可能为空，
|
|
# 思考过程放在 reasoning_content 里（常包含最终 JSON），用它兜底
|
|
reasoning = msg.get('reasoning_content') or ''
|
|
if reasoning.strip():
|
|
return reasoning
|
|
return content
|
|
except (KeyError, IndexError, TypeError):
|
|
raise LLMError(f'LLM 响应异常: {str(data)[:300]}')
|
|
|
|
|
|
def _trim_content(text, limit=300):
|
|
"""截断过长的模型输出（避免超长思考内容反复进入上下文）"""
|
|
text = (text or '').strip()
|
|
return text if len(text) <= limit else text[:limit] + '...(已截断)'
|
|
|
|
|
|
def chat_json(messages, cfg=None, temperature=None, max_tokens=None, retries=2):
|
|
"""调用 LLM 并强制解析 JSON，失败重试"""
|
|
last_err = None
|
|
last_content = ''
|
|
for i in range(retries + 1):
|
|
try:
|
|
last_content = chat(messages, cfg=cfg, temperature=temperature,
|
|
max_tokens=max_tokens)
|
|
obj = _extract_json(last_content)
|
|
if obj is not None:
|
|
return obj
|
|
last_err = f'无法从输出解析 JSON: {last_content[:200]}'
|
|
except LLMError as e:
|
|
last_err = str(e)
|
|
if i < retries:
|
|
messages = messages + [
|
|
{'role': 'assistant', 'content': _trim_content(last_content)},
|
|
{'role': 'user',
|
|
'content': f'刚才的输出不是合法 JSON，请只输出严格的 JSON 对象，不要输出思考过程和代码块。错误: {last_err}'},
|
|
]
|
|
raise LLMError(f'LLM JSON 解析失败: {last_err}')
|