Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from datetime import datetime
-from prompt_builder import build_prompt, bmi_category  # Import from prompt_builder
-from model_api import query_model  # Import the model API
+from prompt_builder import build_prompt
+from model_api import query_model, test_api_connection
 
 st.set_page_config(
     page_title="FIT Plan AI - Milestone 1",
@@ -12,161 +12,140 @@ st.set_page_config(
 st.title("💪 FIT Plan AI – Personalized Fitness Profile")
 st.markdown("---")
 
+# Test API connection on startup (optional)
+if "api_status" not in st.session_state:
+    with st.spinner("Checking API connection..."):
+        api_ok, api_message = test_api_connection()
+        st.session_state.api_status = api_ok
+        st.session_state.api_message = api_message
+
+# Show API status in sidebar
+with st.sidebar:
+    st.header("🔧 System Status")
+    if st.session_state.get("api_status", False):
+        st.success("✅ API Connected")
+    else:
+        st.error(f"❌ API Error: {st.session_state.get('api_message', 'Unknown error')}")
+        st.info("Please set your HF_TOKEN environment variable")
+
 # Initialize session state
 if "form_submitted" not in st.session_state:
     st.session_state.form_submitted = False
 if "workout_plan" not in st.session_state:
     st.session_state.workout_plan = None
+if "active_tab" not in st.session_state:
+    st.session_state.active_tab = "Profile"
 
-with st.form("fitness_profile_form"):
-    st.header("📋 Your Fitness Profile")
-    st.subheader("👤 Personal Information")
-
-    col1, col2 = st.columns(2)
-
-    with col1:
-        name = st.text_input("Full Name *")
-        height = st.number_input("Height (cm) *", min_value=1.0, max_value=300.0, value=170.0)
-
-    with col2:
-        age = st.number_input("Age", min_value=10, max_value=120, value=25)
-        weight = st.number_input("Weight (kg) *", min_value=1.0, max_value=500.0, value=70.0)
-
-    gender = st.selectbox("Gender", ["Male", "Female", "Other"])
-
-    st.subheader("🎯 Fitness Details")
-
-    goal = st.selectbox(
-        "Fitness Goal *",
-        ["Weight Loss", "Build Muscle", "Strength Gain", "Abs Building", "Flexibility"]
-    )
-
-    equipment = st.multiselect(
-        "Available Equipment *",
-        [
-            "Dumbbells",
-            "Resistance Bands",
-            "Barbell",
-            "Pull-up Bar",
-            "Treadmill",
-            "Kettlebells",
-            "Jump Rope",
-            "Yoga Mat",
-            "No Equipment (Bodyweight only)"
-        ],
-        default=["No Equipment (Bodyweight only)"]
-    )
-
-    fitness_level = st.select_slider(
-        "Fitness Level *",
-        options=["Beginner", "Intermediate", "Advanced"],
-        value="Beginner"
-    )
-
-    st.markdown("---")
-    submit = st.form_submit_button("Submit Profile", use_container_width=True)
-
-if submit:
-    if not name:
-        st.error("Please enter your name.")
-        st.stop()
-
-    if not equipment:
-        st.error("Please select at least one equipment option.")
-        st.stop()
-
-    # Build prompt using the imported function
-    prompt, bmi, bmi_status = build_prompt(
-        name=name,
-        gender=gender,
-        height=height,
-        weight=weight,
-        goal=goal,
-        fitness_level=fitness_level,
-        equipment=equipment
-    )
-
-    st.success("✅ Profile Submitted Successfully!")
-
-    st.session_state.update({
-        "form_submitted": True,
-        "name": name,
-        "bmi": bmi,
-        "bmi_status": bmi_status,
-        "age": age,
-        "goal": goal,
-        "equipment": equipment,
-        "fitness_level": fitness_level,
-        "height": height,
-        "weight": weight,
-        "gender": gender,
-        "prompt": prompt  # Store the generated prompt
-    })
-
-if st.session_state.form_submitted:
-    st.markdown("---")
-    st.success(f"Welcome, **{st.session_state.name}**!")
-
-    col1, col2, col3 = st.columns(3)
-    col1.metric("📊 BMI", f"{st.session_state.bmi:.2f}")
-    col2.metric("🏷 Category", st.session_state.bmi_status)
-    col3.metric("🏋️ Level", st.session_state.fitness_level)
-
-    # Generate Workout Plan Button
-    st.markdown("---")
-    st.subheader("🏋️ Generate Your Personalized Workout Plan")
-    
-    if st.button("Generate Workout Plan", type="primary", use_container_width=True):
-        with st.spinner("Creating your personalized workout plan... This may take a moment."):
-            # Call the model API
-            workout_plan = query_model(st.session_state.prompt)
-            st.session_state.workout_plan = workout_plan
-        
-    # Display workout plan if generated
-    if st.session_state.workout_plan:
-        st.markdown("---")
-        st.subheader("📋 Your 5-Day Workout Plan")
-        
-        with st.container():
-            st.markdown(st.session_state.workout_plan)
-        
-        # Download button for workout plan
-        workout_text = f"""
-FIT PLAN AI - PERSONALIZED WORKOUT PLAN
-Generated for: {st.session_state.name}
-Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+# Create tabs for different sections
+tab1, tab2 = st.tabs(["📋 Profile Setup", "🏋️ Workout Plan"])
 
-Profile Summary:
-- BMI: {st.session_state.bmi:.2f} ({st.session_state.bmi_status})
-- Goal: {st.session_state.goal}
-- Level: {st.session_state.fitness_level}
-- Equipment: {', '.join(st.session_state.equipment)}
-
-{st.session_state.workout_plan}
+# Tab 1: Profile Setup
+with tab1:
+    with st.form("fitness_profile_form"):
+        st.header("Your Fitness Profile")
+        st.subheader("👤 Personal Information")
 
----
-Stay consistent and trust the process!
-"""
-        
         col1, col2 = st.columns(2)
+
         with col1:
-            st.download_button(
-                "📥 Download Workout Plan",
-                data=workout_text,
-                file_name=f"FITPlanAI_{st.session_state.name}_Workout.txt",
-                mime="text/plain",
-                use_container_width=True
-            )
-        
+            name = st.text_input("Full Name *")
+            height = st.number_input("Height (cm) *", min_value=1.0, max_value=300.0, value=170.0)
+
         with col2:
-            if st.button("Generate New Plan", use_container_width=True):
-                st.session_state.workout_plan = None
-                st.rerun()
-
-    # Profile Report Section
-    st.markdown("---")
-    with st.expander("📄 View Profile Report"):
-        report_text = f"""
+            age = st.number_input("Age *", min_value=10, max_value=120, value=25)
+            weight = st.number_input("Weight (kg) *", min_value=1.0, max_value=500.0, value=70.0)
+
+        gender = st.selectbox("Gender", ["Male", "Female", "Other"])
+
+        st.subheader("🎯 Fitness Details")
+
+        goal = st.selectbox(
+            "Fitness Goal *",
+            ["Weight Loss", "Build Muscle", "Strength Gain", "Abs Building", "Flexibility"]
+        )
+
+        equipment = st.multiselect(
+            "Available Equipment *",
+            [
+                "Dumbbells",
+                "Resistance Bands",
+                "Barbell",
+                "Pull-up Bar",
+                "Treadmill",
+                "Kettlebells",
+                "Jump Rope",
+                "Yoga Mat",
+                "No Equipment (Bodyweight only)"
+            ],
+            default=["No Equipment (Bodyweight only)"]
+        )
+
+        fitness_level = st.select_slider(
+            "Fitness Level *",
+            options=["Beginner", "Intermediate", "Advanced"],
+            value="Beginner"
+        )
+
+        st.markdown("---")
+        submit = st.form_submit_button("Save Profile", use_container_width=True)
+
+    if submit:
+        if not name:
+            st.error("Please enter your name.")
+            st.stop()
+
+        if not equipment:
+            st.error("Please select at least one equipment option.")
+            st.stop()
+
+        # Build prompt using the imported function
+        prompt, bmi, bmi_status = build_prompt(
+            name=name,
+            age=age,
+            gender=gender,
+            height=height,
+            weight=weight,
+            goal=goal,
+            fitness_level=fitness_level,
+            equipment=equipment
+        )
+
+        st.success("✅ Profile Saved Successfully!")
+
+        st.session_state.update({
+            "form_submitted": True,
+            "name": name,
+            "bmi": bmi,
+            "bmi_status": bmi_status,
+            "age": age,
+            "goal": goal,
+            "equipment": equipment,
+            "fitness_level": fitness_level,
+            "height": height,
+            "weight": weight,
+            "gender": gender,
+            "prompt": prompt
+        })
+
+    # Show profile summary if form is submitted
+    if st.session_state.form_submitted:
+        st.markdown("---")
+        st.success(f"Welcome, **{st.session_state.name}**!")
+
+        col1, col2, col3 = st.columns(3)
+        col1.metric("📊 BMI", f"{st.session_state.bmi:.2f}")
+        col2.metric("🏷 Category", st.session_state.bmi_status)
+        col3.metric("🏋️ Level", st.session_state.fitness_level)
+        
+        # Quick action button to switch to workout plan tab
+        if st.button("👉 Go to Workout Plan Tab", type="primary", use_container_width=True):
+            st.session_state.active_tab = "Workout Plan"
+            st.rerun()
+
+        # Profile Report Section
+        st.markdown("---")
+        with st.expander("📄 View Profile Report"):
+            report_text = f"""
 FIT PLAN AI - PROFILE REPORT
 Generated on: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
 
@@ -183,16 +162,150 @@ Goal: {st.session_state.goal}
 Level: {st.session_state.fitness_level}
 Equipment: {', '.join(st.session_state.equipment)}
 
-Stay consistent and trust the process!
+Go to the Workout Plan tab to generate your personalized 5-day workout plan!
 """
 
-        st.download_button(
-            "📥 Download Profile Report",
-            data=report_text,
-            file_name=f"FITPlanAI_{st.session_state.name}_Profile.txt",
-            mime="text/plain",
-            use_container_width=True
-        )
+            st.download_button(
+                "📥 Download Profile Report",
+                data=report_text,
+                file_name=f"FITPlanAI_{st.session_state.name}_Profile.txt",
+                mime="text/plain",
+                use_container_width=True
+            )
 
+# Tab 2: Workout Plan
+with tab2:
+    st.header("🏋️ Your Personalized Workout Plan")
+    
+    # Check if profile is submitted
+    if not st.session_state.form_submitted:
+        st.warning("⚠️ Please set up your profile first in the 'Profile Setup' tab!")
+        if st.button("Go to Profile Setup"):
+            st.session_state.active_tab = "Profile"
+            st.rerun()
+    else:
+        # Check API status
+        if not st.session_state.get("api_status", False):
+            st.error("⚠️ API is not connected. Please check your HF_TOKEN environment variable.")
+            st.stop()
+        
+        # Show profile summary
+        with st.container():
+            st.markdown("### Profile Summary")
+            col1, col2, col3, col4 = st.columns(4)
+            col1.metric("Name", st.session_state.name)
+            col2.metric("Age", st.session_state.age)
+            col3.metric("BMI", f"{st.session_state.bmi:.1f}")
+            col4.metric("Goal", st.session_state.goal)
+        
+        st.markdown("---")
+        
+        # Generate Workout Plan Button
+        st.subheader("Generate Your 5-Day Workout Plan")
+        
+        col1, col2, col3 = st.columns([1, 2, 1])
+        with col2:
+            if st.button("🚀 Generate Workout Plan", type="primary", use_container_width=True):
+                with st.spinner("Creating your personalized workout plan... This may take a moment."):
+                    # Call the model API from model_api.py
+                    workout_plan = query_model(st.session_state.prompt)
+                    st.session_state.workout_plan = workout_plan
+                
+        # Display workout plan if generated
+        if st.session_state.workout_plan:
+            st.markdown("---")
+            st.subheader("📋 Your 5-Day Workout Plan")
+            
+            # Check if plan is complete
+            is_complete, message = verify_workout_plan(st.session_state.workout_plan)
+            if not is_complete:
+                st.warning(f"⚠️ {message}. You can regenerate the plan if needed.")
+            
+            # Display the workout plan
+            with st.container():
+                st.markdown(st.session_state.workout_plan)
+            
+            st.markdown("---")
+            
+            # Action buttons
+            col1, col2, col3 = st.columns(3)
+            
+            with col1:
+                # Download button for workout plan
+                workout_text = f"""
+FIT PLAN AI - PERSONALIZED WORKOUT PLAN
+Generated for: {st.session_state.name}
+Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+
+PROFILE SUMMARY:
+- Age: {st.session_state.age}
+- BMI: {st.session_state.bmi:.2f} ({st.session_state.bmi_status})
+- Goal: {st.session_state.goal}
+- Level: {st.session_state.fitness_level}
+- Equipment: {', '.join(st.session_state.equipment)}
+
+{st.session_state.workout_plan}
+
+---
+Stay consistent and trust the process!
+Generated by FIT Plan AI
+"""
+                
+                st.download_button(
+                    "📥 Download Plan",
+                    data=workout_text,
+                    file_name=f"FITPlanAI_{st.session_state.name}_Workout.txt",
+                    mime="text/plain",
+                    use_container_width=True
+                )
+            
+            with col2:
+                if st.button("🔄 Regenerate Plan", use_container_width=True):
+                    with st.spinner("Regenerating your workout plan..."):
+                        workout_plan = query_model(st.session_state.prompt)
+                        st.session_state.workout_plan = workout_plan
+                        st.rerun()
+            
+            with col3:
+                if st.button("✏️ Edit Profile", use_container_width=True):
+                    st.session_state.active_tab = "Profile"
+                    st.rerun()
+        
+        # Tips section
+        with st.expander("💡 Workout Tips & Guidelines"):
+            st.markdown("""
+            ### 📌 Tips for Best Results:
+            1. **Warm-up** before each workout (5-10 minutes of light cardio and dynamic stretches)
+            2. **Stay hydrated** during your workout
+            3. **Maintain proper form** over lifting heavy weights
+            4. **Rest adequately** between sets as specified
+            5. **Cool down** and stretch after each session
+            6. **Be consistent** - follow the plan regularly
+            7. **Listen to your body** and adjust intensity if needed
+            
+            ### ⚠️ Safety First:
+            - Consult with a healthcare professional before starting any new exercise program
+            - Stop immediately if you feel sharp pain or discomfort
+            - Start with lighter weights to master the form
+            """)
+
+# Footer
 st.markdown("---")
-st.caption("FIT Plan AI - Milestone 1 | Powered by Mistral-7B")
+st.caption("FIT Plan AI - Milestone 1 | Powered by Mistral-7B")
+
+# Helper function to verify workout plan
+def verify_workout_plan(plan_text):
+    """Check if the workout plan contains all 5 days"""
+    if not plan_text or plan_text.startswith("Error"):
+        return False, "No valid plan generated"
+    
+    days_present = []
+    for i in range(1, 6):
+        if f"Day {i}" in plan_text:
+            days_present.append(i)
+    
+    if len(days_present) == 5:
+        return True, "Complete 5-day plan"
+    else:
+        missing_days = set(range(1,6)) - set(days_present)
+        return False, f"Missing day(s): {', '.join(map(str, missing_days))}"