Get data you can trust from an API: schemas, validation and retries

import datetime, json, re
import anthropic
from jsonschema import Draft202012Validator

MODEL = "claude-haiku-5-5"
NULLABLE_STR = {"type": ["string", "null"]}
SCHEMA = {
    "type": "object",
    "properties": {
        "invoice_number": NULLABLE_STR,
        "invoice_number_quote": NULLABLE_STR,
        "invoice_date": {"type": ["string", "null"], "description": "YYYY-MM-DD"},
        "invoice_date_quote": NULLABLE_STR,
        "currency": {"type": "string", "enum": ["USD", "EUR", "GBP", "OTHER"]},
        "total": {"type": ["number", "null"]},
        "total_quote": NULLABLE_STR,
        "line_items": {"type": "array", "items": {
            "type": "object",
            "properties": {"description": {"type": "string"}, "amount": {"type": "number"}},
            "required": ["description", "amount"],
            "additionalProperties": False}},
    },
    "required": ["invoice_number", "invoice_number_quote", "invoice_date",
                 "invoice_date_quote", "currency", "total", "total_quote", "line_items"],
    "additionalProperties": False,
}
# Rules we check here ourselves (minimum is on the API's unsupported list; we check the date pattern locally too).
LOCAL_SCHEMA = json.loads(json.dumps(SCHEMA))
LOCAL_SCHEMA["properties"]["total"]["minimum"] = 0
LOCAL_SCHEMA["properties"]["invoice_date"]["pattern"] = r"^\d{4}-\d{2}-\d{2}$"
VALIDATOR = Draft202012Validator(LOCAL_SCHEMA)

PROMPT = ("Extract the invoice fields from the document below. Use null when a value "
          "is not in the document; never guess. For each *_quote field, copy the exact "
          "text from the document that shows the value.\n\n<document>\n{doc}\n</document>")


def squash(text):
    return re.sub(r"\s+", " ", text).strip()


def check(document, raw):
    """Return (data, errors). An empty error list means every check passed."""
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        return None, [f"not valid JSON: {e}"]
    if isinstance(data, dict) and isinstance(data.get("currency"), str):
        data["currency"] = data["currency"].upper()  # enum case is not guaranteed
    errors = [f"{'/'.join(map(str, e.path)) or 'root'}: {e.message}"
              for e in VALIDATOR.iter_errors(data)]
    if errors:
        return data, errors
    doc = squash(document)
    for field in ("invoice_number", "invoice_date", "total"):
        value, quote = data[field], data[field + "_quote"]
        if value is not None and not quote:
            errors.append(f"{field}: value given without a quote")
        if quote and squash(quote) not in doc:
            errors.append(f"{field}_quote: not found word for word in the document")
    if data["invoice_number"] and data["invoice_number_quote"] \
            and data["invoice_number"] not in data["invoice_number_quote"]:
        errors.append("invoice_number: not in its own quote")
    if data["total"] is not None and data["total_quote"]:
        numbers = [float(n.replace(",", "")) for n in re.findall(r"\d[\d,]*\.?\d*", data["total_quote"])]
        if data["total"] not in numbers:
            errors.append("total: not in its own quote")
    if data["total"] is not None and data["line_items"]:
        items = round(sum(i["amount"] for i in data["line_items"]), 2)
        if abs(items - data["total"]) > 0.01:
            errors.append(f"total: {data['total']} but line items add up to {items}")
    if data["invoice_date"]:
        try:
            day = datetime.date.fromisoformat(data["invoice_date"])
            if not datetime.date(2020, 1, 1) <= day <= datetime.date.today():
                errors.append(f"invoice_date: {day} is outside the expected range")
        except ValueError:
            errors.append(f"invoice_date: {data['invoice_date']} is not a real date")
    return data, errors


def ask(client, document, feedback=None):
    messages = [{"role": "user", "content": PROMPT.format(doc=document)}]
    if feedback:
        messages += [{"role": "assistant", "content": feedback["raw"]},
                     {"role": "user", "content": "Your answer failed these checks:\n- "
                      + "\n- ".join(feedback["errors"]) + "\nReturn a corrected answer."}]
    response = client.messages.create(
        model=MODEL, max_tokens=1024, messages=messages,
        output_config={"format": {"type": "json_schema", "schema": SCHEMA}})
    if response.stop_reason in ("refusal", "max_tokens"):
        return None, [f"stop_reason: {response.stop_reason}"]
    return next(b.text for b in response.content if b.type == "text"), []


def extract(document, doc_id, client=None, log_path="extract_failures.jsonl"):
    client = client or anthropic.Anthropic()  # reads ANTHROPIC_API_KEY
    feedback = None
    for attempt in (1, 2):  # one try, one retry, never a loop
        raw, errors = ask(client, document, feedback)
        data, errors = check(document, raw) if raw else (None, errors)
        if not errors:
            return {"status": "ok", "data": data, "attempts": attempt}
        with open(log_path, "a") as log:
            log.write(json.dumps({"doc_id": doc_id, "attempt": attempt, "model": MODEL,
                                  "at": datetime.datetime.now().isoformat(),
                                  "errors": errors, "raw": raw}) + "\n")
        if raw is None:
            break  # refusal or cut off: feedback will not fix it (for a cut-off, raise max_tokens)
        feedback = {"raw": raw, "errors": errors}
    return {"status": "needs_person", "errors": errors}