Update app.py

app.py

@@ -65,6 +65,7 @@ if prompt := st.chat_input("Type your message..."):
         with st.spinner("Generating response..."):
             full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
             
+            # Make API request
             response = requests.post(
                 API_URL,
                 headers=headers,
@@ -77,12 +78,34 @@ if prompt := st.chat_input("Type your message..."):
                         "return_full_text": False
                     }
                 }
-            ).json()
+            )
             
-            if isinstance(response, list) and len(response) > 0:
-                assistant_response = response[0].get('generated_text', '')
+            # Handle API errors
+            if response.status_code != 200:
+                error_msg = response.json().get('error', 'Unknown API error')
+                st.error(f"API Error: {error_msg}")
+                if "loading" in error_msg.lower():
+                    st.info("Please wait a moment and try again. The model might be loading.")
+                return
+            
+            # Process successful response
+            result = response.json()
+            
+            if isinstance(result, list):
+                # Handle normal response format
+                assistant_response = result[0].get('generated_text', 'No response generated')
+                
+                # Clean up response
+                if "Assistant:" in assistant_response:
+                    assistant_response = assistant_response.split("Assistant:")[-1].strip()
+                
+            elif isinstance(result, dict) and 'error' in result:
+                # Handle error format
+                st.error(f"API Error: {result['error']}")
+                return
             else:
-                assistant_response = "Error: Unexpected API response"
+                st.error("Unexpected response format from API")
+                return
 
         with st.chat_message("assistant"):
             st.markdown(assistant_response)
@@ -90,4 +113,4 @@ if prompt := st.chat_input("Type your message..."):
         st.session_state.messages.append({"role": "assistant", "content": assistant_response})
 
     except Exception as e:
-        st.error(f"API Error: {str(e)}")
+        st.error(f"Application Error: {str(e)}")