cad-agent/dwg_review/llm_client.py

"""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


def image_size_b64(b64, mime="image/png"):
    """解码图片并返回 `(宽, 高)`；缺 Pillow / 解码失败返回 None。

    用于瓦片规划前测量**绘图区实际像素**（决定单张截图能覆盖多大世界范围）。
    """
    if not HAS_PIL or not b64:
        return None
    try:
        img = Image.open(io.BytesIO(base64.b64decode(b64)))
        w, h = img.size
        return (int(w), int(h))
    except Exception:
        return None


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}"},
        }