Update app.py

app.py

@@ -16,7 +16,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_post, decode_post_predictions, build_post_prompt, parse_debrief,
     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
@@ -171,8 +171,11 @@ class PostWorkoutBertLabels(BaseModel):
 
 
 class PostWorkoutSessionResponse(BaseModel):
-    bert_labels: PostWorkoutBertLabels
-    debrief:     Optional[str] = None
+    bert_labels:     PostWorkoutBertLabels
+    acknowledgement: Optional[str] = None
+    highlights:      Optional[str] = None
+    next_session:    Optional[str] = None
+    raw_debrief:     Optional[str] = None   # full unmodified response — fallback
 
 
 class HealthResponse(BaseModel):
@@ -229,7 +232,7 @@ def post_classify_session(req: PostWorkoutRequest):
     )
 
     # ── Step 3: Optionally generate Claude debrief ────────────
-    debrief = None
+    parsed = {}
     if req.generate_debrief:
         prompt = build_post_prompt(
             bert_labels=bert_labels,
@@ -244,14 +247,18 @@ def post_classify_session(req: PostWorkoutRequest):
                 max_tokens=400,
                 messages=[{"role": "user", "content": prompt}],
             )
-            debrief = message.content[0].text
+            raw_debrief = message.content[0].text
+            parsed      = parse_debrief(raw_debrief)
         except Exception as e:
             logger.error(f"Claude API error (post-workout): {e}")
-            debrief = None
+            parsed = {}
 
     return PostWorkoutSessionResponse(
         bert_labels=PostWorkoutBertLabels(**bert_labels),
-        debrief=debrief,
+        acknowledgement=parsed.get("acknowledgement"),
+        highlights=parsed.get("highlights"),
+        next_session=parsed.get("next_session"),
+        raw_debrief=parsed.get("raw"),
     )