Update app.py

app.py

@@ -25,23 +25,23 @@ if "messages" not in st.session_state:
 with st.sidebar:
     st.header("Model Configuration")
     st.markdown("[Get HuggingFace Token](https://huggingface.co/settings/tokens)")
-    
+
     system_message = st.text_area(
         "System Message",
         value="You are a friendly Chatbot created by ruslanmv.com",
         height=100
     )
-    
+
     max_tokens = st.slider(
         "Max Tokens",
         1, 4000, 512
     )
-    
+
     temperature = st.slider(
         "Temperature",
         0.1, 4.0, 0.7
     )
-    
+
     top_p = st.slider(
         "Top-p",
         0.1, 1.0, 0.9
@@ -59,7 +59,7 @@ for message in st.session_state.messages:
 # Handle input
 if prompt := st.chat_input("Type your message..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
-    
+
     with st.chat_message("user"):
         st.markdown(prompt)
 
@@ -75,21 +75,21 @@ if prompt := st.chat_input("Type your message..."):
                     "return_full_text": False
                 }
             }
-            
+
             # Query the Hugging Face API
             output = query(payload)
-            
+
             # Handle API response
             if isinstance(output, list) and len(output) > 0 and 'generated_text' in output[0]:
                 assistant_response = output[0]['generated_text']
             else:
                 st.error("Error: Unable to generate a response. Please try again.")
-                return
+                return  # Return from the function to stop further execution
+
+            with st.chat_message("assistant"):
+                st.markdown(assistant_response)
 
-        with st.chat_message("assistant"):
-            st.markdown(assistant_response)
-        
-        st.session_state.messages.append({"role": "assistant", "content": assistant_response})
+            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
 
     except Exception as e:
         st.error(f"Application Error: {str(e)}")