90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""LLM 客户端（OpenAI 兼容接口）"""
|
|
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 _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, temperature=None, max_tokens=None, timeout=None):
|
|
"""调用 OpenAI 兼容 chat/completions，返回 content 字符串"""
|
|
url = f'{LLM_BASE_URL}/chat/completions'
|
|
body = {
|
|
'model': LLM_MODEL,
|
|
'messages': messages,
|
|
'temperature': temperature if temperature is not None else LLM_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 {LLM_API_KEY}',
|
|
},
|
|
method='POST',
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=timeout or LLM_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:
|
|
return data['choices'][0]['message']['content']
|
|
except (KeyError, IndexError, TypeError):
|
|
raise LLMError(f'LLM 响应异常: {str(data)[:300]}')
|
|
|
|
|
|
def chat_json(messages, temperature=None, max_tokens=None, retries=2):
|
|
"""调用 LLM 并强制解析 JSON，失败重试"""
|
|
last_err = None
|
|
for i in range(retries + 1):
|
|
try:
|
|
content = chat(messages, temperature=temperature, max_tokens=max_tokens)
|
|
obj = _extract_json(content)
|
|
if obj is not None:
|
|
return obj
|
|
last_err = f'无法从输出解析 JSON: {content[:200]}'
|
|
except LLMError as e:
|
|
last_err = str(e)
|
|
if i < retries:
|
|
messages = messages + [
|
|
{'role': 'assistant', 'content': content if 'content' in dir() else ''},
|
|
{'role': 'user', 'content': f'刚才的输出不是合法 JSON，请只输出严格的 JSON 对象。错误: {last_err}'},
|
|
]
|
|
raise LLMError(f'LLM JSON 解析失败: {last_err}')
|