streamlit_app.py

# File : streamlit_app.py
# Purpose : Streamlit frontend for Banking RAG Platform

import streamlit as st
import requests

# ─── Config ───────────────────────────────────────────────
API_URL = "http://localhost:8000"

st.set_page_config(
    page_title="Banking RAG Platform",
    page_icon="🏦",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ─── Session State ────────────────────────────────────────
if "token" not in st.session_state:
    st.session_state.token = None
if "username" not in st.session_state:
    st.session_state.username = None
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# ─── Helpers ──────────────────────────────────────────────
def auth_headers():
    return {"Authorization": f"Bearer {st.session_state.token}"}


def api_post(endpoint, payload=None, files=None, auth=False):
    headers = auth_headers() if auth else {}
    try:
        if files:
            r = requests.post(f"{API_URL}{endpoint}", headers=headers, files=files, timeout=120)
        else:
            r = requests.post(f"{API_URL}{endpoint}", headers=headers, json=payload, timeout=120)
        return r
    except requests.exceptions.ConnectionError:
        st.error("Cannot connect to backend. Please try again.")
        return None
    except requests.exceptions.Timeout:
        st.error("Request timed out. Please try again.")
        return None


def api_get(endpoint, auth=False):
    headers = auth_headers() if auth else {}
    try:
        r = requests.get(f"{API_URL}{endpoint}", headers=headers, timeout=30)
        return r
    except requests.exceptions.ConnectionError:
        st.error("Cannot connect to backend.")
        return None


# ─── Sidebar ──────────────────────────────────────────────
with st.sidebar:
    st.title("🏦 Banking RAG")
    st.markdown("---")

    if st.session_state.token:
        st.success(f"✅ Logged in as **{st.session_state.username}**")
        if st.button("Logout", use_container_width=True):
            st.session_state.token = None
            st.session_state.username = None
            st.session_state.chat_history = []
            st.rerun()
        st.markdown("---")
        page = st.radio(
            "Navigate",
            ["💬 Ask Documents", "📁 My Documents", "📤 Upload"],
            label_visibility="collapsed",
        )
    else:
        page = "🔐 Login"

    st.markdown("---")
    st.caption("Built with LangChain + Groq + ChromaDB")


# ─── Login / Register Page ────────────────────────────────
if not st.session_state.token:
    st.title("🏦 Banking RAG Platform")
    st.markdown("Enterprise AI assistant for RBI and Finance documents.")
    st.markdown("---")

    tab1, tab2 = st.tabs(["Login", "Register"])

    with tab1:
        st.subheader("Login")
        username = st.text_input("Username", key="login_user")
        password = st.text_input("Password", type="password", key="login_pass")
        if st.button("Login", use_container_width=True, type="primary"):
            if username and password:
                with st.spinner("Logging in..."):
                    r = api_post("/auth/login", {"username": username, "password": password})
                if r and r.status_code == 200:
                    data = r.json()
                    st.session_state.token = data["access_token"]
                    st.session_state.username = username
                    st.success(f"✅ Login Successful! Welcome {username}!")
                    st.balloons()
                    st.rerun()
                elif r:
                    st.error(f"❌ {r.json().get('detail', 'Login failed')}")
            else:
                st.warning("Enter username and password")

    with tab2:
        st.subheader("Register")
        reg_email = st.text_input("Email", key="reg_email")
        reg_username = st.text_input("Username", key="reg_user")
        reg_password = st.text_input("Password", type="password", key="reg_pass")
        if st.button("Register", use_container_width=True, type="primary"):
            if reg_email and reg_username and reg_password:
                with st.spinner("Creating account..."):
                    r = api_post("/auth/register", {
                        "email": reg_email,
                        "username": reg_username,
                        "password": reg_password,
                    })
                if r and r.status_code == 201:
                    st.success("✅ Registration Successful! Please go to Login tab.")
                    st.balloons()
                elif r:
                    st.error(f"❌ {r.json().get('detail', 'Registration failed')}")
            else:
                st.warning("Fill all fields")


# ─── Ask Documents Page ───────────────────────────────────
elif page == "💬 Ask Documents":
    st.title("💬 Ask Your Documents")
    st.markdown("Ask questions about your uploaded RBI and Finance documents.")
    st.markdown("---")

    # Chat history
    for msg in st.session_state.chat_history:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])
            if msg["role"] == "assistant" and "citations" in msg:
                if msg["citations"]:
                    with st.expander("📄 Sources"):
                        for i, cite in enumerate(msg["citations"], 1):
                            st.markdown(f"**[{i}]** Page {cite.get('page_number', 'N/A')} — {cite.get('content_preview', '')}")

    # Query input
    query = st.chat_input("Ask something about your documents...")
    if query:
        st.session_state.chat_history.append({"role": "user", "content": query})
        with st.chat_message("user"):
            st.markdown(query)

        with st.chat_message("assistant"):
            with st.spinner("Searching documents..."):
                r = api_post("/rag/query", {"query": query}, auth=True)
                if r and r.status_code == 200:
                    data = r.json()
                    answer = data.get("answer", "No answer found.")
                    citations = data.get("citations", [])
                    from_cache = data.get("from_cache", False)

                    st.markdown(answer)

                    if from_cache:
                        st.caption("⚡ From cache")

                    if citations:
                        with st.expander("📄 Sources"):
                            for i, cite in enumerate(citations, 1):
                                st.markdown(f"**[{i}]** Page {cite.get('page_number', 'N/A')} — {cite.get('content_preview', '')}")

                    st.session_state.chat_history.append({
                        "role": "assistant",
                        "content": answer,
                        "citations": citations,
                    })
                elif r:
                    err = r.json().get("detail", "Query failed")
                    st.error(f"❌ {err}")
                    st.session_state.chat_history.append({
                        "role": "assistant",
                        "content": f"Error: {err}",
                        "citations": [],
                    })

    if st.session_state.chat_history:
        if st.button("🗑️ Clear Chat"):
            st.session_state.chat_history = []
            st.rerun()


# ─── Upload Page ──────────────────────────────────────────
elif page == "📤 Upload":
    st.title("📤 Upload Documents")
    st.markdown("Upload PDF, DOCX, or TXT files to query with AI.")
    st.markdown("---")

    uploaded_file = st.file_uploader(
        "Choose a file",
        type=["pdf", "docx", "txt"],
        help="Max 50MB per file",
    )

    if uploaded_file:
        st.info(f"📄 Selected: **{uploaded_file.name}** ({round(uploaded_file.size / 1024, 1)} KB)")
        if st.button("Upload & Process", type="primary", use_container_width=True):
            with st.spinner("Uploading and processing document... this may take a minute."):
                files = {"file": (uploaded_file.name, uploaded_file.getvalue(), uploaded_file.type)}
                r = api_post("/documents/upload", files=files, auth=True)
                if r and r.status_code == 201:
                    data = r.json()
                    st.success("✅ Document uploaded and processed successfully!")
                    st.json({
                        "filename": data["filename"],
                        "total_chunks": data["total_chunks"],
                        "status": data["status"],
                    })
                elif r:
                    st.error(f"❌ {r.json().get('detail', 'Upload failed')}")


# ─── My Documents Page ────────────────────────────────────
elif page == "📁 My Documents":
    st.title("📁 My Documents")
    st.markdown("---")

    with st.spinner("Loading documents..."):
        r = api_get("/documents/list", auth=True)

    if r and r.status_code == 200:
        docs = r.json()
        if not docs:
            st.info("No documents uploaded yet. Go to Upload to add documents.")
        else:
            st.success(f"✅ {len(docs)} document(s) found")
            for doc in docs:
                with st.expander(f"📄 {doc['filename']}"):
                    col1, col2, col3 = st.columns(3)
                    col1.metric("Chunks", doc["total_chunks"])
                    col2.metric("Type", doc["file_type"].upper())
                    col3.metric("Status", doc["status"].capitalize())
                    st.caption(f"Uploaded: {doc['created_at']}")
    elif r:
        st.error("Failed to load documents")