Update app.py

app.py

@@ -28,19 +28,7 @@ if "messages" not in st.session_state:
 
 def get_ai_response(prompt, chat_history):
     st.session_state.is_streaming = True
-
-    with stylable_container(
-            key="stop_generating",
-            css_styles="""
-                button {
-                    position: fixed;
-                    bottom: 100px;
-                    left: 50%;
-                    transform: translateX(-50%);
-                }
-                """,
-        ):
-            st.button("Stop generating")
+    st.session_state.response = ""
     
     with st.chat_message("ai"):
         penalty_kwargs = {
@@ -58,8 +46,22 @@ def get_ai_response(prompt, chat_history):
             **penalty_kwargs
         )
 
-        st.session_state.response = ""
         placeholder = st.empty()
+
+        with stylable_container(
+            key="stop_generating",
+            css_styles="""
+                button {
+                    position: fixed;
+                    bottom: 100px;
+                    left: 50%;
+                    transform: translateX(-50%);
+                    z-index: 1;
+                }
+                """,
+        ):
+            st.button("Stop generating")
+        
         for event in stream:
             if event.event_type == "text-generation":
                 content = event.text