fix: path to session.jsonl

Signed-off-by: Rohan R. Arora <rohan.arora@ibm.com>

analysis_src/utils.py

@@ -1,4 +1,4 @@
-analysis_src/utils.pyimport json
+import json
 from pathlib import Path
 
 # Model display names (short for figures)
@@ -152,21 +152,24 @@ def filter_scenarios_with_min_runs(scenario_data: dict[str, list], min_runs_requ
     }
 
 def find_latest_rollout_file(trial_dir: Path) -> Path:
-    """Find the latest rollout file in a trial's sessions directory."""
+    """Find the latest rollout file in a trial's sessions directory or session.jsonl."""
+    # First check for session.jsonl directly in trial_dir (new structure)
+    session_file = trial_dir / "session.jsonl"
+    if session_file.exists():
+        return session_file
+
+    # Fall back to sessions/rollout-*.jsonl (old structure)
     sessions_dir = trial_dir / "sessions"
     if not sessions_dir.exists():
         return None
-    
+
     rollout_files = []
     for rollout_file in sessions_dir.rglob("rollout-*.jsonl"):
         rollout_files.append(rollout_file)
-    
+
     if not rollout_files:
         return None
-    
-    # Sort by filename (which includes timestamp) and return the latest
-    # rollout_files.sort(key=lambda f: f.name, reverse=True)
-    # return rollout_files[0]
 
+    # Sort by modification time and return the latest
     return max(rollout_files, key=lambda p: p.stat().st_mtime)