cad-agent/tests/test_llm_client.py

"""`dwg_review.llm_client` 单元测试（绝不真实调用 LLM）。"""

import base64
import io

import pytest

from dwg_review import config
from dwg_review import llm_client as llm_module
from dwg_review.llm_client import LlmClient, LlmError, downsample_image_b64
from tests.fakes import llm_text


# ---------------------------------------------------------------------------
#  假 OpenAI SDK
# ---------------------------------------------------------------------------
class _Completions:
    def __init__(self, outer):
        self.outer = outer

    def create(self, **kwargs):
        self.outer.kwargs.append(kwargs)
        if self.outer.raises:
            raise RuntimeError("upstream boom")
        return self.outer.result


class _Chat:
    def __init__(self, outer):
        self.completions = _Completions(outer)


class _FakeOpenAI:
    def __init__(self, **kwargs):
        self.init_kwargs = kwargs
        self.kwargs = []
        self.raises = False
        self.result = None
        self.chat = _Chat(self)


@pytest.fixture
def fake_openai(monkeypatch):
    monkeypatch.setattr(llm_module, "OpenAI", _FakeOpenAI)
    monkeypatch.setattr(llm_module.time, "sleep", lambda *_: None)
    return _FakeOpenAI


def build_client(fake_openai, **kw):
    return LlmClient(api_key="sk-test", **kw)


# ---------------------------------------------------------------------------
#  构造
# ---------------------------------------------------------------------------
def test_init_requires_api_key(monkeypatch):
    monkeypatch.setattr(config, "API_KEY", "")
    with pytest.raises(LlmError):
        LlmClient()


def test_init_passes_key_and_base_url(fake_openai):
    client = LlmClient(api_key="sk-abc", base_url="https://example.test")
    assert client._client.init_kwargs["api_key"] == "sk-abc"
    assert client._client.init_kwargs["base_url"] == "https://example.test"


def test_init_uses_config_defaults(fake_openai, monkeypatch):
    monkeypatch.setattr(config, "MODEL_NAME", "test-model")
    monkeypatch.setattr(config, "BASE_URL", "https://cfg.test")
    client = LlmClient(api_key="sk-abc")
    assert client.model == "test-model"
    assert client._client.init_kwargs["base_url"] == "https://cfg.test"


# ---------------------------------------------------------------------------
#  chat / chat_text
# ---------------------------------------------------------------------------
def test_chat_returns_raw_response(fake_openai):
    client = build_client(fake_openai)
    client._client.result = llm_text("hi")
    resp = client.chat([{"role": "user", "content": "x"}])
    assert resp.choices[0].message.content == "hi"


def test_chat_includes_tools_only_when_present(fake_openai):
    client = build_client(fake_openai)
    client._client.result = llm_text("x")
    client.chat([{"role": "user", "content": "x"}])
    assert "tools" not in client._client.kwargs[0]

    client.chat([{"role": "user", "content": "x"}], tools=[{"type": "function"}])
    assert client._client.kwargs[1]["tools"] == [{"type": "function"}]


def test_chat_wraps_exception_into_llmerror(fake_openai):
    client = build_client(fake_openai)
    client._client.raises = True
    with pytest.raises(LlmError):
        client.chat([{"role": "user", "content": "x"}], retries=0)


def test_chat_retries_then_succeeds(fake_openai):
    client = build_client(fake_openai)

    class Flaky(_Completions):
        def __init__(self, outer):
            super().__init__(outer)
            self.n = 0

        def create(self, **kwargs):
            self.n += 1
            if self.n < 3:
                raise RuntimeError("transient")
            return llm_text("ok")

    client._client.chat.completions = Flaky(client._client)
    resp = client.chat([{"role": "user", "content": "x"}], retries=2)
    assert resp.choices[0].message.content == "ok"
    assert client._client.chat.completions.n == 3


def test_chat_text_returns_assistant_text(fake_openai):
    client = build_client(fake_openai)
    client._client.result = llm_text("报告正文")
    assert client.chat_text("sys", "user") == "报告正文"


def test_image_content_format():
    block = LlmClient.image_content("QUJD", "image/jpeg")
    assert block == {
        "type": "image_url",
        "image_url": {"url": "data:image/jpeg;base64,QUJD"},
    }


# ---------------------------------------------------------------------------
#  截图下采样
# ---------------------------------------------------------------------------
def test_downsample_passthrough_without_pil(monkeypatch):
    monkeypatch.setattr(llm_module, "HAS_PIL", False)
    assert downsample_image_b64("QUJD", "image/png", max_edge=10) == ("QUJD", "image/png")


def test_downsample_passthrough_when_max_edge_zero():
    assert downsample_image_b64("QUJD", "image/png", max_edge=0) == ("QUJD", "image/png")


def test_downsample_passthrough_when_no_resize_needed():
    pytest.importorskip("PIL")
    from PIL import Image

    buf = io.BytesIO()
    Image.new("RGB", (100, 80), "white").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    out, mime = downsample_image_b64(b64, "image/png", max_edge=1280)
    assert out == b64 and mime == "image/png"


def test_downsample_resizes_large_image():
    pytest.importorskip("PIL")
    from PIL import Image

    buf = io.BytesIO()
    Image.new("RGB", (2000, 1000), "white").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    out, mime = downsample_image_b64(b64, "image/png", max_edge=500)
    assert mime == "image/png"
    img = Image.open(io.BytesIO(base64.b64decode(out)))
    assert max(img.size) <= 500


def test_downsample_converts_to_jpeg():
    pytest.importorskip("PIL")
    from PIL import Image

    buf = io.BytesIO()
    Image.new("RGB", (2000, 1000), "white").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    out, mime = downsample_image_b64(b64, "image/png", max_edge=400, fmt="jpeg")
    assert mime == "image/jpeg"
    assert Image.open(io.BytesIO(base64.b64decode(out))).format == "JPEG"


def test_downsample_bad_base64_returns_original():
    pytest.importorskip("PIL")
    bad = base64.b64encode(b"this is not an image").decode("ascii")
    assert downsample_image_b64(bad, "image/png", max_edge=10) == (bad, "image/png")