Parse both trace formats: Claude-Code (sessions/<date>/) and pi-sessions (root-level message events with toolCall/toolResult roles)

analyze.py

@@ -8,6 +8,8 @@ from concurrent.futures import ThreadPoolExecutor
 
 from huggingface_hub import InferenceClient
 
+from extract import event_role, event_tool_names
+
 MODEL = "Qwen/Qwen3.6-35B-A3B"
 
 _NO_THINK = {"chat_template_kwargs": {"enable_thinking": False}}
@@ -90,27 +92,21 @@ def _parse_ts(ts: str) -> dt.datetime | None:
 
 
 def compute_stats(sessions: list[tuple[str, list[dict]]]) -> dict:
-    """Count user turns, distinct tool names, and the first→last timestamp span."""
+    """Count user turns, distinct tool names, and the first→last timestamp span.
+    Format-agnostic (Claude-Code style and pi-sessions style both handled)."""
     turns = 0
     tools: set[str] = set()
     timestamps: list[dt.datetime] = []
     for _path, events in sessions:
         for ev in events:
-            if ev.get("type") == "user":
+            if event_role(ev) == "user":
                 turns += 1
             ts = ev.get("timestamp")
             if isinstance(ts, str):
                 parsed = _parse_ts(ts)
                 if parsed:
                     timestamps.append(parsed)
-            msg = ev.get("message") or {}
-            content = msg.get("content")
-            if isinstance(content, list):
-                for block in content:
-                    if isinstance(block, dict) and block.get("type") == "tool_use":
-                        name = block.get("name")
-                        if isinstance(name, str) and name:
-                            tools.add(name)
+            tools.update(event_tool_names(ev))
 
     span = ""
     if timestamps:

dataset.py

@@ -1,4 +1,10 @@
-"""Hugging Face Hub I/O for agent-trace JSONL session files."""
+"""Hugging Face Hub I/O for agent-trace JSONL session files.
+
+Accepts both layouts:
+- `sessions/<YYYY-MM-DD>/<uuid>.jsonl` (e.g. `merve/ml-intern-sessions`)
+- `<YYYY-MM-DDTHH-MM-SS>_<uuid>.jsonl` at the dataset root or under subdirs
+  (e.g. `julien-c/pi-sessions`)
+"""
 
 import json
 import re
@@ -6,14 +12,35 @@ from pathlib import Path
 
 from huggingface_hub import HfApi, hf_hub_download
 
-_SESSION_RE = re.compile(r"^sessions/(\d{4}-\d{2}-\d{2})/[^/]+\.jsonl$")
+_DATE_RE = re.compile(r"(\d{4}-\d{2}-\d{2})")
+
+
+def _sort_key(path: str) -> tuple:
+    """Sortable key: (date prefix found anywhere in path, full path)."""
+    m = _DATE_RE.search(path)
+    return (m.group(1) if m else "0000-00-00", path)
 
 
 def list_sessions(repo_id: str) -> list[str]:
-    """Return JSONL session paths from a dataset repo, newest date first."""
+    """Return JSONL session paths from a dataset repo, newest first.
+
+    Any `.jsonl` file in the repo is treated as a session candidate. Duplicates
+    that share a basename (some repos mirror the same file under subdirs) are
+    deduped, keeping the first occurrence.
+    """
     info = HfApi().dataset_info(repo_id)
-    paths = [s.rfilename for s in info.siblings if _SESSION_RE.match(s.rfilename)]
-    paths.sort(key=lambda p: _SESSION_RE.match(p).group(1), reverse=True)
+    seen: set[str] = set()
+    paths: list[str] = []
+    for s in info.siblings:
+        name = s.rfilename
+        if not name.endswith(".jsonl"):
+            continue
+        base = name.rsplit("/", 1)[-1]
+        if base in seen:
+            continue
+        seen.add(base)
+        paths.append(name)
+    paths.sort(key=_sort_key, reverse=True)
     return paths
 
 

extract.py

@@ -1,8 +1,49 @@
-"""Pure transforms on agent-trace event lists. No I/O."""
+"""Pure transforms on agent-trace event lists. No I/O.
+
+Supports two on-disk formats:
+1. Claude-Code style — `{type: "user"|"assistant", message: {role, content}}`.
+   Example dataset: `merve/ml-intern-sessions`.
+2. pi-sessions style — `{type: "message", message: {role: "user"|"assistant"|"toolResult", content: [...]}}`.
+   Example dataset: `julien-c/pi-sessions`. Tool calls use `toolCall` blocks;
+   tool outputs come back as role=toolResult messages which we drop.
+"""
 
 from typing import Any
 
 
+def event_role(ev: dict) -> str | None:
+    """Normalised role of a trace event. Returns 'user' / 'assistant' or None
+    for non-content events (session metadata, tool results, model_change, etc.)."""
+    t = ev.get("type")
+    if t in ("user", "assistant"):
+        return t
+    if t == "message":
+        msg = ev.get("message") or {}
+        role = msg.get("role")
+        if role in ("user", "assistant"):
+            return role
+        return None
+    return None
+
+
+def event_tool_names(ev: dict) -> list[str]:
+    """Extract tool names invoked in this event, across both formats."""
+    msg = ev.get("message") or {}
+    content = msg.get("content")
+    if not isinstance(content, list):
+        return []
+    out: list[str] = []
+    for block in content:
+        if not isinstance(block, dict):
+            continue
+        bt = block.get("type")
+        if bt in ("tool_use", "toolCall"):
+            name = block.get("name")
+            if isinstance(name, str) and name:
+                out.append(name)
+    return out
+
+
 def _user_content_to_text(content: Any) -> str:
     if isinstance(content, str):
         return content
@@ -22,6 +63,7 @@ def _user_content_to_text(content: Any) -> str:
 
 
 def _assistant_content_to_text(content: Any) -> str:
+    """Concatenate text blocks; drop thinking / tool_use / toolCall blocks."""
     if isinstance(content, str):
         return content
     if isinstance(content, list):
@@ -36,16 +78,19 @@ def _assistant_content_to_text(content: Any) -> str:
 def events_to_transcript(events: list[dict]) -> str:
     lines: list[str] = []
     for ev in events:
+        role = event_role(ev)
+        if role not in ("user", "assistant"):
+            continue
         msg = ev.get("message") or {}
         content = msg.get("content")
-        if ev.get("type") == "user":
-            text = _user_content_to_text(content).strip()
-            if text:
-                lines.append(f"User: {text}")
-        elif ev.get("type") == "assistant":
-            text = _assistant_content_to_text(content).strip()
-            if text:
-                lines.append(f"Assistant: {text}")
+        text = (
+            _user_content_to_text(content)
+            if role == "user"
+            else _assistant_content_to_text(content)
+        ).strip()
+        if text:
+            label = "User" if role == "user" else "Assistant"
+            lines.append(f"{label}: {text}")
     return "\n\n".join(lines)