Update main.py

main.py

@@ -9,6 +9,7 @@ import asyncio
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from typing import Optional
 from models.model import (
     load_models,
     run_inference,
@@ -56,8 +57,6 @@ app.add_middleware(
 
 # ─────────────────────────────────────────────
 # REQUEST / RESPONSE SCHEMAS
-# Pydantic models define exactly what SwiftUI
-# sends and receives — easy to match in Swift
 # ─────────────────────────────────────────────
 
 class UserProfile(BaseModel):
@@ -66,8 +65,11 @@ class UserProfile(BaseModel):
 
 
 class PredictRequest(BaseModel):
-    text:         str           # raw user post-workout summary text
-    user_profile: UserProfile
+    text:             str               # raw user post-workout text — passed to BERT only
+    duration_minutes: int               # actual session duration
+    workout_type:     str               # e.g. "strength", "cardio", "mobility"
+    user_profile:     UserProfile
+    planned_workout:  Optional[str] = None  # Claude output from pre-workout coach, if used
 
 
 class LabelResult(BaseModel):
@@ -90,45 +92,51 @@ class MentalOutput(BaseModel):
 
 
 class PredictResponse(BaseModel):
-    physical:       PhysicalOutput      # physical BERT output
-    mental:         MentalOutput        # mental BERT output
-    coach_response: str                 # Claude's workout summary + tips
-    enriched_text:  str                 # text sent to BERT (for debugging)
+    physical:       PhysicalOutput
+    mental:         MentalOutput
+    coach_response: str             # Claude's workout summary + tips
+    enriched_text:  str             # text sent to BERT (for debugging)
 
 
 # ─────────────────────────────────────────────
 # PROMPT BUILDER
 # ─────────────────────────────────────────────
 
-def build_claude_prompt(goal, modifiers, physical, mental, planned_workout):
-    planned_section = f"Planned session: {planned_workout}" if planned_workout else "Planned session: none — user worked out independently"
+def build_claude_prompt(goal, modifiers, duration_minutes, workout_type, physical, mental, planned_workout):
+    planned_section = (
+        f"Planned session: {planned_workout}"
+        if planned_workout
+        else "Planned session: none — user worked out independently"
+    )
 
     return f"""
 You are a personal fitness coach reviewing a completed training session with your athlete.
 
-User goal    : {goal}
-Constraints  : {', '.join(modifiers) if modifiers else 'none'}
+User goal       : {goal}
+Constraints     : {', '.join(modifiers) if modifiers else 'none'}
+Session duration: {duration_minutes} minutes
+Workout type    : {workout_type}
 {planned_section}
 
 Physical state:
-- Pain        : {physical['pain_label']['label']}
-- Completion  : {physical['completion_label']['label']}
-- Fatigue     : {physical['fatigue_label']['label']}
-- Recovery need: {physical['recovery_need_label']['label']}
+- Pain          : {physical['pain_label']['label']}
+- Completion    : {physical['completion_label']['label']}
+- Fatigue       : {physical['fatigue_label']['label']}
+- Recovery need : {physical['recovery_need_label']['label']}
 
 Mental state:
-- Performance : {mental['performance_label']['label']}
-- Satisfaction: {mental['satisfaction_label']['label']}
-- PR achieved : {mental['pr_achieved_label']['label']}
-- Motivation  : {mental['motivation_label']['label']}
-
-REASONING RULES:
-1. If no planned session was provided, skip the completion comparison entirely. 
-   Do not reference or imply there was a plan. Focus only on what the labels 
+- Performance   : {mental['performance_label']['label']}
+- Satisfaction  : {mental['satisfaction_label']['label']}
+- PR achieved   : {mental['pr_achieved_label']['label']}
+- Motivation    : {mental['motivation_label']['label']}
+
+REASONING RULES — apply before writing:
+1. If no planned session was provided, skip the completion comparison entirely.
+   Do not reference or imply there was a plan. Focus only on what the labels
    reveal about how the session went.
 2. PR handling: if pr_achieved is true, open with a genuine, specific celebration tied to the user's goal — not just "great job." Explain what this PR signals about their progress.
 3. Recovery routing:
-   - recovery_need = rest_day → tomorrow's advice must be strictly non-training: sleep, nutrition, or mobility only. Do not suggest any workout.
+   - recovery_need = rest_day  → tomorrow's advice must be strictly non-training: sleep, nutrition, or mobility only. Do not suggest any workout.
    - recovery_need = light_day → suggest one low-intensity option only (walk, yoga, or easy swim).
    - recovery_need = normal or ready → give a forward-looking tip about the next session.
 4. Pain flag: if pain label is moderate or high, add a one-line note recommending the user monitor that area and consider professional advice if it persists.
@@ -153,33 +161,36 @@ async def predict(request: PredictRequest):
     Main endpoint — called by SwiftUI after user submits post-workout summary.
 
     Flow:
-      1. Enriches text with goal
-      2. Runs physical BERT model
-      3. Runs mental BERT model
-      4. Builds Claude prompt from all labels
-      5. Calls Claude API
-      6. Returns classifications + coaching response to SwiftUI
+      1. Enriches user text with goal (for BERT only)
+      2. Runs physical and mental BERT models concurrently
+      3. Builds Claude prompt from labels + session metadata + optional planned workout
+      4. Calls Claude API
+      5. Returns classifications + coaching response to SwiftUI
     """
     try:
-        # 1. Enrich text with goal
+        # 1. Enrich text with goal for BERT — duration, workout_type, and
+        #    planned_workout are kept separate so they don't pollute the BERT
+        #    input distribution
         enriched_text = f"Goal: {request.user_profile.primary_goal}. {request.text}"
 
-        # 2. Run physical and mental model using asyncio
+        # 2. Run physical and mental models concurrently
         physical_output, mental_output = await asyncio.gather(
             asyncio.to_thread(run_inference, ml_models["physical_model"], ml_models["tokenizer"], enriched_text, PHYSICAL_LABEL_COLS, PHYSICAL_DECODERS),
             asyncio.to_thread(run_inference, ml_models["mental_model"],   ml_models["tokenizer"], enriched_text, MENTAL_LABEL_COLS,   MENTAL_DECODERS)
         )
 
-        # 3. Build Claude prompt
+        # 3. Build Claude prompt with all available context
         prompt = build_claude_prompt(
             request.user_profile.primary_goal,
             request.user_profile.modifiers,
+            request.duration_minutes,
+            request.workout_type,
             physical_output,
-            mental_output
+            mental_output,
+            request.planned_workout
         )
 
         # 4. Call Claude API
-        # ANTHROPIC_API_KEY must be set as a HuggingFace Space secret
         claude_client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
         message       = claude_client.messages.create(
             model="claude-sonnet-4-20250514",