AIM-4180: Wave 3 — Run pattern extraction on 50K traces + generate enhanced skills

analysis/loader.py

@@ -16,20 +16,10 @@
     AnalysisConfig,
 )
 
-# ── Trace type ───────────────────────────────────────────────────
-
 TraceDict = Dict[str, Any]
-"""A single parsed trace with fields: uid, source_file, session, model,
-context, cot, output_type, output, completion, origin."""
-
-# ── Dataset loading ──────────────────────────────────────────────
 
 
 def get_dataset_info(name: str) -> dict[str, Any]:
-    """Check if a dataset exists and return its metadata.
-
-    Raises RuntimeError if the dataset is not accessible.
-    """
     try:
         builder = load_dataset_builder(name)
         return {
@@ -46,14 +36,8 @@ def get_dataset_info(name: str) -> dict[str, Any]:
 def load_dataset_simple(
     config: AnalysisConfig,
 ) -> Iterator[Dataset | IterableDataset]:
-    """Load dataset from HuggingFace with streaming.
-
-    Returns an iterator over batched Dataset slices (each of size
-    config.batch_size). Uses streaming to avoid loading the full
-    2M-row dataset into memory.
-    """
-    import time
-    start = time.time()
+    import time as _time
+    start = _time.time()
     dataset_name = config.resolve_dataset()
 
     try:
@@ -63,11 +47,10 @@ def load_dataset_simple(
             streaming=False,
             cache_dir=config.cache_dir,
         )
-        elapsed_s = time.time() - start
-        if isinstance(ds, Dataset):
-            import rich
-            rich.print(f"[dim]Dataset loaded: {len(ds)} rows in {elapsed_s:.0f}s (non-streaming)[/]")
-            return _iter_batches(ds, config)
+        elapsed_s = _time.time() - start
+        import rich
+        rich.print(f"[dim]Dataset loaded: {len(ds)} rows in {elapsed_s:.0f}s (non-streaming)[/]")
+        return _iter_batches_nonstreaming(ds, config)
     except Exception:
         pass
 
@@ -85,74 +68,66 @@ def load_dataset_simple(
             f"Failed to load dataset '{dataset_name}': {exc}"
         ) from exc
 
-    return _iter_batches(ds, config)
+    return _iter_batches_streaming(ds, config)
 
 
 def _try_fallback(
     config: AnalysisConfig,
 ) -> Iterator[Dataset | IterableDataset]:
-    """Attempt to load from the fallback dataset."""
     try:
         ds = load_dataset(
             FALLBACK_DATASET,
             split="train",
             streaming=True,
             cache_dir=config.cache_dir,
         )
-        return _iter_batches(ds, config)
+        return _iter_batches_streaming(ds, config)
     except Exception as exc:
         raise RuntimeError(
             f"Primary and fallback datasets unavailable: {exc}"
         ) from exc
 
 
-def _iter_batches(
+def _iter_batches_nonstreaming(
+    ds: Dataset,
+    config: AnalysisConfig,
+) -> Generator[Dataset, None, None]:
+    total = len(ds)
+    n = min(total, config.max_samples) if config.max_samples > 0 else total
+    for i in range(0, n, config.batch_size):
+        end = min(i + config.batch_size, n)
+        yield ds.select(range(i, end))
+
+
+def _iter_batches_streaming(
     ds: IterableDataset,
     config: AnalysisConfig,
 ) -> Generator[Dataset, None, None]:
-    """Yield rows from iterable dataset in batches of batch_size."""
+    from datasets import Dataset as BatchDataset
     batch: list[dict[str, Any]] = []
     count = 0
 
     for row in ds:
-        batch.append(row)  # type: ignore[arg-type]
+        batch.append(row)
         if len(batch) >= config.batch_size:
-            # Build a temporary Dataset slice
-            from datasets import Dataset as BatchDataset  # noqa: PLC0415
-
-            yield BatchDataset.from_list(batch)  # type: ignore[no-untyped-call]
+            yield BatchDataset.from_list(batch)
             batch.clear()
             count += config.batch_size
 
         if config.max_samples > 0 and count + len(batch) >= config.max_samples:
             break
 
     if batch:
-        from datasets import Dataset as BatchDataset  # noqa: PLC0415
-
-        yield BatchDataset.from_list(batch)  # type: ignore[no-untyped-call]
-
-
-# ── Trace extraction ─────────────────────────────────────────────
+        yield BatchDataset.from_list(batch)
 
 
 def extract_trace(row: dict[str, Any], is_wrapper: bool) -> TraceDict | None:
-    """Extract a parsed trace dict from a single row.
-
-    Handles both wrapper format (Crownelius) where trace data is
-    a JSON string in the ``row_json`` column, and raw format
-    (Glint-Research) where fields are direct columns.
-
-    Returns None if the row cannot be parsed (malformed JSON, missing
-    required fields).
-    """
     if is_wrapper:
         return _extract_wrapper(row)
     return _extract_raw(row)
 
 
 def _extract_wrapper(row: dict[str, Any]) -> TraceDict | None:
-    """Parse a wrapper-format row (Crownelius)."""
     raw_json = row.get("row_json")
     if not raw_json:
         return None
@@ -167,23 +142,16 @@ def _extract_wrapper(row: dict[str, Any]) -> TraceDict | None:
     else:
         return None
 
-    # Normalize field names — row_json may use different casing
     return _normalize_trace(parsed)
 
 
 def _extract_raw(row: dict[str, Any]) -> TraceDict | None:
-    """Parse a raw-format row (Glint-Research)."""
     return _normalize_trace(row)
 
 
 def _normalize_trace(data: dict[str, Any]) -> TraceDict | None:
-    """Map various field name conventions to canonical names.
-
-    The row_json field may have keys like 'cot', 'CoT', 'chain_of_thought',
-    'context', 'Context', 'instruction', etc.
-    """
     aliases: dict[str, list[str]] = {
-        "uid": ["uid", "id", "trace_id", "ID"],
+        "uid": ["uid", "id", "trace_id", "ID", "leafUuid", "leaf_uuid"],
         "source_file": ["source_file", "source", "file"],
         "session": ["session", "session_id", "Session"],
         "model": ["model", "Model", "model_name", "model_id"],
@@ -239,40 +207,22 @@ def _normalize_trace(data: dict[str, Any]) -> TraceDict | None:
                 result[canonical] = data[key]
                 break
 
-    if not result.get("uid"):
-        return None
-
-    # Ensure cot exists (even if empty)
     if "cot" not in result:
         result["cot"] = ""
 
-    return result
-
-
-# ── Generator API ────────────────────────────────────────────────
+    return result if result.get("uid") or result.get("cot") is not None else None
 
 
 def iter_traces(
     config: AnalysisConfig | None = None,
 ) -> Generator[TraceDict, None, None]:
-    """Generator that yields parsed trace dicts one at a time.
-
-    Handles both wrapper and raw formats automatically based on the
-    dataset configured. Uses streaming to avoid loading the full
-    dataset into memory.
-
-    Usage::
-
-        for trace in iter_traces():
-            print(trace["cot"][:100])
-    """
     cfg = config or AnalysisConfig()
     is_wrapper = cfg.is_wrapper_format
     count = 0
 
     for batch in load_dataset_simple(cfg):
-        for row in batch:  # type: ignore[union-attr]
-            trace = extract_trace(row, is_wrapper)  # type: ignore[arg-type]
+        for row in batch:
+            trace = extract_trace(row, is_wrapper)
             if trace is not None:
                 yield trace
                 count += 1
@@ -282,10 +232,6 @@ def iter_traces(
 
 
 def count_traces(config: AnalysisConfig | None = None) -> int:
-    """Count available traces in the dataset (up to max_samples).
-
-    Runs as a lightweight pass without storing results.
-    """
     count = 0
     for _ in iter_traces(config):
         count += 1

analysis/patterns/architect_patterns.yaml

@@ -1,10 +1,10 @@
 skill: architect
-total_traces: 80
+total_traces: 271
 stats:
   cot:
-    total_traces: 80
+    total_traces: 271
     cot_present: 80
-    cot_rate: 1.0
+    cot_rate: 0.2952
     avg_tokens: 368.29
     avg_paragraphs: 5.54
     avg_sentences: 16.25
@@ -32,20 +32,38 @@ stats:
     - because
     - hence
   tool_usage:
-    total_traces: 80
-    traces_with_tools: 0
+    total_traces: 271
+    traces_with_tools: 70
     tool_calls_per_trace:
-      '0': 1.0
-    tool_type_frequency: {}
-    top_tool_calls: []
+      '0': 0.7417
+      '1': 0.2583
+    tool_type_frequency:
+      Bash: 0.1439
+      Read: 0.0923
+      Search: 0.0074
+      Web: 0.0074
+      Glob: 0.0074
+    top_tool_calls:
+    - 'sh '
+    - read
+    - shell
+    - view
+    - 'ls '
+    - terminal
+    - search
+    - http
+    - look at
+    - web
+    - glob
+    - bash
     transition_matrix: {}
     read_before_edit_rate: 0.0
     verify_after_action_rate: 0.0
-    tool_to_text_ratio: 0.0
-    avg_tool_calls: 0
-    max_tool_calls: 0
+    tool_to_text_ratio: 0.7692
+    avg_tool_calls: 0.26
+    max_tool_calls: 1
   behaviors:
-    total_traces: 80
+    total_traces: 271
     self_correction_rate: 0.925
     avg_self_corrections: 5.9625
     hypothesis_driven_rate: 0.425
@@ -86,7 +104,7 @@ stats:
 patterns:
 - name: common-openers
   description: 'Frequent utterance starters: The, Alright, I’ve, Okay, I need to'
-  frequency: 1.0
+  frequency: 0.2952
 - name: self-correction
   description: Frequently corrects reasoning mid-turn
   frequency: 0.925

analysis/patterns/code_patterns.yaml

@@ -1,10 +1,10 @@
 skill: code
-total_traces: 3203
+total_traces: 6835
 stats:
   cot:
-    total_traces: 3203
+    total_traces: 6835
     cot_present: 3203
-    cot_rate: 1.0
+    cot_rate: 0.4686
     avg_tokens: 413.82
     avg_paragraphs: 7.32
     avg_sentences: 17.08
@@ -35,20 +35,39 @@ stats:
     - therefore
     - given that
   tool_usage:
-    total_traces: 3203
-    traces_with_tools: 0
+    total_traces: 6835
+    traces_with_tools: 96
     tool_calls_per_trace:
-      '0': 1.0
-    tool_type_frequency: {}
-    top_tool_calls: []
+      '0': 0.986
+      '1': 0.014
+    tool_type_frequency:
+      Bash: 0.0067
+      Read: 0.006
+      Search: 0.0006
+      Web: 0.0003
+      Edit: 0.0003
+      Glob: 0.0001
+    top_tool_calls:
+    - 'sh '
+    - read
+    - view
+    - shell
+    - fetch
+    - search
+    - write
+    - bash
+    - grep
+    - execute
+    - 'ls '
+    - find
     transition_matrix: {}
     read_before_edit_rate: 0.0
     verify_after_action_rate: 0.0
-    tool_to_text_ratio: 0.0
-    avg_tool_calls: 0
-    max_tool_calls: 0
+    tool_to_text_ratio: 0.1486
+    avg_tool_calls: 0.01
+    max_tool_calls: 1
   behaviors:
-    total_traces: 3203
+    total_traces: 6835
     self_correction_rate: 0.9756
     avg_self_corrections: 6.1714
     hypothesis_driven_rate: 0.2966
@@ -113,7 +132,7 @@ stats:
 patterns:
 - name: common-openers
   description: 'Frequent utterance starters: Alright, The, Okay, I’ve, I need to'
-  frequency: 1.0
+  frequency: 0.4686
 - name: self-correction
   description: Frequently corrects reasoning mid-turn
   frequency: 0.9756