Update app.py

app.py

@@ -17,7 +17,7 @@ import logging
 from model import PostWorkoutDistilBERT, load_post_model, PreWorkoutDistilBERT, load_pre_model
 from inference import (
     predict_post, decode_post_predictions, build_post_prompt,
-    predict_pre, decode_pre_predictions, build_pre_prompt,
+    predict_pre, decode_pre_predictions, build_pre_prompt, parse_workout_plan,
 )
 
 # ─────────────────────────────────────────────
@@ -97,7 +97,7 @@ async def lifespan(app: FastAPI):
     # ── Post-workout model ────────────────────────────────────
     logger.info("Loading post-workout model...")
     post_model, post_tokenizer = load_post_model(
-        model_path=os.getenv("MODEL_PATH", "post_best_overall_model.pt"),
+        model_path=os.getenv("MODEL_PATH", "best_overall_model.pt"),
         device=device,
     )
     app_state["post_model"]     = post_model
@@ -198,7 +198,7 @@ def health():
     }
 
 
-@app.post("/post-classify", response_model=PostWorkoutSessionResponse,
+@app.post("/classify", response_model=PostWorkoutSessionResponse,
           dependencies=[Depends(verify_key)])
 def post_classify_session(req: PostWorkoutRequest):
     """
@@ -255,7 +255,7 @@ def post_classify_session(req: PostWorkoutRequest):
     )
 
 
-@app.post("/post-classify/labels-only", response_model=PostWorkoutBertLabels,
+@app.post("/classify/labels-only", response_model=PostWorkoutBertLabels,
           dependencies=[Depends(verify_key)])
 def post_classify_labels_only(req: PostWorkoutRequest):
     """
@@ -308,8 +308,12 @@ class PreWorkoutBertLabels(BaseModel):
 
 
 class PreWorkoutResponse(BaseModel):
-    bert_labels:  PreWorkoutBertLabels
-    workout_plan: Optional[str] = None
+    bert_labels:   PreWorkoutBertLabels
+    warm_up:       Optional[str] = None
+    main_workout:  Optional[str] = None
+    cool_down:     Optional[str] = None
+    coaching_note: Optional[str] = None
+    raw_plan:      Optional[str] = None   # full unmodified response — fallback
 
 
 # ─────────────────────────────────────────────
@@ -348,7 +352,7 @@ def pre_classify_session(req: PreWorkoutRequest):
     )
 
     # ── Step 3: Optionally generate Claude workout plan ───────
-    workout_plan = None
+    parsed = {}
     if req.generate_plan:
         prompt = build_pre_prompt(
             bert_labels=bert_labels,
@@ -364,14 +368,19 @@ def pre_classify_session(req: PreWorkoutRequest):
                 max_tokens=800,
                 messages=[{"role": "user", "content": prompt}],
             )
-            workout_plan = message.content[0].text
+            raw_plan = message.content[0].text
+            parsed   = parse_workout_plan(raw_plan)
         except Exception as e:
             logger.error(f"Claude API error (pre-workout): {e}")
-            workout_plan = None
+            parsed = {}
 
     return PreWorkoutResponse(
         bert_labels=PreWorkoutBertLabels(**bert_labels),
-        workout_plan=workout_plan,
+        warm_up=parsed.get("warm_up"),
+        main_workout=parsed.get("main_workout"),
+        cool_down=parsed.get("cool_down"),
+        coaching_note=parsed.get("coaching_note"),
+        raw_plan=parsed.get("raw"),
     )