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,54 @@
|
||||
"""Shared pytest fixtures."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from odoo_ocr.clients.base import BaseVLMClient
|
||||
from odoo_ocr.config import Settings
|
||||
from odoo_ocr.schemas import VLMResponse
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings() -> Settings:
|
||||
"""Default test settings with cache disabled."""
|
||||
return Settings(
|
||||
ollama_base_url="http://test-ollama:11435",
|
||||
cache=Settings().model_dump()["cache"] | {"enabled": False},
|
||||
)
|
||||
|
||||
|
||||
class FakeVLMClient(BaseVLMClient):
|
||||
"""Test double that returns canned responses based on prompt content."""
|
||||
|
||||
def __init__(self, responses: dict[str, str] | None = None) -> None:
|
||||
super().__init__(Settings())
|
||||
self.responses = responses or {}
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
system_prompt: str,
|
||||
user_prompt: str,
|
||||
images: list[Image.Image],
|
||||
response_format: type[Any] | None = None,
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 4096,
|
||||
model: str | None = None,
|
||||
) -> VLMResponse:
|
||||
self.calls.append(
|
||||
{
|
||||
"system_prompt": system_prompt,
|
||||
"user_prompt": user_prompt,
|
||||
"image_count": len(images),
|
||||
"model": model,
|
||||
}
|
||||
)
|
||||
for key, value in self.responses.items():
|
||||
if key in system_prompt or key in user_prompt:
|
||||
return VLMResponse(content=json.dumps(value), model="fake")
|
||||
return VLMResponse(content=json.dumps({}), model="fake")
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Tests for the document classifier."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from odoo_ocr.pipeline.classifier import classify_document
|
||||
from odoo_ocr.schemas.document import DocumentClass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def blank_image() -> Image.Image:
|
||||
return Image.new("RGB", (100, 100), color="white")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heuristic_digital_pdf(tmp_path: Path) -> None:
|
||||
pdf_path = tmp_path / "invoice.pdf"
|
||||
# Minimal valid PDF with text via pymupdf
|
||||
import fitz
|
||||
|
||||
doc = fitz.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((50, 50), "Invoice number 12345\nVendor: Acme")
|
||||
doc.save(str(pdf_path))
|
||||
doc.close()
|
||||
|
||||
from odoo_ocr.config import Settings
|
||||
from tests.conftest import FakeVLMClient
|
||||
|
||||
client = FakeVLMClient()
|
||||
result = await classify_document(pdf_path, [], client, Settings())
|
||||
assert result.category == DocumentClass.digital_pdf
|
||||
assert result.confidence >= 0.90
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heuristic_scanned_print(blank_image: Image.Image) -> None:
|
||||
# Blank image has low handwriting score
|
||||
from odoo_ocr.config import Settings
|
||||
from tests.conftest import FakeVLMClient
|
||||
|
||||
client = FakeVLMClient()
|
||||
result = await classify_document(Path("scan.jpg"), [blank_image], client, Settings())
|
||||
assert result.category == DocumentClass.scanned_print
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vlm_classifier_fallback(blank_image: Image.Image) -> None:
|
||||
|
||||
from odoo_ocr.config import Settings
|
||||
from tests.conftest import FakeVLMClient
|
||||
|
||||
client = FakeVLMClient(
|
||||
responses={
|
||||
"document classifier": {
|
||||
"category": "handwritten",
|
||||
"confidence": 0.91,
|
||||
"reasoning": "handwritten text visible",
|
||||
}
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"odoo_ocr.pipeline.classifier._heuristic_classify", return_value=None
|
||||
):
|
||||
result = await classify_document(Path("unknown.png"), [blank_image], client, Settings())
|
||||
assert result.category == DocumentClass.handwritten
|
||||
assert result.confidence == 0.91
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Tests for the structured extraction stage."""
|
||||
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from odoo_ocr.config import Settings
|
||||
from odoo_ocr.pipeline.extractor import extract_invoice
|
||||
from odoo_ocr.schemas import ExtractedInvoice
|
||||
from tests.conftest import FakeVLMClient
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_invoice() -> None:
|
||||
client = FakeVLMClient(
|
||||
responses={
|
||||
"invoice data extraction": {
|
||||
"vendor_name": "Acme",
|
||||
"vendor_address": None,
|
||||
"vendor_vat": None,
|
||||
"invoice_number": "INV-1",
|
||||
"invoice_date": "2024-01-01",
|
||||
"due_date": None,
|
||||
"currency": "EUR",
|
||||
"payment_terms": None,
|
||||
"line_items": [
|
||||
{
|
||||
"description": "Widget",
|
||||
"quantity": 2.0,
|
||||
"unit_price": 50.0,
|
||||
"total_price": 100.0,
|
||||
"tax_rate": 0.0,
|
||||
}
|
||||
],
|
||||
"subtotal": 100.0,
|
||||
"tax_total": 0.0,
|
||||
"total": 100.0,
|
||||
"iban": None,
|
||||
"raw_ocr_text": "",
|
||||
}
|
||||
}
|
||||
)
|
||||
result = await extract_invoice("raw text", [Image.new("RGB", (10, 10))], client, Settings())
|
||||
assert isinstance(result, ExtractedInvoice)
|
||||
assert result.vendor_name == "Acme"
|
||||
assert result.line_items_sum() == 100
|
||||
@@ -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)
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Tests for the Odoo XML builder."""
|
||||
|
||||
from decimal import Decimal
|
||||
|
||||
from odoo_ocr.config import Settings
|
||||
from odoo_ocr.pipeline.xml_builder import build_odoo_xml
|
||||
from odoo_ocr.schemas import ExtractedInvoice, InvoiceLineItem, ReviewResult
|
||||
|
||||
|
||||
def test_build_odoo_xml() -> None:
|
||||
settings = Settings()
|
||||
invoice = ExtractedInvoice(
|
||||
vendor_name="Acme Supplies",
|
||||
vendor_vat="GB123456789",
|
||||
invoice_number="INV-2024-001",
|
||||
invoice_date="2024-05-01",
|
||||
due_date="2024-06-01",
|
||||
currency="EUR",
|
||||
payment_terms="Net 30",
|
||||
line_items=[
|
||||
InvoiceLineItem(
|
||||
description="Consulting services",
|
||||
quantity=Decimal("10.0"),
|
||||
unit_price=Decimal("100.00"),
|
||||
total_price=Decimal("1000.00"),
|
||||
tax_rate=Decimal("20"),
|
||||
)
|
||||
],
|
||||
subtotal=Decimal("1000.00"),
|
||||
tax_total=Decimal("200.00"),
|
||||
total=Decimal("1200.00"),
|
||||
)
|
||||
review = ReviewResult(valid=True, confidence=0.95)
|
||||
|
||||
xml = build_odoo_xml(invoice, review, settings)
|
||||
|
||||
assert xml.startswith("<?xml version=")
|
||||
assert 'model="res.partner"' in xml
|
||||
assert 'model="account.move"' in xml
|
||||
assert 'model="account.move.line"' in xml
|
||||
assert "Acme Supplies" in xml
|
||||
assert "INV-2024-001" in xml
|
||||
assert "in_invoice" in xml
|
||||
|
||||
|
||||
def test_xml_adds_review_comment_for_low_confidence() -> None:
|
||||
settings = Settings()
|
||||
invoice = ExtractedInvoice(
|
||||
vendor_name="X",
|
||||
invoice_number="1",
|
||||
invoice_date="2024-01-01",
|
||||
line_items=[],
|
||||
subtotal=Decimal("0"),
|
||||
tax_total=Decimal("0"),
|
||||
total=Decimal("0"),
|
||||
)
|
||||
review = ReviewResult(valid=False, confidence=0.80)
|
||||
|
||||
xml = build_odoo_xml(invoice, review, settings)
|
||||
assert "HUMAN REVIEW ADVISED" in xml
|
||||
Reference in New Issue
Block a user