streamlit_app.py

"""
streamlit_app.py
------------------
Streamlit front-end for the RAG project.

Wires together every pipeline stage (01 -> 07) into an interactive app:
    - Upload documents (or use whatever is already in data/)
    - Build/refresh the vector store
    - Ask a question
    - See the generated answer, its sources, and the raw retrieved chunks

API key handling: this file reads OPENROUTER_API_KEY / OPENROUTER_MODEL from
Streamlit secrets (when deployed on Streamlit Cloud) and injects them into
the "rag" module (07_prompting.py) at runtime. No real key is ever written
into any Python file.
"""

from __future__ import annotations

import importlib.util
import sys
import os
import time

import streamlit as st

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))


def _import_module_from_file(module_name: str, filename: str):
    """Import a numbered pipeline file (e.g. '07_prompting.py') as a module."""
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


# ---------------------------------------------------------------------------
# Load pipeline stage modules
# ---------------------------------------------------------------------------
documents_stage = _import_module_from_file("documents_stage", "01_documents.py")
preprocessing_stage = _import_module_from_file("preprocessing_stage", "02_preprocessing.py")
chunking_stage = _import_module_from_file("chunking_stage", "03_chunking.py")
store_stage = _import_module_from_file("store_stage", "05_create_chroma_store.py")
rag = _import_module_from_file("rag", "07_prompting.py")  # stage 6 is imported inside stage 7

# ---------------------------------------------------------------------------
# Read API secrets (Streamlit Cloud) without ever hardcoding the real key.
# Mirrors the pattern required by the project instructions.
# ---------------------------------------------------------------------------
try:
    if not rag.OPENROUTER_API_KEY:
        rag.OPENROUTER_API_KEY = st.secrets.get("OPENROUTER_API_KEY", "")
    rag.OPENROUTER_MODEL = st.secrets.get("OPENROUTER_MODEL", rag.OPENROUTER_MODEL)
except Exception:
    # st.secrets raises if no secrets.toml exists at all (e.g. local run
    # without secrets configured) -- fall back silently to env vars already
    # loaded inside 07_prompting.py.
    pass


# ---------------------------------------------------------------------------
# Streamlit page setup
# ---------------------------------------------------------------------------
st.set_page_config(
    page_title="RAG Assistant",
    page_icon="📚",
    layout="wide",
)

if "store_ready" not in st.session_state:
    st.session_state.store_ready = False
if "chunk_count" not in st.session_state:
    st.session_state.chunk_count = 0

st.title("📚 RAG Assistant")
st.caption(
    "Documents → Preprocessing → Chunking → Embeddings → Chroma → Retrieval → Prompting → Answer"
)

# ---------------------------------------------------------------------------
# Sidebar: knowledge base management
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("📂 Knowledge Base")

    uploaded_files = st.file_uploader(
        "Upload documents (.txt, .md, .pdf)",
        type=["txt", "md", "pdf"],
        accept_multiple_files=True,
    )

    st.divider()

    build_clicked = st.button("🔧 Build / Rebuild Vector Store", use_container_width=True)

    if build_clicked:
        with st.spinner("Running the pipeline: loading → cleaning → chunking → embedding → storing..."):
            try:
                all_raw_docs = documents_stage.load_documents()

                if uploaded_files:
                    all_raw_docs += documents_stage.load_documents_from_uploads(uploaded_files)

                if not all_raw_docs:
                    st.error(
                        "No documents found. Upload files above, or add some to the "
                        "local data/ folder, then try again."
                    )
                else:
                    cleaned_docs = preprocessing_stage.preprocess_documents(all_raw_docs)
                    chunks = chunking_stage.chunk_documents(cleaned_docs)
                    collection = store_stage.build_store(chunks, rebuild=True)

                    st.session_state.store_ready = True
                    st.session_state.chunk_count = collection.count()
                    st.success(
                        f"Vector store built successfully with "
                        f"{st.session_state.chunk_count} chunk(s) from "
                        f"{len(all_raw_docs)} document(s)."
                    )
            except Exception as exc:
                st.error(f"Failed to build the vector store: {exc}")

    if st.session_state.store_ready:
        st.info(f"✅ Store ready — {st.session_state.chunk_count} chunk(s) indexed.")
    else:
        st.warning("⚠️ Vector store not built yet. Click the button above first.")

    st.divider()
    top_k = st.slider("Number of chunks to retrieve (top_k)", min_value=1, max_value=10, value=4)

# ---------------------------------------------------------------------------
# Main area: question input + answer
# ---------------------------------------------------------------------------
question = st.text_input(
    "Ask a question about your documents:",
    placeholder="e.g. What are the main requirements described in the document?",
)

ask_clicked = st.button("🔍 Get Answer", type="primary")

if ask_clicked:
    if not question or not question.strip():
        st.error("Please enter a question first.")
    elif not st.session_state.store_ready:
        st.error("Please build the vector store first using the sidebar.")
    elif not rag.OPENROUTER_API_KEY:
        st.error(
            "No OpenRouter API key configured. Set OPENROUTER_API_KEY as an "
            "environment variable locally, or add it to Streamlit secrets when deployed."
        )
    else:
        with st.spinner("Retrieving relevant context and generating an answer..."):
            try:
                start_time = time.time()
                result = rag.answer_question(question, top_k=top_k)
                elapsed = time.time() - start_time
            except Exception as exc:
                st.error(f"Something went wrong while generating the answer: {exc}")
                result = None

        if result is not None:
            st.subheader("💬 Answer")
            st.write(result.answer)

            if result.used_context:
                st.caption(f"Answered in {elapsed:.1f}s using retrieved context.")

                st.subheader("📌 Sources")
                for source in result.sources:
                    st.markdown(f"- `{source}`")

                with st.expander("🔎 Retrieved chunks (raw context)"):
                    for i, chunk in enumerate(result.retrieved_chunks, start=1):
                        source = chunk.metadata.get("filename", chunk.metadata.get("doc_id", "unknown"))
                        st.markdown(f"**{i}. Source: `{source}` — distance: {chunk.distance:.4f}**")
                        st.text(chunk.text)
                        st.divider()
            else:
                st.warning(
                    "No relevant context was found in the knowledge base, so no "
                    "sources are shown. Try rephrasing your question or adding "
                    "more documents."
                )

st.divider()
st.caption(
    "Built with Python, ChromaDB, sentence-transformers, and OpenRouter — "
    "following the documents → preprocessing → chunking → vector representation → "
    "vector store → retrieval → prompting → Streamlit UI pipeline."
)