77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
LLM 客户端（OpenAI 兼容 /chat/completions）。
|
||
- base_url / model / key 由环境变量配置（默认智谱 GLM）
|
||
- 出站安全：仅 http/https、显式拒绝 localhost、解析 IP 拒绝环回/私有/保留段、
|
||
禁用重定向（防 DNS rebinding 绕过），值全部走 JSON 序列化，密钥不落日志
|
||
"""
|
||
import ipaddress
|
||
import json
|
||
import os
|
||
from urllib.parse import urlparse
|
||
|
||
import requests
|
||
|
||
|
||
class LlmError(Exception):
|
||
pass
|
||
|
||
|
||
def _validated_url(base_url: str) -> str:
|
||
u = urlparse(base_url)
|
||
if u.scheme not in ('http', 'https'):
|
||
raise LlmError('LLM base_url 仅允许 http/https')
|
||
host = u.hostname or ''
|
||
if not host or host.lower() in ('localhost', 'localhost.localdomain'):
|
||
raise LlmError('LLM base_url 拒绝 localhost')
|
||
port = u.port or (443 if u.scheme == 'https' else 80)
|
||
try:
|
||
infos = socket.getaddrinfo(host, port)
|
||
except socket.gaierror as e:
|
||
raise LlmError('LLM base_url 域名解析失败: {}'.format(host))
|
||
for info in infos:
|
||
ip = ipaddress.ip_address(info[4][0])
|
||
if (ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_reserved
|
||
or ip.is_multicast or ip.is_unspecified):
|
||
raise LlmError('LLM base_url 拒绝非公网地址: {}'.format(ip))
|
||
return '{}://{}{}'.format(u.scheme, u.netloc, u.path)
|
||
|
||
|
||
class LlmClient:
|
||
|
||
def __init__(self):
|
||
self.api_key = os.environ.get('JQUANT_LLM_API_KEY', '').strip()
|
||
self.base_url = (os.environ.get('JQUANT_LLM_BASE_URL', '').strip()
|
||
or 'https://open.bigmodel.cn/api/paas/v4')
|
||
self.model = os.environ.get('JQUANT_LLM_MODEL', '').strip() or 'glm-4-flash'
|
||
self.temperature = float(os.environ.get('JQUANT_LLM_TEMPERATURE', '0.2'))
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
return bool(self.api_key)
|
||
|
||
def chat(self, system_prompt: str, user_prompt: str) -> str:
|
||
if not self.enabled:
|
||
raise LlmError('未配置 LLM API Key（JQUANT_LLM_API_KEY），请在服务环境变量中设置')
|
||
url = _validated_url(self.base_url.rstrip('/')) + '/chat/completions'
|
||
body = {
|
||
'model': self.model,
|
||
'temperature': self.temperature,
|
||
'messages': [
|
||
{'role': 'system', 'content': system_prompt},
|
||
{'role': 'user', 'content': user_prompt},
|
||
],
|
||
}
|
||
# 校验与请求紧邻；禁重定向防 DNS rebinding 绕过 IP 校验
|
||
r = requests.post(url, json=body, timeout=90, allow_redirects=False,
|
||
headers={'Authorization': 'Bearer ' + self.api_key,
|
||
'Content-Type': 'application/json'})
|
||
if r.status_code in (301, 302, 303, 307, 308):
|
||
raise LlmError('LLM 端点发生重定向，已拒绝（防 SSRF 绕过）')
|
||
if r.status_code != 200:
|
||
raise LlmError('LLM HTTP {}: {}'.format(r.status_code, r.text[:200]))
|
||
content = r.json().get('choices', [{}])[0].get('message', {}).get('content')
|
||
if not content:
|
||
raise LlmError('LLM 响应缺少 content')
|
||
return content
|