Returning suitable exercises

app.py

@@ -2,6 +2,7 @@ import os
 
 from fastapi import FastAPI
 from pydantic import BaseModel
+from supabase import create_client
 
 import torch
 import torch.nn as nn
@@ -70,6 +71,19 @@ class PredictResponse(BaseModel):
     mood_conf: float
     soreness: str
     soreness_conf: float
+
+def get_suitable_exercises(workout_type: int, mood: int, soreness: int):
+    supabase = create_client(os.getenv('SUPA_URL'), os.getenv('SUPA_KEY'))
+    
+    response = (
+        supabase.table('exercises')
+        .select('*')
+        .eq('workout_type', workout_type)
+        .contains('suitable_moods',    [mood])
+        .contains('suitable_soreness', [soreness])
+        .execute()
+    )
+    return response.data
     
 @app.get("/")
 def greet_json():
@@ -119,7 +133,12 @@ def predict(request: PredictRequest):
     predicted_mood = mood_label_map[mood_logits.argmax().item()]
     predicted_soreness = soreness_label_map[soreness_logits.argmax().item()]
 
-    
+    # Fetch suitable exercises from Supabase
+    suitable_exercises = get_suitable_exercises(
+        workout_type = workout_logits.argmax().item(),
+        mood         = mood_logits.argmax().item(),
+        soreness     = soreness_logits.argmax().item()
+    )    
     return PredictResponse(
         workout = predicted_workout,
         workout_conf = round(workout_conf.item()  * 100, 1),
@@ -127,5 +146,6 @@ def predict(request: PredictRequest):
         mood_conf = round(mood_conf.item()  * 100, 1),
         soreness = predicted_soreness,
         soreness_conf = round(soreness_conf.item() * 100, 1),
+        exercises = suitable_exercises
     )