fegger
7e706793fa
- Add AGENTS.md and project-specific Zed skills (odoo-ocr-pipeline, odoo-xml-import, local-vlm-client). - Implement Pydantic schemas for documents, invoices, review results, and VLM responses. - Add unified BaseVLMClient with Ollama implementation and llama.cpp stub. - Build pipeline stages: loader, classifier, digital_pdf/scanned_print/handwritten/mixed_unknown branches, extractor, reviewer, xml_builder. - Add CLI entry point with sidecar JSON and confidence-gated XML output. - Include prompts for classifier, OCR, extraction, and review models. - Add tests with FakeVLMClient; pytest, ruff, and mypy all pass.
75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
"""Tests for the Ollama VLM client."""
|
|
|
|
import base64
|
|
import json
|
|
from io import BytesIO
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from odoo_ocr.clients.base import BaseVLMClient
|
|
from odoo_ocr.clients.ollama import OllamaClient
|
|
from odoo_ocr.config import Settings
|
|
from odoo_ocr.schemas import ExtractedInvoice
|
|
|
|
|
|
@pytest.fixture
|
|
def settings() -> Settings:
|
|
return Settings(
|
|
ollama_base_url="http://test-ollama:11434",
|
|
cache=Settings().model_dump()["cache"] | {"enabled": False},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_complete_parses_response(settings: Settings) -> None:
|
|
client = OllamaClient(settings)
|
|
fake_response = {
|
|
"model": "glm-ocr",
|
|
"message": {"role": "assistant", "content": '{"text": "Invoice 123"}'},
|
|
"done": True,
|
|
"prompt_eval_count": 100,
|
|
"eval_count": 20,
|
|
}
|
|
|
|
resp = MagicMock()
|
|
resp.json.return_value = fake_response
|
|
resp.raise_for_status = lambda: None
|
|
client.http = AsyncMock()
|
|
client.http.post.return_value = resp
|
|
|
|
img = Image.new("RGB", (50, 50), color="red")
|
|
response = await client.complete(
|
|
system_prompt="ocr",
|
|
user_prompt="read",
|
|
images=[img],
|
|
)
|
|
|
|
assert response.content == '{"text": "Invoice 123"}'
|
|
assert response.model == "glm-ocr"
|
|
assert response.prompt_tokens == 100
|
|
assert response.completion_tokens == 20
|
|
|
|
|
|
def test_parse_json_extracted_invoice() -> None:
|
|
data = {
|
|
"vendor_name": "Acme",
|
|
"invoice_number": "1",
|
|
"invoice_date": "2024-01-01",
|
|
"line_items": [],
|
|
"subtotal": "100",
|
|
"tax_total": "20",
|
|
"total": "120",
|
|
}
|
|
result = BaseVLMClient.parse_json(json.dumps(data), ExtractedInvoice)
|
|
assert isinstance(result, ExtractedInvoice)
|
|
assert result.vendor_name == "Acme"
|
|
|
|
|
|
def test_encode_image_roundtrip() -> None:
|
|
img = Image.new("RGB", (10, 10), color="blue")
|
|
encoded = OllamaClient._encode_image(img) # type: ignore[attr-defined]
|
|
decoded = Image.open(BytesIO(base64.b64decode(encoded)))
|
|
assert decoded.size == (10, 10)
|