File size: 2,338 Bytes
2edb151 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | 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
|