Merge pull request #10 from gabacode/feat/refactor-csv

src/chatbot_csv.py

@@ -1,221 +1,64 @@
 import os
-import pickle
 import streamlit as st
-import tempfile
-import pandas as pd
 import asyncio
+from dotenv import load_dotenv
 
-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.csv_loader import CSVLoader
-from langchain.vectorstores import FAISS
-from langchain.prompts.prompt import PromptTemplate
 
+from modules.history import ChatHistory
+from modules.layout import Layout
+from modules.utils import Utilities
+from modules.sidebar import Sidebar
 
-st.set_page_config(layout="wide", page_icon="💬", page_title="ChatBot-CSV")
 
-st.markdown(
-    "<h1 style='text-align: center;'>ChatBot-CSV, Talk with your  csv-data ! 💬</h1>",
-    unsafe_allow_html=True)
+def init():
+    load_dotenv()
+    st.set_page_config(layout="wide", page_icon="💬", page_title="ChatBot-CSV")
 
-user_api_key = st.sidebar.text_input(
-    label="#### Your OpenAI API key 👇",
-    placeholder="Paste your openAI API key, sk-",
-    type="password")
 
 async def main():
-    
-    if user_api_key == "":
-        
-        st.markdown(
-            "<div style='text-align: center;'><h4>Enter your OpenAI API key to start chatting 😉</h4></div>",
-            unsafe_allow_html=True)
-        
+    init()
+    layout, sidebar, utils = Layout(), Sidebar(), Utilities()
+    layout.show_header()
+    user_api_key = utils.load_api_key()
+
+    if not user_api_key:
+        layout.show_api_key_missing()
     else:
         os.environ["OPENAI_API_KEY"] = user_api_key
-        
-        uploaded_file = st.sidebar.file_uploader("upload", type="csv", label_visibility="hidden")
-        
-        if uploaded_file is not None:
-            def show_user_file(uploaded_file):
-                file_container = st.expander("Your CSV file :")
-                shows = pd.read_csv(uploaded_file)
-                uploaded_file.seek(0)
-                file_container.write(shows)
-                
-            show_user_file(uploaded_file)
-            
-        else :
-            st.sidebar.info(
-            "👆 Upload your CSV file to get started, "
-            "sample for try : [fishfry-locations.csv](https://drive.google.com/file/d/18i7tN2CqrmoouaSqm3hDfAk17hmWx94e/view?usp=sharing)" 
-            )
-    
-        if uploaded_file :
-            try :
-                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 CSV file using Langchain
-                    loader = CSVLoader(file_path=tmp_file_path, encoding="utf-8")
-                    data = loader.load()
-
-                    embeddings = OpenAIEmbeddings()
-                    
-                    vectors = FAISS.from_documents(data, embeddings)
-                    os.remove(tmp_file_path)
-
-                    with open(filename + ".pkl", "wb") as f:
-                        pickle.dump(vectors, f)
-                    
-                async def getDocEmbeds(file, filename):
-                    
-                    if not os.path.isfile(filename + ".pkl"):
-                        # If not, store the vectors using the storeDocEmbeds function
-                        await storeDocEmbeds(file, filename)
-                    
-                    with open(filename + ".pkl", "rb") as f:
-                        #global vectors
-                        vectors = pickle.load(f)
-                        
-                    return vectors
-
-                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"]))
-                    
-                    # You can print the chat history for debugging :
-                    #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','gpt-4'])
-
-                if 'history' not in st.session_state:
-                    st.session_state['history'] = []
-
-                if 'ready' not in st.session_state:
-                    st.session_state['ready'] = False
-                    
-                if 'reset_chat' not in st.session_state:
-                    st.session_state['reset_chat'] = False
-                
-                if uploaded_file is not None:
-
-                    # Display a spinner while processing the file
-                    with st.spinner("Processing..."):
-
-                        uploaded_file.seek(0)
-                        file = uploaded_file.read()
-                        
-                        # Generate embeddings vectors for the file
-                        vectors = await getDocEmbeds(file, uploaded_file.name)
-
-                        _template = """Given the following conversation and a follow-up question, rephrase the follow-up question to be a stand-alone question.
-                        You can assume that the question is about the information in a CSV file.
-                        Chat History:
-                        {chat_history}
-                        Follow-up entry: {question}
-                        Standalone question:"""
-                        CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
-
-                        qa_template = """"You are an AI conversational assistant to answer questions based on information from a csv file.
-                        You are given data from a csv file and a question, you must help the user find the information they need. 
-                        Only give responses for information you know about. Don't try to make up an answer.
-                        Your answers should be short,friendly, in the same language.
-                        question: {question}
-                        =========
-                        {context}
-                        =======
-                        """
-                        QA_PROMPT = PromptTemplate(template=qa_template, input_variables=["question", "context"])
-
-                        chain = ConversationalRetrievalChain.from_llm(llm = ChatOpenAI(temperature=0.0,model_name=MODEL),
-                                                                      condense_question_prompt=CONDENSE_QUESTION_PROMPT,qa_prompt=QA_PROMPT,retriever=vectors.as_retriever())
-
-                    # Set the "ready" flag to True now that the chatbot is ready to chat
-                    st.session_state['ready'] = True
-
-                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 " + uploaded_file.name + " 🤗"]
-
-                    if 'past' not in st.session_state:
-                        st.session_state['past'] = ["Hey ! 👋"]
-
-                    #container for displaying the chat history
-                    response_container = st.container()
-                    
-                    #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_input("Query:", placeholder="Talk about your csv data here (:", key='input')
-                            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 " + uploaded_file.name + " 🤗"]
-                                response_container.empty()
-                                st.session_state['reset_chat'] = False
-
-                        if submit_button and user_input:
-                            
-                            # Generate a response using the Langchain ConversationalRetrievalChain
-                            output = await conversational_chat(user_input)
-                            
-                            # Add the user's input and the chatbot's output to the chat history
-                            st.session_state['past'].append(user_input)
-                            st.session_state['generated'].append(output)
-
-                    if st.session_state['generated']:
-                        
-                        # Display the chat history
-                        with response_container:
-                            
-                            for i in range(len(st.session_state['generated'])):
-                                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="big-smile")
-                                message(st.session_state["generated"][i], key=str(i), avatar_style="thumbs")
-                
+        uploaded_file = utils.handle_upload()
+
+        if uploaded_file:
+            history = ChatHistory()
+            sidebar.show_options()
+
+            try:
+                chatbot = await utils.setup_chatbot(
+                    uploaded_file, st.session_state["model"], st.session_state["temperature"]
+                )
+                st.session_state["chatbot"] = chatbot
+
+                if st.session_state["ready"]:
+                    response_container, prompt_container = st.container(), st.container()
+
+                    with prompt_container:
+                        is_ready, user_input = layout.prompt_form()
+
+                        history.initialize(uploaded_file)
+                        if st.session_state["reset_chat"]:
+                            history.reset(uploaded_file)
+
+                        if is_ready:
+                            history.append("user", user_input)
+                            output = await st.session_state["chatbot"].conversational_chat(user_input)
+                            history.append("assistant", output)
+
+                    history.generate_messages(response_container)
 
             except Exception as e:
                 st.error(f"Error: {str(e)}")
 
-    about = st.sidebar.expander("About 🤖")
-    about.write("#### ChatBot-CSV is an AI chatbot featuring conversational memory, designed to enable users to discuss their CSV data in a more intuitive manner. 📄")
-    about.write("#### He employs large language models to provide users with seamless, context-aware natural language interactions for a better understanding of their CSV 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 : [yvann-hub/ChatBot-CSV](https://github.com/yvann-hub/ChatBot-CSV)")
+    sidebar.about()
+
 
-#Run the main function using asyncio
 if __name__ == "__main__":
     asyncio.run(main())
-

src/modules/chatbot.py

@@ -0,0 +1,49 @@
+import streamlit as st
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import ConversationalRetrievalChain
+from langchain.prompts.prompt import PromptTemplate
+
+
+class Chatbot:
+    _template = """Given the following conversation and a follow-up question, rephrase the follow-up question to be a stand-alone question.
+    You can assume that the question is about the information in a CSV file.
+    Chat History:
+    {chat_history}
+    Follow-up entry: {question}
+    Standalone question:"""
+
+    CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
+
+    qa_template = """"You are an AI conversational assistant to answer questions based on information from a csv file.
+    You are given data from a csv file and a question, you must help the user find the information they need. 
+    Only give responses for information you know about. Don't try to make up an answer.
+    Your answers should be short,friendly, in the same language.
+    question: {question}
+    =========
+    {context}
+    =======
+    """
+
+    QA_PROMPT = PromptTemplate(template=qa_template, input_variables=["question", "context"])
+
+    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),
+            condense_question_prompt=self.CONDENSE_QUESTION_PROMPT,
+            qa_prompt=self.QA_PROMPT,
+            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"]

src/modules/embedder.py

@@ -0,0 +1,58 @@
+import os
+import pickle
+import tempfile
+from langchain.document_loaders.csv_loader import CSVLoader
+from langchain.vectorstores import FAISS
+from langchain.embeddings.openai import OpenAIEmbeddings
+
+
+class Embedder:
+    def __init__(self):
+        self.PATH = "embeddings"
+        self.createEmbeddingsDir()
+
+    def createEmbeddingsDir(self):
+        """
+        Creates a directory to store the embeddings vectors
+        """
+        if not os.path.exists(self.PATH):
+            os.mkdir(self.PATH)
+
+    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 = CSVLoader(file_path=tmp_file_path, encoding="utf-8")
+        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(f"{self.PATH}/{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(f"{self.PATH}/{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(f"{self.PATH}/{filename}.pkl", "rb") as f:
+            vectors = pickle.load(f)
+
+        return vectors

src/modules/history.py

@@ -0,0 +1,57 @@
+import os
+import streamlit as st
+from streamlit_chat import message
+
+
+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, topic):
+        return f"Hello ! Ask me anything about {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(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 generate_messages(self, container):
+        if st.session_state["assistant"]:
+            with container:
+                for i in range(len(st.session_state["assistant"])):
+                    message(
+                        st.session_state["user"][i],
+                        is_user=True,
+                        key=f"{i}_user",
+                        avatar_style="big-smile",
+                    )
+                    message(st.session_state["assistant"][i], key=str(i), avatar_style="thumbs")
+
+    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))

src/modules/layout.py

@@ -0,0 +1,42 @@
+import streamlit as st
+
+
+class Layout:
+    def show_header(self):
+        """
+        Displays the header of the app
+        """
+        st.markdown(
+            """
+            <h1 style='text-align: center;'>ChatBot-CSV, Talk with your  csv-data ! 💬</h1>
+            """,
+            unsafe_allow_html=True,
+        )
+
+    def show_api_key_missing(self):
+        """
+        Displays a message if the user has not entered an API key
+        """
+        st.markdown(
+            """
+            <div style='text-align: center;'>
+                <h4>Enter your <a href="https://platform.openai.com/account/api-keys" target="_blank">OpenAI API key</a> to start chatting  😉</h4>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
+
+    def prompt_form(self):
+        """
+        Displays the prompt form
+        """
+        with st.form(key="my_form", clear_on_submit=True):
+            user_input = st.text_area(
+                "Query:",
+                placeholder="Ask me anything about the document...",
+                key="input",
+                label_visibility="collapsed",
+            )
+            submit_button = st.form_submit_button(label="Send")
+            is_ready = submit_button and user_input
+        return is_ready, user_input

src/modules/sidebar.py

@@ -0,0 +1,49 @@
+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.0
+    TEMPERATURE_STEP = 0.01
+
+    @staticmethod
+    def about():
+        about = st.sidebar.expander("About 🤖")
+        sections = [
+            "#### ChatBot-CSV is an AI chatbot featuring conversational memory, designed to enable users to discuss their CSV data in a more intuitive manner. 📄",
+            "#### He employs large language models to provide users with seamless, context-aware natural language interactions for a better understanding of their CSV 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 : [yvann-hub/ChatBot-CSV](https://github.com/yvann-hub/ChatBot-CSV)",
+        ]
+        for section in sections:
+            about.write(section)
+
+    @staticmethod
+    def reset_chat_button():
+        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 show_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)

src/modules/utils.py

@@ -0,0 +1,62 @@
+import os
+import pandas as pd
+import streamlit as st
+
+from modules.chatbot import Chatbot
+from modules.embedder import Embedder
+
+
+class Utilities:
+    @staticmethod
+    def load_api_key():
+        """
+        Loads the OpenAI API key from the .env file or from the user's input
+        and returns it
+        """
+        if os.path.exists(".env") and os.environ.get("OPENAI_API_KEY") is not None:
+            user_api_key = os.environ["OPENAI_API_KEY"]
+            st.sidebar.success("API key loaded from .env", icon="🚀")
+        else:
+            user_api_key = st.sidebar.text_input(
+                label="#### Your OpenAI API key 👇", placeholder="Paste your openAI API key, sk-", type="password"
+            )
+            if user_api_key:
+                st.sidebar.success("API key loaded", icon="🚀")
+        return user_api_key
+
+    @staticmethod
+    def handle_upload():
+        """
+        Handles the file upload and displays the uploaded file
+        """
+        uploaded_file = st.sidebar.file_uploader("upload", type="csv", label_visibility="collapsed")
+        if uploaded_file is not None:
+
+            def show_user_file(uploaded_file):
+                file_container = st.expander("Your CSV file :")
+                shows = pd.read_csv(uploaded_file)
+                uploaded_file.seek(0)
+                file_container.write(shows)
+
+            show_user_file(uploaded_file)
+        else:
+            st.sidebar.info(
+                "👆 Upload your CSV file to get started, "
+                "sample for try : [fishfry-locations.csv](https://drive.google.com/file/d/18i7tN2CqrmoouaSqm3hDfAk17hmWx94e/view?usp=sharing)"
+            )
+            st.session_state["reset_chat"] = True
+        return uploaded_file
+
+    @staticmethod
+    async def setup_chatbot(uploaded_file, model, temperature):
+        """
+        Sets up the chatbot with the uploaded file, model, and 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