Update app.py

app.py

@@ -19,16 +19,23 @@ session.headers = {
     "Referer": "https://bard.google.com/",
 }
 session.cookies.set("__Secure-1PSID", os.getenv("_BARD_API_KEY"))
-with st.sidebar:
-    # Add a selector in the sidebar using the dictionary's keys
-    selected_language_name = st.sidebar.selectbox("Select Language", list(GOOGLE_LANGUAGES_TO_CODES.keys()))
-    code_interpreter = st.sidebar.checkbox("Code Interpreter", value=True)
-    
+
+# Set up the sidebar with language selection and code interpreter checkbox
+selected_language_name = st.sidebar.selectbox(
+    "Select Language", list(GOOGLE_LANGUAGES_TO_CODES.keys()), index=0
+)
+code_interpreter = st.sidebar.checkbox("Code Interpreter", value=True)
+
 # Retrieve the corresponding language code from the dictionary
 selected_language_code = GOOGLE_LANGUAGES_TO_CODES[selected_language_name]
 
 # Initialize Bard with the selected language code
-bard = Bard(token=os.getenv("_BARD_API_KEY"), language=selected_language_code, session=session, timeout=30)
+bard = Bard(
+    token=os.getenv("_BARD_API_KEY"),
+    language=selected_language_code,
+    session=session,
+    timeout=30,
+)
 
 TITLE = "Palm 2🌴 Chatbot"
 DESCRIPTION = """
@@ -40,46 +47,64 @@ st.write(DESCRIPTION)
 
 # Prediction function
 def predict(message):
-    with st.status("Requesting Palm-2🌴..."):
+    with st.spinner("Requesting Palm-2🌴..."):
         st.write("Requesting API...")
-        response = bard.get_answer(message if not code_interpreter else message + "Rule 1: If User requires a code snippet, write each only one code snippet and only in that way that it would run in streamlit app, and but don't output anything if it requires some additional libraries.")
+        response = bard.get_answer(
+            message if not code_interpreter else message + " Rule 1: If User requires a code snippet, write only one code snippet that would run in the Streamlit app without requiring additional libraries."
+        )
         st.write("Done...")
-        
-        st.write("Checking images...")
-    for i in response['images']:
-        st.image(i)
-    
-    return response
+
+        if "images" in response:
+            st.write("Checking images...")
+            for i in response["images"]:
+                st.image(i)
+
+        return response
+
+# Create a class to handle the chat messages
+class ChatMessage:
+    def __init__(self, role, content):
+        self.role = role
+        self.content = content
+
+    def display(self):
+        if self.role == "human":
+            st.text(self.content)
+        else:
+            with st.echo():
+                exec(self.content)
 
 # Display chat messages from history on app rerun
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+st.session_state.messages = st.session_state.get("messages", [])
 
 for message in st.session_state.messages:
-    with st.chat_message(message["role"], avatar=("🧑‍💻" if message["role"] == 'human' else '🌴')):
-        st.markdown(message["content"])
+    chat_message = ChatMessage(message["role"], message["content"])
+    chat_message.display()
 
 # React to user input
-if prompt := st.chat_input("Ask Palm 2 anything..."):
-    st.chat_message("human", avatar="🧑‍💻").markdown(prompt)
-    # st.session_state.messages.append({"role": "human", "content": prompt})
+if prompt := st.text_input("Ask Palm 2 anything..."):
+    chat_message = ChatMessage("human", prompt)
+    chat_message.display()
 
     response = predict(prompt)
-    with st.chat_message("assistant", avatar='🌴'):
-        st.markdown(response['content'])
-    
-    if response['code']:
-        with st.status("Exporting replit..."):
+    chat_message = ChatMessage("assistant", response["content"])
+    chat_message.display()
+
+    if response.get("code"):
+        with st.spinner("Exporting to repl.it..."):
             url = bard.export_replit(
-                code=response['code'],
-                program_lang=response['program_lang'],
-            )['url']
-        st.title('Export to repl.it')
-        st.markdown(f'[link]({url})')
+                code=response["code"], program_lang=response["program_lang"]
+            )["url"]
+        st.title("Export to repl.it")
+        st.markdown(f"[link]({url})")
         if code_interpreter:
             try:
-                exec(response['code'])
+                exec(response["code"])
             except Exception as e:
-                st.write(f"ERROR {e}...")
-    
-    # st.session_state.messages.append({"role": "assistant", "content": response['content']})
+                st.error(f"Error: {e}")
+
+    # Append chat messages to the session state
+    st.session_state.messages.append(
+        {"role": "human", "content": prompt},
+        {"role": "assistant", "content": response["content"]},
+    )