gjm
1cbcbb92f4
- 视觉 token 成本由像素尺寸决定，故仅在长边超过 IMAGE_MAX_EDGE(默认1280) 时才缩放 - 不需要缩放且格式未变时原样返回，不重新编码，避免 CAD 线稿的无谓压缩损失 - 缩放后默认输出无损 PNG（IMAGE_FORMAT=png），可设 jpeg 换取更小体积 - 缺少 Pillow 时自动降级为原图直传
120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
"""LLM 封装（文本 + 视觉），统一超时 / 重试 / 截图预处理。
|
|
|
|
默认走 DeepSeek（视觉模型 deepseek-flash）。密钥从环境变量读取，禁止硬编码。
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
import time
|
|
|
|
from openai import OpenAI
|
|
|
|
from . import config
|
|
|
|
try: # 可选依赖：仅用于截图下采样
|
|
from PIL import Image
|
|
|
|
HAS_PIL = True
|
|
except ImportError: # pragma: no cover
|
|
HAS_PIL = False
|
|
|
|
|
|
class LlmError(Exception):
|
|
pass
|
|
|
|
|
|
def downsample_image_b64(b64, mime="image/png", max_edge=None, fmt=None, quality=None):
|
|
"""缩放截图，降低视觉 token 成本。
|
|
|
|
策略（视觉 token 成本主要由**像素尺寸**决定，而非文件大小）：
|
|
- 仅当长边超过 `max_edge` 时才缩放；
|
|
- 既不需要缩放、目标格式又与原格式一致时，**原样返回**（不重新编码，
|
|
避免给 CAD 线稿引入无谓的压缩损失）；
|
|
- 需要缩放时按 `fmt` 编码：`png`（默认，无损，适合线稿）/ `jpeg`（体积更小）。
|
|
|
|
缺少 Pillow 或 max_edge <= 0 时原样返回 (b64, mime)。
|
|
"""
|
|
max_edge = config.IMAGE_MAX_EDGE if max_edge is None else max_edge
|
|
fmt = (config.IMAGE_FORMAT if fmt is None else fmt).strip().lower()
|
|
quality = config.IMAGE_JPEG_QUALITY if quality is None else quality
|
|
if not HAS_PIL or max_edge <= 0:
|
|
return b64, mime
|
|
|
|
target_mime = "image/jpeg" if fmt == "jpeg" else "image/png"
|
|
try:
|
|
img = Image.open(io.BytesIO(base64.b64decode(b64)))
|
|
w, h = img.size
|
|
scale = max_edge / float(max(w, h))
|
|
need_resize = scale < 1.0
|
|
|
|
if not need_resize and mime == target_mime:
|
|
return b64, mime
|
|
|
|
if need_resize:
|
|
img = img.resize(
|
|
(max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS
|
|
)
|
|
|
|
buf = io.BytesIO()
|
|
if target_mime == "image/jpeg":
|
|
if img.mode not in ("RGB", "L"):
|
|
img = img.convert("RGB")
|
|
img.save(buf, format="JPEG", quality=quality)
|
|
else:
|
|
img.save(buf, format="PNG", optimize=True)
|
|
return base64.b64encode(buf.getvalue()).decode("ascii"), target_mime
|
|
except Exception:
|
|
return b64, mime
|
|
|
|
|
|
class LlmClient:
|
|
"""统一的 LLM 调用入口。"""
|
|
|
|
def __init__(self, api_key=None, base_url=None, model=None):
|
|
key = config.API_KEY if api_key is None else api_key
|
|
if not key:
|
|
raise LlmError("未配置 DEEPSEEK_API_KEY 环境变量")
|
|
self.model = model or config.MODEL_NAME
|
|
self._client = OpenAI(api_key=key, base_url=base_url or config.BASE_URL)
|
|
|
|
def chat(self, messages, *, tools=None, timeout=None, retries=None):
|
|
"""发起一次对话补全，返回原始 response。"""
|
|
retries = config.LLM_RETRIES if retries is None else retries
|
|
last = None
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
kwargs = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"timeout": timeout or config.LLM_TIMEOUT,
|
|
}
|
|
if tools:
|
|
kwargs["tools"] = tools
|
|
return self._client.chat.completions.create(**kwargs)
|
|
except Exception as exc: # noqa: BLE001 - 统一包装为 LlmError
|
|
last = exc
|
|
if attempt < retries:
|
|
time.sleep(0.5 * (2 ** attempt))
|
|
continue
|
|
raise LlmError(f"LLM 调用失败: {exc}") from exc
|
|
raise LlmError(f"LLM 调用失败: {last}")
|
|
|
|
def chat_text(self, system, user, **kwargs):
|
|
"""纯文本对话，返回助手文本。"""
|
|
resp = self.chat(
|
|
[
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": user},
|
|
],
|
|
**kwargs,
|
|
)
|
|
return resp.choices[0].message.content
|
|
|
|
@staticmethod
|
|
def image_content(b64, mime="image/png"):
|
|
"""构造 OpenAI/DeepSeek 视觉消息里的图片内容块。"""
|
|
return {
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:{mime};base64,{b64}"},
|
|
}
|