docqa/tests/corpus.py

"""Synthetic invoice/receipt corpus with verified ground truth.

Ground truth is written into the PDF itself, so every expected value is known
exactly without human labelling. Layout variety is deliberate: these models are
layout-aware, so a single tidy template would overstate accuracy.

Vendors carry both US (EIN) and Indian (GSTIN) identifiers because the target
extraction keys in the brief include `GSTIN` and `Vendor Id`.
"""

from __future__ import annotations

import json
import random
from dataclasses import asdict, dataclass, field
from pathlib import Path

from reportlab.lib import colors
from reportlab.lib.pagesizes import A4, LETTER
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import mm
from reportlab.platypus import (BaseDocTemplate, Frame, PageTemplate, Paragraph,
                                Spacer, Table, TableStyle)

# --------------------------------------------------------------------------
VENDORS = [
    dict(name="Northwind Logistics LLC",
         address="4820 Harborview Parkway, Suite 300, Seattle, WA 98109",
         country="US", tax_id="88-4122371", gstin="27AAPFU0939F1ZV",
         vat="VAT: US-8841-223"),
    dict(name="Contoso Manufacturing Pvt Ltd",
         address="1180 Industrial Loop, Building 7, Austin, TX 78744",
         country="IN", tax_id="74-2930118", gstin="29AAGCC7391J1Z2",
         vat="GSTIN: 29AAGCC7391J1Z2"),
    dict(name="Fabrikam Consulting Group",
         address="915 Ocean Drive, Floor 12, Miami, FL 33139",
         country="US", tax_id="55-9014423", gstin="07AAACF9876K1ZQ",
         vat="VAT: US-5512-009"),
    dict(name="Tailspin Tooling Co.",
         address="77 Ironworks Road, Detroit, MI 48226",
         country="US", tax_id="38-2210547", gstin="36AABCT1234M1ZP",
         vat="Tax ID: 38-2210547"),
    dict(name="Adventure Works Supply",
         address="3301 Market Street, San Francisco, CA 94114",
         country="US", tax_id="94-3301876", gstin="33AACCA9876T1Z9",
         vat="EIN: 94-3301876"),
    dict(name="Litware Cloud Services",
         address="2200 Telegraph Avenue, Denver, CO 80202",
         country="US", tax_id="91-2200981", gstin="19AAHCL4321K1ZR",
         vat="VAT: US-7719-556"),
    dict(name="Proseware Data Systems",
         address="58 Birchwood Court, Boston, MA 02110",
         country="IN", tax_id="04-7123309", gstin="27AABPD5678L1ZB",
         vat="GSTIN: 27AABPD5678L1ZB"),
    dict(name="Wide World Importers",
         address="9910 Commerce Way, Seattle, WA 98108",
         country="IN", tax_id="55-1102874", gstin="24AACCW8765M1ZN",
         vat="GSTIN: 24AACCW8765M1ZN"),
]

PEOPLE = [
    ("Dana Whitfield", "Accounts Payable Manager"),
    ("Marcus Oyelaran", "Procurement Lead"),
    ("Priya Raghunathan", "Finance Director"),
    ("Tomas Lindqvist", "Operations Director"),
    ("Alicia Moreau", "Controller"),
    ("Ravi Deshpande", "Purchasing Analyst"),
]

ITEMS = [
    ("Warehouse management services", "WMS-ENT", 2450.00),
    ("Pallet racking, steel, 96in", "RCK-96", 189.50),
    ("Forklift rental, monthly", "FLT-MTH", 1150.00),
    ("Customs brokerage service", "CUS-BRK", 210.00),
    ("Data migration consulting", "CONS-MIG", 195.00),
    ("Annual support renewal", "SUP-ANN", 8800.00),
    ("Freight forwarding, LTL", "FRT-LTL", 342.75),
    ("Cold chain packaging supplies", "PKG-CC", 3.20),
    ("Safety compliance audit", "AUD-SFT", 1450.00),
    ("API integration development", "DEV-API", 168.00),
    ("Diesel fuel surcharge", "SUR-DSL", 612.40),
    ("Warehouse safety training", "TRN-SFT", 275.00),
]

TERMS = [
    ("Net 30", "Payment due within 30 days of invoice date."),
    ("Net 45", "Payment due within 45 days of invoice date."),
    ("Net 15", "Payment due within 15 days of invoice date."),
    ("Due on receipt", "Payment due immediately upon receipt."),
    ("Net 60", "Payment due within 60 days of invoice date."),
]

TAXES = [("Sales Tax 8.5%", 0.085), ("VAT 20%", 0.20),
         ("GST 18%", 0.18), ("Sales Tax 6.25%", 0.0625), ("IGST 5%", 0.05)]
CURRENCIES = ["USD", "USD", "USD", "EUR", "INR"]
DATES = ["January 12, 2025", "February 3, 2025", "March 18, 2025",
         "April 7, 2025", "May 22, 2025", "June 30, 2025",
         "August 15, 2025", "September 2, 2025", "October 9, 2025",
         "November 21, 2025", "December 5, 2025"]
DUE_DATES = ["February 11, 2025", "March 20, 2025", "May 5, 2025",
             "June 14, 2025", "August 1, 2025", "September 30, 2025",
             "December 4, 2025", "January 23, 2026"]

SYM = {"USD": "$", "EUR": "\u20ac", "INR": "\u20b9"}


def money(value: float, currency: str = "USD") -> str:
    return f"{SYM.get(currency, '')}{value:,.2f}"


@dataclass
class InvoiceTruth:
    """Every value written onto the rendered invoice."""

    document_id: str
    path: str
    variant: str
    invoice_number: str
    invoice_date: str
    due_date: str
    po_number: str
    payment_terms: str
    currency: str
    vendor_name: str
    vendor_address: str
    vendor_tax_id: str
    gstin: str
    bill_to: str
    subtotal: str
    tax_amount: str
    total: str
    tax_label: str
    page_count: int = 1
    extra: dict = field(default_factory=dict)


class InvoiceBuilder:
    def __init__(self, seed: int, out: Path, variant: str,
                 currency: str = "USD", pages: int = 1):
        self.rng = random.Random(seed)
        self.out = out
        self.variant = variant
        self.currency = currency
        self.pages = pages
        self.vendor = self.rng.choice(VENDORS)
        self.bill_to = self.rng.choice([v for v in VENDORS
                                        if v["name"] != self.vendor["name"]])
        self.contact, self.contact_title = self.rng.choice(PEOPLE)
        self.po = f"PO-{self.rng.randint(10000, 99999)}"
        self.invoice_no = f"INV-{self.rng.randint(100000, 999999)}"
        self.date = self.rng.choice(DATES)
        self.due = self.rng.choice(DUE_DATES)
        self.terms, self.terms_note = self.rng.choice(TERMS)
        self.tax_label, self.tax_rate = self.rng.choice(TAXES)
        self.lines = [(desc, sku, self.rng.randint(1, 20), price)
                      for desc, sku, price in
                      self.rng.sample(ITEMS, self.rng.randint(3, 8))]
        self.subtotal = round(sum(q * p for _, _, q, p in self.lines), 2)
        self.discount_rate = self.rng.choice([0.0, 0.0, 0.0, 0.02, 0.05])
        self.discount = round(self.subtotal * self.discount_rate, 2)
        self.tax = round((self.subtotal - self.discount) * self.tax_rate, 2)
        self.total = round(self.subtotal - self.discount + self.tax, 2)

    # ------------------------------------------------------------------
    def build(self) -> InvoiceTruth:
        s = getSampleStyleSheet()
        small = ParagraphStyle("sm", parent=s["BodyText"], fontSize=8.5, leading=11)
        tiny = ParagraphStyle("ti", parent=s["BodyText"], fontSize=6.5, leading=8)
        h1 = ParagraphStyle("h1", parent=s["Title"], fontSize=19, leading=23,
                            textColor=colors.HexColor("#1a1a2e"))
        h2 = ParagraphStyle("h2", parent=s["Heading2"], fontSize=12, leading=15,
                            textColor=colors.HexColor("#1a1a2e"))

        pagesize = LETTER if self.variant == "letter" else A4
        doc = BaseDocTemplate(
            str(self.out), pagesize=pagesize,
            leftMargin=16 * mm, rightMargin=16 * mm,
            topMargin=14 * mm, bottomMargin=14 * mm,
            title=f"Invoice {self.invoice_no}", author=self.vendor["name"])

        frame = Frame(doc.leftMargin, doc.bottomMargin, doc.width, doc.height,
                      id="f")
        doc.addPageTemplates([PageTemplate(id="main", frames=[frame])])

        story: list = []
        v = self.variant
        if v == "minimal":
            story += self._minimal(s)
        elif v == "boxed":
            story += self._boxed(s, small)
        elif v == "two_column":
            story += self._two_column(s)
        elif v == "dense":
            story += self._dense(tiny)
        else:
            story += self._classic(s, h1, h2, small)

        story += [Spacer(1, 7 * mm), Paragraph(self.terms_note, tiny),
                  Paragraph(
                      f"Remit to: {self.vendor['name']} — "
                      f"{self.vendor['vat']}. Reference {self.po} with all "
                      f"payments. Queries to {self.contact}.", tiny)]

        # extra pages force multi-page handling in the service
        for p in range(1, self.pages):
            story += [Spacer(1, 150 * mm),
                      Paragraph(f"Continued — page {p + 1}", small),
                      Paragraph(
                          f"This page continues invoice {self.invoice_no}. "
                          f"Total amount due remains {money(self.total, self.currency)}.",
                          small)]

        doc.build(story)
        return self.truth()

    # ---------------- layouts ----------------
    def _items_table(self, fs: ParagraphStyle) -> Table:
        rows = [[Paragraph("<b>Description</b>", fs), Paragraph("<b>SKU</b>", fs),
                 Paragraph("<b>Qty</b>", fs),
                 Paragraph("<b>Unit Price</b>", fs),
                 Paragraph("<b>Amount</b>", fs)]]
        for desc, sku, qty, price in self.lines:
            rows.append([Paragraph(desc, fs), Paragraph(sku, fs),
                         Paragraph(str(qty), fs),
                         Paragraph(money(price, self.currency), fs),
                         Paragraph(money(qty * price, self.currency), fs)])
        t = Table(rows, colWidths=[76 * mm, 26 * mm, 14 * mm, 30 * mm, 32 * mm])
        t.setStyle(TableStyle([
            ("GRID", (0, 0), (-1, -1), 0.4, colors.HexColor("#999999")),
            ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#eeeeee")),
            ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
        ]))
        return t

    def _totals(self, fs: ParagraphStyle) -> Table:
        rows = [["Subtotal", money(self.subtotal, self.currency)]]
        if self.discount:
            rows.append([f"Discount ({self.discount_rate*100:.0f}%)",
                         f"-{money(self.discount, self.currency)}"])
        rows += [[self.tax_label, money(self.tax, self.currency)],
                 ["Total Amount Due", money(self.total, self.currency)]]
        t = Table(rows, colWidths=[60 * mm, 42 * mm], hAlign="RIGHT")
        t.setStyle(TableStyle([
            ("FONTNAME", (0, -1), (-1, -1), "Helvetica-Bold"),
            ("LINEABOVE", (0, -1), (-1, -1), 0.8, colors.black),
            ("FONTSIZE", (0, 0), (-1, -1), 9.5),
            ("ALIGN", (1, 0), (1, -1), "RIGHT"),
        ]))
        return t

    def _identity(self, fs: ParagraphStyle) -> str:
        return (f"<b>{self.vendor['name']}</b><br/>{self.vendor['address']}"
                f"<br/>{self.vendor['vat']}<br/>GSTIN: {self.vendor['gstin']}")

    def _classic(self, s, h1, h2, small):
        left = Paragraph(self._identity(small), small)
        right = Paragraph(
            f"<b>INVOICE</b><br/>Invoice Number: {self.invoice_no}<br/>"
            f"Invoice Date: {self.date}<br/>Due Date: {self.due}", small)
        head = Table([[left, right]], colWidths=[89 * mm, 89 * mm])
        head.setStyle(TableStyle([("VALIGN", (0, 0), (-1, -1), "TOP")]))

        meta = Table([
            ["Bill To:", f"{self.bill_to['name']}<br/>{self.bill_to['address']}"],
            ["Attention:", f"{self.contact}, {self.contact_title}"],
            ["PO Number:", self.po],
            ["Terms:", self.terms],
            ["Currency:", self.currency],
            ["Vendor Id:", self.vendor["tax_id"]],
        ], colWidths=[32 * mm, 120 * mm])
        meta.setStyle(TableStyle([
            ("FONTNAME", (0, 0), (0, -1), "Helvetica-Bold"),
            ("FONTSIZE", (0, 0), (-1, -1), 9.5),
            ("VALIGN", (0, 0), (-1, -1), "TOP"),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
        ]))
        return [head, Spacer(1, 5 * mm), meta, Spacer(1, 5 * mm),
                Paragraph("Line Items", h2), Spacer(1, 2 * mm),
                self._items_table(small), Spacer(1, 4 * mm),
                self._totals(small)]

    def _minimal(self, s):
        small = s["BodyText"]
        info = Table([
            [Paragraph(f"<b>Invoice No.</b> {self.invoice_no}", small),
             Paragraph(f"<b>Date</b> {self.date}", small),
             Paragraph(f"<b>Due</b> {self.due}", small)],
            [Paragraph(f"<b>Bill to</b> {self.bill_to['name']}", small),
             Paragraph(f"<b>PO</b> {self.po}", small),
             Paragraph(f"<b>Terms</b> {self.terms}", small)],
            [Paragraph(f"<b>Vendor</b> {self.vendor['name']}", small),
             Paragraph(f"<b>Vendor Id</b> {self.vendor['tax_id']}", small),
             Paragraph(f"<b>GSTIN</b> {self.vendor['gstin']}", small)],
        ], colWidths=[62 * mm, 58 * mm, 58 * mm])
        info.setStyle(TableStyle([("TOPPADDING", (0, 0), (-1, -1), 4),
                                  ("BOTTOMPADDING", (0, 0), (-1, -1), 4)]))
        return [Spacer(1, 8 * mm), Paragraph(self.vendor["name"], h1_small(s)),
                Paragraph(self.vendor["address"], small), Spacer(1, 6 * mm),
                info, Spacer(1, 5 * mm), self._items_table(small),
                Spacer(1, 4 * mm), self._totals(small)]

    def _boxed(self, s, small):
        head = Table([[Paragraph(self._identity(small), small),
                       Paragraph(f"<b>INVOICE {self.invoice_no}</b><br/>"
                                 f"Date: {self.date}<br/>Due: {self.due}", small)]],
                     colWidths=[89 * mm, 89 * mm])
        head.setStyle(TableStyle([("VALIGN", (0, 0), (-1, -1), "TOP")]))
        box = Table([[Paragraph(
            f"<b>SHIP TO / BILL TO</b><br/>{self.bill_to['name']}<br/>"
            f"{self.bill_to['address']}<br/>Attn: {self.contact}, "
            f"{self.contact_title}<br/>PO: {self.po}<br/>"
            f"Currency: {self.currency}", small)]], colWidths=[90 * mm])
        box.setStyle(TableStyle([
            ("BOX", (0, 0), (-1, -1), 1.2, colors.HexColor("#1a1a2e")),
            ("BACKGROUND", (0, 0), (-1, -1), colors.HexColor("#f4f4f8")),
            ("LEFTPADDING", (0, 0), (-1, -1), 8),
            ("TOPPADDING", (0, 0), (-1, -1), 6),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 6),
        ]))
        return [head, Spacer(1, 5 * mm), box, Spacer(1, 5 * mm),
                self._items_table(small), Spacer(1, 4 * mm),
                self._totals(small)]

    def _two_column(self, s):
        small = s["BodyText"]
        left = [Paragraph(f"<b>{self.vendor['name']}</b>", small),
                Paragraph(self.vendor["address"], small),
                Paragraph(self.vendor["vat"], small),
                Paragraph(f"GSTIN: {self.vendor['gstin']}", small),
                Paragraph(f"Vendor Id: {self.vendor['tax_id']}", small),
                Spacer(1, 4 * mm),
                Paragraph(f"<b>Bill To</b><br/>{self.bill_to['name']}<br/>"
                          f"{self.bill_to['address']}", small),
                Paragraph(f"Attn: {self.contact}", small)]
        right = [Paragraph(f"<b>INVOICE {self.invoice_no}</b>", small),
                 Paragraph(f"Date: {self.date}<br/>Due: {self.due}<br/>"
                           f"PO: {self.po}<br/>Terms: {self.terms}<br/>"
                           f"Currency: {self.currency}", small)]
        t = Table([[left, right]], colWidths=[89 * mm, 89 * mm])
        t.setStyle(TableStyle([("VALIGN", (0, 0), (-1, -1), "TOP")]))
        return [t, Spacer(1, 5 * mm), self._items_table(small),
                Spacer(1, 4 * mm), self._totals(small)]

    def _dense(self, tiny):
        rows = []
        for i in range(26):
            desc, sku, price = ITEMS[i % len(ITEMS)]
            qty = (i % 9) + 1
            rows.append([Paragraph(desc, tiny), Paragraph(sku, tiny),
                         Paragraph(str(qty), tiny),
                         Paragraph(money(price, self.currency), tiny),
                         Paragraph(money(qty * price, self.currency), tiny)])
        rows.insert(0, [Paragraph(
            f"<b>{self.vendor['name']}</b> &mdash; INVOICE {self.invoice_no}"
            f" &mdash; Date {self.date} &mdash; Due {self.due}"
            f" &mdash; PO {self.po} &mdash; Terms {self.terms}"
            f" &mdash; Currency {self.currency} &mdash; Vendor Id "
            f"{self.vendor['tax_id']} &mdash; GSTIN {self.vendor['gstin']}",
            tiny)])
        t = Table(rows, colWidths=[84 * mm, 22 * mm, 14 * mm, 28 * mm, 30 * mm])
        t.setStyle(TableStyle([
            ("GRID", (0, 1), (-1, -1), 0.3, colors.HexColor("#aaaaaa")),
            ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#e8e8f0")),
        ]))
        return [t, Spacer(1, 3 * mm), self._totals(tiny),
                Paragraph(f"Bill to: {self.bill_to['name']}", tiny),
                Paragraph(f"Attention: {self.contact}", tiny)]

    def truth(self) -> InvoiceTruth:
        return InvoiceTruth(
            document_id=self.out.stem, path=str(self.out.resolve()),
            variant=self.variant, invoice_number=self.invoice_no,
            invoice_date=self.date, due_date=self.due, po_number=self.po,
            payment_terms=self.terms, currency=self.currency,
            vendor_name=self.vendor["name"], vendor_address=self.vendor["address"],
            vendor_tax_id=self.vendor["tax_id"], gstin=self.vendor["gstin"],
            bill_to=self.bill_to["name"],
            subtotal=money(self.subtotal, self.currency),
            tax_amount=money(self.tax, self.currency),
            total=money(self.total, self.currency),
            tax_label=self.tax_label, page_count=self.pages,
        )


def h1_small(s):
    st = ParagraphStyle("h1s", parent=s["Title"], fontSize=17, leading=21,
                        alignment=1, textColor=colors.HexColor("#1a1a2e"))
    return st


VARIANTS = ["classic", "minimal", "boxed", "two_column", "dense", "letter"]


def generate(out_dir: Path, n: int = 12, seed: int = 7) -> list[InvoiceTruth]:
    out_dir.mkdir(parents=True, exist_ok=True)
    truths: list[InvoiceTruth] = []
    for i in range(n):
        variant = VARIANTS[i % len(VARIANTS)]
        currency = CURRENCIES[i % len(CURRENCIES)] if variant != "dense" else "USD"
        pages = 2 if i % 5 == 4 else 1
        builder = InvoiceBuilder(seed=seed * 1000 + i,
                                 out=out_dir / f"invoice_{i:02d}_{variant}.pdf",
                                 variant=variant, currency=currency, pages=pages)
        truths.append(builder.build())
    manifest = out_dir / "ground_truth.json"
    manifest.write_text(json.dumps([asdict(t) for t in truths], indent=2),
                        encoding="utf-8")
    return truths


if __name__ == "__main__":
    import sys
    target = Path(sys.argv[1] if len(sys.argv) > 1 else "corpus")
    count = int(sys.argv[2]) if len(sys.argv) > 2 else 12
    rows = generate(target, count)
    print(f"generated {len(rows)} invoices -> {target}")