Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from datetime import datetime
 from prompt_builder import build_prompt
-from model_api import query_model, test_api_connection
+from model_api import query_model, test_api_connection, switch_model, set_ollama_mode
 
 st.set_page_config(
     page_title="FIT Plan AI - Milestone 2",
@@ -36,7 +36,7 @@ if "api_status" not in st.session_state:
         st.session_state.api_status = api_ok
         st.session_state.api_message = api_message
 
-# Show API status in sidebar
+# Show API status and model settings in sidebar
 with st.sidebar:
     st.header("🔧 System Status")
     if st.session_state.get("api_status", False):
@@ -44,6 +44,45 @@ with st.sidebar:
     else:
         st.error(f"❌ API Error: {st.session_state.get('api_message', 'Unknown error')}")
         st.info("Please set your HF_TOKEN environment variable")
+    
+    # Add model selection
+    st.markdown("---")
+    st.header("🤖 Model Settings")
+    
+    # Model selection dropdown
+    model_options = {
+        "Llama 3.2 3B (Recommended)": "meta-llama/Llama-3.2-3B-Instruct",
+        "Llama 3.2 1B (Lightweight)": "meta-llama/Llama-3.2-1B-Instruct",
+        "Llama 3.3 70B (Best Quality)": "meta-llama/Llama-3.3-70B-Instruct",
+        "Fitness Assistant (Specialized)": "Soorya03/Llama-3.2-1B-FitnessAssistant",
+        "Mistral-7B (Original)": "mistralai/Mistral-7B-Instruct-v0.2"
+    }
+    
+    selected_model = st.selectbox(
+        "Select Model",
+        options=list(model_options.keys()),
+        index=0
+    )
+    
+    # Mode selection (API vs Local)
+    use_ollama = st.checkbox("Use Local Ollama (No API key needed)", value=False)
+    
+    if st.button("Apply Settings"):
+        switch_model(model_options[selected_model])
+        set_ollama_mode(use_ollama)
+        # Retest connection with new settings
+        api_ok, api_message = test_api_connection()
+        st.session_state.api_status = api_ok
+        st.session_state.api_message = api_message
+        st.success(f"Switched to {selected_model}")
+        st.rerun()
+    
+    if use_ollama:
+        st.info("📌 Make sure Ollama is running locally: 'ollama run llama3.2:3b'")
+    
+    # Show current model info
+    st.markdown("---")
+    st.caption(f"Current: {selected_model}")
 
 # Initialize session state
 if "form_submitted" not in st.session_state:
@@ -203,7 +242,7 @@ with tab2:
     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.error("⚠️ API is not connected. Please check your HF_TOKEN environment variable or use local Ollama.")
             st.stop()
         
         # Show profile summary
@@ -223,7 +262,7 @@ with tab2:
         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."):
+                with st.spinner("Creating your personalized workout plan with Llama... 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
@@ -234,12 +273,12 @@ with tab2:
             st.markdown("---")
             st.subheader("📋 Your 5-Day Workout Plan")
             
-            # Check if plan is complete - Now verify_workout_plan is defined
+            # 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
+            # Display the workout plan in a nice container
             with st.container():
                 st.markdown(st.session_state.workout_plan)
             
@@ -254,6 +293,7 @@ with tab2:
 FIT PLAN AI - PERSONALIZED WORKOUT PLAN
 Generated for: {st.session_state.name}
 Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+Model: {selected_model}
 
 PROFILE SUMMARY:
 - Age: {st.session_state.age}
@@ -279,7 +319,7 @@ Generated by FIT Plan AI
             
             with col2:
                 if st.button("🔄 Regenerate Plan", use_container_width=True):
-                    with st.spinner("Regenerating your workout plan..."):
+                    with st.spinner("Regenerating your workout plan with Llama..."):
                         workout_plan = query_model(st.session_state.prompt)
                         st.session_state.workout_plan = workout_plan
                         st.rerun()
@@ -309,4 +349,4 @@ Generated by FIT Plan AI
 
 # Footer
 st.markdown("---")
-st.caption("FIT Plan AI - Milestone 2 | Powered by Mistral-7B")
+st.caption("FIT Plan AI - Milestone 2 | Powered by Llama Models")