app.py

import os
import streamlit as st
from dotenv import load_dotenv

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter 
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import ChatPromptTemplate

import tempfile


load_dotenv()

st.set_page_config(
    page_title="Chatbot with Memory + PDF Q&A",
    page_icon="🤖",
    layout="wide"
)

st.title("**Chatbot with Memory**")
st.write("Upload a PDF and ask questions. The chatbot remembers the conversation.")


if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

if "vectorstore" not in st.session_state:
    st.session_state.vectorstore = None

if "qa_chain" not in st.session_state:
    st.session_state.qa_chain = None


GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

if not GOOGLE_API_KEY:
    st.error("Missing GOOGLE_API_KEY. Add it inside your .env file.")
    st.stop()


llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    google_api_key=GOOGLE_API_KEY,
    temperature=0.3
)

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)


def process_pdf(uploaded_file):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        temp_file.write(uploaded_file.read())
        temp_path = temp_file.name

    loader = PyPDFLoader(temp_path)
    documents = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )

    chunks = text_splitter.split_documents(documents)

    vectorstore = FAISS.from_documents(
        documents=chunks,
        embedding=embeddings
    )

    os.remove(temp_path)

    return vectorstore


def create_qa_chain(vectorstore):
    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
    
    prompt = ChatPromptTemplate.from_template(
        """Answer the question based on the following context:

Context: {context}

Question: {question}

Answer:"""
    )
    
    def format_docs(docs):
        return "\n".join([doc.page_content for doc in docs])
    
    chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
    )
    
    return {"chain": chain, "retriever": retriever}


with st.sidebar:
    st.header("Upload PDF")

    uploaded_file = st.file_uploader(
        "Choose a PDF file",
        type=["pdf"]
    )

    if uploaded_file is not None:
        if st.button("Process PDF"):
            with st.spinner("Processing PDF..."):
                vectorstore = process_pdf(uploaded_file)
                st.session_state.vectorstore = vectorstore
                st.session_state.qa_chain = create_qa_chain(vectorstore)
                st.success("PDF processed successfully!")

    if st.button("Clear Chat"):
        st.session_state.chat_history = []
        st.session_state.qa_chain = None
        st.session_state.vectorstore = None
        st.success("Chat cleared.")


for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.write(message["content"])


user_question = st.chat_input("Ask something...")

if user_question:
    st.session_state.chat_history.append({
        "role": "user",
        "content": user_question
    })

    with st.chat_message("user"):
        st.write(user_question)

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):

            if st.session_state.qa_chain:
                chain = st.session_state.qa_chain["chain"]
                retriever = st.session_state.qa_chain["retriever"]
                
                result = chain.invoke(user_question)
                answer = result.content

                st.write(answer)

                with st.expander("Sources"):
                    docs = retriever.invoke(user_question)
                    for i, doc in enumerate(docs):
                        st.write(f"Source {i + 1}")
                        st.write(doc.page_content[:500])

            else:
                normal_response = llm.invoke(user_question)
                answer = normal_response.content
                st.write(answer)

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