app.py

"""Gradio app: OCR documents into structured notification/data records.

Upload scanned PDFs or images (construction reports, compliance
certificates, invoices, etc.), scan the extracted text for domain
"notification" keywords, pull out labeled data fields (standard IDs/dates/
amounts plus any custom labels the user names), and export everything to
CSV as tidy (one-fact-per-row) tables, alongside an audit record of the
settings that produced the run.
"""
import json
import os
import tempfile
import uuid
from datetime import datetime

import gradio as gr
import pandas as pd

import chat_config
import extractors
import ocr_engine

CATEGORY_CHOICES = chat_config.CATEGORY_CHOICES

DEFAULT_SETTINGS = {
    "categories": CATEGORY_CHOICES,
    "custom_keywords": [],
    "custom_fields": [],
    "only_matches": True,
}

FIELDS_COLUMNS = [
    "run_id", "file_name", "page", "field_type", "field_name", "field_value",
    "text_source", "ocr_confidence",
]
NOTIFICATIONS_COLUMNS = [
    "run_id", "file_name", "page", "category", "keyword", "context",
    "text_source", "ocr_confidence",
]
METADATA_COLUMNS = [
    "run_id", "timestamp", "files", "categories", "custom_keywords",
    "custom_fields", "only_matches_filter", "chat_model", "chat_transcript",
]


def _file_path(f):
    return f if isinstance(f, str) else f.name


def _parse_comma_list(s):
    return [k.strip() for k in s.split(",") if k.strip()] if s and s.strip() else []


def _format_transcript(history):
    if not history:
        return ""
    lines = []
    for turn in history:
        role = turn.get("role", "?")
        content = turn.get("content", "")
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


def process_files(files, categories, custom_keywords_str, custom_fields_str, only_matches, chat_history):
    empty = (pd.DataFrame(columns=FIELDS_COLUMNS), pd.DataFrame(columns=NOTIFICATIONS_COLUMNS))
    if not files:
        return "Upload at least one PDF or image file.", empty[0], empty[1], None, "", {}

    custom_keywords = _parse_comma_list(custom_keywords_str)
    custom_field_labels = _parse_comma_list(custom_fields_str)

    run_id = uuid.uuid4().hex[:8]
    field_rows = []
    notification_rows = []
    full_text_log = []
    warnings = []
    total_pages = 0
    file_names = []

    for f in files:
        path = _file_path(f)
        file_name = os.path.basename(path)
        file_names.append(file_name)
        try:
            pages = ocr_engine.extract_pages(path)
        except Exception as exc:
            warnings.append(f"{file_name}: {exc}")
            continue

        for page_data in pages:
            total_pages += 1
            text = page_data["text"]
            source = page_data["source"]
            confidence = page_data["ocr_confidence"]

            notifications = extractors.find_notifications(text, categories, custom_keywords)
            std_fields = extractors.extract_fields(text, categories)
            custom_field_values = extractors.extract_custom_fields(text, custom_field_labels)

            full_text_log.append(f"--- {file_name} | page {page_data['page']} | {source} ---\n{text}\n")

            has_content = bool(notifications) or any(std_fields.values()) or any(custom_field_values.values())
            if only_matches and not has_content:
                continue

            for field_name, values in std_fields.items():
                for value in values:
                    field_rows.append({
                        "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                        "field_type": "standard", "field_name": field_name, "field_value": value,
                        "text_source": source, "ocr_confidence": confidence,
                    })
            for field_name, values in custom_field_values.items():
                for value in values:
                    field_rows.append({
                        "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                        "field_type": "custom", "field_name": field_name, "field_value": value,
                        "text_source": source, "ocr_confidence": confidence,
                    })
            for n in notifications:
                notification_rows.append({
                    "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                    "category": n["category"], "keyword": n["keyword"], "context": n["context"],
                    "text_source": source, "ocr_confidence": confidence,
                })

    fields_df = pd.DataFrame(field_rows, columns=FIELDS_COLUMNS)
    notifications_df = pd.DataFrame(notification_rows, columns=NOTIFICATIONS_COLUMNS)

    metadata_row = {
        "run_id": run_id,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "files": "; ".join(file_names),
        "categories": "; ".join(categories),
        "custom_keywords": "; ".join(custom_keywords),
        "custom_fields": "; ".join(custom_field_labels),
        "only_matches_filter": only_matches,
        "chat_model": chat_config.DEFAULT_MODEL if chat_history else "",
        "chat_transcript": _format_transcript(chat_history),
    }
    metadata_df = pd.DataFrame([metadata_row], columns=METADATA_COLUMNS)

    csv_paths = []
    if not fields_df.empty or not notifications_df.empty:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        tmp_dir = tempfile.gettempdir()
        fields_path = os.path.join(tmp_dir, f"ocr_fields_{timestamp}_{run_id}.csv")
        notifications_path = os.path.join(tmp_dir, f"ocr_notifications_{timestamp}_{run_id}.csv")
        metadata_path = os.path.join(tmp_dir, f"ocr_run_metadata_{timestamp}_{run_id}.csv")
        fields_df.to_csv(fields_path, index=False)
        notifications_df.to_csv(notifications_path, index=False)
        metadata_df.to_csv(metadata_path, index=False)
        csv_paths = [fields_path, notifications_path, metadata_path]

    status_lines = [
        f"Processed {len(files)} file(s), {total_pages} page(s) total.",
        f"{len(fields_df)} field value(s), {len(notifications_df)} notification hit(s)"
        + (" (matches-only pages)." if only_matches else " (all pages)."),
    ]
    if warnings:
        status_lines.append("Warnings: " + "; ".join(warnings))
    status = "\n".join(status_lines)

    return status, fields_df, notifications_df, csv_paths, "\n".join(full_text_log), metadata_row


def handle_chat(history, message, settings, hf_token):
    history = history or []
    if not message or not message.strip():
        return (
            history, "", settings, settings,
            settings["categories"], ", ".join(settings["custom_keywords"]),
            ", ".join(settings["custom_fields"]), settings["only_matches"],
        )

    reply, new_settings = chat_config.chat_update(history, message.strip(), settings, hf_token)
    history = history + [
        {"role": "user", "content": message.strip()},
        {"role": "assistant", "content": reply},
    ]
    return (
        history, "", new_settings, new_settings,
        new_settings["categories"], ", ".join(new_settings["custom_keywords"]),
        ", ".join(new_settings["custom_fields"]), new_settings["only_matches"],
    )


with gr.Blocks(title="OCR Notification & Data Extractor") as demo:
    gr.Markdown(
        "# OCR Notification & Data Extractor\n"
        "Upload scanned PDFs or images (construction reports, compliance "
        "certificates, invoices, SAP printouts, etc.). The app OCRs each "
        "page, flags notification-worthy keywords per domain, and pulls out "
        "labeled data fields - standard ones scoped to the categories you "
        "pick, plus any custom labels you name (e.g. \"gross weight\", "
        "\"delivery date\"). Every run exports three tidy CSVs: extracted "
        "fields, notification hits, and a metadata/audit record of exactly "
        "what settings produced them."
    )

    settings_state = gr.State(dict(DEFAULT_SETTINGS))

    with gr.Tab("1. Tailor extraction (chat)"):
        gr.Markdown(
            "Describe what you want **flagged** (notifications) vs. **pulled "
            "as data** (fields), in plain language, e.g. *\"flag overdue "
            "invoices and non-compliance notices, and pull the reference "
            "number, gross weight, and delivery date\"*. This updates the "
            "settings used in **Upload & Run** - you can still edit them by "
            "hand there too."
        )
        hf_token_input = gr.Textbox(
            label="Hugging Face token (optional if HF_TOKEN is set as a Space secret)",
            type="password",
            placeholder="hf_...",
        )
        chatbot = gr.Chatbot(height=350, label="Extraction assistant")
        with gr.Row():
            chat_input = gr.Textbox(
                label="Message", scale=4,
                placeholder="e.g. Flag overdue invoices; pull reference number, gross weight, delivery date",
            )
            chat_send = gr.Button("Send", scale=1, variant="primary")
        settings_display = gr.JSON(label="Active extraction settings", value=DEFAULT_SETTINGS)

    with gr.Tab("2. Upload & Run"):
        with gr.Row():
            with gr.Column(scale=1):
                files_input = gr.Files(
                    label="PDF or image files",
                    file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"],
                )
                categories_input = gr.CheckboxGroup(
                    choices=CATEGORY_CHOICES,
                    value=CATEGORY_CHOICES,
                    label="Notification categories",
                    info="Also scopes which standard fields (PO/SAP/invoice/cert) get extracted - see CATEGORY_FIELD_MAP.",
                )
                custom_keywords_input = gr.Textbox(
                    label="Custom notification keywords (comma-separated)",
                    placeholder="e.g. re-inspection, warranty claim, hold payment",
                    info="Phrases that just raise a flag - not extracted as values.",
                )
                custom_fields_input = gr.Textbox(
                    label="Custom fields to extract as data (comma-separated labels)",
                    placeholder="e.g. reference number, gross weight, delivery date",
                    info="Labels looked up in the document; the value after each label is captured.",
                )
                only_matches_input = gr.Checkbox(
                    value=True,
                    label="Only export pages with a notification hit or extracted field",
                )
                run_button = gr.Button("Run OCR & Extract", variant="primary")

            with gr.Column(scale=2):
                status_output = gr.Textbox(label="Status", lines=3, interactive=False)
                fields_table_output = gr.Dataframe(label="Extracted fields (tidy)", headers=FIELDS_COLUMNS, wrap=True)
                notifications_table_output = gr.Dataframe(
                    label="Notification hits (tidy)", headers=NOTIFICATIONS_COLUMNS, wrap=True
                )
                csv_output = gr.Files(label="Download CSVs (fields, notifications, run metadata)")
                run_metadata_output = gr.JSON(label="This run's settings (audit record)")
                with gr.Accordion("Full OCR text (debug)", open=False):
                    text_output = gr.Textbox(label="Raw extracted text", lines=20)

    chat_outputs = [
        chatbot, chat_input, settings_state, settings_display,
        categories_input, custom_keywords_input, custom_fields_input, only_matches_input,
    ]
    chat_send.click(fn=handle_chat, inputs=[chatbot, chat_input, settings_state, hf_token_input], outputs=chat_outputs)
    chat_input.submit(fn=handle_chat, inputs=[chatbot, chat_input, settings_state, hf_token_input], outputs=chat_outputs)

    run_button.click(
        fn=process_files,
        inputs=[files_input, categories_input, custom_keywords_input, custom_fields_input, only_matches_input, chatbot],
        outputs=[status_output, fields_table_output, notifications_table_output, csv_output, text_output, run_metadata_output],
    )

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