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_content_and_thinking():
    resp = {"choices": [{"message": {"content": "hello world", "thinking": "reasoning"}}]}
    content, thinking = llm._extract_message(resp)
    assert content == "hello world"
    assert thinking == "reasoning"


def test_extract_message_empty_content():
    resp = {"choices": [{"message": {"content": "", "thinking": None}}]}
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


def test_extract_message_dict_no_choices():
    resp = {"not_choices": []}
    content, thinking = llm._extract_message(resp)
    assert content == ""
    assert thinking == ""


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


def test_extract_delta_text_from_chunk():
    chunk = {"choices": [{"delta": {"content": "text"}}]}
    assert llm._extract_delta_text(chunk) == "text"


def test_extract_delta_thinking_from_chunk():
    chunk = {"choices": [{"delta": {"thinking": "thought"}}]}
    assert llm._extract_delta_thinking(chunk) == "thought"


def test_call_ollama_no_system(monkeypatch):
    captured = {}

    async def fake_post(url, json=None):
        captured["json"] = json

        class FakeResp:
            def raise_for_status(self): pass
            def json(self): return {"choices": [{"message": {"content": "ok"}}]}

        return FakeResp()

    async def fake_client(*args, **kwargs):
        class Ctx:
            async def __aenter__(self2): return self2
            async def __aexit__(*a): pass
            post = fake_post
        return Ctx()

    monkeypatch.setattr(llm.httpx, "AsyncClient", fake_client)
    monkeypatch.setattr(llm.asyncio, "wait_for", lambda coro, **kw: coro)

    result = asyncio.run(
        llm.call_ollama("user prompt", system_prompt=None, tag="no-system")
    )

    assert result["content"] == "ok"
    # Should only have user message, no system
    assert len(captured["json"]["messages"]) == 1


def test_call_ollama_with_system(monkeypatch):
    captured = {}

    async def fake_post(url, json=None):
        captured["json"] = json

        class FakeResp:
            def raise_for_status(self): pass
            def json(self): return {"choices": [{"message": {"content": "ok"}}]}

        return FakeResp()

    async def fake_client(*args, **kwargs):
        class Ctx:
            async def __aenter__(self2): return self2
            async def __aexit__(*a): pass
            post = fake_post
        return Ctx()

    monkeypatch.setattr(llm.httpx, "AsyncClient", fake_client)
    monkeypatch.setattr(llm.asyncio, "wait_for", lambda coro, **kw: coro)

    result = asyncio.run(
        llm.call_ollama("user prompt", system_prompt="sys prompt", tag="with-system")
    )

    assert result["content"] == "ok"
    # Should have both system and user messages
    msgs = captured["json"]["messages"]
    assert len(msgs) == 2
    assert msgs[0]["role"] == "system"


def test_call_ollama_with_custom_model(monkeypatch):
    captured = {}

    async def fake_post(url, json=None):
        captured["json"] = json

        class FakeResp:
            def raise_for_status(self): pass
            def json(self): return {"choices": [{"message": {"content": "ok"}}]}

        return FakeResp()

    async def fake_client(*args, **kwargs):
        class Ctx:
            async def __aenter__(self2): return self2
            async def __aexit__(*a): pass
            post = fake_post
        return Ctx()

    monkeypatch.setattr(llm.httpx, "AsyncClient", fake_client)
    monkeypatch.setattr(llm.asyncio, "wait_for", lambda coro, **kw: coro)

    result = asyncio.run(
        llm.call_ollama("prompt", model="custom-model")
    )

    assert captured["json"]["model"] == "custom-model"


def test_stream_ollama_events_error_handling(monkeypatch):
    def make_lines():
        lines_iter = iter([
            'data: {"error": "model not found"}',
        ])

        class LineIterator:
            async def __anext__(self):
                try:
                    return next(lines_iter)
                except StopIteration:
                    raise StopAsyncIteration()

        class Response:
            def __init__(self2): self2._lines = LineIterator()

            async def raise_for_status(self2): pass
            async def aiter_lines(self2): return self2._lines

        class StreamCtx:
            async def __aenter__(self2): return Response()
            async def __aexit__(*a): pass

        class Client:
            stream = lambda self2, *args, **kw: StreamCtx()

        return Client()

    async def fake_client(*args, **kwargs):
        return make_lines()

    monkeypatch.setattr(llm.httpx, "AsyncClient", fake_client)
    monkeypatch.setattr(llm.asyncio, "wait_for", lambda coro, **kw: coro)

    async def collect():
        try:
            async for _ in llm.stream_ollama_events("prompt", tag="err"):
                pass
        except RuntimeError as e:
            return str(e)

    result = asyncio.run(collect())
    assert "model not found" in str(result)


def test_call_vlm_ocr_payload_format(monkeypatch):
    captured = {}

    async def fake_post(url, json=None):
        captured["json"] = json

        class FakeResp:
            def raise_for_status(self): pass
            def json(self): return {"choices": [{"message": {"content": "ocr result"}}]}

        return FakeResp()

    async def fake_client(*args, **kwargs):
        class Ctx:
            async def __aenter__(self2): return self2
            async def __aexit__(*a): pass
            post = fake_post
        return Ctx()

    monkeypatch.setattr(llm.httpx, "AsyncClient", fake_client)
    monkeypatch.setattr(llm.asyncio, "wait_for", lambda coro, **kw: coro)

    result = asyncio.run(llm.call_vlm_ocr(b"image"))
    assert result == "ocr result"

    # Verify vision format: image_url content part with base64
    msgs = captured["json"]["messages"]
    assert len(msgs) == 1
    content_parts = msgs[0]["content"]
    image_part = [p for p in content_parts if p.get("type") == "image_url"]
    assert len(image_part) == 1