Mistral OCR

MISTRAL_API_KEY=your-key-here

-- models/raw/invoices.sql
-- @kind: tracked
-- @group_key: source_file

SELECT
    source_file,
    source_hash,
    invoice_number,
    vendor,
    total::DECIMAL AS total,
    date
FROM mistral_ocr('data/invoices/*.pdf')

ondatrasql run raw.invoices

-- models/raw/invoices.sql
-- @kind: tracked
-- @group_key: source_file

SELECT
    source_file,
    source_hash,
    invoice_number,
    vendor,
    total::DECIMAL AS total,
    date,
    line_items::JSON AS line_items
FROM mistral_ocr('data/invoices/*.pdf')

-- models/staging/invoice_lines.sql
-- @kind: table

SELECT
    invoice_number,
    vendor,
    date,
    j.value->>'item' AS item,
    (j.value->>'quantity')::INTEGER AS qty,
    (j.value->>'price')::DECIMAL AS price
FROM raw.invoices,
LATERAL json_each(line_items) AS j

-- models/mart/invoice_report.sql
-- @kind: table

SELECT
    r.invoice_number,
    r.vendor,
    r.total AS invoice_total,
    r.date,
    l.item,
    l.qty,
    l.price,
    l.qty * l.price AS line_total
FROM raw.invoices r
JOIN staging.invoice_lines l ON r.invoice_number = l.invoice_number

SELECT invoice_number, vendor, date, x.*
FROM raw.invoices,
LATERAL (SELECT unnest(json_transform(line_items,
    '[{"item":"VARCHAR","quantity":"INTEGER","price":"DOUBLE"}]')) AS x)

SELECT
    invoice_number,
    line_items->>0->>'item' AS first_item,
    (line_items->>0->>'price')::DECIMAL AS first_price
FROM raw.invoices

-- source_hash is required for incremental skip logic
SELECT source_file, source_hash,
    store, date, total::DECIMAL AS total, payment_method
FROM mistral_ocr('data/receipts/*.pdf')

-- source_hash is required for incremental skip logic
SELECT source_file, source_hash,
    parties, effective_date, termination_date,
    value::DECIMAL AS value
FROM mistral_ocr('data/contracts/*.pdf')

RESERVED_COLUMNS = ("source_file", "source_hash")

API = {
    "base_url": "https://api.mistral.ai",
    "auth": {"env": "MISTRAL_API_KEY"},
    "timeout": 120,
    "retry": 2,
    "fetch": {
        "args": ["pattern"],
        "page_size": 10,
        "supported_kinds": ["tracked"],
    },
}

def fetch(pattern, page, columns=[], target="", is_backfill=True):
    files = glob(pattern)
    schema = _build_schema(columns)

    # Index-based pagination. Assumes the file set is stable for the
    # duration of the pipeline run.
    start = int(page.cursor) if page.cursor != None else 0
    batch = files[start:start + page.size]
    known = _load_known_hashes(target, batch, is_backfill)
    rows = []

    for filepath in batch:
        content_hash = md5_file(filepath)
        if known.get(filepath, None) == content_hash:
            print("skip (unchanged): " + filepath)
            continue

        print("ocr: " + filepath)
        file_id = _upload(filepath)
        resp = http.post("/v1/ocr", json={
            "model": "mistral-ocr-latest",
            "document": {"type": "file", "file_id": file_id},
            "document_annotation_format": {
                "type": "json_schema",
                "json_schema": {"name": "extraction", "schema": schema},
            },
        })
        _cleanup(file_id)

        if not resp.ok:
            fail("ocr " + filepath + ": " + resp.text)

        # Mistral returns annotation as JSON-encoded string inside JSON.
        # Sometimes wraps response in a JSON Schema envelope — unwrap if detected.
        annotation = json.decode(resp.json["document_annotation"])
        if "properties" in annotation and "type" in annotation:
            annotation = annotation["properties"]

        row = {}
        for col in columns:
            row[col["name"]] = annotation.get(col["name"], None)
        row["source_file"] = filepath
        row["source_hash"] = content_hash
        rows.append(row)

    next_start = start + page.size
    next_cursor = next_start if next_start < len(files) else None

    return {"rows": rows, "next": next_cursor}


def _build_schema(columns):
    """Build JSON Schema from typed columns using json_schema_type from runtime."""
    properties = {}
    for col in columns:
        name = col["name"]
        if name in RESERVED_COLUMNS:
            continue
        properties[name] = {"type": col.get("json_schema_type", "string")}
    return {"type": "object", "properties": properties}


def _load_known_hashes(target, batch, is_backfill):
    if is_backfill or not target or not batch:
        return {}
    return lookup(table=target, key="source_file", value="source_hash", where=batch)


def _upload(filepath):
    resp = http.upload("/v1/files", file=filepath, fields={"purpose": "ocr"})
    if not resp.ok:
        fail("upload " + filepath + ": " + resp.text)
    return resp.json["id"]


def _cleanup(file_id):
    """Best-effort cleanup — silent on failure."""
    http.delete("/v1/files/" + file_id)