async_batch.py

"""A folder of long PDFs, submitted to the queue and collected as they finish.

Two things make this different from a loop of analyze() calls:

1. Nothing is held open for 60 seconds. A 50-page contract can take minutes, and a
   synchronous request would be killed at 60 s with a 504 sync_timeout.
2. A stable idempotency_key per file means re-running the script after a crash replays the
   stored responses instead of paying for the batch twice.

    python examples/async_batch.py ./contracts 4
"""

import datetime as dt
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from visionapi import TaskFailedError, VisionAPI, unwrap

folder = Path(sys.argv[1] if len(sys.argv) > 1 else "./contracts")
workers = int(sys.argv[2]) if len(sys.argv) > 2 else 4

vision = VisionAPI()
BATCH = dt.date.today().isoformat()  # one idempotency namespace per day


def process(path: Path) -> dict:
    ref = vision.analyze_async(
        file=path,
        preset="contract",
        pages="1-50",
        # Derived from the file, not random: a re-run of the same batch must not re-charge.
        idempotency_key=f"{BATCH}:{path.name}",
    )
    try:
        task = vision.wait_for_task(ref["task_id"], poll_interval=2.0, max_wait=20 * 60)
    except TaskFailedError as err:
        # A failed task costs 0 credits — the reservation is released in full.
        return {"path": path, "error": err.code}
    return {
        "path": path,
        "credits": task.get("credits_used", 0),
        "data": unwrap(task.get("result"), drop_null=True),
    }


# The rate limit is 60 requests/minute per key, and polling counts. Four concurrent files
# at one poll every 2 s is comfortably inside it.
with ThreadPoolExecutor(max_workers=workers) as pool:
    results = list(pool.map(process, sorted(folder.glob("*.pdf"))))

for r in results:
    if "error" in r:
        print(f"✗ {r['path'].name}: {r['error']}")
    else:
        print(f"✓ {r['path'].name}: {r['credits']} credits")

print(f"\n{len(results)} files, {sum(r.get('credits', 0) for r in results)} credits")