Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import streamlit as st
 import pandas as pd
+from utils.helper import *
 
 # Set page layout to wide mode
 st.set_page_config(layout="wide")
@@ -62,5 +63,58 @@ if st.session_state.logged_in:
 
     # Render the DataFrame with the URL as a hyperlink
     st.markdown(df.to_markdown(index=False), unsafe_allow_html=True)
+
+    # Chatbot
+    with st.sidebar:
+        # Add a button to clear the session state
+        if st.button("Clear Session"):
+            st.session_state.messages = []
+            st.experimental_rerun()
+
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    # Ensure messages are a list of dictionaries
+    if not isinstance(st.session_state.messages, list):
+        st.session_state.messages = []
+    if not all(isinstance(msg, dict) for msg in st.session_state.messages):
+        st.session_state.messages = []
+    
+    # Display chat messages from history on app rerun
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+    
+    # React to user input
+    if prompt := st.chat_input("😉 What GPU shall I use?"):
+    
+        # Display user message in chat message container
+        st.chat_message("user").markdown(prompt)
+    
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "system", "content": f"""
+        You are a helpful assistant assiting users on GPU selections.
+        Here's the data:
+        {df.to_markdown(index=False)}
+        
+        User may ask what is the best GPU selection.
+        You will need to ask user: 1) type of task, 2) size of data, 3) size of models.
+        You will then make a suggestion of what type of GPU or instance is the best for the user.
+        """})
+        st.session_state.messages.append({"role": "user", "content": prompt})
+    
+        # API Call
+        bot = ChatBot()
+        bot.history = st.session_state.messages.copy()  # Update history from messages
+        response = bot.generate_response(prompt)
+    
+        # Display assistant response in chat message container
+        with st.chat_message("assistant"):
+            st.markdown(response)
+    
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
 else:
     st.info("👉 Please log in to view the tasks.")