Update app.py

app.py

@@ -154,13 +154,17 @@ def main():
         if submit_clicked:
             output_container = output_container.container()
             output_container.chat_message("user").write(user_input)
+            with st.chat_message("assistant"):
+                st_callback = StreamlitCallbackHandler(st.container())
+                response = agent.run(user_input, callbacks=[st_callback])
+                st.write(response)
 
-            answer_container = output_container.chat_message("assistant", avatar="🦜")
-            st_callback = StreamlitCallbackHandler(answer_container)
+            # answer_container = output_container.chat_message("assistant", avatar="🦜")
+            # st_callback = StreamlitCallbackHandler(answer_container,)
 
-            answer = agent.run(user_input, callbacks=[st_callback])
+            # answer = agent.run(user_input, callbacks=[st_callback])
 
-            answer_container = output_container.container()
+            # answer_container = output_container.container()
             # answer_container.chat_message("assistant").write(answer)