feat: refactoring and optimizations (#1)

.streamlit/config.toml

@@ -1,4 +1,3 @@
 [theme]
 base="light"
-backgroundColor="#FFF1F9"
-secondaryBackgroundColor="#FFDCF1"
+primaryColor="#0098ff"

main.py

@@ -1,249 +1,123 @@
 import os
-import pickle
 import streamlit as st
-import tempfile
 import asyncio
 from dotenv import load_dotenv
 
-# Import modules needed for building the chatbot application
 from streamlit_chat import message
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.chat_models import ChatOpenAI
-from langchain.chains import ConversationalRetrievalChain
-from langchain.document_loaders import PyPDFLoader
-from langchain.vectorstores import FAISS
 
-# Load the environment variables from the .env file
-load_dotenv()
+from modules.chatbot import Chatbot
+from modules.embedder import Embedder
+from modules.history import ChatHistory
+from modules.layout import Layout
+from modules.sidebar import Sidebar
 
-# Set the Streamlit page configuration, including the layout and page title/icon
-st.set_page_config(layout="wide", page_icon="💬", page_title="ChatBot-PDF")
 
-# Display the header for the application using HTML markdown
-st.markdown("<h1 style='text-align: center;'>ChatBot-PDF, Talk with your documents ! 💬</h1>", unsafe_allow_html=True)
+load_dotenv()
 
-# Get the OpenAI API key from an environment variable if present
-user_api_key = os.getenv("OPENAI_API_KEY")
+st.set_page_config(layout="wide", page_icon="💬", page_title="ChatBot-PDF")
 
-# Allow the user to enter their OpenAI API key if it's not present in the environment variables
-if not user_api_key:
-    user_api_key = st.sidebar.text_input(
-        label="#### Your OpenAI API key 👇", placeholder="Paste your openAI API key, sk-", type="password"
-    )
-else:
-    st.sidebar.success("API key loaded from .env", icon="🚀")
+# Load the OpenAI API key from the .env file or from the user's input
+def load_api_key():
+    user_api_key = os.getenv("OPENAI_API_KEY")
+    if not user_api_key:
+        user_api_key = st.sidebar.text_input(
+            label="#### Your OpenAI API key 👇", placeholder="Paste your openAI API key, sk-", type="password"
+        )
+    else:
+        st.sidebar.success("API key loaded from .env", icon="🚀")
+    return user_api_key
 
 
-async def main():
-    # Check if the user has entered an OpenAI API key
-    if user_api_key == "":
-        # Display a message asking the user to enter their API key
-        st.markdown(
-            "<div style='text-align: center;'><h4>Enter your OpenAI API key to start chatting 😉</h4></div>",
-            unsafe_allow_html=True,
+# Handle the file upload and display the uploaded file
+def handle_upload():
+    uploaded_file = st.sidebar.file_uploader("upload", type="pdf", label_visibility="collapsed")
+    if uploaded_file is not None:
+        file_container = st.expander("Your PDF file :")
+        file_container.write(uploaded_file)
+    else:
+        st.sidebar.info(
+            "👆 Upload your PDF file to get started, "
+            "sample for try : [file.pdf](https://github.com/gabacode/chatPDF/blob/main/file.pdf)"
         )
+        st.session_state["reset_chat"] = True
+    return uploaded_file
 
-    else:
-        # Set the OpenAI API key as an environment variable
-        os.environ["OPENAI_API_KEY"] = user_api_key
 
-        # Allow the user to upload a file
-        uploaded_file = st.sidebar.file_uploader("upload", type="pdf", label_visibility="collapsed")
+# Set up the chatbot with the uploaded file, model, and temperature
+async def setup_chatbot(uploaded_file, model, temperature):
+    embeds = Embedder()
+    with st.spinner("Processing..."):
+        uploaded_file.seek(0)
+        file = uploaded_file.read()
+        vectors = await embeds.getDocEmbeds(file, uploaded_file.name)
+        chatbot = Chatbot(model, temperature, vectors)
+    st.session_state["ready"] = True
+    return chatbot
 
-        # If the user has uploaded a file, display it in an expander
-        if uploaded_file is not None:
 
-            def show_user_file(uploaded_file):
-                file_container = st.expander("Your PDF file :")
-                file_container.write(uploaded_file)
+async def main():
 
-            show_user_file(uploaded_file)
+    layout = Layout()
+    sidebar = Sidebar()
 
-        # If the user has not uploaded a file, display a message asking them to do so
-        else:
-            st.sidebar.info(
-                "👆 Upload your PDF file to get started, "
-                "sample for try : [file.pdf](https://github.com/gabacode/chatPDF/blob/main/file.pdf)"
-            )
+    layout.display_header()
+    user_api_key = load_api_key()
 
-        if uploaded_file:
-            try:
-                # Define an asynchronous function for storing document embeddings using Langchain and FAISS
-                async def storeDocEmbeds(file, filename):
-                    # Write the uploaded file to a temporary file
-                    with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tmp_file:
-                        tmp_file.write(file)
-                        tmp_file_path = tmp_file.name
-
-                    # Load the data from the file using Langchain
-                    loader = PyPDFLoader(file_path=tmp_file_path)
-                    data = loader.load_and_split()
-
-                    # Create an embeddings object using Langchain
-                    embeddings = OpenAIEmbeddings()
-
-                    # Store the embeddings vectors using FAISS
-                    vectors = FAISS.from_documents(data, embeddings)
-                    os.remove(tmp_file_path)
-
-                    # Save the vectors to a pickle file
-                    with open(filename + ".pkl", "wb") as f:
-                        pickle.dump(vectors, f)
-
-                # Define an asynchronous function for retrieving document embeddings
-                async def getDocEmbeds(file, filename):
-                    # Check if embeddings vectors have already been stored in a pickle file
-                    if not os.path.isfile(filename + ".pkl"):
-                        # If not, store the vectors using the storeDocEmbeds function
-                        await storeDocEmbeds(file, filename)
-
-                    # Load the vectors from the pickle file
-                    with open(filename + ".pkl", "rb") as f:
-                        global vectors
-                        vectors = pickle.load(f)
-
-                    return vectors
-
-                # Define an asynchronous function for conducting conversational chat using Langchain
-                async def conversational_chat(query):
-                    # Use the Langchain ConversationalRetrievalChain to generate a response to the user's query
-                    result = chain({"question": query, "chat_history": st.session_state["history"]})
-
-                    # Add the user's query and the chatbot's response to the chat history
-                    st.session_state["history"].append((query, result["answer"]))
-
-                    # Print the chat history for debugging purposes
-                    print("Log: ")
-                    print(st.session_state["history"])
-
-                    return result["answer"]
-
-                # Set up sidebar with various options
-                with st.sidebar.expander("🛠️ Settings", expanded=False):
-                    # Add a button to reset the chat history
-                    if st.button("Reset Chat"):
-                        st.session_state["reset_chat"] = True
-
-                    # Allow the user to select a chatbot model to use
-                    MODEL = st.selectbox(label="Model", options=["gpt-3.5-turbo"])
-
-                    # Allow the user to change the model temperature
-                    TEMPERATURE = st.slider(label="Temperature", min_value=0.0, max_value=1.0, value=0.618, step=0.01)
-
-                # If the chat history has not yet been initialized, do so now
-                if "history" not in st.session_state:
-                    st.session_state["history"] = []
-
-                # If the chatbot is not yet ready to chat, set the "ready" flag to False
-                if "ready" not in st.session_state:
-                    st.session_state["ready"] = False
-
-                # If the "reset_chat" flag has not been set, set it to False
-                if "reset_chat" not in st.session_state:
-                    st.session_state["reset_chat"] = False
-
-                    # If a PDF file has been uploaded
-                if uploaded_file is not None:
-                    # Display a spinner while processing the file
-                    with st.spinner("Processing..."):
-                        # Read the uploaded PDF file
-                        uploaded_file.seek(0)
-                        file = uploaded_file.read()
-
-                        # Generate embeddings vectors for the file
-                        vectors = await getDocEmbeds(file, uploaded_file.name)
-
-                        # Use the Langchain ConversationalRetrievalChain to set up the chatbot
-                        chain = ConversationalRetrievalChain.from_llm(
-                            llm=ChatOpenAI(temperature=TEMPERATURE, model_name=MODEL), retriever=vectors.as_retriever()
-                        )
-
-                    # Set the "ready" flag to True now that the chatbot is ready to chat
-                    st.session_state["ready"] = True
-
-                # If the chatbot is ready to chat
-                if st.session_state["ready"]:
-                    # If the chat history has not yet been initialized, initialize it now
-                    if "generated" not in st.session_state:
-                        st.session_state["generated"] = [
-                            "Hello ! Ask me anything about the document " + uploaded_file.name + " 🤗"
-                        ]
+    if user_api_key == "":
+        layout.show_api_key_error()
+    else:
+        os.environ["OPENAI_API_KEY"] = user_api_key
+        uploaded_file = handle_upload()
 
-                    if "past" not in st.session_state:
-                        st.session_state["past"] = ["Hey ! 👋"]
+        if uploaded_file is not None:
+            history = ChatHistory()
+            sidebar.options()
+            try:
+                chatbot = await setup_chatbot(uploaded_file, st.session_state["model"], st.session_state["temperature"])
+                st.session_state["chatbot"] = chatbot
 
-                    # Create a container for displaying the chat history
+                if st.session_state["ready"]:
+                    # Create a containers for displaying the chat history
                     response_container = st.container()
-
-                    # Create a container for the user's text input
                     container = st.container()
 
                     with container:
-                        # Create a form for the user to enter their query
                         with st.form(key="my_form", clear_on_submit=True):
                             user_input = st.text_area(
                                 "Query:",
-                                placeholder="Talk about your data here (:",
+                                placeholder="Ask me anything about the document...",
                                 key="input",
                                 label_visibility="collapsed",
                             )
                             submit_button = st.form_submit_button(label="Send")
 
-                            # If the "reset_chat" flag has been set, reset the chat history and generated messages
-                            if st.session_state["reset_chat"]:
-                                st.session_state["history"] = []
-                                st.session_state["past"] = ["Hey ! 👋"]
-                                st.session_state["generated"] = [
-                                    "Hello ! Ask me anything about the document " + uploaded_file.name + " 🤗"
-                                ]
-                                response_container.empty()
-                                st.session_state["reset_chat"] = False
+                        if st.session_state["reset_chat"]:
+                            history.reset(uploaded_file)
+
+                        history.initialize(uploaded_file)
 
                         # If the user has submitted a query
                         if submit_button and user_input:
-                            # Add the user's input to the chat history
-                            st.session_state["past"].append(user_input)
-
-                            # Generate a response using the Langchain ConversationalRetrievalChain
-                            output = await conversational_chat(user_input)
-
-                            # Add the user's chatbot's output to the chat history
-                            st.session_state["generated"].append(output)
+                            history.append("user", user_input)
+                            output = await st.session_state["chatbot"].conversational_chat(user_input)
+                            history.append("assistant", output)
 
                     # If there are generated messages to display
-                    if st.session_state["generated"]:
-                        # Display the chat history
+                    if st.session_state["assistant"]:
                         with response_container:
-                            for i in range(len(st.session_state["generated"])):
+                            for i in range(len(st.session_state["assistant"])):
                                 message(
-                                    st.session_state["past"][i],
+                                    st.session_state["user"][i],
                                     is_user=True,
-                                    key=str(i) + "_user",
+                                    key=f"{i}_user",
                                     avatar_style="big-smile",
                                 )
-                                message(st.session_state["generated"][i], key=str(i), avatar_style="thumbs")
-                # st.write(chain)
+                                message(st.session_state["assistant"][i], key=str(i), avatar_style="thumbs")
 
             except Exception as e:
                 st.error(f"Error: {str(e)}")
 
-    # Create an expander for the "About" section
-    about = st.sidebar.expander("About 🤖")
-
-    # Write information about the chatbot in the "About" section
-    about.write(
-        "#### ChatBot-PDF is an AI chatbot featuring conversational memory, designed to enable users to discuss their PDF data in a more intuitive manner. 📄"
-    )
-    about.write(
-        "#### This is a fork of [ChatBot-CSV](https://github.com/yvann-hub/ChatBot-CSV) by [yvann-hub](https://github.com/yvann-hub), many thanks to him for his work. 🤗"
-    )
-    about.write(
-        "#### He employs large language models to provide users with seamless, context-aware natural language interactions for a better understanding of their data. 🌐"
-    )
-    about.write(
-        "#### Powered by [Langchain](https://github.com/hwchase17/langchain), [OpenAI](https://platform.openai.com/docs/models/gpt-3-5) and [Streamlit](https://github.com/streamlit/streamlit) ⚡"
-    )
-    about.write("#### Source code : [gabacode/ChatBot-PDF](https://github.com/gabacode/ChatBot-PDF)")
+    sidebar.about()
 
 
 # Run the main function using asyncio

modules/chatbot.py

@@ -0,0 +1,25 @@
+import streamlit as st
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import ConversationalRetrievalChain
+
+
+class Chatbot:
+    def __init__(self, model_name, temperature, vectors):
+        self.model_name = model_name
+        self.temperature = temperature
+        self.vectors = vectors
+
+    async def conversational_chat(self, query):
+        """
+        Starts a conversational chat with a model via Langchain
+        """
+
+        chain = ConversationalRetrievalChain.from_llm(
+            llm=ChatOpenAI(model_name=self.model_name, temperature=self.temperature),
+            retriever=self.vectors.as_retriever(),
+        )
+        result = chain({"question": query, "chat_history": st.session_state["history"]})
+
+        st.session_state["history"].append((query, result["answer"]))
+
+        return result["answer"]

modules/embedder.py

@@ -0,0 +1,51 @@
+import os
+import pickle
+import tempfile
+from langchain.document_loaders import PyPDFLoader
+from langchain.vectorstores import FAISS
+from langchain.embeddings.openai import OpenAIEmbeddings
+
+
+class Embedder:
+    def __init__(self):
+        pass
+
+    async def storeDocEmbeds(self, file, filename):
+        """
+        Stores document embeddings using Langchain and FAISS
+        """
+        # Write the uploaded file to a temporary file
+        with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tmp_file:
+            tmp_file.write(file)
+            tmp_file_path = tmp_file.name
+
+        # Load the data from the file using Langchain
+        loader = PyPDFLoader(file_path=tmp_file_path)
+        data = loader.load_and_split()
+
+        # Create an embeddings object using Langchain
+        embeddings = OpenAIEmbeddings()
+
+        # Store the embeddings vectors using FAISS
+        vectors = FAISS.from_documents(data, embeddings)
+        os.remove(tmp_file_path)
+
+        # Save the vectors to a pickle file
+        with open(filename + ".pkl", "wb") as f:
+            pickle.dump(vectors, f)
+
+    async def getDocEmbeds(self, file, filename):
+        """
+        Retrieves document embeddings
+        """
+        # Check if embeddings vectors have already been stored in a pickle file
+        if not os.path.isfile(filename + ".pkl"):
+            # If not, store the vectors using the storeDocEmbeds function
+            await self.storeDocEmbeds(file, filename)
+
+        # Load the vectors from the pickle file
+        with open(filename + ".pkl", "rb") as f:
+            global vectors
+            vectors = pickle.load(f)
+
+        return vectors

modules/history.py

@@ -0,0 +1,44 @@
+import os
+import streamlit as st
+
+
+class ChatHistory:
+    def __init__(self):
+        self.history = st.session_state.get("history", [])
+        st.session_state["history"] = self.history
+
+    def default_greeting(self):
+        return "Hey ! 👋"
+
+    def default_prompt(self, thingy, topic):
+        return f'Hello! Ask me anything about the {thingy} "{topic}" 🤗'
+
+    def initialize_user_history(self):
+        st.session_state["user"] = [self.default_greeting()]
+
+    def initialize_assistant_history(self, uploaded_file):
+        st.session_state["assistant"] = [self.default_prompt("document", uploaded_file.name)]
+
+    def initialize(self, uploaded_file):
+        if "assistant" not in st.session_state:
+            self.initialize_assistant_history(uploaded_file)
+        if "user" not in st.session_state:
+            self.initialize_user_history()
+
+    def reset(self, uploaded_file):
+        st.session_state["history"] = []
+        self.initialize_user_history()
+        self.initialize_assistant_history(uploaded_file)
+        st.session_state["reset_chat"] = False
+
+    def append(self, mode, message):
+        st.session_state[mode].append(message)
+
+    def load(self):
+        if os.path.exists(self.history_file):
+            with open(self.history_file, "r") as f:
+                self.history = f.read().splitlines()
+
+    def save(self):
+        with open(self.history_file, "w") as f:
+            f.write("\n".join(self.history))

modules/layout.py

@@ -0,0 +1,20 @@
+import streamlit as st
+
+
+class Layout:
+    def display_header(self):
+        """
+        Displays the header of the app
+        """
+        st.markdown(
+            "<h1 style='text-align: center;'>ChatBot-PDF, Talk with your documents ! 💬</h1>", unsafe_allow_html=True
+        )
+
+    def show_api_key_error(self):
+        """
+        Displays an error message if the user has not entered an API key
+        """
+        st.markdown(
+            "<div style='text-align: center;'><h4>Enter your OpenAI API key to start chatting 😉</h4></div>",
+            unsafe_allow_html=True,
+        )

modules/sidebar.py

@@ -0,0 +1,51 @@
+import streamlit as st
+
+
+class Sidebar:
+    MODEL_OPTIONS = ["gpt-3.5-turbo"]
+    TEMPERATURE_MIN_VALUE = 0.0
+    TEMPERATURE_MAX_VALUE = 1.0
+    TEMPERATURE_DEFAULT_VALUE = 0.618
+    TEMPERATURE_STEP = 0.01
+
+    def __init__(self):
+        pass
+
+    def about(self):
+        about = st.sidebar.expander("About 🤖")
+        sections = [
+            "#### ChatBot-PDF is an AI chatbot featuring conversational memory, designed to enable users to discuss their PDF data in a more intuitive manner. 📄",
+            "#### This is a fork of [ChatBot-CSV](https://github.com/yvann-hub/ChatBot-CSV) by [yvann-hub](https://github.com/yvann-hub), many thanks to him for his work. 🤗",
+            "#### It employs large language models to provide users with seamless, context-aware natural language interactions for a better understanding of their data. 🌐",
+            "#### Powered by [Langchain](https://github.com/hwchase17/langchain), [OpenAI](https://platform.openai.com/docs/models/gpt-3-5) and [Streamlit](https://github.com/streamlit/streamlit) ⚡",
+            "#### Source code : [gabacode/ChatBot-PDF](https://github.com/gabacode/ChatBot-PDF)",
+        ]
+        for section in sections:
+            about.write(section)
+
+    def reset_chat_button(self):
+        if st.button("Reset chat"):
+            st.session_state["reset_chat"] = True
+        st.session_state.setdefault("reset_chat", False)
+
+    def model_selector(self):
+        model = st.selectbox(label="Model", options=self.MODEL_OPTIONS)
+        st.session_state["model"] = model
+
+    def temperature_slider(self):
+        temperature = st.slider(
+            label="Temperature",
+            min_value=self.TEMPERATURE_MIN_VALUE,
+            max_value=self.TEMPERATURE_MAX_VALUE,
+            value=self.TEMPERATURE_DEFAULT_VALUE,
+            step=self.TEMPERATURE_STEP,
+        )
+        st.session_state["temperature"] = temperature
+
+    def options(self):
+        with st.sidebar.expander("🛠️ Settings", expanded=False):
+            self.reset_chat_button()
+            self.model_selector()
+            self.temperature_slider()
+            st.session_state.setdefault("model", self.MODEL_OPTIONS[0])
+            st.session_state.setdefault("temperature", self.TEMPERATURE_DEFAULT_VALUE)