streamlit_app_FastAPI.py

"""
streamlit_app_FastAPI.py — Insights Assistant UI (RAG + FastAPI + optional LLM + File Preview)

What this app does:
- Upload PDFs/CSVs and send their file paths to a FastAPI server for ingestion.
- Ask questions against a Chroma vector DB via the server's /tools/search_docs endpoint.
- Optionally summarize retrieved snippets using an LLM (Ollama/OpenAI) with citations.
- Preview the exact PDF page text or CSV row when the source file is available locally.

Prereqs:
- Place this file alongside: FastAPI_client.py and summarizer.py
- Run the server: uvicorn FastAPI_server:app --port 8799 --reload
- Start Streamlit: streamlit run streamlit_app_FastAPI.py
"""

from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())  # loads ${workspaceFolder}/.env by default

import os
import pandas as pd
import streamlit as st

# Local modules (pulls your local HTTP client for the FastAPI server -> FastAPI_client.py)
from FastAPI_client import (
    FastAPIClient,
    AuthError,
    RateLimitError,
    FastAPIHTTPError,
    FastAPIConnectionError,
)
from summarizer import summarize  # Provider-agnostic LLM helper (Ollama/OpenAI)


# ------------------------- Streamlit page config -------------------------
st.set_page_config(page_title="Insights Assistant", layout="wide")
st.title("Insights Assistant — RAG + FastAPI")


# ------------------------- Sidebar: Settings -------------------------
st.sidebar.header("Settings")

# URL for your FastAPI server exposing /tools endpoints
server_url = st.sidebar.text_input(
    "FastAPI Server URL",
    value=os.getenv("FASTAPI_SERVER_URL", "http://127.0.0.1:8799"),
    help="FastAPI server that exposes /tools/ingest_pdf, /tools/ingest_csv, /tools/search_docs",
)

# Optional Bearer token (if you later add auth to the server)
token = st.sidebar.text_input(
    "Bearer token (optional)",
    value=os.getenv("FASTAPI_TOKEN", ""),
    type="password",
)

# Retrieval top-k (how many passages to fetch per query)
top_k = st.sidebar.slider("Top-K passages", 1, 20, 5, 1)

st.sidebar.markdown("---")

# Optional LLM summarization switch + provider settings
use_llm = st.sidebar.checkbox("Use LLM summarizer", value=False)
llm_provider = st.sidebar.selectbox(
    "LLM provider", ["ollama", "openai"], index=0, help="Which backend to use for the summary."
)
ollama_url = st.sidebar.text_input(
    "Ollama URL",
    value=os.getenv("OLLAMA_URL", "http://127.0.0.1:11434"),
    help="Local Ollama endpoint.",
)
ollama_model = st.sidebar.text_input(
    "Ollama model",
    value=os.getenv("OLLAMA_MODEL", "llama3.1:8b"),
    help="e.g., llama3.1:8b, mistral, qwen2.5, etc.",
)
openai_model = st.sidebar.text_input(
    "OpenAI model",
    value=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
    help="Only used if provider is 'openai' and OPENAI_API_KEY is set.",
)

# Where uploaded files are stored so the server can read them by path
upload_dir = os.path.abspath(os.path.join(os.getcwd(), "uploads"))
os.makedirs(upload_dir, exist_ok=True)


def make_host() -> FastAPIClient:
    """
    Create a FastAPIClient with the sidebar settings.
    We also sync env vars so other local code using defaults stays consistent.
    """
    os.environ["FASTAPI_SERVER_URL"] = server_url
    if token:
        os.environ["FASTAPI_TOKEN"] = token
    return FastAPIClient(base_url=server_url, token=token)


def _coerce_hit_shape(hit: dict) -> dict:
    """
    Normalize server result shapes:
    - content: prefer 'content', else fall back to 'text'
    - source: prefer 'source', else convert 'metadata' to a similar shape
    - score: pass through if present
    """
    content = hit.get("content")
    if content is None:
        content = hit.get("text")

    # Source/metadata normalization
    src = hit.get("source")
    if src is None:
        meta = hit.get("metadata") or {}
        # Map common fields if present
        src = {
            "doc_type": meta.get("doc_type"),
            "page": meta.get("page"),
            "row": meta.get("row"),
            "source": meta.get("source") or meta.get("path"),
        }

    return {
        "content": content or "",
        "source": src or {},
        "score": hit.get("score", 0.0),
    }


def show_error(e: Exception) -> None:
    """
    Friendly error messages that map exception types (from FastAPI_client.py)
    to clear user-facing guidance.
    """
    if isinstance(e, AuthError):
        st.error("Authorization error. Re-authenticate or request the missing scope/role.")
    elif isinstance(e, RateLimitError):
        wait = getattr(e, "retry_after", None) or "a few"
        st.warning(f"Rate limited (429). Wait {wait} seconds and try again. Consider lowering Top-K.")
    elif isinstance(e, FastAPIConnectionError):
        st.error("Cannot reach the server. Is it running? `uvicorn FastAPI_server:app --port 8799 --reload`")
    elif isinstance(e, FastAPIHTTPError):
        st.error(f"HTTP error {e.status}. Check parameters or server logs.")
    else:
        st.error(f"Unexpected error: {e}")


# Quick connectivity check button
with st.sidebar:
    if st.button("Test connection"):
        try:
            _ = make_host().search_docs("ping", k=1)
            st.success("Server reachable ✅")
        except Exception as e:
            show_error(e)


# ------------------------- Cached helpers for previews -------------------------
@st.cache_data(show_spinner=False)
def _load_csv(path: str) -> pd.DataFrame:
    """Cached CSV loader for preview panel."""
    return pd.read_csv(path)


@st.cache_data(show_spinner=False)
def _pdf_page_text(path: str, page: int) -> str | None:
    """Extract a single PDF page's text (best-effort)."""
    try:
        from pypdf import PdfReader  # Lazy import
        reader = PdfReader(path)
        if page is None or page < 0 or page >= len(reader.pages):
            return None
        return (reader.pages[page].extract_text() or "").strip()
    except Exception:
        return None


# ------------------------- Main UI: Tabs -------------------------
tab_ingest, tab_ask = st.tabs(["📥 Ingest", "❓ Ask"])


# ------------------------- Ingest Tab -------------------------
with tab_ingest:
    st.subheader("Ingest documents into RAG")

    # Choose between PDF / CSV ingestion
    choice = st.radio("Type", ["PDF", "CSV"], horizontal=True)

    # --- PDF ingest ---
    if choice == "PDF":
        pdf = st.file_uploader("Upload PDF", type=["pdf"])
        doc_type = st.text_input("doc_type", value="pdf", help="Tag for downstream filtering (e.g., qna, whitepaper).")

        if st.button("Ingest PDF", disabled=pdf is None, type="primary"):
            if pdf:
                save_path = os.path.join(upload_dir, pdf.name)
                with open(save_path, "wb") as f:
                    f.write(pdf.getbuffer())
                try:
                    with st.spinner("Ingesting PDF…"):
                        res = make_host().ingest_pdf(save_path, doc_type=doc_type)
                    st.success(res)
                except Exception as e:
                    show_error(e)

    # --- CSV ingest ---
    if choice == "CSV":
        csv = st.file_uploader("Upload CSV", type=["csv"])
        text_col = st.text_input(
            "Text column (optional)",
            help="If empty, the loader may concatenate row cells into a single text chunk.",
        )
        doc_type = st.text_input("doc_type", value="csv", key="csv_doc_type")

        if st.button("Ingest CSV", disabled=csv is None, type="primary"):
            if csv:
                save_path = os.path.join(upload_dir, csv.name)
                with open(save_path, "wb") as f:
                    f.write(csv.getbuffer())
                try:
                    with st.spinner("Ingesting CSV…"):
                        res = make_host().ingest_csv(save_path, text_col=text_col or None, doc_type=doc_type)
                    st.success(res)
                except Exception as e:
                    show_error(e)


# ------------------------- Ask Tab -------------------------
with tab_ask:
    st.subheader("Ask a question")
    query = st.text_input(
        "Your question",
        placeholder="e.g., What does the sample say about market growth?",
        help="Query is sent to /tools/search_docs; results are summarized locally (optional LLM).",
    )

    # Trigger the search action
    if st.button("Search", type="primary") and query.strip():
        try:
            # 1) Fetch results from the server
            with st.spinner("Searching…"):
                res = make_host().search_docs(query.strip(), k=top_k)

            # Normalize response schema across server variants—prefer `results` (MCP-style), fallback to `hits` (legacy FastAPI).
            raw_hits = res.get("results") or res.get("hits") or []
            hits = [_coerce_hit_shape(h) for h in raw_hits]

            if not hits:
                st.info("No matches found. Ingest more docs or broaden your query.")
            else:
                # 2) Build snippets + citations for optional LLM summarization
                snippets = [h["content"] for h in hits]
                citations = [str(h["source"] or {}) for h in hits]

                # 3) Try to generate LLM summary if enabled; otherwise show extractive
                answer = None
                if use_llm:
                    with st.spinner("Generating LLM summary…"):
                        answer = summarize(
                            query.strip(),
                            snippets,
                            citations,
                            provider=llm_provider,
                            ollama_url=ollama_url,
                            ollama_model=ollama_model,
                            openai_model=openai_model,
                        )

                if answer:
                    st.markdown("**Answer (LLM, cited):**")
                    st.write(answer)
                else:
                    # Extractive fallback
                    joined = " ".join(snippets)[:1500]
                    preview = (joined[:800] + " …") if len(joined) > 800 else joined
                    st.markdown("**Answer (extractive):**")
                    st.write(preview)

                # 4) Show result rows in a grid (with indices for selection)
                rows = []
                for i, h in enumerate(hits):
                    src = h.get("source", {}) or {}
                    rows.append({
                        "idx": i,  # keep an index so we can reference the original hit
                        "score": round(h.get("score", 0.0), 3),
                        "content": (h.get("content", "")[:200].replace("\n", " ")
                                    + (" …" if len(h.get("content", "")) > 200 else "")),
                        "doc_type": src.get("doc_type"),
                        "page": src.get("page"),
                        "row": src.get("row"),
                        "source_path": src.get("source"),
                    })
                df = pd.DataFrame(rows)
                st.dataframe(df, use_container_width=True, hide_index=True)

                # 5) Preview panel to inspect the selected hit's original content (PDF page / CSV row)
                st.markdown("### Preview selected hit")

                # Build user-friendly labels: "<idx> — <doc_type> p.<page> row <row>"
                options = [
                    f"{r['idx']} — {r['doc_type'] or 'unknown'}"
                    + (f" p.{int(r['page'])}" if pd.notna(r['page']) and r['page'] is not None else "")
                    + (f" row {int(r['row'])}" if pd.notna(r['row']) and r['row'] is not None else "")
                    for _, r in df.iterrows()
                ]
                default_idx = 0
                selected_label = st.selectbox("Select a hit", options, index=default_idx if options else 0)

                if options:
                    # Parse the selected index (string before the " — ")
                    sel_idx = int(selected_label.split(" — ")[0])
                    hit = hits[sel_idx]
                    src = hit.get("source", {}) or {}
                    path = src.get("source")
                    doc_type = src.get("doc_type")
                    page = src.get("page")
                    row_idx = src.get("row")

                    st.write(f"**Source path:** `{path}`")
                    st.write(f"**Type:** `{doc_type}`  |  **Page:** `{page}`  |  **Row:** `{row_idx}`")

                    # We can preview only if the server-ingested file is accessible to this UI process
                    if path and os.path.exists(path):
                        # --- Preview PDF page text ---
                        if (str(path).lower().endswith(".pdf") or (doc_type == "pdf")) and page is not None:
                            try:
                                page_int = int(page)
                            except Exception:
                                page_int = None

                            text = _pdf_page_text(path, page_int) if page_int is not None else None
                            if text:
                                st.markdown("**PDF page text:**")
                                st.code(text[:4000], language="markdown")
                            else:
                                st.info("Could not extract PDF page text (scanned PDF or no text on page).")

                        # --- Preview CSV row ---
                        elif (str(path).lower().endswith(".csv") or (doc_type == "csv")) and row_idx is not None:
                            try:
                                row_int = int(row_idx)
                            except Exception:
                                row_int = None

                            if row_int is None:
                                st.info("Row index not an integer; showing snippet instead.")
                                st.write(hit.get("content", "")[:800])
                            else:
                                try:
                                    df_src = _load_csv(path)
                                    if 0 <= row_int < len(df_src):
                                        st.markdown("**CSV row preview:**")
                                        st.dataframe(
                                            df_src.iloc[[row_int]],
                                            use_container_width=True,
                                            hide_index=True,
                                        )
                                    else:
                                        st.info("Row index out of range in CSV; showing snippet instead.")
                                        st.write(hit.get("content", "")[:800])
                                except Exception as e:
                                    st.warning(f"Could not load CSV for preview: {e}")
                        else:
                            # Unknown file type or insufficient metadata for precise preview
                            st.info("Preview available only for PDF (page text) and CSV (row). Showing snippet instead.")
                            st.write(hit.get("content", "")[:800])
                    else:
                        # Path might be remote (e.g., server ingested from a different machine)
                        st.info("Source file not accessible from this machine. Showing snippet instead.")
                        st.write(hit.get("content", "")[:800])

        except Exception as e:
            show_error(e)