llm-in-text/backend/tests/test_llm_extended.py

import asyncio
import importlib
import sys
from pathlib import Path
import pytest


BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

try:
    llm = importlib.import_module("llm")
except ModuleNotFoundError:
    pytest.skip("llm module dependencies are not available", allow_module_level=True)


def test_extract_message_with_object_message_content_and_thinking():
    class Msg:
        def __init__(self, content, thinking):
            self.content = content
            self.thinking = thinking

    class Resp:
        def __init__(self, message):
            self.message = message

    resp = Resp(Msg("hello world", "thinking about it"))
    content, thinking = llm._extract_message(resp)
    assert content == "hello world"
    assert thinking == "thinking about it"


def test_extract_message_with_object_message_empty_content():
    class Msg:
        def __init__(self, content, thinking):
            self.content = content
            self.thinking = thinking

    class Resp:
        def __init__(self, message):
            self.message = message

    resp = Resp(Msg("", None))
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


def test_extract_message_with_dict_message():
    resp = {"message": {"content": "ok", "thinking": "calc"}}
    content, thinking = llm._extract_message(resp)
    assert content == "ok"
    assert thinking == "calc"


def test_extract_message_dict_no_message_key():
    resp = {"not_message": {"content": "irrelevant"}}
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


def test_extract_message_dict_message_content_none_and_thinking_none():
    resp = {"message": {"content": None, "thinking": None}}
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


def test_extract_message_dict_message_thinking_none():
    resp = {"message": {"content": "val", "thinking": None}}
    content, thinking = llm._extract_message(resp)
    assert content == "val"
    assert thinking == ""


def test_extract_message_empty_dict():
    resp = {}
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


def test_call_ollama_no_system_message(monkeypatch):
    captured = {}

    async def fake_generate(**kwargs):
        captured["kwargs"] = kwargs
        return {"response": "ok"}

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    result = asyncio.run(
        llm.call_ollama("user prompt body", system_prompt=None, tag="no-system", temperature=0.1)
    )
    assert result["content"] == "ok"
    assert captured["kwargs"]["prompt"] == "user prompt body"
    assert captured["kwargs"]["raw"] is True


def test_call_ollama_whitespace_system_message(monkeypatch):
    captured = {}

    async def fake_generate(**kwargs):
        captured["kwargs"] = kwargs
        return {"response": "ok"}

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    result = asyncio.run(
        llm.call_ollama("user prompt", system_prompt="   ", tag="whitespace-system", temperature=0.1)
    )
    assert result["content"] == "ok"
    assert captured["kwargs"]["prompt"] == "user prompt"
    assert captured["kwargs"]["raw"] is True


def test_call_ollama_thinking_in_kwargs(monkeypatch):
    captured = {}

    async def fake_generate(**kwargs):
        captured.update(kwargs)
        return {"response": "ok", "message": {"thinking": "boom"}} # keep thinking for backward test though it might not be perfect

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    res = asyncio.run(
        llm.call_ollama("prompt", thinking="boom", tag="think-flag", temperature=0.7)
    )
    assert res["content"] == "ok" and res["think"] == "boom"
    assert captured.get("think") == "boom"


def test_call_ollama_cancelled_reraises(monkeypatch):
    async def fake_generate(**kwargs):
        raise asyncio.CancelledError

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    with pytest.raises(asyncio.CancelledError):
        asyncio.run(
            llm.call_ollama("prompt", system_prompt=None, tag="cancel", temperature=0.7)
        )


def test_call_ollama_chat_raises_rethrows(monkeypatch):
    async def fake_generate(**kwargs):
        raise ValueError("boom")

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    with pytest.raises(ValueError):
        asyncio.run(
            llm.call_ollama("prompt", system_prompt=None, tag="exception", temperature=0.7)
        )


def test_call_ollama_returns_content_and_think_from_response(monkeypatch):
    async def fake_generate(**kwargs):
        return {"response": "final", "message": {"thinking": "process"}}

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    res = asyncio.run(
        llm.call_ollama("prompt", system_prompt=None, tag="return", temperature=0.7)
    )
    assert res["content"] == "final" and res["think"] == "process"


def test_stream_ollama_uses_requested_model_and_yields_chunks(monkeypatch):
    captured = {}

    class FakeStream:
        def __init__(self, chunks):
            self._chunks = iter(chunks)

        def __aiter__(self):
            return self

        async def __anext__(self):
            try:
                return next(self._chunks)
            except StopIteration:
                raise StopAsyncIteration

    async def fake_generate(**kwargs):
        captured["kwargs"] = kwargs
        return FakeStream([
            {"response": "深度"},
            {"response": "回答"},
        ])

    monkeypatch.setattr(llm.client, "generate", fake_generate)

    chunks = []

    async def collect_stream():
        async for chunk in llm.stream_ollama(
            "prompt",
            system_prompt="system",
            tag="stream",
            temperature=0.8,
            model="pro-model",
            use_pro_model=True,
        ):
            chunks.append(chunk)

    asyncio.run(collect_stream())

    assert "".join(chunks) == "深度回答"
    assert captured["kwargs"]["model"] == "pro-model"
    assert captured["kwargs"]["stream"] is True


def test_call_vlm_ocr_passes_image_and_prompt(monkeypatch):
    image_bytes = b"image-bytes"
    called = {}
    monkeypatch.setattr(llm, "get_vlm_ocr_prompt", lambda: "OCR PROMPT")

    async def fake_chat(**kwargs):
        called["kwargs"] = kwargs
        return {"message": {"content": "ocr result", "thinking": ""}}

    monkeypatch.setattr(llm.client, "chat", fake_chat)

    result = asyncio.run(llm.call_vlm_ocr(image_bytes, language="auto"))

    messages = called["kwargs"].get("messages", [])
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == "OCR PROMPT"
    assert messages[0]["images"] == [image_bytes]
    assert result == "ocr result"


def test_call_vlm_ocr_chat_raises_rethrows(monkeypatch):
    image_bytes = b"image-bytes"
    monkeypatch.setattr(llm, "get_vlm_ocr_prompt", lambda: "OCR PROMPT")

    async def fake_chat(**kwargs):
        raise RuntimeError("ocr fail")
    monkeypatch.setattr(llm.client, "chat", fake_chat)
    with pytest.raises(RuntimeError):
        asyncio.run(llm.call_vlm_ocr(image_bytes))


def test_call_vlm_ocr_returns_content_from_response(monkeypatch):
    image_bytes = b"img"
    monkeypatch.setattr(llm, "get_vlm_ocr_prompt", lambda: "OCR PROMPT")

    async def fake_chat(**kwargs):
        return {"message": {"content": "ocr text", "thinking": ""}}
    monkeypatch.setattr(llm.client, "chat", fake_chat)
    content = asyncio.run(llm.call_vlm_ocr(image_bytes))
    assert content == "ocr text"