04_RAG_Chatbot_v2.py

import os
import time
import shutil
import streamlit as st
from langchain.chains import RetrievalQA
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.callbacks.manager import CallbackManager
from langchain_community.llms import Ollama
from langchain_community.embeddings.ollama import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory

# Define models as constants for easy configuration
EMBEDDING_MODEL = "nomic-embed-text:latest"
LLM_MODEL = "llama2:latest"
BASE_URL = "http://localhost:11434"

class PDFChatbot:
    def __init__(self):
        self.setup_directories()
        self.setup_session_state()
        self.display_title()

    def setup_directories(self):
        # Ensure that the necessary directories exist
        os.makedirs('files', exist_ok=True)
        os.makedirs('vector_database', exist_ok=True)

    def setup_session_state(self):
        # Initialize session state variables if they do not exist
        if 'template' not in st.session_state:
            st.session_state.template = """You are a knowledgeable chatbot, here to help with questions of the user. Your tone should be professional and informative.

            Context: {context}
            History: {history}

            User: {question}
            Chatbot:"""

        if 'prompt' not in st.session_state:
            st.session_state.prompt = PromptTemplate(
                input_variables=["history", "context", "question"],
                template=st.session_state.template,
            )

        if 'memory' not in st.session_state:
            st.session_state.memory = ConversationBufferMemory(
                memory_key="history",
                return_messages=True,
                input_key="question"
            )

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

        if 'llm' not in st.session_state:
            st.session_state.llm = Ollama(base_url=BASE_URL,
                                          model=LLM_MODEL,
                                          verbose=True,
                                          callback_manager=CallbackManager(
                                              [StreamingStdOutCallbackHandler()]),
                                          )

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

    def display_title(self):
        # Display the title of the Streamlit application
        st.title("PDF Chatbot")

    def clear_vectorstore(self):
        # Clear the vector database directory
        if os.path.exists('vector_database'):
            shutil.rmtree('vector_database')
            os.mkdir('vector_database')
        # Clear the session state variables
        st.session_state.vectorstore = None
        if 'retriever' in st.session_state:
            del st.session_state['retriever']
        if 'qa_chain' in st.session_state:
            del st.session_state['qa_chain']
        st.session_state.chat_history = []

    def handle_upload(self, uploaded_file):
        if uploaded_file is not None:
            # Only process the file if it's different from the current one
            if 'current_pdf' not in st.session_state or st.session_state.current_pdf != uploaded_file.name:
                self.clear_vectorstore()  # Clear existing vectorstore if a new file is uploaded
                st.session_state.current_pdf = uploaded_file.name

                st.info("Uploading and processing your PDF...")

                start_time = time.time()

                # Read the bytes of the uploaded file
                bytes_data = uploaded_file.read()
                file_path = f"files/{uploaded_file.name}"
                with open(file_path, "wb") as f:
                    f.write(bytes_data)  # Write the uploaded file to disk

                upload_time = time.time() - start_time
                st.success(f"PDF uploaded and saved in {upload_time:.2f} seconds.")

                st.info("Analyzing your document...")

                start_time = time.time()

                # Load and process the PDF
                loader = PyPDFLoader(file_path)
                data = loader.load()

                # Split the document into chunks
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1500,
                    chunk_overlap=200,
                    length_function=len
                )
                all_splits = text_splitter.split_documents(data)

                # Create a new vectorstore
                st.session_state.vectorstore = Chroma.from_documents(
                    documents=all_splits,
                    embedding=OllamaEmbeddings(base_url=BASE_URL, model=EMBEDDING_MODEL),
                    persist_directory='vector_database'
                )
                st.session_state.vectorstore.persist()  # Persist the vectorstore

                processing_time = time.time() - start_time
                st.success(f"Document analyzed and embeddings created in {processing_time:.2f} seconds.")

                # Set up the retriever and QA chain
                st.session_state.retriever = st.session_state.vectorstore.as_retriever()
                st.session_state.qa_chain = RetrievalQA.from_chain_type(
                    llm=st.session_state.llm,
                    chain_type='stuff',
                    retriever=st.session_state.retriever,
                    verbose=True,
                    chain_type_kwargs={
                        "verbose": True,
                        "prompt": st.session_state.prompt,
                        "memory": st.session_state.memory,
                    }
                )

    def display_chat_history(self):
        # Display the chat history in the Streamlit app
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.markdown(message["message"])

    def handle_user_input(self):
        # Handle user input from the chat interface
        user_input = st.chat_input("You:", key="user_input")
        if user_input:
            user_message = {"role": "user", "message": user_input}
            st.session_state.chat_history.append(user_message)  # Append user message to chat history
            with st.chat_message("user"):
                st.markdown(user_input)
            with st.chat_message("assistant"):
                with st.spinner("Assistant is typing..."):
                    start_time = time.time()
                    # Query the QA chain with user input
                    response = st.session_state.qa_chain(user_input)
                    response_time = time.time() - start_time

                message_placeholder = st.empty()
                full_response = ""
                words = response['result'].split()

                for word in words:
                    full_response += word + " "
                    message_placeholder.markdown(full_response + "▌")
                    time.sleep(0.05)  # Simulate typing effect
                message_placeholder.markdown(full_response)

            st.success(f"Response generated in {response_time:.2f} seconds.")

            chatbot_message = {"role": "assistant", "message": response['result']}
            st.session_state.chat_history.append(chatbot_message)  # Append chatbot response to chat history

    def run(self):
        # Main method to run the chatbot
        uploaded_file = st.file_uploader("Upload your PDF", type='pdf')
        self.handle_upload(uploaded_file)
        self.display_chat_history()
        self.handle_user_input()

if __name__ == "__main__":
    chatbot = PDFChatbot()
    chatbot.run()