Extract Documents from Databricks with Unsiloed

"""Turn documents in a Unity Catalog volume into a queryable Delta table with Unsiloed.

Create an ETL pipeline in Databricks, point it at this file, and click Run.
New files dropped into the volume are picked up on the next run.
"""
import json
import mimetypes
import time

import requests
from pyspark import pipelines as dp
from pyspark.sql.functions import col, element_at, split, udf

# --- Configure -------------------------------------------------------------
VOLUME = "/Volumes/workspace/unsiloed/docs"
MAX_FILE_BYTES = 50 * 1024 * 1024
MAX_CONCURRENT_EXTRACTIONS = 4

# The fields to pull from each document. Each description is an instruction to
# the model, so be specific.
FIELDS = {
    "vendor_name":    "Company that issued the invoice",
    "invoice_number": "Invoice number or ID",
    "invoice_date":   "Issue date as YYYY-MM-DD",
    "total_amount":   "Grand total payable",
}
# ---------------------------------------------------------------------------

API_KEY = dbutils.secrets.get("unsiloed", "api_key")
BASE = "https://prod.visionapi.unsiloed.ai"
SCHEMA = json.dumps({
    "type": "object",
    "properties": {k: {"type": "string", "description": v} for k, v in FIELDS.items()},
})


@udf("string")
def unsiloed_extract(file_name, content):
    """Send one document to Unsiloed and return its result as JSON text."""
    if content is None:
        return json.dumps({"_error": "The document has no content"})
    if len(content) > MAX_FILE_BYTES:
        return json.dumps({"_error": f"File exceeds the {MAX_FILE_BYTES}-byte limit"})

    try:
        content_type = mimetypes.guess_type(file_name)[0] or "application/octet-stream"
        with requests.Session() as session:
            submit = session.post(
                f"{BASE}/v2/extract",
                headers={"api-key": API_KEY},
                files={"pdf_file": (file_name, bytes(content), content_type)},
                data={"schema_data": SCHEMA, "model": "gamma", "enable_citations": "true"},
                timeout=180,
            )
            if not submit.ok:
                detail = submit.text.replace("\n", " ")[:240]
                return json.dumps({"_error": f"Submit HTTP {submit.status_code}: {detail}"})
            job_id = submit.json()["job_id"]

            deadline = time.monotonic() + 240
            while time.monotonic() < deadline:
                remaining = deadline - time.monotonic()
                poll_response = session.get(
                    f"{BASE}/extract/{job_id}",
                    headers={"api-key": API_KEY},
                    timeout=max(1, min(30, remaining)),
                )
                if poll_response.status_code == 429:
                    time.sleep(min(8, max(0, deadline - time.monotonic())))
                    continue
                poll_response.raise_for_status()
                poll = poll_response.json()
                if poll.get("status") in ("completed", "review"):
                    return json.dumps(poll.get("result") or {})
                if poll.get("status") == "failed":
                    return json.dumps({"_error": f"Job {job_id} failed: {json.dumps(poll)[:240]}"})
                time.sleep(min(4, max(0, deadline - time.monotonic())))
            return json.dumps({"_error": f"Job {job_id} timed out after 240 seconds"})
    except Exception as e:  # one bad document must not fail the pipeline
        return json.dumps({"_error": f"{type(e).__name__}: {e}"})


unsiloed_extract = unsiloed_extract.asNondeterministic()


@dp.table(
    name="extractions",
    comment="One row per document, containing the Unsiloed result object",
)
def extractions():
    # Auto Loader tracks which files it has already seen, so each run only
    # extracts documents that are new since last time.
    return (
        spark.readStream.format("cloudFiles")
        .option("cloudFiles.format", "binaryFile")
        .load(VOLUME)
        .repartition(MAX_CONCURRENT_EXTRACTIONS)
        .withColumn("file_name", element_at(split(col("path"), "/"), -1))
        .withColumn("result_json", unsiloed_extract(col("file_name"), col("content")))
        .selectExpr(
            "replace(path, 'dbfs:', '') AS path",
            "file_name",
            "length AS size_bytes",
            "try_parse_json(result_json) AS result",
            "current_timestamp() AS extracted_at")
    )


@dp.table(
    name="extracted_fields",
    comment="One row per extracted field, so you can triage on confidence",
)
def extracted_fields():
    return spark.sql("""
        SELECT path, file_name, key AS field,
               value:value::string                  AS value,
               value:score.extraction_score::double AS extraction_score,
               value:score.grounding_score::double  AS grounding_score,
               value:citation.page::int             AS citation_page,
               to_json(value:citation.bbox)          AS citation_bbox,
               value:citation.page_width::double     AS citation_page_width,
               value:citation.page_height::double    AS citation_page_height
        FROM STREAM(extractions), LATERAL variant_explode(result)
        WHERE result:_error IS NULL
    """)


@dp.table(
    name="extraction_errors",
    comment="One row per document that Unsiloed could not extract",
)
def extraction_errors():
    return spark.sql("""
        SELECT path, file_name, result:_error::string AS error, extracted_at
        FROM STREAM(extractions)
        WHERE result:_error IS NOT NULL
    """)