app.py

"""
Smart Warehouse AI Assistant
=============================
A Daifuku-style intralogistics AI copilot demo, built for a Hugging Face
Space. Combines:

  1. An LLM-powered assistant (RAG: TF-IDF retrieval + hosted LLM via the
     HF Inference API) for natural-language warehouse operations Q&A.
  2. A TF-IDF + Logistic Regression intent classifier that routes queries
     into 8 operational categories.
  3. A lightweight NL -> structured query layer over synthetic inventory /
     order tables.
  4. An Isolation Forest anomaly detector for conveyor/crane sensor
     streams (predictive maintenance).
  5. A Model Evaluation tab reporting real accuracy/F1/ROC-AUC metrics
     computed by build_artifacts.py.

Author: (your name here) -- built as an application project for Daifuku Co., Ltd.
"""

import json
import os

import gradio as gr
import pandas as pd

from src.anomaly_model import FEATURES as SENSOR_FEATURES
from src.anomaly_model import load_artifacts as load_anomaly_artifacts
from src.anomaly_model import score_reading
from src.data_generation import generate_inventory_db, generate_orders_db
from src.intent_model import INTENT_DESCRIPTIONS, load_pipeline, predict as intent_predict
from src.inventory_db import query_inventory, query_orders
from src.llm_client import answer_query
from src.retriever import KBRetriever

ROOT = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(ROOT, "models")
DATA_DIR = os.path.join(ROOT, "data")
ASSETS_DIR = os.path.join(ROOT, "assets")

# --------------------------------------------------------------------------
# Load pre-trained artifacts (fast: no training happens at Space startup
# under normal conditions). If unpickling fails -- e.g. the models were
# built with a different scikit-learn version than the one installed in
# this container -- we transparently rebuild everything from scratch. This
# is fast (a few seconds, see build_artifacts.py) and makes the app immune
# to sklearn version-pinning mismatches between build time and deploy time.
# --------------------------------------------------------------------------

def _load_or_rebuild_artifacts():
    intent_path = os.path.join(MODELS_DIR, "intent_pipeline.joblib")
    anomaly_model_path = os.path.join(MODELS_DIR, "anomaly_iforest.joblib")
    anomaly_scaler_path = os.path.join(MODELS_DIR, "anomaly_scaler.joblib")

    def _try_load():
        pipeline = load_pipeline(intent_path)
        model, scaler = load_anomaly_artifacts(anomaly_model_path, anomaly_scaler_path)
        # Smoke-test the loaded pipeline against the exact code path the app
        # uses at request time. On a scikit-learn version mismatch, sklearn
        # sometimes unpickles "successfully" but throws later on first real
        # use (e.g. LogisticRegression missing an internal attribute) --
        # catching that here, at import time, is what makes this self-healing.
        intent_predict(pipeline, "healthcheck")
        score_reading(model, scaler, {"motor_temp_c": 55, "vibration_mm_s": 2.2, "current_amps": 12, "belt_speed_mps": 1.5})
        return pipeline, model, scaler

    try:
        return _try_load()
    except Exception as e:  # noqa: BLE001 -- any load/version issue triggers a rebuild
        print(f"[startup] Could not load pre-built model artifacts ({type(e).__name__}: {e}). "
              f"Rebuilding from scratch with the installed scikit-learn version...")
        import build_artifacts
        build_artifacts.build_intent_classifier()
        build_artifacts.build_anomaly_detector()
        build_artifacts.build_retrieval_eval()
        build_artifacts.build_inventory_and_orders()
        print("[startup] Rebuild complete.")
        return _try_load()


intent_pipeline, anomaly_model, anomaly_scaler = _load_or_rebuild_artifacts()
retriever = KBRetriever()

# Prefer the pre-generated CSVs (so demo state matches the eval run); fall back to
# regenerating in-memory if they're missing for some reason.
try:
    inventory_df = pd.read_csv(os.path.join(DATA_DIR, "inventory.csv"))
    orders_df = pd.read_csv(os.path.join(DATA_DIR, "orders.csv"))
except FileNotFoundError:
    inventory_df = generate_inventory_db()
    orders_df = generate_orders_db()


def _load_json(name):
    path = os.path.join(DATA_DIR, name)
    if os.path.exists(path):
        with open(path) as f:
            return json.load(f)
    return {}


intent_eval = _load_json("intent_eval.json")
anomaly_eval = _load_json("anomaly_eval.json")
retrieval_eval = _load_json("retrieval_eval.json")
latency_eval = _load_json("latency_eval.json")

HF_TOKEN_SET = bool(os.environ.get("HF_TOKEN"))

# --------------------------------------------------------------------------
# ZeroGPU compatibility shim
# --------------------------------------------------------------------------
# This app is CPU-only by design (scikit-learn locally, LLM calls go to the
# remote HF Inference API). Some Spaces accounts, however, only offer the
# free "ZeroGPU" hardware tier, which requires at least one function
# decorated with `@spaces.GPU` to be present or the platform's startup
# check fails with "No @spaces.GPU function detected". This is a harmless,
# unused health-check function that satisfies that requirement without
# changing any real behaviour -- it is never called on the request path.
try:
    import spaces

    @spaces.GPU(duration=5)
    def _zerogpu_healthcheck():
        return True

except ImportError:
    # Running locally / on CPU-basic hardware where `spaces` isn't installed.
    def _zerogpu_healthcheck():
        return True


# ==========================================================================
# TAB 1 -- AI Assistant (LLM + RAG + intent routing)
# ==========================================================================

def chat_fn(message, history):
    intent = intent_predict(intent_pipeline, message)
    response = answer_query(message, retriever)

    intent_label = INTENT_DESCRIPTIONS.get(intent.intent, intent.intent)
    meta_lines = [f"**Detected intent:** {intent_label} ({intent.confidence:.0%} confidence)"]
    if response.sources:
        src_str = ", ".join(f"{s.title} ({s.score:.2f})" for s in response.sources)
        meta_lines.append(f"**Retrieved context:** {src_str}")
    meta_lines.append(
        f"**Generation:** {'LLM (' + response.model_id + ')' if response.used_llm else 'retrieval-only fallback'}"
        f" · {response.latency_s * 1000:.0f} ms"
    )

    full_reply = response.answer + "\n\n---\n" + "\n".join(meta_lines)
    return full_reply


ASSISTANT_EXAMPLES = [
    "The conveyor belt in Zone C is making noise",
    "How many units of SKU-1042 are in Zone B?",
    "What's the difference between an AGV and an AMR?",
    "A forklift near-miss was reported in Zone A",
    "What's the fastest picking route for a high volume order?",
    "Is Crane-03 operational?",
]


# ==========================================================================
# TAB 2 -- Inventory & Order Query
# ==========================================================================

def inventory_query_fn(text):
    intent = intent_predict(intent_pipeline, text)
    if intent.intent == "order_status":
        result = query_orders(orders_df, text)
        note = "Interpreted as an **order status** query."
    else:
        result = query_inventory(inventory_df, text)
        note = "Interpreted as an **inventory** query."
    if result.empty:
        result = pd.DataFrame({"message": ["No matching records found for this query."]})
    return note, result


# ==========================================================================
# TAB 3 -- Predictive Maintenance / Anomaly Detection
# ==========================================================================

def anomaly_fn(motor_temp, vibration, current, belt_speed):
    reading = {
        "motor_temp_c": motor_temp,
        "vibration_mm_s": vibration,
        "current_amps": current,
        "belt_speed_mps": belt_speed,
    }
    result = score_reading(anomaly_model, anomaly_scaler, reading)
    verdict = "🔴 ANOMALY DETECTED" if result.is_anomaly else "🟢 Normal operating range"
    detail = (
        f"### {verdict}\n\n"
        f"**Anomaly score:** {result.anomaly_score:.2f} / 1.00\n\n"
        f"| Sensor | Value | Typical normal range |\n"
        f"|---|---|---|\n"
        f"| Motor temperature | {motor_temp:.1f} °C | 47–63 °C |\n"
        f"| Vibration | {vibration:.2f} mm/s | 1.2–3.2 mm/s |\n"
        f"| Motor current | {current:.1f} A | 9–15 A |\n"
        f"| Belt speed | {belt_speed:.2f} m/s | 1.2–1.8 m/s |\n"
    )
    if result.is_anomaly:
        detail += (
            "\n**Recommended action:** Flag for inspection. Elevated temperature + "
            "vibration + current with reduced belt speed typically indicates bearing "
            "wear, belt misalignment, or motor overload -- schedule maintenance before "
            "the next shift to avoid an unplanned stoppage."
        )
    return detail


ANOMALY_PRESETS = {
    "Normal reading": (55.0, 2.2, 12.0, 1.5),
    "Early bearing wear": (68.0, 3.8, 15.5, 1.3),
    "Severe fault (imminent failure)": (85.0, 6.5, 21.0, 0.6),
}


def load_preset(name):
    return ANOMALY_PRESETS[name]


# ==========================================================================
# TAB 4 -- Model Evaluation
# ==========================================================================

def build_evaluation_markdown():
    cls_report = intent_eval.get("classification_report", {})
    per_class_rows = []
    for cls in intent_eval.get("classes", []):
        stats = cls_report.get(cls, {})
        per_class_rows.append(
            f"| {cls} | {stats.get('precision', 0):.2f} | {stats.get('recall', 0):.2f} | "
            f"{stats.get('f1-score', 0):.2f} | {int(stats.get('support', 0))} |"
        )
    per_class_table = "\n".join(per_class_rows)

    retrieval_rows = "\n".join(
        f"| {r['query']} | {r['expected']} | {r['retrieved_top1']} | "
        f"{'✅' if r['hit@1'] else ('〰️' if r['hit@2'] else '❌')} | {r['top1_score']:.2f} |"
        for r in retrieval_eval.get("rows", [])
    )

    md = f"""
## 1. Intent Classifier (TF-IDF + Logistic Regression)

Trained on {intent_eval.get('n_train', '?')} examples, evaluated on a held-out
stratified test split of {intent_eval.get('n_test', '?')} examples across
{intent_eval.get('n_classes', '?')} intent classes.

| Metric | Score |
|---|---|
| **Accuracy** | **{intent_eval.get('accuracy', 0):.2%}** |
| **Macro F1** | **{intent_eval.get('macro_f1', 0):.2%}** |

**Per-class performance:**

| Intent | Precision | Recall | F1 | Support |
|---|---|---|---|---|
{per_class_table}

![Confusion Matrix](assets/intent_confusion_matrix.png)

---

## 2. Predictive Maintenance Anomaly Detector (Isolation Forest)

Trained unsupervised on scaled sensor features ({', '.join(SENSOR_FEATURES)}),
evaluated against held-out ground-truth anomaly labels ({anomaly_eval.get('n_test', '?')} test readings,
{anomaly_eval.get('test_anomaly_rate', 0):.1%} true anomaly rate).

| Metric | Score |
|---|---|
| **Precision** | **{anomaly_eval.get('precision', 0):.2%}** |
| **Recall** | **{anomaly_eval.get('recall', 0):.2%}** |
| **F1 Score** | **{anomaly_eval.get('f1', 0):.2%}** |
| **ROC-AUC** | **{anomaly_eval.get('roc_auc', 0):.3f}** |
| Accuracy | {anomaly_eval.get('accuracy', 0):.2%} |

![Anomaly Confusion Matrix](assets/anomaly_confusion_matrix.png)
![ROC Curve](assets/anomaly_roc_curve.png)

---

## 3. Retrieval (RAG) Evaluation

Hit-rate of the TF-IDF retriever against a hand-labelled query -> expected-document
evaluation set (higher is better; hit@1 = correct doc ranked first, hit@2 = correct
doc within top 2).

| Metric | Score |
|---|---|
| **Hit Rate @ 1** | **{retrieval_eval.get('hit_rate_at_1', 0):.0%}** |
| **Hit Rate @ 2** | **{retrieval_eval.get('hit_rate_at_2', 0):.0%}** |

| Query | Expected Doc | Retrieved (top-1) | Hit | Score |
|---|---|---|---|---|
{retrieval_rows}

---

## 4. Latency Benchmark (per-request, CPU)

| Component | Avg. latency |
|---|---|
| Intent classification | {latency_eval.get('intent_classifier_ms', '?')} ms |
| Anomaly scoring | {latency_eval.get('anomaly_detector_ms', '?')} ms |
| KB retrieval (TF-IDF) | {latency_eval.get('kb_retrieval_ms', '?')} ms |
| LLM generation | Depends on hosted Inference API (measured live per-request in the Assistant tab) |

---

### Evaluation methodology notes

- All datasets are **synthetically generated** (see `src/data_generation.py`) using
  templated-but-varied natural language and randomised sensor distributions with a
  fixed seed, so results are fully reproducible via `python build_artifacts.py`.
- The intent classifier and anomaly detector are evaluated on a **held-out test
  split** they never saw during training (stratified, 25% / 30% respectively).
- The anomaly detector itself is trained **unsupervised** (Isolation Forest never
  sees the `label` column during `.fit()`); labels are used only to *evaluate* it,
  mirroring how you'd validate an anomaly model against a small set of confirmed
  historical incidents in production.
- In a production deployment, all three components would be continuously
  re-evaluated against real WMS/WCS/sensor logs rather than synthetic data.
"""
    return md


# ==========================================================================
# TAB 5 -- About
# ==========================================================================

ABOUT_MD = f"""
# 🏭 Smart Warehouse AI Assistant

**A portfolio project demonstrating an applied-AI approach to intralogistics
operations, built as part of a job application to Daifuku Co., Ltd.**

## What this demonstrates

Daifuku builds material handling and automation systems -- AS/RS, conveyors
and sortation, AGVs/AMRs, and the software (WMS/WCS) that orchestrates them.
This project is a compact but end-to-end example of how an AI layer can sit
on top of that kind of system:

| Capability | Where |
|---|---|
| **LLM-powered natural-language assistant**, grounded with retrieval (RAG) so it answers from real warehouse-ops knowledge rather than hallucinating | *AI Assistant* tab |
| **Intent classification** to route free-text requests (maintenance, safety, navigation, inventory, etc.) the way a real ops system would triage tickets | *AI Assistant* / *Inventory* tabs |
| **Predictive maintenance** via unsupervised anomaly detection on conveyor/crane sensor streams -- catching bearing wear or misalignment before an unplanned stoppage | *Predictive Maintenance* tab |
| **NL-to-structured-query** over inventory/order data, a lightweight stand-in for a WMS query tool | *Inventory & Order Query* tab |
| **Rigorous, reproducible evaluation** of every ML component (accuracy, F1, ROC-AUC, retrieval hit-rate, latency) rather than just a demo that "looks like it works" | *Model Evaluation* tab |

## Tech stack

- **UI / deployment:** [Gradio](https://gradio.app) on Hugging Face Spaces
- **LLM:** Hosted instruct model via the Hugging Face **Inference API**
  (`huggingface_hub.InferenceClient`), configurable via the `LLM_MODEL_ID`
  env var. Falls back gracefully to a retrieval-only answer if no `HF_TOKEN`
  is configured, so the public demo never breaks.
- **Retrieval:** TF-IDF + cosine similarity over a small hand-written
  warehouse-operations knowledge base (simple, fast, fully local RAG).
- **Intent classification:** TF-IDF + Logistic Regression (scikit-learn) --
  chosen deliberately over a heavier transformer classifier because it
  trains in under a second and comfortably reaches **{intent_eval.get('accuracy', 0):.0%}
  accuracy** on this task; right-sizing the model to the problem.
- **Anomaly detection:** Isolation Forest (scikit-learn), trained
  unsupervised on scaled sensor features.
- **Evaluation:** scikit-learn metrics + matplotlib, all computed by
  `build_artifacts.py` and saved as static artifacts the app loads at
  startup (fast, reproducible Space boot).

## Architecture

```
                    ┌───────────────────────┐
   User query  ───► │  Intent Classifier     │  (TF-IDF + LogisticRegression)
                    └──────────┬────────────┘
                               │ intent label
                               ▼
                    ┌───────────────────────┐
                    │  KB Retriever (RAG)    │  (TF-IDF cosine similarity)
                    └──────────┬────────────┘
                               │ top-k passages
                               ▼
                    ┌───────────────────────┐
                    │  Hosted LLM            │  (HF Inference API)
                    │  (or extractive        │
                    │   fallback if offline) │
                    └──────────┬────────────┘
                               ▼
                          Grounded answer

   Sensor stream ───► StandardScaler ───► IsolationForest ───► anomaly / normal
```

## Why this matters for Daifuku

Modern intralogistics platforms generate huge volumes of operational data --
equipment telemetry, WMS transactions, safety logs. The value of AI here isn't
a flashy chatbot; it's **routing, grounding, and reliability**: correctly
triaging a request, answering from real operational context instead of
guessing, and flagging equipment problems before they cause downtime. This
project tries to demonstrate that mindset in miniature, with honest,
reproducible evaluation numbers rather than cherry-picked demo runs.

## Limitations & next steps

- All data here is **synthetic**, for portfolio/demo purposes -- a production
  version would connect to real WMS/WCS APIs and historical sensor logs.
- The intent set (8 classes) and knowledge base (10 articles) are intentionally
  small to keep the demo fast and auditable; both are easy to extend.
- The anomaly detector uses 4 hand-picked features; a production system would
  likely use a richer multivariate sensor set and a supervised or
  semi-supervised model once labelled failure data is available.

---
*Built as an application project. Source code available on request / in the
linked repository. Feedback welcome.*
"""


# ==========================================================================
# GRADIO APP
# ==========================================================================

CUSTOM_CSS = """
#title-banner { text-align: center; margin-bottom: 0.5em; }
.gradio-container { max-width: 1150px !important; margin: auto; }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="Smart Warehouse AI Assistant") as demo:
    gr.Markdown(
        "<h1 id='title-banner'>🏭 Smart Warehouse AI Assistant</h1>"
        "<p style='text-align:center; color:gray;'>LLM-powered intralogistics copilot · "
        "intent routing · predictive maintenance · retrieval-grounded Q&A</p>"
    )
    if not HF_TOKEN_SET:
        gr.Markdown(
            "> ⚠️ **No `HF_TOKEN` secret detected.** The AI Assistant tab will run in "
            "**retrieval-only fallback mode** (still functional, just not LLM-generated "
            "prose). Add an `HF_TOKEN` secret in *Space settings → Variables and secrets* "
            "to enable full LLM responses."
        )

    with gr.Tab("💬 AI Assistant"):
        gr.Markdown(
            "Ask about equipment status, maintenance, safety, inventory, order status, "
            "AGV routing, picking strategy, or general warehouse-automation concepts. "
            "Answers are grounded (RAG) in a small warehouse-operations knowledge base."
        )
        gr.ChatInterface(
            fn=chat_fn,
            type="messages",
            examples=ASSISTANT_EXAMPLES,
            chatbot=gr.Chatbot(height=430, label="Warehouse Assistant", type="messages"),
            textbox=gr.Textbox(placeholder="e.g. The conveyor belt in Zone C is making noise"),
        )

    with gr.Tab("📦 Inventory & Order Query"):
        gr.Markdown(
            "Type a natural-language inventory or order question. The intent classifier "
            "decides whether to query the inventory table or the orders table, then "
            "extracts SKU / order-id / zone slots to filter the result."
        )
        with gr.Row():
            inv_input = gr.Textbox(label="Query", placeholder="How many units of SKU-1042 are in Zone B?", scale=4)
            inv_btn = gr.Button("Search", variant="primary", scale=1)
        inv_note = gr.Markdown()
        inv_result = gr.Dataframe(label="Results", wrap=True)
        inv_btn.click(inventory_query_fn, inputs=inv_input, outputs=[inv_note, inv_result])
        inv_input.submit(inventory_query_fn, inputs=inv_input, outputs=[inv_note, inv_result])

        gr.Examples(
            examples=[
                "How many units of SKU-1042 are in Zone B?",
                "What's the status of order #10007?",
                "Show me low stock items",
                "Any delayed orders?",
            ],
            inputs=inv_input,
        )
        with gr.Accordion("Browse full tables", open=False):
            gr.Markdown("**Inventory** (synthetic)")
            gr.Dataframe(inventory_df, wrap=True)
            gr.Markdown("**Orders** (synthetic)")
            gr.Dataframe(orders_df, wrap=True)

    with gr.Tab("⚠️ Predictive Maintenance"):
        gr.Markdown(
            "Enter live (or hypothetical) conveyor/crane motor sensor readings to check "
            "for anomalous behaviour using an Isolation Forest model trained on "
            "historical sensor patterns."
        )
        preset_dropdown = gr.Dropdown(
            choices=list(ANOMALY_PRESETS.keys()), label="Load a preset reading", value="Normal reading"
        )
        with gr.Row():
            motor_temp_in = gr.Slider(30, 100, value=55, step=0.5, label="Motor temperature (°C)")
            vibration_in = gr.Slider(0, 10, value=2.2, step=0.1, label="Vibration (mm/s)")
        with gr.Row():
            current_in = gr.Slider(5, 30, value=12, step=0.5, label="Motor current (A)")
            belt_speed_in = gr.Slider(0.1, 2.5, value=1.5, step=0.05, label="Belt speed (m/s)")
        check_btn = gr.Button("Check for anomaly", variant="primary")
        anomaly_output = gr.Markdown()

        preset_dropdown.change(
            load_preset, inputs=preset_dropdown,
            outputs=[motor_temp_in, vibration_in, current_in, belt_speed_in],
        )
        check_btn.click(
            anomaly_fn,
            inputs=[motor_temp_in, vibration_in, current_in, belt_speed_in],
            outputs=anomaly_output,
        )

    with gr.Tab("📊 Model Evaluation"):
        gr.Markdown(build_evaluation_markdown())

    with gr.Tab("ℹ️ About"):
        gr.Markdown(ABOUT_MD)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))