PaperPod/paperpod/llm/vision.py

"""Identify documents with a local vision LLM served by Ollama.

Everything stays on-machine: images go to localhost:11434, never the cloud.
A vision model (default: minicpm-v) reads the whole document at once, which
is far more robust than Tesseract + regex for curled receipts, hands in
frame, and stylized logos — it correctly reads "COSTCO WHOLESALE" where
Tesseract produced "Cosrco". It also reports quality issues (fingers, blur,
glare) so bad captures can be flagged for re-scanning.

minicpm-v was chosen over qwen2.5vl after A/B on overhead statement scans:
~2x faster and far fewer false "blank" labels on dense mono bank pages.
Dates can still hallucinate — the pipeline keeps OCR as date/total fill-in.

Trade-off: ~8-12s per document on an M-series Mac, vs <2s for Tesseract.
The pipeline uses the LLM for naming when available and falls back to the
OCR heuristics when Ollama isn't running.
"""

from __future__ import annotations

import base64
import json
import re
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from datetime import datetime

import cv2
import numpy as np

_PROMPT = """Look at this image of a document. Reply with ONLY a JSON object, no other text:
{"doc_type": "receipt|tax_form|letter|invoice|blank|other",
 "vendor": "business or organization name, or null",
 "party": "account holder / customer / SHORTNAME / payee printed on the document, or null",
 "form_code": "tax form code like T4/T4A/T5008 if this is a tax form, else null",
 "date": "document/transaction date as YYYY-MM-DD, or null",
 "time": "transaction time as 24-hour HH:MM if printed on the document, else null",
 "tax_year": "tax year like 2023 if this is a tax form, else null",
 "total": "total amount like 103.60, or null",
 "description": "3-6 word summary",
 "issues": ["any of: fingers_visible, blurry, cut_off, upside_down, glare, partial_document"]}
Use doc_type "blank" if the page has no meaningful printed content (blank back of a page, empty sheet).
Set vendor to null unless you can actually read a real business/organization name — never output garbled text.
Set party to the person/entity the document is about (e.g. SHORTNAME on a bank statement, payee on a slip) when clearly printed; otherwise null."""


@dataclass
class LlmAnalysis:
    doc_type: str
    vendor: str | None
    party: str | None
    form_code: str | None
    date: str | None  # normalized YYYY-MM-DD
    time: str | None  # normalized HH:MM (24h)
    tax_year: str | None
    total: str | None
    description: str
    issues: list[str] = field(default_factory=list)


def ollama_available(base_url: str = "http://localhost:11434", timeout: float = 2.0) -> bool:
    try:
        with urllib.request.urlopen(f"{base_url}/api/tags", timeout=timeout):
            return True
    except (urllib.error.URLError, OSError):
        return False


def _encode_image(image_bgr: np.ndarray, max_width: int = 1200) -> str:
    """JPEG-encode (downscaled) for the API; PNG of a 4K crop is wastefully big."""
    h, w = image_bgr.shape[:2]
    if w > max_width:
        image_bgr = cv2.resize(image_bgr, (max_width, round(h * max_width / w)))
    ok, buf = cv2.imencode(".jpg", image_bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
    if not ok:
        raise ValueError("Could not JPEG-encode image")
    return base64.b64encode(buf.tobytes()).decode()


def _normalize_date(raw: str | None) -> str | None:
    if not raw:
        return None
    match = re.search(r"(\d{4})[-/](\d{1,2})[-/](\d{1,2})", str(raw))
    if not match:
        return None
    y, m, d = (int(g) for g in match.groups())
    if not (1 <= m <= 12 and 1 <= d <= 31 and 1990 <= y <= datetime.now().year + 1):
        return None
    return f"{y:04d}-{m:02d}-{d:02d}"


def _clean_str(value) -> str | None:
    if value is None:
        return None
    text = str(value).strip()
    return text if text and text.lower() not in ("null", "none", "n/a", "unknown") else None


def _normalize_time(raw: str | None) -> str | None:
    if not raw:
        return None
    match = re.search(r"([01]?\d|2[0-3]):([0-5]\d)", str(raw))
    if not match:
        return None
    return f"{int(match.group(1)):02d}:{match.group(2)}"


def analyze_document(
    image_bgr: np.ndarray,
    model: str = "minicpm-v",
    base_url: str = "http://localhost:11434",
    timeout: float = 180.0,
) -> LlmAnalysis | None:
    """Ask the local vision model what this document is. None on any failure."""
    try:
        body = json.dumps(
            {
                "model": model,
                "prompt": _PROMPT,
                "images": [_encode_image(image_bgr)],
                "stream": False,
                "format": "json",
                "options": {"temperature": 0},
            }
        ).encode()
        req = urllib.request.Request(
            f"{base_url}/api/generate", data=body, headers={"Content-Type": "application/json"}
        )
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            payload = json.load(resp)
        data = json.loads(payload["response"])
    except (urllib.error.URLError, OSError, KeyError, ValueError):
        return None

    issues = data.get("issues") or []
    if not isinstance(issues, list):
        issues = [str(issues)]

    tax_year = _clean_str(data.get("tax_year"))
    if tax_year:
        year_match = re.search(r"(19|20)\d{2}", tax_year)
        tax_year = year_match.group(0) if year_match else None

    return LlmAnalysis(
        doc_type=_clean_str(data.get("doc_type")) or "other",
        vendor=_clean_str(data.get("vendor")),
        party=_clean_str(data.get("party")),
        form_code=_clean_str(data.get("form_code")),
        date=_normalize_date(data.get("date")),
        time=_normalize_time(data.get("time")),
        tax_year=tax_year,
        total=_clean_str(data.get("total")),
        description=_clean_str(data.get("description")) or "",
        issues=[str(i) for i in issues],
    )