Land the SetFit/ONNX intent router and data pipeline

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: ci
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: "3.12"
+
+      - name: Create venv and install base deps
+        run: uv sync --no-dev
+
+      - name: Lint
+        run: uvx ruff check
+
+      - name: Unit tests
+        run: uv run python -m unittest discover -s tests -v

.gitignore

@@ -16,6 +16,7 @@ checkpoints/
 *.bin
 *.gguf
 data/*.jsonl
+data/raw/
 !examples/*.jsonl
 
 .env

README.md

@@ -15,16 +15,43 @@ The frozen intent set and labeling rules live in [docs/taxonomy.md](docs/taxonom
 
 ## Layout
 
+The pipeline runs **data → train → export → evaluate**. Two files are the spine
+everything else imports: `data.py` (the `Intent` label set and JSONL loaders)
+and `preprocess.py` (the redaction/normalization applied identically at training
+and inference, mirrored in the Rust consumer).
+
 ```
 src/routelet/
-  data.py       label schema and train/test prep
-  train.py      fine-tune a small model on the labels
-  serve.py      FastAPI endpoint: text in, intent out
-  evaluate.py   routelet vs the LLM baseline (accuracy, latency, cost)
-examples/
-  intents.sample.jsonl   example label format
+  data.py            Intent enum (the 5 labels) + JSONL load/split. Source of truth.
+  preprocess.py      mask secrets/emails/digits. Same pass at train and inference.
+
+  augment.py         clean rows -> voice-like disfluent variants (data/augmented.jsonl)
+  checkdata.py       dataset health: per-class counts, train/eval leakage
+
+  train.py           TF-IDF + logistic-regression baseline (the floor to beat)
+  train_setfit.py    the shipped model: SetFit fine-tune of bge-small + LR head,
+                     then temperature calibration -> models/setfit
+  setfit_baseline.py benchmark harness that trains+discards SetFit (not the ship path)
+  export_onnx.py     models/setfit -> embedder.onnx + tokenizer.json + head.json (int8)
+
+  evaluate.py        the Claude LLM baseline (the accuracy ceiling), one call per row
+  teacher.py         shared taxonomy prompt + schema + classify(), used by eval + ingest
+  serve.py           optional FastAPI endpoint (stub; Aegis runs the ONNX in-process)
+
+Scripts/             runnable tools (not library code)
+  pull_samples.py    pull redacted samples from the proxy's R2 -> one JSONL
+  ingest_samples.py  label them (free Claude label or teacher), dedup -> data/collected.jsonl
+  refresh_haiku_baseline.py   re-run the paid Haiku eval -> report/baselines.json
+
+report/              report.py scores the models on the holdout -> figures + metrics.json
+evals/               frozen eval data (holdout.jsonl). Never trained on.
+experiments/         throwaway prototypes, not part of the library (see its README)
 ```
 
+The data-collection loop closes back on itself: Aegis logs redacted samples ->
+proxy -> R2, then `Scripts/pull_samples.py` + `ingest_samples.py` turn those into
+new labeled training data.
+
 ## Data format
 
 One JSON object per line:

Scripts/ingest_samples.py

@@ -0,0 +1,158 @@
+"""Turn pulled distillation samples into labeled training data.
+
+Reads the raw samples produced by ``pull_samples.py`` and resolves a single
+training label for each:
+
+  * If the sample carries ``claude_label`` (the Claude fallback fired on-device
+    that turn), use it. Free teacher signal, already paid for.
+  * Otherwise call the offline Claude teacher to label it (docs/distillation.md).
+
+Then it normalizes the text through ``preprocess`` (the same pass training and
+inference use, so there is no skew), drops anything already in the training pool
+or in the eval set (keeping the eval honest), and appends the survivors to
+``data/collected.jsonl`` as ``{"text", "intent"}`` lines that ``data.load`` reads.
+
+Needs ANTHROPIC_API_KEY only when some samples lack a ``claude_label``. The
+teacher model is claude-opus-4-8 by default; override with ROUTELET_TEACHER_MODEL.
+
+Usage:
+    uv run python Scripts/ingest_samples.py
+    uv run python Scripts/ingest_samples.py --dry-run
+    uv run python Scripts/ingest_samples.py --raw data/raw/today.jsonl --limit 200
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT / "src"))
+
+from routelet.data import Intent, load, load_dir  # noqa: E402
+from routelet.preprocess import preprocess  # noqa: E402
+
+DEFAULT_RAW = PROJECT_ROOT / "data" / "raw" / "samples.jsonl"
+DEFAULT_OUT = PROJECT_ROOT / "data" / "collected.jsonl"
+TRAIN_DIR = PROJECT_ROOT / "data"
+EVAL_FILE = PROJECT_ROOT / "evals" / "holdout.jsonl"
+
+
+def _norm(text: str) -> str:
+    """The dedup key: normalized text, case-folded and stripped. Mirrors the
+    leakage check in checkdata.py so 'same command' means the same thing here."""
+    return preprocess(text).lower().strip()
+
+
+def load_raw(path: Path) -> list[dict]:
+    rows: list[dict] = []
+    with path.open() as f:
+        for n, line in enumerate(f, start=1):
+            if not line.strip():
+                continue
+            try:
+                rows.append(json.loads(line))
+            except json.JSONDecodeError:
+                print(f"  skip {path}:{n}: bad json", file=sys.stderr)
+    return rows
+
+
+def existing_keys() -> tuple[set[str], set[str]]:
+    """Texts already in the training pool and in the eval set, both normalized.
+    A sample matching either is dropped: the first is a duplicate, the second
+    would leak the eval into training."""
+    train = {_norm(e.text) for e in load_dir(TRAIN_DIR)}
+    evalset = {_norm(e.text) for e in load(EVAL_FILE)} if EVAL_FILE.exists() else set()
+    return train, evalset
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--raw", type=Path, default=DEFAULT_RAW, help="raw samples JSONL")
+    parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="labeled output JSONL")
+    parser.add_argument("--limit", type=int, help="stop after this many raw samples")
+    parser.add_argument("--model", help="teacher model override (else ROUTELET_TEACHER_MODEL)")
+    parser.add_argument("--dry-run", action="store_true", help="don't call the teacher or write")
+    args = parser.parse_args()
+
+    if not args.raw.exists():
+        sys.exit(f"no raw samples at {args.raw}; run pull_samples.py first")
+
+    raw = load_raw(args.raw)
+    if args.limit:
+        raw = raw[: args.limit]
+
+    train_keys, eval_keys = existing_keys()
+    # Dedup against the existing pool, the output file (so re-runs are idempotent
+    # even when --out lives outside data/), and within this batch.
+    seen = set(train_keys)
+    if args.out.exists():
+        seen |= {_norm(r["text"]) for r in load_raw(args.out) if r.get("text")}
+
+    client = None  # created lazily; only needed for teacher labeling
+    labeled: list[dict] = []
+    stats = {"client_label": 0, "teacher_label": 0, "dup": 0, "leak": 0, "empty": 0, "bad": 0}
+
+    for row in raw:
+        text = (row.get("text") or "").strip()
+        if not text:
+            stats["empty"] += 1
+            continue
+
+        key = _norm(text)
+        if key in eval_keys:
+            stats["leak"] += 1
+            continue
+        if key in seen:
+            stats["dup"] += 1
+            continue
+
+        # Prefer the free in-turn Claude label; fall back to the offline teacher.
+        label_str = row.get("claude_label")
+        if label_str:
+            try:
+                intent = Intent(label_str)
+            except ValueError:
+                stats["bad"] += 1
+                continue
+            stats["client_label"] += 1
+        elif args.dry_run:
+            # Count what a real run would send to the teacher without spending.
+            stats["teacher_label"] += 1
+            seen.add(key)
+            continue
+        else:
+            if client is None:
+                import anthropic
+
+                client = anthropic.Anthropic()
+            from routelet.teacher import classify
+
+            intent = classify(client, preprocess(text), model=args.model)
+            stats["teacher_label"] += 1
+
+        seen.add(key)
+        labeled.append({"text": preprocess(text), "intent": intent.value})
+
+    if args.dry_run:
+        would = stats["client_label"] + stats["teacher_label"]
+        print(f"dry run: {would} new samples ({stats['teacher_label']} need the teacher)")
+        print(f"  skipped: {stats['dup']} dup, {stats['leak']} eval-leak, "
+              f"{stats['empty']} empty, {stats['bad']} bad-label")
+        return
+
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    with args.out.open("a") as f:
+        for rec in labeled:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+    print(f"appended {len(labeled)} samples to {args.out}")
+    print(f"  {stats['client_label']} client-labeled, {stats['teacher_label']} teacher-labeled")
+    print(f"  skipped: {stats['dup']} dup, {stats['leak']} eval-leak, "
+          f"{stats['empty']} empty, {stats['bad']} bad-label")
+
+
+if __name__ == "__main__":
+    main()

Scripts/pull_samples.py

@@ -0,0 +1,107 @@
+"""Pull redacted distillation samples from the proxy's R2 bucket into one local
+JSONL file.
+
+The aegis proxy stores one immutable JSON object per sample under
+``samples/<date>/<device>/<ts>-<uuid>.json`` (see aegis proxy/handlers/routelet.ts).
+R2 has no append, so the per-sample-object layout sidesteps write races across
+devices; compaction into a single batched JSONL happens here, at pull time.
+
+This is a full snapshot: every object under the prefix is read and written out,
+deduped by R2 key. Re-running overwrites the output. The raw text was already
+redacted on-device and scrubbed again by the proxy, but it is still real user
+input, so the output is gitignored (see data/raw/).
+
+Reads R2 over the S3-compatible API. Set in the environment (or .env):
+    R2_ACCOUNT_ID          Cloudflare account id
+    R2_ACCESS_KEY_ID       R2 API token access key
+    R2_SECRET_ACCESS_KEY   R2 API token secret
+    R2_BUCKET              bucket name (default: aegis-routelet-samples)
+
+Create a read-only R2 API token in the Cloudflare dashboard
+(R2 > Manage API Tokens). boto3 is a tooling dep installed into the venv, not a
+project dependency.
+
+Usage:
+    uv run python Scripts/pull_samples.py
+    uv run python Scripts/pull_samples.py --prefix samples/2026-05-29/ --out data/raw/today.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import boto3
+from botocore.config import Config
+
+PROJECT_ROOT = Path(__file__).parent.parent
+DEFAULT_OUT = PROJECT_ROOT / "data" / "raw" / "samples.jsonl"
+DEFAULT_BUCKET = "aegis-routelet-samples"
+DEFAULT_PREFIX = "samples/"
+
+
+def _require_env(name: str) -> str:
+    val = os.environ.get(name)
+    if not val:
+        sys.exit(
+            f"missing ${name}. Set R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, and "
+            f"R2_SECRET_ACCESS_KEY (read-only R2 API token)."
+        )
+    return val
+
+
+def r2_client():
+    account_id = _require_env("R2_ACCOUNT_ID")
+    return boto3.client(
+        "s3",
+        endpoint_url=f"https://{account_id}.r2.cloudflarestorage.com",
+        aws_access_key_id=_require_env("R2_ACCESS_KEY_ID"),
+        aws_secret_access_key=_require_env("R2_SECRET_ACCESS_KEY"),
+        # R2 ignores the region but boto3 requires one; signature v4 is required.
+        region_name="auto",
+        config=Config(signature_version="s3v4"),
+    )
+
+
+def pull(bucket: str, prefix: str, out: Path) -> tuple[int, int]:
+    """Read every object under prefix and write one JSON line per valid sample.
+    Returns (written, skipped)."""
+    client = r2_client()
+    paginator = client.get_paginator("list_objects_v2")
+
+    out.parent.mkdir(parents=True, exist_ok=True)
+    written = 0
+    skipped = 0
+    with out.open("w") as f:
+        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
+            for obj in page.get("Contents", []):
+                key = obj["Key"]
+                body = client.get_object(Bucket=bucket, Key=key)["Body"].read()
+                try:
+                    sample = json.loads(body)
+                except json.JSONDecodeError:
+                    print(f"  skip (bad json): {key}", file=sys.stderr)
+                    skipped += 1
+                    continue
+                # One object is one sample; the key already dedupes, so write as-is.
+                f.write(json.dumps(sample, ensure_ascii=False) + "\n")
+                written += 1
+    return written, skipped
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--bucket", default=DEFAULT_BUCKET)
+    parser.add_argument("--prefix", default=DEFAULT_PREFIX, help="R2 key prefix to pull")
+    parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="output JSONL path")
+    args = parser.parse_args()
+
+    written, skipped = pull(args.bucket, args.prefix, args.out)
+    print(f"wrote {written} samples to {args.out} ({skipped} skipped)")
+
+
+if __name__ == "__main__":
+    main()