feat(agent): --distill flag recording LLM calls as a ChatML NDJSON fine-tuning dataset

docs/agent.md

@@ -217,6 +217,32 @@ endpoint; neither is required):
   the threshold turns that measurement into a terminal gate. See
   [Biolink validity](#biolink-validity) below.
 
+### Distilling a fine-tuning dataset (`--distill`)
+
+`--distill` (short: `-d`, `-dt`) records **every LLM call of the run** — the inner agent's
+multi-turn conversations, plus the judge and reflexion calls when those gates are enabled — as one
+ChatML JSON object per line, appended to `<state-dir>/distill/records.ndjson`:
+
+```json
+{"messages": [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}], "purpose": "agent", "model_id": "...", "pmc_id": "PMC11708054", "call_index": 0, "timestamp": "...", "token_usage": {"input_tokens": 0, "output_tokens": 0}}
+```
+
+The file is **append-only**: every `--distill` invocation keeps adding to the same dataset, so a
+corpus accumulates over many batches. The `messages` column is plain ChatML, which Unsloth Studio
+auto-detects on JSONL upload (no column mapping needed); the metadata columns (`purpose`,
+`pmc_id`, `call_index`, `timestamp`, `token_usage`) ride along for filtering — e.g. join on
+`pmc_id` against `state.json` to keep only `MAPPED` runs, or keep each run's highest `call_index`
+for the most complete conversation. Recording is zero-dependency and never breaks a run: a failed
+write is logged, not raised. `--distill` is not supported with `--optimize` (the GEPA path
+bypasses the recording seam).
+
+To convert the NDJSON into an on-disk Hugging Face dataset, use
+[`tablassert distill-export`](cli.md#distill-export) (requires the `[distill]` extra):
+
+```bash
+tablassert distill-export --distill-dir .tablassert/agent/distill --out ./hf-dataset
+```
+
 ### Biolink validity
 
 Coverage answers *did the terms resolve?* It says nothing about whether the resulting records are

docs/cli.md

@@ -1,7 +1,7 @@
 # CLI Reference
 
 Tablassert extracts knowledge assertions from tabular data into KGX NDJSON. The `tablassert` app
-exposes **five subcommands**: `agent`, `build-fullmap`, `build-kg`, `validate`,
+exposes **six subcommands**: `agent`, `build-fullmap`, `build-kg`, `distill-export`, `validate`,
 and `validate-kgx`, plus an app-level `--version` flag. Run `tablassert --help` (or `<command> --help`)
 for the live surface.
 
@@ -12,6 +12,7 @@ for the live surface.
 | [`agent`](#agent) | Autonomously derive, build, audit, and improve KG configs from PMC articles |
 | [`build-fullmap`](#build-fullmap) | Build the embedded fullmap redb used for entity resolution |
 | [`build-kg`](#build-kg) | Build a KGX NDJSON knowledge graph from a YAML configuration |
+| [`distill-export`](#distill-export) | Export a recorded distillation NDJSON dataset to an on-disk Hugging Face dataset |
 | [`validate`](#validate) | Validate a graph or table configuration without executing it |
 | [`validate-kgx`](#validate-kgx) | Validate built KGX NDJSON against the Biolink Model |
 
@@ -68,6 +69,7 @@ page lists the flags; see
 | `--biolink-threshold` | float | No | `0.0` | Minimum Biolink pass rate of the built KGX for MAPPED; `0.0` reports the rate without gating |
 | `--local`, `-l` | list[str] | No | `None` | Local payload: one DIR for all ids, or `PMCid=DIR` mappings; skips the PMC-AWS fetch (exit 2 on a missing DIR) |
 | `--optimize`, `-o` | bool | No | `False` | Run GEPA prompt optimization and persist optimized instructions instead of running the supervisor |
+| `--distill`, `-d`, `-dt` | bool | No | `False` | Record every LLM call (agent, judge, reflexion) as ChatML NDJSON under `<state-dir>/distill/records.ndjson` for fine-tuning; not supported with `--optimize` |
 | `--instructions-file` | Path | No | `None` | Load GEPA-optimized instructions from a prior `--optimize` run |
 | `--instructions-out` | Path | No | `None` | Where `--optimize` writes optimized instructions (default `<state-dir>/optimized_instructions.yaml`) |
 | `--max-metric-calls` | int | No | `8` | GEPA metric-call budget for `--optimize` |
@@ -88,6 +90,25 @@ tablassert agent PMC11708054 -f ./graph.yaml
 
 ---
 
+## distill-export
+
+Use this to convert a distillation dataset recorded with
+[`agent --distill`](#agent) into an on-disk Hugging Face dataset (`save_to_disk`). Requires the
+`[distill]` extra (`pip install "tablassert[distill]"`, pulls `datasets`). The raw NDJSON already
+loads directly in Unsloth Studio and via `datasets.load_dataset("json", ...)` — this export is
+only needed for `datasets`-native workflows.
+
+```bash
+tablassert distill-export --distill-dir .tablassert/agent/distill --out ./hf-dataset
+```
+
+| Option | Type | Required | Default | Description |
+| --- | --- | --- | --- | --- |
+| `--distill-dir`, `-dd` | Path | Yes | n/a | Directory holding the recorded `*.ndjson` files (exit 2 when empty) |
+| `--out`, `-o` | Path | Yes | n/a | Destination directory for the `save_to_disk` dataset |
+
+---
+
 ## build-fullmap
 
 Use this to obtain the embedded `fullmap.redb` entity-resolution database. By default it first tries

pyproject.toml

@@ -95,6 +95,9 @@ agent = [
 optimize = [
     "dspy>=3.2.1",
 ]
+distill = [
+    "datasets>=3.0.0",
+]
 log = [
     "loguru>=0.7.3",
 ]

src/tablassert/agent.py

@@ -3232,6 +3232,56 @@ def generate(
     return FakeModel()
 
 
+def make_distilling_model(model: object, recorder: object, *, purpose: str, meta: dict[str, object] | None = None) -> object:
+    """Wrap a smolagents model so every ``generate`` call is recorded by ``recorder``.
+
+    Returns a ``Model`` subclass (same shape as :func:`make_fake_model`) that delegates
+    ``generate`` to the wrapped model and appends one ChatML NDJSON record per call, tagged with
+    ``purpose`` (``"agent"``/``"judge"``/``"reflexion"``) plus ``meta`` (e.g. ``pmc_id``). The
+    wrapped model's own ``model_id`` is added to the record when recoverable. Attribute access
+    falls through to the wrapped model (``__getattr__``) so smolagents sees the real model's
+    metadata. Recording failures never propagate — the response is returned untouched.
+    """
+    _require("smolagents")
+    from smolagents.models import ChatMessage, Model  # local import keeps module import lazy  # pyright: ignore[reportMissingImports]
+
+    extra_meta: dict[str, object] = dict(meta or {})
+    if "model_id" not in extra_meta:
+        model_id: object = getattr(model, "model_id", None)
+        if model_id is not None:
+            extra_meta["model_id"] = str(model_id)
+
+    class DistillingModel(Model):  # pyright: ignore[reportMissingImports]
+        def __init__(self) -> None:
+            with contextlib.suppress(Exception):
+                super().__init__()
+            self._wrapped: object = model
+
+        def __getattr__(self, name: str) -> object:
+            # Only fires for attributes Model does not define; delegates model_id & friends.
+            return getattr(self._wrapped, name)
+
+        def generate(
+            self,
+            messages: list[ChatMessage],
+            stop_sequences: list[str] | None = None,
+            response_format: dict[str, str] | None = None,
+            tools_to_call_from: object = None,
+            **kwargs: Any,
+        ) -> ChatMessage:
+            response: ChatMessage = cast(
+                ChatMessage,
+                self._wrapped.generate(  # pyright: ignore[reportAttributeAccessIssue]
+                    messages, stop_sequences=stop_sequences, response_format=response_format, tools_to_call_from=tools_to_call_from, **kwargs
+                ),
+            )
+            with contextlib.suppress(Exception):  # recording must never break the run
+                recorder.record(purpose, messages, response, **extra_meta)  # pyright: ignore[reportAttributeAccessIssue]
+            return response
+
+    return DistillingModel()
+
+
 # --------------------------------------------------------------------------- #
 # US-009: outer DETERMINISTIC supervisor + monotonic improve loop + checkpoint/resume
 #
@@ -3433,6 +3483,11 @@ def pmc_build_dir(root: Path, pmc_id: str) -> Path:
     return root / "builds" / pmc_id
 
 
+def distill_dir(root: Path) -> Path:
+    """Return ``<root>/distill`` (the distillation dataset dir); pure, no mkdir."""
+    return root / "distill"
+
+
 @dataclass
 class ConfigRecord:
     """Per-PMC supervisor record: status, derived/best config paths, and coverage history.
@@ -3616,6 +3671,7 @@ def run_supervisor(
     instructions: str | None = None,
     derive_mode: DeriveMode = "full",
     min_rows: int = MIN_TABLE_ROWS,
+    distill_recorder: object | None = None,
 ) -> dict[str, object]:
     """Run the deterministic supervisor over a batch of PMC ids with checkpoint/resume.
 
@@ -3652,7 +3708,9 @@ def run_supervisor(
     The whole per-pmc body is wrapped in try/except: ANY failure marks that record SKIPPED with the
     reason and advances (one bad pmc never aborts the batch). ``build_model_factory`` is a zero-arg
     callable returning a configured model so tests inject a FakeModel and the real CLI keeps secrets
-    out of this signature. Returns ``{"state", "records", "metrics"}`` after a final checkpoint.
+    out of this signature. ``distill_recorder`` (optional) wraps each article's model via
+    :func:`make_distilling_model` so every LLM call is appended to the distillation NDJSON dataset.
+    Returns ``{"state", "records", "metrics"}`` after a final checkpoint.
     """
     if min_rows < 0:
         raise ValueError("min_rows must be non-negative")
@@ -3749,8 +3807,13 @@ def audit_config(config_yaml: str, **kwargs: Any) -> dict[str, object]:
             article_xml: Path | None = next((path for path in files if path.suffix.lower() in {".xml", ".nxml"}), None)
 
             metrics: dict[str, object] = {}
+            model: object = build_model_factory()
+            if distill_recorder is not None:
+                # Distillation capture: wrap so every generate() call lands in the NDJSON dataset,
+                # tagged with this article's id for later filtering against state.json status.
+                model = make_distilling_model(model, distill_recorder, purpose="agent", meta={"pmc_id": pmc_id})
             agent: object = build_agent(
-                model=build_model_factory(),
+                model=model,
                 tools=make_tools(
                     graph=target_graph,
                     fullmap=effective_fullmap,

src/tablassert/cli.py

@@ -776,6 +776,7 @@ def agent(
     biolink_threshold: Annotated[float, cyclopts.Parameter(name=["--biolink-threshold"])] = 0.0,
     local: Annotated[list[str] | None, cyclopts.Parameter(name=["--local", "-l"])] = None,
     optimize: Annotated[bool, cyclopts.Parameter(name=["--optimize", "-o"], negative="")] = False,
+    distill: Annotated[bool, cyclopts.Parameter(name=["--distill", "-d", "-dt"], negative="")] = False,
     instructions_file: Annotated[Path | None, cyclopts.Parameter(name=["--instructions-file"])] = None,
     instructions_out: Annotated[Path | None, cyclopts.Parameter(name=["--instructions-out"])] = None,
     max_metric_calls: Annotated[int, cyclopts.Parameter(name=["--max-metric-calls"])] = 8,
@@ -825,6 +826,9 @@ def agent(
             one or more ``PMCid=DIR`` mappings (per-article). Fails loud (exit 2) if a DIR does not exist.
         optimize: Run GEPA prompt optimization over the model config and persist optimized instructions
             (instead of running the supervisor); use ``--instructions-out`` to choose the output file.
+        distill: Record every LLM call of the run (inner agent, judge, reflexion) as ChatML NDJSON
+            under ``<state-dir>/distill/records.ndjson`` for fine-tuning (Unsloth Studio / QLoRA).
+            Zero extra dependencies; ``tablassert distill-export`` converts it to an on-disk HF dataset.
         instructions_file: Load GEPA-optimized instructions (from a prior ``--optimize`` run) for this run.
         instructions_out: Where ``--optimize`` writes optimized instructions (default
             ``<state-dir>/optimized_instructions.yaml``).
@@ -877,6 +881,12 @@ def agent(
         print("tablassert agent: --gepa-threads must be a positive integer.", file=sys.stderr)
         raise SystemExit(2)
 
+    # --distill records the SUPERVISOR's model calls; the --optimize path returns early below and
+    # GEPA's dspy LM bypasses the recording seam, so the combination would silently record nothing.
+    if distill and optimize:
+        print("tablassert agent: --distill records supervisor LLM calls and is not supported with --optimize.", file=sys.stderr)
+        raise SystemExit(2)
+
     # Preflight the extras once the flags are known to be valid and BEFORE any model is
     # built or any article fetched. smolagents is otherwise only required per-article
     # (inside build_agent) and dspy only once GEPA starts, so an absent extra would
@@ -885,6 +895,17 @@ def agent(
     if optimize:
         extras.require("optimize", required_by="tablassert agent --optimize")
 
+    # Distillation capture (optional, zero-dep): every LLM call is appended as one ChatML NDJSON
+    # record. The inner agent's model is wrapped per-article inside run_supervisor (which knows the
+    # pmc_id); the judge and reflexion models are wrapped at their construction sites below.
+    distill_recorder: object | None = None
+    if distill:
+        from tablassert import distill as distill_mod
+
+        distill_path: Path = agent_mod.distill_dir(state_dir) / distill_mod.RECORDS_FILENAME
+        distill_recorder = distill_mod.DistillRecorder(distill_path)
+        print(f"tablassert agent: distilling LLM calls -> {distill_path}")
+
     def build_model_factory() -> object:
         return agent_mod.build_model(resolved_id, resolved_base, resolved_key, backend=backend)
 
@@ -893,14 +914,20 @@ def build_model_factory() -> object:
     if reflexion:
 
         def _make_reflexion() -> object:
-            return agent_mod.make_prompt_callable(agent_mod.build_model(resolved_id, resolved_base, resolved_key, backend=backend))
+            reflexion_model: object = agent_mod.build_model(resolved_id, resolved_base, resolved_key, backend=backend)
+            if distill_recorder is not None:
+                reflexion_model = agent_mod.make_distilling_model(reflexion_model, distill_recorder, purpose="reflexion")
+            return agent_mod.make_prompt_callable(reflexion_model)
 
         reflexion_factory = _make_reflexion
 
     # Semantic judge (optional): a prompt-callable over the judge model (same api_base/api_key).
     judge: object | None = None
     if judge_model is not None:
-        judge = agent_mod.make_prompt_callable(agent_mod.build_model(judge_model, resolved_base, resolved_key, backend=backend))
+        judge_base_model: object = agent_mod.build_model(judge_model, resolved_base, resolved_key, backend=backend)
+        if distill_recorder is not None:
+            judge_base_model = agent_mod.make_distilling_model(judge_base_model, distill_recorder, purpose="judge")
+        judge = agent_mod.make_prompt_callable(judge_base_model)
 
     # Local payload (optional, W4): a DIR for all ids, or PMCid=DIR mappings; fail loud on a missing dir.
     def parse_local(specs: list[str] | None) -> dict[str, Path] | Path | None:
@@ -990,6 +1017,7 @@ def parse_local(specs: list[str] | None) -> dict[str, Path] | Path | None:
         biolink_threshold=biolink_threshold,
         local=local_payload,
         instructions=run_instructions,
+        distill_recorder=distill_recorder,
     )
 
     metrics_raw: object = result.get("metrics")
@@ -1012,6 +1040,38 @@ def metric(key: str, default: float) -> float:
     )
 
 
+@APP.command(name="distill-export")
+def distill_export(
+    *, distill_dir: Annotated[Path, cyclopts.Parameter(name=["--distill-dir", "-dd"])], out: Annotated[Path, cyclopts.Parameter(name=["--out", "-o"])]
+) -> None:
+    """Export a recorded distillation NDJSON dataset to an on-disk Hugging Face dataset.
+
+    Loads every ``*.ndjson`` under ``--distill-dir`` (the ChatML records written by ``tablassert
+    agent --distill``) with ``datasets.load_dataset("json", ...)`` and writes the result with
+    ``save_to_disk`` to ``--out``. Requires the ``distill`` extra (``pip install
+    "tablassert[distill]"``). The raw NDJSON also loads directly in Unsloth Studio — this export
+    is only needed for ``datasets``-native workflows.
+
+    Args:
+        distill_dir: Directory holding the recorded ``*.ndjson`` files (default output of
+            ``tablassert agent --distill`` is ``<state-dir>/distill``).
+        out: Destination directory for the ``save_to_disk`` dataset.
+    """
+    # Input validation precedes the extras preflight: a missing directory is the user's typo, an
+    # absent extra is their environment, and the typo is the faster loop to close first.
+    files: list[Path] = sorted(distill_dir.glob("*.ndjson"))
+    if not files:
+        print(f"tablassert distill-export: no .ndjson records under {distill_dir} — run tablassert agent --distill first.", file=sys.stderr)
+        raise SystemExit(2)
+    extras.require("distill", required_by="tablassert distill-export")
+    from datasets import load_dataset  # local import keeps the CLI import-light  # pyright: ignore[reportMissingImports]
+
+    dataset: object = load_dataset("json", data_files=[str(path) for path in files], split="train")
+    out.parent.mkdir(parents=True, exist_ok=True)
+    dataset.save_to_disk(str(out))  # pyright: ignore[reportAttributeAccessIssue]
+    print(f"tablassert distill-export: {len(dataset)} record(s) from {len(files)} file(s) -> {out}")  # pyright: ignore[reportArgumentType]
+
+
 class PrebuiltFullmapUnavailable(Exception):
     """A prebuilt fullmap could not be fetched or extracted.