Receipt & Invoice Extractor

import json
import re

# How many extracted rows are shown to the model when answering questions.
MAX_ROWS_IN_PROMPT = [[[MAX_ROWS_IN_PROMPT|60]]]
# Hard cap on the facts sheet handed to the model (small local context windows).
MAX_SHEET_CHARS = [[[MAX_SHEET_CHARS|6000]]]
# How much receipt text the vendor/category classifier sees.
CLASSIFY_CHARS = [[[CLASSIFY_CHARS|700]]]

CATEGORIES = [
    "groceries", "restaurant", "fuel", "travel", "lodging", "software",
    "office", "utilities", "health", "other",
]

# One dict per extracted receipt (see _scan_receipt for the shape).
_RECEIPTS = []

_CURRENCY_SYMBOLS = {"$": "USD", "\u20ac": "EUR", "\u00a3": "GBP",
                     "\u20ba": "TRY", "\u00a5": "JPY"}
_CURRENCY_CODE_RE = re.compile(
    r"\b(USD|EUR|GBP|TRY|CHF|JPY|CAD|AUD|SEK|NOK|DKK|PLN)\b"
)
# Money, and only money: 1.234,56 / 1,234.56 / 249.00 / $254.
#
# A bare integer is deliberately NOT an amount unless a currency symbol is
# attached to it. Receipts are full of long digit strings that are not money —
# UPCs, product codes, store and phone numbers — and browser testing on real
# photographed receipts showed exactly that failure: "#60101" and
# "04900005375" were picked up as the largest amount on the page and reported
# as the total. Requiring a two-digit decimal, thousands grouping, or an
# adjacent symbol is what separates "$2.18" from a barcode.
_AMOUNT_RE = re.compile(
    r"(?<![0-9])"
    r"(?:"
    r"[$\u20ac\u00a3\u20ba\u00a5]\s?[0-9]{1,6}(?![0-9.,])"  # symbol + whole units (JPY)
    r"|[0-9]{1,3}(?:[.,\s][0-9]{3})+(?:[.,][0-9]{2})?"  # grouped thousands
    r"|[0-9]{1,6}[.,][0-9]{2}"  # plain decimal amount
    r")"
    r"(?![0-9])"
)
_TOTAL_RE = re.compile(
    r"grand total|total due|amount due|balance due|amount paid|to pay|\btotal\b"
    r"|\btoplam\b|\bgesamt\b|\bimporte\b", re.I
)
_SUBTOTAL_RE = re.compile(r"sub\s?-?total|net total|\bnet\b", re.I)
_TAX_RE = re.compile(r"\b(vat|tax|gst|hst|kdv|iva|mwst|tva)\b", re.I)
_PERCENT_RE = re.compile(r"[0-9]+(?:[.,][0-9]+)?\s*%")
_MONTHS = {
    "jan": 1, "feb": 2, "mar": 3, "apr": 4, "may": 5, "jun": 6,
    "jul": 7, "aug": 8, "sep": 9, "oct": 10, "nov": 11, "dec": 12,
}
_ISO_DATE_RE = re.compile(r"(?<![0-9])(20[0-9]{2})[-/.]([01]?[0-9])[-/.]([0-3]?[0-9])(?![0-9])")
_DMY_DATE_RE = re.compile(r"(?<![0-9])([0-3]?[0-9])[-/.]([0-3]?[0-9])[-/.]((?:20)?[0-9]{2})(?![0-9])")
_NAMED_DATE_RE = re.compile(
    r"(?<![0-9])([0-3]?[0-9])\s+([A-Za-z]{3,9})\.?\s+(20[0-9]{2})(?![0-9])"
)


def _parse_amount(raw):
    """Parse one money-ish token into a float, or None.

    Handles both separator conventions: whichever of "." / "," comes last and
    is followed by one or two digits is the decimal point; everything else is
    thousands grouping.
    """
    # Drop the currency symbol the amount pattern may have captured, plus any
    # stray spacing; what is left must be digits and separators only.
    s = re.sub(r"[^0-9.,]", "", str(raw))
    if not s:
        return None
    decimal_at = max(s.rfind("."), s.rfind(","))
    if decimal_at != -1 and len(s) - decimal_at - 1 in (1, 2):
        whole = re.sub(r"[.,]", "", s[:decimal_at])
        s = whole + "." + s[decimal_at + 1:]
    else:
        s = re.sub(r"[.,]", "", s)
    try:
        return round(float(s), 2)
    except ValueError:
        return None


def _amounts_in(line):
    """Every money-ish amount on one line.

    Percentages are stripped first, so "VAT 8.1% 1.26" yields the 1.26 charged
    and not the 8.1 rate.
    """
    found = []
    for match in _AMOUNT_RE.finditer(_PERCENT_RE.sub(" ", line)):
        value = _parse_amount(match.group(0))
        if value is not None:
            found.append(value)
    return found


def _normalize_date(text):
    """Return YYYY-MM-DD when the date is unambiguous, else the matched text.

    Deliberately does NOT guess between 03/04/2026 readings — an expense table
    that silently swaps day and month is worse than one that shows what it saw.
    """
    match = _ISO_DATE_RE.search(text)
    if match:
        year, month, day = (int(g) for g in match.groups())
        if 1 <= month <= 12 and 1 <= day <= 31:
            return "%04d-%02d-%02d" % (year, month, day)
    match = _NAMED_DATE_RE.search(text)
    if match:
        day, name, year = match.group(1), match.group(2)[:3].lower(), match.group(3)
        if name in _MONTHS:
            return "%s-%02d-%02d" % (year, _MONTHS[name], int(day))
    match = _DMY_DATE_RE.search(text)
    if match:
        first, second, year = (int(g) for g in match.groups())
        if year < 100:
            year += 2000
        if first > 12 and second <= 12:
            return "%04d-%02d-%02d" % (year, second, first)
        if second > 12 and first <= 12:
            return "%04d-%02d-%02d" % (year, first, second)
        return match.group(0)  # ambiguous — report it verbatim
    return ""


def _detect_currency(text):
    match = _CURRENCY_CODE_RE.search(text)
    if match:
        return match.group(1).upper()
    for symbol, code in _CURRENCY_SYMBOLS.items():
        if symbol in text:
            return code
    return ""


def _guess_vendor(lines):
    """Receipts put the merchant at the top — take the first line that reads
    like a name rather than an address, a date or a price."""
    for line in lines[:6]:
        letters = sum(character.isalpha() for character in line)
        if letters < 3 or letters < len(line) / 3:
            continue
        if _TOTAL_RE.search(line) or _normalize_date(line):
            continue
        return line[:60]
    return ""


def _find_total(lines):
    """(amount, how it was found) for the bottom-line total.

    A line that names a total wins over a big number found anywhere, and the
    LAST such line wins over earlier ones (grand totals come last).
    """
    candidate = None
    for line in lines:
        if _SUBTOTAL_RE.search(line) or not _TOTAL_RE.search(line):
            continue
        amounts = _amounts_in(line)
        if amounts:
            candidate = max(amounts)
    if candidate is not None:
        return candidate, "total line"
    everything = [value for line in lines for value in _amounts_in(line)]
    if everything:
        return max(everything), "largest amount"
    return None, ""


def _find_tax(lines):
    tax = None
    for line in lines:
        if not _TAX_RE.search(line):
            continue
        amounts = _amounts_in(line)
        if amounts:
            tax = max(amounts)
    return tax


def _scan_receipt(text):
    """Deterministic extraction of one receipt's text. No model involved."""
    lines = [" ".join(raw.split()) for raw in text.splitlines()]
    lines = [line for line in lines if line]
    total, total_source = _find_total(lines)
    return {
        "vendor": _guess_vendor(lines),
        "date": _normalize_date(text),
        "currency": _detect_currency(text),
        "total": total,
        "total_source": total_source,
        "tax": _find_tax(lines),
        "category": "",
        "lines": len(lines),
    }


def _first_json_object(reply):
    """Pull the first {...} block out of a model reply, or return None."""
    start = reply.find("{")
    end = reply.rfind("}")
    if start == -1 or end <= start:
        return None
    try:
        parsed = json.loads(reply[start:end + 1])
    except (ValueError, TypeError):
        return None
    return parsed if isinstance(parsed, dict) else None


async def _classify(row, text):
    """Ask the model for the two fuzzy fields only: vendor name and category.

    Best-effort by construction — any failure keeps the deterministic guess.
    Amounts are never sent back through the model.
    """
    prompt = (
        "Read this receipt text and reply with ONLY a JSON object, no prose:\n"
        '{"vendor": "<merchant name>", "category": "<one of: '
        + ", ".join(CATEGORIES) + '>"}\n\nRECEIPT TEXT:\n'
        + text[:CLASSIFY_CHARS]
    )
    reply = await agentop_llm.generate(prompt, "You extract fields as strict JSON.")
    parsed = _first_json_object(reply or "")
    if not parsed:
        return row["vendor"], "other"
    vendor = str(parsed.get("vendor") or row["vendor"]).strip()[:60]
    category = str(parsed.get("category") or "").strip().lower()
    return vendor or row["vendor"], category if category in CATEGORIES else "other"


async def _extract_receipt(source, text):
    """(JS-callable) Turn one receipt's text into a table row.

    Underscore-prefixed so it is never exposed to the LLM as a tool. Returns
    the row as JSON for the UI.
    """
    row = _scan_receipt(text)
    row["source"] = source
    try:
        row["vendor"], row["category"] = await _classify(row, text)
    except Exception:  # noqa: BLE001 — enrichment is optional, never fatal
        row["category"] = row["category"] or "other"
    _RECEIPTS.append(row)
    return json.dumps(row)


def _rows_json():
    """(JS-callable) Every extracted row, for the table and the CSV export."""
    return json.dumps(_RECEIPTS)


def _summary():
    """(JS-callable) Totals per currency and per category, for the summary bar."""
    by_currency, by_category = {}, {}
    for row in _RECEIPTS:
        if row.get("total") is None:
            continue
        currency = row.get("currency") or "?"
        by_currency[currency] = round(by_currency.get(currency, 0) + row["total"], 2)
        category = row.get("category") or "other"
        by_category[category] = round(by_category.get(category, 0) + row["total"], 2)
    return json.dumps(
        {"count": len(_RECEIPTS), "by_currency": by_currency, "by_category": by_category}
    )


def _reset():
    """(JS-callable) Clear the table."""
    _RECEIPTS.clear()
    return "0"


def _facts_sheet():
    """The extracted table as text, capped for small context windows."""
    rows = _RECEIPTS[-MAX_ROWS_IN_PROMPT:]
    out = ["#. vendor | date | category | total | currency | tax | file"]
    for index, row in enumerate(rows, 1):
        out.append(
            "%d. %s | %s | %s | %s | %s | %s | %s"
            % (
                index,
                row.get("vendor") or "(unknown)",
                row.get("date") or "(no date)",
                row.get("category") or "other",
                "(not found)" if row.get("total") is None else row["total"],
                row.get("currency") or "?",
                "-" if row.get("tax") is None else row["tax"],
                row.get("source") or "",
            )
        )
    return "\n".join(out)[:MAX_SHEET_CHARS]


async def process_user_query(query):
    """Answer questions about the extracted receipts, and only about them.

    Overrides the default router: the data is a small structured table, so it
    is handed over whole rather than retrieved (no embeddings needed).
    """
    if not _RECEIPTS:
        return (
            "No receipts yet. Add a photo or a PDF invoice above (or press "
            "Try a sample receipt) and I will pull out the vendor, date and total."
        )
    grounded_prompt = (
        "Answer the QUESTION using ONLY the EXPENSE TABLE below, which was "
        "extracted from the user's own receipts. Amounts in the table are exact "
        "- never invent or re-estimate one, and say so if a figure is missing. "
        "When you add figures up, show the arithmetic briefly.\n\n"
        "EXPENSE TABLE:\n" + _facts_sheet() + "\n\nQUESTION: " + query
    )
    return await agentop_llm.generate(
        grounded_prompt, globals().get("TEMPLATE_SYSTEM_PROMPT", "")
    )


async def _ocr_line(image_url):
    """(JS-callable) Read one cropped line strip with the on-device OCR model."""
    return (await agentop_ml.ocr(image_url)) or ""


def _extract_pdf_text(pdf_bytes):
    """(JS-callable) Extract text from a digital invoice PDF via pypdf."""
    import io
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(pdf_bytes))
    return "\n\n".join((page.extract_text() or "") for page in reader.pages)