app.py

"""Spend Elegy — receipt -> categorized spend, a bar chart, and a short elegy.

Pipeline:
  input (paste | .txt | .pdf | image)
    -> extract text (direct | PDF text-layer | MiniCPM-V-4.6 OCR for images)
    -> text model -> strict categorized JSON (pipeline.py); non-receipt input is
       rejected -> outputs: summary, line-items table, per-category bar chart,
       per-category totals, an elegy (a 2nd text-model call), and the raw JSON.

Deployment: a Hugging Face ZeroGPU Gradio Space. Models load at module scope
(CUDA on the Space) and generation runs inside @spaces.GPU.

Local dev (Apple Silicon / no CUDA) — Nemotron-3 Nano's Mamba kernels are
CUDA-only, so use an MPS-friendly text model and skip the vision model:

    NORMALIZER_MODEL_ID=Qwen/Qwen2.5-3B-Instruct LOAD_VISION=0 python app.py
"""

from __future__ import annotations

import html
import json
import os

import gradio as gr
import pandas as pd
import spaces

import extraction
import models
import pipeline

GPU_DURATION = int(os.environ.get("GPU_DURATION", "180"))

CSS = """
#se-elegy .se-card {
    background: linear-gradient(160deg, #2a2440 0%, #1c2030 100%);
    color: #e8e4f2; border: 1px solid #3b3656; border-radius: 14px;
    padding: 20px 22px; min-height: 180px;
    font-family: Georgia, "Times New Roman", serif; font-style: italic;
    font-size: 15px; line-height: 1.75; text-align: center;
    box-shadow: inset 0 0 50px rgba(0,0,0,.28);
}
#se-elegy .se-card .se-candle { font-style: normal; font-size: 20px; opacity: .8; }
.se-working { color: #7c83d8; font-weight: 600; animation: se-pulse 1.2s ease-in-out infinite; }
@keyframes se-pulse { 0%,100% { opacity: .4 } 50% { opacity: 1 } }
"""

PLACEHOLDER = "GREENLEAF MARKET\nOrganic Bananas   1.29\nWhole Milk 2L     2.49\n..."


def _elegy_html(elegy: str) -> str:
    text = (elegy or "").replace("—", ", ").replace("–", "-").strip()
    body = html.escape(text).replace("\n", "<br>")
    if not body:
        return ""
    return f"<div class='se-card'><div class='se-candle'>🕯️</div>{body}</div>"


def _notice(message: str):
    """Cleared outputs with a notice in the summary (errors / non-receipts)."""
    empty_chart = pd.DataFrame(columns=["Category", "Total"])
    return message, [], empty_chart, [], "", ""


@spaces.GPU(duration=GPU_DURATION)
def run_inference(images, receipt_text: str) -> dict:
    """GPU stage: OCR (if images) -> categorized JSON -> elegy."""
    if images:
        try:
            ocr_text = extraction.ocr_images(images)
        except Exception as exc:
            raise RuntimeError(f"couldn't read the image/PDF (OCR failed): {exc}")
        receipt_text = f"{receipt_text}\n\n{ocr_text}".strip() if receipt_text else ocr_text

    raw = models.text_model.generate(pipeline.build_messages(receipt_text))
    try:
        data = pipeline.parse_json(raw)
    except Exception:
        raw = models.text_model.generate(
            pipeline.build_messages(receipt_text, remind=True)
        )
        data = pipeline.parse_json(raw)  # may raise -> handled by caller

    record = pipeline.normalize_record(data)
    if not record["is_receipt"] or not record["line_items"]:
        return {"record": record, "elegy": None, "is_receipt": False}

    elegy = models.text_model.generate(pipeline.build_elegy_messages(record)).strip()
    return {"record": record, "elegy": elegy, "is_receipt": True}


def parse_receipt(file_path: str | None, pasted_text: str | None):
    # A file takes precedence over pasted text (they are not combined).
    if file_path:
        try:
            text, images = extraction.extract_from_file(file_path)
        except ValueError as exc:  # unsupported file type
            return _notice(f"⚠️ {exc}")
    else:
        text, images = (pasted_text or "").strip(), []

    if not text and not images:
        return _notice("⚠️ Upload a .txt / .pdf / image receipt, or paste some text first.")

    try:
        result = run_inference(images, text)
    except Exception as exc:  # OCR error, or JSON unparseable after retry
        return _notice(f"⚠️ Sorry, couldn't read this one. ({exc})")

    if not result.get("is_receipt", True):
        return _notice(
            "🤔 That doesn't look like a receipt, bill, or statement. Try a grocery "
            "bill, a bank/card statement, or a note about money you spent."
        )

    record = result["record"]
    items = record["line_items"]
    chart_df = pd.DataFrame(pipeline.category_totals(items), columns=["Category", "Total"])
    return (
        pipeline.summary_markdown(record),
        pipeline.items_table(items),
        chart_df,
        pipeline.category_totals(items),
        _elegy_html(result["elegy"]),
        json.dumps(record, indent=2, ensure_ascii=False),
    )


def _show_working():
    return "<span class='se-working'>⏳ Reading your receipt and composing its elegy…</span>"


def _on_file_change(file_path):
    # File takes precedence: clear + disable the paste box while a file is uploaded.
    if file_path:
        return gr.update(
            value="", interactive=False,
            placeholder="Using the uploaded file — remove it to paste text instead.",
        )
    return gr.update(interactive=True, placeholder=PLACEHOLDER)


# --- UI ---------------------------------------------------------------------

theme = gr.themes.Soft(primary_hue="orange", neutral_hue="slate")

with gr.Blocks(title="Spend Elegy", theme=theme, css=CSS) as demo:
    gr.Markdown(
        "# 🧾 Spend Elegy\n"
        "Turn a receipt, bill, or statement into a categorized spending breakdown "
        "— with a little elegy for your money."
    )

    with gr.Row(equal_height=True):
        with gr.Column():  # input
            text_input = gr.Textbox(
                label="Paste receipt / statement text", lines=12, placeholder=PLACEHOLDER
            )
            file_input = gr.File(
                label="…or upload a file (.txt, .pdf, image)",
                file_types=[".txt", ".pdf", ".png", ".jpg", ".jpeg", ".webp"],
                type="filepath",
                height=100,
            )
            with gr.Row():
                parse_button = gr.Button("Parse receipt", variant="primary", scale=2)
                gr.ClearButton([text_input, file_input], scale=1)
        with gr.Column():  # headline results
            summary_output = gr.Markdown()
            elegy_output = gr.HTML(elem_id="se-elegy")

    with gr.Row(equal_height=True):
        chart_output = gr.BarPlot(x="Category", y="Total", title="Spend by category")
        items_output = gr.Dataframe(
            headers=["Item", "Qty", "Amount", "Category"],
            datatype=["str", "number", "number", "str"],
            label="Line items",
            wrap=True,
        )

    with gr.Accordion("Per-category totals & raw JSON", open=False):
        with gr.Row(equal_height=True):
            category_output = gr.Dataframe(
                headers=["Category", "Total"],
                datatype=["str", "number"],
                label="Per-category totals",
            )
            json_output = gr.Code(language="json", label="Raw JSON")

    file_input.change(_on_file_change, inputs=file_input, outputs=text_input)

    outputs = [
        summary_output,
        items_output,
        chart_output,
        category_output,
        elegy_output,
        json_output,
    ]
    parse_button.click(_show_working, None, summary_output).then(
        parse_receipt, inputs=[file_input, text_input], outputs=outputs
    )

if __name__ == "__main__":
    demo.launch()