Update app.py

app.py

@@ -1,55 +1,22 @@
 from datetime import datetime
-
 import streamlit as st
 import os
 from mistralai import Mistral
 
-
-api_key = os.environ["MISTRAL_API_KEY"]
-
-class ChatBot:
-    def __init__(self, api_key):
-        self.client = Mistral(api_key=api_key)
-        self.history = []
-
-    def generate_response(self, agent_id, user_message):
-        # Add the user's message to the history
-        self.history.append({
-            "role": "user",
-            "content": user_message
-        })
-
-        # Send the message to the API and get the response
-        chat_response = self.client.agents.complete(
-            agent_id=agent_id,
-            messages=self.history,
-        )
-
-        # Extract the AI's response and add it to the history
-        ai_response = chat_response.choices[0].message.content
-        self.history.append({
-            "role": "assistant",
-            "content": ai_response
-        })
-
-        # Return the AI's response
-        return ai_response
-
-    def get_history(self):
-        return self.history
-
-
-# Credit: Time
+# Function to get the current year
 def current_year():
     now = datetime.now()
     return now.year
 
-
 st.set_page_config(layout="wide")
 st.title("Just chat! 🤖")
 
-
+# Sidebar inputs for MISTRAL_API_ID and MISTRAL_API_KEY
 with st.sidebar:
+    st.header("Mistral API Settings")
+    mistral_api_id = st.text_input("Enter your MISTRAL API ID", placeholder="Enter API ID")
+    api_key = st.text_input("Enter your MISTRAL API Key", type="password", placeholder="Enter API Key")
+
     with st.expander("Instruction Manual"):
         st.markdown("""
             ## Mistral AI Agent 🤖 Chatbot
@@ -63,27 +30,16 @@ with st.sidebar:
             Enjoy chatting with Mistral AI agent!
         """)
 
-    # Example:
     with st.expander("Examples"):
         st.success("Example: Write a simple hello world flask app for me.")
         st.success("Example: Write a program to simulate the value of pi.")
         st.success("Example: Write a 3-layer neural network for me?")
 
-    # Add a button to clear the session state
+    # Button to clear session state
     if st.button("Clear Session"):
         st.session_state.messages = []
         st.rerun()
 
-    # Credit:
-    # current_year = current_year()  # This will print the current year
-    # st.markdown(
-    #     f"""
-    #         <h6 style='text-align: left;'>Copyright © 2010-{current_year} Present Yiqiao Yin</h6>
-    #     """,
-    #     unsafe_allow_html=True,
-    # )
-
-
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -106,20 +62,23 @@ if prompt := st.chat_input("😉 Ask any question or feel free to use the exampl
     st.chat_message("user").markdown(prompt)
 
     # Add user message to chat history
-    st.session_state.messages.append({"role": "assistant", "content": f"You are a helpful assistant. Year now is {current_year}"})
+    st.session_state.messages.append({"role": "assistant", "content": f"You are a helpful assistant. Year now is {current_year()}"})
     st.session_state.messages.append({"role": "user", "content": prompt})
 
     # API Call
-    bot = ChatBot(api_key)
-    bot.history = st.session_state.messages.copy()  # Update history from messages
-    response = bot.generate_response(
-        agent_id="ag:bfb4e4d9:20240809:coding-agent:b1d09feb",
-        user_message=prompt
-    )
+    if mistral_api_id and api_key:  # Ensure both API ID and Key are entered
+        bot = ChatBot(api_key)
+        bot.history = st.session_state.messages.copy()  # Update history from messages
+        response = bot.generate_response(
+            agent_id=mistral_api_id,
+            user_message=prompt
+        )
 
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        st.markdown(response)
+        # 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})
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+    else:
+        st.error("Please enter both MISTRAL API ID and API Key in the sidebar.")