Update model_api.py

model_api.py

@@ -2,25 +2,91 @@ from huggingface_hub import InferenceClient
 import os
 
 def query_model(prompt):
+    """
+    Query the Mistral-7B model with the given prompt
+    """
     try:
         HF_TOKEN = os.getenv("HF_TOKEN")
-
+        
+        if not HF_TOKEN:
+            return "Error: HF_TOKEN not found. Please set your Hugging Face token in environment variables."
+        
+        # Initialize the client
         client = InferenceClient(
             model="mistralai/Mistral-7B-Instruct-v0.2",
             token=HF_TOKEN
         )
-
+        
+        # Enhanced system prompt for better responses
+        system_prompt = """You are a certified professional fitness trainer with expertise in creating personalized workout plans. 
+        Always provide complete, detailed workout plans with:
+        - Clear day-by-day structure
+        - Specific exercises with sets, reps, and rest periods
+        - Warm-up and cool-down recommendations
+        - Safety considerations based on user's profile
+        When asked for a 5-day plan, ensure ALL 5 days are included with clear day headers."""
+        
+        # Make the API call
         response = client.chat_completion(
             messages=[
-                {"role": "system", "content": "You are a certified professional fitness trainer."},
+                {"role": "system", "content": system_prompt},
                 {"role": "user", "content": prompt}
             ],
-            max_tokens=2500,#BEFORE 600
-            temperature=0.7
+            max_tokens=2500,  # Increased for complete 5-day plan
+            temperature=0.7,
+            top_p=0.95
         )
+        
+        # Extract and return the response
+        workout_plan = response.choices[0].message.content
+        
+        # Verify if the response contains all 5 days
+        days_found = sum([f"Day {i}" in workout_plan for i in range(1, 6)])
+        
+        if days_found < 5:
+            # If incomplete, try one more time with more explicit instruction
+            retry_prompt = prompt + "\n\nIMPORTANT: The previous response was incomplete. Please ensure ALL 5 days (Day 1 through Day 5) are included in the plan. Each day should be clearly marked with 'Day X' header and include 4-6 exercises."
+            
+            retry_response = client.chat_completion(
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": retry_prompt}
+                ],
+                max_tokens=2500,
+                temperature=0.7
+            )
+            workout_plan = retry_response.choices[0].message.content
+        
+        return workout_plan
+        
+    except Exception as e:
+        return f"Error generating workout plan: {str(e)}"
 
-        return response.choices[0].message.content
-
+def test_api_connection():
+    """
+    Test function to verify API connection
+    """
+    try:
+        HF_TOKEN = os.getenv("HF_TOKEN")
+        if not HF_TOKEN:
+            return False, "HF_TOKEN not found"
+        
+        client = InferenceClient(
+            model="mistralai/Mistral-7B-Instruct-v0.2",
+            token=HF_TOKEN
+        )
+        
+        # Simple test prompt
+        response = client.chat_completion(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Say 'API connection successful' if you can read this."}
+            ],
+            max_tokens=50,
+            temperature=0.1
+        )
+        
+        return True, "API connection successful"
+        
     except Exception as e:
-        return f"Error: {str(e)}"
-        
+        return False, f"API connection failed: {str(e)}"