Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 
 from fastapi import FastAPI
 from pydantic import BaseModel
+from typing import List 
 from supabase import create_client
 
 import torch
@@ -64,6 +65,15 @@ class MultiHeadDistilBERT(nn.Module):
 class PredictRequest(BaseModel):
     user_input: str
 
+class ExerciseResponse(BaseModel):
+    id:                int
+    name:              str
+    workout_type:      int
+    difficulty:        str
+    notes:             str
+    suitable_moods:    List[int]
+    suitable_soreness: List[int]
+    
 class PredictResponse(BaseModel):
     workout: str
     workout_conf: float
@@ -71,27 +81,21 @@ class PredictResponse(BaseModel):
     mood_conf: float
     soreness: str
     soreness_conf: float
+    exercises: List[ExerciseResponse]
+
 
 def get_suitable_exercises(workout_type: int, mood: int, soreness: int):
     supabase = create_client(os.getenv('SUPA_URL'), os.getenv('SUPA_KEY'))
 
-    mood = str(mood)
-    soreness = str(soreness)
-
-    print(f"Workout type {workout_type}")
-    print(f"mood type {mood}")
-
-    print(f"sorness type {soreness}")
-
     response = (
         supabase.table('exerciseai')
         .select('*')
         .eq('workout_type', workout_type)
-        .contains('suitable_moods',    [mood])
-        .contains('suitable_soreness', [soreness])
+        .contains('suitable_moods',    [str(mood)])
+        .contains('suitable_soreness', [str(soreness)])
         .execute()
     )
-    print(f"Response data is {response.data}")
+    
     return response.data
     
 @app.get("/")