app.py

import streamlit as st
from rag_pipeline import load_model, load_retriever, get_response
import tempfile
import os


@st.cache_resource
def get_model():
    return load_model()


@st.cache_resource(show_spinner=False)
def get_retriever_cached(file_name: str, file_size: int, tmp_path: str):
    return load_retriever(tmp_path)


st.set_page_config(page_title="RAG CHATBOT", page_icon=":shark:")
st.title(":shark: RAG CHATBOT")

uploaded_file = st.file_uploader(label="Upload a PDF", type=["pdf"])

if uploaded_file:
    if st.session_state.get("current_file") != uploaded_file.name:
        st.session_state.chat_history = []
        st.session_state.current_file = uploaded_file.name

    if (
        "tmp_path" not in st.session_state
        or st.session_state.current_file != uploaded_file.name
    ):
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded_file.read())
            st.session_state.tmp_path = tmp.name

    with st.spinner("Processing PDF... ⏳"):
        retriever = get_retriever_cached(
            file_name=uploaded_file.name,
            file_size=uploaded_file.size,
            tmp_path=st.session_state.tmp_path,
        )

    st.success(f"Loaded: {uploaded_file.name}")

    for turn in st.session_state.chat_history:
        with st.chat_message("user"):
            st.write(turn["user"])
        with st.chat_message("assistant"):
            st.write(turn["assistant"])

    if question := st.chat_input("Ask something about the document...."):
        with st.chat_message("user"):
            st.write(question)

        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = get_response(
                    question=question,
                    retriever=retriever,
                    model=get_model(),
                    chat_history=st.session_state.chat_history,
                )
            st.write(response)

        st.session_state.chat_history.append({"user": question, "assistant": response})

else:
    if "tmp_path" in st.session_state and os.path.exists(st.session_state.tmp_path):
        os.unlink(st.session_state.tmp_path)
        del st.session_state.tmp_path

    st.info("Please Upload a PDF to start chatting")