tests/test_extract_vision.py

from __future__ import annotations

import json

import httpx

from app.config import Settings
from app.extract import extract_receipt
from backends.gemma import GemmaLLM
from backends.ollama import OllamaLLM
from tests.conftest import tiny_jpeg

EXTRACT = {
    "doc_kind": "receipt",
    "category": "dining",
    "vendor": "Cafe",
    "date": "2026-08-21",
    "tax": 0.5,
    "total": 8.0,
    "currency": "USD",
    "line_items": [],
}


def _chat_handler(sink: list[dict]):
    def handler(request: httpx.Request) -> httpx.Response:
        payload = json.loads(request.content)
        sink.append(payload)
        return httpx.Response(
            200,
            json={"choices": [{"message": {"content": json.dumps(EXTRACT)}}]},
        )

    return handler


def test_gemma_sends_image_url(settings: Settings) -> None:
    sink: list[dict] = []
    client = httpx.Client(
        transport=httpx.MockTransport(_chat_handler(sink)),
        base_url="http://llm.test/v1",
    )
    llm = GemmaLLM(settings, client=client)
    extract = extract_receipt(
        llm, settings=settings, image_jpeg=tiny_jpeg(), ocr_text=None
    )
    assert extract.vendor == "Cafe"
    content = sink[0]["messages"][1]["content"]
    assert isinstance(content, list)
    kinds = {part["type"] for part in content}
    assert "image_url" in kinds
    url = next(part["image_url"]["url"] for part in content if part["type"] == "image_url")
    assert url.startswith("data:image/jpeg;base64,")


def test_lightning_never_sends_image(settings: Settings) -> None:
    settings = settings.model_copy(
        update={
            "llm_backend": "ollama",
            "llm_model": "nemotron-3.5-lightning",
            "llm_accepts_images": False,
        }
    )
    sink: list[dict] = []
    client = httpx.Client(
        transport=httpx.MockTransport(_chat_handler(sink)),
        base_url="http://llm.test/v1",
    )
    llm = OllamaLLM(settings, client=client)
    assert llm.accepts_images is False
    extract = extract_receipt(
        llm, settings=settings, image_jpeg=tiny_jpeg(), ocr_text="Cafe 8.00"
    )
    assert extract.total is not None
    content = sink[0]["messages"][1]["content"]
    assert isinstance(content, str)
    dumped = json.dumps(sink[0])
    assert "image_url" not in dumped
    assert "data:image" not in dumped