app.py

"""
Project: AskMe-FAQ-Bot
Author: Devdatta Supnekar
Date: 2025-09-21
Notes:

Features
--------
- Clean, readable chat UI (each turn groups user → assistant)
- Avatars, timestamps, and "Clear chat" button
- Model Management tab: upload CSV and retrain (merge) or train new (replace)
- Visible progress, status messages, and index stats
- Fallback CSV download
- Centralized logging via logger.py

Run:
    streamlit run streamlit_app.py
"""

from __future__ import annotations

import os
import time
from datetime import datetime
from typing import Optional, List, Dict

import pandas as pd
import streamlit as st

from logger import setup_logging, get_logger
from config import (
    EMBED_MODEL, GENERATE_MODEL, THRESHOLD, HYBRID_ALPHA, USE_SPELLCHECK,
    FALLBACK_PATH,
)

# Initialize logging early; utility import will use this logger
setup_logging()
log = get_logger()

# Import after logging is set up
from utility import (
    store,
    load_index_if_exists,
    build_index_from_df,
    load_base_csv,
    merge_and_rebuild,
    replace_and_rebuild,
    answer_query,
)


# ----------------------------- Helpers & State -----------------------------

def _now_str() -> str:
    """Return a compact UTC timestamp for message captions."""
    return datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")


def _init_state() -> None:
    """Initialize Streamlit session state for chat and readiness flags."""
    if "chat" not in st.session_state:
        # [{'role': 'user'|'assistant', 'content': str}, ...]
        st.session_state.chat: List[Dict[str, str]] = []
    if "ready" not in st.session_state:
        st.session_state.ready = False


# ---------------------------- Startup & Indexing ---------------------------

def ensure_index_on_start() -> None:
    """
    Ensure an index is available in-memory.

    Load order:
      1) Try to load an existing FAISS index from disk.
      2) If not found, and base CSV exists, build a fresh index.
      3) If neither exists, show a friendly error.

    This avoids unnecessary re-indexing when an index already exists.
    """
    try:
        # 1) Prefer loading an existing index from disk
        load_index_if_exists()
        if store.index is not None and store.meta is not None:
            st.session_state.ready = True
            log.info(
                "[Startup] Loaded existing index. d=%s, items=%s",
                store.index.d, store.index.ntotal
            )
            return

        # 2) If no index, but base CSV exists, build it once
        base = load_base_csv()
        if base is not None and not base.empty:
            log.info("[Startup] No on-disk index. Building from base CSV (rows=%d)…", len(base))
            with st.status("Building initial index from base CSV…", expanded=True) as s:
                s.write("Embedding and indexing questions…")
                build_index_from_df(base)
                s.update(label="Initial index built.", state="complete")
            st.session_state.ready = True
            return

        # 3) Nothing to load or build
        log.warning("[Startup] No index on disk and no base CSV. UI will run without retrieval.")
        st.session_state.ready = True

    except Exception as e:
        log.exception("Startup indexing/load failed: %s", e)
        st.error(f"Startup failed: {e}")



def index_stats_text() -> str:
    """
    Compose a short markdown string describing FAISS index status.

    Returns
    -------
    str : Markdown showing FAISS dimensionality and item count, or a friendly
          message if not available.
    """
    try:
        if store.index is not None and store.meta is not None:
            return f"**Index:** d={store.index.d}, items={store.index.ntotal}"
        return "**Index:** not loaded"
    except Exception as e:
        log.exception("Failed to compute index stats: %s", e)
        return "**Index:** error (see logs)"


# -------------------------------- Sidebar ---------------------------------

def sidebar_info() -> None:
    """
    Render the sidebar with configuration, index status, and fallback download.

    Notes
    -----
    - Uses `st.download_button` to provide `fallback.csv` when present.
    - Provides a quick "Refresh Index Stats" button to re-render the sidebar.
    """
    with st.sidebar:
        st.header("Settings", anchor=False)
        col1, col2 = st.columns(2)
        with col1:
            st.caption("Embed model")
            st.code(EMBED_MODEL, language=None)
        with col2:
            st.caption("Generate model")
            st.code(GENERATE_MODEL, language=None)

        c1, c2, c3 = st.columns(3)
        c1.metric("Threshold", f"{THRESHOLD}")
        c2.metric("α (hybrid)", f"{HYBRID_ALPHA}")
        c3.metric("Spell", "On" if USE_SPELLCHECK else "Off")

        st.divider()
        st.subheader("Index", anchor=False)
        st.markdown(index_stats_text())
        if st.button("Refresh Index Stats", use_container_width=True):
            log.info("[UI] Index stats refresh requested.")
            st.rerun()

        st.divider()
        st.subheader("Fallbacks", anchor=False)
        if os.path.exists(FALLBACK_PATH):
            try:
                with open(FALLBACK_PATH, "rb") as f:
                    st.download_button(
                        label="Download fallback.csv",
                        data=f.read(),
                        file_name="fallback.csv",
                        mime="text/csv",
                        use_container_width=True,
                    )
            except Exception as e:
                log.exception("Failed to open fallback.csv for download: %s", e)
                st.caption("Could not open fallback.csv (see logs).")
        else:
            st.caption("No fallback.csv yet.")


# --------------------------------- Chat -----------------------------------

def chat_panel() -> None:
    """
    Render a clean, turn-based chat interface (default theme).

    Layout
    ------
    - Each turn shows the user message followed by the assistant reply.
    - Avatars: 🧑‍💻 (user), 🤖 (assistant)
    - Subtle `st.divider()` between turns for readability.

    Logging
    -------
    - Logs the incoming user query (truncated) and the answer length.
    - Logs the latency for each turn in ms.
    """
    st.subheader("Chatbot Interface", anchor=False)

    # --- Render existing history as turns (user -> assistant) ---
    msgs = st.session_state.chat
    i = 0
    while i < len(msgs):
        # user message
        if i < len(msgs) and msgs[i]["role"] == "user":
            with st.chat_message("user", avatar="🧑‍💻"):
                st.markdown(msgs[i]["content"])
                st.caption(f"sent • {_now_str()}")
        # assistant reply (if present)
        if i + 1 < len(msgs) and msgs[i + 1]["role"] == "assistant":
            with st.chat_message("assistant", avatar="🤖"):
                st.markdown(msgs[i + 1]["content"])
                st.caption(f"answered • {_now_str()}")
            st.divider()
            i += 2
        else:
            # if the stored history is unpaired, move forward safely
            i += 1

    # --- Input box at the bottom ---
    prompt = st.chat_input("Type your question…")
    if not prompt:
        return

    # Append user message, render immediately
    st.session_state.chat.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar="🧑‍💻"):
        st.markdown(prompt)
        st.caption(f"sent • {_now_str()}")

    # Generate assistant answer
    with st.chat_message("assistant", avatar="🤖"):
        with st.spinner("Thinking…"):
            t0 = time.perf_counter()
            try:
                log.info("Chat query received: '%s'", (prompt[:120] + "…") if len(prompt) > 120 else prompt)
                answer = answer_query(prompt)
                log.info("Chat answer produced (len=%d chars)", len(answer))
            except Exception as e:
                log.exception("Error in chat answer for query='%s': %s", prompt, e)
                answer = "Sorry, something went wrong while processing your question."
            latency_ms = (time.perf_counter() - t0) * 1000.0
            log.info("Chat turn latency: %.1f ms", latency_ms)
            st.markdown(answer)
            st.caption(f"answered • {_now_str()}")

    st.session_state.chat.append({"role": "assistant", "content": answer})
    st.divider()


# ------------------------------ Management -------------------------------

def management_panel() -> None:
    """
    Render the Model Management UI (upload, train, index status).

    Behavior
    --------
    - Validates schema (question_id, question, answer).
    - For "Retrain (merge with base)" → merges with current base, then rebuilds index.
    - For "Train New (ignore base)" → replaces base with uploaded file, then rebuilds index.
    - Shows progress and status messages and logs each step.
    """
    st.subheader("Model Management", anchor=False)
    st.caption("Upload a CSV with columns: **question_id, question, answer**")

    up = st.file_uploader("Upload CSV", type=["csv"], accept_multiple_files=False)
    mode = st.radio(
        "Training Mode",
        ["Retrain (merge with base)", "Train New (ignore base)"],
        index=0,
        horizontal=True,
    )
    run = st.button("Run Training", type="primary", use_container_width=True)

    if not run:
        return
    if not up:
        st.warning("Please upload a CSV first.")
        return

    # Read uploaded CSV
    try:
        df = pd.read_csv(up)
        log.info("[Train] Uploaded CSV loaded: rows=%d, cols=%d", df.shape[0], df.shape[1])
    except Exception as e:
        log.exception("Could not read uploaded CSV: %s", e)
        st.error(f"Could not read the uploaded CSV: {e}")
        return

    # Validate schema
    df.columns = [c.strip().lower() for c in df.columns]
    required = {"question_id", "question", "answer"}
    if not required.issubset(set(df.columns)):
        st.error(f"CSV must include columns: {required}. Received: {set(df.columns)}")
        log.warning("[Train] Invalid CSV schema: %s", set(df.columns))
        return

    # Progress
    progress = st.progress(0.0, text="Preparing…")
    status = st.empty()

    try:
        progress.progress(0.15, text="Validating & preparing data…")
        status.info("CSV loaded. Preparing data…")

        if mode.startswith("Retrain"):
            status.info("Merging with existing base…")
            progress.progress(0.35, text="Merging with base…")
            log.info("[Train] Merging with base + rebuilding index…")
            merge_and_rebuild(df)
        else:
            status.info("Replacing base with uploaded data…")
            progress.progress(0.35, text="Replacing base…")
            log.info("[Train] Replacing base + rebuilding index…")
            replace_and_rebuild(df)

        progress.progress(0.95, text="Finalizing…")
        status.success("Training complete. Index rebuilt.")
        progress.progress(1.0, text="Done")

        st.success(index_stats_text())
        log.info("[Train] Training flow completed successfully.")
    except Exception as e:
        log.exception("[Train] Exception during training: %s", e)
        status.error(f"Training failed: {e}")


# ----------------------------------- Main ---------------------------------

def main() -> None:
    """
    Render the Streamlit app using the default theme.

    Adds a small "Clear chat" button near the title to reset the session
    chat history cleanly.
    """
    st.set_page_config(page_title="AskMe-FAQ-Bot", page_icon="💬", layout="centered")
    _init_state()

    # Header + Clear chat button
    col_title, col_btn = st.columns([1, 0.18])
    with col_title:
        st.title("💬 AskMe-FAQ-Bot")
        st.caption("Professional FAQ chatbot with semantic + lexical matching and admin retraining panel.")
    with col_btn:
        if st.button("Clear chat", help="Remove chat history for this session"):
            log.info("[UI] Clear chat requested.")
            st.session_state.chat = []
            st.rerun()

    st.divider()
    sidebar_info()

    if not st.session_state.ready:
        ensure_index_on_start()

    tab_chat, tab_mgmt = st.tabs(["Chat", "Model Management"])
    with tab_chat:
        chat_panel()
    with tab_mgmt:
        management_panel()


if __name__ == "__main__":
    log.info(
        "Launching Streamlit UI | Embed=%s | Generate=%s | TH=%.2f | ALPHA=%.2f | Spell=%s",
        EMBED_MODEL, GENERATE_MODEL, THRESHOLD, HYBRID_ALPHA, USE_SPELLCHECK
    )
    main()