Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,13 +8,23 @@ st.set_page_config(page_title="FitPlan AI", layout="centered")
 # ---------------------------------------------------
 
 API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+
 headers = {
-    "Authorization": f"Bearer {st.secrets['HF_TOKEN']}"
+    "Authorization": f"Bearer {st.secrets['HF_TOKEN']}",
+    "Content-Type": "application/json"
 }
 
 def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+    try:
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
+
+        if response.status_code != 200:
+            return {"error": f"API Error {response.status_code}: {response.text}"}
+
+        return response.json()
+
+    except requests.exceptions.RequestException as e:
+        return {"error": f"Request failed: {str(e)}"}
 
 # ---------------------------------------------------
 # TITLE
@@ -174,16 +184,20 @@ You are a certified professional fitness trainer.
 
 Generate a structured 5-day workout plan.
 
-Return ONLY the workout plan.
-Do not add explanations.
+Rules:
+- Clearly label Day 1 to Day 5
+- Include exercises
+- Include sets and reps
+- Include short rest guidance
+- Return only the workout plan
 
+User Details:
 Gender: {gender}
 BMI: {bmi:.2f} ({bmi_status})
 Goal: {goal}
 Fitness Level: {fitness_level}
-Equipment: {equipment_list}
-[/INST]
-"""
+Equipment Available: {equipment_list}
+[/INST]"""
 
         with st.spinner("Generating your AI workout plan..."):
 
@@ -191,14 +205,22 @@ Equipment: {equipment_list}
                 "inputs": prompt,
                 "parameters": {
                     "max_new_tokens": 600,
-                    "temperature": 0.3
+                    "temperature": 0.3,
+                    "return_full_text": False
                 }
             })
 
-            if isinstance(output, list):
-                response = output[0]["generated_text"]
+            if isinstance(output, dict) and "error" in output:
+                response = f"⚠️ {output['error']}\n\nPlease try again."
+
+            elif isinstance(output, list):
+                response = output[0]["generated_text"].strip()
+
+                if "[/INST]" in response:
+                    response = response.split("[/INST]")[-1].strip()
+
             else:
-                response = "Error generating plan. Please try again."
+                response = "Unexpected response from model."
 
         st.subheader("🏋️ Your Personalized Workout Plan")
         st.markdown(response)