Initial scaffold: local invoice OCR pipeline with Ollama, classifier branches, structured extraction, review, and Odoo XML export
- 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.
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user