app.py

"""Invoice extraction model worker on HF ZeroGPU.

Public Space, but it exposes only a status panel. All invoice data lives in the
owner's Cloudflare D1/R2; this worker claims jobs from the invoice-api Worker's
internal API (secret-authenticated), extracts with LFM2-1.2B-Extract on a ZeroGPU
slice, and posts results back. Space secrets: WORKER_URL, INTERNAL_KEY.
"""
import json
import logging
import os
import re
import threading
import time

import fitz  # PyMuPDF
import gradio as gr
import httpx
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
log = logging.getLogger("worker")

MODEL_ID = "LiquidAI/LFM2-1.2B-Extract"
POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "3"))
MAX_NEW_TOKENS = 2048

_TEMPLATE = """{
  "invoice_number": null,
  "invoice_date": null,
  "due_date": null,
  "vendor": {"name": null, "address": null, "gstin": null, "pan": null, "email": null, "phone": null},
  "buyer": {"name": null, "address": null, "gstin": null, "pan": null},
  "service_description": null,
  "billing_period": null,
  "sac_code": null,
  "work_order": null,
  "amounts": {"taxable_value": null, "cgst": null, "sgst": null, "igst": null, "grand_total": null},
  "amount_in_words": null,
  "bank_details": {"bank_name": null, "account_number": null, "ifsc": null, "micr": null},
  "employees": [
    {"name": null, "monthly_billing": null, "payable_days": null, "amount": null, "gst_amount": null, "total": null}
  ],
  "other_details": null
}"""


def build_system_prompt() -> str:
    if VENDORS:
        vendor_context = "\n".join(
            f"{i}. {v.get('name')} - GSTIN {v.get('gstin')}, PAN {v.get('pan')}. Bill numbers start with {v.get('prefix')}/."
            for i, v in enumerate(VENDORS, 1)
        )
    else:
        vendor_context = os.environ.get("VENDOR_CONTEXT", "").strip()
    vc = f"\n\nKnown vendors (the invoice is always ISSUED BY exactly one of these):\n{vendor_context}\n" if vendor_context else ""
    return f"""You are an expert data extraction engine for Indian GST invoices.
Fill this exact JSON template from the invoice text. Respond with ONLY the completed JSON object.

{_TEMPLATE}
{vc}
Field guide:
- "invoice_number": the value printed after "Bill No.:" (e.g. "XXX/003/26-27"). Never an employee code.
- "vendor": the party that ISSUED the invoice (its name is in the letterhead at the very top).
  The GSTIN/PAN printed on the RIGHT side near "Bill No." belong to the VENDOR.
- "buyer": the party being billed - the name/address block at the top LEFT.
  The GSTIN/PAN printed directly under the buyer's address belong to the BUYER.
- "amounts": rupee amounts as plain JSON numbers, no commas, no symbols. Taxes (CGST/SGST/IGST)
  are amounts, not percentages; use null for any tax type not charged. "grand_total" is the
  final payable total (equals the amount in words).
- "amount_in_words": the sentence after "Rupees:" exactly as printed.
- "bank_details": the vendor's bank account printed under "Bank Details".
- "employees": one entry per employee row in the annexure table at the end (ignore the TOTAL row).
  IMPORTANT: the first column is the manager ("Kind Attention Person") - the employee's own name
  is in the "Employee Name" column. "monthly_billing" = the Billing column (monthly rate),
  "payable_days" = Total Payable Days, "amount" = Total Payable Billing, "gst_amount" = the row's
  GST/IGST amount if shown, "total" = the last number in the row (row grand total).
  Use [] if there is no employee table.
- "other_details": PO/work-order references, declarations, PF/ESIC numbers, or anything else notable.
- Copy values exactly as printed; use null when absent. NEVER invent or guess values.
- GSTIN is 15 characters; PAN is 10 characters."""

state = {
    "started_at": int(time.time()),
    "last_poll": None,
    "processed": 0,
    "failed": 0,
    "current_job": None,
    "last_error": None,
}

from anchors import VENDORS, deterministic_fields, merge_result  # noqa: E402


class NoTextLayerError(Exception):
    pass


# ---------- model ----------

log.info("Loading %s ...", MODEL_ID)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16)
model.to("cuda")  # ZeroGPU: safe at startup, GPU attaches inside @spaces.GPU calls
log.info("Model ready.")


@spaces.GPU(duration=120)
def llm_generate(text: str, company_hint: str | None) -> str:
    hint = f"(Hint: this invoice relates to {company_hint}.)\n\n" if company_hint else ""
    messages = [
        {"role": "system", "content": build_system_prompt()},
        {"role": "user", "content": f"{hint}Invoice text:\n\n{text}"},
    ]
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)
    with torch.inference_mode():
        out = model.generate(
            inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True)


_NUMERIC_KEYS = {
    "taxable_value", "cgst", "sgst", "igst", "grand_total",
    "monthly_billing", "payable_days", "amount", "gst_amount", "total",
    "quantity", "rate",
}


def _coerce_numbers(obj, key=None):
    """Recursively convert '1,26,166' style strings to numbers on numeric fields."""
    if isinstance(obj, dict):
        return {k: _coerce_numbers(v, k) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_coerce_numbers(v, key) for v in obj]
    if key in _NUMERIC_KEYS and isinstance(obj, str):
        cleaned = obj.replace(",", "").replace("₹", "").replace("Rs.", "").strip()
        try:
            n = float(cleaned)
            return int(n) if n.is_integer() else n
        except ValueError:
            return obj
    return obj


def parse_json(raw: str) -> dict:
    raw = raw.strip()
    raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw)
    start, end = raw.find("{"), raw.rfind("}")
    if start == -1 or end == -1:
        raise ValueError(f"model returned no JSON object: {raw[:200]!r}")
    raw = raw[start : end + 1]
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        import json_repair
        data = json_repair.loads(raw)
        if not isinstance(data, dict):
            raise ValueError(f"model returned unrepairable JSON: {raw[:200]!r}")
    return _coerce_numbers(data)


def pdf_to_text(pdf_bytes: bytes) -> str:
    with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
        text = "\n\n".join(page.get_text("text", sort=True) for page in doc).strip()
    if len(text) < 50:
        raise NoTextLayerError(
            "PDF contains no usable text layer (scanned image?). OCR is not supported."
        )
    return text


# ---------- poller ----------

def _base() -> str:
    return os.environ["WORKER_URL"].rstrip("/")


def _headers():
    return {"x-internal-key": os.environ["INTERNAL_KEY"]}


def process_one(client: httpx.Client) -> bool:
    """Claim and process one job. Returns False when the queue is empty."""
    resp = client.post(f"{_base()}/internal/claim", headers=_headers(), timeout=30)
    if resp.status_code == 204:
        return False
    resp.raise_for_status()
    job = resp.json()
    job_id = job["id"]
    started = time.time()
    state["current_job"] = job_id
    log.info("Processing job %s (%s)", job_id, job.get("filename"))
    try:
        pdf_resp = client.get(f"{_base()}/internal/pdf/{job_id}", headers=_headers(), timeout=120)
        if pdf_resp.status_code == 404:
            raise FileNotFoundError("PDF not found in storage (expired after 24h?)")
        pdf_resp.raise_for_status()
        text = pdf_to_text(pdf_resp.content)
        result = parse_json(llm_generate(text, job.get("company_hint")))
        result = merge_result(result, deterministic_fields(text))
        payload = {"result": result, "latency_s": round(time.time() - started, 1)}
        state["processed"] += 1
        log.info("Job %s done in %.1fs", job_id, time.time() - started)
    except (NoTextLayerError, FileNotFoundError, ValueError, json.JSONDecodeError) as e:
        payload = {"error": str(e), "latency_s": round(time.time() - started, 1)}
        state["failed"] += 1
        log.warning("Job %s failed: %s", job_id, e)
    except Exception as e:
        # Transient (GPU quota, network): don't complete the job; the Worker
        # re-queues it automatically after the stale timeout.
        state["last_error"] = f"{type(e).__name__}: {e}"
        state["current_job"] = None
        log.exception("Transient error on job %s; leaving it for stale re-queue", job_id)
        time.sleep(60)
        return True
    client.post(
        f"{_base()}/internal/jobs/{job_id}/complete",
        headers=_headers(),
        json=payload,
        timeout=60,
    ).raise_for_status()
    state["current_job"] = None
    return True


def poll_loop():
    log.info("Poller started against %s", _base())
    with httpx.Client() as client:
        while True:
            try:
                state["last_poll"] = int(time.time())
                if process_one(client):
                    continue  # drain queue without sleeping
            except Exception as e:
                state["last_error"] = f"{type(e).__name__}: {e}"
                log.exception("Poll loop error; backing off")
                time.sleep(15)
                continue
            time.sleep(POLL_INTERVAL)


if os.environ.get("WORKER_URL") and os.environ.get("INTERNAL_KEY"):
    threading.Thread(target=poll_loop, daemon=True, name="poller").start()
else:
    log.error("WORKER_URL / INTERNAL_KEY not set; poller not started")


# ---------- minimal public UI (status only, no data) ----------

def status():
    s = dict(state)
    s["uptime_s"] = int(time.time()) - s.pop("started_at")
    if s["last_poll"]:
        s["seconds_since_poll"] = int(time.time()) - s["last_poll"]
    return s


with gr.Blocks(title="invoice-processor worker") as demo:
    gr.Markdown(
        "# 🧾 Invoice processor — model worker\n"
        "Private job worker for its owner's invoice pipeline. No data is accessible here.\n"
        "Powered by [LiquidAI/LFM2-1.2B-Extract](https://huggingface.co/LiquidAI/LFM2-1.2B-Extract)."
    )
    out = gr.JSON(label="worker status")
    demo.load(status, outputs=out)
    gr.Timer(30).tick(status, outputs=out)

demo.launch()