Update app.py

app.py

@@ -33,9 +33,8 @@ def predict(message):
         st.write("Checking images...")
     for i in response['images']:
         st.image(i)
-    if response['code']:
-        exec(response['code'])
-    return response['content']
+    
+    return response
 
 # Streamlit UI
 st.title(TITLE)
@@ -60,6 +59,17 @@ if prompt := st.chat_input("Ask Palm 2 anything..."):
     response = predict(prompt)
     # Display assistant response in chat message container
     with st.chat_message("assistant", avatar='🌴'):
-        st.markdown(response)
+        st.markdown(response['content'])
+    
+    if response['code']:
+        if st.button('Run Code')
+            with st.status("Running code..."):
+                try:
+                    exec(response['code'])
+                except Exception as e:
+                    st.write(f"ERROR {e}...")
+                    
+            
+            
     # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": response})