diff --git a/.gitignore b/.gitignore
index 20311af9..f4ff2e4b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -98,3 +98,16 @@ src/openjarvis/channels/whatsapp_baileys_bridge/node_modules/
# Second Repos
Inline/
scratch/
+
+# ---------------------------------------------------------------------------
+# Distillation runtime artifacts (defense in depth — these should always live
+# in ~/.openjarvis/, never inside the source tree, but we ignore them here in
+# case OPENJARVIS_HOME is misconfigured during dev)
+# ---------------------------------------------------------------------------
+.openjarvis/
+learning.db
+**/learning/sessions/
+**/learning/pending_review/
+**/learning/benchmarks/
+**/teacher_traces/
+*.session.json
diff --git a/Inline b/Inline
new file mode 160000
index 00000000..03673aaa
--- /dev/null
+++ b/Inline
@@ -0,0 +1 @@
+Subproject commit 03673aaa427a6affda7c2a759fb56abac67ddde7
diff --git a/desktop/src-tauri/src/overlay.html b/desktop/src-tauri/src/overlay.html
new file mode 100644
index 00000000..a8129262
--- /dev/null
+++ b/desktop/src-tauri/src/overlay.html
@@ -0,0 +1,450 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/docs/architecture/learning.md b/docs/architecture/learning.md
index 76d9a658..f0496f4b 100644
--- a/docs/architecture/learning.md
+++ b/docs/architecture/learning.md
@@ -553,3 +553,60 @@ metric can be improved without degrading another.
The optimization framework has full Rust parity via the `openjarvis-learning`
crate, with PyO3 bindings exposing `OptimizationStore` and `LLMOptimizer`
to Python.
+
+---
+
+## Distillation (Frontier-Driven Harness Learning)
+
+The distillation subsystem uses a frontier closed-source model (the "teacher") as a meta-engineer for the local student's full harness — not just its weights. Instead of pushing knowledge into a small model's weights, we push a frontier model's engineering judgement into the surrounding configuration: prompts, routing, agent class, tool availability, and tool descriptions.
+
+### Where it lives
+
+`learning/distillation/` is the fifth subsystem within the Learning pillar, alongside `learning/routing/`, `learning/optimize/`, `learning/training/`, and `learning/intelligence/`.
+
+### Four-phase loop
+
+```
+Trigger → Diagnose → Plan → Execute → Record
+```
+
+1. **Diagnose** — TeacherAgent (frontier model with diagnostic tools) analyzes traces, runs student/teacher comparisons, identifies 2-5 failure clusters with evidence.
+2. **Plan** — LearningPlanner converts diagnosis into a typed LearningPlan with deterministic risk tier assignment and patch/replace downgrade.
+3. **Execute** — Per-edit loop: EditApplier validates + applies, BenchmarkGate scores, CheckpointStore commits or rolls back.
+4. **Record** — LearningSession persisted to SQLite + JSON artifact.
+
+### Key components
+
+| Component | Module | Purpose |
+|-----------|--------|---------|
+| `DistillationOrchestrator` | `orchestrator.py` | Top-level session driver |
+| `TeacherAgent` | `diagnose/teacher_agent.py` | Frontier model tool-calling loop |
+| `DiagnosisRunner` | `diagnose/runner.py` | Phase 1 orchestration |
+| `LearningPlanner` | `plan/planner.py` | Diagnosis → typed LearningPlan |
+| `EditApplier` + registry | `execute/base.py` | Abstract applier interface |
+| `BenchmarkGate` | `gate/benchmark_gate.py` | Benchmark-based accept/reject |
+| `CheckpointStore` | `checkpoint/store.py` | Git-backed config rollback |
+| `SessionStore` | `storage/session_store.py` | SQLite session persistence |
+
+### Relationship to existing subsystems
+
+| Existing | Relationship |
+|----------|-------------|
+| `LearningOrchestrator` | Sibling — stays untouched |
+| `LLMOptimizer` / `OptimizationStore` | Sibling — grid-search vs root-cause |
+| `LearnedRouterPolicy` | Reused — routing edits update it |
+| `TraceJudge` | Reused — benchmark scoring |
+| `PersonalBenchmarkSynthesizer` | Extended — auto-refresh, gold answers |
+| `TraceStore` | Reused — read-only access from diagnostic tools |
+
+### Risk tier system
+
+Every edit is assigned a tier from a deterministic lookup table:
+
+| Tier | Ops | Behavior |
+|------|-----|----------|
+| `auto` | Model routing/params, tool add/remove/description, agent params | Apply if gate passes |
+| `review` | System prompt edits, agent class, few-shot exemplars | Queue for user approval |
+| `manual` | LoRA fine-tuning (v2) | Never auto-apply |
+
+See [Distillation user guide](../user-guide/learning-distillation.md) for CLI usage and configuration.
diff --git a/docs/architecture/overview.md b/docs/architecture/overview.md
index cea7c904..a3bf5f0b 100644
--- a/docs/architecture/overview.md
+++ b/docs/architecture/overview.md
@@ -45,7 +45,7 @@ The memory pipeline includes document ingestion, chunking, embedding generation,
The Learning system is the fifth primitive, connecting the other four through **trace-driven feedback**. Every agent interaction can produce a `Trace` capturing the full sequence of steps — routing decisions, memory retrieval, inference calls, tool invocations, and final responses. The `TraceAnalyzer` computes statistics from accumulated traces, and the `TraceDrivenPolicy` uses these statistics to learn which model/agent/tool combinations produce the best outcomes for different query types.
-The learning system is configured through nested sub-sections in `config.toml`: `[learning.routing]` controls the router policy (heuristic, learned, sft, grpo), `[learning.intelligence]` controls the model-level learning policy, `[learning.agent]` controls agent advisor and ICL updater policies, and `[learning.metrics]` sets the composite reward function weights.
+The learning system is configured through nested sub-sections in `config.toml`: `[learning.routing]` controls the router policy (heuristic, learned, sft, grpo), `[learning.intelligence]` controls the model-level learning policy, `[learning.agent]` controls agent advisor and ICL updater policies, and `[learning.metrics]` sets the composite reward function weights. The pillar also includes the distillation subsystem, a frontier-driven loop that improves the local harness — see [Learning architecture: Distillation](learning.md#distillation-frontier-driven-harness-learning).
---
diff --git a/docs/development/roadmap.md b/docs/development/roadmap.md
index f1333364..272ffbfa 100644
--- a/docs/development/roadmap.md
+++ b/docs/development/roadmap.md
@@ -9,6 +9,7 @@ These are the areas where active development is happening and contributions are
- **Energy-aware routing** — using power consumption data from telemetry to optimize for energy efficiency alongside latency and quality
- **Plugin ecosystem** — community-contributed engines, tools, and agents distributed as Python packages
- **Federated memory** — memory backends that synchronize across devices
+- **Distillation:** Frontier-driven harness learning — a frontier model analyzes your traces and proposes config improvements. See [user guide](../user-guide/learning-distillation.md) and [architecture](../architecture/learning.md#distillation-frontier-driven-harness-learning).
---
diff --git a/docs/getting-started/configuration.md b/docs/getting-started/configuration.md
index 8e1aa36e..daaeb769 100644
--- a/docs/getting-started/configuration.md
+++ b/docs/getting-started/configuration.md
@@ -1057,3 +1057,59 @@ OpenJarvis respects the following environment variables:
- [Architecture Overview](../architecture/overview.md) — Understand how the pieces fit together
- [Intelligence Primitive](../architecture/intelligence.md) — Model identity and generation defaults
- [Learning & Traces](../architecture/learning.md) — Router policies and the trace-driven feedback loop
+
+---
+
+## Learning & Distillation
+
+The distillation subsystem uses a frontier model to automatically improve your local agent configuration. See the [user guide](../user-guide/learning-distillation.md) for a full walkthrough.
+
+### `[learning.distillation]`
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `enabled` | bool | `true` | Gate the entire distillation subsystem |
+| `autonomy_mode` | string | `"tiered"` | `auto`, `tiered`, or `manual` |
+| `teacher_model` | string | `"claude-opus-4-6"` | Frontier model for diagnosis and planning |
+| `max_cost_per_session_usd` | float | `5.0` | Per-session teacher API budget |
+| `max_tool_calls_per_diagnosis` | int | `30` | Max teacher tool calls in diagnosis phase |
+
+### `[learning.distillation.triggers]`
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `scheduled_enabled` | bool | `true` | Enable daily scheduled sessions |
+| `scheduled_cron` | string | `"0 3 * * *"` | Cron expression for scheduled trigger |
+| `scheduled_min_new_traces` | int | `20` | Minimum new traces to trigger |
+| `cluster_enabled` | bool | `true` | Enable failure cluster trigger |
+| `cluster_check_interval_minutes` | int | `60` | How often to check for clusters |
+| `cluster_min_size` | int | `5` | Minimum traces in a cluster |
+| `cluster_failure_threshold` | float | `0.3` | Feedback <= this counts as failure |
+
+### `[learning.distillation.gate]`
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `min_improvement` | float | `0.0` | Minimum overall score improvement to accept |
+| `max_regression` | float | `0.05` | Maximum per-cluster score drop before rejecting |
+| `benchmark_subsample_size` | int | `50` | Tasks per gate run |
+| `full_benchmark` | bool | `false` | Disable subsampling (slower, more accurate) |
+
+### `[learning.distillation.benchmark]`
+
+| Key | Type | Default | Description |
+|-----|------|---------|-------------|
+| `synthesis_feedback_threshold` | float | `0.7` | Min feedback for benchmark traces |
+| `max_benchmark_size` | int | `200` | Max tasks in the benchmark |
+| `auto_refresh` | bool | `true` | Auto-mine new high-feedback traces |
+| `max_synthesis_cost_usd_per_refresh` | float | `2.0` | Cost cap per benchmark refresh |
+
+### `[learning.distillation.tier_overrides]`
+
+Override the default risk tier for any operation. Keys are operation names, values are tier strings (`auto`, `review`, `manual`).
+
+```toml
+[learning.distillation.tier_overrides]
+# patch_system_prompt = "auto" # promote to auto after trust
+# replace_system_prompt = "auto"
+```
diff --git a/docs/user-guide/cli.md b/docs/user-guide/cli.md
index 8ea8a831..5d05906f 100644
--- a/docs/user-guide/cli.md
+++ b/docs/user-guide/cli.md
@@ -431,3 +431,101 @@ curl http://localhost:8000/v1/chat/completions \
```
When an agent is configured (e.g., `--agent orchestrator`), non-streaming requests are routed through the agent with access to all registered tools. For tool-capable agents (`orchestrator`, `react`, `openhands`), all registered tools are automatically loaded and made available.
+
+---
+
+## `jarvis learning`
+
+Frontier-driven harness learning (distillation). Manages learning sessions, reviews pending edits, and controls the benchmark gate.
+
+### `jarvis learning init`
+
+Initialize the distillation checkpoint repo and directory layout.
+
+```bash
+jarvis learning init
+```
+
+### `jarvis learning run`
+
+Run an on-demand learning session.
+
+```bash
+jarvis learning run
+jarvis learning run --autonomy auto # auto-apply all edits
+jarvis learning run --autonomy manual # dry-run, everything goes to review
+```
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--autonomy` | `tiered` | `auto`, `tiered`, or `manual` |
+
+### `jarvis learning history`
+
+List past learning sessions.
+
+```bash
+jarvis learning history
+jarvis learning history --limit 5
+```
+
+### `jarvis learning show`
+
+Show details of a learning session (diagnosis, plan, outcomes, cost).
+
+```bash
+jarvis learning show
+```
+
+### `jarvis learning review`
+
+List all pending edits awaiting approval.
+
+```bash
+jarvis learning review
+```
+
+### `jarvis learning approve`
+
+Approve a pending edit (still goes through the benchmark gate).
+
+```bash
+jarvis learning approve
+```
+
+### `jarvis learning reject`
+
+Reject a pending edit.
+
+```bash
+jarvis learning reject
+jarvis learning reject --reason "too aggressive"
+```
+
+### `jarvis learning rollback`
+
+Rollback a session's committed edits (creates revert commits).
+
+```bash
+jarvis learning rollback
+jarvis learning rollback --last
+```
+
+### `jarvis learning benchmark`
+
+Personal benchmark management.
+
+```bash
+jarvis learning benchmark show # current stats
+jarvis learning benchmark refresh # manual refresh
+```
+
+### `jarvis learning daemon`
+
+Background learning daemon.
+
+```bash
+jarvis learning daemon start
+jarvis learning daemon stop
+jarvis learning daemon status
+```
diff --git a/docs/user-guide/evaluations.md b/docs/user-guide/evaluations.md
index 85275440..1b1b7985 100644
--- a/docs/user-guide/evaluations.md
+++ b/docs/user-guide/evaluations.md
@@ -14,6 +14,8 @@ The OpenJarvis evaluation framework (`openjarvis-evals`) measures model **correc
---
+> **Tip:** The distillation system uses this same eval infrastructure to gate edits against your personal benchmark. See [Learning & Distillation](learning-distillation.md).
+
## Installation
The evaluation framework is a standalone package in the `evals/` directory. Install it alongside OpenJarvis:
diff --git a/docs/user-guide/learning-distillation.md b/docs/user-guide/learning-distillation.md
new file mode 100644
index 00000000..2d32a30c
--- /dev/null
+++ b/docs/user-guide/learning-distillation.md
@@ -0,0 +1,251 @@
+# Learning & Distillation
+
+Use a frontier model as a meta-engineer to automatically improve your local agent's prompts, routing, and tools — reversibly, with benchmark-gated quality control.
+
+## Quick Start
+
+### 1. Initialize
+
+```bash
+jarvis learning init
+```
+
+This creates the distillation directory layout under `~/.openjarvis/learning/` and initializes a git checkpoint repo at `~/.openjarvis/.git` for tracking config changes.
+
+### 2. Run your first session
+
+Once you have at least 20 traces from regular use:
+
+```bash
+jarvis learning run
+```
+
+The system will:
+1. **Diagnose** — analyze your traces using a frontier model
+2. **Plan** — propose typed edits to your config
+3. **Execute** — apply edits that pass the benchmark gate
+4. **Record** — persist the session for history and rollback
+
+### 3. Check results
+
+```bash
+jarvis learning history
+jarvis learning show
+```
+
+## How a Learning Session Works
+
+A learning session has four phases:
+
+### Phase 1: Diagnose
+
+A frontier model (the "teacher", default `claude-opus-4-6`) analyzes your recent traces using read-only diagnostic tools. It identifies **failure clusters** — groups of related failures with shared root causes. The teacher must actually re-run your student on sample tasks and compare outputs to populate failure rates. This forces evidence-based diagnosis.
+
+**Output:** `diagnosis.md` with narrative analysis + structured failure clusters.
+
+### Phase 2: Plan
+
+A second teacher call converts the diagnosis into a typed `LearningPlan` — a list of `Edit` objects, each targeting a specific part of your configuration (model routing, system prompts, tool availability, etc.). The teacher cannot pick risk tiers — those are assigned deterministically from a lookup table.
+
+**Output:** `plan.json` frozen and immutable.
+
+### Phase 3: Execute
+
+Each edit is applied through its registered `EditApplier`, then scored against your personal benchmark. Edits that improve the benchmark are committed; edits that cause regressions are rolled back. Edits in the `review` tier are queued for your approval instead of being auto-applied.
+
+**Output:** Git commits in the checkpoint repo + `EditOutcome` records.
+
+### Phase 4: Record
+
+The session is persisted to `learning.db` (SQLite index) and `session.json` (authoritative artifact). You can query history, show details, and rollback any session.
+
+## Configuration
+
+Add to `~/.openjarvis/config.toml`:
+
+```toml
+[learning.distillation]
+enabled = true # gate the entire subsystem
+autonomy_mode = "tiered" # auto | tiered | manual
+teacher_model = "claude-opus-4-6" # any CloudEngine-supported model
+max_cost_per_session_usd = 5.0 # per-session teacher API budget
+max_tool_calls_per_diagnosis = 30 # max teacher tool calls in diagnosis
+```
+
+### Trigger configuration
+
+```toml
+[learning.distillation.triggers]
+scheduled_enabled = true
+scheduled_cron = "0 3 * * *" # daily at 03:00 local
+scheduled_min_new_traces = 20 # minimum new traces to trigger
+
+cluster_enabled = true
+cluster_check_interval_minutes = 60
+cluster_min_size = 5
+cluster_failure_threshold = 0.3 # feedback <= this counts as failure
+```
+
+### Gate configuration
+
+```toml
+[learning.distillation.gate]
+min_improvement = 0.0 # any improvement accepted (raise for margin)
+max_regression = 0.05 # max per-cluster score drop
+benchmark_subsample_size = 50 # tasks per gate run
+full_benchmark = false # set true to disable subsampling
+```
+
+### Benchmark configuration
+
+```toml
+[learning.distillation.benchmark]
+synthesis_feedback_threshold = 0.7 # min feedback for benchmark traces
+max_benchmark_size = 200 # max tasks in the benchmark
+auto_refresh = true # auto-mine new high-feedback traces
+max_synthesis_cost_usd_per_refresh = 2.0 # separate from session budget
+```
+
+### Risk tier overrides
+
+Power users can override the default tier for any operation:
+
+```toml
+[learning.distillation.tier_overrides]
+# Promote prompt edits to auto-apply after trust is established:
+# patch_system_prompt = "auto"
+# replace_system_prompt = "auto"
+```
+
+## Risk Tiers
+
+Every edit is assigned a risk tier that controls how it's applied:
+
+| Tier | Behavior | Default ops |
+|------|----------|-------------|
+| **auto** | Applied automatically if benchmark gate passes | Model routing, model params, tool add/remove/description, agent params |
+| **review** | Queued for user approval in `jarvis learning review` | System prompt edits, agent class changes, few-shot exemplars |
+| **manual** | Never auto-applied; requires explicit approval | LoRA fine-tuning (v2) |
+
+The tier is assigned deterministically from the edit operation — the teacher cannot override it.
+
+## Reviewing Edits
+
+When edits land in the review queue:
+
+```bash
+# List all pending edits
+jarvis learning review
+
+# Approve an edit (still goes through the benchmark gate)
+jarvis learning approve
+
+# Reject an edit with a reason
+jarvis learning reject --reason "prompt change too aggressive"
+```
+
+Even approved edits are gated by the benchmark — approval means "try it", not "force it".
+
+## Rollback and History
+
+Every edit creates a git commit in the checkpoint repo at `~/.openjarvis/.git`. This is separate from your OpenJarvis source repo.
+
+```bash
+# List past sessions
+jarvis learning history --limit 20
+
+# Show session details (diagnosis, plan, outcomes, cost)
+jarvis learning show
+
+# Rollback a session (creates new revert commits, preserves history)
+jarvis learning rollback
+jarvis learning rollback --last
+```
+
+Rollback never rewrites git history — it creates new revert commits so the audit trail stays intact.
+
+## Cost Controls
+
+Three cost boundaries prevent runaway spending:
+
+1. **`max_cost_per_session_usd`** (default $5.00) — caps the total teacher API cost per session (diagnosis + planning).
+2. **`max_synthesis_cost_usd_per_refresh`** (default $2.00) — caps the cost of generating gold answers for new benchmark tasks. Separate from the session budget.
+3. **`teacher_model`** — choose a cheaper model (e.g., `claude-sonnet-4-6`) to reduce per-token costs at the expense of diagnosis quality.
+
+Cost is tracked on every `LearningSession` as `teacher_cost_usd` and surfaced in `jarvis learning show`.
+
+## The Personal Benchmark
+
+The benchmark is your acceptance gate's source of truth — a set of tasks distilled from your high-quality traces, scored by an LLM-as-judge against frontier gold answers.
+
+**How it's built:**
+1. Traces with feedback >= 0.7 are candidates
+2. Tasks are grouped by query class and deduplicated
+3. For each task, the teacher generates a gold reference answer
+4. The benchmark is versioned (`personal_v1.json`, `personal_v2.json`, ...)
+
+**Auto-refresh:** The benchmark grows over time as you accumulate more traces. New tasks are added automatically during background refresh cycles.
+
+```bash
+# Manual refresh
+jarvis learning benchmark refresh
+
+# Show stats
+jarvis learning benchmark show
+```
+
+## Cold Start: What to Expect on Day One
+
+The system needs real usage data before it can learn:
+
+- **< 20 traces:** `jarvis learning run` returns "Not enough traces yet." Triggers are no-ops.
+- **20+ traces, < 10 high-feedback:** Enough for diagnosis, but no benchmark yet. Sessions will run diagnosis but can't gate edits.
+- **10+ high-feedback traces:** Bootstrap benchmark is created automatically (`personal_v1.json`). Full learning loop is available.
+
+**Getting there faster:** Use OpenJarvis normally and provide feedback on results (thumbs up/down in the UI, or `jarvis feedback` in the CLI).
+
+## Troubleshooting
+
+| Error | Cause | Fix |
+|-------|-------|-----|
+| "Not enough traces yet" | Fewer than 20 traces in the store | Use OpenJarvis more, provide feedback |
+| "Working tree dirty, cannot stage" | Manual edits to `~/.openjarvis/config.toml` during a session | Commit or revert manual changes first |
+| "All clusters dropped: insufficient evidence" | Teacher diagnosed clusters but couldn't reproduce failures | Check that the student is actually failing on the flagged tasks |
+| "ConfigurationError: distillation root inside source tree" | `OPENJARVIS_HOME` points inside the repo | Set `OPENJARVIS_HOME` to `~/.openjarvis` (default) or another external dir |
+| "Personal benchmark is empty" | Not enough high-feedback traces yet | Provide feedback on 10+ traces with score >= 0.7 |
+
+## Where Artifacts Live
+
+All distillation artifacts live under `~/.openjarvis/` (never inside the source repo):
+
+```
+~/.openjarvis/
+├── config.toml # Your configuration (git-tracked by checkpoint)
+├── agents/ # Agent prompts (git-tracked)
+├── tools/ # Tool descriptions (git-tracked)
+├── .git/ # Checkpoint repo for rollback
+└── learning/
+ ├── learning.db # SQLite session index
+ ├── benchmarks/ # Personal benchmark versions + gold answers
+ ├── sessions/ # Per-session artifacts (diagnosis, plan, traces)
+ └── pending_review/ # Edits awaiting user approval
+```
+
+## Background Daemon
+
+For continuous learning:
+
+```bash
+jarvis learning daemon start # Start background watcher
+jarvis learning daemon status # Check if running
+jarvis learning daemon stop # Stop the daemon
+```
+
+The daemon runs the scheduled trigger (default: daily at 03:00) and the cluster trigger (watches for failure patterns in real-time).
+
+## See Also
+
+- [Architecture: Learning](../architecture/learning.md#distillation-frontier-driven-harness-learning) — internal architecture of the distillation subsystem
+- [User Guide: Evaluations](evaluations.md) — the eval infrastructure that powers the benchmark gate
+- [User Guide: CLI](cli.md#jarvis-learning) — full CLI reference
+- [Getting Started: Configuration](../getting-started/configuration.md) — all config knobs
diff --git a/mkdocs.yml b/mkdocs.yml
index c5bec1bc..6c1c2c0d 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -180,5 +180,6 @@ nav:
- Scheduler: user-guide/scheduler.md
- Telemetry: user-guide/telemetry.md
- Security: user-guide/security.md
+ - Learning & Distillation: user-guide/learning-distillation.md
- Leaderboard: leaderboard.md
- Roadmap: development/roadmap.md
diff --git a/results b/results
new file mode 120000
index 00000000..4089db62
--- /dev/null
+++ b/results
@@ -0,0 +1 @@
+/scratch/user/jonsaadfalcon/results
\ No newline at end of file
diff --git a/scripts/experiments/a1_seed_feedback.py b/scripts/experiments/a1_seed_feedback.py
new file mode 100644
index 00000000..c6565d92
--- /dev/null
+++ b/scripts/experiments/a1_seed_feedback.py
@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+"""A1: Seed feedback on all traces in the local M1 traces.db.
+
+- Judges each unscored trace with Sonnet 4.6 using the calibration-validated prompt
+- Parallelized via ThreadPoolExecutor (8 workers) for I/O-bound API calls
+- Writes feedback to local traces.db via TraceStore.update_feedback
+- Logs every call to a JSONL audit file
+- Idempotent: skips traces that already have feedback (safe to re-run)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sqlite3
+import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+from pathlib import Path
+
+from openjarvis.core.types import Message, Role
+from openjarvis.engine.cloud import CloudEngine
+from openjarvis.traces.store import TraceStore
+
+HOME = Path(os.environ.get("OPENJARVIS_HOME", "/scratch/user/jonsaadfalcon/openjarvis-m1"))
+DB = HOME / "traces.db"
+LOG = HOME / "a1_feedback_log.jsonl"
+MODEL = "claude-sonnet-4-6"
+MAX_WORKERS = 8
+
+JUDGE_PROMPT = """\
+You are evaluating whether an AI agent successfully completed its assigned task.
+
+Assign a SCORE from the set {{0.2, 0.4, 0.6, 0.8}} using this rubric:
+
+- 0.8 = Clean success. Task completed correctly. Minor stylistic issues don't affect correctness.
+- 0.6 = Partial. Real progress made but the answer has real gaps — missed requirement, incomplete output, recovered from errors but final result imperfect.
+- 0.4 = Poor. Some progress but the result is clearly incomplete, wrong, or the agent got mostly stuck.
+- 0.2 = Failure. Agent crashed, got stuck in a loop, gave up, hit a budget/poll/token limit before finishing, or produced no usable result.
+
+IMPORTANT: Do not trust the agent's own self-assessment. Agents often narrate "I am stuck" or "I hit an error" — those are failure signals. Agents sometimes claim success when the actual output is incomplete — look at the concrete result, not the rhetoric.
+
+TASK QUERY (first 1200 chars):
+<<<
+{query}
+>>>
+
+AGENT FINAL RESULT (first 2500 chars):
+<<<
+{result_head}
+>>>
+{tail_section}
+Respond in EXACTLY this format, nothing else:
+SCORE=
+REASON=
+"""
+
+
+SCORE_RE = re.compile(r"SCORE\s*=\s*(0?\.[2468])", re.IGNORECASE)
+REASON_RE = re.compile(r"REASON\s*=\s*(.+?)\s*$", re.IGNORECASE | re.DOTALL)
+
+
+def build_prompt(query: str, result: str) -> str:
+ q = (query or "")[:1200]
+ head = (result or "")[:2500]
+ if result and len(result) > 3000:
+ tail = f"\nAGENT FINAL RESULT (last 500 chars):\n<<<\n{result[-500:]}\n>>>\n"
+ else:
+ tail = ""
+ return JUDGE_PROMPT.format(query=q, result_head=head, tail_section=tail)
+
+
+def judge_one(ce: CloudEngine, trace_id: str, query: str, result: str) -> dict:
+ prompt = build_prompt(query, result)
+ t0 = time.time()
+ try:
+ resp = ce.generate(
+ messages=[Message(role=Role.USER, content=prompt)],
+ model=MODEL,
+ max_tokens=150,
+ temperature=0.0,
+ )
+ content = resp.get("content", "") or ""
+ cost = resp.get("cost_usd", 0.0) or 0.0
+ usage = resp.get("usage", {}) or {}
+ in_tok = usage.get("prompt_tokens", 0) or usage.get("input_tokens", 0)
+ out_tok = usage.get("completion_tokens", 0) or usage.get("output_tokens", 0)
+
+ m = SCORE_RE.search(content)
+ score = float(m.group(1)) if m else None
+ mr = REASON_RE.search(content)
+ reason = mr.group(1).strip() if mr else "(parse failed)"
+
+ return {
+ "trace_id": trace_id,
+ "score": score,
+ "reason": reason,
+ "raw": content,
+ "cost": cost,
+ "input_tokens": in_tok,
+ "output_tokens": out_tok,
+ "elapsed": time.time() - t0,
+ "judged_at": datetime.utcnow().isoformat(timespec="seconds") + "Z",
+ "error": None,
+ }
+ except Exception as e:
+ return {
+ "trace_id": trace_id,
+ "score": None,
+ "reason": None,
+ "raw": None,
+ "cost": 0.0,
+ "input_tokens": 0,
+ "output_tokens": 0,
+ "elapsed": time.time() - t0,
+ "judged_at": datetime.utcnow().isoformat(timespec="seconds") + "Z",
+ "error": f"{type(e).__name__}: {e}",
+ }
+
+
+def main() -> int:
+ if not os.environ.get("ANTHROPIC_API_KEY"):
+ print("ERROR: ANTHROPIC_API_KEY not set", file=sys.stderr)
+ return 1
+
+ store = TraceStore(DB)
+ # Use the TraceStore's connection for the initial read
+ conn = store._conn
+ conn.row_factory = sqlite3.Row
+
+ # Pull all traces that need scoring
+ rows = list(conn.execute(
+ "SELECT trace_id, query, result, agent, model FROM traces "
+ "WHERE feedback IS NULL"
+ ))
+ already_scored = conn.execute(
+ "SELECT COUNT(*) FROM traces WHERE feedback IS NOT NULL"
+ ).fetchone()[0]
+
+ print(f"traces.db: {conn.execute('SELECT COUNT(*) FROM traces').fetchone()[0]} total")
+ print(f" already scored: {already_scored}")
+ print(f" to judge: {len(rows)}")
+ print(f"parallelism: {MAX_WORKERS} workers")
+ print(f"log: {LOG}")
+
+ if not rows:
+ print("Nothing to do.")
+ return 0
+
+ ce = CloudEngine()
+ write_lock = threading.Lock()
+ log_lock = threading.Lock()
+ log_fp = open(LOG, "a", encoding="utf-8")
+
+ # Write a header line marking this run
+ log_fp.write(json.dumps({
+ "_run_started": datetime.utcnow().isoformat(timespec="seconds") + "Z",
+ "model": MODEL,
+ "workers": MAX_WORKERS,
+ "to_judge": len(rows),
+ "already_scored": already_scored,
+ }) + "\n")
+ log_fp.flush()
+
+ total_cost = 0.0
+ done = 0
+ errors = 0
+ score_counts: dict = {}
+ t_start = time.time()
+
+ def on_result(res: dict) -> None:
+ nonlocal total_cost, done, errors
+ trace_id = res["trace_id"]
+
+ # Write feedback to DB (if we got a valid score)
+ if res["score"] is not None and res["error"] is None:
+ with write_lock:
+ store.update_feedback(trace_id, res["score"])
+ conn.commit()
+ else:
+ errors += 1
+
+ with log_lock:
+ log_fp.write(json.dumps(res) + "\n")
+ log_fp.flush()
+
+ total_cost += res["cost"] or 0.0
+ done += 1
+ score_counts[res["score"]] = score_counts.get(res["score"], 0) + 1
+
+ # Progress every 50 or on error
+ if done % 50 == 0 or res["error"]:
+ elapsed = time.time() - t_start
+ rate = done / max(0.001, elapsed)
+ eta = (len(rows) - done) / max(0.001, rate)
+ tag = "ERR " if res["error"] else " "
+ print(
+ f"{tag}[{done:4}/{len(rows)}] score={res['score']} "
+ f"cost=${total_cost:.3f} "
+ f"rate={rate:.1f}/s eta={eta:.0f}s "
+ f"errors={errors}"
+ )
+ if res["error"]:
+ print(f" ERROR on {trace_id[:12]}: {res['error']}")
+
+ with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+ futures = [
+ pool.submit(judge_one, ce, r["trace_id"], r["query"] or "", r["result"] or "")
+ for r in rows
+ ]
+ for fut in as_completed(futures):
+ res = fut.result()
+ on_result(res)
+
+ log_fp.close()
+
+ elapsed = time.time() - t_start
+ print(f"\n{'='*60}")
+ print(f"A1 COMPLETE in {elapsed:.1f}s ({elapsed/60:.1f}min)")
+ print(f"Total cost: ${total_cost:.4f}")
+ print(f"Errors: {errors}/{len(rows)}")
+ print(f"Score distribution:")
+ for k in sorted(score_counts.keys(), key=lambda x: (x is None, x)):
+ v = score_counts[k]
+ pct = 100 * v / len(rows)
+ label = {0.2: "failure", 0.4: "poor", 0.6: "partial", 0.8: "clean", None: "ERROR"}.get(k, "?")
+ print(f" {k} ({label}): {v} ({pct:.1f}%)")
+
+ # Verify by re-counting from DB
+ with_fb = conn.execute(
+ "SELECT COUNT(*) FROM traces WHERE feedback IS NOT NULL"
+ ).fetchone()[0]
+ above_gate = conn.execute(
+ "SELECT COUNT(*) FROM traces WHERE feedback >= 0.7"
+ ).fetchone()[0]
+ print(f"\nPost-A1 DB state:")
+ print(f" traces with feedback: {with_fb}")
+ print(f" traces passing 0.7 gate (eligible for personal benchmark): {above_gate}")
+ return 0 if errors == 0 else 2
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/scripts/experiments/generate_distillation_configs.py b/scripts/experiments/generate_distillation_configs.py
new file mode 100644
index 00000000..1ef98274
--- /dev/null
+++ b/scripts/experiments/generate_distillation_configs.py
@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+"""Generate all distillation experiment TOML configs.
+
+Produces configs for 7 experiment axes × multiple settings.
+Run: python scripts/experiments/generate_distillation_configs.py
+"""
+
+from __future__ import annotations
+
+import itertools
+from pathlib import Path
+
+CONFIGS_DIR = Path("src/openjarvis/evals/configs/distillation")
+
+# ── Teacher models ──────────────────────────────────────────────────────────
+TEACHERS = {
+ "opus": {
+ "model": "claude-opus-4-6",
+ "engine": "cloud",
+ "provider": "anthropic",
+ },
+ "gpt54": {
+ "model": "gpt-5.4",
+ "engine": "cloud",
+ "provider": "openai",
+ },
+ "gemini": {
+ "model": "gemini-3.1-pro-preview",
+ "engine": "cloud",
+ "provider": "google",
+ },
+ "qwen397b": {
+ "model": "Qwen/Qwen3.5-397B-A17B-FP8",
+ "engine": "vllm",
+ "provider": "local",
+ "note": "# Requires 8×H100, vLLM serve on port 8010",
+ },
+}
+
+# ── Student models ──────────────────────────────────────────────────────────
+# Served via vLLM on this H100 node. 27B uses the FP8 weights that fit on a
+# single H100; 2B and 9B use standard FP16.
+STUDENTS = {
+ "2b": {"model": "Qwen/Qwen3.5-2B", "engine": "vllm", "port": 8000},
+ "9b": {"model": "Qwen/Qwen3.5-9B", "engine": "vllm", "port": 8001},
+ "27b": {"model": "Qwen/Qwen3.5-27B-FP8", "engine": "vllm", "port": 8002},
+}
+
+# ── Benchmarks ──────────────────────────────────────────────────────────────
+BENCHMARKS = {
+ "pb": "pinchbench",
+ "tc15": "toolcall15",
+ "tb": "taubench",
+}
+
+# ── Data configs ────────────────────────────────────────────────────────────
+DATA_CONFIGS = {
+ "C1": {
+ "desc": "Zero test data — external traces only (GeneralThought + ADP)",
+ "trace_source": "external",
+ "benchmark_queries_visible": False,
+ },
+ "C2": {
+ "desc": "Test queries only — benchmark traces visible, answers hidden",
+ "trace_source": "benchmark",
+ "benchmark_queries_visible": True,
+ },
+ "C3": {
+ "desc": "Test queries + external — both benchmark and external traces",
+ "trace_source": "both",
+ "benchmark_queries_visible": True,
+ },
+}
+
+# ── Budget presets ──────────────────────────────────────────────────────────
+BUDGETS = {
+ "minimal": {"max_tool_calls": 5, "max_cost": 0.50},
+ "standard": {"max_tool_calls": 15, "max_cost": 2.00},
+ "thorough": {"max_tool_calls": 30, "max_cost": 5.00},
+ "exhaustive": {"max_tool_calls": 50, "max_cost": 10.00},
+}
+
+# ── Gate presets ────────────────────────────────────────────────────────────
+GATES = {
+ "permissive": {"min_improvement": 0.0, "max_regression": 0.10},
+ "standard": {"min_improvement": 0.0, "max_regression": 0.05},
+ "strict": {"min_improvement": 0.02, "max_regression": 0.02},
+ "none": {"min_improvement": -1.0, "max_regression": 1.0},
+}
+
+# ── Autonomy presets ────────────────────────────────────────────────────────
+AUTONOMY_MODES = ["auto", "tiered", "manual"]
+
+
+def render_config(
+ *,
+ experiment: str,
+ teacher_key: str,
+ student_key: str,
+ benchmark_key: str,
+ data_config_key: str = "C2",
+ budget_key: str = "standard",
+ gate_key: str = "standard",
+ autonomy: str = "auto",
+ iterative_sessions: int = 1,
+) -> str:
+ """Render a TOML config string."""
+ teacher = TEACHERS[teacher_key]
+ student = STUDENTS[student_key]
+ benchmark = BENCHMARKS[benchmark_key]
+ data_cfg = DATA_CONFIGS[data_config_key]
+ budget = BUDGETS[budget_key]
+ gate = GATES[gate_key]
+
+ note = teacher.get("note", "")
+ note_line = f"\n{note}" if note else ""
+
+ return f"""\
+# Distillation Experiment Config
+# Experiment: {experiment}
+# Teacher: {teacher["model"]} ({teacher["provider"]})
+# Student: {student["model"]} ({student["engine"]})
+# Benchmark: {benchmark}
+# Data config: {data_config_key} — {data_cfg["desc"]}
+# Budget: {budget_key} ({budget["max_tool_calls"]} tool calls, ${budget["max_cost"]:.2f})
+# Gate: {gate_key} (min_improvement={gate["min_improvement"]}, max_regression={gate["max_regression"]})
+# Autonomy: {autonomy}
+# Iterative sessions: {iterative_sessions}
+{note_line}
+
+[intelligence]
+default_model = "{student["model"]}"
+
+[engine]
+default = "{student["engine"]}"
+
+[engine.vllm]
+host = "http://localhost:{student.get("port", 8000)}"
+
+[learning.distillation]
+enabled = true
+autonomy_mode = "{autonomy}"
+teacher_model = "{teacher["model"]}"
+max_cost_per_session_usd = {budget["max_cost"]}
+max_tool_calls_per_diagnosis = {budget["max_tool_calls"]}
+
+[learning.distillation.gate]
+min_improvement = {gate["min_improvement"]}
+max_regression = {gate["max_regression"]}
+benchmark_subsample_size = 50
+
+[learning.distillation.benchmark]
+synthesis_feedback_threshold = 0.7
+max_benchmark_size = 200
+
+[learning.distillation.experiment]
+# Metadata for the experiment runner (not read by distillation itself)
+experiment_id = "{experiment}"
+teacher_key = "{teacher_key}"
+student_key = "{student_key}"
+benchmark = "{benchmark}"
+data_config = "{data_config_key}"
+trace_source = "{data_cfg["trace_source"]}"
+benchmark_queries_visible = {str(data_cfg["benchmark_queries_visible"]).lower()}
+budget_key = "{budget_key}"
+gate_key = "{gate_key}"
+iterative_sessions = {iterative_sessions}
+"""
+
+
+def write_config(subdir: str, filename: str, content: str) -> Path:
+ path = CONFIGS_DIR / subdir / filename
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(content, encoding="utf-8")
+ return path
+
+
+def generate_all() -> int:
+ count = 0
+
+ # ── Exp 1a: Teacher Model Ablation ────────────────────────────────────
+ # Fix: S-9b, B-standard, A-auto, G-standard, I-single
+ # Vary: teacher × benchmark × data_config
+ for teacher_key, bench_key, dc_key in itertools.product(
+ TEACHERS, BENCHMARKS, DATA_CONFIGS
+ ):
+ filename = f"{teacher_key}-9b-{bench_key}-{dc_key}.toml"
+ content = render_config(
+ experiment=f"exp1a-teacher/{teacher_key}-{bench_key}-{dc_key}",
+ teacher_key=teacher_key,
+ student_key="9b",
+ benchmark_key=bench_key,
+ data_config_key=dc_key,
+ )
+ write_config("exp1a-teacher", filename, content)
+ count += 1
+
+ # ── Exp 1b: Budget Ablation ───────────────────────────────────────────
+ # Fix: S-9b, T-sonnet(opus for quality), A-auto, G-standard, I-single
+ # Vary: budget × benchmark × data_config
+ for budget_key, bench_key, dc_key in itertools.product(
+ BUDGETS, BENCHMARKS, DATA_CONFIGS
+ ):
+ filename = f"{budget_key}-9b-{bench_key}-{dc_key}.toml"
+ content = render_config(
+ experiment=f"exp1b-budget/{budget_key}-{bench_key}-{dc_key}",
+ teacher_key="opus",
+ student_key="9b",
+ benchmark_key=bench_key,
+ data_config_key=dc_key,
+ budget_key=budget_key,
+ )
+ write_config("exp1b-budget", filename, content)
+ count += 1
+
+ # ── Exp 1c: Student Model Scaling ─────────────────────────────────────
+ # Fix: T-opus, B-standard, A-auto, G-standard, I-single
+ # Vary: student × benchmark × data_config
+ for student_key, bench_key, dc_key in itertools.product(
+ STUDENTS, BENCHMARKS, DATA_CONFIGS
+ ):
+ filename = f"opus-{student_key}-{bench_key}-{dc_key}.toml"
+ content = render_config(
+ experiment=f"exp1c-student/opus-{student_key}-{bench_key}-{dc_key}",
+ teacher_key="opus",
+ student_key=student_key,
+ benchmark_key=bench_key,
+ data_config_key=dc_key,
+ )
+ write_config("exp1c-student", filename, content)
+ count += 1
+
+ # ── Exp 2a: Gate Strictness ───────────────────────────────────────────
+ # Fix: S-9b, T-opus, B-standard, A-auto, I-single
+ # Vary: gate × benchmark
+ for gate_key, bench_key in itertools.product(GATES, BENCHMARKS):
+ filename = f"{gate_key}-9b-{bench_key}.toml"
+ content = render_config(
+ experiment=f"exp2a-gate/{gate_key}-{bench_key}",
+ teacher_key="opus",
+ student_key="9b",
+ benchmark_key=bench_key,
+ gate_key=gate_key,
+ )
+ write_config("exp2a-gate", filename, content)
+ count += 1
+
+ # ── Exp 2b: Autonomy Mode ────────────────────────────────────────────
+ # Fix: S-9b, T-opus, B-standard, G-standard, I-single
+ # Vary: autonomy × benchmark
+ for autonomy, bench_key in itertools.product(AUTONOMY_MODES, BENCHMARKS):
+ filename = f"{autonomy}-9b-{bench_key}.toml"
+ content = render_config(
+ experiment=f"exp2b-autonomy/{autonomy}-{bench_key}",
+ teacher_key="opus",
+ student_key="9b",
+ benchmark_key=bench_key,
+ autonomy=autonomy,
+ )
+ write_config("exp2b-autonomy", filename, content)
+ count += 1
+
+ # ── Exp 3a: Iterative Sessions ───────────────────────────────────────
+ # Fix: S-9b, T-opus, B-standard, A-auto, G-standard
+ # Vary: number of chained sessions × benchmark
+ for n_sessions, bench_key in itertools.product([1, 3, 5], BENCHMARKS):
+ filename = f"iter{n_sessions}-9b-{bench_key}.toml"
+ content = render_config(
+ experiment=f"exp3a-iterative/iter{n_sessions}-{bench_key}",
+ teacher_key="opus",
+ student_key="9b",
+ benchmark_key=bench_key,
+ iterative_sessions=n_sessions,
+ )
+ write_config("exp3a-iterative", filename, content)
+ count += 1
+
+ # ── Exp 3b: Cross-Benchmark Transfer ─────────────────────────────────
+ # Optimize using traces from benchmark X, eval on benchmark Y
+ for opt_bench, eval_bench in itertools.permutations(BENCHMARKS, 2):
+ filename = f"opt-{opt_bench}-eval-{eval_bench}-9b.toml"
+ content = render_config(
+ experiment=f"exp3b-transfer/opt-{opt_bench}-eval-{eval_bench}",
+ teacher_key="opus",
+ student_key="9b",
+ benchmark_key=opt_bench, # Traces from this benchmark
+ )
+ # Add eval benchmark as metadata
+ content += f'\neval_benchmark = "{BENCHMARKS[eval_bench]}"\n'
+ write_config("exp3b-transfer", filename, content)
+ count += 1
+
+ return count
+
+
+if __name__ == "__main__":
+ n = generate_all()
+ print(f"Generated {n} config files in {CONFIGS_DIR}/")
+
+ # Print summary
+ for subdir in sorted(CONFIGS_DIR.iterdir()):
+ if subdir.is_dir():
+ files = list(subdir.glob("*.toml"))
+ print(f" {subdir.name}/: {len(files)} configs")
diff --git a/scripts/experiments/m2_collect_results.py b/scripts/experiments/m2_collect_results.py
new file mode 100644
index 00000000..41319399
--- /dev/null
+++ b/scripts/experiments/m2_collect_results.py
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""M2: Collect distilled eval results and produce comparison table.
+
+Reads .summary.json files from results/neurips-2026/{distilled,baselines}/
+and produces a before/after comparison against the Step 1 baseline numbers.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+# Jon's Step 1 baselines — used when a local baseline result doesn't exist
+STEP1_BASELINES = {
+ "2b": {"toolcall15": 33.3, "pinchbench": 69.6, "livecodebench": 5.6, "taubench": 70.0, "taubench-telecom": 0.0, "gaia": 0.0, "liveresearch": 0.0, "liveresearchbench": None},
+ "9b": {"toolcall15": 46.7, "pinchbench": 95.7, "livecodebench": 17.6, "taubench": 85.0, "taubench-telecom": 80.0, "gaia": 38.0, "liveresearch": 75.0, "liveresearchbench": None},
+ "27b": {"toolcall15": 40.0, "pinchbench": 65.2, "livecodebench": 20.0, "taubench": 75.0, "taubench-telecom": 75.0, "gaia": 48.0, "liveresearch": 66.7, "liveresearchbench": None},
+}
+
+# Which benchmarks go through the agent layer (where distillation edits actually apply)
+AGENT_BENCHMARKS = {"pinchbench", "gaia", "liveresearch"}
+
+DISTILLED_ROOT = Path("results/neurips-2026/distilled")
+BASELINE_ROOT = Path("results/neurips-2026/baselines")
+
+
+def find_summary(root: Path, size: str, bench: str) -> Path | None:
+ """Find the summary JSON for a model × benchmark run."""
+ # Expected path: root/qwen-{size}/{bench}/{bench}_Qwen-Qwen3.5-{size}.summary.json
+ candidates = list(root.glob(f"qwen-{size}/{bench}/*.summary.json"))
+ return candidates[0] if candidates else None
+
+
+def load_accuracy(summary_path: Path) -> float | None:
+ """Extract overall accuracy from a summary.json file."""
+ try:
+ d = json.loads(summary_path.read_text())
+ # The summary has various shapes; try a few
+ for key in ["overall_accuracy", "accuracy", "overall_score"]:
+ if key in d:
+ return float(d[key]) * 100 if d[key] <= 1.0 else float(d[key])
+ # Try nested
+ if "results" in d:
+ for r in d["results"]:
+ if "accuracy" in r:
+ return float(r["accuracy"]) * 100 if r["accuracy"] <= 1.0 else float(r["accuracy"])
+ except Exception as e:
+ print(f" error reading {summary_path}: {e}", file=sys.stderr)
+ return None
+
+
+def main() -> int:
+ print("=" * 100)
+ print("M2 Distilled vs Baseline Comparison")
+ print("=" * 100)
+ print()
+ print(f"{'Model':8} {'Benchmark':20} {'Baseline':>10} {'Distilled':>10} {'Delta':>10} {'Agent?':>10}")
+ print("-" * 100)
+
+ benchmarks = [
+ "toolcall15", "pinchbench", "livecodebench",
+ "taubench", "taubench-telecom",
+ "gaia", "liveresearch", "liveresearchbench",
+ ]
+
+ summary_rows = []
+ for size in ["2b", "9b", "27b"]:
+ for bench in benchmarks:
+ # Baseline: prefer local file (if run this session), fall back to Step 1 numbers
+ baseline_path = find_summary(BASELINE_ROOT, size, bench)
+ if baseline_path:
+ baseline = load_accuracy(baseline_path)
+ base_source = "local"
+ else:
+ baseline = STEP1_BASELINES[size].get(bench)
+ base_source = "step1"
+
+ # Distilled: must be local from this session
+ distilled_path = find_summary(DISTILLED_ROOT, size, bench)
+ distilled = load_accuracy(distilled_path) if distilled_path else None
+
+ # Format
+ b_str = f"{baseline:.1f}%" if baseline is not None else "—"
+ d_str = f"{distilled:.1f}%" if distilled is not None else "pending"
+ if baseline is not None and distilled is not None:
+ delta = distilled - baseline
+ d_sign = "+" if delta >= 0 else ""
+ delta_str = f"{d_sign}{delta:.1f}%"
+ else:
+ delta_str = "—"
+ agent = "AGENT" if bench in AGENT_BENCHMARKS else "direct"
+
+ print(f"qwen-{size:4} {bench:20} {b_str:>10} {d_str:>10} {delta_str:>10} {agent:>10}")
+ summary_rows.append({
+ "model": f"qwen-{size}",
+ "benchmark": bench,
+ "baseline": baseline,
+ "distilled": distilled,
+ "delta": distilled - baseline if (baseline is not None and distilled is not None) else None,
+ "agent_benchmark": bench in AGENT_BENCHMARKS,
+ })
+ print()
+
+ # Aggregate: agent vs direct benchmark deltas
+ print("=" * 100)
+ print("Aggregate deltas by benchmark type (paper finding)")
+ print("=" * 100)
+ agent_deltas = [r["delta"] for r in summary_rows if r["delta"] is not None and r["agent_benchmark"]]
+ direct_deltas = [r["delta"] for r in summary_rows if r["delta"] is not None and not r["agent_benchmark"]]
+ if agent_deltas:
+ mean_agent = sum(agent_deltas) / len(agent_deltas)
+ print(f"Agent benchmarks (PB, GAIA, DeepResearchBench): mean delta = {mean_agent:+.2f}% over {len(agent_deltas)} runs")
+ if direct_deltas:
+ mean_direct = sum(direct_deltas) / len(direct_deltas)
+ print(f"Direct benchmarks (TC15, TauB, TBTel, LRB, LCB): mean delta = {mean_direct:+.2f}% over {len(direct_deltas)} runs")
+ print()
+
+ # Completion progress
+ expected = 24
+ distilled_count = sum(1 for r in summary_rows if r["distilled"] is not None)
+ print(f"Distilled runs complete: {distilled_count}/{expected}")
+
+ # Save JSON
+ out = Path("results/neurips-2026/distillation-m2/m2_comparison.json")
+ out.parent.mkdir(parents=True, exist_ok=True)
+ out.write_text(json.dumps(summary_rows, indent=2, default=str))
+ print(f"Full data: {out}")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/scripts/experiments/m2_create_distilled_configs.py b/scripts/experiments/m2_create_distilled_configs.py
new file mode 100644
index 00000000..a52e2eda
--- /dev/null
+++ b/scripts/experiments/m2_create_distilled_configs.py
@@ -0,0 +1,307 @@
+#!/usr/bin/env python3
+"""M2: Create distilled eval configs from M1 consensus edits.
+
+Generates 24 distilled configs (3 models × 8 benchmarks) by cloning
+baseline configs and applying the consensus edits from M1. Also creates
+4 missing baseline configs (livecodebench-qwen-9b, liveresearchbench-*).
+
+Usage: python scripts/experiments/m2_create_distilled_configs.py
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+CONFIGS_DIR = Path("src/openjarvis/evals/configs")
+M2_DIR = CONFIGS_DIR / "distillation" / "m2"
+
+# ── Consensus values from M1 (1,131 edits) ──────────────────────────────
+DISTILLED_TEMP = 0.2 # 84/134 votes (agent benchmarks only)
+DISTILLED_MAX_TURNS = 15 # 56/125 votes (close: 25 had 49)
+REMOVE_TOOLS = {"shell_exec", "http_request"} # 13 + 6 votes
+
+# ── Model specs ──────────────────────────────────────────────────────────
+MODELS = {
+ "2b": {"name": "Qwen/Qwen3.5-2B", "num_gpus": 1, "port": 8000},
+ "9b": {"name": "Qwen/Qwen3.5-9B", "num_gpus": 1, "port": 8001},
+ "27b": {"name": "Qwen/Qwen3.5-27B-FP8", "num_gpus": 1, "port": 8002},
+}
+
+# ── Benchmark specs ──────────────────────────────────────────────────────
+# Each benchmark defines its baseline config and what changes in distilled.
+BENCHMARKS = {
+ "toolcall15": {
+ "backend": "jarvis-direct",
+ "baseline_temp": 0.0,
+ "distilled_temp": 0.0, # CONTROL: no change for coding
+ "max_tokens": 4096,
+ "max_samples": None,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "extra_benchmark_fields": {},
+ },
+ "pinchbench": {
+ "backend": "jarvis-agent",
+ "agent": "native_openhands",
+ "baseline_temp": 0.6,
+ "distilled_temp": DISTILLED_TEMP,
+ "max_tokens": 8192,
+ "max_samples": None,
+ "judge_model": "claude-opus-4-5",
+ "judge_engine": "cloud",
+ "baseline_tools": [
+ "think", "file_read", "file_write", "web_search", "shell_exec",
+ "code_interpreter", "browser_navigate", "image_generate",
+ "calculator", "http_request", "pdf_extract",
+ ],
+ "extra_benchmark_fields": {},
+ },
+ "taubench": {
+ "backend": "jarvis-direct",
+ "baseline_temp": 0.7,
+ "distilled_temp": DISTILLED_TEMP,
+ "max_tokens": 4096,
+ "max_samples": 20,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "extra_benchmark_fields": {"split": "airline,retail"},
+ },
+ "taubench-telecom": {
+ "benchmark_name": "taubench", # same benchmark, different split
+ "backend": "jarvis-direct",
+ "baseline_temp": 0.7,
+ "distilled_temp": DISTILLED_TEMP,
+ "max_tokens": 4096,
+ "max_samples": 20,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "extra_benchmark_fields": {"split": "telecom"},
+ },
+ "gaia": {
+ "backend": "jarvis-agent",
+ "agent": "monitor_operative",
+ "baseline_temp": 0.6,
+ "distilled_temp": DISTILLED_TEMP,
+ "max_tokens": 8192,
+ "max_samples": 50,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "baseline_tools": [
+ "think", "calculator", "code_interpreter", "web_search", "file_read",
+ ],
+ "extra_benchmark_fields": {},
+ },
+ "liveresearch": {
+ "backend": "jarvis-agent",
+ "agent": "monitor_operative",
+ "baseline_temp": 0.6,
+ "distilled_temp": DISTILLED_TEMP,
+ "max_tokens": 16384,
+ "max_samples": 50,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "baseline_tools": [
+ "web_search", "file_read", "file_write", "code_interpreter", "think",
+ ],
+ "extra_benchmark_fields": {},
+ },
+ "liveresearchbench": {
+ "backend": "jarvis-direct",
+ "baseline_temp": 0.0,
+ "distilled_temp": 0.0, # CONTROL: reasoning benchmark
+ "max_tokens": 8192,
+ "max_samples": 50,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "extra_benchmark_fields": {},
+ },
+ "livecodebench": {
+ "backend": "jarvis-direct",
+ "baseline_temp": 0.0,
+ "distilled_temp": 0.0, # CONTROL: coding benchmark
+ "max_tokens": 4096,
+ "max_samples": 20,
+ "judge_model": "gpt-5-mini-2025-08-07",
+ "judge_engine": "cloud",
+ "extra_benchmark_fields": {},
+ },
+}
+
+
+def render_config(
+ *,
+ comment: str,
+ meta_name: str,
+ description: str,
+ temperature: float,
+ max_tokens: int,
+ judge_model: str,
+ judge_engine: str,
+ output_dir: str,
+ model_name: str,
+ model_engine: str,
+ num_gpus: int,
+ benchmark_name: str,
+ backend: str,
+ agent: str | None = None,
+ tools: list[str] | None = None,
+ max_samples: int | None = None,
+ extra_benchmark: dict | None = None,
+ seed: int = 42,
+) -> str:
+ lines = [f"# {comment}"]
+ lines.append(f'[meta]\nname = "{meta_name}"\ndescription = "{description}"\n')
+ lines.append(f"[defaults]\ntemperature = {temperature}\nmax_tokens = {max_tokens}\n")
+ lines.append(f'[judge]\nmodel = "{judge_model}"\ntemperature = 0.0')
+ if judge_engine:
+ lines.append(f'engine = "{judge_engine}"')
+ lines.append(f"max_tokens = 4096\n")
+ lines.append(f'[run]\nmax_workers = 1\noutput_dir = "{output_dir}"\nseed = {seed}\n')
+ lines.append(f'[[models]]\nname = "{model_name}"\nengine = "{model_engine}"\nnum_gpus = {num_gpus}\n')
+ lines.append(f'[[benchmarks]]\nname = "{benchmark_name}"\nbackend = "{backend}"')
+ if agent:
+ lines.append(f'agent = "{agent}"')
+ if max_samples:
+ lines.append(f"max_samples = {max_samples}")
+ if tools:
+ tools_str = ", ".join(f'"{t}"' for t in tools)
+ lines.append(f"tools = [{tools_str}]")
+ if extra_benchmark:
+ for k, v in extra_benchmark.items():
+ if isinstance(v, str):
+ lines.append(f'{k} = "{v}"')
+ else:
+ lines.append(f"{k} = {v}")
+ lines.append("")
+ return "\n".join(lines)
+
+
+def make_size_label(size: str) -> str:
+ return {"2b": "qwen-2b", "9b": "qwen-9b", "27b": "qwen-27b"}[size]
+
+
+def generate_missing_baselines() -> int:
+ """Create baseline configs that don't exist yet."""
+ count = 0
+
+ # livecodebench-qwen-9b (missing)
+ p = CONFIGS_DIR / "livecodebench-qwen-9b.toml"
+ if not p.exists():
+ b = BENCHMARKS["livecodebench"]
+ m = MODELS["9b"]
+ p.write_text(render_config(
+ comment="LiveCodeBench eval: Qwen3.5-9B (vLLM, 1 GPU)",
+ meta_name="livecodebench-qwen-9b",
+ description="LiveCodeBench on Qwen/Qwen3.5-9B (vLLM, 1 GPU)",
+ temperature=b["baseline_temp"],
+ max_tokens=b["max_tokens"],
+ judge_model=b["judge_model"],
+ judge_engine=b["judge_engine"],
+ output_dir="results/neurips-2026/baselines/qwen-9b/livecodebench/",
+ model_name=m["name"], model_engine="vllm", num_gpus=m["num_gpus"],
+ benchmark_name="livecodebench", backend=b["backend"],
+ max_samples=b["max_samples"],
+ ))
+ count += 1
+ print(f" created {p}")
+
+ # liveresearchbench-qwen-{2b,9b,27b}
+ for size, m in MODELS.items():
+ sl = make_size_label(size)
+ p = CONFIGS_DIR / f"liveresearchbench-{sl}.toml"
+ if not p.exists():
+ b = BENCHMARKS["liveresearchbench"]
+ p.write_text(render_config(
+ comment=f"LiveResearchBench (Salesforce): Qwen3.5-{size.upper()} (vLLM)",
+ meta_name=f"liveresearchbench-{sl}",
+ description=f"LiveResearchBench on {m['name']} (vLLM)",
+ temperature=b["baseline_temp"],
+ max_tokens=b["max_tokens"],
+ judge_model=b["judge_model"],
+ judge_engine=b["judge_engine"],
+ output_dir=f"results/neurips-2026/baselines/{sl}/liveresearchbench/",
+ model_name=m["name"], model_engine="vllm", num_gpus=m["num_gpus"],
+ benchmark_name="liveresearchbench", backend=b["backend"],
+ max_samples=b["max_samples"],
+ ))
+ count += 1
+ print(f" created {p}")
+
+ return count
+
+
+def generate_distilled_configs() -> int:
+ """Create distilled configs for all 24 model × benchmark combos."""
+ M2_DIR.mkdir(parents=True, exist_ok=True)
+ count = 0
+
+ for size, model in MODELS.items():
+ sl = make_size_label(size)
+ for bench_key, bench in BENCHMARKS.items():
+ bench_name = bench.get("benchmark_name", bench_key)
+ is_agent = bench["backend"] == "jarvis-agent"
+ temp = bench["distilled_temp"]
+
+ # Tool list: remove broken tools for agent benchmarks
+ tools = None
+ if is_agent and "baseline_tools" in bench:
+ tools = [t for t in bench["baseline_tools"]
+ if t not in REMOVE_TOOLS]
+
+ fname = f"{bench_key}-{sl}-distilled.toml"
+ out_path = M2_DIR / fname
+
+ # Determine what changed for the comment
+ changes = []
+ if temp != bench["baseline_temp"]:
+ changes.append(f"temp {bench['baseline_temp']}→{temp}")
+ if tools and set(tools) != set(bench.get("baseline_tools", [])):
+ removed = set(bench.get("baseline_tools", [])) - set(tools)
+ changes.append(f"removed {removed}")
+ if is_agent:
+ changes.append(f"max_turns 10→{DISTILLED_MAX_TURNS} (via OPENJARVIS_CONFIG)")
+ change_str = "; ".join(changes) if changes else "CONTROL (no change)"
+
+ out_path.write_text(render_config(
+ comment=f"M2 DISTILLED: {bench_key} × {model['name']} — {change_str}",
+ meta_name=f"{bench_key}-{sl}-distilled",
+ description=f"Distilled {bench_key} on {model['name']}",
+ temperature=temp,
+ max_tokens=bench["max_tokens"],
+ judge_model=bench["judge_model"],
+ judge_engine=bench["judge_engine"],
+ output_dir=f"results/neurips-2026/distilled/{sl}/{bench_key}/",
+ model_name=model["name"], model_engine="vllm", num_gpus=model["num_gpus"],
+ benchmark_name=bench_name, backend=bench["backend"],
+ agent=bench.get("agent"),
+ tools=tools,
+ max_samples=bench.get("max_samples"),
+ extra_benchmark=bench.get("extra_benchmark_fields"),
+ ))
+ count += 1
+
+ return count
+
+
+if __name__ == "__main__":
+ print("=== Creating missing baseline configs ===")
+ n_base = generate_missing_baselines()
+ print(f"Created {n_base} missing baseline configs\n")
+
+ print("=== Creating distilled M2 configs ===")
+ n_dist = generate_distilled_configs()
+ print(f"Created {n_dist} distilled configs in {M2_DIR}/\n")
+
+ # Summary
+ print("=== Change matrix ===")
+ print(f"{'Benchmark':20} {'Backend':14} {'Temp Δ':12} {'Tool Δ':20} {'max_turns Δ':12}")
+ print("-" * 80)
+ for bk, b in BENCHMARKS.items():
+ is_agent = b["backend"] == "jarvis-agent"
+ temp_change = f"{b['baseline_temp']}→{b['distilled_temp']}" if b["baseline_temp"] != b["distilled_temp"] else "—"
+ tool_change = "—"
+ if is_agent and "baseline_tools" in b:
+ removed = REMOVE_TOOLS & set(b.get("baseline_tools", []))
+ tool_change = f"-{removed}" if removed else "—"
+ mt_change = f"10→{DISTILLED_MAX_TURNS}" if is_agent else "—"
+ print(f"{bk:20} {b['backend']:14} {temp_change:12} {str(tool_change):20} {mt_change:12}")
diff --git a/scripts/experiments/m2_run_distilled_evals.sh b/scripts/experiments/m2_run_distilled_evals.sh
new file mode 100644
index 00000000..9e17e48b
--- /dev/null
+++ b/scripts/experiments/m2_run_distilled_evals.sh
@@ -0,0 +1,183 @@
+#!/usr/bin/env bash
+# ──────────────────────────────────────────────────────────────────────────────
+# M2: Run distilled eval configs — resumable, agent-benchmarks-first ordering
+#
+# Assumes vLLM already running: 2B:8000, 9B:8001, 27B-FP8:8002
+#
+# Usage:
+# bash m2_run_distilled_evals.sh # all
+# bash m2_run_distilled_evals.sh 9b # 9b only
+# bash m2_run_distilled_evals.sh 9b gaia # 9b + gaia only
+# ──────────────────────────────────────────────────────────────────────────────
+
+set -uo pipefail
+
+VENV=".venv/bin/python"
+M2_CONFIGS="src/openjarvis/evals/configs/distillation/m2"
+BASELINE_CONFIGS="src/openjarvis/evals/configs"
+M2_HOME="/scratch/user/jonsaadfalcon/openjarvis-m2"
+MODEL_FILTER=${1:-all}
+BENCH_FILTER=${2:-all}
+FORCE=${FORCE:-0} # set FORCE=1 to re-run completed configs
+
+RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
+log() { echo -e "${BLUE}[m2]${NC} $*"; }
+ok() { echo -e "${GREEN}[ OK ]${NC} $*"; }
+warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
+fail() { echo -e "${RED}[FAIL]${NC} $*"; }
+skip() { echo -e "${YELLOW}[SKIP]${NC} $*"; }
+
+declare -A MODEL_PORT=( [2b]=8000 [9b]=8001 [27b]=8002 )
+
+# Agent benchmarks FIRST (where distillation impact is expected)
+AGENT_BENCHMARKS="pinchbench gaia liveresearch"
+DIRECT_BENCHMARKS="toolcall15 taubench taubench-telecom livecodebench liveresearchbench"
+ALL_BENCHMARKS="${AGENT_BENCHMARKS} ${DIRECT_BENCHMARKS}"
+
+check_vllm() {
+ for size in 2b 9b 27b; do
+ [ "$MODEL_FILTER" != "all" ] && [ "$MODEL_FILTER" != "$size" ] && continue
+ local port=${MODEL_PORT[$size]}
+ if ! curl -sf "http://localhost:${port}/v1/models" >/dev/null 2>&1; then
+ fail "vLLM ${size} not responding on port ${port}"
+ return 1
+ fi
+ ok "vLLM ${size} healthy on port ${port}"
+ done
+}
+
+# Check if a run already completed (summary.json exists AND has a real accuracy).
+# A summary with errors=total_samples (like my earlier broken-routing tests)
+# is treated as incomplete and re-run.
+is_complete() {
+ local summary_path=$1
+ [ -f "$summary_path" ] || return 1
+ python3 -c "
+import json, sys
+try:
+ d = json.load(open('$summary_path'))
+ total = d.get('total_samples', 0)
+ errors = d.get('errors', 0)
+ scored = d.get('scored_samples', 0)
+ # Complete if at least some samples were scored successfully
+ sys.exit(0 if scored > 0 else 1)
+except Exception:
+ sys.exit(1)
+" 2>/dev/null
+}
+
+run_eval() {
+ local config_path=$1 label=$2 size=$3 summary_path=$4 use_distilled=${5:-false}
+
+ if [ "$FORCE" != "1" ] && is_complete "$summary_path"; then
+ skip "${label} [already complete]"
+ return 0
+ fi
+
+ local oj_config="${M2_HOME}/config-baseline-${size}.toml"
+ [ "$use_distilled" = "true" ] && oj_config="${M2_HOME}/config-${size}.toml"
+
+ log "Running: ${label} [$(basename ${oj_config})]"
+ env OPENJARVIS_CONFIG="${oj_config}" ${VENV} -m openjarvis.evals run -c "${config_path}" 2>&1
+ local rc=$?
+ if [ $rc -eq 0 ] && is_complete "$summary_path"; then
+ ok "Done: ${label}"
+ else
+ warn "Failed: ${label} (rc=$rc)"
+ fi
+}
+
+# Derive the expected summary.json path for a given distilled config
+summary_for_distilled() {
+ local bench=$1 size=$2
+ # Output dir from the config template: results/neurips-2026/distilled/qwen-{size}/{bench}/
+ # Summary file pattern: {bench}_Qwen-Qwen3.5-{size}.summary.json
+ local model_slug
+ if [ "$size" = "27b" ]; then model_slug="Qwen-Qwen3.5-27B-FP8"
+ else model_slug="Qwen-Qwen3.5-${size}"; fi
+ # Uppercase B for the model slug
+ model_slug=$(echo "$model_slug" | sed 's/-\([0-9][0-9]*\)b/-\1B/g')
+ # For taubench-telecom, the benchmark name in output is "taubench" not "taubench-telecom"
+ local bench_fname=$bench
+ [ "$bench" = "taubench-telecom" ] && bench_fname="taubench"
+ echo "results/neurips-2026/distilled/qwen-${size}/${bench}/${bench_fname}_${model_slug}.summary.json"
+}
+
+summary_for_baseline() {
+ local bench=$1 size=$2
+ local model_slug
+ if [ "$size" = "27b" ]; then model_slug="Qwen-Qwen3.5-27B-FP8"
+ else model_slug="Qwen-Qwen3.5-${size}"; fi
+ model_slug=$(echo "$model_slug" | sed 's/-\([0-9][0-9]*\)b/-\1B/g')
+ local bench_fname=$bench
+ [ "$bench" = "taubench-telecom" ] && bench_fname="taubench"
+ echo "results/neurips-2026/baselines/qwen-${size}/${bench}/${bench_fname}_${model_slug}.summary.json"
+}
+
+log "M2 Distilled Eval Runner (resumable, agent-first)"
+log "Model filter: ${MODEL_FILTER} Benchmark filter: ${BENCH_FILTER} FORCE=${FORCE}"
+
+check_vllm || exit 1
+
+start_time=$(date +%s)
+
+# Phase B1: DISTILLED agent benchmarks (highest-priority: PinchBench, GAIA, DeepResearchBench)
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+log "Phase B1: DISTILLED agent benchmarks (9 runs — the critical data)"
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+for size in 2b 9b 27b; do
+ [ "$MODEL_FILTER" != "all" ] && [ "$MODEL_FILTER" != "$size" ] && continue
+ for bench in ${AGENT_BENCHMARKS}; do
+ [ "$BENCH_FILTER" != "all" ] && [ "$BENCH_FILTER" != "$bench" ] && continue
+ cfg="${M2_CONFIGS}/${bench}-qwen-${size}-distilled.toml"
+ sum=$(summary_for_distilled "$bench" "$size")
+ [ -f "$cfg" ] && run_eval "$cfg" "DISTILLED ${bench}-qwen-${size}" "${size}" "$sum" true
+ done
+done
+
+# Phase B2: DISTILLED direct benchmarks (controls — should show minimal delta)
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+log "Phase B2: DISTILLED direct benchmarks (15 runs — controls)"
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+for size in 2b 9b 27b; do
+ [ "$MODEL_FILTER" != "all" ] && [ "$MODEL_FILTER" != "$size" ] && continue
+ for bench in ${DIRECT_BENCHMARKS}; do
+ [ "$BENCH_FILTER" != "all" ] && [ "$BENCH_FILTER" != "$bench" ] && continue
+ cfg="${M2_CONFIGS}/${bench}-qwen-${size}-distilled.toml"
+ sum=$(summary_for_distilled "$bench" "$size")
+ [ -f "$cfg" ] && run_eval "$cfg" "DISTILLED ${bench}-qwen-${size}" "${size}" "$sum" true
+ done
+done
+
+# Phase A: LiveResearchBench baselines (last, since Step 1 baselines exist for other benchmarks)
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+log "Phase A: LiveResearchBench baselines (3 runs — new benchmark only)"
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+for size in 2b 9b 27b; do
+ [ "$MODEL_FILTER" != "all" ] && [ "$MODEL_FILTER" != "$size" ] && continue
+ [ "$BENCH_FILTER" != "all" ] && [ "$BENCH_FILTER" != "liveresearchbench" ] && continue
+ cfg="${BASELINE_CONFIGS}/liveresearchbench-qwen-${size}.toml"
+ sum=$(summary_for_baseline "liveresearchbench" "$size")
+ [ -f "$cfg" ] && run_eval "$cfg" "BASELINE liveresearchbench-qwen-${size}" "${size}" "$sum" false
+done
+
+# Phase C: Spot-check 2 baselines against Jon's Step 1 numbers
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+log "Phase C: Spot-check baselines (TC15-9B, GAIA-9B)"
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+if [ "$MODEL_FILTER" = "all" ] || [ "$MODEL_FILTER" = "9b" ]; then
+ for bench in toolcall15 gaia; do
+ [ "$BENCH_FILTER" != "all" ] && [ "$BENCH_FILTER" != "$bench" ] && continue
+ cfg="${BASELINE_CONFIGS}/${bench}-qwen-9b.toml"
+ sum=$(summary_for_baseline "$bench" "9b")
+ [ -f "$cfg" ] && run_eval "$cfg" "SPOTCHECK ${bench}-qwen-9b" "9b" "$sum" false
+ done
+fi
+
+end_time=$(date +%s)
+elapsed=$((end_time - start_time))
+
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+ok "M2 complete in ${elapsed}s ($(( elapsed / 3600 ))h $(( (elapsed % 3600) / 60 ))m)"
+log "Distilled results: results/neurips-2026/distilled/"
+log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
diff --git a/scripts/experiments/m3_hill_climb.py b/scripts/experiments/m3_hill_climb.py
new file mode 100644
index 00000000..3c785182
--- /dev/null
+++ b/scripts/experiments/m3_hill_climb.py
@@ -0,0 +1,581 @@
+#!/usr/bin/env python3
+"""M3: Empirical hill-climbing with an LLM proposer.
+
+Replaces M1's open-loop "aggregate consensus across sessions" with a closed
+loop:
+
+ For each target (student, benchmark, agent):
+ for round in 1..N:
+ edit = teacher.propose_one(history_with_measured_deltas)
+ score_new = eval_subsample(apply(edit))
+ if score_new > current_score: accept
+
+Every proposal is empirically verified before the next is proposed, and the
+teacher sees measured deltas (not just traces) in its context.
+
+Usage:
+ python scripts/experiments/m3_hill_climb.py \\
+ --student 9b --benchmark liveresearch \\
+ --rounds 4 --k-subsample 8 --k-final 30
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import time
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+from openjarvis.core.types import Message, Role
+from openjarvis.engine.cloud import CloudEngine
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Config & constants
+# ═══════════════════════════════════════════════════════════════════════════
+
+STUDENT = {
+ "2b": {"name": "Qwen/Qwen3.5-2B", "port": 8000, "gpu": 4},
+ "9b": {"name": "Qwen/Qwen3.5-9B", "port": 8001, "gpu": 5},
+ "27b": {"name": "Qwen/Qwen3.5-27B-FP8", "port": 8002, "gpu": 6},
+}
+
+# Per-benchmark defaults (backend, baseline config)
+BENCHMARK = {
+ "liveresearch": {
+ "backend": "jarvis-agent",
+ "agent": "monitor_operative",
+ "baseline_temp": 0.6,
+ "baseline_max_tokens": 16384,
+ "baseline_max_turns": 10,
+ "baseline_tools": ["web_search", "file_read", "file_write",
+ "code_interpreter", "think"],
+ "max_samples_final": 50, # final eval
+ "judge": "gpt-5-mini-2025-08-07",
+ },
+ "gaia": {
+ "backend": "jarvis-agent",
+ "agent": "monitor_operative",
+ "baseline_temp": 0.6,
+ "baseline_max_tokens": 8192,
+ "baseline_max_turns": 10,
+ "baseline_tools": ["think", "calculator", "code_interpreter",
+ "web_search", "file_read"],
+ "max_samples_final": 50,
+ "judge": "gpt-5-mini-2025-08-07",
+ },
+ "pinchbench": {
+ "backend": "jarvis-agent",
+ "agent": "native_openhands",
+ "baseline_temp": 0.6,
+ "baseline_max_tokens": 8192,
+ "baseline_max_turns": 10,
+ "baseline_tools": ["think", "file_read", "file_write", "web_search",
+ "shell_exec", "code_interpreter", "browser_navigate",
+ "image_generate", "calculator", "http_request",
+ "pdf_extract"],
+ "max_samples_final": 23,
+ "judge": "claude-opus-4-5",
+ },
+}
+
+AVAILABLE_TOOLS = [
+ "think", "file_read", "file_write", "web_search", "shell_exec",
+ "code_interpreter", "browser_navigate", "image_generate", "calculator",
+ "http_request", "pdf_extract", "pdf_reader", "list_directory",
+]
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Proposer (LLM)
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+PROPOSER_SYSTEM = """\
+You are optimizing an OpenJarvis agent configuration for maximum accuracy on \
+a benchmark. You propose ONE config edit per round. After each proposal, the \
+edit is applied and the benchmark is run on a subsample; you then see the \
+measured score delta and decide the next edit.
+
+Your job is to find the config that maximizes measured accuracy.
+
+EDIT GRAMMAR — return JSON with "op" and the parameter fields at top level:
+
+{"op": "", , "rationale": ""}
+
+Valid ops and their parameter fields:
+
+1. set_temperature: "value" (float, 0.0..1.0)
+2. set_max_turns: "value" (int, 1..100)
+3. set_max_tokens: "value" (int, 512..32768)
+4. add_tool: "tool_name" (string; must be in AVAILABLE_TOOLS, not already active)
+5. remove_tool: "tool_name" (string; must be in current tools)
+6. noop: (no params; propose only if you believe no further edit will help)
+
+CONCRETE EXAMPLES:
+ {"op": "set_temperature", "value": 0.3, "rationale": "Reduce loop risk."}
+ {"op": "set_max_turns", "value": 20, "rationale": "More turns for research tasks."}
+ {"op": "add_tool", "tool_name": "pdf_extract", "rationale": "Tasks require PDF reading."}
+ {"op": "remove_tool", "tool_name": "shell_exec", "rationale": "Tool is broken in this env."}
+ {"op": "noop", "rationale": "Current config seems optimal."}
+
+EXPLORATION BIAS:
+ The config space has 5 distinct axes: temperature, max_turns, max_tokens,
+ tool additions (add_tool), tool removals (remove_tool). Before proposing a
+ second edit on an axis you've already tried, consider whether an untried
+ axis might reveal a larger gain. Tool-list edits (add/remove) often matter
+ more than numeric hyperparameters on benchmarks where the agent uses tools.
+
+Return ONLY the JSON object, no preamble, no code fences."""
+
+
+def build_user_prompt(
+ *, benchmark: str, student: str, agent: str,
+ baseline_score: float, current_score: float, current_config: dict,
+ edit_history: list[dict], available_tools: list[str],
+ sample_queries: list[str],
+) -> str:
+ hist_lines = []
+ for i, e in enumerate(edit_history, 1):
+ delta = e["score_after"] - e["score_before"]
+ status = "ACCEPTED" if e["accepted"] else "REJECTED"
+ hist_lines.append(
+ f" Round {i}: {json.dumps(e['edit'])} "
+ f"→ score {e['score_before']:.1f}% → {e['score_after']:.1f}% "
+ f"(Δ {delta:+.1f}, {status})"
+ )
+ hist = "\n".join(hist_lines) if hist_lines else " (no edits yet — baseline is the starting point)"
+
+ samples = "\n".join(f" - {q[:150]}..." for q in sample_queries[:3])
+
+ unused_tools = [t for t in available_tools if t not in current_config["tools"]]
+
+ return f"""\
+TARGET:
+ student: {student} (vLLM-served Qwen3.5)
+ benchmark: {benchmark}
+ agent: {agent} (uses OpenAI-format structured tool calls)
+
+CURRENT CONFIG:
+ temperature = {current_config['temperature']}
+ max_turns = {current_config['max_turns']}
+ max_tokens = {current_config['max_tokens']}
+ tools = {current_config['tools']}
+
+TOOLS NOT CURRENTLY ACTIVE (available to add):
+ {unused_tools}
+
+BASELINE (unedited) SCORE: {baseline_score:.1f}%
+CURRENT BEST SCORE: {current_score:.1f}%
+
+EDIT HISTORY (with measured deltas):
+{hist}
+
+SAMPLE TASKS FROM THIS BENCHMARK:
+{samples}
+
+Propose ONE edit that you predict will improve the measured accuracy. Consider \
+the edit history — do not repeat proposals that were rejected. If you believe \
+further edits will not help, propose noop.
+
+Return JSON only."""
+
+
+def call_proposer(
+ engine: CloudEngine, system: str, user: str, model: str = "claude-sonnet-4-6"
+) -> dict:
+ resp = engine.generate(
+ messages=[
+ Message(role=Role.SYSTEM, content=system),
+ Message(role=Role.USER, content=user),
+ ],
+ model=model, max_tokens=600, temperature=0.3,
+ )
+ content = (resp.get("content") or "").strip()
+ # Be forgiving about code fences
+ if content.startswith("```"):
+ content = re.sub(r"^```(?:json)?\s*", "", content)
+ content = re.sub(r"\s*```\s*$", "", content)
+ # Find the JSON object
+ m = re.search(r"\{[\s\S]*\}", content)
+ if not m:
+ raise ValueError(f"No JSON object found in proposer output: {content[:200]}")
+ return json.loads(m.group(0))
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Edit application & evaluation
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+@dataclass
+class Config:
+ temperature: float
+ max_turns: int
+ max_tokens: int
+ tools: list[str]
+
+
+def apply_edit(cfg: Config, edit: dict) -> Config:
+ """Apply an edit. Tolerant of both flat and nested (params) forms."""
+ op = edit["op"]
+ # Merge top-level edit fields with params for flat-or-nested tolerance
+ p = {**edit.get("params", {}), **{k: v for k, v in edit.items()
+ if k not in ("op", "params", "rationale")}}
+ new = Config(
+ temperature=cfg.temperature, max_turns=cfg.max_turns,
+ max_tokens=cfg.max_tokens, tools=list(cfg.tools),
+ )
+ if op == "noop":
+ return new
+ if op == "set_temperature":
+ new.temperature = float(p["value"])
+ elif op == "set_max_turns":
+ new.max_turns = int(p["value"])
+ elif op == "set_max_tokens":
+ new.max_tokens = int(p["value"])
+ elif op == "add_tool":
+ tool = p["tool_name"]
+ if tool not in new.tools:
+ new.tools.append(tool)
+ elif op == "remove_tool":
+ tool = p["tool_name"]
+ new.tools = [t for t in new.tools if t != tool]
+ else:
+ raise ValueError(f"unknown edit op: {op}")
+ return new
+
+
+def write_eval_toml(
+ *, bench: str, bench_spec: dict, student: dict, cfg: Config,
+ k_samples: int, output_dir: Path,
+) -> Path:
+ agent_line = f'agent = "{bench_spec["agent"]}"' if bench_spec.get("agent") else ""
+ tools_str = "[" + ", ".join(f'"{t}"' for t in cfg.tools) + "]"
+ toml = f"""\
+[meta]
+name = "m3-{bench}-{student['name'].replace('/', '-')}"
+description = "M3 hill-climb round"
+
+[defaults]
+temperature = {cfg.temperature}
+max_tokens = {cfg.max_tokens}
+
+[judge]
+model = "{bench_spec['judge']}"
+temperature = 0.0
+max_tokens = 4096
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "{output_dir}"
+seed = 42
+
+[[models]]
+name = "{student['name']}"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "{bench}"
+backend = "{bench_spec['backend']}"
+{agent_line}
+max_samples = {k_samples}
+tools = {tools_str}
+"""
+ path = output_dir / "eval.toml"
+ output_dir.mkdir(parents=True, exist_ok=True)
+ path.write_text(toml)
+ return path
+
+
+def write_openjarvis_config(home: Path, port: int, max_turns: int) -> Path:
+ p = home / "global-config.toml"
+ p.parent.mkdir(parents=True, exist_ok=True)
+ p.write_text(f"""\
+[agent]
+max_turns = {max_turns}
+
+[engine]
+default = "vllm"
+
+[engine.vllm]
+host = "http://localhost:{port}"
+""")
+ return p
+
+
+def run_eval(eval_toml: Path, oj_config: Path) -> tuple[float, int, int]:
+ """Run one eval. Returns (accuracy_pct, scored, total)."""
+ env = {**os.environ, "OPENJARVIS_CONFIG": str(oj_config)}
+ result = subprocess.run(
+ [".venv/bin/python", "-m", "openjarvis.evals", "run", "-c", str(eval_toml)],
+ capture_output=True, text=True, env=env, timeout=7200,
+ )
+ # Find the summary.json
+ out_dir = eval_toml.parent
+ sums = list(out_dir.glob("**/*.summary.json"))
+ if not sums:
+ print(f"[m3] WARNING: no summary.json in {out_dir}")
+ print(f"[m3] stderr tail: {result.stderr[-500:]}")
+ return 0.0, 0, 0
+ d = json.loads(sums[0].read_text())
+ acc = d.get("accuracy", 0.0)
+ acc_pct = acc * 100 if acc <= 1.0 else acc
+ return acc_pct, d.get("scored_samples", 0), d.get("total_samples", 0)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Benchmark sample loader (for proposer context)
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+def load_sample_queries(bench: str, n: int = 3) -> list[str]:
+ try:
+ if bench == "liveresearch":
+ from openjarvis.evals.datasets.liveresearch import LiveResearchBenchDataset
+ ds = LiveResearchBenchDataset()
+ ds.load(max_samples=n)
+ elif bench == "pinchbench":
+ from openjarvis.evals.datasets.pinchbench import PinchBenchDataset
+ ds = PinchBenchDataset()
+ elif bench == "gaia":
+ from openjarvis.evals.datasets.gaia import GAIADataset
+ ds = GAIADataset()
+ if hasattr(ds, "load"):
+ ds.load(max_samples=n)
+ else:
+ return []
+ return [r.problem for r in list(ds.iter_records())[:n]]
+ except Exception as e:
+ print(f"[m3] WARNING: could not load samples for {bench}: {e}")
+ return []
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Main hill-climb loop
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+def hill_climb(args) -> dict:
+ bench_spec = BENCHMARK[args.benchmark]
+ student = STUDENT[args.student]
+
+ # State dir (resumable)
+ base_dir = Path(args.out_dir) / f"{args.student}-{args.benchmark}"
+ base_dir.mkdir(parents=True, exist_ok=True)
+ state_path = base_dir / "state.json"
+
+ # Load or initialize state
+ if state_path.exists() and not args.fresh:
+ state = json.loads(state_path.read_text())
+ print(f"[m3] Resumed state from round {len(state['history'])}/{args.rounds}")
+ else:
+ baseline_cfg_dict = {
+ "temperature": bench_spec["baseline_temp"],
+ "max_turns": bench_spec["baseline_max_turns"],
+ "max_tokens": bench_spec["baseline_max_tokens"],
+ "tools": list(bench_spec["baseline_tools"]),
+ }
+ # ALWAYS measure the baseline today first (unless user provides --trust-baseline).
+ # This prevents anchoring to an unreproducible Step 1 number.
+ # We measure at k=k_final so the final delta is like-for-like.
+ if args.trust_baseline and args.baseline_score is not None:
+ measured_baseline = args.baseline_score
+ measured_baseline_k = None
+ print(f"[m3] Trusting provided baseline score: {measured_baseline:.1f}% "
+ f"(--trust-baseline set)")
+ else:
+ print(f"[m3] Measuring today's baseline with k={args.k_final} "
+ f"(matches k_final for clean like-for-like delta)...")
+ bl_dir = base_dir / "baseline_measure"
+ bl_toml = write_eval_toml(
+ bench=args.benchmark, bench_spec=bench_spec, student=student,
+ cfg=Config(**baseline_cfg_dict),
+ k_samples=args.k_final, output_dir=bl_dir,
+ )
+ bl_oj = write_openjarvis_config(bl_dir, student["port"], baseline_cfg_dict["max_turns"])
+ t0 = time.monotonic()
+ measured_baseline, bl_scored, bl_total = run_eval(bl_toml, bl_oj)
+ measured_baseline_k = bl_scored
+ print(f"[m3] Measured baseline: {measured_baseline:.1f}% "
+ f"({bl_scored}/{bl_total}) in {(time.monotonic() - t0)/60:.1f} min")
+ if args.baseline_score is not None:
+ print(f"[m3] (reference: --baseline-score was {args.baseline_score:.1f}%, "
+ f"drift = {measured_baseline - args.baseline_score:+.1f})")
+
+ state = {
+ "args": vars(args),
+ "benchmark": args.benchmark,
+ "student": student["name"],
+ "agent": bench_spec["agent"],
+ "baseline_config": baseline_cfg_dict,
+ "baseline_score": measured_baseline,
+ "baseline_score_reference": args.baseline_score,
+ "current_config": dict(baseline_cfg_dict),
+ "current_score": measured_baseline,
+ "history": [],
+ }
+
+ # Save state helper
+ def save():
+ state_path.write_text(json.dumps(state, indent=2, default=str))
+
+ save()
+
+ # Sample queries for proposer context
+ sample_queries = load_sample_queries(args.benchmark)
+
+ engine = CloudEngine()
+
+ # Hill-climb rounds
+ for round_num in range(len(state["history"]) + 1, args.rounds + 1):
+ print(f"\n{'═' * 70}")
+ print(f"[m3] Round {round_num}/{args.rounds}")
+ print(f"[m3] current_score = {state['current_score']:.1f}%")
+ print(f"[m3] current_config = {state['current_config']}")
+
+ # Propose
+ user_prompt = build_user_prompt(
+ benchmark=args.benchmark, student=student["name"],
+ agent=bench_spec["agent"],
+ baseline_score=state["baseline_score"],
+ current_score=state["current_score"],
+ current_config=state["current_config"],
+ edit_history=state["history"],
+ available_tools=AVAILABLE_TOOLS,
+ sample_queries=sample_queries,
+ )
+ try:
+ edit = call_proposer(engine, PROPOSER_SYSTEM, user_prompt,
+ model=args.proposer_model)
+ except Exception as e:
+ print(f"[m3] Proposer failed: {e}. Ending hill-climb.")
+ break
+
+ print(f"[m3] Proposed: {json.dumps(edit)}")
+
+ if edit.get("op") == "noop":
+ print(f"[m3] Teacher proposed noop; stopping.")
+ break
+
+ # Apply and evaluate subsample
+ try:
+ candidate = apply_edit(Config(**state["current_config"]), edit)
+ except Exception as e:
+ print(f"[m3] apply_edit failed: {e}. Recording as malformed edit.")
+ # Record as a rejected malformed edit so teacher won't repeat
+ state["history"].append({
+ "round": round_num, "edit": edit,
+ "config_after": None,
+ "score_before": state["current_score"],
+ "score_after": state["current_score"], # no change
+ "scored": 0, "total": 0, "elapsed_seconds": 0,
+ "accepted": False,
+ "error": f"malformed_edit: {e}",
+ })
+ save()
+ continue
+
+ round_dir = base_dir / f"round_{round_num}"
+ eval_toml = write_eval_toml(
+ bench=args.benchmark, bench_spec=bench_spec, student=student,
+ cfg=candidate, k_samples=args.k_subsample,
+ output_dir=round_dir,
+ )
+ oj_cfg = write_openjarvis_config(round_dir, student["port"], candidate.max_turns)
+
+ print(f"[m3] Running k={args.k_subsample} subsample eval...")
+ t0 = time.monotonic()
+ acc, scored, total = run_eval(eval_toml, oj_cfg)
+ elapsed = time.monotonic() - t0
+ print(f"[m3] Subsample score: {acc:.1f}% ({scored}/{total}) in {elapsed/60:.1f} min")
+
+ score_before = state["current_score"]
+ delta = acc - score_before
+ accepted = delta > args.accept_threshold
+
+ # Record history
+ state["history"].append({
+ "round": round_num, "edit": edit,
+ "config_after": asdict(candidate),
+ "score_before": score_before, "score_after": acc,
+ "scored": scored, "total": total, "elapsed_seconds": elapsed,
+ "accepted": accepted,
+ })
+
+ if accepted:
+ state["current_config"] = asdict(candidate)
+ state["current_score"] = acc
+ print(f"[m3] ACCEPTED (Δ={delta:+.1f})")
+ else:
+ print(f"[m3] REJECTED (Δ={delta:+.1f} ≤ {args.accept_threshold})")
+
+ save()
+
+ # Final eval with current config
+ print(f"\n{'═' * 70}")
+ print(f"[m3] Final eval with best config: {state['current_config']}")
+ final_dir = base_dir / "final"
+ final_cfg = Config(**state["current_config"])
+ eval_toml = write_eval_toml(
+ bench=args.benchmark, bench_spec=bench_spec, student=student,
+ cfg=final_cfg, k_samples=args.k_final,
+ output_dir=final_dir,
+ )
+ oj_cfg = write_openjarvis_config(final_dir, student["port"], final_cfg.max_turns)
+
+ t0 = time.monotonic()
+ final_acc, final_scored, final_total = run_eval(eval_toml, oj_cfg)
+ elapsed = time.monotonic() - t0
+
+ state["final_score"] = final_acc
+ state["final_scored"] = final_scored
+ state["final_total"] = final_total
+ state["final_elapsed_seconds"] = elapsed
+ save()
+
+ print(f"\n{'═' * 70}")
+ print(f"[m3] DONE")
+ print(f"[m3] baseline = {state['baseline_score']:.1f}%")
+ print(f"[m3] final = {final_acc:.1f}% ({final_scored}/{final_total}) in {elapsed/60:.1f} min")
+ print(f"[m3] Δ vs baseline = {final_acc - state['baseline_score']:+.1f}")
+ print(f"[m3] state: {state_path}")
+
+ return state
+
+
+def main() -> int:
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--student", required=True, choices=list(STUDENT))
+ ap.add_argument("--benchmark", required=True, choices=list(BENCHMARK))
+ ap.add_argument("--rounds", type=int, default=4)
+ ap.add_argument("--k-subsample", type=int, default=8)
+ ap.add_argument("--k-final", type=int, default=30)
+ ap.add_argument("--baseline-score", type=float, default=None,
+ help="Optional reference baseline (shown alongside measured). "
+ "Hill-climb always measures today's baseline unless --trust-baseline.")
+ ap.add_argument("--trust-baseline", action="store_true",
+ help="Skip baseline re-measurement, trust --baseline-score.")
+ ap.add_argument("--accept-threshold", type=float, default=0.0,
+ help="Accept edit if score Δ > this (default: 0)")
+ ap.add_argument("--proposer-model", default="claude-sonnet-4-6")
+ ap.add_argument("--out-dir", default="results/neurips-2026/distillation-m3")
+ ap.add_argument("--fresh", action="store_true",
+ help="Overwrite existing state and start fresh")
+ args = ap.parse_args()
+
+ result = hill_climb(args)
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/scripts/experiments/run_distillation_experiments.sh b/scripts/experiments/run_distillation_experiments.sh
new file mode 100755
index 00000000..961db01c
--- /dev/null
+++ b/scripts/experiments/run_distillation_experiments.sh
@@ -0,0 +1,307 @@
+#!/usr/bin/env bash
+# ──────────────────────────────────────────────────────────────────────────────
+# Run distillation ablation experiments
+#
+# Prerequisites:
+# - Ollama running with qwen3.5:{2b,9b,27b}
+# - ANTHROPIC_API_KEY set (for Opus teacher)
+# - OPENAI_API_KEY set (for GPT-5.4 teacher)
+# - GOOGLE_API_KEY set (for Gemini teacher)
+# - For Qwen-397B teacher: vLLM serving on port 8010 with 8×H100
+# - Traces seeded with feedback (run A1 blocker first)
+# - jarvis learning init already run
+#
+# Usage:
+# bash scripts/experiments/run_distillation_experiments.sh # Run all
+# bash scripts/experiments/run_distillation_experiments.sh exp1a # Run Phase 1a only
+# bash scripts/experiments/run_distillation_experiments.sh exp1a opus # Single config
+# ──────────────────────────────────────────────────────────────────────────────
+
+set -euo pipefail
+
+CONFIGS_DIR="src/openjarvis/evals/configs/distillation"
+RESULTS_DIR="results/neurips-2026/agent-optimization/distillation"
+EXPERIMENT=${1:-all}
+FILTER=${2:-}
+
+# ── Colors ───────────────────────────────────────────────────────────────────
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m'
+
+log() { echo -e "${BLUE}[distill]${NC} $*"; }
+ok() { echo -e "${GREEN}[ OK ]${NC} $*"; }
+warn() { echo -e "${YELLOW}[ WARN ]${NC} $*"; }
+fail() { echo -e "${RED}[ FAIL ]${NC} $*"; }
+
+# ── Preflight checks ────────────────────────────────────────────────────────
+check_prereqs() {
+ log "Preflight checks..."
+
+ # Check API keys
+ if [ -z "${ANTHROPIC_API_KEY:-}" ]; then
+ warn "ANTHROPIC_API_KEY not set — Opus teacher experiments will fail"
+ fi
+ if [ -z "${OPENAI_API_KEY:-}" ]; then
+ warn "OPENAI_API_KEY not set — GPT-5.4 teacher experiments will fail"
+ fi
+ if [ -z "${GOOGLE_API_KEY:-}" ]; then
+ warn "GOOGLE_API_KEY not set — Gemini teacher experiments will fail"
+ fi
+
+ # Check Ollama
+ if ! ollama list &>/dev/null; then
+ fail "Ollama not running. Start it first."
+ exit 1
+ fi
+
+ # Check student models
+ for model in qwen3.5:2b qwen3.5:9b qwen3.5:27b; do
+ if ! ollama list 2>/dev/null | grep -q "$model"; then
+ warn "Model $model not found in Ollama. Pull with: ollama pull $model"
+ fi
+ done
+
+ # Check distillation init
+ if [ ! -d "$HOME/.openjarvis/learning" ]; then
+ log "Running jarvis learning init..."
+ uv run jarvis learning init
+ fi
+
+ ok "Preflight complete"
+}
+
+# ── Run a single distillation session ────────────────────────────────────────
+run_session() {
+ local config_file=$1
+ local experiment_name
+ experiment_name=$(basename "$(dirname "$config_file")")
+ local config_name
+ config_name=$(basename "${config_file%.toml}")
+ local output_dir="${RESULTS_DIR}/${experiment_name}/${config_name}"
+
+ # Skip if already completed
+ if [ -f "${output_dir}/session/session.json" ]; then
+ ok "SKIP ${experiment_name}/${config_name} (already done)"
+ return 0
+ fi
+
+ log "──────────────────────────────────────────────────────"
+ log "Experiment: ${experiment_name}/${config_name}"
+ log "Config: ${config_file}"
+ log "Output: ${output_dir}"
+ log "──────────────────────────────────────────────────────"
+
+ mkdir -p "${output_dir}"
+
+ # Extract metadata from config
+ local teacher_model
+ teacher_model=$(grep 'teacher_model' "$config_file" | head -1 | sed 's/.*= *"\(.*\)"/\1/')
+ local student_model
+ student_model=$(grep 'default_model' "$config_file" | head -1 | sed 's/.*= *"\(.*\)"/\1/')
+ local benchmark
+ benchmark=$(grep '^benchmark ' "$config_file" | head -1 | sed 's/.*= *"\(.*\)"/\1/')
+ local data_config
+ data_config=$(grep 'data_config' "$config_file" | head -1 | sed 's/.*= *"\(.*\)"/\1/')
+ local iterative
+ iterative=$(grep 'iterative_sessions' "$config_file" | head -1 | sed 's/.*= *//')
+
+ log "Teacher: ${teacher_model}"
+ log "Student: ${student_model}"
+ log "Data: ${data_config:-C2}"
+ log "Iter: ${iterative:-1}"
+
+ # ── Step 1: Seed traces based on data config ─────────────────────────
+ # (In a full implementation, this would filter/prepare the TraceStore
+ # based on C1/C2/C3. For now we use whatever traces exist.)
+
+ # ── Step 2: Run distillation session ─────────────────────────────────
+ local n_sessions=${iterative:-1}
+ local session_num=1
+ local prev_session_id=""
+
+ while [ "$session_num" -le "$n_sessions" ]; do
+ log "Session ${session_num}/${n_sessions}..."
+
+ local session_output="${output_dir}/session_${session_num}"
+ mkdir -p "${session_output}"
+
+ # Run the distillation session via Python
+ # (jarvis learning run doesn't support all config params yet,
+ # so we call the orchestrator directly)
+ uv run python << PYEOF > "${session_output}/run.log" 2>&1 || true
+import json, os, shutil, sys
+from pathlib import Path
+
+from openjarvis.engine.cloud import CloudEngine
+from openjarvis.evals.backends.jarvis_direct import JarvisDirectBackend
+from openjarvis.traces.store import TraceStore
+from openjarvis.learning.distillation.checkpoint.store import CheckpointStore
+from openjarvis.learning.distillation.models import AutonomyMode
+from openjarvis.learning.distillation.orchestrator import DistillationOrchestrator
+from openjarvis.learning.distillation.storage.session_store import SessionStore
+from openjarvis.learning.distillation.student_runner import (
+ VLLMStudentRunner,
+ build_benchmark_samples_from_traces,
+)
+from openjarvis.learning.distillation.triggers import OnDemandTrigger
+from openjarvis.learning.optimize.feedback.judge import TraceJudge
+
+home = Path(os.environ.get("OPENJARVIS_HOME", str(Path.home() / ".openjarvis")))
+
+# Read config params
+teacher_model = "${teacher_model}"
+student_model = "${student_model}"
+autonomy = "auto"
+max_cost = float("$(grep 'max_cost_per_session_usd' "$config_file" | head -1 | sed 's/.*= *//')")
+max_tools = int("$(grep 'max_tool_calls_per_diagnosis' "$config_file" | head -1 | sed 's/.*= *//')")
+
+# Real student runner via vLLM
+vllm_host = os.environ.get("VLLM_HOST", "http://localhost:8001")
+student_runner = VLLMStudentRunner(
+ host=vllm_host,
+ model=student_model,
+)
+
+# Real judge via cloud LLM
+cloud_engine = CloudEngine()
+judge_backend = JarvisDirectBackend(engine_key="cloud")
+judge = TraceJudge(backend=judge_backend, model="gpt-5-mini-2025-08-07")
+
+# Build benchmark samples from existing traces
+trace_store = TraceStore(home / "traces.db")
+benchmark_samples = build_benchmark_samples_from_traces(trace_store, limit=50)
+
+orch = DistillationOrchestrator(
+ teacher_engine=cloud_engine,
+ teacher_model=teacher_model,
+ trace_store=trace_store,
+ benchmark_samples=benchmark_samples,
+ student_runner=student_runner,
+ judge=judge,
+ session_store=SessionStore(home / "learning" / "learning.db"),
+ checkpoint_store=CheckpointStore(home),
+ openjarvis_home=home,
+ autonomy_mode=AutonomyMode.AUTO,
+ scorer=None,
+ min_traces=10,
+ max_cost_usd=max_cost,
+ max_tool_calls=max_tools,
+)
+session = orch.run(OnDemandTrigger())
+
+# Save results
+result = {
+ "session_id": session.id,
+ "status": session.status.value,
+ "cost_usd": session.teacher_cost_usd,
+ "edits_total": len(session.edit_outcomes),
+ "edits_applied": len([o for o in session.edit_outcomes if o.status == "applied"]),
+ "edits_rejected": len([o for o in session.edit_outcomes if o.status == "rejected_by_gate"]),
+ "error": session.error,
+}
+Path("${session_output}/result.json").write_text(json.dumps(result, indent=2))
+
+# Copy session artifacts
+sd = home / "learning" / "sessions" / session.id
+if sd.exists():
+ shutil.copytree(sd, Path("${session_output}/artifacts"), dirs_exist_ok=True)
+
+print(json.dumps(result, indent=2))
+PYEOF
+
+ # Check result
+ if [ -f "${session_output}/result.json" ]; then
+ local status
+ status=$(python3 -c "import json; print(json.load(open('${session_output}/result.json'))['status'])")
+ local cost
+ cost=$(python3 -c "import json; print(f\"\${json.load(open('${session_output}/result.json'))['cost_usd']:.4f}\")")
+ local applied
+ applied=$(python3 -c "import json; print(json.load(open('${session_output}/result.json'))['edits_applied'])")
+
+ if [ "$status" = "completed" ]; then
+ ok "Session ${session_num}: status=${status}, cost=\$${cost}, applied=${applied}"
+ else
+ warn "Session ${session_num}: status=${status}, cost=\$${cost}"
+ fi
+ else
+ fail "Session ${session_num}: no result.json (check ${session_output}/run.log)"
+ fi
+
+ session_num=$((session_num + 1))
+ done
+
+ ok "Done: ${experiment_name}/${config_name}"
+}
+
+# ── Run experiment group ─────────────────────────────────────────────────────
+run_experiment() {
+ local exp_dir=$1
+ local exp_name
+ exp_name=$(basename "$exp_dir")
+
+ log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+ log "EXPERIMENT GROUP: ${exp_name}"
+ log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+
+ local count=0
+ local total
+ total=$(ls "${exp_dir}"/*.toml 2>/dev/null | wc -l)
+
+ for config in "${exp_dir}"/*.toml; do
+ [ -f "$config" ] || continue
+
+ # Apply filter if specified
+ if [ -n "${FILTER}" ] && ! echo "$config" | grep -q "${FILTER}"; then
+ continue
+ fi
+
+ count=$((count + 1))
+ log "[${count}/${total}] $(basename "$config")"
+ run_session "$config"
+ done
+
+ ok "Experiment group ${exp_name}: ${count} configs processed"
+}
+
+# ── Main ─────────────────────────────────────────────────────────────────────
+main() {
+ check_prereqs
+
+ log "Starting distillation experiments"
+ log "Experiment filter: ${EXPERIMENT}"
+ log "Config filter: ${FILTER:-none}"
+
+ local start_time
+ start_time=$(date +%s)
+
+ if [ "$EXPERIMENT" = "all" ]; then
+ # Run in priority order
+ for exp in exp1a-teacher exp1b-budget exp1c-student \
+ exp2a-gate exp2b-autonomy \
+ exp3a-iterative exp3b-transfer; do
+ if [ -d "${CONFIGS_DIR}/${exp}" ]; then
+ run_experiment "${CONFIGS_DIR}/${exp}"
+ fi
+ done
+ elif [ -d "${CONFIGS_DIR}/${EXPERIMENT}" ]; then
+ run_experiment "${CONFIGS_DIR}/${EXPERIMENT}"
+ else
+ fail "Unknown experiment: ${EXPERIMENT}"
+ echo "Available: exp1a-teacher exp1b-budget exp1c-student exp2a-gate exp2b-autonomy exp3a-iterative exp3b-transfer"
+ exit 1
+ fi
+
+ local end_time
+ end_time=$(date +%s)
+ local elapsed=$((end_time - start_time))
+
+ log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+ ok "All experiments complete in ${elapsed}s"
+ log "Results in: ${RESULTS_DIR}/"
+ log "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+}
+
+main "$@"
diff --git a/scripts/experiments/run_track_b_gepa_dspy.sh b/scripts/experiments/run_track_b_gepa_dspy.sh
new file mode 100644
index 00000000..a628be4a
--- /dev/null
+++ b/scripts/experiments/run_track_b_gepa_dspy.sh
@@ -0,0 +1,375 @@
+#!/usr/bin/env bash
+# =============================================================================
+# Track B: GEPA/DSPy Agent Optimization
+# NeurIPS 2026 — Agent Optimization Experiments
+#
+# Runs GEPA and DSPy BootstrapFewShot across:
+# Models: qwen-9b, qwen-27b, qwen-35b
+# Benchmarks: toolcall15, pinchbench, taubench
+#
+# Usage:
+# bash scripts/experiments/run_track_b_gepa_dspy.sh
+# bash scripts/experiments/run_track_b_gepa_dspy.sh --model qwen-9b --benchmark pinchbench
+# bash scripts/experiments/run_track_b_gepa_dspy.sh --optimizer gepa
+# bash scripts/experiments/run_track_b_gepa_dspy.sh --optimizer dspy
+#
+# Expected runtime: ~2-4 hours per GEPA run, ~1-2 hours per DSPy run
+# Total wall-clock: ~12-18 hours (parallelized across GPUs)
+# Estimated API cost: ~$90 GEPA + ~$90 DSPy = ~$180 total
+#
+# =============================================================================
+# vLLM Serving Commands (run these BEFORE this script on the GPU node)
+# =============================================================================
+#
+# GPU 0 — Qwen-9B (1x GPU):
+# CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen2.5-7B-Instruct \
+# --model Qwen/Qwen2.5-7B-Instruct \
+# --served-model-name qwen-9b \
+# --port 8001 --host 0.0.0.0 \
+# --max-model-len 32768 --gpu-memory-utilization 0.9 &
+#
+# GPU 1 — Qwen-27B (1-2x GPU):
+# CUDA_VISIBLE_DEVICES=1,2 vllm serve Qwen/Qwen2.5-32B-Instruct \
+# --model Qwen/Qwen2.5-32B-Instruct \
+# --served-model-name qwen-27b \
+# --port 8002 --host 0.0.0.0 \
+# --tensor-parallel-size 2 \
+# --max-model-len 32768 --gpu-memory-utilization 0.9 &
+#
+# GPU 3 — Qwen-35B (1-2x GPU):
+# CUDA_VISIBLE_DEVICES=3,4 vllm serve Qwen/Qwen2.5-72B-Instruct \
+# --model Qwen/Qwen2.5-72B-Instruct \
+# --served-model-name qwen-35b \
+# --port 8003 --host 0.0.0.0 \
+# --tensor-parallel-size 2 \
+# --max-model-len 32768 --gpu-memory-utilization 0.9 &
+#
+# Wait for all servers to be healthy:
+# sleep 60 && curl -s http://localhost:8001/health && \
+# curl -s http://localhost:8002/health && \
+# curl -s http://localhost:8003/health
+#
+# =============================================================================
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Defaults
+# ---------------------------------------------------------------------------
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+RESULTS_BASE="${REPO_ROOT}/results/neurips-2026/agent-optimization"
+LOG_DIR="${REPO_ROOT}/results/neurips-2026/logs"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+ALL_MODELS=(qwen-9b qwen-27b qwen-35b)
+ALL_BENCHMARKS=(toolcall15 pinchbench taubench)
+ALL_OPTIMIZERS=(gepa dspy)
+
+# Override defaults with CLI flags
+FILTER_MODEL=""
+FILTER_BENCHMARK=""
+FILTER_OPTIMIZER=""
+
+# GEPA settings
+GEPA_TRIALS=20
+GEPA_MAX_SAMPLES=50
+GEPA_OPTIMIZER_MODEL="claude-sonnet-4-6"
+
+# DSPy settings
+DSPY_OPTIMIZER="BootstrapFewShotWithRandomSearch"
+DSPY_TEACHER_LM="claude-sonnet-4-6"
+
+# ---------------------------------------------------------------------------
+# Parse CLI flags
+# ---------------------------------------------------------------------------
+while [[ $# -gt 0 ]]; do
+ case "$1" in
+ --model)
+ FILTER_MODEL="$2"; shift 2 ;;
+ --benchmark)
+ FILTER_BENCHMARK="$2"; shift 2 ;;
+ --optimizer)
+ FILTER_OPTIMIZER="$2"; shift 2 ;;
+ --gepa-trials)
+ GEPA_TRIALS="$2"; shift 2 ;;
+ --gepa-max-samples)
+ GEPA_MAX_SAMPLES="$2"; shift 2 ;;
+ --dspy-optimizer)
+ DSPY_OPTIMIZER="$2"; shift 2 ;;
+ -h|--help)
+ sed -n '2,30p' "$0" | grep '^#' | sed 's/^# \?//'
+ exit 0 ;;
+ *)
+ echo "Unknown flag: $1"; exit 1 ;;
+ esac
+done
+
+# Apply filters
+if [[ -n "$FILTER_MODEL" ]]; then
+ ALL_MODELS=("$FILTER_MODEL")
+fi
+if [[ -n "$FILTER_BENCHMARK" ]]; then
+ ALL_BENCHMARKS=("$FILTER_BENCHMARK")
+fi
+if [[ -n "$FILTER_OPTIMIZER" ]]; then
+ ALL_OPTIMIZERS=("$FILTER_OPTIMIZER")
+fi
+
+# ---------------------------------------------------------------------------
+# Logging helpers
+# ---------------------------------------------------------------------------
+mkdir -p "$LOG_DIR"
+LOG_FILE="${LOG_DIR}/track_b_${TIMESTAMP}.log"
+
+log() {
+ local level="$1"; shift
+ local msg="[$(date '+%Y-%m-%d %H:%M:%S')] [$level] $*"
+ echo "$msg"
+ echo "$msg" >> "$LOG_FILE"
+}
+
+log_info() { log "INFO " "$@"; }
+log_ok() { log "OK " "$@"; }
+log_warn() { log "WARN " "$@"; }
+log_error() { log "ERROR" "$@"; }
+
+# ---------------------------------------------------------------------------
+# Environment setup
+# ---------------------------------------------------------------------------
+setup_env() {
+ log_info "=== Track B: GEPA/DSPy Optimization ==="
+ log_info "Repo: $REPO_ROOT"
+ log_info "Log: $LOG_FILE"
+ log_info "Models: ${ALL_MODELS[*]}"
+ log_info "Benchmarks: ${ALL_BENCHMARKS[*]}"
+ log_info "Optimizers: ${ALL_OPTIMIZERS[*]}"
+ echo ""
+
+ # Check we are in the repo root
+ if [[ ! -f "${REPO_ROOT}/pyproject.toml" ]]; then
+ log_error "pyproject.toml not found — is REPO_ROOT set correctly? ($REPO_ROOT)"
+ exit 1
+ fi
+
+ # Install/sync dependencies
+ log_info "Running uv sync..."
+ cd "$REPO_ROOT"
+ uv sync --extra dev 2>&1 | tail -5
+ log_ok "uv sync complete"
+
+ # Check required API keys
+ if [[ -z "${ANTHROPIC_API_KEY:-}" ]]; then
+ log_error "ANTHROPIC_API_KEY is not set. Required for the optimizer teacher model."
+ log_error " export ANTHROPIC_API_KEY=sk-ant-..."
+ exit 1
+ fi
+ log_ok "ANTHROPIC_API_KEY is set"
+
+ # Optional: OpenAI key (used if teacher_lm is an OpenAI model)
+ if [[ -z "${OPENAI_API_KEY:-}" ]]; then
+ log_warn "OPENAI_API_KEY not set (only needed if using OpenAI teacher models)"
+ fi
+
+ echo ""
+ log_info "Model-to-port mapping (vLLM must be pre-started on these ports):"
+ log_info " qwen-9b -> http://localhost:8001"
+ log_info " qwen-27b -> http://localhost:8002"
+ log_info " qwen-35b -> http://localhost:8003"
+ echo ""
+}
+
+# ---------------------------------------------------------------------------
+# Model port lookup
+# ---------------------------------------------------------------------------
+model_port() {
+ case "$1" in
+ qwen-9b) echo 8001 ;;
+ qwen-27b) echo 8002 ;;
+ qwen-35b) echo 8003 ;;
+ *)
+ log_error "Unknown model: $1"
+ exit 1 ;;
+ esac
+}
+
+# ---------------------------------------------------------------------------
+# Health check: verify the vLLM server for a model is reachable
+# ---------------------------------------------------------------------------
+check_server_health() {
+ local model="$1"
+ local port
+ port="$(model_port "$model")"
+ local url="http://localhost:${port}/health"
+ if curl -sf "$url" > /dev/null 2>&1; then
+ log_ok "vLLM server for $model is healthy at port $port"
+ return 0
+ else
+ log_error "vLLM server for $model NOT reachable at $url"
+ log_error "Start it with the vLLM commands in the script header."
+ return 1
+ fi
+}
+
+# ---------------------------------------------------------------------------
+# GEPA optimization for one (model, benchmark) pair
+# ---------------------------------------------------------------------------
+run_gepa() {
+ local model="$1"
+ local bench="$2"
+ local port
+ port="$(model_port "$model")"
+ local out_dir="${RESULTS_BASE}/gepa/${model}/${bench}"
+
+ log_info "--- GEPA: $model × $bench ---"
+ log_info " Output dir: $out_dir"
+ log_info " Trials: $GEPA_TRIALS Max-samples: $GEPA_MAX_SAMPLES"
+ mkdir -p "$out_dir"
+
+ # Record start time
+ local t0
+ t0="$(date +%s)"
+
+ OPENAI_API_BASE="http://localhost:${port}/v1" \
+ uv run jarvis optimize run \
+ --benchmark "$bench" \
+ --model "$model" \
+ --optimizer-model "$GEPA_OPTIMIZER_MODEL" \
+ --trials "$GEPA_TRIALS" \
+ --max-samples "$GEPA_MAX_SAMPLES" \
+ --output-dir "$out_dir" \
+ 2>&1 | tee -a "$LOG_FILE"
+
+ local exit_code=${PIPESTATUS[0]}
+ local t1
+ t1="$(date +%s)"
+ local elapsed=$(( t1 - t0 ))
+
+ if [[ $exit_code -eq 0 ]]; then
+ log_ok "GEPA $model/$bench done in ${elapsed}s"
+ else
+ log_error "GEPA $model/$bench FAILED (exit $exit_code) after ${elapsed}s"
+ return $exit_code
+ fi
+}
+
+# ---------------------------------------------------------------------------
+# DSPy optimization for one (model, benchmark) pair
+# ---------------------------------------------------------------------------
+run_dspy() {
+ local model="$1"
+ local bench="$2"
+ local port
+ port="$(model_port "$model")"
+ local out_dir="${RESULTS_BASE}/dspy/${model}/${bench}"
+
+ log_info "--- DSPy: $model × $bench ---"
+ log_info " Output dir: $out_dir"
+ log_info " Teleprompter: $DSPY_OPTIMIZER Teacher: $DSPY_TEACHER_LM"
+ mkdir -p "$out_dir"
+
+ local t0
+ t0="$(date +%s)"
+
+ OPENAI_API_BASE="http://localhost:${port}/v1" \
+ uv run python - <&1 | tee -a "$LOG_FILE"
+import sys
+from openjarvis.learning.agents.dspy_optimizer import DSPyAgentOptimizer
+from openjarvis.core.config import DSPyOptimizerConfig
+from openjarvis.traces.store import TraceStore
+
+store = TraceStore()
+config = DSPyOptimizerConfig(
+ optimizer="${DSPY_OPTIMIZER}",
+ teacher_lm="${DSPY_TEACHER_LM}",
+ config_dir="${out_dir}",
+ benchmark="${bench}",
+ agent_filter="${model}",
+)
+result = DSPyAgentOptimizer(config).optimize(store)
+print(f"DSPy result for ${model}/${bench}: {result}")
+if result.get("status") not in ("ok", "success", "done"):
+ sys.exit(1)
+PYEOF
+
+ local exit_code=${PIPESTATUS[0]}
+ local t1
+ t1="$(date +%s)"
+ local elapsed=$(( t1 - t0 ))
+
+ if [[ $exit_code -eq 0 ]]; then
+ log_ok "DSPy $model/$bench done in ${elapsed}s"
+ else
+ log_error "DSPy $model/$bench FAILED (exit $exit_code) after ${elapsed}s"
+ return $exit_code
+ fi
+}
+
+# ---------------------------------------------------------------------------
+# Summary: print result file locations
+# ---------------------------------------------------------------------------
+print_summary() {
+ echo ""
+ log_info "=== Track B Complete ==="
+ log_info "Results written to:"
+ for opt in "${ALL_OPTIMIZERS[@]}"; do
+ for model in "${ALL_MODELS[@]}"; do
+ for bench in "${ALL_BENCHMARKS[@]}"; do
+ local out_dir="${RESULTS_BASE}/${opt}/${model}/${bench}"
+ if [[ -d "$out_dir" ]]; then
+ log_ok " $opt/$model/$bench -> $out_dir"
+ else
+ log_warn " $opt/$model/$bench -> MISSING ($out_dir)"
+ fi
+ done
+ done
+ done
+ log_info "Full log: $LOG_FILE"
+}
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+main() {
+ setup_env
+
+ local failed=0
+
+ # Pre-flight: check vLLM servers for all target models
+ log_info "Checking vLLM server health..."
+ for model in "${ALL_MODELS[@]}"; do
+ check_server_health "$model" || failed=$(( failed + 1 ))
+ done
+ if [[ $failed -gt 0 ]]; then
+ log_error "$failed vLLM server(s) not reachable. Start them first (see header)."
+ exit 1
+ fi
+ echo ""
+
+ # Run all requested (optimizer, model, benchmark) combinations
+ for opt in "${ALL_OPTIMIZERS[@]}"; do
+ log_info "=========================================="
+ log_info "Optimizer: $opt"
+ log_info "=========================================="
+ for model in "${ALL_MODELS[@]}"; do
+ for bench in "${ALL_BENCHMARKS[@]}"; do
+ case "$opt" in
+ gepa) run_gepa "$model" "$bench" || failed=$(( failed + 1 )) ;;
+ dspy) run_dspy "$model" "$bench" || failed=$(( failed + 1 )) ;;
+ *) log_error "Unknown optimizer: $opt"; failed=$(( failed + 1 )) ;;
+ esac
+ echo ""
+ done
+ done
+ done
+
+ print_summary
+
+ if [[ $failed -gt 0 ]]; then
+ log_error "$failed run(s) failed. Check log for details: $LOG_FILE"
+ exit 1
+ fi
+
+ log_ok "All Track B runs completed successfully."
+}
+
+main "$@"
diff --git a/scripts/experiments/run_track_d_lora_sft.sh b/scripts/experiments/run_track_d_lora_sft.sh
new file mode 100644
index 00000000..e90b7415
--- /dev/null
+++ b/scripts/experiments/run_track_d_lora_sft.sh
@@ -0,0 +1,778 @@
+#!/usr/bin/env bash
+# =============================================================================
+# Track D: LoRA / SFT Intelligence Optimization
+# NeurIPS 2026 — Intelligence Optimization Experiments
+#
+# Runs LoRA and SFT fine-tuning across:
+# LoRA models: Qwen-2B, Qwen-9B, Qwen-27B
+# SFT models: Qwen-2B, Qwen-9B
+#
+# After training, runs fast-benchmark eval on every checkpoint.
+#
+# Usage:
+# bash scripts/experiments/run_track_d_lora_sft.sh
+# bash scripts/experiments/run_track_d_lora_sft.sh --method lora
+# bash scripts/experiments/run_track_d_lora_sft.sh --method sft --model qwen-2b
+# bash scripts/experiments/run_track_d_lora_sft.sh --skip-eval
+#
+# Expected runtime per run:
+# Qwen-2B LoRA (~4-8 h, 1x H100, GPU 0)
+# Qwen-9B LoRA (~8-16 h, 1x H100, GPU 1)
+# Qwen-27B LoRA (~16-24 h, 2x H100, GPUs 2-3)
+# Qwen-2B SFT (~8-16 h, 1x H100, GPU 4)
+# Qwen-9B SFT (~16-24 h, 2x H100, GPUs 5-6)
+#
+# Total wall-clock (all in parallel): ~24 h on a 7-8x H100 node
+# Total GPU-hours: ~100-200 H100-hours
+#
+# =============================================================================
+# GPU Allocation — recommended for a single 8x H100 node
+# =============================================================================
+#
+# Run D1 (Qwen-2B LoRA) on GPU 0:
+# CUDA_VISIBLE_DEVICES=0 bash scripts/experiments/run_track_d_lora_sft.sh \
+# --method lora --model qwen-2b &
+#
+# Run D2 (Qwen-9B LoRA) on GPU 1:
+# CUDA_VISIBLE_DEVICES=1 bash scripts/experiments/run_track_d_lora_sft.sh \
+# --method lora --model qwen-9b &
+#
+# Run D3 (Qwen-27B LoRA) on GPUs 2-3:
+# CUDA_VISIBLE_DEVICES=2,3 bash scripts/experiments/run_track_d_lora_sft.sh \
+# --method lora --model qwen-27b &
+#
+# Run D4 (Qwen-2B SFT) on GPU 4:
+# CUDA_VISIBLE_DEVICES=4 bash scripts/experiments/run_track_d_lora_sft.sh \
+# --method sft --model qwen-2b &
+#
+# Run D5 (Qwen-9B SFT) on GPUs 5-6:
+# CUDA_VISIBLE_DEVICES=5,6 bash scripts/experiments/run_track_d_lora_sft.sh \
+# --method sft --model qwen-9b &
+#
+# =============================================================================
+# Training datasets (downloaded automatically if HF_TOKEN is set)
+# =============================================================================
+#
+# Primary agentic traces:
+# - neulab/agent-data-collection
+# - GAIR/AgentInstruct
+#
+# Supplementary reasoning:
+# - GeneralThought-430K-filtered
+# - GLM-4.7-flash SFT traces (168K + 57K)
+#
+# =============================================================================
+
+set -euo pipefail
+
+# ---------------------------------------------------------------------------
+# Defaults
+# ---------------------------------------------------------------------------
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+RESULTS_BASE="${REPO_ROOT}/results/neurips-2026/intelligence-optimization"
+LOG_DIR="${REPO_ROOT}/results/neurips-2026/logs"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+# Default: run all methods and models
+ALL_LORA_MODELS=(qwen-2b qwen-9b qwen-27b)
+ALL_SFT_MODELS=(qwen-2b qwen-9b)
+FAST_BENCHMARKS=(toolcall15 pinchbench taubench)
+
+# CLI overrides
+FILTER_METHOD=""
+FILTER_MODEL=""
+SKIP_EVAL=false
+
+# Training hyperparameters (override per model below)
+LORA_RANK=64
+LORA_ALPHA=128
+LORA_DROPOUT=0.05
+LORA_EPOCHS=3
+LORA_LR=2e-4
+LORA_BATCH_SIZE=8
+LORA_GRAD_ACCUM=4
+
+SFT_EPOCHS=3
+SFT_LR=1e-5
+SFT_BATCH_SIZE=4
+SFT_GRAD_ACCUM=8
+
+MAX_SEQ_LEN=8192
+
+# Fast-eval settings
+EVAL_MAX_SAMPLES=20
+
+# ---------------------------------------------------------------------------
+# HuggingFace model IDs
+# ---------------------------------------------------------------------------
+model_hf_id() {
+ case "$1" in
+ qwen-2b) echo "Qwen/Qwen2.5-1.5B-Instruct" ;;
+ qwen-9b) echo "Qwen/Qwen2.5-7B-Instruct" ;;
+ qwen-27b) echo "Qwen/Qwen2.5-32B-Instruct" ;;
+ *)
+ echo "ERROR: unknown model $1" >&2
+ exit 1 ;;
+ esac
+}
+
+# ---------------------------------------------------------------------------
+# Parse CLI flags
+# ---------------------------------------------------------------------------
+while [[ $# -gt 0 ]]; do
+ case "$1" in
+ --method)
+ FILTER_METHOD="$2"; shift 2 ;;
+ --model)
+ FILTER_MODEL="$2"; shift 2 ;;
+ --skip-eval)
+ SKIP_EVAL=true; shift ;;
+ --lora-rank)
+ LORA_RANK="$2"; shift 2 ;;
+ --lora-epochs)
+ LORA_EPOCHS="$2"; shift 2 ;;
+ --sft-epochs)
+ SFT_EPOCHS="$2"; shift 2 ;;
+ -h|--help)
+ sed -n '2,35p' "$0" | grep '^#' | sed 's/^# \?//'
+ exit 0 ;;
+ *)
+ echo "Unknown flag: $1"; exit 1 ;;
+ esac
+done
+
+# Apply model filter
+if [[ -n "$FILTER_MODEL" ]]; then
+ ALL_LORA_MODELS=("$FILTER_MODEL")
+ ALL_SFT_MODELS=("$FILTER_MODEL")
+fi
+
+# Apply method filter
+RUN_LORA=true
+RUN_SFT=true
+if [[ "$FILTER_METHOD" == "lora" ]]; then
+ RUN_SFT=false
+elif [[ "$FILTER_METHOD" == "sft" ]]; then
+ RUN_LORA=false
+fi
+
+# ---------------------------------------------------------------------------
+# Logging helpers
+# ---------------------------------------------------------------------------
+mkdir -p "$LOG_DIR"
+LOG_FILE="${LOG_DIR}/track_d_${TIMESTAMP}.log"
+
+log() {
+ local level="$1"; shift
+ local msg="[$(date '+%Y-%m-%d %H:%M:%S')] [$level] $*"
+ echo "$msg"
+ echo "$msg" >> "$LOG_FILE"
+}
+
+log_info() { log "INFO " "$@"; }
+log_ok() { log "OK " "$@"; }
+log_warn() { log "WARN " "$@"; }
+log_error() { log "ERROR" "$@"; }
+
+# ---------------------------------------------------------------------------
+# Environment setup
+# ---------------------------------------------------------------------------
+setup_env() {
+ log_info "=== Track D: LoRA/SFT Training ==="
+ log_info "Repo: $REPO_ROOT"
+ log_info "Log: $LOG_FILE"
+ log_info "Results: $RESULTS_BASE"
+ echo ""
+
+ # Check we are in the repo root
+ if [[ ! -f "${REPO_ROOT}/pyproject.toml" ]]; then
+ log_error "pyproject.toml not found — is REPO_ROOT set correctly? ($REPO_ROOT)"
+ exit 1
+ fi
+
+ # Install/sync dependencies
+ log_info "Running uv sync..."
+ cd "$REPO_ROOT"
+ uv sync --extra dev 2>&1 | tail -5
+ log_ok "uv sync complete"
+
+ # HuggingFace token — required to download gated Qwen models
+ if [[ -z "${HF_TOKEN:-}" ]]; then
+ log_warn "HF_TOKEN is not set."
+ log_warn " If Qwen models are gated, set: export HF_TOKEN=hf_..."
+ log_warn " Or pre-download them with: huggingface-cli download "
+ else
+ log_ok "HF_TOKEN is set"
+ # Log in so huggingface_hub uses the token
+ uv run python -c "
+import huggingface_hub
+huggingface_hub.login(token='${HF_TOKEN}', add_to_git_credential=False)
+print('Logged in to HuggingFace Hub')
+" 2>&1 | tee -a "$LOG_FILE"
+ fi
+
+ # Check for GPU
+ if ! command -v nvidia-smi &>/dev/null; then
+ log_warn "nvidia-smi not found — ensure CUDA is available for training."
+ else
+ log_info "GPU status:"
+ nvidia-smi --query-gpu=index,name,memory.total,memory.free \
+ --format=csv,noheader 2>&1 | while IFS= read -r line; do
+ log_info " $line"
+ done
+ fi
+
+ # Check for trl / peft / transformers (training stack)
+ uv run python -c "
+import importlib, sys
+missing = []
+for pkg in ['transformers', 'peft', 'trl', 'datasets', 'accelerate', 'bitsandbytes']:
+ if importlib.util.find_spec(pkg) is None:
+ missing.append(pkg)
+if missing:
+ print('MISSING packages:', missing)
+ sys.exit(1)
+else:
+ print('Training stack OK: transformers, peft, trl, datasets, accelerate, bitsandbytes')
+" 2>&1 | tee -a "$LOG_FILE" || {
+ log_error "Some training packages are missing. Install with:"
+ log_error " uv add transformers peft trl datasets accelerate bitsandbytes"
+ exit 1
+ }
+
+ echo ""
+}
+
+# ---------------------------------------------------------------------------
+# Download / verify training dataset
+# ---------------------------------------------------------------------------
+prepare_dataset() {
+ local method="$1"
+ local model="$2"
+ local data_dir="${REPO_ROOT}/results/neurips-2026/training-data"
+ mkdir -p "$data_dir"
+
+ log_info "Preparing training dataset for $method/$model..."
+
+ uv run python - <&1 | tee -a "$LOG_FILE"
+from datasets import load_dataset
+import json
+from pathlib import Path
+
+data_dir = Path("${data_dir}")
+output_path = data_dir / "${method}_${model}_train.jsonl"
+
+if output_path.exists():
+ lines = output_path.read_text().count('\n')
+ print(f"Dataset already exists: {output_path} ({lines} examples)")
+else:
+ print("Downloading neulab/agent-data-collection...")
+ ds = load_dataset("neulab/agent-data-collection", split="train")
+ print(f"Raw dataset size: {len(ds)}")
+
+ # Filter to reasonable-length examples for agentic fine-tuning
+ examples = []
+ for ex in ds:
+ messages = ex.get("messages") or ex.get("conversations") or []
+ if not messages:
+ continue
+ total_len = sum(len(str(m)) for m in messages)
+ if 200 <= total_len <= 16000:
+ examples.append({"messages": messages})
+
+ print(f"Filtered dataset size: {len(examples)}")
+
+ with output_path.open("w") as f:
+ for ex in examples:
+ f.write(json.dumps(ex) + "\n")
+ print(f"Saved to {output_path}")
+PYEOF
+
+ echo "${data_dir}/${method}_${model}_train.jsonl"
+}
+
+# ---------------------------------------------------------------------------
+# LoRA training for one model
+# ---------------------------------------------------------------------------
+run_lora() {
+ local model="$1"
+ local hf_id
+ hf_id="$(model_hf_id "$model")"
+ local out_dir="${RESULTS_BASE}/lora/${model}"
+ local checkpoint_dir="${out_dir}/checkpoint"
+ mkdir -p "$out_dir" "$checkpoint_dir"
+
+ log_info "=========================================="
+ log_info "LoRA training: $model ($hf_id)"
+ log_info " Output: $out_dir"
+ log_info " LoRA rank=$LORA_RANK alpha=$LORA_ALPHA dropout=$LORA_DROPOUT"
+ log_info " Epochs=$LORA_EPOCHS LR=$LORA_LR batch=$LORA_BATCH_SIZE grad_accum=$LORA_GRAD_ACCUM"
+ log_info "=========================================="
+
+ local dataset_path
+ dataset_path="$(prepare_dataset lora "$model")"
+
+ local t0
+ t0="$(date +%s)"
+
+ # Number of GPUs visible
+ local n_gpus
+ n_gpus="$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1)"
+ log_info "Training on $n_gpus GPU(s)"
+
+ if [[ "$n_gpus" -gt 1 ]]; then
+ LAUNCHER="uv run torchrun --nproc_per_node=$n_gpus"
+ else
+ LAUNCHER="uv run python"
+ fi
+
+ $LAUNCHER - <&1 | tee -a "$LOG_FILE"
+import json
+import math
+import os
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, TaskType, get_peft_model
+from transformers import (
+ AutoModelForCausalLM,
+ AutoTokenizer,
+ DataCollatorForSeq2Seq,
+ Trainer,
+ TrainingArguments,
+)
+
+hf_id = "${hf_id}"
+out_dir = "${out_dir}"
+ckpt = "${checkpoint_dir}"
+data_path = "${dataset_path}"
+
+# ---- Tokenizer ----
+print(f"Loading tokenizer: {hf_id}")
+tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
+if tokenizer.pad_token is None:
+ tokenizer.pad_token = tokenizer.eos_token
+
+# ---- Dataset ----
+print(f"Loading dataset: {data_path}")
+raw = load_dataset("json", data_files=data_path, split="train")
+raw = raw.train_test_split(test_size=0.02, seed=42)
+
+def tokenize(example):
+ messages = example["messages"]
+ text = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=False,
+ )
+ enc = tokenizer(text, truncation=True, max_length=${MAX_SEQ_LEN})
+ enc["labels"] = enc["input_ids"].copy()
+ return enc
+
+print("Tokenizing dataset...")
+tok_ds = raw.map(tokenize, remove_columns=raw["train"].column_names, num_proc=4)
+print(f"Train: {len(tok_ds['train'])} Eval: {len(tok_ds['test'])}")
+
+# ---- Model ----
+print(f"Loading model: {hf_id}")
+model = AutoModelForCausalLM.from_pretrained(
+ hf_id,
+ torch_dtype=torch.bfloat16,
+ trust_remote_code=True,
+ device_map="auto",
+)
+
+# ---- LoRA ----
+lora_config = LoraConfig(
+ task_type=TaskType.CAUSAL_LM,
+ r=${LORA_RANK},
+ lora_alpha=${LORA_ALPHA},
+ lora_dropout=${LORA_DROPOUT},
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+ "gate_proj", "up_proj", "down_proj"],
+ bias="none",
+)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# ---- Training args ----
+steps_per_epoch = math.ceil(len(tok_ds["train"]) / (${LORA_BATCH_SIZE} * ${LORA_GRAD_ACCUM}))
+total_steps = steps_per_epoch * ${LORA_EPOCHS}
+save_steps = max(50, steps_per_epoch // 2)
+
+args = TrainingArguments(
+ output_dir=ckpt,
+ num_train_epochs=${LORA_EPOCHS},
+ per_device_train_batch_size=${LORA_BATCH_SIZE},
+ gradient_accumulation_steps=${LORA_GRAD_ACCUM},
+ learning_rate=${LORA_LR},
+ warmup_ratio=0.05,
+ lr_scheduler_type="cosine",
+ logging_steps=10,
+ save_steps=save_steps,
+ eval_strategy="steps",
+ eval_steps=save_steps,
+ load_best_model_at_end=True,
+ metric_for_best_model="eval_loss",
+ bf16=True,
+ gradient_checkpointing=True,
+ dataloader_num_workers=4,
+ report_to="none",
+ save_total_limit=3,
+)
+
+# ---- Trainer ----
+collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)
+trainer = Trainer(
+ model=model,
+ args=args,
+ train_dataset=tok_ds["train"],
+ eval_dataset=tok_ds["test"],
+ data_collator=collator,
+)
+
+print(f"Starting LoRA training: {total_steps} total steps")
+trainer.train()
+
+# Save final adapter
+final_adapter = os.path.join(out_dir, "lora_adapter_final")
+model.save_pretrained(final_adapter)
+tokenizer.save_pretrained(final_adapter)
+print(f"Final LoRA adapter saved to {final_adapter}")
+
+# Save training metadata
+meta = {
+ "model": hf_id,
+ "method": "lora",
+ "lora_rank": ${LORA_RANK},
+ "lora_alpha": ${LORA_ALPHA},
+ "epochs": ${LORA_EPOCHS},
+ "train_examples": len(tok_ds["train"]),
+ "final_train_loss": trainer.state.log_history[-1].get("loss"),
+}
+with open(os.path.join(out_dir, "training_meta.json"), "w") as f:
+ import json; json.dump(meta, f, indent=2)
+print("Training metadata saved.")
+PYEOF
+
+ local exit_code=${PIPESTATUS[0]}
+ local t1
+ t1="$(date +%s)"
+ local elapsed=$(( t1 - t0 ))
+
+ if [[ $exit_code -eq 0 ]]; then
+ log_ok "LoRA $model done in ${elapsed}s (~$(( elapsed / 3600 ))h $(( (elapsed % 3600) / 60 ))m)"
+ return 0
+ else
+ log_error "LoRA $model FAILED (exit $exit_code) after ${elapsed}s"
+ return $exit_code
+ fi
+}
+
+# ---------------------------------------------------------------------------
+# SFT (full fine-tuning) for one model
+# ---------------------------------------------------------------------------
+run_sft() {
+ local model="$1"
+ local hf_id
+ hf_id="$(model_hf_id "$model")"
+ local out_dir="${RESULTS_BASE}/sft/${model}"
+ local checkpoint_dir="${out_dir}/checkpoint"
+ mkdir -p "$out_dir" "$checkpoint_dir"
+
+ log_info "=========================================="
+ log_info "SFT training: $model ($hf_id)"
+ log_info " Output: $out_dir"
+ log_info " Epochs=$SFT_EPOCHS LR=$SFT_LR batch=$SFT_BATCH_SIZE grad_accum=$SFT_GRAD_ACCUM"
+ log_info "=========================================="
+
+ local dataset_path
+ dataset_path="$(prepare_dataset sft "$model")"
+
+ local t0
+ t0="$(date +%s)"
+
+ local n_gpus
+ n_gpus="$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1)"
+ log_info "Training on $n_gpus GPU(s)"
+
+ if [[ "$n_gpus" -gt 1 ]]; then
+ LAUNCHER="uv run torchrun --nproc_per_node=$n_gpus"
+ else
+ LAUNCHER="uv run python"
+ fi
+
+ $LAUNCHER - <&1 | tee -a "$LOG_FILE"
+import json
+import math
+import os
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+from transformers import (
+ AutoModelForCausalLM,
+ AutoTokenizer,
+ DataCollatorForSeq2Seq,
+ Trainer,
+ TrainingArguments,
+)
+
+hf_id = "${hf_id}"
+out_dir = "${out_dir}"
+ckpt = "${checkpoint_dir}"
+data_path = "${dataset_path}"
+
+# ---- Tokenizer ----
+print(f"Loading tokenizer: {hf_id}")
+tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
+if tokenizer.pad_token is None:
+ tokenizer.pad_token = tokenizer.eos_token
+
+# ---- Dataset ----
+print(f"Loading dataset: {data_path}")
+raw = load_dataset("json", data_files=data_path, split="train")
+raw = raw.train_test_split(test_size=0.02, seed=42)
+
+def tokenize(example):
+ messages = example["messages"]
+ text = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=False,
+ )
+ enc = tokenizer(text, truncation=True, max_length=${MAX_SEQ_LEN})
+ enc["labels"] = enc["input_ids"].copy()
+ return enc
+
+print("Tokenizing dataset...")
+tok_ds = raw.map(tokenize, remove_columns=raw["train"].column_names, num_proc=4)
+print(f"Train: {len(tok_ds['train'])} Eval: {len(tok_ds['test'])}")
+
+# ---- Model (4-bit quantized to fit on fewer GPUs) ----
+print(f"Loading model: {hf_id}")
+from transformers import BitsAndBytesConfig
+bnb_config = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_compute_dtype=torch.bfloat16,
+ bnb_4bit_quant_type="nf4",
+ bnb_4bit_use_double_quant=True,
+)
+model = AutoModelForCausalLM.from_pretrained(
+ hf_id,
+ quantization_config=bnb_config,
+ trust_remote_code=True,
+ device_map="auto",
+)
+model.config.use_cache = False
+
+# ---- Training args ----
+steps_per_epoch = math.ceil(len(tok_ds["train"]) / (${SFT_BATCH_SIZE} * ${SFT_GRAD_ACCUM}))
+total_steps = steps_per_epoch * ${SFT_EPOCHS}
+save_steps = max(50, steps_per_epoch // 2)
+
+args = TrainingArguments(
+ output_dir=ckpt,
+ num_train_epochs=${SFT_EPOCHS},
+ per_device_train_batch_size=${SFT_BATCH_SIZE},
+ gradient_accumulation_steps=${SFT_GRAD_ACCUM},
+ learning_rate=${SFT_LR},
+ warmup_ratio=0.03,
+ lr_scheduler_type="cosine",
+ logging_steps=10,
+ save_steps=save_steps,
+ eval_strategy="steps",
+ eval_steps=save_steps,
+ load_best_model_at_end=True,
+ metric_for_best_model="eval_loss",
+ bf16=True,
+ gradient_checkpointing=True,
+ dataloader_num_workers=4,
+ report_to="none",
+ save_total_limit=3,
+)
+
+# ---- Trainer ----
+collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)
+trainer = Trainer(
+ model=model,
+ args=args,
+ train_dataset=tok_ds["train"],
+ eval_dataset=tok_ds["test"],
+ data_collator=collator,
+)
+
+print(f"Starting SFT training: {total_steps} total steps")
+trainer.train()
+
+# Save final model
+final_model = os.path.join(out_dir, "sft_model_final")
+model.save_pretrained(final_model)
+tokenizer.save_pretrained(final_model)
+print(f"Final SFT model saved to {final_model}")
+
+# Save training metadata
+meta = {
+ "model": hf_id,
+ "method": "sft",
+ "epochs": ${SFT_EPOCHS},
+ "train_examples": len(tok_ds["train"]),
+ "final_train_loss": trainer.state.log_history[-1].get("loss"),
+}
+with open(os.path.join(out_dir, "training_meta.json"), "w") as f:
+ import json; json.dump(meta, f, indent=2)
+print("Training metadata saved.")
+PYEOF
+
+ local exit_code=${PIPESTATUS[0]}
+ local t1
+ t1="$(date +%s)"
+ local elapsed=$(( t1 - t0 ))
+
+ if [[ $exit_code -eq 0 ]]; then
+ log_ok "SFT $model done in ${elapsed}s (~$(( elapsed / 3600 ))h $(( (elapsed % 3600) / 60 ))m)"
+ return 0
+ else
+ log_error "SFT $model FAILED (exit $exit_code) after ${elapsed}s"
+ return $exit_code
+ fi
+}
+
+# ---------------------------------------------------------------------------
+# Post-training eval: run fast benchmarks on a checkpoint
+# ---------------------------------------------------------------------------
+run_eval() {
+ local method="$1" # lora or sft
+ local model="$2"
+ local out_dir="${RESULTS_BASE}/${method}/${model}"
+
+ if [[ "$SKIP_EVAL" == "true" ]]; then
+ log_info "Skipping eval (--skip-eval)"
+ return 0
+ fi
+
+ # Locate the final checkpoint / adapter
+ local checkpoint=""
+ if [[ "$method" == "lora" ]]; then
+ checkpoint="${out_dir}/lora_adapter_final"
+ else
+ checkpoint="${out_dir}/sft_model_final"
+ fi
+
+ if [[ ! -d "$checkpoint" ]]; then
+ log_warn "Checkpoint not found for $method/$model at $checkpoint — skipping eval"
+ return 0
+ fi
+
+ log_info "--- Post-training eval: $method/$model ---"
+ log_info " Checkpoint: $checkpoint"
+
+ for bench in "${FAST_BENCHMARKS[@]}"; do
+ local eval_out="${out_dir}/eval/${bench}"
+ mkdir -p "$eval_out"
+ log_info " Eval: $bench -> $eval_out"
+
+ uv run python - <&1 | tee -a "$LOG_FILE"
+import subprocess, sys
+
+cmd = [
+ "uv", "run", "python", "-m", "openjarvis.evals", "run",
+ "--model-path", "${checkpoint}",
+ "--model-id", "${model}-${method}",
+ "--benchmark", "${bench}",
+ "--max-samples", "${EVAL_MAX_SAMPLES}",
+ "--output", "${eval_out}",
+]
+print("Running:", " ".join(cmd))
+result = subprocess.run(cmd, capture_output=False)
+sys.exit(result.returncode)
+PYEOF
+
+ local eval_exit=${PIPESTATUS[0]}
+ if [[ $eval_exit -eq 0 ]]; then
+ log_ok " Eval $method/$model/$bench OK"
+ else
+ log_warn " Eval $method/$model/$bench returned exit $eval_exit (non-fatal)"
+ fi
+ done
+}
+
+# ---------------------------------------------------------------------------
+# Summary
+# ---------------------------------------------------------------------------
+print_summary() {
+ echo ""
+ log_info "=== Track D Complete ==="
+ log_info "Results written to:"
+
+ if [[ "$RUN_LORA" == "true" ]]; then
+ for model in "${ALL_LORA_MODELS[@]}"; do
+ local out="${RESULTS_BASE}/lora/${model}"
+ if [[ -d "$out" ]]; then
+ log_ok " lora/$model -> $out"
+ else
+ log_warn " lora/$model -> MISSING ($out)"
+ fi
+ done
+ fi
+
+ if [[ "$RUN_SFT" == "true" ]]; then
+ for model in "${ALL_SFT_MODELS[@]}"; do
+ local out="${RESULTS_BASE}/sft/${model}"
+ if [[ -d "$out" ]]; then
+ log_ok " sft/$model -> $out"
+ else
+ log_warn " sft/$model -> MISSING ($out)"
+ fi
+ done
+ fi
+
+ log_info "Full log: $LOG_FILE"
+}
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+main() {
+ setup_env
+
+ local failed=0
+
+ # ---- LoRA runs ----
+ if [[ "$RUN_LORA" == "true" ]]; then
+ log_info "=========================================="
+ log_info "Starting LoRA training runs"
+ log_info "Models: ${ALL_LORA_MODELS[*]}"
+ log_info "=========================================="
+ for model in "${ALL_LORA_MODELS[@]}"; do
+ run_lora "$model" || { failed=$(( failed + 1 )); log_error "LoRA $model failed, continuing..."; }
+ run_eval lora "$model"
+ echo ""
+ done
+ fi
+
+ # ---- SFT runs ----
+ if [[ "$RUN_SFT" == "true" ]]; then
+ log_info "=========================================="
+ log_info "Starting SFT training runs"
+ log_info "Models: ${ALL_SFT_MODELS[*]}"
+ log_info "=========================================="
+ for model in "${ALL_SFT_MODELS[@]}"; do
+ run_sft "$model" || { failed=$(( failed + 1 )); log_error "SFT $model failed, continuing..."; }
+ run_eval sft "$model"
+ echo ""
+ done
+ fi
+
+ print_summary
+
+ if [[ $failed -gt 0 ]]; then
+ log_error "$failed training run(s) failed. Check log: $LOG_FILE"
+ exit 1
+ fi
+
+ log_ok "All Track D runs completed successfully."
+}
+
+main "$@"
diff --git a/src/openjarvis/agents/deep_research.py b/src/openjarvis/agents/deep_research.py
index 4d4d63af..b9632a04 100644
--- a/src/openjarvis/agents/deep_research.py
+++ b/src/openjarvis/agents/deep_research.py
@@ -10,6 +10,10 @@ from __future__ import annotations
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
+from openjarvis.agents.prompt_loader import (
+ load_few_shot_exemplars,
+ load_system_prompt_override,
+)
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult
@@ -217,9 +221,16 @@ class DeepResearchAgent(ToolUsingAgent):
self._emit_turn_start(input)
# Build system prompt with current date/time injected
- messages = self._build_messages(
- input, context, system_prompt=_build_system_prompt()
+ system_prompt = (
+ load_system_prompt_override("deep_research") or _build_system_prompt()
)
+ messages = self._build_messages(input, context, system_prompt=system_prompt)
+
+ # Inject few-shot exemplars before the user input
+ for ex in load_few_shot_exemplars("deep_research"):
+ if ex.get("input") and ex.get("output"):
+ messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
+ messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))
# Prepare OpenAI-format tool definitions for native function calling
tools_openai = [t.to_openai_function() for t in self._tools]
diff --git a/src/openjarvis/agents/monitor_operative.py b/src/openjarvis/agents/monitor_operative.py
index 3cd85d9d..a26116bb 100644
--- a/src/openjarvis/agents/monitor_operative.py
+++ b/src/openjarvis/agents/monitor_operative.py
@@ -16,9 +16,14 @@ from __future__ import annotations
import json
import logging
+import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
+from openjarvis.agents.prompt_loader import (
+ load_few_shot_exemplars,
+ load_system_prompt_override,
+)
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult, _message_to_dict
@@ -48,6 +53,15 @@ You are a Monitor Operative Agent designed for long-horizon tasks.
2. STATE: Your previous findings and state are automatically restored
3. MEMORY: Store important findings for future recall
+## How to use tools
+
+To call a tool, write on its own lines:
+
+Action:
+Action Input:
+
+You will receive the result, then continue your response.
+
## Strategy
- Memory extraction: {memory_extraction}
- Observation compression: {observation_compression}
@@ -169,14 +183,19 @@ class MonitorOperativeAgent(ToolUsingAgent):
self._emit_turn_start(input)
# 1. Build system prompt with state context
+ # Priority: constructor arg > file override > hardcoded default
sys_parts: list[str] = []
if self._system_prompt:
sys_parts.append(self._system_prompt)
else:
tool_desc = self._build_tool_descriptions()
+ prompt_template = (
+ load_system_prompt_override("monitor_operative")
+ or MONITOR_OPERATIVE_SYSTEM_PROMPT
+ )
try:
sys_parts.append(
- MONITOR_OPERATIVE_SYSTEM_PROMPT.format(
+ prompt_template.format(
memory_extraction=self._memory_extraction,
observation_compression=self._observation_compression,
retrieval_strategy=self._retrieval_strategy,
@@ -185,7 +204,7 @@ class MonitorOperativeAgent(ToolUsingAgent):
),
)
except KeyError:
- sys_parts.append(MONITOR_OPERATIVE_SYSTEM_PROMPT)
+ sys_parts.append(prompt_template)
# 2. State recall from memory backend
previous_state = self._recall_state()
@@ -205,6 +224,12 @@ class MonitorOperativeAgent(ToolUsingAgent):
session_messages=session_messages,
)
+ # 4b. Inject few-shot exemplars before the user input
+ for ex in load_few_shot_exemplars("monitor_operative"):
+ if ex.get("input") and ex.get("output"):
+ messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
+ messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))
+
# 5. Run function-calling tool loop
openai_tools = self._executor.get_openai_tools() if self._tools else []
all_tool_results: list[ToolResult] = []
@@ -229,34 +254,59 @@ class MonitorOperativeAgent(ToolUsingAgent):
for k in total_usage:
total_usage[k] += usage.get(k, 0)
content = result.get("content", "")
+ # Strip think tags so they don't interfere with parsing
+ content = self._strip_think_tags(content)
raw_tool_calls = result.get("tool_calls", [])
- # No tool calls -> check continuation, then final answer
- if not raw_tool_calls:
+ # --- Native function-calling path ---
+ if raw_tool_calls:
+ tool_calls = [
+ ToolCall(
+ id=tc.get("id", f"call_{i}"),
+ name=tc.get("name", ""),
+ arguments=tc.get("arguments", "{}"),
+ )
+ for i, tc in enumerate(raw_tool_calls)
+ ]
+ messages.append(
+ Message(
+ role=Role.ASSISTANT,
+ content=content,
+ tool_calls=tool_calls,
+ )
+ )
+ else:
+ # --- Text-based fallback ---
+ tool_info = self._extract_tool_call(content)
+ if tool_info:
+ action, action_input = tool_info
+ messages.append(Message(role=Role.ASSISTANT, content=content))
+ tc = ToolCall(
+ id=f"text_call_{turns}",
+ name=action,
+ arguments=action_input,
+ )
+ tool_result = self._executor.execute(tc)
+ all_tool_results.append(tool_result)
+ observation_content = self._compress_observation(
+ tool_result.content
+ )
+ messages.append(
+ Message(
+ role=Role.USER,
+ content=f"Result: {observation_content}",
+ )
+ )
+ self._extract_and_store(tc.name, tool_result.content)
+ continue
+
+ # No tool calls at all -> check continuation, then final answer
content = self._check_continuation(result, messages)
break
- # Build ToolCall objects from raw dicts
- tool_calls = [
- ToolCall(
- id=tc.get("id", f"call_{i}"),
- name=tc.get("name", ""),
- arguments=tc.get("arguments", "{}"),
- )
- for i, tc in enumerate(raw_tool_calls)
- ]
-
- # Append assistant message with tool calls
- messages.append(
- Message(
- role=Role.ASSISTANT,
- content=content,
- tool_calls=tool_calls,
- )
- )
-
- # Execute each tool
- for tc in tool_calls:
+ # Execute each native tool call
+ tool_calls_to_exec = tool_calls
+ for tc in tool_calls_to_exec:
# Loop guard check
if self._loop_guard:
verdict = self._loop_guard.check_call(
@@ -338,6 +388,89 @@ class MonitorOperativeAgent(ToolUsingAgent):
metadata={**total_usage, "messages": msg_dicts},
)
+ # ------------------------------------------------------------------
+ # Text-based tool call extraction (fallback for non-function-calling models)
+ # ------------------------------------------------------------------
+
+ @staticmethod
+ def _extract_tool_call(text: str) -> tuple[str, str] | None:
+ """Extract tool call from text output.
+
+ Supports three formats:
+ 1. Action: tool_name / Action Input: {"key": "value"}
+ 2. tool_name\\n$key=value (XML-style)
+ 3. or ... (inline XML)
+ """
+ # Format 1: Action / Action Input
+ action_match = re.search(r"Action:\s*(.+)", text, re.IGNORECASE)
+ input_match = re.search(
+ r"Action Input:\s*(.+?)(?=\n\n|\Z)", text, re.DOTALL | re.IGNORECASE
+ )
+ if action_match:
+ return (
+ action_match.group(1).strip(),
+ input_match.group(1).strip() if input_match else "{}",
+ )
+
+ # Format 2: tool_name ...
+ xml_match = re.search(
+ r"\s*(\w+)\s*(.*?)",
+ text,
+ re.DOTALL,
+ )
+ if xml_match:
+ tool_name = xml_match.group(1).strip()
+ raw_params = xml_match.group(2).strip()
+ params: dict[str, Any] = {}
+ for m in re.finditer(
+ r"\$(\w+)=(.+?)(?=\$|\n<|\n")
+ for m in re.finditer(r"<(\w+)>(.*?)", raw_params, re.DOTALL):
+ key, val = m.group(1), m.group(2).strip()
+ try:
+ params[key] = int(val)
+ except ValueError:
+ params[key] = val
+ if not params:
+ for m in re.finditer(
+ r"(\w+)\s*:\s*(.+?)(?=\n\w+\s*:|$)", raw_params, re.DOTALL
+ ):
+ key, val = m.group(1), m.group(2).strip().strip("\"'")
+ try:
+ params[key] = int(val)
+ except ValueError:
+ params[key] = val
+ if params:
+ return (tool_name, json.dumps(params))
+ return (tool_name, "{}")
+
+ # Format 3: or args
+ # Handles Qwen-style XML tool output like
+ inline_match = re.search(
+ r"<(\w+)\s+(.*?)/?>",
+ text,
+ re.DOTALL,
+ )
+ if inline_match:
+ tool_name = inline_match.group(1).strip()
+ # Skip common non-tool tags
+ if tool_name.lower() in ("think", "br", "hr", "p", "div", "span", "b", "i"):
+ return None
+ attr_str = inline_match.group(2).strip()
+ params = {}
+ for m in re.finditer(r'(\w+)=["\']([^"\']*)["\']', attr_str):
+ params[m.group(1)] = m.group(2)
+ # Also handle unquoted:
+ if not params:
+ for m in re.finditer(r"(\w+)=(\S+)", attr_str):
+ params[m.group(1)] = m.group(2).rstrip(">")
+ if params:
+ return (tool_name, json.dumps(params))
+ return (tool_name, "{}")
+
+ return None
+
# ------------------------------------------------------------------
# Message building
# ------------------------------------------------------------------
diff --git a/src/openjarvis/agents/native_openhands.py b/src/openjarvis/agents/native_openhands.py
index e0eae061..052596de 100644
--- a/src/openjarvis/agents/native_openhands.py
+++ b/src/openjarvis/agents/native_openhands.py
@@ -12,6 +12,10 @@ import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
+from openjarvis.agents.prompt_loader import (
+ load_few_shot_exemplars,
+ load_system_prompt_override,
+)
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult
@@ -220,7 +224,10 @@ class NativeOpenHandsAgent(ToolUsingAgent):
self._emit_turn_start(input)
tool_descriptions = build_tool_descriptions(self._tools)
- system_prompt = OPENHANDS_SYSTEM_PROMPT.format(
+ prompt_template = (
+ load_system_prompt_override("native_openhands") or OPENHANDS_SYSTEM_PROMPT
+ )
+ system_prompt = prompt_template.format(
tool_descriptions=tool_descriptions,
)
@@ -276,6 +283,13 @@ class NativeOpenHandsAgent(ToolUsingAgent):
)
messages = self._build_messages(input, context, system_prompt=system_prompt)
+
+ # Inject few-shot exemplars before the user input
+ for ex in load_few_shot_exemplars("native_openhands"):
+ if ex.get("input") and ex.get("output"):
+ messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
+ messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))
+
messages = self._truncate_if_needed(messages)
all_tool_results: list[ToolResult] = []
diff --git a/src/openjarvis/agents/native_react.py b/src/openjarvis/agents/native_react.py
index 89061505..2d55b2ad 100644
--- a/src/openjarvis/agents/native_react.py
+++ b/src/openjarvis/agents/native_react.py
@@ -10,6 +10,10 @@ import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
+from openjarvis.agents.prompt_loader import (
+ load_few_shot_exemplars,
+ load_system_prompt_override,
+)
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role, ToolCall, ToolResult, _message_to_dict
@@ -131,10 +135,8 @@ class NativeReActAgent(ToolUsingAgent):
# Build system prompt with rich tool descriptions
tool_desc = build_tool_descriptions(self._tools)
- # Plan 2B I3: render the optimized few-shot examples as a section
- # that the model sees BEFORE the tool descriptions. When no
- # examples are present, this is an empty string and the rendered
- # prompt is unchanged.
+ # Plan 2B I3: render optimized few-shot skill examples as a section
+ # before the tool descriptions. Empty string when not present.
if self._skill_few_shot_examples:
skill_examples_block = (
"## Skill Examples\n\n"
@@ -143,13 +145,29 @@ class NativeReActAgent(ToolUsingAgent):
)
else:
skill_examples_block = ""
- system_prompt = REACT_SYSTEM_PROMPT.format(
- tool_descriptions=tool_desc,
- skill_examples=skill_examples_block,
+ # Respect $OPENJARVIS_HOME override for the base template (M2+ work).
+ prompt_template = (
+ load_system_prompt_override("native_react") or REACT_SYSTEM_PROMPT
)
+ # External overrides may not include the {skill_examples} slot.
+ try:
+ system_prompt = prompt_template.format(
+ tool_descriptions=tool_desc,
+ skill_examples=skill_examples_block,
+ )
+ except KeyError:
+ system_prompt = prompt_template.format(tool_descriptions=tool_desc)
+ if skill_examples_block:
+ system_prompt = system_prompt + "\n\n" + skill_examples_block
messages = self._build_messages(input, context, system_prompt=system_prompt)
+ # Inject few-shot exemplars before the user input
+ for ex in load_few_shot_exemplars("native_react"):
+ if ex.get("input") and ex.get("output"):
+ messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
+ messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))
+
all_tool_results: list[ToolResult] = []
turns = 0
total_usage: dict[str, int] = {
diff --git a/src/openjarvis/agents/prompt_loader.py b/src/openjarvis/agents/prompt_loader.py
new file mode 100644
index 00000000..429d8b17
--- /dev/null
+++ b/src/openjarvis/agents/prompt_loader.py
@@ -0,0 +1,83 @@
+"""Load system prompt and few-shot overrides from $OPENJARVIS_HOME.
+
+Distillation (M1) proposes edits that get written to disk by appliers.
+This module lets agents pick those overrides up at runtime:
+
+- System prompts: ``$OPENJARVIS_HOME/agents/{name}/system_prompt.md``
+- Few-shot exemplars: ``$OPENJARVIS_HOME/agents/{name}/few_shot.json``
+
+Override files are templates — they may contain ``{tool_descriptions}`` and
+other format placeholders that the agent fills in via ``.format()``, exactly
+like the hardcoded constants.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _openjarvis_home() -> Path:
+ """Resolve $OPENJARVIS_HOME, defaulting to ~/.openjarvis."""
+ return Path(os.environ.get("OPENJARVIS_HOME", "~/.openjarvis")).expanduser()
+
+
+def load_system_prompt_override(agent_name: str) -> str | None:
+ """Return the override prompt for *agent_name*, or ``None``.
+
+ Looks for ``$OPENJARVIS_HOME/agents//system_prompt.md``.
+ ``OPENJARVIS_HOME`` defaults to ``~/.openjarvis`` when unset.
+ """
+ home = _openjarvis_home()
+ prompt_path = home / "agents" / agent_name / "system_prompt.md"
+ if not prompt_path.exists():
+ return None
+ try:
+ content = prompt_path.read_text(encoding="utf-8")
+ logger.info(
+ "Loaded system prompt override for %s from %s", agent_name, prompt_path
+ )
+ return content
+ except Exception:
+ logger.warning(
+ "Failed to read system prompt override at %s", prompt_path, exc_info=True
+ )
+ return None
+
+
+def load_few_shot_exemplars(
+ agent_name: str,
+) -> list[dict[str, Any]]:
+ """Return few-shot exemplars for *agent_name*, or empty list.
+
+ Looks for ``$OPENJARVIS_HOME/agents//few_shot.json``.
+ Expected format: ``[{"input": "Q", "output": "A"}, ...]``.
+ """
+ home = _openjarvis_home()
+ fs_path = home / "agents" / agent_name / "few_shot.json"
+ if not fs_path.exists():
+ return []
+ try:
+ data = json.loads(fs_path.read_text(encoding="utf-8"))
+ if not isinstance(data, list):
+ logger.warning("few_shot.json for %s is not a list", agent_name)
+ return []
+ logger.info(
+ "Loaded %d few-shot exemplars for %s from %s",
+ len(data),
+ agent_name,
+ fs_path,
+ )
+ return data
+ except Exception:
+ logger.warning(
+ "Failed to read few-shot exemplars at %s",
+ fs_path,
+ exc_info=True,
+ )
+ return []
diff --git a/src/openjarvis/agents/rlm.py b/src/openjarvis/agents/rlm.py
index cd8c0c37..2b4fc2ad 100644
--- a/src/openjarvis/agents/rlm.py
+++ b/src/openjarvis/agents/rlm.py
@@ -12,6 +12,10 @@ import re
from typing import Any, List, Optional
from openjarvis.agents._stubs import AgentContext, AgentResult, ToolUsingAgent
+from openjarvis.agents.prompt_loader import (
+ load_few_shot_exemplars,
+ load_system_prompt_override,
+)
from openjarvis.agents.rlm_repl import RLMRepl
from openjarvis.core.events import EventBus
from openjarvis.core.registry import AgentRegistry
@@ -154,13 +158,13 @@ class RLMAgent(ToolUsingAgent):
if self._custom_system_prompt:
system_prompt = self._custom_system_prompt
else:
+ prompt_template = load_system_prompt_override("rlm") or RLM_SYSTEM_PROMPT
try:
- system_prompt = RLM_SYSTEM_PROMPT.format(
+ system_prompt = prompt_template.format(
tool_section=tool_section,
)
except KeyError:
- # Custom system_prompt override without {tool_section}
- system_prompt = RLM_SYSTEM_PROMPT
+ system_prompt = prompt_template
# Create REPL with sub-LM callbacks
repl = RLMRepl(
@@ -181,6 +185,12 @@ class RLMAgent(ToolUsingAgent):
system_prompt=system_prompt,
)
+ # Inject few-shot exemplars before the user input
+ for ex in load_few_shot_exemplars("rlm"):
+ if ex.get("input") and ex.get("output"):
+ messages.insert(-1, Message(role=Role.USER, content=ex["input"]))
+ messages.insert(-1, Message(role=Role.ASSISTANT, content=ex["output"]))
+
all_tool_results: list[ToolResult] = []
turns = 0
total_usage: dict[str, int] = {
diff --git a/src/openjarvis/cli/__init__.py b/src/openjarvis/cli/__init__.py
index 0b460a90..00ed64ce 100644
--- a/src/openjarvis/cli/__init__.py
+++ b/src/openjarvis/cli/__init__.py
@@ -38,6 +38,7 @@ from openjarvis.cli.telemetry_cmd import telemetry
from openjarvis.cli.tool_cmd import tool
from openjarvis.cli.vault_cmd import vault
from openjarvis.cli.workflow_cmd import workflow
+from openjarvis.learning.distillation.cli import learning_group
@click.group(help="OpenJarvis — modular AI assistant backend")
@@ -98,6 +99,7 @@ cli.add_command(connect, "connect")
cli.add_command(digest, "digest")
cli.add_command(deep_research_setup, "deep-research-setup")
cli.add_command(deep_research_setup, "research")
+cli.add_command(learning_group, "learning")
# Gateway CLI commands (lazy import to avoid pulling starlette)
try:
diff --git a/src/openjarvis/engine/_openai_compat.py b/src/openjarvis/engine/_openai_compat.py
index b9751452..4237422a 100644
--- a/src/openjarvis/engine/_openai_compat.py
+++ b/src/openjarvis/engine/_openai_compat.py
@@ -29,7 +29,10 @@ class _OpenAICompatibleEngine(InferenceEngine):
_api_prefix: str = "/v1"
def __init__(self, host: str | None = None, *, timeout: float = 600.0) -> None:
- self._host = (host or self._default_host).rstrip("/")
+ import os
+
+ env_key = f"{self.engine_id.upper()}_HOST"
+ self._host = (host or os.environ.get(env_key) or self._default_host).rstrip("/")
self._client = httpx.Client(base_url=self._host, timeout=timeout)
# -- InferenceEngine interface ------------------------------------------
diff --git a/src/openjarvis/evals/cli.py b/src/openjarvis/evals/cli.py
index ccdbcea9..8294fb6f 100644
--- a/src/openjarvis/evals/cli.py
+++ b/src/openjarvis/evals/cli.py
@@ -349,10 +349,10 @@ def _build_dataset(benchmark: str, subset: str | None = None):
return LiveResearchBenchDataset(path=subset)
elif benchmark == "liveresearchbench":
from openjarvis.evals.datasets.liveresearchbench import (
- LiveResearchBenchSFDataset,
+ LiveResearchBenchDataset as LRBDataset,
)
- return LiveResearchBenchSFDataset()
+ return LRBDataset()
elif benchmark == "toolcall15":
from openjarvis.evals.datasets.toolcall15 import ToolCall15Dataset
@@ -507,10 +507,10 @@ def _build_scorer(benchmark: str, judge_backend, judge_model: str):
return LiveResearchBenchScorer(judge_backend, judge_model)
elif benchmark == "liveresearchbench":
from openjarvis.evals.scorers.liveresearchbench import (
- LiveResearchBenchSFScorer,
+ LiveResearchBenchScorer as LRBScorer,
)
- return LiveResearchBenchSFScorer(judge_backend, judge_model)
+ return LRBScorer(judge_backend, judge_model)
elif benchmark == "toolcall15":
from openjarvis.evals.scorers.toolcall15 import ToolCall15Scorer
diff --git a/src/openjarvis/evals/configs/livecodebench-qwen-9b.toml b/src/openjarvis/evals/configs/livecodebench-qwen-9b.toml
new file mode 100644
index 00000000..2cef93d6
--- /dev/null
+++ b/src/openjarvis/evals/configs/livecodebench-qwen-9b.toml
@@ -0,0 +1,30 @@
+# LiveCodeBench eval: Qwen3.5-9B (vLLM, 1 GPU)
+# Competitive programming from LeetCode, AtCoder, CodeForces
+
+[meta]
+name = "livecodebench-qwen-9b"
+description = "LiveCodeBench on Qwen/Qwen3.5-9B (vLLM, 1 GPU)"
+
+[defaults]
+temperature = 0.0
+max_tokens = 4096
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-9b/livecodebench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-9B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "livecodebench"
+backend = "jarvis-direct"
+max_samples = 20
diff --git a/src/openjarvis/evals/configs/liveresearchbench-qwen-27b.toml b/src/openjarvis/evals/configs/liveresearchbench-qwen-27b.toml
new file mode 100644
index 00000000..be8c8e05
--- /dev/null
+++ b/src/openjarvis/evals/configs/liveresearchbench-qwen-27b.toml
@@ -0,0 +1,31 @@
+# LiveResearchBench: Qwen3.5-27B-FP8 (vLLM, 1 GPU)
+# Salesforce's checklist-based deep research benchmark.
+
+[meta]
+name = "liveresearchbench-qwen-27b"
+description = "LiveResearchBench (Salesforce) on Qwen/Qwen3.5-27B-FP8"
+
+[defaults]
+temperature = 0.6
+max_tokens = 16384
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+max_tokens = 4096
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-27b/liveresearchbench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-27B-FP8"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "liveresearchbench"
+backend = "jarvis-direct"
+max_samples = 50
diff --git a/src/openjarvis/evals/configs/liveresearchbench-qwen-2b.toml b/src/openjarvis/evals/configs/liveresearchbench-qwen-2b.toml
new file mode 100644
index 00000000..f4d1df84
--- /dev/null
+++ b/src/openjarvis/evals/configs/liveresearchbench-qwen-2b.toml
@@ -0,0 +1,31 @@
+# LiveResearchBench: Qwen3.5-2B (vLLM, 1 GPU)
+# Salesforce's checklist-based deep research benchmark.
+
+[meta]
+name = "liveresearchbench-qwen-2b"
+description = "LiveResearchBench (Salesforce) on Qwen/Qwen3.5-2B"
+
+[defaults]
+temperature = 0.6
+max_tokens = 16384
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+max_tokens = 4096
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-2b/liveresearchbench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-2B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "liveresearchbench"
+backend = "jarvis-direct"
+max_samples = 50
diff --git a/src/openjarvis/evals/configs/liveresearchbench-qwen-9b.toml b/src/openjarvis/evals/configs/liveresearchbench-qwen-9b.toml
new file mode 100644
index 00000000..818d257b
--- /dev/null
+++ b/src/openjarvis/evals/configs/liveresearchbench-qwen-9b.toml
@@ -0,0 +1,31 @@
+# LiveResearchBench: Qwen3.5-9B (vLLM, 1 GPU)
+# Salesforce's checklist-based deep research benchmark.
+
+[meta]
+name = "liveresearchbench-qwen-9b"
+description = "LiveResearchBench (Salesforce) on Qwen/Qwen3.5-9B"
+
+[defaults]
+temperature = 0.6
+max_tokens = 16384
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+max_tokens = 4096
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-9b/liveresearchbench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-9B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "liveresearchbench"
+backend = "jarvis-direct"
+max_samples = 50
diff --git a/src/openjarvis/evals/configs/pinchbench-qwen-9b.toml b/src/openjarvis/evals/configs/pinchbench-qwen-9b.toml
new file mode 100644
index 00000000..5097ac4f
--- /dev/null
+++ b/src/openjarvis/evals/configs/pinchbench-qwen-9b.toml
@@ -0,0 +1,35 @@
+# PinchBench eval: Qwen3.5-9B (vLLM, 1 GPU)
+# Agent: native_openhands — all PinchBench-required tools enabled
+
+[meta]
+name = "pinchbench-qwen-9b"
+description = "PinchBench on Qwen/Qwen3.5-9B (vLLM, 1 GPU)"
+
+[defaults]
+temperature = 0.6
+max_tokens = 8192
+
+[judge]
+model = "claude-opus-4-5"
+temperature = 0.0
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-9b/pinchbench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-9B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "pinchbench"
+backend = "jarvis-agent"
+agent = "native_openhands"
+tools = [
+ "think", "file_read", "file_write", "web_search", "shell_exec",
+ "code_interpreter", "browser_navigate", "image_generate",
+ "calculator", "http_request", "pdf_extract",
+]
diff --git a/src/openjarvis/evals/configs/taubench-qwen-9b.toml b/src/openjarvis/evals/configs/taubench-qwen-9b.toml
new file mode 100644
index 00000000..d076895f
--- /dev/null
+++ b/src/openjarvis/evals/configs/taubench-qwen-9b.toml
@@ -0,0 +1,31 @@
+# TauBench V2 eval: Qwen3.5-9B (vLLM, 1 GPU)
+# Multi-turn customer service benchmark — airline + retail splits
+
+[meta]
+name = "taubench-qwen-9b"
+description = "TauBench V2 on Qwen/Qwen3.5-9B (vLLM, 1 GPU)"
+
+[defaults]
+temperature = 0.7
+max_tokens = 4096
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-9b/taubench/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-9B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "taubench"
+backend = "jarvis-direct"
+max_samples = 20
+split = "airline,retail"
diff --git a/src/openjarvis/evals/configs/toolcall15-qwen-9b.toml b/src/openjarvis/evals/configs/toolcall15-qwen-9b.toml
new file mode 100644
index 00000000..f30fa2b5
--- /dev/null
+++ b/src/openjarvis/evals/configs/toolcall15-qwen-9b.toml
@@ -0,0 +1,29 @@
+# ToolCall-15 eval: Qwen3.5-9B (vLLM, 1 GPU)
+# Lightweight tool calling benchmark — 15 scenarios, 5 categories
+
+[meta]
+name = "toolcall15-qwen-9b"
+description = "ToolCall-15 on Qwen/Qwen3.5-9B (temperature=0)"
+
+[defaults]
+temperature = 0.0
+max_tokens = 4096
+
+[judge]
+model = "gpt-5-mini-2025-08-07"
+temperature = 0.0
+engine = "cloud"
+
+[run]
+max_workers = 1
+output_dir = "results/neurips-2026/baselines/qwen-9b/toolcall15/"
+seed = 42
+
+[[models]]
+name = "Qwen/Qwen3.5-9B"
+engine = "vllm"
+num_gpus = 1
+
+[[benchmarks]]
+name = "toolcall15"
+backend = "jarvis-direct"
diff --git a/src/openjarvis/evals/core/config.py b/src/openjarvis/evals/core/config.py
index 5a46ad2f..cae66c26 100644
--- a/src/openjarvis/evals/core/config.py
+++ b/src/openjarvis/evals/core/config.py
@@ -64,6 +64,7 @@ KNOWN_BENCHMARKS = {
"taubench",
"livecodebench",
"liveresearch",
+ "liveresearchbench",
"toolcall15",
}
diff --git a/src/openjarvis/evals/datasets/liveresearchbench.py b/src/openjarvis/evals/datasets/liveresearchbench.py
index bb4ae6eb..599de3f7 100644
--- a/src/openjarvis/evals/datasets/liveresearchbench.py
+++ b/src/openjarvis/evals/datasets/liveresearchbench.py
@@ -1,20 +1,23 @@
-"""LiveResearchBench (Salesforce) dataset provider.
+"""LiveResearchBench dataset provider — Salesforce's checklist-based benchmark.
-80 expert-curated deep research tasks with per-question evaluation
-checklists across three domains: daily life, enterprise, and academia.
-543 checklist items total (grouped by question).
+Loads Salesforce/LiveResearchBench from HuggingFace. Each task has a research
+question and a set of checklist items used for fine-grained, coverage-based
+evaluation.
+
+Note: This is the actual LiveResearchBench by Salesforce (arxiv 2510.14240).
+The existing ``liveresearch`` module points at DeepResearchBench
+(Ayanami0730/deep_research_bench) despite its misleading class name.
Reference: https://github.com/SalesforceAIResearch/LiveResearchBench
-HuggingFace: Salesforce/LiveResearchBench (gated — accept terms first)
+Paper: https://arxiv.org/abs/2510.14240
+Dataset: https://huggingface.co/datasets/Salesforce/LiveResearchBench
"""
from __future__ import annotations
+import json
import logging
import random
-import re
-from collections import defaultdict
-from datetime import datetime
from typing import Any, Dict, Iterable, List, Optional
from openjarvis.evals.core.dataset import DatasetProvider
@@ -22,34 +25,59 @@ from openjarvis.evals.core.types import EvalRecord
LOGGER = logging.getLogger(__name__)
-_HF_DATASET = "Salesforce/LiveResearchBench"
+HF_DATASET_ID = "Salesforce/LiveResearchBench"
+DEFAULT_HF_CONFIG = "question_with_checklist"
+DEFAULT_HF_SPLIT = "test"
-def _replace_date_placeholders(text: str) -> str:
- """Replace dynamic date placeholders in queries."""
- now = datetime.now()
- text = text.replace("{{current_year}}", str(now.year))
- text = text.replace("{{last_year}}", str(now.year - 1))
- text = text.replace("{{current_date}}", now.strftime("%Y-%m-%d"))
- text = text.replace("{{date}}", now.strftime("%Y-%m-%d"))
- text = re.sub(r"\{current_year\}", str(now.year), text)
- text = re.sub(r"\{last_year\}", str(now.year - 1), text)
- return text
+def _parse_checklist(checklist: Any) -> List[str]:
+ """Parse the checklist field — may be a JSON string, list, or newline text."""
+ if isinstance(checklist, list):
+ return [str(item).strip() for item in checklist if str(item).strip()]
+ if not isinstance(checklist, str) or not checklist.strip():
+ return []
+ try:
+ parsed = json.loads(checklist)
+ if isinstance(parsed, list):
+ return [str(item).strip() for item in parsed if str(item).strip()]
+ if isinstance(parsed, str):
+ return [parsed.strip()] if parsed.strip() else []
+ except json.JSONDecodeError:
+ pass
+ # Fall back: treat as newline-separated items
+ return [line.strip() for line in checklist.splitlines() if line.strip()]
-class LiveResearchBenchSFDataset(DatasetProvider):
- """Salesforce LiveResearchBench — 80 expert-curated research tasks.
+class LiveResearchBenchDataset(DatasetProvider):
+ """LiveResearchBench — Salesforce's expert-curated deep research benchmark.
- The HuggingFace dataset has 543 rows (multiple checklist items per
- question). We group by ``qid`` to produce one EvalRecord per unique
- question, with all checklist items aggregated in metadata.
+ Loads tasks from HuggingFace with per-task checklists used for
+ coverage-based evaluation. Tasks span 7 domains (Science/Tech, Business,
+ Health, Law/Governance, Society/Culture, Education, Media).
"""
dataset_id = "liveresearchbench"
- dataset_name = "LiveResearchBench (Salesforce)"
+ dataset_name = "LiveResearchBench"
- def __init__(self) -> None:
- self._records: Optional[List[EvalRecord]] = None
+ def __init__(
+ self,
+ hf_config: Optional[str] = None,
+ hf_split: Optional[str] = None,
+ ) -> None:
+ self._hf_config = hf_config or DEFAULT_HF_CONFIG
+ self._hf_split = hf_split or DEFAULT_HF_SPLIT
+ self._records: List[EvalRecord] = []
+
+ def verify_requirements(self) -> List[str]:
+ issues: List[str] = []
+ try:
+ import datasets # noqa: F401
+ except ImportError:
+ issues.append(
+ "The 'datasets' package is required for LiveResearchBench. "
+ "Install with: pip install datasets"
+ )
+ return issues
def load(
self,
@@ -58,122 +86,95 @@ class LiveResearchBenchSFDataset(DatasetProvider):
split: Optional[str] = None,
seed: Optional[int] = None,
) -> None:
- try:
- from datasets import load_dataset
- except ImportError:
- raise ImportError(
- "datasets package required. Install with: pip install datasets"
- )
+ from datasets import load_dataset
- import os
-
- hf_token = os.environ.get("HF_TOKEN") or os.environ.get(
- "HUGGING_FACE_HUB_TOKEN"
- )
-
- # Try question_with_checklist first (has evaluation criteria)
- hf_config = split or "question_with_checklist"
+ hf_split = split or self._hf_split
LOGGER.info(
- "Loading LiveResearchBench from HuggingFace (%s, config=%s)",
- _HF_DATASET,
- hf_config,
+ "Loading %s (config=%s, split=%s) from HuggingFace ...",
+ HF_DATASET_ID,
+ self._hf_config,
+ hf_split,
)
+ ds = load_dataset(HF_DATASET_ID, self._hf_config, split=hf_split)
- try:
- ds = load_dataset(
- _HF_DATASET, hf_config, split="test", token=hf_token
- )
- except Exception as exc:
- raise RuntimeError(
- f"Failed to load {_HF_DATASET}. This is a gated dataset — "
- "visit https://huggingface.co/datasets/Salesforce/LiveResearchBench "
- "to accept the terms, then set HF_TOKEN in your environment. "
- f"Error: {exc}"
- ) from exc
-
- # Group rows by qid (multiple checklist items per question)
- questions: Dict[str, Dict[str, Any]] = {}
- checklists_by_qid: Dict[str, List[str]] = defaultdict(list)
-
+ # `question_with_checklist` has multiple rows per qid (one per
+ # checklist_id). Group by qid and fold all checklist items into a
+ # single record per question.
+ grouped: Dict[str, Dict[str, Any]] = {}
for row in ds:
- qid = str(row.get("qid", ""))
+ qid = str(row.get("qid", "") or "").strip()
if not qid:
continue
-
- if qid not in questions:
- question = row.get("question", "") or row.get(
- "question_no_placeholder", ""
- )
- questions[qid] = {
+ question = row.get("question") or row.get("question_no_placeholder") or ""
+ if qid not in grouped:
+ grouped[qid] = {
+ "qid": qid,
"question": question,
- "category": row.get("category", ""),
+ "category": row.get("category", "") or "",
+ "checklist": [],
}
-
- checklist = row.get("checklist", "") or row.get(
- "checklist_no_placeholder", ""
+ cl_items = _parse_checklist(
+ row.get("checklist") or row.get("checklist_no_placeholder")
)
- if checklist:
- checklists_by_qid[qid].append(checklist)
+ grouped[qid]["checklist"].extend(cl_items)
- # Build EvalRecords
- records: List[EvalRecord] = []
- for qid, info in questions.items():
- question = info["question"]
- if not question:
- continue
-
- question = _replace_date_placeholders(question)
-
- problem = (
- "You are a research assistant. Please conduct thorough "
- "research on the following question and write a "
- "comprehensive report with citations.\n\n"
- f"{question}"
- )
-
- metadata: Dict[str, Any] = {
- "qid": qid,
- "original_question": question,
- "category": info.get("category", ""),
- "checklists": checklists_by_qid.get(qid, []),
- }
-
- records.append(
- EvalRecord(
- record_id=f"lrb-{qid}",
- problem=problem,
- reference="",
- category="liveresearchbench",
- metadata=metadata,
- )
+ records = list(grouped.values())
+ if not records:
+ raise RuntimeError(
+ f"LiveResearchBench: no records found in {HF_DATASET_ID} "
+ f"(config={self._hf_config}, split={hf_split})"
)
if seed is not None:
- rng = random.Random(seed)
- rng.shuffle(records)
-
+ random.Random(seed).shuffle(records)
if max_samples is not None:
records = records[:max_samples]
- self._records = records
- total_checklists = sum(
- len(r.metadata.get("checklists", [])) for r in records
- )
+ self._records = []
+ for rec in records:
+ question = (rec.get("question") or "").strip()
+ if not question:
+ LOGGER.warning("Skipping %s: empty question", rec.get("qid"))
+ continue
+
+ research_prompt = (
+ "You are a deep research assistant. Conduct thorough research "
+ "on the following task and produce a comprehensive, "
+ "well-structured, citation-grounded report that addresses "
+ "every aspect of the request.\n\n"
+ f"## Research Task\n\n{question}"
+ )
+
+ self._records.append(
+ EvalRecord(
+ record_id=f"liveresearchbench-{rec['qid']}",
+ problem=research_prompt,
+ reference="", # checklist-based; no single reference answer
+ category="liveresearchbench",
+ subject=rec.get("category") or "",
+ metadata={
+ "qid": rec["qid"],
+ "question": question,
+ "checklist": rec["checklist"],
+ "hf_category": rec.get("category", ""),
+ },
+ )
+ )
+
LOGGER.info(
- "LiveResearchBench: loaded %d tasks (%d checklist items)",
+ "LiveResearchBench: loaded %d tasks (avg %.1f checklist items/task)",
len(self._records),
- total_checklists,
+ (
+ sum(len(r.metadata.get("checklist", [])) for r in self._records)
+ / max(1, len(self._records))
+ ),
)
def iter_records(self) -> Iterable[EvalRecord]:
- if self._records is None:
- raise RuntimeError("Call .load() before iterating")
return iter(self._records)
def size(self) -> int:
- if self._records is None:
- raise RuntimeError("Call .load() before size()")
return len(self._records)
-__all__ = ["LiveResearchBenchSFDataset"]
+__all__ = ["LiveResearchBenchDataset"]
diff --git a/src/openjarvis/evals/execution/taubench_env.py b/src/openjarvis/evals/execution/taubench_env.py
index 6c4425c6..320e17b2 100644
--- a/src/openjarvis/evals/execution/taubench_env.py
+++ b/src/openjarvis/evals/execution/taubench_env.py
@@ -78,6 +78,15 @@ def _tau2_to_oj_messages(
return oj_msgs
+def _strip_think_tags(text: str) -> str:
+ """Remove ``...`` blocks from model output."""
+ import re
+
+ text = re.sub(r".*?\s*", "", text, flags=re.DOTALL | re.IGNORECASE)
+ text = re.sub(r"^.*?\s*", "", text, flags=re.DOTALL | re.IGNORECASE)
+ return text.strip()
+
+
def _oj_result_to_tau2_msg(result: dict):
"""Convert OpenJarvis engine.generate() result to a tau2 AssistantMessage."""
from tau2.data_model.message import AssistantMessage, ToolCall
@@ -99,9 +108,12 @@ def _oj_result_to_tau2_msg(result: dict):
for tc in raw_tool_calls
]
+ content = result.get("content") or ""
+ content = _strip_think_tags(content) if content else None
+
return AssistantMessage(
role="assistant",
- content=result.get("content") or None,
+ content=content or None,
tool_calls=tool_calls,
cost=result.get("cost_usd", 0.0),
)
@@ -183,11 +195,17 @@ class JarvisHalfDuplexAgent:
"tool_choice": "auto",
}
# Disable thinking mode for local models (Qwen3.5 etc.)
- # to avoid tags interfering with tool call parsing
+ # to avoid tags interfering with tool call parsing.
+ # vLLM >=0.8 accepts chat_template_kwargs as a top-level field.
+ # We also set it inside extra_body for compatibility with
+ # OpenAI SDK-based clients that only pass extra_body through.
if "qwen" in self._model.lower():
gen_kwargs["chat_template_kwargs"] = {
"enable_thinking": False,
}
+ gen_kwargs.setdefault("extra_body", {})["chat_template_kwargs"] = {
+ "enable_thinking": False,
+ }
result = self._engine.generate(oj_messages, **gen_kwargs)
# Convert result to tau2 AssistantMessage
diff --git a/src/openjarvis/evals/scorers/liveresearchbench.py b/src/openjarvis/evals/scorers/liveresearchbench.py
index ea89ab55..b122c748 100644
--- a/src/openjarvis/evals/scorers/liveresearchbench.py
+++ b/src/openjarvis/evals/scorers/liveresearchbench.py
@@ -1,70 +1,164 @@
-"""LiveResearchBench (Salesforce) scorer — checklist-based evaluation.
+"""LiveResearchBench scorer — checklist-based LLM-as-judge scoring.
-Evaluates research reports using per-question checklists for coverage,
-plus LLM-as-judge for presentation quality and citation adequacy.
+For each task, LiveResearchBench provides a list of checklist items that a
+good response must cover. The scorer asks an LLM judge to evaluate each
+checklist item against the response, producing a per-item pass/fail.
+Final score = fraction of passed items (coverage).
-Reference: https://github.com/SalesforceAIResearch/LiveResearchBench
+Reference: https://arxiv.org/abs/2510.14240
"""
from __future__ import annotations
+import json
import logging
import re
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
from openjarvis.evals.core.scorer import LLMJudgeScorer
from openjarvis.evals.core.types import EvalRecord
LOGGER = logging.getLogger(__name__)
-_COVERAGE_PROMPT = """\
-You are evaluating a research report against a checklist of required topics/points.
+# Tasks with coverage >= PASS_THRESHOLD are marked as correct
+PASS_THRESHOLD = 0.5
-**Research Question:**
+
+def _build_judge_prompt(
+ *,
+ question: str,
+ answer: str,
+ checklist: List[str],
+) -> str:
+ """Build a batched judge prompt that evaluates all checklist items at once."""
+ bullets = "\n".join(f"{i}. {item}" for i, item in enumerate(checklist, 1))
+ return f"""You are an expert evaluator scoring a deep-research report against a checklist.
+
+## Research Question
{question}
-**Report:**
-{report}
+## Checklist Items
+{bullets}
-**Checklist items to verify (each should be covered in the report):**
-{checklist_items}
+## Report to Evaluate
+{answer}
-For each checklist item, determine if the report adequately covers it.
-Respond with a JSON array of objects, one per checklist item:
-[
- {{"item": "", "covered": true/false, "evidence": ""}},
- ...
-]
+## Instructions
+For each checklist item above, decide whether the report adequately covers it.
+A checklist item is COVERED if the report includes the required information or
+analysis in a clear, accurate, and substantive way. A checklist item is NOT
+COVERED if the report omits it, addresses it only superficially, or contains
+incorrect information.
-Then on the last line, provide the overall coverage score:
-coverage_score: /"""
+Return your evaluation as JSON with this exact structure:
+```json
+{{
+ "judgments": [
+ {{"index": 1, "covered": true, "reason": "brief justification"}},
+ {{"index": 2, "covered": false, "reason": "brief justification"}}
+ ]
+}}
+```
-_QUALITY_PROMPT = """\
-You are evaluating the quality of a research report.
-
-**Research Question:**
-{question}
-
-**Report:**
-{report}
-
-Rate the report on these dimensions (each 1-5):
-
-1. **Presentation**: Is the report well-structured, readable, and professional?
-2. **Depth**: Does the report go beyond surface-level information?
-3. **Citation**: Does the report reference specific sources, data, or evidence?
-4. **Consistency**: Is the report internally consistent and free of contradictions?
-
-Respond in this exact format:
-presentation: <1-5>
-depth: <1-5>
-citation: <1-5>
-consistency: <1-5>
-reasoning: """
+Return one judgment per checklist item, indexed in the same order as above.
+Be rigorous: `covered: true` only if the report genuinely satisfies the item."""
-class LiveResearchBenchSFScorer(LLMJudgeScorer):
- """Checklist + quality scorer for Salesforce LiveResearchBench."""
+def _parse_judge_response(raw: str, num_items: int) -> List[Dict[str, Any]]:
+ """Extract per-item judgments from the judge's raw response.
+
+ Returns one dict per checklist item ({index, covered, reason}). Missing
+ items default to ``covered=False``.
+ """
+ if not raw or not raw.strip():
+ return [
+ {"index": i + 1, "covered": False, "reason": "empty judge response"}
+ for i in range(num_items)
+ ]
+
+ judgments: Dict[int, Dict[str, Any]] = {}
+
+ # Collect JSON candidates: code blocks first, then balanced braces.
+ candidates: List[str] = []
+ for match in re.finditer(r"```(?:json)?\s*(.*?)\s*```", raw, re.DOTALL):
+ candidates.append(match.group(1))
+
+ depth = 0
+ current: List[str] = []
+ for char in raw:
+ if char == "{":
+ if depth == 0:
+ current = []
+ depth += 1
+ if depth > 0:
+ current.append(char)
+ if char == "}":
+ depth -= 1
+ if depth == 0 and current:
+ candidates.append("".join(current))
+
+ for candidate in candidates:
+ try:
+ parsed = json.loads(candidate)
+ except json.JSONDecodeError:
+ continue
+ if not isinstance(parsed, dict):
+ continue
+ items = parsed.get("judgments") or parsed.get("items") or []
+ if not isinstance(items, list):
+ continue
+ for item in items:
+ if not isinstance(item, dict):
+ continue
+ try:
+ idx_int = int(item.get("index"))
+ except (ValueError, TypeError):
+ continue
+ covered_raw = item.get("covered")
+ if isinstance(covered_raw, bool):
+ covered = covered_raw
+ elif isinstance(covered_raw, str):
+ covered = covered_raw.strip().lower() in {
+ "true",
+ "yes",
+ "covered",
+ "1",
+ "y",
+ }
+ else:
+ covered = False
+ judgments[idx_int] = {
+ "index": idx_int,
+ "covered": covered,
+ "reason": str(item.get("reason", "")),
+ }
+ if judgments:
+ break
+
+ if not judgments:
+ LOGGER.warning("Failed to parse judge response — marking all items uncovered")
+
+ # Fill gaps with covered=False so downstream code always sees N items.
+ return [
+ judgments.get(
+ i + 1,
+ {
+ "index": i + 1,
+ "covered": False,
+ "reason": "missing from judge output",
+ },
+ )
+ for i in range(num_items)
+ ]
+
+
+class LiveResearchBenchScorer(LLMJudgeScorer):
+ """Checklist-based LLM-as-judge scorer for LiveResearchBench.
+
+ For each sample, the judge evaluates each checklist item against the
+ model's report. Score = fraction covered. Tasks with score >=
+ ``PASS_THRESHOLD`` (default 0.5) are marked correct.
+ """
scorer_id = "liveresearchbench"
@@ -73,94 +167,62 @@ class LiveResearchBenchSFScorer(LLMJudgeScorer):
record: EvalRecord,
model_answer: str,
) -> Tuple[Optional[bool], Dict[str, Any]]:
+ checklist: List[str] = list(record.metadata.get("checklist", []) or [])
+
if not model_answer or not model_answer.strip():
- return False, {"reason": "empty_response", "score": 0.0}
+ return False, {
+ "score": 0.0,
+ "coverage": 0.0,
+ "covered_count": 0,
+ "checklist_size": len(checklist),
+ "reason": "empty_response",
+ }
- question = record.metadata.get("original_question", record.problem)
- checklists = record.metadata.get("checklists", [])
-
- meta: Dict[str, Any] = {}
-
- # Phase 1: Checklist coverage (if available)
- coverage_score = 0.0
- if checklists:
- try:
- checklist_text = "\n".join(
- f"- {item}" for item in checklists
- )
- prompt = _COVERAGE_PROMPT.format(
- question=question,
- report=model_answer[:8000], # Truncate long reports
- checklist_items=checklist_text,
- )
- raw = self._ask_judge(
- prompt, temperature=1.0, max_tokens=4096
- )
-
- # Parse coverage_score from last line
- match = re.search(
- r"coverage_score:\s*(\d+)\s*/\s*(\d+)", raw
- )
- if match:
- covered = int(match.group(1))
- total = int(match.group(2))
- coverage_score = covered / total if total > 0 else 0.0
- meta["coverage_covered"] = covered
- meta["coverage_total"] = total
- else:
- # Fallback: count "covered": true occurrences
- covered = raw.lower().count('"covered": true') + raw.lower().count('"covered":true')
- total = len(checklists)
- coverage_score = covered / total if total > 0 else 0.0
- meta["coverage_covered"] = covered
- meta["coverage_total"] = total
-
- meta["coverage_score"] = coverage_score
- meta["coverage_raw"] = raw[:500]
- except Exception as exc:
- LOGGER.warning(
- "Coverage scoring failed for %s: %s",
- record.record_id,
- exc,
- )
- meta["coverage_error"] = str(exc)
-
- # Phase 2: Quality dimensions
- quality_score = 0.0
- try:
- prompt = _QUALITY_PROMPT.format(
- question=question,
- report=model_answer[:8000],
- )
- raw = self._ask_judge(prompt, temperature=1.0, max_tokens=1024)
-
- dims = {}
- for dim in ["presentation", "depth", "citation", "consistency"]:
- match = re.search(rf"{dim}:\s*(\d)", raw)
- if match:
- dims[dim] = int(match.group(1))
-
- if dims:
- quality_score = sum(dims.values()) / (5 * len(dims))
- meta["quality_dims"] = dims
- meta["quality_score"] = quality_score
- meta["quality_raw"] = raw[:500]
- except Exception as exc:
+ if not checklist:
LOGGER.warning(
- "Quality scoring failed for %s: %s", record.record_id, exc
+ "No checklist attached to %s — scoring cannot produce a coverage number",
+ record.record_id,
)
- meta["quality_error"] = str(exc)
+ return None, {
+ "score": 0.0,
+ "coverage": 0.0,
+ "covered_count": 0,
+ "checklist_size": 0,
+ "reason": "no_checklist",
+ }
- # Final score: weighted average of coverage and quality
- if checklists:
- final_score = 0.6 * coverage_score + 0.4 * quality_score
- else:
- final_score = quality_score
+ question = record.metadata.get("question") or record.problem
- meta["final_score"] = final_score
- is_correct = final_score >= 0.5
+ prompt = _build_judge_prompt(
+ question=question,
+ answer=model_answer,
+ checklist=checklist,
+ )
- return is_correct, meta
+ try:
+ raw = self._ask_judge(prompt, temperature=0.0, max_tokens=4096)
+ except Exception as exc:
+ LOGGER.error("LLM judge call failed for %s: %s", record.record_id, exc)
+ return None, {
+ "score": 0.0,
+ "error": str(exc),
+ "checklist_size": len(checklist),
+ }
+
+ judgments = _parse_judge_response(raw, num_items=len(checklist))
+ covered_count = sum(1 for j in judgments if j["covered"])
+ coverage = covered_count / len(checklist)
+ is_correct = coverage >= PASS_THRESHOLD
+
+ metadata: Dict[str, Any] = {
+ "score": coverage,
+ "coverage": coverage,
+ "covered_count": covered_count,
+ "checklist_size": len(checklist),
+ "judgments": judgments,
+ "raw_judge_output": raw,
+ }
+ return is_correct, metadata
-__all__ = ["LiveResearchBenchSFScorer"]
+__all__ = ["LiveResearchBenchScorer"]
diff --git a/src/openjarvis/learning/distillation/__init__.py b/src/openjarvis/learning/distillation/__init__.py
new file mode 100644
index 00000000..7e5f027b
--- /dev/null
+++ b/src/openjarvis/learning/distillation/__init__.py
@@ -0,0 +1 @@
+"""Harness distillation: frontier-driven learning subsystem."""
diff --git a/src/openjarvis/learning/distillation/checkpoint/__init__.py b/src/openjarvis/learning/distillation/checkpoint/__init__.py
new file mode 100644
index 00000000..c4d10c97
--- /dev/null
+++ b/src/openjarvis/learning/distillation/checkpoint/__init__.py
@@ -0,0 +1 @@
+"""Git-backed checkpoint store for config rollback."""
diff --git a/src/openjarvis/learning/distillation/checkpoint/store.py b/src/openjarvis/learning/distillation/checkpoint/store.py
new file mode 100644
index 00000000..deda782a
--- /dev/null
+++ b/src/openjarvis/learning/distillation/checkpoint/store.py
@@ -0,0 +1,255 @@
+"""Git-backed checkpoint store for distillation config rollback.
+
+A thin wrapper over a local git repository at ``/.git``.
+The repo tracks ``config.toml``, ``agents/``, and ``tools/`` so that the
+diff between two commits captures the harness state at any point in time.
+The repo does NOT track ``learning/`` (sessions are append-only artifacts,
+not config state).
+
+The wrapper shells out to ``git`` via ``subprocess`` rather than depending
+on a third-party library — keeps the dependency surface zero.
+
+See spec §7.4 and §7.6.
+"""
+
+from __future__ import annotations
+
+import logging
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+
+from openjarvis.learning.distillation.storage.paths import (
+ ConfigurationError,
+ _find_source_root,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# Files and directories the checkpoint repo tracks.
+_TRACKED_PATHS = ("config.toml", "agents", "tools")
+
+# Marker line in baseline commits so we can detect them on re-init.
+_BASELINE_COMMIT_MESSAGE = "learning: checkpoint baseline"
+
+
+class DirtyWorkingTreeError(RuntimeError):
+ """Raised when begin_stage is called on a working tree with uncommitted changes."""
+
+
+@dataclass(frozen=True)
+class StageHandle:
+ """Snapshot of repo state captured at the start of a staging operation."""
+
+ edit_id: str
+ pre_stage_sha: str
+
+
+class CheckpointStore:
+ """Thin git wrapper for the distillation checkpoint repo.
+
+ Parameters
+ ----------
+ root :
+ The directory that *contains* the checkpoint repo (the repo's
+ ``.git`` lives at ``root / ".git"``). For production use this is
+ ``~/.openjarvis/`` (or ``$OPENJARVIS_HOME``); for tests it's a
+ ``tmp_path`` subdirectory.
+ """
+
+ def __init__(self, root: Path) -> None:
+ self._root = Path(root).resolve()
+
+ @property
+ def root(self) -> Path:
+ return self._root
+
+ # ------------------------------------------------------------------
+ # init
+ # ------------------------------------------------------------------
+
+ def init(self) -> None:
+ """Initialize the checkpoint repo if it doesn't exist.
+
+ Refuses to initialize if ``self.root`` is inside the OpenJarvis source
+ tree — this is the same defense-in-depth check as
+ ``resolve_distillation_root``: we never want a stray git repo writing
+ config snapshots into the working copy.
+
+ Idempotent: if ``.git`` already exists and contains a baseline commit,
+ does nothing.
+ """
+ source_root = _find_source_root()
+ if source_root is not None:
+ try:
+ self._root.relative_to(source_root)
+ except ValueError:
+ pass
+ else:
+ raise ConfigurationError(
+ f"CheckpointStore root ({self._root}) is inside the "
+ f"OpenJarvis source tree ({source_root}). Refusing to "
+ "initialize a checkpoint repo there."
+ )
+
+ self._root.mkdir(parents=True, exist_ok=True)
+
+ if (self._root / ".git").exists() and self._has_baseline_commit():
+ return
+
+ if not (self._root / ".git").exists():
+ self._git("init", "-q")
+ self._git("config", "user.email", "distillation@openjarvis.local")
+ self._git("config", "user.name", "OpenJarvis Distillation")
+
+ # Stage whatever tracked paths currently exist (it's OK if some
+ # don't yet — the user may not have agents or tools dirs at first
+ # init time, in which case the baseline commit is empty).
+ for rel in _TRACKED_PATHS:
+ target = self._root / rel
+ if target.exists():
+ self._git("add", rel)
+
+ # Allow empty so init succeeds even on a brand-new openjarvis home.
+ self._git(
+ "commit",
+ "--allow-empty",
+ "-q",
+ "-m",
+ _BASELINE_COMMIT_MESSAGE,
+ )
+
+ # ------------------------------------------------------------------
+ # current_sha
+ # ------------------------------------------------------------------
+
+ def current_sha(self) -> str:
+ """Return the abbreviated sha of HEAD."""
+ return self._git("rev-parse", "--short", "HEAD")
+
+ # ------------------------------------------------------------------
+ # Staging primitives
+ # ------------------------------------------------------------------
+
+ def begin_stage(self, edit_id: str) -> StageHandle:
+ """Capture HEAD sha and assert the working tree is clean.
+
+ Raises
+ ------
+ DirtyWorkingTreeError
+ If there are uncommitted changes to tracked files. The
+ orchestrator should never start a stage on a dirty tree —
+ it indicates the user has manual edits in progress.
+ """
+ if self._working_tree_dirty():
+ raise DirtyWorkingTreeError(
+ f"Cannot begin stage for {edit_id}: working tree has "
+ "uncommitted changes. Commit or stash them first."
+ )
+ return StageHandle(
+ edit_id=edit_id,
+ pre_stage_sha=self.current_sha(),
+ )
+
+ def commit_stage(
+ self,
+ handle: StageHandle,
+ *,
+ message: str,
+ session_id: str,
+ risk_tier: str,
+ ) -> str:
+ """Stage all tracked-path changes and create a commit.
+
+ The commit message has the form::
+
+
+
+ Edit-ID:
+ Session-ID:
+ Risk-Tier:
+
+ Returns the new commit sha (abbreviated).
+ """
+ for rel in _TRACKED_PATHS:
+ target = self._root / rel
+ if target.exists():
+ self._git("add", rel)
+
+ full_message = (
+ f"{message}\n"
+ "\n"
+ f"Edit-ID: {handle.edit_id}\n"
+ f"Session-ID: {session_id}\n"
+ f"Risk-Tier: {risk_tier}\n"
+ )
+ self._git("commit", "-q", "-m", full_message)
+ return self.current_sha()
+
+ def discard_stage(self, handle: StageHandle) -> None:
+ """Restore the working tree to ``handle.pre_stage_sha``."""
+ for rel in _TRACKED_PATHS:
+ target = self._root / rel
+ if target.exists():
+ self._git("checkout", handle.pre_stage_sha, "--", rel)
+ # Sanity check: HEAD must still equal the pre-stage sha.
+ if self.current_sha() != handle.pre_stage_sha:
+ raise RuntimeError(
+ "discard_stage left HEAD at unexpected sha — something "
+ "committed during the stage. Aborting for safety."
+ )
+
+ # ------------------------------------------------------------------
+ # Session-level rollback
+ # ------------------------------------------------------------------
+
+ def revert_session(self, session_id: str) -> list[str]:
+ """Revert all commits tagged ``Session-ID: ``.
+
+ Reverts are applied in reverse chronological order (newest first),
+ each producing a *new* commit so history is preserved (no rewriting).
+ Returns the list of new revert commit shas.
+ """
+ log = self._git(
+ "log",
+ "--format=%H",
+ f"--grep=Session-ID: {session_id}",
+ )
+ if not log.strip():
+ return []
+
+ # log is newest-first, which is the order we want for revert.
+ target_shas = log.strip().splitlines()
+ new_shas: list[str] = []
+ for sha in target_shas:
+ self._git("revert", "--no-edit", sha)
+ new_shas.append(self.current_sha())
+ return new_shas
+
+ # ------------------------------------------------------------------
+ # Internals
+ # ------------------------------------------------------------------
+
+ def _git(self, *args: str) -> str:
+ """Run a git command in the repo and return stripped stdout."""
+ result = subprocess.run(
+ ["git", *args],
+ cwd=self._root,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ return result.stdout.strip()
+
+ def _has_baseline_commit(self) -> bool:
+ try:
+ log = self._git("log", "--format=%s")
+ except subprocess.CalledProcessError:
+ return False
+ return _BASELINE_COMMIT_MESSAGE in log
+
+ def _working_tree_dirty(self) -> bool:
+ """Return True if there are uncommitted changes to tracked files."""
+ status = self._git("status", "--porcelain")
+ return bool(status.strip())
diff --git a/src/openjarvis/learning/distillation/cli.py b/src/openjarvis/learning/distillation/cli.py
new file mode 100644
index 00000000..0daf61e4
--- /dev/null
+++ b/src/openjarvis/learning/distillation/cli.py
@@ -0,0 +1,132 @@
+"""``jarvis learning`` — distillation learning CLI subcommands.
+
+See spec §12.
+"""
+
+from __future__ import annotations
+
+import click
+from rich.console import Console
+
+console = Console()
+
+
+@click.group("learning")
+def learning_group() -> None:
+ """Frontier-driven harness learning (distillation)."""
+
+
+@learning_group.command("init")
+def learning_init() -> None:
+ """Initialize the distillation checkpoint repo and directory layout."""
+ from openjarvis.learning.distillation.checkpoint.store import CheckpointStore
+ from openjarvis.learning.distillation.storage.paths import (
+ ensure_distillation_dirs,
+ resolve_distillation_root,
+ )
+
+ root = resolve_distillation_root()
+ ensure_distillation_dirs()
+ home = root.parent # ~/.openjarvis
+ store = CheckpointStore(home)
+ store.init()
+ console.print(f"[green]Initialized distillation at {root}[/green]")
+
+
+@learning_group.command("run")
+@click.option(
+ "--autonomy",
+ type=click.Choice(["auto", "tiered", "manual"]),
+ default="tiered",
+)
+def learning_run(autonomy: str) -> None:
+ """Run an on-demand learning session."""
+ console.print("[yellow]On-demand session started.[/yellow]")
+ console.print("[dim]Use 'jarvis learning history' to check results.[/dim]")
+ console.print(f"[dim]Autonomy mode: {autonomy}[/dim]")
+ # Full wiring deferred to integration — this registers the CLI surface.
+ console.print("[dim]Full orchestration requires configured teacher engine.[/dim]")
+
+
+@learning_group.command("history")
+@click.option("--limit", type=int, default=10, help="Max sessions to show.")
+def learning_history(limit: int) -> None:
+ """List past learning sessions."""
+ console.print(f"[dim]Showing last {limit} sessions (requires learning.db).[/dim]")
+
+
+@learning_group.command("show")
+@click.argument("session_id")
+def learning_show(session_id: str) -> None:
+ """Show details of a learning session."""
+ console.print(f"[dim]Session: {session_id}[/dim]")
+
+
+@learning_group.command("review")
+def learning_review() -> None:
+ """Review pending edits awaiting approval."""
+ console.print("[dim]Pending review queue.[/dim]")
+
+
+@learning_group.command("approve")
+@click.argument("edit_id")
+def learning_approve(edit_id: str) -> None:
+ """Approve a pending edit."""
+ console.print(f"[dim]Approving edit: {edit_id}[/dim]")
+
+
+@learning_group.command("reject")
+@click.argument("edit_id")
+@click.option("--reason", type=str, default="", help="Rejection reason.")
+def learning_reject(edit_id: str, reason: str) -> None:
+ """Reject a pending edit."""
+ console.print(f"[dim]Rejecting edit: {edit_id}[/dim]")
+
+
+@learning_group.command("rollback")
+@click.argument("session_id", required=False)
+@click.option("--last", is_flag=True, help="Rollback the most recent session.")
+def learning_rollback(session_id: str | None, last: bool) -> None:
+ """Rollback a learning session's commits."""
+ target = session_id or ("last session" if last else "none")
+ console.print(f"[dim]Rolling back: {target}[/dim]")
+
+
+@learning_group.group("benchmark")
+def benchmark_group() -> None:
+ """Personal benchmark management."""
+
+
+@benchmark_group.command("refresh")
+def benchmark_refresh() -> None:
+ """Manually refresh the personal benchmark."""
+ console.print("[dim]Refreshing personal benchmark.[/dim]")
+
+
+@benchmark_group.command("show")
+def benchmark_show() -> None:
+ """Show current benchmark statistics."""
+ console.print("[dim]Benchmark stats.[/dim]")
+
+
+@learning_group.group("daemon")
+def daemon_group() -> None:
+ """Background learning daemon."""
+
+
+@daemon_group.command("start")
+def daemon_start() -> None:
+ """Start the learning daemon."""
+ console.print("[dim]Starting daemon.[/dim]")
+
+
+@daemon_group.command("stop")
+def daemon_stop() -> None:
+ """Stop the learning daemon."""
+ console.print("[dim]Stopping daemon.[/dim]")
+
+
+@daemon_group.command("status")
+def daemon_status() -> None:
+ """Check daemon status."""
+ console.print("[dim]Daemon status.[/dim]")
diff --git a/src/openjarvis/learning/distillation/diagnose/__init__.py b/src/openjarvis/learning/distillation/diagnose/__init__.py
new file mode 100644
index 00000000..9959de16
--- /dev/null
+++ b/src/openjarvis/learning/distillation/diagnose/__init__.py
@@ -0,0 +1 @@
+"""Diagnose phase: teacher-driven failure analysis."""
diff --git a/src/openjarvis/learning/distillation/diagnose/runner.py b/src/openjarvis/learning/distillation/diagnose/runner.py
new file mode 100644
index 00000000..b0d34616
--- /dev/null
+++ b/src/openjarvis/learning/distillation/diagnose/runner.py
@@ -0,0 +1,318 @@
+"""DiagnosisRunner: orchestrates phase 1 of the distillation loop.
+
+Builds diagnostic tools, runs the TeacherAgent, parses failure clusters
+from the teacher's output, and persists artifacts.
+
+See spec §5.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+)
+from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+)
+from openjarvis.learning.distillation.diagnose.types import ToolCallRecord
+from openjarvis.learning.distillation.models import FailureCluster
+
+logger = logging.getLogger(__name__)
+
+_SYSTEM_PROMPT = """\
+You are a meta-engineer analyzing the performance of a local AI assistant \
+called OpenJarvis. Your job is to diagnose why the local student model fails \
+on certain tasks and identify root cause patterns.
+
+IMPORTANT — OUTPUT REQUIREMENT: You MUST end your response with a JSON array \
+of failure clusters inside a ```json code fence. This is required. Example:
+
+```json
+[
+ {{
+ "id": "cluster-001",
+ "description": "Short description of this failure pattern",
+ "sample_trace_ids": ["trace-abc", "trace-def", "trace-ghi"],
+ "student_failure_rate": 0.75,
+ "teacher_success_rate": 0.95,
+ "skill_gap": "Explanation of the skill gap between student and teacher"
+ }}
+]
+```
+
+You have access to diagnostic tools that let you:
+- Browse and search the student's trace history
+- Read the student's current configuration, prompts, and tools
+- Re-run the student on benchmark tasks
+- Run yourself on the same tasks for comparison
+- Compare student vs teacher outputs
+
+Your analysis should:
+1. Identify 2-5 failure clusters — groups of related failures with shared root causes.
+2. For each cluster, actually run the student and yourself on at least 3 sample tasks \
+to populate student_failure_rate and teacher_success_rate with real data.
+3. Describe the skill gap for each cluster.
+
+Each cluster object in the final JSON array MUST have these fields:
+- id (string)
+- description (string)
+- sample_trace_ids (list of strings)
+- student_failure_rate (float 0-1)
+- teacher_success_rate (float 0-1)
+- skill_gap (string)
+
+Budget: max ~{max_turns} tool calls, max ${max_cost_usd:.2f} USD.
+
+Remember: You MUST end your response with the ```json ... ``` block described above.
+"""
+
+
+@dataclass
+class DiagnosisResult:
+ """The output of a diagnosis run."""
+
+ diagnosis_md: str
+ clusters: list[FailureCluster] = field(default_factory=list)
+ cost_usd: float = 0.0
+ tool_call_records: list[ToolCallRecord] = field(default_factory=list)
+
+
+class DiagnosisRunner:
+ """Orchestrates phase 1 of the distillation loop.
+
+ Parameters
+ ----------
+ teacher_engine :
+ The CloudEngine (or mock) for teacher inference.
+ teacher_model :
+ Frontier model id (e.g. "claude-opus-4-6").
+ trace_store :
+ TraceStore for reading student traces.
+ benchmark_samples :
+ List of PersonalBenchmarkSample objects.
+ student_runner :
+ Callable to re-execute the student on a task.
+ judge :
+ TraceJudge for comparing outputs.
+ session_dir :
+ Path where session artifacts are written.
+ session_id :
+ Current session id.
+ config :
+ Dict with config_path and openjarvis_home.
+ max_turns :
+ Max teacher tool calls (default 30).
+ max_cost_usd :
+ Max teacher API cost (default 5.0).
+ """
+
+ def __init__(
+ self,
+ *,
+ teacher_engine: Any,
+ teacher_model: str,
+ trace_store: Any,
+ benchmark_samples: list,
+ student_runner: Any,
+ judge: Any,
+ session_dir: Path,
+ session_id: str,
+ config: dict[str, Any],
+ max_turns: int = 30,
+ max_cost_usd: float = 5.0,
+ ) -> None:
+ self._teacher_engine = teacher_engine
+ self._teacher_model = teacher_model
+ self._trace_store = trace_store
+ self._benchmark_samples = benchmark_samples
+ self._student_runner = student_runner
+ self._judge = judge
+ self._session_dir = Path(session_dir)
+ self._session_id = session_id
+ self._config = config
+ self._max_turns = max_turns
+ self._max_cost_usd = max_cost_usd
+
+ def run(self) -> DiagnosisResult:
+ """Execute the diagnosis phase.
+
+ Returns
+ -------
+ DiagnosisResult
+ Contains the diagnosis markdown, parsed clusters, cost, and
+ tool call records.
+ """
+ # Ensure session directory exists
+ self._session_dir.mkdir(parents=True, exist_ok=True)
+
+ # Build diagnostic tools
+ tools = build_diagnostic_tools(
+ trace_store=self._trace_store,
+ config=self._config,
+ benchmark_samples=self._benchmark_samples,
+ student_runner=self._student_runner,
+ teacher_engine=self._teacher_engine,
+ teacher_model=self._teacher_model,
+ judge=self._judge,
+ session_id=self._session_id,
+ )
+
+ # Build system prompt with budget hints
+ system_prompt = _SYSTEM_PROMPT.format(
+ max_turns=self._max_turns,
+ max_cost_usd=self._max_cost_usd,
+ )
+
+ # Run the teacher
+ agent = TeacherAgent(
+ engine=self._teacher_engine,
+ model=self._teacher_model,
+ tools=tools,
+ max_turns=self._max_turns,
+ max_cost_usd=self._max_cost_usd,
+ )
+ agent_result = agent.run(
+ "Analyze the student's recent trace history, identify failure patterns, "
+ "and produce a structured diagnosis with failure clusters.",
+ system_prompt=system_prompt,
+ )
+
+ # Persist diagnosis.md
+ diagnosis_path = self._session_dir / "diagnosis.md"
+ diagnosis_path.write_text(agent_result.content, encoding="utf-8")
+
+ # Persist teacher traces JSONL
+ traces_dir = self._session_dir / "teacher_traces"
+ traces_dir.mkdir(parents=True, exist_ok=True)
+ jsonl_path = traces_dir / "diagnose.jsonl"
+ with jsonl_path.open("w", encoding="utf-8") as f:
+ for record in agent_result.tool_call_records:
+ f.write(json.dumps(record.to_jsonl_dict()) + "\n")
+
+ # Parse failure clusters from the diagnosis content
+ clusters = _parse_clusters(agent_result.content)
+
+ # Fallback: if no clusters parsed, ask the teacher to emit only the JSON
+ if not clusters:
+ logger.warning(
+ "No clusters in primary diagnosis (%d chars, %d tool calls). "
+ "Attempting fallback extraction.",
+ len(agent_result.content),
+ len(agent_result.tool_call_records),
+ )
+ clusters = self._fallback_extract_clusters(agent_result.content)
+
+ return DiagnosisResult(
+ diagnosis_md=agent_result.content,
+ clusters=clusters,
+ cost_usd=agent_result.total_cost_usd,
+ tool_call_records=agent_result.tool_call_records,
+ )
+
+ def _fallback_extract_clusters(self, diagnosis: str) -> list[FailureCluster]:
+ """One-shot fallback: ask the teacher to emit only the JSON array.
+
+ Makes a single no-tools call with max 1 turn. If the response still
+ does not contain valid clusters, returns an empty list.
+
+ Parameters
+ ----------
+ diagnosis :
+ The full diagnosis text from the primary teacher run.
+ """
+ fallback_prompt = (
+ "Your diagnosis did not include the required JSON cluster array. "
+ "Here is your diagnosis:\n\n"
+ f"{diagnosis[:3000]}\n\n"
+ "Now output ONLY a raw JSON array of failure clusters. "
+ "Each object must have: id, description, sample_trace_ids, "
+ "student_failure_rate (float 0-1), teacher_success_rate (float 0-1), "
+ "skill_gap. Output ONLY the JSON array — no markdown, no code "
+ "fences, no explanation, no other text."
+ )
+
+ fallback_agent = TeacherAgent(
+ engine=self._teacher_engine,
+ model=self._teacher_model,
+ tools=[], # no tools — single generation call
+ max_turns=1,
+ max_cost_usd=self._max_cost_usd,
+ )
+ try:
+ fallback_result = fallback_agent.run(fallback_prompt)
+ except Exception:
+ logger.exception("Fallback extraction TeacherAgent call failed")
+ return []
+
+ clusters = _parse_clusters(fallback_result.content)
+ if not clusters:
+ logger.warning("Fallback extraction also produced no clusters")
+ return clusters
+
+
+def _parse_clusters(content: str) -> list[FailureCluster]:
+ """Extract failure clusters from teacher diagnosis output.
+
+ Looks for a JSON array inside a ```json code fence. Falls back to
+ searching for any JSON array in the content. Returns an empty list
+ if no valid clusters are found.
+ """
+ # Try to find JSON in a code fence first
+ fence_match = re.search(r"```json\s*\n(.*?)\n```", content, re.DOTALL)
+ if fence_match:
+ try:
+ return _parse_cluster_list(fence_match.group(1))
+ except Exception as e:
+ logger.warning("Failed to parse clusters from JSON code fence: %s", e)
+
+ # Fallback: try to find any JSON array in the content
+ # Use greedy match to capture the full array (not just the first [...])
+ for match in re.finditer(r"\[[\s\S]*\]", content):
+ try:
+ return _parse_cluster_list(match.group(0))
+ except Exception:
+ continue
+
+ logger.warning("No failure clusters found in diagnosis output")
+ return []
+
+
+def _parse_cluster_list(json_str: str) -> list[FailureCluster]:
+ """Parse a JSON string into a list of FailureCluster."""
+ data = json.loads(json_str)
+ if not isinstance(data, list):
+ return []
+ clusters = []
+ for item in data:
+ try:
+ # Clamp rates to 0-1 range (teacher sometimes outputs percentages)
+ failure_rate = float(item.get("student_failure_rate", 0))
+ success_rate = float(item.get("teacher_success_rate", 0))
+ if failure_rate > 1.0:
+ failure_rate = failure_rate / 100.0
+ if success_rate > 1.0:
+ success_rate = success_rate / 100.0
+ failure_rate = max(0.0, min(1.0, failure_rate))
+ success_rate = max(0.0, min(1.0, success_rate))
+
+ clusters.append(
+ FailureCluster(
+ id=str(item.get("id", f"cluster-{len(clusters) + 1}")),
+ description=str(item.get("description", "")),
+ sample_trace_ids=[str(t) for t in item.get("sample_trace_ids", [])],
+ student_failure_rate=failure_rate,
+ teacher_success_rate=success_rate,
+ skill_gap=str(item.get("skill_gap", "")),
+ )
+ )
+ except Exception as e:
+ logger.warning("Skipping invalid cluster: %s", e)
+ continue
+ return clusters
diff --git a/src/openjarvis/learning/distillation/diagnose/teacher_agent.py b/src/openjarvis/learning/distillation/diagnose/teacher_agent.py
new file mode 100644
index 00000000..c919ced6
--- /dev/null
+++ b/src/openjarvis/learning/distillation/diagnose/teacher_agent.py
@@ -0,0 +1,269 @@
+"""TeacherAgent: frontier model as a tool-calling meta-engineer.
+
+NOT registered in ``AgentRegistry``. NOT a subclass of ``BaseAgent``.
+This is a standalone tool-calling loop that wraps ``CloudEngine`` with
+diagnostic tools for the diagnose phase.
+
+The teacher:
+- Uses a frontier model (default ``claude-opus-4-6``) regardless of the
+ user's local intelligence config.
+- Has its own tool set (diagnostic tools) that user agents do not have.
+- Tracks cost and stops when the budget is exhausted.
+- Logs every tool call to a list of ``ToolCallRecord``.
+
+See spec §5.1.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from openjarvis.core.types import Message, Role, ToolCall
+from openjarvis.learning.distillation.diagnose.types import (
+ DiagnosticTool,
+ ToolCallRecord,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TeacherAgentResult:
+ """The result of a TeacherAgent.run() call."""
+
+ content: str
+ turns: int
+ total_cost_usd: float
+ total_tokens: int
+ tool_call_records: list[ToolCallRecord] = field(default_factory=list)
+
+
+class TeacherAgent:
+ """Frontier model tool-calling loop for the diagnose phase.
+
+ Parameters
+ ----------
+ engine :
+ A ``CloudEngine`` (or mock) that provides ``generate()``.
+ model :
+ The frontier model id (e.g. ``"claude-opus-4-6"``).
+ tools :
+ Diagnostic tools exposed to the teacher.
+ max_turns :
+ Maximum number of generate() calls before stopping.
+ max_cost_usd :
+ Maximum accumulated cost before stopping.
+ """
+
+ def __init__(
+ self,
+ engine: Any,
+ model: str,
+ tools: list[DiagnosticTool],
+ max_turns: int = 30,
+ max_cost_usd: float = 5.0,
+ max_tokens: int = 8192,
+ ) -> None:
+ self._engine = engine
+ self._model = model
+ self._tools = {t.name: t for t in tools}
+ self._tool_specs = [t.to_openai_function() for t in tools]
+ self._max_turns = max_turns
+ self._max_cost_usd = max_cost_usd
+ self._max_tokens = max_tokens
+
+ def run(
+ self,
+ user_prompt: str,
+ *,
+ system_prompt: str | None = None,
+ ) -> TeacherAgentResult:
+ """Run the teacher tool-calling loop.
+
+ Parameters
+ ----------
+ user_prompt :
+ The instruction to the teacher (e.g. "Analyze student failures").
+ system_prompt :
+ Optional system prompt explaining the teacher's role.
+
+ Returns
+ -------
+ TeacherAgentResult
+ The teacher's final content, cost, and tool call records.
+ """
+ messages: list[Message] = []
+ if system_prompt:
+ messages.append(Message(role=Role.SYSTEM, content=system_prompt))
+ messages.append(Message(role=Role.USER, content=user_prompt))
+
+ total_cost = 0.0
+ total_tokens = 0
+ tool_call_records: list[ToolCallRecord] = []
+ final_content = ""
+
+ gen_kwargs: dict[str, Any] = {}
+ if self._tool_specs:
+ gen_kwargs["tools"] = self._tool_specs
+
+ for turn in range(1, self._max_turns + 1):
+ # Budget pre-check: don't start a new turn if we've already
+ # exceeded the budget. Without this, the `if not raw_tool_calls`
+ # early-return path below bypasses the post-check entirely,
+ # meaning the final (terminal) turn runs unchecked. Pre-checking
+ # here closes that hole.
+ #
+ # Note: a single turn that ITSELF exceeds the remaining budget
+ # still overshoots. With Opus + ~8k max_tokens + growing input
+ # context, a late-in-conversation final turn can cost $2-4. To
+ # bound this further, lower `max_turns` (caps context growth)
+ # or lower `max_cost_usd` in the config.
+ if total_cost >= self._max_cost_usd:
+ logger.warning(
+ "Teacher budget exceeded before turn %d: $%.2f >= $%.2f",
+ turn,
+ total_cost,
+ self._max_cost_usd,
+ )
+ return TeacherAgentResult(
+ content=final_content,
+ turns=turn - 1,
+ total_cost_usd=total_cost,
+ total_tokens=total_tokens,
+ tool_call_records=tool_call_records,
+ )
+
+ result = self._engine.generate(
+ messages=messages,
+ model=self._model,
+ max_tokens=self._max_tokens,
+ **gen_kwargs,
+ )
+
+ cost = result.get("cost_usd", 0.0)
+ total_cost += cost
+ usage = result.get("usage", {})
+ total_tokens += usage.get("total_tokens", 0)
+
+ content = result.get("content", "")
+ raw_tool_calls = result.get("tool_calls", [])
+
+ if not raw_tool_calls:
+ final_content = content
+ return TeacherAgentResult(
+ content=final_content,
+ turns=turn,
+ total_cost_usd=total_cost,
+ total_tokens=total_tokens,
+ tool_call_records=tool_call_records,
+ )
+
+ # Convert raw tool calls to ToolCall objects
+ tool_call_objs = []
+ for tc in raw_tool_calls:
+ tc_obj = ToolCall(
+ id=tc["id"] if isinstance(tc, dict) else tc.id,
+ name=tc["name"] if isinstance(tc, dict) else tc.name,
+ arguments=(
+ tc.get("arguments", "{}")
+ if isinstance(tc, dict)
+ else tc.arguments
+ ),
+ )
+ tool_call_objs.append(tc_obj)
+
+ # Append assistant message with tool calls
+ messages.append(
+ Message(
+ role=Role.ASSISTANT,
+ content=content,
+ tool_calls=tool_call_objs,
+ )
+ )
+
+ # Execute each tool call
+ for tc in tool_call_objs:
+ tc_name = tc.name
+ tc_id = tc.id
+ tc_args_str = tc.arguments
+
+ try:
+ tc_args = json.loads(tc_args_str)
+ except json.JSONDecodeError:
+ tc_args = {}
+
+ tool = self._tools.get(tc_name)
+ start_time = time.monotonic()
+ if tool is not None:
+ try:
+ tool_result = tool.fn(**tc_args)
+ except Exception as e:
+ tool_result = json.dumps({"error": str(e)})
+ logger.warning("Tool %s raised: %s", tc_name, e)
+ else:
+ tool_result = json.dumps({"error": f"Unknown tool: {tc_name}"})
+ elapsed_ms = (time.monotonic() - start_time) * 1000
+
+ # Safety cap on tool results to bound context growth.
+ # Primary truncation should happen inside each tool, but this
+ # is defense-in-depth for tools that return too much. 20KB
+ # per call × 30 max_turns = 600KB total, well under Opus's
+ # 1M-token context window.
+ MAX_TOOL_RESULT_CHARS = 20000
+ if len(tool_result) > MAX_TOOL_RESULT_CHARS:
+ tool_result = (
+ tool_result[:MAX_TOOL_RESULT_CHARS]
+ + f"\n...[tool result truncated from {len(tool_result)} chars]"
+ )
+
+ tool_call_records.append(
+ ToolCallRecord(
+ timestamp=datetime.now(timezone.utc),
+ tool=tc_name,
+ args=tc_args,
+ result=tool_result[:8000], # Shorter cap for audit log
+ latency_ms=elapsed_ms,
+ cost_usd=0.0, # Tool calls themselves are free
+ )
+ )
+
+ # Append tool result message (uses the safety-capped value)
+ messages.append(
+ Message(
+ role=Role.TOOL,
+ content=tool_result,
+ tool_call_id=tc_id,
+ name=tc_name,
+ )
+ )
+
+ # Check cost budget
+ if total_cost >= self._max_cost_usd:
+ logger.warning(
+ "Teacher cost budget exceeded: $%.2f >= $%.2f",
+ total_cost,
+ self._max_cost_usd,
+ )
+ final_content = content
+ return TeacherAgentResult(
+ content=final_content,
+ turns=turn,
+ total_cost_usd=total_cost,
+ total_tokens=total_tokens,
+ tool_call_records=tool_call_records,
+ )
+
+ # Exhausted max_turns
+ logger.warning("Teacher exhausted max_turns=%d", self._max_turns)
+ return TeacherAgentResult(
+ content=final_content,
+ turns=self._max_turns,
+ total_cost_usd=total_cost,
+ total_tokens=total_tokens,
+ tool_call_records=tool_call_records,
+ )
diff --git a/src/openjarvis/learning/distillation/diagnose/tools.py b/src/openjarvis/learning/distillation/diagnose/tools.py
new file mode 100644
index 00000000..24feda78
--- /dev/null
+++ b/src/openjarvis/learning/distillation/diagnose/tools.py
@@ -0,0 +1,543 @@
+"""Diagnostic tools exposed to the teacher in the diagnose phase.
+
+All tools are **read-only** relative to the user's config. They do not mutate
+``~/.openjarvis/config.toml``, agent prompts, or tool descriptions.
+
+Tools that execute code (``run_student_on_task``, ``run_self_on_task``) append
+new traces to ``TraceStore`` as a side effect. These traces are tagged with
+``source=distillation_session:`` so they can be excluded from future
+learning input.
+
+See spec §5.2.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from openjarvis.learning.distillation.diagnose.types import DiagnosticTool
+
+
+def build_diagnostic_tools(
+ *,
+ trace_store: Any,
+ config: dict[str, Any],
+ benchmark_samples: list,
+ student_runner: Any,
+ teacher_engine: Any,
+ teacher_model: str,
+ judge: Any,
+ session_id: str,
+) -> list[DiagnosticTool]:
+ """Build all diagnostic tools as closures over shared dependencies.
+
+ Parameters
+ ----------
+ trace_store :
+ A ``TraceStore`` instance for trace queries.
+ config :
+ Dict with ``config_path`` (Path) and ``openjarvis_home`` (Path).
+ benchmark_samples :
+ List of ``PersonalBenchmarkSample`` objects.
+ student_runner :
+ Callable that re-executes the student on a task.
+ teacher_engine :
+ The ``CloudEngine`` instance used by the teacher.
+ teacher_model :
+ Model id for teacher inference (e.g. "claude-opus-4-6").
+ judge :
+ A ``TraceJudge`` instance for comparing outputs.
+ session_id :
+ Current session id for tagging traces.
+ """
+ openjarvis_home = config["openjarvis_home"]
+
+ # ------------------------------------------------------------------
+ # list_traces
+ # ------------------------------------------------------------------
+ def _list_traces(
+ limit: int = 20,
+ agent: str | None = None,
+ outcome: str | None = None,
+ min_feedback: float | None = None,
+ max_feedback: float | None = None,
+ ) -> str:
+ kwargs: dict[str, Any] = {"limit": limit}
+ if agent:
+ kwargs["agent"] = agent
+ if outcome:
+ kwargs["outcome"] = outcome
+ traces = trace_store.list_traces(**kwargs)
+ metas = []
+ for t in traces:
+ fb = getattr(t, "feedback", None)
+ if min_feedback is not None and (fb is None or fb < min_feedback):
+ continue
+ if max_feedback is not None and (fb is None or fb > max_feedback):
+ continue
+ metas.append(
+ {
+ "trace_id": t.trace_id,
+ "query": t.query[:200],
+ "agent": t.agent,
+ "model": t.model,
+ "outcome": t.outcome,
+ "feedback": fb,
+ "started_at": t.started_at,
+ }
+ )
+ return json.dumps(metas, default=str)
+
+ # ------------------------------------------------------------------
+ # get_trace
+ # ------------------------------------------------------------------
+ def _get_trace(trace_id: str) -> str:
+ trace = trace_store.get(trace_id)
+ if trace is None:
+ return json.dumps({"error": f"Trace {trace_id} not found"})
+ # Truncate large free-text fields so a single get_trace call can't
+ # blow out the teacher's context window. monitor_operative traces
+ # can have 50-95K tokens of result text each.
+ query = trace.query or ""
+ result = trace.result or ""
+ max_query, max_result = 2000, 6000
+ query_display = query[:max_query] + (
+ f"...[query truncated from {len(query)} chars]"
+ if len(query) > max_query
+ else ""
+ )
+ result_display = result[:max_result] + (
+ f"...[result truncated from {len(result)} chars]"
+ if len(result) > max_result
+ else ""
+ )
+ return json.dumps(
+ {
+ "trace_id": trace.trace_id,
+ "query": query_display,
+ "agent": trace.agent,
+ "model": trace.model,
+ "outcome": trace.outcome,
+ "feedback": getattr(trace, "feedback", None),
+ "result": result_display,
+ "total_tokens": trace.total_tokens,
+ "total_latency_seconds": trace.total_latency_seconds,
+ "steps": [str(s)[:200] for s in getattr(trace, "steps", [])[:20]],
+ },
+ default=str,
+ )
+
+ # ------------------------------------------------------------------
+ # search_traces
+ # ------------------------------------------------------------------
+ def _search_traces(query: str, limit: int = 20) -> str:
+ # Cap limit to avoid pathological calls that would OOM the context.
+ # 20 x 95K-token traces = ~1.9M tokens, well over Opus's 1M window.
+ limit = max(1, min(limit, 20))
+ try:
+ results = trace_store.search(query, limit=limit)
+ except Exception as e:
+ # SQLite FTS5 raises "unknown special query" for patterns with
+ # unescaped special chars (- / . etc). Return a structured error
+ # so the teacher can retry with a cleaner query instead of the
+ # call being swallowed and context still growing.
+ return json.dumps(
+ {
+ "error": f"search_traces failed: {e}",
+ "hint": (
+ "FTS5 special chars (- / . : etc) need escaping. "
+ "Try a simpler keyword query."
+ ),
+ }
+ )
+ # Truncate per-trace text fields. Each trace's `result` column can be
+ # 50-95K tokens; returning raw content would blow the context window.
+ max_query_per_trace = 300
+ max_result_per_trace = 1500
+ for r in results:
+ q = r.get("query") or ""
+ if len(q) > max_query_per_trace:
+ r["query"] = q[:max_query_per_trace] + f"...[{len(q)} chars]"
+ res = r.get("result") or ""
+ if len(res) > max_result_per_trace:
+ r["result"] = res[:max_result_per_trace] + f"...[{len(res)} chars]"
+ return json.dumps(results, default=str)
+
+ # ------------------------------------------------------------------
+ # get_current_config
+ # ------------------------------------------------------------------
+ def _get_current_config() -> str:
+ config_path = config["config_path"]
+ try:
+ return config_path.read_text(encoding="utf-8")
+ except FileNotFoundError:
+ return "No config.toml found."
+
+ # ------------------------------------------------------------------
+ # get_agent_prompt
+ # ------------------------------------------------------------------
+ def _get_agent_prompt(agent_name: str) -> str:
+ prompt_path = openjarvis_home / "agents" / agent_name / "system_prompt.md"
+ try:
+ return prompt_path.read_text(encoding="utf-8")
+ except FileNotFoundError:
+ return f"No system prompt found for agent '{agent_name}'."
+
+ # ------------------------------------------------------------------
+ # get_tool_description
+ # ------------------------------------------------------------------
+ def _get_tool_description(tool_name: str) -> str:
+ desc_path = openjarvis_home / "tools" / "descriptions.toml"
+ try:
+ content = desc_path.read_text(encoding="utf-8")
+ # Simple TOML parsing for the description field
+ in_section = False
+ for line in content.splitlines():
+ if line.strip() == f"[{tool_name}]":
+ in_section = True
+ continue
+ if in_section and line.startswith("["):
+ break
+ if in_section and "description" in line:
+ return line.split("=", 1)[1].strip().strip('"')
+ return f"Tool '{tool_name}' found but no description field."
+ except FileNotFoundError:
+ return "No descriptions.toml found."
+
+ # ------------------------------------------------------------------
+ # list_available_tools
+ # ------------------------------------------------------------------
+ def _list_available_tools() -> str:
+ # Read from the on-disk descriptions.toml
+ desc_path = openjarvis_home / "tools" / "descriptions.toml"
+ tools_list = []
+ try:
+ content = desc_path.read_text(encoding="utf-8")
+ current_tool = None
+ for line in content.splitlines():
+ if line.startswith("[") and line.endswith("]"):
+ current_tool = line[1:-1]
+ tools_list.append(
+ {
+ "name": current_tool,
+ "description": "",
+ "category": "general",
+ "agents": [],
+ }
+ )
+ elif current_tool and "description" in line and "=" in line:
+ desc = line.split("=", 1)[1].strip().strip('"')
+ if tools_list:
+ tools_list[-1]["description"] = desc
+ except FileNotFoundError:
+ pass
+ return json.dumps(tools_list, default=str)
+
+ # ------------------------------------------------------------------
+ # list_personal_benchmark
+ # ------------------------------------------------------------------
+ def _list_personal_benchmark(limit: int = 50) -> str:
+ tasks = []
+ for sample in benchmark_samples[:limit]:
+ tasks.append(
+ {
+ "task_id": sample.trace_id,
+ "query": sample.query,
+ "reference_answer": sample.reference_answer[:500],
+ "category": getattr(sample, "category", "chat"),
+ }
+ )
+ return json.dumps(tasks, default=str)
+
+ # ------------------------------------------------------------------
+ # run_student_on_task
+ # ------------------------------------------------------------------
+ def _run_student_on_task(task_id: str) -> str:
+ sample = next((s for s in benchmark_samples if s.trace_id == task_id), None)
+ if sample is None:
+ return json.dumps({"error": f"Task {task_id} not found in benchmark"})
+ result = student_runner(sample.query, session_id=session_id)
+ return json.dumps(
+ {
+ "task_id": task_id,
+ "output": str(getattr(result, "content", result)),
+ "score": getattr(result, "score", 0.0),
+ "trace_id": getattr(result, "trace_id", ""),
+ "latency_seconds": getattr(result, "latency_seconds", 0.0),
+ "tokens_used": getattr(result, "tokens_used", 0),
+ },
+ default=str,
+ )
+
+ # ------------------------------------------------------------------
+ # run_self_on_task
+ # ------------------------------------------------------------------
+ def _run_self_on_task(task_id: str, max_tokens: int = 2048) -> str:
+ sample = next((s for s in benchmark_samples if s.trace_id == task_id), None)
+ if sample is None:
+ return json.dumps({"error": f"Task {task_id} not found in benchmark"})
+ response = teacher_engine.generate(
+ messages=[{"role": "user", "content": sample.query}],
+ model=teacher_model,
+ max_tokens=max_tokens,
+ )
+ return json.dumps(
+ {
+ "task_id": task_id,
+ "output": response.get("content", ""),
+ "reasoning": "",
+ "cost_usd": response.get("cost_usd", 0.0),
+ "tokens_used": response.get("usage", {}).get("total_tokens", 0),
+ },
+ default=str,
+ )
+
+ # ------------------------------------------------------------------
+ # compare_outputs
+ # ------------------------------------------------------------------
+ def _compare_outputs(student_output: str, teacher_output: str, task: str) -> str:
+ score, reasoning = judge.score_trace(
+ type(
+ "FakeTrace",
+ (),
+ {
+ "query": task,
+ "result": student_output,
+ "steps": [],
+ "messages": [],
+ "agent": "student",
+ "model": "local",
+ "total_tokens": 0,
+ "total_latency_seconds": 0,
+ "metadata": {},
+ },
+ )()
+ )
+ teacher_score, _ = judge.score_trace(
+ type(
+ "FakeTrace",
+ (),
+ {
+ "query": task,
+ "result": teacher_output,
+ "steps": [],
+ "messages": [],
+ "agent": "teacher",
+ "model": "frontier",
+ "total_tokens": 0,
+ "total_latency_seconds": 0,
+ "metadata": {},
+ },
+ )()
+ )
+ return json.dumps(
+ {
+ "task_id": "",
+ "student_score": score,
+ "teacher_score": teacher_score,
+ "judge_reasoning": reasoning,
+ },
+ default=str,
+ )
+
+ # ------------------------------------------------------------------
+ # Assemble
+ # ------------------------------------------------------------------
+ return [
+ DiagnosticTool(
+ name="list_traces",
+ description=(
+ "Browse traces by agent, outcome, feedback range."
+ " Returns a JSON list of trace summaries."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "limit": {
+ "type": "integer",
+ "description": "Max results",
+ "default": 20,
+ },
+ "agent": {
+ "type": "string",
+ "description": "Filter by agent name",
+ },
+ "outcome": {
+ "type": "string",
+ "description": "Filter by outcome: success/failure",
+ },
+ "min_feedback": {
+ "type": "number",
+ "description": "Minimum feedback score",
+ },
+ "max_feedback": {
+ "type": "number",
+ "description": "Maximum feedback score",
+ },
+ },
+ },
+ fn=_list_traces,
+ ),
+ DiagnosticTool(
+ name="get_trace",
+ description=(
+ "Read a single trace including query, result, steps, and metrics."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "trace_id": {
+ "type": "string",
+ "description": "The trace ID to retrieve",
+ },
+ },
+ "required": ["trace_id"],
+ },
+ fn=_get_trace,
+ ),
+ DiagnosticTool(
+ name="search_traces",
+ description=(
+ "Full-text search across traces. Returns matching trace summaries."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "query": {"type": "string", "description": "Search query"},
+ "limit": {
+ "type": "integer",
+ "description": "Max results",
+ "default": 20,
+ },
+ },
+ "required": ["query"],
+ },
+ fn=_search_traces,
+ ),
+ DiagnosticTool(
+ name="get_current_config",
+ description="Read the current OpenJarvis config.toml.",
+ parameters={"type": "object", "properties": {}},
+ fn=_get_current_config,
+ ),
+ DiagnosticTool(
+ name="get_agent_prompt",
+ description="Read the current system prompt for a named agent.",
+ parameters={
+ "type": "object",
+ "properties": {
+ "agent_name": {
+ "type": "string",
+ "description": "Agent name (e.g. 'simple')",
+ },
+ },
+ "required": ["agent_name"],
+ },
+ fn=_get_agent_prompt,
+ ),
+ DiagnosticTool(
+ name="get_tool_description",
+ description="Read the LM-facing description of a tool.",
+ parameters={
+ "type": "object",
+ "properties": {
+ "tool_name": {
+ "type": "string",
+ "description": "Tool name (e.g. 'web_search')",
+ },
+ },
+ "required": ["tool_name"],
+ },
+ fn=_get_tool_description,
+ ),
+ DiagnosticTool(
+ name="list_available_tools",
+ description="List all tools and which agents currently have them enabled.",
+ parameters={"type": "object", "properties": {}},
+ fn=_list_available_tools,
+ ),
+ DiagnosticTool(
+ name="list_personal_benchmark",
+ description=(
+ "Browse personal benchmark tasks with their reference answers."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "limit": {
+ "type": "integer",
+ "description": "Max tasks to return",
+ "default": 50,
+ },
+ },
+ },
+ fn=_list_personal_benchmark,
+ ),
+ DiagnosticTool(
+ name="run_student_on_task",
+ description=(
+ "Re-execute the local student agent on a benchmark task."
+ " Returns the student's output and score."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "task_id": {
+ "type": "string",
+ "description": "Benchmark task ID",
+ },
+ },
+ "required": ["task_id"],
+ },
+ fn=_run_student_on_task,
+ ),
+ DiagnosticTool(
+ name="run_self_on_task",
+ description=(
+ "Run yourself (the teacher) on a benchmark task"
+ " to produce a reference answer."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "task_id": {
+ "type": "string",
+ "description": "Benchmark task ID",
+ },
+ "max_tokens": {
+ "type": "integer",
+ "description": "Max tokens for response",
+ "default": 2048,
+ },
+ },
+ "required": ["task_id"],
+ },
+ fn=_run_self_on_task,
+ ),
+ DiagnosticTool(
+ name="compare_outputs",
+ description=(
+ "Compare student and teacher outputs on a task using the TraceJudge."
+ ),
+ parameters={
+ "type": "object",
+ "properties": {
+ "student_output": {
+ "type": "string",
+ "description": "The student's output",
+ },
+ "teacher_output": {
+ "type": "string",
+ "description": "The teacher's output",
+ },
+ "task": {
+ "type": "string",
+ "description": "The original task/query",
+ },
+ },
+ "required": ["student_output", "teacher_output", "task"],
+ },
+ fn=_compare_outputs,
+ ),
+ ]
diff --git a/src/openjarvis/learning/distillation/diagnose/types.py b/src/openjarvis/learning/distillation/diagnose/types.py
new file mode 100644
index 00000000..40669bc6
--- /dev/null
+++ b/src/openjarvis/learning/distillation/diagnose/types.py
@@ -0,0 +1,129 @@
+"""Data types for the diagnose phase.
+
+Lightweight dataclasses used as return types by diagnostic tools and
+as internal data carriers. These are NOT pydantic models — they don't
+need validation or JSON schema generation.
+
+See spec §5.2 for the tool return type rationale.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Optional
+
+
+@dataclass(slots=True)
+class TraceMeta:
+ """Lightweight summary of a trace for browsing."""
+
+ trace_id: str
+ query: str
+ agent: str
+ model: str
+ outcome: Optional[str]
+ feedback: Optional[float]
+ started_at: float
+
+
+@dataclass(slots=True)
+class BenchmarkTask:
+ """One task from the personal benchmark."""
+
+ task_id: str
+ query: str
+ reference_answer: str
+ category: str = "chat"
+
+
+@dataclass(slots=True)
+class StudentRun:
+ """Result of re-executing the local student on a benchmark task."""
+
+ task_id: str
+ output: str
+ score: float
+ trace_id: str
+ latency_seconds: float
+ tokens_used: int
+
+
+@dataclass(slots=True)
+class TeacherRun:
+ """Result of the teacher running itself on a benchmark task."""
+
+ task_id: str
+ output: str
+ reasoning: str
+ cost_usd: float
+ tokens_used: int
+
+
+@dataclass(slots=True)
+class ComparisonResult:
+ """Structured comparison between student and teacher outputs."""
+
+ task_id: str
+ student_score: float
+ teacher_score: float
+ judge_reasoning: str
+
+
+@dataclass(slots=True)
+class ToolMeta:
+ """Metadata about a tool in the ToolRegistry."""
+
+ name: str
+ description: str
+ category: str
+ agents: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class DiagnosticTool:
+ """A tool exposed to the teacher in the diagnose phase.
+
+ Unlike ``BaseTool``, these are not registered in ``ToolRegistry``.
+ They are lightweight wrappers: a name, description, JSON schema
+ for parameters, and a callable that implements the tool.
+ """
+
+ name: str
+ description: str
+ parameters: dict[str, Any]
+ fn: Callable[..., Any]
+
+ def to_openai_function(self) -> dict[str, Any]:
+ """Convert to OpenAI function-calling format."""
+ return {
+ "type": "function",
+ "function": {
+ "name": self.name,
+ "description": self.description,
+ "parameters": self.parameters,
+ },
+ }
+
+
+@dataclass(slots=True)
+class ToolCallRecord:
+ """One teacher tool call, persisted to the JSONL log."""
+
+ timestamp: datetime
+ tool: str
+ args: dict[str, Any]
+ result: str
+ latency_ms: float
+ cost_usd: float
+
+ def to_jsonl_dict(self) -> dict[str, Any]:
+ """Serialize to a JSON-safe dict for JSONL output."""
+ return {
+ "timestamp": self.timestamp.isoformat(),
+ "tool": self.tool,
+ "args": self.args,
+ "result": self.result,
+ "latency_ms": self.latency_ms,
+ "cost_usd": self.cost_usd,
+ }
diff --git a/src/openjarvis/learning/distillation/execute/__init__.py b/src/openjarvis/learning/distillation/execute/__init__.py
new file mode 100644
index 00000000..bcf1a45d
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/__init__.py
@@ -0,0 +1 @@
+"""Execute phase: apply edits to the harness configuration."""
diff --git a/src/openjarvis/learning/distillation/execute/appliers/__init__.py b/src/openjarvis/learning/distillation/execute/appliers/__init__.py
new file mode 100644
index 00000000..6f698731
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/appliers/__init__.py
@@ -0,0 +1,6 @@
+"""Concrete EditApplier implementations.
+
+Importing this package triggers registration of all appliers
+in the EditApplierRegistry. Applier modules are imported in Task 6
+once all implementations exist.
+"""
diff --git a/src/openjarvis/learning/distillation/execute/appliers/agent.py b/src/openjarvis/learning/distillation/execute/appliers/agent.py
new file mode 100644
index 00000000..00714e5d
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/appliers/agent.py
@@ -0,0 +1,226 @@
+"""Agent-pillar appliers: prompts, class, params, few-shot.
+
+See spec §4.1 op semantics for agent ops.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ ApplyResult,
+ EditApplier,
+ ValidationResult,
+)
+from openjarvis.learning.distillation.models import Edit, EditOp
+from openjarvis.learning.distillation.plan.prompt_diff import apply_unified_diff
+
+logger = logging.getLogger(__name__)
+
+
+def _agent_prompt_path(ctx: ApplyContext, agent_name: str) -> Path:
+ return ctx.agents_dir / agent_name / "system_prompt.md"
+
+
+def _extract_agent_name(edit: Edit) -> str:
+ """Extract agent name from edit target or payload."""
+ if "agent" in edit.payload:
+ return edit.payload["agent"]
+ # Try from target: "agents.simple.system_prompt" -> "simple"
+ parts = edit.target.split(".")
+ if len(parts) >= 2:
+ return parts[1]
+ return "default"
+
+
+class ReplaceSystemPromptApplier(EditApplier):
+ """Overwrite an agent's entire system prompt."""
+
+ op = EditOp.REPLACE_SYSTEM_PROMPT
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ if "new_content" not in edit.payload:
+ return ValidationResult(ok=False, reason="Missing new_content in payload")
+ agent = _extract_agent_name(edit)
+ prompt_path = _agent_prompt_path(ctx, agent)
+ if not prompt_path.parent.exists():
+ return ValidationResult(
+ ok=False, reason=f"Agent directory not found: {prompt_path.parent}"
+ )
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = _extract_agent_name(edit)
+ prompt_path = _agent_prompt_path(ctx, agent)
+ prompt_path.write_text(edit.payload["new_content"], encoding="utf-8")
+ return ApplyResult(changed_files=[str(prompt_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class PatchSystemPromptApplier(EditApplier):
+ """Apply a unified diff to an agent's system prompt."""
+
+ op = EditOp.PATCH_SYSTEM_PROMPT
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ if "diff" not in edit.payload:
+ return ValidationResult(ok=False, reason="Missing diff in payload")
+ agent = _extract_agent_name(edit)
+ prompt_path = _agent_prompt_path(ctx, agent)
+ if not prompt_path.exists():
+ return ValidationResult(
+ ok=False, reason=f"Prompt file not found: {prompt_path}"
+ )
+ original = prompt_path.read_text(encoding="utf-8")
+ patched = apply_unified_diff(original, edit.payload["diff"])
+ if patched is None:
+ return ValidationResult(
+ ok=False, reason="Diff cannot be applied to current prompt"
+ )
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = _extract_agent_name(edit)
+ prompt_path = _agent_prompt_path(ctx, agent)
+ original = prompt_path.read_text(encoding="utf-8")
+ patched = apply_unified_diff(original, edit.payload["diff"])
+ if patched is None:
+ raise RuntimeError(f"Failed to apply diff to {prompt_path}")
+ prompt_path.write_text(patched, encoding="utf-8")
+ return ApplyResult(changed_files=[str(prompt_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class SetAgentClassApplier(EditApplier):
+ """Change which agent class is used."""
+
+ op = EditOp.SET_AGENT_CLASS
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ if "new_class" not in edit.payload:
+ return ValidationResult(ok=False, reason="Missing new_class in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = _extract_agent_name(edit)
+ new_class = edit.payload["new_class"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ ) # noqa: E501
+
+ section = f"[agent.{agent}]"
+ if section in content:
+ lines = content.splitlines()
+ in_section = False
+ found = False
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ break
+ if in_section and line.strip().startswith("class"):
+ lines[i] = f'class = "{new_class}"'
+ found = True
+ break
+ if not found:
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ lines.insert(i + 1, f'class = "{new_class}"')
+ break
+ content = "\n".join(lines) + "\n"
+ else:
+ content += f'\n{section}\nclass = "{new_class}"\n'
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class SetAgentParamApplier(EditApplier):
+ """Update an agent parameter."""
+
+ op = EditOp.SET_AGENT_PARAM
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ for key in ("agent", "param", "value"):
+ if key not in edit.payload:
+ return ValidationResult(ok=False, reason=f"Missing {key} in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = edit.payload["agent"]
+ param = edit.payload["param"]
+ value = edit.payload["value"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ ) # noqa: E501
+
+ section = f"[agent.{agent}]"
+ if section in content:
+ lines = content.splitlines()
+ in_section = False
+ found = False
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ lines.insert(i, f"{param} = {value}")
+ found = True
+ break
+ if in_section and line.strip().startswith(f"{param}"):
+ lines[i] = f"{param} = {value}"
+ found = True
+ break
+ if not found:
+ lines.append(f"{param} = {value}")
+ content = "\n".join(lines) + "\n"
+ else:
+ content += f"\n{section}\n{param} = {value}\n"
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class EditFewShotExemplarsApplier(EditApplier):
+ """Write few-shot exemplars to an agent's directory."""
+
+ op = EditOp.EDIT_FEW_SHOT_EXEMPLARS
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ if "exemplars" not in edit.payload:
+ return ValidationResult(ok=False, reason="Missing exemplars in payload")
+ agent = _extract_agent_name(edit)
+ agent_dir = ctx.agents_dir / agent
+ if not agent_dir.exists():
+ return ValidationResult(
+ ok=False, reason=f"Agent directory not found: {agent_dir}"
+ )
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = _extract_agent_name(edit)
+ fs_path = ctx.agents_dir / agent / "few_shot.json"
+ fs_path.write_text(
+ json.dumps(edit.payload["exemplars"], indent=2),
+ encoding="utf-8",
+ )
+ return ApplyResult(changed_files=[str(fs_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
diff --git a/src/openjarvis/learning/distillation/execute/appliers/intelligence.py b/src/openjarvis/learning/distillation/execute/appliers/intelligence.py
new file mode 100644
index 00000000..44b4d142
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/appliers/intelligence.py
@@ -0,0 +1,117 @@
+"""Intelligence-pillar appliers: model routing and parameters.
+
+See spec §4.1 op semantics for SET_MODEL_FOR_QUERY_CLASS and SET_MODEL_PARAM.
+"""
+
+from __future__ import annotations
+
+from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ ApplyResult,
+ EditApplier,
+ ValidationResult,
+)
+from openjarvis.learning.distillation.models import Edit, EditOp
+
+
+class SetModelForQueryClassApplier(EditApplier):
+ """Update the routing policy map for a query class."""
+
+ op = EditOp.SET_MODEL_FOR_QUERY_CLASS
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ if "query_class" not in edit.payload or "model" not in edit.payload:
+ return ValidationResult(
+ ok=False, reason="Missing query_class or model in payload"
+ )
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ query_class = edit.payload["query_class"]
+ model = edit.payload["model"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ )
+
+ # Check if the policy_map section exists
+ if "[learning.routing.policy_map]" in content:
+ # Check if this query_class already has a line
+ lines = content.splitlines()
+ found = False
+ for i, line in enumerate(lines):
+ stripped = line.strip()
+ if stripped.startswith(f"{query_class}") and "=" in stripped:
+ lines[i] = f'{query_class} = "{model}"'
+ found = True
+ break
+ if not found:
+ # Append after the [learning.routing.policy_map] header
+ for i, line in enumerate(lines):
+ if line.strip() == "[learning.routing.policy_map]":
+ lines.insert(i + 1, f'{query_class} = "{model}"')
+ break
+ content = "\n".join(lines) + "\n"
+ else:
+ # Append the section
+ content += f'\n[learning.routing.policy_map]\n{query_class} = "{model}"\n'
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass # Delegated to CheckpointStore.discard_stage() in the execution loop
+
+
+class SetModelParamApplier(EditApplier):
+ """Update a model parameter in config."""
+
+ op = EditOp.SET_MODEL_PARAM
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ for key in ("model", "param", "value"):
+ if key not in edit.payload:
+ return ValidationResult(ok=False, reason=f"Missing {key} in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ model = edit.payload["model"]
+ param = edit.payload["param"]
+ value = edit.payload["value"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ )
+
+ # Sanitize model name for TOML section header (replace : with -)
+ section_key = model.replace(":", "-")
+ section_header = f"[models.{section_key}]"
+
+ if section_header in content:
+ lines = content.splitlines()
+ in_section = False
+ found = False
+ for i, line in enumerate(lines):
+ if line.strip() == section_header:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ # Hit next section — insert before it
+ lines.insert(i, f"{param} = {value}")
+ found = True
+ break
+ if in_section and line.strip().startswith(f"{param}"):
+ lines[i] = f"{param} = {value}"
+ found = True
+ break
+ if not found:
+ lines.append(f"{param} = {value}")
+ content = "\n".join(lines) + "\n"
+ else:
+ content += f"\n{section_header}\n{param} = {value}\n"
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass # Delegated to CheckpointStore.discard_stage()
diff --git a/src/openjarvis/learning/distillation/execute/appliers/lora_stub.py b/src/openjarvis/learning/distillation/execute/appliers/lora_stub.py
new file mode 100644
index 00000000..6c30e3c8
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/appliers/lora_stub.py
@@ -0,0 +1,38 @@
+"""LoRA fine-tuning stub — deferred to v2.
+
+The planner can emit LORA_FINETUNE edits so the diagnosis surfaces
+"this should be a weight update" pressure to the user, but the executor
+refuses them in v1.
+
+See spec §4.1.
+"""
+
+from __future__ import annotations
+
+from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ ApplyResult,
+ EditApplier,
+ ValidationResult,
+)
+from openjarvis.learning.distillation.models import Edit, EditOp
+
+
+class LoraStubApplier(EditApplier):
+ """Refuses LORA_FINETUNE with a clear v2 message."""
+
+ op = EditOp.LORA_FINETUNE
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ return ValidationResult(
+ ok=False,
+ reason="LORA_FINETUNE is deferred to v2. "
+ "The planner emitted this edit to signal that weight updates "
+ "would help, but the executor cannot apply them yet.",
+ )
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ raise NotImplementedError("LORA_FINETUNE deferred to v2")
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
diff --git a/src/openjarvis/learning/distillation/execute/appliers/tools.py b/src/openjarvis/learning/distillation/execute/appliers/tools.py
new file mode 100644
index 00000000..9ec3aca8
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/appliers/tools.py
@@ -0,0 +1,159 @@
+"""Tools-pillar appliers: add/remove tools, edit descriptions.
+
+See spec §4.1 op semantics for tool ops.
+"""
+
+from __future__ import annotations
+
+import re
+
+from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ ApplyResult,
+ EditApplier,
+ ValidationResult,
+)
+from openjarvis.learning.distillation.models import Edit, EditOp
+
+
+class AddToolToAgentApplier(EditApplier):
+ """Add a tool to an agent's tool list in config."""
+
+ op = EditOp.ADD_TOOL_TO_AGENT
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ for key in ("agent", "tool_name"):
+ if key not in edit.payload:
+ return ValidationResult(ok=False, reason=f"Missing {key} in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = edit.payload["agent"]
+ tool_name = edit.payload["tool_name"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ )
+
+ section = f"[agent.{agent}]"
+ if section in content:
+ lines = content.splitlines()
+ in_section = False
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ break
+ if in_section and line.strip().startswith("tools"):
+ # Parse existing tool list and add new tool
+ match = re.search(r"\[([^\]]*)\]", line)
+ if match:
+ existing = match.group(1)
+ if tool_name not in existing:
+ new_list = existing.rstrip() + f', "{tool_name}"'
+ lines[i] = f"tools = [{new_list}]"
+ break
+ content = "\n".join(lines) + "\n"
+ else:
+ content += f'\n{section}\ntools = ["{tool_name}"]\n'
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class RemoveToolFromAgentApplier(EditApplier):
+ """Remove a tool from an agent's tool list in config."""
+
+ op = EditOp.REMOVE_TOOL_FROM_AGENT
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ for key in ("agent", "tool_name"):
+ if key not in edit.payload:
+ return ValidationResult(ok=False, reason=f"Missing {key} in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ agent = edit.payload["agent"]
+ tool_name = edit.payload["tool_name"]
+ config_path = ctx.config_path
+ content = (
+ config_path.read_text(encoding="utf-8") if config_path.exists() else ""
+ )
+
+ section = f"[agent.{agent}]"
+ if section in content:
+ lines = content.splitlines()
+ in_section = False
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ break
+ if in_section and line.strip().startswith("tools"):
+ # Remove the tool from the list
+ line = re.sub(rf',?\s*"{re.escape(tool_name)}"', "", line)
+ line = re.sub(rf'"{re.escape(tool_name)}"\s*,?\s*', "", line)
+ lines[i] = line
+ break
+ content = "\n".join(lines) + "\n"
+
+ config_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(config_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+
+class EditToolDescriptionApplier(EditApplier):
+ """Update a tool's LM-facing description in descriptions.toml."""
+
+ op = EditOp.EDIT_TOOL_DESCRIPTION
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ for key in ("tool_name", "new_description"):
+ if key not in edit.payload:
+ return ValidationResult(ok=False, reason=f"Missing {key} in payload")
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ tool_name = edit.payload["tool_name"]
+ new_desc = edit.payload["new_description"]
+ desc_path = ctx.tools_dir / "descriptions.toml"
+ desc_path.parent.mkdir(parents=True, exist_ok=True)
+
+ content = desc_path.read_text(encoding="utf-8") if desc_path.exists() else ""
+
+ section = f"[{tool_name}]"
+ if section in content:
+ lines = content.splitlines()
+ in_section = False
+ found = False
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ in_section = True
+ continue
+ if in_section and line.strip().startswith("["):
+ break
+ if in_section and line.strip().startswith("description"):
+ lines[i] = f'description = "{new_desc}"'
+ found = True
+ break
+ if not found:
+ for i, line in enumerate(lines):
+ if line.strip() == section:
+ lines.insert(i + 1, f'description = "{new_desc}"')
+ break
+ content = "\n".join(lines) + "\n"
+ else:
+ content += f'\n{section}\ndescription = "{new_desc}"\n'
+
+ desc_path.write_text(content, encoding="utf-8")
+ return ApplyResult(changed_files=[str(desc_path)])
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
diff --git a/src/openjarvis/learning/distillation/execute/base.py b/src/openjarvis/learning/distillation/execute/base.py
new file mode 100644
index 00000000..62d8dc85
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/base.py
@@ -0,0 +1,95 @@
+"""EditApplier ABC, registry, and context types for the execute phase.
+
+Each concrete applier implements validate/apply/rollback for a single EditOp.
+Appliers are registered in an EditApplierRegistry keyed by EditOp.
+
+See spec §7.1.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import ClassVar
+
+from openjarvis.learning.distillation.models import Edit, EditOp
+
+
+@dataclass
+class ApplyContext:
+ """Shared context passed to all appliers."""
+
+ openjarvis_home: Path
+ session_id: str
+
+ @property
+ def config_path(self) -> Path:
+ return self.openjarvis_home / "config.toml"
+
+ @property
+ def agents_dir(self) -> Path:
+ return self.openjarvis_home / "agents"
+
+ @property
+ def tools_dir(self) -> Path:
+ return self.openjarvis_home / "tools"
+
+
+@dataclass
+class ValidationResult:
+ """Result of EditApplier.validate()."""
+
+ ok: bool
+ reason: str = ""
+
+
+@dataclass
+class ApplyResult:
+ """Result of EditApplier.apply()."""
+
+ changed_files: list[str] = field(default_factory=list)
+
+
+class EditApplier(ABC):
+ """Abstract base for edit appliers.
+
+ Each subclass handles one EditOp. It validates the edit against
+ the current config state, applies the mutation, and can roll back.
+ """
+
+ op: ClassVar[EditOp]
+
+ @abstractmethod
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ """Check if the edit can be applied to the current config."""
+ ...
+
+ @abstractmethod
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ """Mutate the config. Must be idempotent."""
+ ...
+
+ @abstractmethod
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ """Restore pre-edit state. Most appliers delegate to git checkout."""
+ ...
+
+
+class EditApplierRegistry:
+ """Registry of EditApplier instances keyed by EditOp."""
+
+ def __init__(self) -> None:
+ self._appliers: dict[EditOp, EditApplier] = {}
+
+ def register(self, applier: EditApplier) -> None:
+ """Register an applier instance."""
+ self._appliers[applier.op] = applier
+
+ def get(self, op: EditOp) -> EditApplier:
+ """Return the applier for the given op. Raises KeyError if not found."""
+ return self._appliers[op]
+
+ def is_supported(self, op: EditOp) -> bool:
+ """Return True if an applier is registered for the op."""
+ return op in self._appliers
diff --git a/src/openjarvis/learning/distillation/execute/loop.py b/src/openjarvis/learning/distillation/execute/loop.py
new file mode 100644
index 00000000..d30fea33
--- /dev/null
+++ b/src/openjarvis/learning/distillation/execute/loop.py
@@ -0,0 +1,194 @@
+"""Per-edit execution loop for the distillation execute phase.
+
+Iterates over a plan's edits, handles tier routing, validates, and applies.
+Does NOT include the benchmark gate — that's wired in M5.
+
+See spec §7.2.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+
+from openjarvis.learning.distillation.execute.appliers.agent import (
+ EditFewShotExemplarsApplier,
+ PatchSystemPromptApplier,
+ ReplaceSystemPromptApplier,
+ SetAgentClassApplier,
+ SetAgentParamApplier,
+)
+from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelForQueryClassApplier,
+ SetModelParamApplier,
+)
+from openjarvis.learning.distillation.execute.appliers.lora_stub import (
+ LoraStubApplier,
+)
+from openjarvis.learning.distillation.execute.appliers.tools import (
+ AddToolToAgentApplier,
+ EditToolDescriptionApplier,
+ RemoveToolFromAgentApplier,
+)
+from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ EditApplierRegistry,
+)
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ Edit,
+ EditOutcome,
+ EditRiskTier,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _build_registry() -> EditApplierRegistry:
+ """Build and populate the default applier registry."""
+ registry = EditApplierRegistry()
+ registry.register(SetModelForQueryClassApplier())
+ registry.register(SetModelParamApplier())
+ registry.register(PatchSystemPromptApplier())
+ registry.register(ReplaceSystemPromptApplier())
+ registry.register(SetAgentClassApplier())
+ registry.register(SetAgentParamApplier())
+ registry.register(EditFewShotExemplarsApplier())
+ registry.register(AddToolToAgentApplier())
+ registry.register(RemoveToolFromAgentApplier())
+ registry.register(EditToolDescriptionApplier())
+ registry.register(LoraStubApplier())
+ return registry
+
+
+def execute_edits(
+ *,
+ edits: list[Edit],
+ ctx: ApplyContext,
+ autonomy_mode: AutonomyMode,
+ registry: EditApplierRegistry | None = None,
+) -> list[EditOutcome]:
+ """Execute a list of edits, returning outcomes for each.
+
+ This loop handles tier routing and validation but does NOT include
+ the benchmark gate (wired in M5).
+
+ Parameters
+ ----------
+ edits :
+ The edits to process.
+ ctx :
+ Shared context with config paths.
+ autonomy_mode :
+ How aggressively to apply edits.
+ registry :
+ Optional pre-built registry. If None, builds the default.
+ """
+ if registry is None:
+ registry = _build_registry()
+
+ outcomes: list[EditOutcome] = []
+
+ for edit in edits:
+ # Manual mode: everything goes to review
+ if autonomy_mode == AutonomyMode.MANUAL:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="pending_review",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Manual tier: always skip
+ if edit.risk_tier == EditRiskTier.MANUAL:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="skipped",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error="manual tier, requires explicit approval",
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Review tier in tiered mode: route to pending
+ if (
+ edit.risk_tier == EditRiskTier.REVIEW
+ and autonomy_mode == AutonomyMode.TIERED
+ ):
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="pending_review",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Check if the op is supported
+ if not registry.is_supported(edit.op):
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="skipped",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=f"op {edit.op.value} not implemented in v1",
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Validate
+ applier = registry.get(edit.op)
+ validation = applier.validate(edit, ctx)
+ if not validation.ok:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="rejected_by_gate",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=validation.reason,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Apply
+ try:
+ applier.apply(edit, ctx)
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="applied",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=datetime.now(timezone.utc),
+ )
+ )
+ except Exception as e:
+ logger.warning("Edit %s failed: %s", edit.id, e)
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="rejected_by_gate",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=str(e),
+ applied_at=None,
+ )
+ )
+
+ return outcomes
diff --git a/src/openjarvis/learning/distillation/gate/__init__.py b/src/openjarvis/learning/distillation/gate/__init__.py
new file mode 100644
index 00000000..95eb4acd
--- /dev/null
+++ b/src/openjarvis/learning/distillation/gate/__init__.py
@@ -0,0 +1 @@
+"""Gate phase: benchmark-based accept/reject for edits."""
diff --git a/src/openjarvis/learning/distillation/gate/benchmark_gate.py b/src/openjarvis/learning/distillation/gate/benchmark_gate.py
new file mode 100644
index 00000000..c3f94b27
--- /dev/null
+++ b/src/openjarvis/learning/distillation/gate/benchmark_gate.py
@@ -0,0 +1,131 @@
+"""BenchmarkGate: accept or reject edits based on benchmark performance.
+
+Runs the personal benchmark via a provided scorer callable, compares
+before/after snapshots, and decides accept/reject based on thresholds.
+
+See spec §7.3.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Callable
+
+from openjarvis.learning.distillation.gate.regression import regression_check
+from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+logger = logging.getLogger(__name__)
+
+# Type for the scorer callable: takes benchmark_version, subsample_size,
+# seed and returns a BenchmarkSnapshot.
+ScorerFn = Callable[..., BenchmarkSnapshot]
+
+
+@dataclass
+class GateResult:
+ """Result of a benchmark gate evaluation."""
+
+ accepted: bool
+ snapshot: BenchmarkSnapshot
+ delta: float
+ reason: str = ""
+
+
+class BenchmarkGate:
+ """Runs the personal benchmark and decides accept/reject.
+
+ Parameters
+ ----------
+ scorer :
+ Callable that runs the benchmark and returns a BenchmarkSnapshot.
+ Signature: ``(benchmark_version, subsample_size, seed) -> BenchmarkSnapshot``.
+ benchmark_version :
+ Which benchmark version to score against (locked per session).
+ min_improvement :
+ Minimum overall score improvement to accept (default 0.0).
+ max_regression :
+ Maximum per-cluster score drop before rejecting (default 0.05).
+ subsample_size :
+ Number of tasks to score per gate run (default 50).
+ """
+
+ def __init__(
+ self,
+ *,
+ scorer: ScorerFn,
+ benchmark_version: str,
+ min_improvement: float = 0.0,
+ max_regression: float = 0.05,
+ subsample_size: int = 50,
+ ) -> None:
+ self._scorer = scorer
+ self._benchmark_version = benchmark_version
+ self._min_improvement = min_improvement
+ self._max_regression = max_regression
+ self._subsample_size = subsample_size
+
+ def evaluate(
+ self,
+ *,
+ before: BenchmarkSnapshot,
+ session_seed: int,
+ ) -> GateResult:
+ """Run the benchmark and compare against the before snapshot.
+
+ Parameters
+ ----------
+ before :
+ Snapshot captured before the edit was applied.
+ session_seed :
+ Deterministic seed for subsampling (same across all gate
+ runs in one session).
+
+ Returns
+ -------
+ GateResult
+ ``accepted`` is True if the edit should be committed.
+ """
+ after = self._scorer(
+ benchmark_version=self._benchmark_version,
+ subsample_size=self._subsample_size,
+ seed=session_seed,
+ )
+
+ delta = after.overall_score - before.overall_score
+
+ # Check regression
+ reg = regression_check(before, after, max_regression=self._max_regression)
+ if reg.has_regression:
+ clusters_str = ", ".join(
+ f"{cid}: {d:+.3f}" for cid, d in reg.regressed_clusters.items()
+ )
+ reason = f"regression in clusters: {clusters_str}"
+ logger.info("Gate rejected: %s", reason)
+ return GateResult(
+ accepted=False,
+ snapshot=after,
+ delta=delta,
+ reason=reason,
+ )
+
+ # Check improvement
+ if delta <= self._min_improvement:
+ reason = (
+ f"no improvement: delta={delta:.4f}, "
+ f"min_improvement={self._min_improvement}"
+ )
+ logger.info("Gate rejected: %s", reason)
+ return GateResult(
+ accepted=False,
+ snapshot=after,
+ delta=delta,
+ reason=reason,
+ )
+
+ logger.info("Gate accepted: delta=%.4f", delta)
+ return GateResult(
+ accepted=True,
+ snapshot=after,
+ delta=delta,
+ )
diff --git a/src/openjarvis/learning/distillation/gate/cold_start.py b/src/openjarvis/learning/distillation/gate/cold_start.py
new file mode 100644
index 00000000..4179d664
--- /dev/null
+++ b/src/openjarvis/learning/distillation/gate/cold_start.py
@@ -0,0 +1,109 @@
+"""Cold start detection and bootstrap for the distillation subsystem.
+
+Day one: no traces, no benchmark. The system must not crash and must
+give the user a clear message about what's needed. This module provides
+readiness checks and the bootstrap logic.
+
+See spec §13.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ReadinessResult:
+ """Result of a readiness check."""
+
+ ready: bool
+ message: str
+ trace_count: int = 0
+ high_feedback_count: int = 0
+
+
+def check_readiness(
+ trace_store: Any,
+ min_traces: int = 20,
+) -> ReadinessResult:
+ """Check if there are enough traces to run a learning session.
+
+ Parameters
+ ----------
+ trace_store :
+ TraceStore instance.
+ min_traces :
+ Minimum total trace count required.
+
+ Returns
+ -------
+ ReadinessResult
+ ``ready`` is True if there are enough traces.
+ """
+ count = trace_store.count()
+ if count < min_traces:
+ return ReadinessResult(
+ ready=False,
+ message=(
+ f"Not enough traces yet to learn from. "
+ f"Have {count}, need at least {min_traces}. "
+ f"Use OpenJarvis for a while and try again."
+ ),
+ trace_count=count,
+ )
+ return ReadinessResult(
+ ready=True,
+ message=f"Ready: {count} traces available.",
+ trace_count=count,
+ )
+
+
+def check_benchmark_ready(
+ trace_store: Any,
+ min_feedback: float = 0.7,
+ min_samples: int = 10,
+) -> ReadinessResult:
+ """Check if there are enough high-feedback traces for a benchmark.
+
+ The bootstrap benchmark needs at least ``min_samples`` traces with
+ feedback >= ``min_feedback``.
+
+ Parameters
+ ----------
+ trace_store :
+ TraceStore instance.
+ min_feedback :
+ Minimum feedback score for benchmark-quality traces.
+ min_samples :
+ Minimum number of high-feedback traces needed.
+ """
+ # Query traces with high feedback
+ traces = trace_store.list_traces(limit=min_samples * 2)
+ high_feedback = [
+ t
+ for t in traces
+ if getattr(t, "feedback", None) is not None and t.feedback >= min_feedback
+ ]
+ count = len(high_feedback)
+
+ if count < min_samples:
+ return ReadinessResult(
+ ready=False,
+ message=(
+ f"Personal benchmark needs at least {min_samples} "
+ f"high-feedback traces (feedback >= {min_feedback}). "
+ f"Have {count} so far. Will be populated automatically."
+ ),
+ trace_count=trace_store.count(),
+ high_feedback_count=count,
+ )
+ return ReadinessResult(
+ ready=True,
+ message=f"Benchmark ready: {count} high-feedback traces available.",
+ trace_count=trace_store.count(),
+ high_feedback_count=count,
+ )
diff --git a/src/openjarvis/learning/distillation/gate/regression.py b/src/openjarvis/learning/distillation/gate/regression.py
new file mode 100644
index 00000000..ef4599f5
--- /dev/null
+++ b/src/openjarvis/learning/distillation/gate/regression.py
@@ -0,0 +1,60 @@
+"""Per-cluster regression detection for the benchmark gate.
+
+After an edit is applied, the gate compares the before and after
+BenchmarkSnapshots. If any cluster's score dropped by more than
+``max_regression``, the edit is rejected.
+
+See spec §7.3.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+
+@dataclass
+class RegressionResult:
+ """Result of a regression check."""
+
+ has_regression: bool
+ regressed_clusters: dict[str, float] = field(default_factory=dict)
+ """Cluster id → negative delta for clusters that regressed beyond threshold."""
+
+
+def regression_check(
+ before: BenchmarkSnapshot,
+ after: BenchmarkSnapshot,
+ max_regression: float = 0.05,
+) -> RegressionResult:
+ """Check if any cluster regressed beyond the threshold.
+
+ Parameters
+ ----------
+ before :
+ Benchmark snapshot before the edit.
+ after :
+ Benchmark snapshot after the edit.
+ max_regression :
+ Maximum allowed per-cluster score drop (default 0.05).
+
+ Returns
+ -------
+ RegressionResult
+ ``has_regression`` is True if any cluster dropped more than
+ ``max_regression``. ``regressed_clusters`` maps cluster ids
+ to their negative deltas.
+ """
+ regressed: dict[str, float] = {}
+
+ for cluster_id, before_score in before.cluster_scores.items():
+ after_score = after.cluster_scores.get(cluster_id, 0.0)
+ delta = after_score - before_score
+ if delta < -max_regression:
+ regressed[cluster_id] = delta
+
+ return RegressionResult(
+ has_regression=bool(regressed),
+ regressed_clusters=regressed,
+ )
diff --git a/src/openjarvis/learning/distillation/models.py b/src/openjarvis/learning/distillation/models.py
new file mode 100644
index 00000000..973d891e
--- /dev/null
+++ b/src/openjarvis/learning/distillation/models.py
@@ -0,0 +1,372 @@
+"""Pydantic models and enums for the distillation subsystem.
+
+This module defines the typed vocabulary used by the diagnose, plan, execute,
+and record phases. Three model families:
+
+- Enums: pillar / risk tier / op / trigger kind / autonomy mode / session status
+- Edit + LearningPlan + FailureCluster: the teacher's frozen output
+- LearningSession + EditOutcome + BenchmarkSnapshot: the durable session record
+
+See spec §4 for the data model rationale.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class EditPillar(str, Enum):
+ """Which OpenJarvis pillar an edit targets."""
+
+ INTELLIGENCE = "intelligence"
+ AGENT = "agent"
+ TOOLS = "tools"
+ ENGINE = "engine" # designed-for; no v1 appliers ship
+
+
+class EditRiskTier(str, Enum):
+ """How an edit gets applied: auto, review queue, or manual-only."""
+
+ AUTO = "auto"
+ REVIEW = "review"
+ MANUAL = "manual"
+
+
+class EditOp(str, Enum):
+ """The set of typed operations a teacher can propose.
+
+ Each op corresponds to one EditApplier in v1 (or a refusing stub for
+ deferred ops). The teacher cannot invent new ops — only choose from this
+ set.
+ """
+
+ # Intelligence
+ SET_MODEL_FOR_QUERY_CLASS = "set_model_for_query_class"
+ SET_MODEL_PARAM = "set_model_param"
+
+ # Agent
+ PATCH_SYSTEM_PROMPT = "patch_system_prompt"
+ REPLACE_SYSTEM_PROMPT = "replace_system_prompt"
+ SET_AGENT_CLASS = "set_agent_class"
+ SET_AGENT_PARAM = "set_agent_param"
+ EDIT_FEW_SHOT_EXEMPLARS = "edit_few_shot_exemplars"
+
+ # Tools
+ ADD_TOOL_TO_AGENT = "add_tool_to_agent"
+ REMOVE_TOOL_FROM_AGENT = "remove_tool_from_agent"
+ EDIT_TOOL_DESCRIPTION = "edit_tool_description"
+
+ # v2 placeholder — planner can emit, executor refuses with NotImplementedError
+ LORA_FINETUNE = "lora_finetune"
+
+
+class TriggerKind(str, Enum):
+ """What kicked off a learning session."""
+
+ SCHEDULED = "scheduled"
+ CLUSTER = "cluster"
+ USER_FLAG = "user_flag"
+ ON_DEMAND = "on_demand"
+
+
+class AutonomyMode(str, Enum):
+ """How aggressively the orchestrator applies edits without review."""
+
+ AUTO = "auto" # all tiers auto-apply, ignore review tier
+ TIERED = "tiered" # default: respect per-edit risk tier
+ MANUAL = "manual" # everything goes to review queue (dry-run mode)
+
+
+class SessionStatus(str, Enum):
+ """Lifecycle states for a LearningSession.
+
+ See spec §7.7 for the transition rules.
+ """
+
+ INITIATED = "initiated"
+ DIAGNOSING = "diagnosing"
+ PLANNING = "planning"
+ EXECUTING = "executing"
+ AWAITING_REVIEW = "awaiting_review"
+ COMPLETED = "completed"
+ FAILED = "failed"
+ ROLLED_BACK = "rolled_back"
+
+
+# ---------------------------------------------------------------------------
+# Edit — atomic unit of change
+# ---------------------------------------------------------------------------
+
+
+class Edit(BaseModel):
+ """One atomic edit to the OpenJarvis harness.
+
+ Emitted by the LearningPlanner, consumed by an EditApplier. The teacher
+ proposes the op, target, payload, rationale, and references; the planner
+ overwrites ``risk_tier`` deterministically from the (pillar, op) lookup
+ table — the teacher cannot pick its own tier.
+
+ See spec §4.1.
+ """
+
+ id: str = Field(
+ ...,
+ description="UUID for this edit; also used as a footer in git commits.",
+ )
+ pillar: EditPillar = Field(..., description="Which OpenJarvis pillar is touched.")
+ op: EditOp = Field(..., description="The typed operation to perform.")
+ target: str = Field(
+ ...,
+ description="Dotted path to the target, e.g. 'agents.simple.system_prompt'.",
+ )
+ payload: dict[str, Any] = Field(
+ default_factory=dict,
+ description="Op-specific arguments. Schema depends on op.",
+ )
+ rationale: str = Field(
+ ...,
+ description="Teacher's natural-language reason for this edit.",
+ )
+ expected_improvement: str = Field(
+ ...,
+ description="Which failure cluster id this edit is intended to address.",
+ )
+ risk_tier: EditRiskTier = Field(
+ ...,
+ description="Set by the planner from a (pillar, op) lookup table.",
+ )
+ references: list[str] = Field(
+ default_factory=list,
+ description="Trace ids or benchmark task ids that justify this edit.",
+ )
+
+
+# ---------------------------------------------------------------------------
+# FailureCluster — a group of related failures with a shared root cause
+# ---------------------------------------------------------------------------
+
+
+class FailureCluster(BaseModel):
+ """A group of failing traces that share a hypothesised root cause.
+
+ Populated by the teacher in the diagnose phase. The student/teacher rates
+ must come from real ``run_student_on_task`` and ``run_self_on_task`` calls
+ against benchmark tasks (see spec §5.3); clusters where both rates are
+ missing or zero are dropped by the planner.
+
+ See spec §4.2.
+ """
+
+ id: str
+ description: str = Field(..., description="Short human description of the cluster.")
+ sample_trace_ids: list[str] = Field(
+ default_factory=list,
+ description="Trace ids that exemplify this cluster (>= 3 typical).",
+ )
+ student_failure_rate: float = Field(
+ ...,
+ ge=0.0,
+ le=1.0,
+ description="Local student's failure rate on benchmark tasks in this cluster.",
+ )
+ teacher_success_rate: float = Field(
+ ...,
+ ge=0.0,
+ le=1.0,
+ description="Frontier teacher's success rate on the same tasks.",
+ )
+ skill_gap: str = Field(
+ ...,
+ description="Teacher's qualitative analysis of what the student is missing.",
+ )
+ addressed_by_edit_ids: list[str] = Field(
+ default_factory=list,
+ description="Ids of edits in the LearningPlan that target this cluster.",
+ )
+
+
+# ---------------------------------------------------------------------------
+# LearningPlan — frozen output of the planning phase
+# ---------------------------------------------------------------------------
+
+
+class LearningPlan(BaseModel):
+ """The teacher's frozen plan of edits for a learning session.
+
+ Once written to ``/plan.json`` this is immutable. The
+ execution layer reads this file and does not re-prompt the teacher.
+
+ See spec §4.2 and §6.
+ """
+
+ session_id: str = Field(..., description="Owning session id.")
+ diagnosis_summary: str = Field(
+ ...,
+ description="Markdown narrative analysis from the teacher (~500-2000 words).",
+ )
+ failure_clusters: list[FailureCluster] = Field(
+ default_factory=list,
+ description="Clusters identified in phase 1.",
+ )
+ edits: list[Edit] = Field(
+ default_factory=list,
+ description="Typed edit list emitted by the planner.",
+ )
+ teacher_model: str = Field(
+ ...,
+ description="Frontier model id used as the teacher (e.g. 'claude-opus-4-6').",
+ )
+ estimated_cost_usd: float = Field(
+ ...,
+ ge=0.0,
+ description="Total teacher API cost estimate for this session.",
+ )
+ created_at: datetime = Field(
+ ...,
+ description="When the plan was finalized.",
+ )
+
+
+# ---------------------------------------------------------------------------
+# BenchmarkSnapshot — one personal-benchmark run result
+# ---------------------------------------------------------------------------
+
+
+class BenchmarkSnapshot(BaseModel):
+ """A point-in-time score from running the personal benchmark.
+
+ Two of these live on every LearningSession: one captured before any edits
+ apply, one captured after. The version is locked at session start so the
+ delta is interpretable even if the benchmark is refreshed mid-session.
+
+ See spec §4.3 and §9.
+ """
+
+ benchmark_version: str = Field(
+ ...,
+ description="Personal benchmark version (e.g. 'personal_v3').",
+ )
+ overall_score: float = Field(
+ ...,
+ ge=0.0,
+ le=1.0,
+ description="Mean per-task score across the benchmark.",
+ )
+ cluster_scores: dict[str, float] = Field(
+ default_factory=dict,
+ description="Mean score per failure cluster.",
+ )
+ task_count: int = Field(
+ ...,
+ ge=0,
+ description="Number of tasks scored in this snapshot.",
+ )
+ elapsed_seconds: float = Field(
+ ...,
+ ge=0.0,
+ description="Wall-clock time of the benchmark run.",
+ )
+
+
+# ---------------------------------------------------------------------------
+# EditOutcome — what happened when one edit was processed
+# ---------------------------------------------------------------------------
+
+
+EditOutcomeStatus = Literal[
+ "applied",
+ "rejected_by_gate",
+ "pending_review",
+ "rejected_by_user",
+ "rolled_back",
+ "skipped",
+]
+
+
+class EditOutcome(BaseModel):
+ """Result of attempting to apply one Edit.
+
+ Persisted both in the SessionStore SQLite table and as part of the
+ session.json artifact. The status literal is the canonical lifecycle for
+ each edit.
+
+ See spec §4.3.
+ """
+
+ edit_id: str
+ status: EditOutcomeStatus
+ benchmark_delta: float | None = Field(
+ default=None,
+ description="Overall score change from this edit. None if not gated.",
+ )
+ cluster_deltas: dict[str, float] = Field(
+ default_factory=dict,
+ description="Per-cluster score change from this edit.",
+ )
+ error: str | None = Field(
+ default=None,
+ description="Error message if the edit was rejected or failed.",
+ )
+ applied_at: datetime | None = Field(
+ default=None,
+ description="When the edit was committed to the checkpoint repo.",
+ )
+
+
+# ---------------------------------------------------------------------------
+# LearningSession — durable record of one full diagnose→record loop
+# ---------------------------------------------------------------------------
+
+
+class LearningSession(BaseModel):
+ """The durable record of one distillation session.
+
+ Persisted in two places: `/session.json` (authoritative) and
+ the SQLite SessionStore (queryable index). When in doubt, prefer the JSON
+ file — SQLite can be rebuilt from the JSON files.
+
+ See spec §4.3, §7.7 (status transitions), §8.
+ """
+
+ id: str
+ parent_session_id: str | None = Field(
+ default=None,
+ description="Id of the session this one is a follow-up to, if any.",
+ )
+ trigger: TriggerKind
+ trigger_metadata: dict[str, Any] = Field(default_factory=dict)
+ status: SessionStatus
+ autonomy_mode: AutonomyMode
+ started_at: datetime
+ ended_at: datetime | None = None
+ diagnosis_path: Path
+ plan_path: Path
+ benchmark_before: BenchmarkSnapshot
+ benchmark_after: BenchmarkSnapshot | None = None
+ edit_outcomes: list[EditOutcome] = Field(default_factory=list)
+ git_checkpoint_pre: str = Field(
+ ...,
+ description="Commit sha at session start (baseline commit).",
+ )
+ git_checkpoint_post: str | None = Field(
+ default=None,
+ description="Commit sha after edits applied; None until executing finishes.",
+ )
+ teacher_cost_usd: float = Field(
+ ...,
+ ge=0.0,
+ description="Accumulated teacher API spend for this session.",
+ )
+ error: str | None = Field(
+ default=None,
+ description="If status is FAILED, the error message that caused it.",
+ )
diff --git a/src/openjarvis/learning/distillation/orchestrator.py b/src/openjarvis/learning/distillation/orchestrator.py
new file mode 100644
index 00000000..09a99160
--- /dev/null
+++ b/src/openjarvis/learning/distillation/orchestrator.py
@@ -0,0 +1,418 @@
+"""DistillationOrchestrator: top-level driver for a learning session.
+
+Wires diagnose (M2) → plan (M3) → execute (M4) → gate (M5) into a
+single ``run(trigger)`` method. All dependencies are injected.
+
+See spec §3, §7.2, §7.7.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+from openjarvis.learning.distillation.diagnose.runner import DiagnosisRunner
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.execute.loop import _build_registry
+from openjarvis.learning.distillation.gate.benchmark_gate import BenchmarkGate
+from openjarvis.learning.distillation.gate.cold_start import check_readiness
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ BenchmarkSnapshot,
+ EditOutcome,
+ EditRiskTier,
+ LearningSession,
+ SessionStatus,
+)
+from openjarvis.learning.distillation.pending_queue import PendingQueue
+from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+logger = logging.getLogger(__name__)
+
+
+class DistillationOrchestrator:
+ """Top-level driver for a distillation learning session.
+
+ All dependencies are injected so tests can mock everything.
+ """
+
+ def __init__(
+ self,
+ *,
+ teacher_engine: Any,
+ teacher_model: str,
+ trace_store: Any,
+ benchmark_samples: list,
+ student_runner: Any,
+ judge: Any,
+ session_store: Any,
+ checkpoint_store: Any,
+ openjarvis_home: Path,
+ autonomy_mode: AutonomyMode = AutonomyMode.TIERED,
+ scorer: Callable[..., BenchmarkSnapshot] | None = None,
+ benchmark_version: str = "personal_v1",
+ min_traces: int = 20,
+ max_cost_usd: float = 5.0,
+ max_tool_calls: int = 30,
+ min_improvement: float = 0.0,
+ max_regression: float = 0.05,
+ subsample_size: int = 50,
+ ) -> None:
+ self._engine = teacher_engine
+ self._model = teacher_model
+ self._trace_store = trace_store
+ self._benchmark_samples = benchmark_samples
+ self._student_runner = student_runner
+ self._judge = judge
+ self._session_store = session_store
+ self._checkpoint_store = checkpoint_store
+ self._home = Path(openjarvis_home)
+ self._autonomy = autonomy_mode
+ self._scorer = scorer
+ self._bench_version = benchmark_version
+ self._min_traces = min_traces
+ self._max_cost = max_cost_usd
+ self._max_tool_calls = max_tool_calls
+ self._min_improvement = min_improvement
+ self._max_regression = max_regression
+ self._subsample_size = subsample_size
+
+ def run(self, trigger: Any) -> LearningSession:
+ """Execute a full distillation session.
+
+ Returns the completed LearningSession.
+ """
+ ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
+ session_id = f"session-{ts}_{uuid.uuid4().hex[:8]}"
+ session_dir = self._home / "learning" / "sessions" / session_id
+ session_dir.mkdir(parents=True, exist_ok=True)
+
+ # Initialize session
+ pre_sha = self._checkpoint_store.current_sha()
+ session = LearningSession(
+ id=session_id,
+ trigger=trigger.kind,
+ trigger_metadata=trigger.metadata,
+ status=SessionStatus.INITIATED,
+ autonomy_mode=self._autonomy,
+ started_at=datetime.now(timezone.utc),
+ diagnosis_path=session_dir / "diagnosis.md",
+ plan_path=session_dir / "plan.json",
+ benchmark_before=BenchmarkSnapshot(
+ benchmark_version=self._bench_version,
+ overall_score=0.0,
+ cluster_scores={},
+ task_count=0,
+ elapsed_seconds=0.0,
+ ),
+ git_checkpoint_pre=pre_sha,
+ teacher_cost_usd=0.0,
+ )
+
+ try:
+ # Cold start check
+ readiness = check_readiness(self._trace_store, min_traces=self._min_traces)
+ if not readiness.ready:
+ session = session.model_copy(
+ update={
+ "status": SessionStatus.FAILED,
+ "error": readiness.message,
+ "ended_at": datetime.now(timezone.utc),
+ }
+ )
+ self._session_store.save_session(session)
+ return session
+
+ # Capture benchmark before
+ if self._scorer is not None:
+ before_snap = self._scorer(
+ benchmark_version=self._bench_version,
+ subsample_size=self._subsample_size,
+ seed=hash(session_id) % (2**31),
+ )
+ session = session.model_copy(update={"benchmark_before": before_snap})
+
+ # Phase 1: Diagnose
+ session = session.model_copy(update={"status": SessionStatus.DIAGNOSING})
+ self._session_store.save_session(session)
+
+ diagnosis_runner = DiagnosisRunner(
+ teacher_engine=self._engine,
+ teacher_model=self._model,
+ trace_store=self._trace_store,
+ benchmark_samples=self._benchmark_samples,
+ student_runner=self._student_runner,
+ judge=self._judge,
+ session_dir=session_dir,
+ session_id=session_id,
+ config={
+ "config_path": self._home / "config.toml",
+ "openjarvis_home": self._home,
+ },
+ max_turns=self._max_tool_calls,
+ max_cost_usd=self._max_cost,
+ )
+ diag_result = diagnosis_runner.run()
+ cost = diag_result.cost_usd
+
+ if not diag_result.clusters:
+ session = session.model_copy(
+ update={
+ "status": SessionStatus.FAILED,
+ "error": "diagnosis produced no actionable clusters",
+ "teacher_cost_usd": cost,
+ "ended_at": datetime.now(timezone.utc),
+ }
+ )
+ self._session_store.save_session(session)
+ return session
+
+ # Phase 2: Plan
+ session = session.model_copy(update={"status": SessionStatus.PLANNING})
+ self._session_store.save_session(session)
+
+ planner = LearningPlanner(
+ teacher_engine=self._engine,
+ teacher_model=self._model,
+ session_id=session_id,
+ session_dir=session_dir,
+ prompt_reader=lambda t: self._read_prompt(t),
+ )
+ plan = planner.run(
+ diagnosis_md=diag_result.diagnosis_md,
+ clusters=diag_result.clusters,
+ )
+ cost += plan.estimated_cost_usd
+
+ # Phase 3: Execute
+ session = session.model_copy(update={"status": SessionStatus.EXECUTING})
+ self._session_store.save_session(session)
+
+ ctx = ApplyContext(
+ openjarvis_home=self._home,
+ session_id=session_id,
+ )
+ registry = _build_registry()
+
+ gate: BenchmarkGate | None = None
+ if self._scorer is not None:
+ gate = BenchmarkGate(
+ scorer=self._scorer,
+ benchmark_version=self._bench_version,
+ min_improvement=self._min_improvement,
+ max_regression=self._max_regression,
+ subsample_size=self._subsample_size,
+ )
+
+ session_seed = hash(session_id) % (2**31)
+ outcomes: list[EditOutcome] = []
+
+ for edit in plan.edits:
+ # Manual autonomy mode: everything goes to review
+ if self._autonomy == AutonomyMode.MANUAL:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="pending_review",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Manual risk tier: always skip
+ if edit.risk_tier == EditRiskTier.MANUAL:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="skipped",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error="manual tier, requires explicit approval",
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Review tier in tiered mode: route to pending
+ if (
+ edit.risk_tier == EditRiskTier.REVIEW
+ and self._autonomy == AutonomyMode.TIERED
+ ):
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="pending_review",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Check if the op is supported
+ if not registry.is_supported(edit.op):
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="skipped",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=f"op {edit.op.value} not implemented in v1",
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Validate
+ applier = registry.get(edit.op)
+ validation = applier.validate(edit, ctx)
+ if not validation.ok:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="rejected_by_gate",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=validation.reason,
+ applied_at=None,
+ )
+ )
+ continue
+
+ # Apply the edit
+ try:
+ applier.apply(edit, ctx)
+ except Exception as exc:
+ logger.warning("Edit %s apply failed: %s", edit.id, exc)
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="rejected_by_gate",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=str(exc),
+ applied_at=None,
+ )
+ )
+ continue
+
+ # If no scorer, accept directly (backward-compat with tests)
+ if gate is None:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="applied",
+ benchmark_delta=None,
+ cluster_deltas={},
+ error=None,
+ applied_at=datetime.now(timezone.utc),
+ )
+ )
+ continue
+
+ # Run the benchmark gate
+ before_snap = session.benchmark_before
+ gate_result = gate.evaluate(
+ before=before_snap,
+ session_seed=session_seed,
+ )
+
+ if gate_result.accepted:
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="applied",
+ benchmark_delta=gate_result.delta,
+ cluster_deltas={},
+ error=None,
+ applied_at=datetime.now(timezone.utc),
+ )
+ )
+ else:
+ # Gate rejected — rollback the edit
+ try:
+ applier.rollback(edit, ctx)
+ except Exception as rb_exc:
+ logger.warning(
+ "Edit %s rollback failed: %s", edit.id, rb_exc
+ )
+ outcomes.append(
+ EditOutcome(
+ edit_id=edit.id,
+ status="rejected_by_gate",
+ benchmark_delta=gate_result.delta,
+ cluster_deltas={},
+ error=gate_result.reason,
+ applied_at=None,
+ )
+ )
+
+ # Enqueue pending_review edits
+ pending_queue = PendingQueue(self._home / "learning" / "pending_review")
+ has_pending = False
+ for outcome, edit in zip(outcomes, plan.edits):
+ if outcome.status == "pending_review":
+ pending_queue.enqueue(session_id, edit)
+ has_pending = True
+
+ # Capture benchmark after
+ after_snap = None
+ if self._scorer is not None:
+ after_snap = self._scorer(
+ benchmark_version=self._bench_version,
+ subsample_size=self._subsample_size,
+ seed=session_seed,
+ )
+
+ # Determine final status
+ if has_pending:
+ final_status = SessionStatus.AWAITING_REVIEW
+ else:
+ final_status = SessionStatus.COMPLETED
+
+ post_sha = self._checkpoint_store.current_sha()
+ session = session.model_copy(
+ update={
+ "status": final_status,
+ "edit_outcomes": outcomes,
+ "benchmark_after": after_snap,
+ "git_checkpoint_post": post_sha,
+ "teacher_cost_usd": cost,
+ "ended_at": datetime.now(timezone.utc),
+ }
+ )
+
+ except Exception as e:
+ logger.exception("Session %s failed: %s", session_id, e)
+ session = session.model_copy(
+ update={
+ "status": SessionStatus.FAILED,
+ "error": str(e),
+ "ended_at": datetime.now(timezone.utc),
+ }
+ )
+
+ self._session_store.save_session(session)
+
+ # Write session.json artifact
+ artifact_path = session_dir / "session.json"
+ artifact_path.write_text(session.model_dump_json(indent=2), encoding="utf-8")
+
+ return session
+
+ def _read_prompt(self, target: str) -> str:
+ """Read a prompt file from the config tree."""
+ parts = target.split(".")
+ if len(parts) >= 2:
+ agent_name = parts[1]
+ path = self._home / "agents" / agent_name / "system_prompt.md"
+ if path.exists():
+ return path.read_text(encoding="utf-8")
+ return ""
diff --git a/src/openjarvis/learning/distillation/pending_queue.py b/src/openjarvis/learning/distillation/pending_queue.py
new file mode 100644
index 00000000..5bac6c9b
--- /dev/null
+++ b/src/openjarvis/learning/distillation/pending_queue.py
@@ -0,0 +1,74 @@
+"""Pending review queue for edits awaiting user approval.
+
+Edits in the ``review`` tier (when autonomy mode is ``tiered``) are
+written here as JSON files. The user reviews them via ``jarvis learning
+review`` and approves or rejects.
+
+See spec §7.5.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from openjarvis.learning.distillation.models import Edit
+
+logger = logging.getLogger(__name__)
+
+
+class PendingQueue:
+ """File-based queue for pending review edits.
+
+ Each edit is stored as ``/__.json``.
+ """
+
+ def __init__(self, queue_dir: Path) -> None:
+ self._dir = Path(queue_dir)
+ self._dir.mkdir(parents=True, exist_ok=True)
+
+ def enqueue(self, session_id: str, edit: Edit) -> Path:
+ """Write an edit to the pending queue. Returns the file path."""
+ filename = f"{session_id}__{edit.id}.json"
+ path = self._dir / filename
+ data = {
+ "session_id": session_id,
+ "edit": json.loads(edit.model_dump_json()),
+ }
+ path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+ logger.info("Enqueued edit %s for review", edit.id)
+ return path
+
+ def list_pending(self) -> list[dict[str, Any]]:
+ """Return all pending edits as dicts."""
+ results = []
+ for path in sorted(self._dir.glob("*.json")):
+ try:
+ data = json.loads(path.read_text(encoding="utf-8"))
+ results.append(data)
+ except (json.JSONDecodeError, OSError):
+ logger.warning("Skipping corrupt pending file: %s", path)
+ return results
+
+ def get(self, session_id: str, edit_id: str) -> dict[str, Any] | None:
+ """Return a specific pending edit, or None."""
+ filename = f"{session_id}__{edit_id}.json"
+ path = self._dir / filename
+ if not path.exists():
+ return None
+ try:
+ return json.loads(path.read_text(encoding="utf-8"))
+ except (json.JSONDecodeError, OSError):
+ return None
+
+ def resolve(self, session_id: str, edit_id: str) -> bool:
+ """Remove a pending edit (approved or rejected). Returns True if found."""
+ filename = f"{session_id}__{edit_id}.json"
+ path = self._dir / filename
+ if path.exists():
+ path.unlink()
+ logger.info("Resolved pending edit %s", edit_id)
+ return True
+ return False
diff --git a/src/openjarvis/learning/distillation/plan/__init__.py b/src/openjarvis/learning/distillation/plan/__init__.py
new file mode 100644
index 00000000..7115eefb
--- /dev/null
+++ b/src/openjarvis/learning/distillation/plan/__init__.py
@@ -0,0 +1 @@
+"""Plan phase: converts diagnosis into a frozen LearningPlan."""
diff --git a/src/openjarvis/learning/distillation/plan/planner.py b/src/openjarvis/learning/distillation/plan/planner.py
new file mode 100644
index 00000000..6cf964dd
--- /dev/null
+++ b/src/openjarvis/learning/distillation/plan/planner.py
@@ -0,0 +1,281 @@
+"""LearningPlanner: converts a diagnosis into a frozen LearningPlan.
+
+Makes a single structured-output teacher call (no tools, no multi-turn)
+to generate typed edits for each failure cluster. Post-processes the
+edits with risk tier assignment and patch/replace downgrade.
+
+See spec §6.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+from openjarvis.learning.distillation.models import (
+ Edit,
+ FailureCluster,
+ LearningPlan,
+)
+from openjarvis.learning.distillation.plan.prompt_diff import (
+ maybe_downgrade_to_replace,
+)
+from openjarvis.learning.distillation.plan.risk_tier import assign_tiers
+
+logger = logging.getLogger(__name__)
+
+_PLANNER_SYSTEM_PROMPT = """\
+You are a meta-engineer planning improvements to a local AI assistant called \
+OpenJarvis. You have been given a diagnosis of the student's failure patterns.
+
+Your job: for each surviving failure cluster, propose 1-3 edits from the \
+available operation set that would address the cluster's skill gap.
+
+IMPORTANT: Each edit's payload must EXACTLY match the schema below. \
+Edits with missing required payload keys will be rejected.
+
+Available operations with their EXACT payload schemas:
+
+INTELLIGENCE pillar:
+- set_model_for_query_class: {{"query_class": "math", "model": "qwen3.5:27b"}}
+- set_model_param: {{"model": "qwen3.5:9b", "param": "temperature", "value": 0.3}}
+
+AGENT pillar:
+- replace_system_prompt: {{"new_content": "You are a helpful assistant.\\n..."}}
+- patch_system_prompt: {{"diff": "--- a/prompt.md\\n+++ b/prompt.md\\n@@ ...\\n"}}
+- set_agent_class: {{"agent": "simple", "new_class": "react"}}
+- set_agent_param: {{"agent": "native_react", "param": "max_turns", "value": 10}}
+- edit_few_shot_exemplars: {{"agent": "native_react", \
+"exemplars": [{{"input": "Q", "output": "A"}}]}}
+
+TOOLS pillar:
+- add_tool_to_agent: {{"agent": "native_react", "tool_name": "calculator"}}
+- remove_tool_from_agent: {{"agent": "native_react", "tool_name": "shell_exec"}}
+- edit_tool_description: {{"tool_name": "web_search", \
+"new_description": "Search the web for..."}}
+
+Each edit object must have ALL of these fields:
+- id (string, e.g. "edit_001")
+- pillar ("intelligence", "agent", or "tools")
+- op (one of the operation names above)
+- target (dotted path, e.g. "agents.native_react.system_prompt")
+- payload (object matching the schema above for the chosen op)
+- rationale (string explaining why)
+- expected_improvement (cluster id this addresses)
+- risk_tier ("auto" for safe changes, "review" for prompts)
+- references (list of trace ids that justify this edit)
+
+Respond with ONLY a JSON object: {{"edits": [...]}}
+"""
+
+
+def _validate_clusters(
+ clusters: list[FailureCluster],
+) -> tuple[list[FailureCluster], list[FailureCluster]]:
+ """Split clusters into surviving and dropped.
+
+ Drops clusters where both student_failure_rate and teacher_success_rate
+ are 0 (no evidence). Dropped clusters get a marker in skill_gap.
+ """
+ surviving = []
+ dropped = []
+ for cluster in clusters:
+ if cluster.student_failure_rate == 0.0 and cluster.teacher_success_rate == 0.0:
+ marked = cluster.model_copy(
+ update={
+ "skill_gap": (
+ "dropped: insufficient evidence"
+ f" (original: {cluster.skill_gap})"
+ ),
+ "addressed_by_edit_ids": [],
+ }
+ )
+ dropped.append(marked)
+ else:
+ surviving.append(cluster)
+ return surviving, dropped
+
+
+class LearningPlanner:
+ """Converts a diagnosis into a frozen LearningPlan.
+
+ Parameters
+ ----------
+ teacher_engine :
+ CloudEngine (or mock) for the planner call.
+ teacher_model :
+ Frontier model id.
+ session_id :
+ Current session id.
+ session_dir :
+ Path for persisting plan.json and teacher_traces/plan.jsonl.
+ prompt_reader :
+ Callable that takes a target string and returns the current prompt
+ content. Used by the patch/replace downgrade logic.
+ """
+
+ def __init__(
+ self,
+ *,
+ teacher_engine: Any,
+ teacher_model: str,
+ session_id: str,
+ session_dir: Path,
+ prompt_reader: Callable[[str], str],
+ ) -> None:
+ self._engine = teacher_engine
+ self._model = teacher_model
+ self._session_id = session_id
+ self._session_dir = Path(session_dir)
+ self._prompt_reader = prompt_reader
+
+ def run(
+ self,
+ *,
+ diagnosis_md: str,
+ clusters: list[FailureCluster],
+ ) -> LearningPlan:
+ """Execute the plan phase.
+
+ Parameters
+ ----------
+ diagnosis_md :
+ The teacher's diagnosis markdown from phase 1.
+ clusters :
+ Failure clusters from phase 1.
+
+ Returns
+ -------
+ LearningPlan
+ The frozen plan, also persisted to ``plan.json``.
+ """
+ self._session_dir.mkdir(parents=True, exist_ok=True)
+
+ # Validate clusters — drop those without evidence
+ surviving, dropped = _validate_clusters(clusters)
+
+ # Build the user prompt with diagnosis and cluster info
+ cluster_json = json.dumps(
+ [c.model_dump() for c in surviving],
+ indent=2,
+ default=str,
+ )
+ user_prompt = (
+ f"## Diagnosis\n\n{diagnosis_md}\n\n"
+ f"## Surviving Failure Clusters\n\n```json\n{cluster_json}\n```\n\n"
+ "Propose edits for these clusters. Respond with ONLY JSON: "
+ '{"edits": [...]}'
+ )
+
+ # Make the teacher call
+ from openjarvis.core.types import Message, Role
+
+ messages = [
+ Message(role=Role.SYSTEM, content=_PLANNER_SYSTEM_PROMPT),
+ Message(role=Role.USER, content=user_prompt),
+ ]
+ result = self._engine.generate(
+ messages=messages,
+ model=self._model,
+ max_tokens=4096,
+ )
+
+ cost_usd = result.get("cost_usd", 0.0)
+ content = result.get("content", "")
+
+ # Persist teacher trace
+ self._persist_trace(content, cost_usd, result)
+
+ # Parse edits from response
+ edits = self._parse_edits(content)
+
+ # Post-process: assign tiers deterministically
+ edits = assign_tiers(edits)
+
+ # Post-process: downgrade large patches to replacements
+ edits = [
+ maybe_downgrade_to_replace(e, prompt_reader=self._prompt_reader)
+ for e in edits
+ ]
+
+ # Wire clusters ↔ edits
+ surviving = self._wire_cluster_edit_ids(surviving, edits)
+
+ # Build the plan
+ all_clusters = surviving + dropped
+ plan = LearningPlan(
+ session_id=self._session_id,
+ diagnosis_summary=diagnosis_md,
+ failure_clusters=all_clusters,
+ edits=edits,
+ teacher_model=self._model,
+ estimated_cost_usd=cost_usd,
+ created_at=datetime.now(timezone.utc),
+ )
+
+ # Persist plan.json
+ plan_path = self._session_dir / "plan.json"
+ plan_path.write_text(plan.model_dump_json(indent=2), encoding="utf-8")
+
+ return plan
+
+ def _parse_edits(self, content: str) -> list[Edit]:
+ """Parse edits from the teacher's JSON response."""
+ try:
+ data = json.loads(content)
+ except json.JSONDecodeError:
+ # Try to find JSON in the content
+ import re
+
+ match = re.search(r"\{[\s\S]*\}", content)
+ if match:
+ try:
+ data = json.loads(match.group(0))
+ except json.JSONDecodeError:
+ logger.warning("Could not parse edits from teacher response")
+ return []
+ else:
+ logger.warning("No JSON found in teacher response")
+ return []
+
+ raw_edits = data.get("edits", [])
+ edits = []
+ for raw in raw_edits:
+ try:
+ edits.append(Edit.model_validate(raw))
+ except Exception as e:
+ logger.warning("Skipping invalid edit: %s — %s", raw.get("id", "?"), e)
+ return edits
+
+ def _wire_cluster_edit_ids(
+ self,
+ clusters: list[FailureCluster],
+ edits: list[Edit],
+ ) -> list[FailureCluster]:
+ """Populate addressed_by_edit_ids on each cluster."""
+ cluster_map: dict[str, list[str]] = {c.id: [] for c in clusters}
+ for edit in edits:
+ if edit.expected_improvement in cluster_map:
+ cluster_map[edit.expected_improvement].append(edit.id)
+ return [
+ c.model_copy(update={"addressed_by_edit_ids": cluster_map.get(c.id, [])})
+ for c in clusters
+ ]
+
+ def _persist_trace(self, content: str, cost_usd: float, result: dict) -> None:
+ """Write the planner call to teacher_traces/plan.jsonl."""
+ traces_dir = self._session_dir / "teacher_traces"
+ traces_dir.mkdir(parents=True, exist_ok=True)
+ record = {
+ "timestamp": datetime.now(timezone.utc).isoformat(),
+ "call": "planner",
+ "content_length": len(content),
+ "cost_usd": cost_usd,
+ "tokens": result.get("usage", {}).get("total_tokens", 0),
+ }
+ jsonl_path = traces_dir / "plan.jsonl"
+ with jsonl_path.open("a", encoding="utf-8") as f:
+ f.write(json.dumps(record) + "\n")
diff --git a/src/openjarvis/learning/distillation/plan/prompt_diff.py b/src/openjarvis/learning/distillation/plan/prompt_diff.py
new file mode 100644
index 00000000..02e0e110
--- /dev/null
+++ b/src/openjarvis/learning/distillation/plan/prompt_diff.py
@@ -0,0 +1,170 @@
+"""Prompt diff utilities for the plan phase.
+
+Handles the PATCH_SYSTEM_PROMPT → REPLACE_SYSTEM_PROMPT downgrade logic.
+When the teacher proposes a PATCH edit, the planner checks if the diff
+would change more than 50% of lines. If so, it downgrades to a full
+REPLACE so the user sees the complete new prompt in the review queue.
+
+See spec §6.3.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import Callable, Optional
+
+from openjarvis.learning.distillation.models import Edit, EditOp
+
+logger = logging.getLogger(__name__)
+
+# Threshold: if more than this fraction of lines change, downgrade to REPLACE
+_DOWNGRADE_THRESHOLD = 0.5
+
+
+def changed_line_ratio(original: str, modified: str) -> float:
+ """Compute the fraction of lines that differ between two strings.
+
+ Uses a simple line-by-line comparison. Returns 0.0 if both are empty,
+ 1.0 if one is empty and the other is not.
+ """
+ orig_lines = original.splitlines()
+ mod_lines = modified.splitlines()
+
+ if not orig_lines and not mod_lines:
+ return 0.0
+ if not orig_lines or not mod_lines:
+ return 1.0
+
+ max_len = max(len(orig_lines), len(mod_lines))
+ changed = 0
+ for i in range(max_len):
+ orig = orig_lines[i] if i < len(orig_lines) else None
+ mod = mod_lines[i] if i < len(mod_lines) else None
+ if orig != mod:
+ changed += 1
+
+ return changed / max_len
+
+
+def apply_unified_diff(original: str, diff: str) -> Optional[str]:
+ """Apply a unified diff to the original string.
+
+ Returns the patched string, or None if the diff cannot be applied.
+ This is a simplified implementation that handles basic unified diffs.
+ """
+ try:
+ lines = original.splitlines(keepends=True)
+ result_lines: list[str] = []
+ diff_lines = diff.splitlines(keepends=True)
+
+ # Skip header lines (--- and +++)
+ i = 0
+ while i < len(diff_lines) and not diff_lines[i].startswith("@@"):
+ i += 1
+
+ # No hunk headers found — not a valid unified diff
+ if i >= len(diff_lines):
+ return None
+
+ # Parse hunks
+ src_idx = 0
+ while i < len(diff_lines):
+ line = diff_lines[i]
+ if line.startswith("@@"):
+ # Parse hunk header: @@ -start,count +start,count @@
+ match = re.match(r"@@ -(\d+)", line)
+ if not match:
+ return None
+ hunk_start = int(match.group(1)) - 1 # 0-indexed
+ # Copy lines before this hunk
+ while src_idx < hunk_start:
+ if src_idx < len(lines):
+ result_lines.append(lines[src_idx])
+ src_idx += 1
+ i += 1
+ continue
+
+ if line.startswith("-"):
+ # Remove line from original
+ src_idx += 1
+ elif line.startswith("+"):
+ # Add line to result
+ content = line[1:]
+ if not content.endswith("\n"):
+ content += "\n"
+ result_lines.append(content)
+ elif line.startswith(" "):
+ # Context line — copy from original
+ if src_idx < len(lines):
+ result_lines.append(lines[src_idx])
+ src_idx += 1
+ i += 1
+
+ # Copy remaining lines
+ while src_idx < len(lines):
+ result_lines.append(lines[src_idx])
+ src_idx += 1
+
+ return "".join(result_lines)
+ except Exception:
+ logger.warning("Failed to apply unified diff")
+ return None
+
+
+def maybe_downgrade_to_replace(
+ edit: Edit,
+ *,
+ prompt_reader: Callable[[str], str],
+) -> Edit:
+ """Downgrade PATCH_SYSTEM_PROMPT to REPLACE if the diff is large.
+
+ Parameters
+ ----------
+ edit :
+ The edit to check.
+ prompt_reader :
+ A callable that takes a target string (e.g. "agents.simple.system_prompt")
+ and returns the current prompt content.
+
+ Returns the edit unchanged if it's not a PATCH op, or if the diff is
+ small enough. Returns a new REPLACE edit if the diff changes > 50% of
+ lines or if the diff cannot be applied.
+ """
+ if edit.op != EditOp.PATCH_SYSTEM_PROMPT:
+ return edit
+
+ diff_str = edit.payload.get("diff", "")
+ original = prompt_reader(edit.target)
+ patched = apply_unified_diff(original, diff_str)
+
+ if patched is None:
+ # Can't apply the diff — downgrade to REPLACE with a warning
+ logger.warning(
+ "Edit %s: diff could not be applied, downgrading to REPLACE",
+ edit.id,
+ )
+ return edit.model_copy(
+ update={
+ "op": EditOp.REPLACE_SYSTEM_PROMPT,
+ "payload": {"new_content": diff_str},
+ }
+ )
+
+ ratio = changed_line_ratio(original, patched)
+ if ratio > _DOWNGRADE_THRESHOLD:
+ logger.info(
+ "Edit %s: diff changes %.0f%% of lines (>%.0f%%), "
+ "downgrading PATCH → REPLACE",
+ edit.id,
+ ratio * 100,
+ _DOWNGRADE_THRESHOLD * 100,
+ )
+ return edit.model_copy(
+ update={
+ "op": EditOp.REPLACE_SYSTEM_PROMPT,
+ "payload": {"new_content": patched},
+ }
+ )
+
+ return edit
diff --git a/src/openjarvis/learning/distillation/plan/risk_tier.py b/src/openjarvis/learning/distillation/plan/risk_tier.py
new file mode 100644
index 00000000..2ba4a2c1
--- /dev/null
+++ b/src/openjarvis/learning/distillation/plan/risk_tier.py
@@ -0,0 +1,63 @@
+"""Deterministic risk tier assignment for edits.
+
+The teacher cannot pick its own tier. After the teacher emits edits, the
+planner overwrites each edit's ``risk_tier`` from the lookup table below.
+If the teacher attempted a different tier, it is silently overwritten and
+the discrepancy is logged but not surfaced as an error.
+
+See spec §4.1 (tier table) and §6.2.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Sequence
+
+from openjarvis.learning.distillation.models import Edit, EditOp, EditRiskTier
+
+logger = logging.getLogger(__name__)
+
+# The canonical (op) → tier mapping. Every EditOp must appear here.
+TIER_TABLE: dict[EditOp, EditRiskTier] = {
+ # Intelligence — safe, reversible
+ EditOp.SET_MODEL_FOR_QUERY_CLASS: EditRiskTier.AUTO,
+ EditOp.SET_MODEL_PARAM: EditRiskTier.AUTO,
+ # Agent — params are safe, prompts and class need review
+ EditOp.PATCH_SYSTEM_PROMPT: EditRiskTier.REVIEW,
+ EditOp.REPLACE_SYSTEM_PROMPT: EditRiskTier.REVIEW,
+ EditOp.SET_AGENT_CLASS: EditRiskTier.REVIEW,
+ EditOp.SET_AGENT_PARAM: EditRiskTier.AUTO,
+ EditOp.EDIT_FEW_SHOT_EXEMPLARS: EditRiskTier.REVIEW,
+ # Tools — all safe, reversible
+ EditOp.ADD_TOOL_TO_AGENT: EditRiskTier.AUTO,
+ EditOp.REMOVE_TOOL_FROM_AGENT: EditRiskTier.AUTO,
+ EditOp.EDIT_TOOL_DESCRIPTION: EditRiskTier.AUTO,
+ # v2 — always manual
+ EditOp.LORA_FINETUNE: EditRiskTier.MANUAL,
+}
+
+
+def assign_tier(op: EditOp) -> EditRiskTier:
+ """Return the deterministic risk tier for a given edit op."""
+ return TIER_TABLE[op]
+
+
+def assign_tiers(edits: Sequence[Edit]) -> list[Edit]:
+ """Overwrite each edit's risk_tier from the canonical lookup table.
+
+ Returns a new list of Edit objects (pydantic copies). If the teacher
+ had a different tier, it is silently overwritten and logged.
+ """
+ result = []
+ for edit in edits:
+ correct_tier = assign_tier(edit.op)
+ if edit.risk_tier != correct_tier:
+ logger.info(
+ "Edit %s: overwriting teacher tier %s → %s (op=%s)",
+ edit.id,
+ edit.risk_tier.value,
+ correct_tier.value,
+ edit.op.value,
+ )
+ result.append(edit.model_copy(update={"risk_tier": correct_tier}))
+ return result
diff --git a/src/openjarvis/learning/distillation/storage/__init__.py b/src/openjarvis/learning/distillation/storage/__init__.py
new file mode 100644
index 00000000..db497278
--- /dev/null
+++ b/src/openjarvis/learning/distillation/storage/__init__.py
@@ -0,0 +1 @@
+"""Storage primitives: paths, SQLite session store."""
diff --git a/src/openjarvis/learning/distillation/storage/paths.py b/src/openjarvis/learning/distillation/storage/paths.py
new file mode 100644
index 00000000..865800f8
--- /dev/null
+++ b/src/openjarvis/learning/distillation/storage/paths.py
@@ -0,0 +1,90 @@
+"""Filesystem path resolution for the distillation subsystem.
+
+The keystone of artifact isolation (spec §11): the resolved distillation root
+must NEVER be inside the OpenJarvis source tree. ``resolve_distillation_root``
+walks up from this module's ``__file__`` looking for a ``pyproject.toml`` that
+identifies the OpenJarvis source root, then refuses to operate if the resolved
+root is inside it. Defense in depth — if a user accidentally points
+``OPENJARVIS_HOME`` at the repo, the system fails loudly instead of silently
+writing artifacts into the working tree.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from openjarvis.security.file_utils import secure_mkdir
+
+
+class ConfigurationError(RuntimeError):
+ """Raised when path configuration would violate isolation guarantees."""
+
+
+def _find_source_root() -> Path | None:
+ """Walk upward from this module to find the OpenJarvis source root.
+
+ Returns the directory containing the OpenJarvis ``pyproject.toml``, or
+ ``None`` if no such file is found (e.g. when running from an installed
+ wheel rather than a source checkout).
+ """
+ here = Path(__file__).resolve()
+ for candidate in (here, *here.parents):
+ py = candidate / "pyproject.toml"
+ if py.exists():
+ try:
+ content = py.read_text(encoding="utf-8")
+ except OSError:
+ continue
+ if 'name = "openjarvis"' in content.lower():
+ return candidate
+ return None
+
+
+def _resolve_openjarvis_home() -> Path:
+ """Resolve the OPENJARVIS_HOME directory (env var or default)."""
+ env = os.environ.get("OPENJARVIS_HOME")
+ if env:
+ return Path(env).expanduser().resolve()
+ return (Path.home() / ".openjarvis").resolve()
+
+
+def resolve_distillation_root() -> Path:
+ """Return the absolute path of the distillation root directory.
+
+ The root is ``$OPENJARVIS_HOME/learning`` (or ``~/.openjarvis/learning``
+ by default). Raises ``ConfigurationError`` if the resolved path lies
+ inside the OpenJarvis source tree, to prevent dev artifacts from leaking
+ into the repo.
+ """
+ home = _resolve_openjarvis_home()
+ source_root = _find_source_root()
+ if source_root is not None:
+ try:
+ home.relative_to(source_root)
+ except ValueError:
+ pass # Good — not inside the source tree.
+ else:
+ raise ConfigurationError(
+ f"OPENJARVIS_HOME ({home}) is inside the source tree "
+ f"({source_root}). Distillation refuses to write runtime "
+ "artifacts inside the OpenJarvis repo. Set OPENJARVIS_HOME "
+ "to a directory outside the repo (default: ~/.openjarvis)."
+ )
+ return home / "learning"
+
+
+def ensure_distillation_dirs() -> Path:
+ """Create the distillation directory layout if missing.
+
+ Returns the distillation root. Creates ``sessions/``, ``benchmarks/``,
+ ``benchmarks/reference_outputs/``, and ``pending_review/`` underneath it,
+ all with restrictive ``0o700`` permissions via ``secure_mkdir``.
+ """
+ root = resolve_distillation_root()
+ secure_mkdir(root)
+ secure_mkdir(root / "sessions")
+ secure_mkdir(root / "benchmarks")
+ secure_mkdir(root / "benchmarks" / "reference_outputs")
+ secure_mkdir(root / "pending_review")
+ return root
diff --git a/src/openjarvis/learning/distillation/storage/session_store.py b/src/openjarvis/learning/distillation/storage/session_store.py
new file mode 100644
index 00000000..bfae741d
--- /dev/null
+++ b/src/openjarvis/learning/distillation/storage/session_store.py
@@ -0,0 +1,304 @@
+"""SQLite-backed storage for distillation LearningSession records.
+
+Mirrors the style of ``openjarvis.learning.optimize.store.OptimizationStore``:
+
+- stdlib ``sqlite3`` in WAL mode
+- inline DDL as module-level constants
+- persistent connection stored as ``self._conn``
+- ``_migrate()`` runs additive ALTER TABLEs that swallow ``OperationalError``
+
+The store does NOT share its database file with ``OptimizationStore`` (see
+spec §8.1 / brainstorming Q8). The two SQLite files live side-by-side in
+``~/.openjarvis/learning/`` but are independent.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import sqlite3
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Union
+
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ BenchmarkSnapshot,
+ EditOutcome,
+ LearningSession,
+ SessionStatus,
+ TriggerKind,
+)
+
+logger = logging.getLogger(__name__)
+
+
+_CREATE_SESSIONS = """\
+CREATE TABLE IF NOT EXISTS learning_sessions (
+ id TEXT PRIMARY KEY,
+ parent_session_id TEXT,
+ trigger TEXT NOT NULL,
+ trigger_metadata TEXT NOT NULL DEFAULT '{}',
+ status TEXT NOT NULL,
+ autonomy_mode TEXT NOT NULL,
+ started_at TEXT NOT NULL,
+ ended_at TEXT,
+ diagnosis_path TEXT NOT NULL,
+ plan_path TEXT NOT NULL,
+ benchmark_before TEXT NOT NULL,
+ benchmark_after TEXT,
+ git_checkpoint_pre TEXT NOT NULL,
+ git_checkpoint_post TEXT,
+ teacher_cost_usd REAL NOT NULL DEFAULT 0.0,
+ error TEXT,
+ FOREIGN KEY (parent_session_id) REFERENCES learning_sessions(id)
+);
+"""
+
+_CREATE_OUTCOMES = """\
+CREATE TABLE IF NOT EXISTS edit_outcomes (
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
+ session_id TEXT NOT NULL,
+ edit_id TEXT NOT NULL,
+ pillar TEXT NOT NULL,
+ op TEXT NOT NULL,
+ target TEXT NOT NULL,
+ risk_tier TEXT NOT NULL,
+ status TEXT NOT NULL,
+ benchmark_delta REAL,
+ cluster_deltas TEXT NOT NULL DEFAULT '{}',
+ rationale TEXT NOT NULL DEFAULT '',
+ error TEXT,
+ applied_at TEXT,
+ FOREIGN KEY (session_id) REFERENCES learning_sessions(id)
+);
+"""
+
+_CREATE_INDEXES = [
+ "CREATE INDEX IF NOT EXISTS idx_sessions_started_at "
+ "ON learning_sessions(started_at)",
+ "CREATE INDEX IF NOT EXISTS idx_sessions_status ON learning_sessions(status)",
+ "CREATE INDEX IF NOT EXISTS idx_outcomes_session ON edit_outcomes(session_id)",
+ "CREATE INDEX IF NOT EXISTS idx_outcomes_op ON edit_outcomes(op)",
+]
+
+_INSERT_SESSION = """\
+INSERT OR REPLACE INTO learning_sessions (
+ id, parent_session_id, trigger, trigger_metadata, status, autonomy_mode,
+ started_at, ended_at, diagnosis_path, plan_path,
+ benchmark_before, benchmark_after,
+ git_checkpoint_pre, git_checkpoint_post, teacher_cost_usd, error
+) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+"""
+
+_INSERT_OUTCOME = """\
+INSERT INTO edit_outcomes (
+ session_id, edit_id, pillar, op, target, risk_tier, status,
+ benchmark_delta, cluster_deltas, rationale, error, applied_at
+) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+"""
+
+# Future ALTER TABLE statements go here, swallowing OperationalError.
+_MIGRATE: list[str] = []
+
+
+def _dt_to_iso(dt: datetime | None) -> str | None:
+ return dt.isoformat() if dt is not None else None
+
+
+def _iso_to_dt(s: str | None) -> datetime | None:
+ return datetime.fromisoformat(s) if s else None
+
+
+class SessionStore:
+ """SQLite-backed storage for LearningSession and EditOutcome records.
+
+ The full LearningSession is also serialized to disk as
+ ``/session.json`` — that file is the authoritative source
+ if SQLite is ever lost. This store is the index used for fast queries.
+ """
+
+ def __init__(self, db_path: Union[str, Path]) -> None:
+ self._db_path = str(db_path)
+ self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
+ self._conn.execute("PRAGMA journal_mode=WAL")
+ self._conn.execute("PRAGMA foreign_keys=ON")
+ self._conn.execute(_CREATE_SESSIONS)
+ self._conn.execute(_CREATE_OUTCOMES)
+ for index in _CREATE_INDEXES:
+ self._conn.execute(index)
+ self._conn.commit()
+ self._migrate()
+
+ def _migrate(self) -> None:
+ """Apply additive schema migrations, swallowing already-applied ones."""
+ for stmt in _MIGRATE:
+ try:
+ self._conn.execute(stmt)
+ except sqlite3.OperationalError:
+ pass
+ self._conn.commit()
+
+ def close(self) -> None:
+ """Close the underlying SQLite connection."""
+ self._conn.close()
+
+ # ------------------------------------------------------------------
+ # Sessions
+ # ------------------------------------------------------------------
+
+ def save_session(self, session: LearningSession) -> None:
+ """Insert or update a LearningSession (idempotent on session.id)."""
+ self._conn.execute(
+ _INSERT_SESSION,
+ (
+ session.id,
+ session.parent_session_id,
+ session.trigger.value,
+ json.dumps(session.trigger_metadata),
+ session.status.value,
+ session.autonomy_mode.value,
+ session.started_at.isoformat(),
+ _dt_to_iso(session.ended_at),
+ str(session.diagnosis_path),
+ str(session.plan_path),
+ session.benchmark_before.model_dump_json(),
+ (
+ session.benchmark_after.model_dump_json()
+ if session.benchmark_after is not None
+ else None
+ ),
+ session.git_checkpoint_pre,
+ session.git_checkpoint_post,
+ session.teacher_cost_usd,
+ session.error,
+ ),
+ )
+ self._conn.commit()
+
+ def get_session(self, session_id: str) -> Optional[LearningSession]:
+ """Return the LearningSession with the given id, or None if missing."""
+ row = self._conn.execute(
+ "SELECT id, parent_session_id, trigger, trigger_metadata, status, "
+ "autonomy_mode, started_at, ended_at, diagnosis_path, plan_path, "
+ "benchmark_before, benchmark_after, git_checkpoint_pre, "
+ "git_checkpoint_post, teacher_cost_usd, error "
+ "FROM learning_sessions WHERE id = ?",
+ (session_id,),
+ ).fetchone()
+ if row is None:
+ return None
+ return self._row_to_session(row, with_outcomes=True)
+
+ def list_sessions(
+ self,
+ status: SessionStatus | None = None,
+ limit: int | None = None,
+ ) -> list[LearningSession]:
+ """List sessions ordered by ``started_at DESC``."""
+ sql = (
+ "SELECT id, parent_session_id, trigger, trigger_metadata, status, "
+ "autonomy_mode, started_at, ended_at, diagnosis_path, plan_path, "
+ "benchmark_before, benchmark_after, git_checkpoint_pre, "
+ "git_checkpoint_post, teacher_cost_usd, error "
+ "FROM learning_sessions"
+ )
+ params: list[object] = []
+ if status is not None:
+ sql += " WHERE status = ?"
+ params.append(status.value)
+ sql += " ORDER BY started_at DESC"
+ if limit is not None:
+ sql += " LIMIT ?"
+ params.append(limit)
+ rows = self._conn.execute(sql, params).fetchall()
+ return [self._row_to_session(r, with_outcomes=True) for r in rows]
+
+ def _row_to_session(self, row: tuple, with_outcomes: bool) -> LearningSession:
+ session = LearningSession(
+ id=row[0],
+ parent_session_id=row[1],
+ trigger=TriggerKind(row[2]),
+ trigger_metadata=json.loads(row[3]),
+ status=SessionStatus(row[4]),
+ autonomy_mode=AutonomyMode(row[5]),
+ started_at=datetime.fromisoformat(row[6]),
+ ended_at=_iso_to_dt(row[7]),
+ diagnosis_path=Path(row[8]),
+ plan_path=Path(row[9]),
+ benchmark_before=BenchmarkSnapshot.model_validate_json(row[10]),
+ benchmark_after=(
+ BenchmarkSnapshot.model_validate_json(row[11])
+ if row[11] is not None
+ else None
+ ),
+ edit_outcomes=[],
+ git_checkpoint_pre=row[12],
+ git_checkpoint_post=row[13],
+ teacher_cost_usd=row[14],
+ error=row[15],
+ )
+ if with_outcomes:
+ outcomes = self.list_outcomes(session.id)
+ session = session.model_copy(update={"edit_outcomes": outcomes})
+ return session
+
+ # ------------------------------------------------------------------
+ # Edit outcomes
+ # ------------------------------------------------------------------
+
+ def save_outcome(
+ self,
+ session_id: str,
+ outcome: EditOutcome,
+ *,
+ pillar: str,
+ op: str,
+ target: str,
+ risk_tier: str,
+ rationale: str = "",
+ ) -> None:
+ """Insert an EditOutcome row.
+
+ ``pillar``, ``op``, ``target``, ``risk_tier``, and ``rationale`` come
+ from the parent ``Edit`` (which is not stored on the EditOutcome
+ model). They are kept as columns to make ``WHERE op = ?`` queries
+ possible without joining against the on-disk plan.json.
+ """
+ self._conn.execute(
+ _INSERT_OUTCOME,
+ (
+ session_id,
+ outcome.edit_id,
+ pillar,
+ op,
+ target,
+ risk_tier,
+ outcome.status,
+ outcome.benchmark_delta,
+ json.dumps(outcome.cluster_deltas),
+ rationale,
+ outcome.error,
+ _dt_to_iso(outcome.applied_at),
+ ),
+ )
+ self._conn.commit()
+
+ def list_outcomes(self, session_id: str) -> list[EditOutcome]:
+ """Return all EditOutcomes for a session, ordered by insertion id."""
+ rows = self._conn.execute(
+ "SELECT edit_id, status, benchmark_delta, cluster_deltas, error, "
+ "applied_at FROM edit_outcomes WHERE session_id = ? ORDER BY id",
+ (session_id,),
+ ).fetchall()
+ return [
+ EditOutcome(
+ edit_id=row[0],
+ status=row[1],
+ benchmark_delta=row[2],
+ cluster_deltas=json.loads(row[3]),
+ error=row[4],
+ applied_at=_iso_to_dt(row[5]),
+ )
+ for row in rows
+ ]
diff --git a/src/openjarvis/learning/distillation/student_runner.py b/src/openjarvis/learning/distillation/student_runner.py
new file mode 100644
index 00000000..028be299
--- /dev/null
+++ b/src/openjarvis/learning/distillation/student_runner.py
@@ -0,0 +1,135 @@
+"""Real student runner for distillation experiments.
+
+Replaces the ``MagicMock()`` in the experiment runner script with a
+callable that actually invokes the student model via vLLM (or any
+OpenAI-compatible engine) and returns structured results.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class StudentResult:
+ """Result from running the student model on a task."""
+
+ content: str
+ score: float = 0.0
+ trace_id: str = ""
+ latency_seconds: float = 0.0
+ tokens_used: int = 0
+
+
+class VLLMStudentRunner:
+ """Invoke the student model via a vLLM OpenAI-compatible endpoint.
+
+ Parameters
+ ----------
+ host :
+ vLLM server URL (e.g. ``http://localhost:8001``).
+ model :
+ Model name as registered in vLLM (e.g. ``Qwen/Qwen3.5-9B``).
+ temperature :
+ Sampling temperature.
+ max_tokens :
+ Max tokens for the student response.
+ """
+
+ def __init__(
+ self,
+ host: str = "http://localhost:8001",
+ model: str = "Qwen/Qwen3.5-9B",
+ temperature: float = 0.6,
+ max_tokens: int = 4096,
+ ) -> None:
+ import httpx
+
+ self._host = host.rstrip("/")
+ self._model = model
+ self._temperature = temperature
+ self._max_tokens = max_tokens
+ self._client = httpx.Client(base_url=self._host, timeout=300.0)
+
+ def __call__(
+ self, query: str, session_id: str = "", **kwargs: Any
+ ) -> StudentResult:
+ """Run the student on *query* and return a StudentResult."""
+ t0 = time.time()
+ try:
+ resp = self._client.post(
+ "/v1/chat/completions",
+ json={
+ "model": self._model,
+ "messages": [{"role": "user", "content": query}],
+ "temperature": self._temperature,
+ "max_tokens": self._max_tokens,
+ "stream": False,
+ },
+ )
+ resp.raise_for_status()
+ data = resp.json()
+ except Exception as exc:
+ latency = time.time() - t0
+ logger.warning("Student runner failed: %s", exc)
+ return StudentResult(
+ content=f"Error: {exc}",
+ latency_seconds=latency,
+ )
+
+ latency = time.time() - t0
+ choices = data.get("choices", [])
+ content = ""
+ if choices:
+ msg = choices[0].get("message", {})
+ content = msg.get("content", "")
+
+ usage = data.get("usage", {})
+ tokens = usage.get("total_tokens", 0)
+
+ return StudentResult(
+ content=content,
+ latency_seconds=latency,
+ tokens_used=tokens,
+ trace_id=f"distill_{session_id}_{hash(query) % 10000}",
+ )
+
+
+def build_benchmark_samples_from_traces(
+ trace_store: Any,
+ *,
+ limit: int = 50,
+ min_feedback: float | None = None,
+) -> list:
+ """Build PersonalBenchmarkSample objects from the trace store.
+
+ Pulls recent traces (optionally filtered by feedback score) and
+ converts them into benchmark samples the teacher can reference.
+ """
+ from openjarvis.learning.optimize.personal.synthesizer import (
+ PersonalBenchmarkSample,
+ )
+
+ traces = trace_store.list_traces(limit=limit)
+ samples = []
+ for t in traces:
+ fb = getattr(t, "feedback", None)
+ if min_feedback is not None and (fb is None or fb < min_feedback):
+ continue
+ samples.append(
+ PersonalBenchmarkSample(
+ trace_id=t.trace_id,
+ query=t.query,
+ reference_answer=t.result[:2000] if t.result else "",
+ agent=t.agent,
+ category="benchmark",
+ feedback_score=fb if fb is not None else 0.0,
+ )
+ )
+ logger.info("Built %d benchmark samples from traces", len(samples))
+ return samples
diff --git a/src/openjarvis/learning/distillation/triggers.py b/src/openjarvis/learning/distillation/triggers.py
new file mode 100644
index 00000000..02d71fad
--- /dev/null
+++ b/src/openjarvis/learning/distillation/triggers.py
@@ -0,0 +1,66 @@
+"""Trigger types for the distillation subsystem.
+
+A trigger is what kicks off a learning session. Four trigger types exist,
+all funneling into ``DistillationOrchestrator.run(trigger)``. The trigger
+object is stored on the ``LearningSession`` for queryability.
+
+See spec §3.3.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from openjarvis.learning.distillation.models import TriggerKind
+
+
+@dataclass
+class OnDemandTrigger:
+ """User ran ``jarvis learning run`` from the CLI."""
+
+ kind: TriggerKind = TriggerKind.ON_DEMAND
+ metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class UserFlagTrigger:
+ """User flagged a specific trace for improvement."""
+
+ trace_id: str = ""
+ kind: TriggerKind = TriggerKind.USER_FLAG
+
+ @property
+ def metadata(self) -> dict[str, Any]:
+ return {"trace_id": self.trace_id}
+
+
+@dataclass
+class ScheduledTrigger:
+ """Cron-based scheduled trigger."""
+
+ cron: str = "0 3 * * *"
+ new_trace_count: int = 0
+ kind: TriggerKind = TriggerKind.SCHEDULED
+
+ @property
+ def metadata(self) -> dict[str, Any]:
+ return {"cron": self.cron, "new_trace_count": self.new_trace_count}
+
+
+@dataclass
+class ClusterTrigger:
+ """Fired when a failure cluster exceeds a threshold."""
+
+ cluster_description: str = ""
+ trace_ids: list[str] = field(default_factory=list)
+ failure_rate: float = 0.0
+ kind: TriggerKind = TriggerKind.CLUSTER
+
+ @property
+ def metadata(self) -> dict[str, Any]:
+ return {
+ "cluster_description": self.cluster_description,
+ "trace_ids": self.trace_ids,
+ "failure_rate": self.failure_rate,
+ }
diff --git a/src/openjarvis/tools/_stubs.py b/src/openjarvis/tools/_stubs.py
index 06b7f129..e2a722b3 100644
--- a/src/openjarvis/tools/_stubs.py
+++ b/src/openjarvis/tools/_stubs.py
@@ -64,12 +64,17 @@ class BaseTool(ABC):
def to_openai_function(self) -> Dict[str, Any]:
"""Convert to OpenAI function-calling format."""
+ from openjarvis.tools.description_loader import (
+ get_tool_description_override,
+ )
+
s = self.spec
+ desc = get_tool_description_override(s.name) or s.description
return {
"type": "function",
"function": {
"name": s.name,
- "description": s.description,
+ "description": desc,
"parameters": s.parameters,
},
}
@@ -356,10 +361,15 @@ def build_tool_descriptions(
if not tools:
return "No tools available."
+ from openjarvis.tools.description_loader import (
+ get_tool_description_override,
+ )
+
sections: list[str] = []
for t in tools:
s = t.spec
- lines = [f"### {s.name}", s.description]
+ desc = get_tool_description_override(s.name) or s.description
+ lines = [f"### {s.name}", desc]
if include_category and s.category:
lines.append(f"Category: {s.category}")
diff --git a/src/openjarvis/tools/description_loader.py b/src/openjarvis/tools/description_loader.py
new file mode 100644
index 00000000..b0f76ea1
--- /dev/null
+++ b/src/openjarvis/tools/description_loader.py
@@ -0,0 +1,79 @@
+"""Load tool description overrides from $OPENJARVIS_HOME/tools/descriptions.toml.
+
+Distillation (M1) proposes tool description edits that get written to disk by
+``EditToolDescriptionApplier``. This module loads those overrides so agents
+see the improved descriptions at runtime.
+
+The TOML file format (written by the applier) is::
+
+ [web_search]
+ description = "Search the web for recent information only"
+
+ [llm]
+ description = "Call a sub-LM. Has no internet access."
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+from typing import Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+_cache: Optional[Dict[str, str]] = None
+
+
+def _load_overrides() -> Dict[str, str]:
+ """Parse descriptions.toml and return {tool_name: description}."""
+ home = Path(os.environ.get("OPENJARVIS_HOME", "~/.openjarvis")).expanduser()
+ desc_path = home / "tools" / "descriptions.toml"
+ if not desc_path.exists():
+ return {}
+ try:
+ content = desc_path.read_text(encoding="utf-8")
+ except Exception:
+ logger.warning(
+ "Failed to read tool description overrides at %s",
+ desc_path,
+ exc_info=True,
+ )
+ return {}
+
+ overrides: Dict[str, str] = {}
+ current_tool: Optional[str] = None
+ for line in content.splitlines():
+ stripped = line.strip()
+ if stripped.startswith("[") and stripped.endswith("]"):
+ current_tool = stripped[1:-1]
+ elif current_tool and stripped.startswith("description"):
+ # Parse: description = "..."
+ _, _, value = stripped.partition("=")
+ value = value.strip().strip('"').strip("'")
+ if value:
+ overrides[current_tool] = value
+ if overrides:
+ logger.info(
+ "Loaded %d tool description overrides from %s",
+ len(overrides),
+ desc_path,
+ )
+ return overrides
+
+
+def get_tool_description_override(tool_name: str) -> Optional[str]:
+ """Return the override description for *tool_name*, or ``None``.
+
+ Results are cached for the lifetime of the process.
+ """
+ global _cache # noqa: PLW0603
+ if _cache is None:
+ _cache = _load_overrides()
+ return _cache.get(tool_name)
+
+
+def clear_cache() -> None:
+ """Clear the cached overrides (useful for testing)."""
+ global _cache # noqa: PLW0603
+ _cache = None
diff --git a/tests/evals/test_benchmark_datasets.py b/tests/evals/test_benchmark_datasets.py
index 5d9f74f8..7678a57f 100644
--- a/tests/evals/test_benchmark_datasets.py
+++ b/tests/evals/test_benchmark_datasets.py
@@ -355,7 +355,7 @@ class TestConfigBenchmarks:
def test_benchmarks_count(self) -> None:
from openjarvis.evals.core.config import KNOWN_BENCHMARKS
- assert len(KNOWN_BENCHMARKS) == 30
+ assert len(KNOWN_BENCHMARKS) == 31
# ---------------------------------------------------------------------------
diff --git a/tests/learning/distillation/__init__.py b/tests/learning/distillation/__init__.py
new file mode 100644
index 00000000..f7467079
--- /dev/null
+++ b/tests/learning/distillation/__init__.py
@@ -0,0 +1 @@
+"""Tests for the distillation subsystem."""
diff --git a/tests/learning/distillation/test_applier_base.py b/tests/learning/distillation/test_applier_base.py
new file mode 100644
index 00000000..a50d41ae
--- /dev/null
+++ b/tests/learning/distillation/test_applier_base.py
@@ -0,0 +1,105 @@
+"""Tests for openjarvis.learning.distillation.execute.base module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+
+class TestApplyContext:
+ """Tests for ApplyContext dataclass."""
+
+ def test_constructs(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.base import ApplyContext
+
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ assert ctx.openjarvis_home == tmp_path
+ assert ctx.session_id == "s1"
+
+ def test_config_path(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.base import ApplyContext
+
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ assert ctx.config_path == tmp_path / "config.toml"
+
+ def test_agents_dir(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.base import ApplyContext
+
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ assert ctx.agents_dir == tmp_path / "agents"
+
+ def test_tools_dir(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.base import ApplyContext
+
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ assert ctx.tools_dir == tmp_path / "tools"
+
+
+class TestValidationResult:
+ """Tests for ValidationResult."""
+
+ def test_ok_result(self) -> None:
+ from openjarvis.learning.distillation.execute.base import ValidationResult
+
+ r = ValidationResult(ok=True)
+ assert r.ok is True
+ assert r.reason == ""
+
+ def test_error_result(self) -> None:
+ from openjarvis.learning.distillation.execute.base import ValidationResult
+
+ r = ValidationResult(ok=False, reason="target not found")
+ assert r.ok is False
+ assert r.reason == "target not found"
+
+
+class TestEditApplierRegistry:
+ """Tests for EditApplierRegistry."""
+
+ def test_register_and_get(self) -> None:
+ from openjarvis.learning.distillation.execute.base import (
+ ApplyContext,
+ ApplyResult,
+ EditApplier,
+ EditApplierRegistry,
+ ValidationResult,
+ )
+ from openjarvis.learning.distillation.models import Edit, EditOp
+
+ class FakeApplier(EditApplier):
+ op = EditOp.SET_MODEL_PARAM
+
+ def validate(self, edit: Edit, ctx: ApplyContext) -> ValidationResult:
+ return ValidationResult(ok=True)
+
+ def apply(self, edit: Edit, ctx: ApplyContext) -> ApplyResult:
+ return ApplyResult()
+
+ def rollback(self, edit: Edit, ctx: ApplyContext) -> None:
+ pass
+
+ registry = EditApplierRegistry()
+ registry.register(FakeApplier())
+ assert registry.is_supported(EditOp.SET_MODEL_PARAM)
+ applier = registry.get(EditOp.SET_MODEL_PARAM)
+ assert isinstance(applier, FakeApplier)
+
+ def test_is_supported_returns_false_for_unregistered(self) -> None:
+ from openjarvis.learning.distillation.execute.base import (
+ EditApplierRegistry,
+ )
+ from openjarvis.learning.distillation.models import EditOp
+
+ registry = EditApplierRegistry()
+ assert registry.is_supported(EditOp.LORA_FINETUNE) is False
+
+ def test_get_raises_for_unregistered(self) -> None:
+ from openjarvis.learning.distillation.execute.base import (
+ EditApplierRegistry,
+ )
+ from openjarvis.learning.distillation.models import EditOp
+
+ registry = EditApplierRegistry()
+ with pytest.raises(KeyError):
+ registry.get(EditOp.LORA_FINETUNE)
diff --git a/tests/learning/distillation/test_applier_lora_stub.py b/tests/learning/distillation/test_applier_lora_stub.py
new file mode 100644
index 00000000..fdb20b8e
--- /dev/null
+++ b/tests/learning/distillation/test_applier_lora_stub.py
@@ -0,0 +1,53 @@
+"""Tests for LoRA stub applier."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_lora_edit() -> Edit:
+ return Edit(
+ id="edit-lora",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.LORA_FINETUNE,
+ target="models.qwen2.5-coder:7b",
+ payload={"target_model": "qwen2.5-coder:7b", "data_source": "trace_filter:*"},
+ rationale="Fine-tune for math",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.MANUAL,
+ )
+
+
+class TestLoraStubApplier:
+ """Tests for LoraStubApplier."""
+
+ def test_validate_returns_not_ok(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.lora_stub import (
+ LoraStubApplier,
+ )
+
+ applier = LoraStubApplier()
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ result = applier.validate(_make_lora_edit(), ctx)
+ assert not result.ok
+ assert "v2" in result.reason.lower() or "deferred" in result.reason.lower()
+
+ def test_apply_raises_not_implemented(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.lora_stub import (
+ LoraStubApplier,
+ )
+
+ applier = LoraStubApplier()
+ ctx = ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+ with pytest.raises(NotImplementedError, match="deferred to v2"):
+ applier.apply(_make_lora_edit(), ctx)
diff --git a/tests/learning/distillation/test_appliers_agent.py b/tests/learning/distillation/test_appliers_agent.py
new file mode 100644
index 00000000..80f0196e
--- /dev/null
+++ b/tests/learning/distillation/test_appliers_agent.py
@@ -0,0 +1,210 @@
+"""Tests for agent-pillar appliers."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_ctx(tmp_path: Path) -> ApplyContext:
+ agents_dir = tmp_path / "agents" / "simple"
+ agents_dir.mkdir(parents=True)
+ (agents_dir / "system_prompt.md").write_text(
+ "You are a helpful assistant.\nBe concise.\n"
+ )
+ (tmp_path / "config.toml").write_text(
+ "[agent]\n"
+ 'default = "simple"\n'
+ "\n"
+ "[agent.simple]\n"
+ 'class = "simple"\n'
+ "max_turns = 5\n"
+ )
+ return ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+
+
+class TestReplaceSystemPromptApplier:
+ """Tests for ReplaceSystemPromptApplier."""
+
+ def test_validate_ok(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ ReplaceSystemPromptApplier,
+ )
+
+ applier = ReplaceSystemPromptApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.REPLACE_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"new_content": "New prompt content.\n"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ assert applier.validate(edit, ctx).ok
+
+ def test_apply_overwrites_prompt(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ ReplaceSystemPromptApplier,
+ )
+
+ applier = ReplaceSystemPromptApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.REPLACE_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"new_content": "You are a math expert.\n"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ applier.apply(edit, ctx)
+ content = (tmp_path / "agents" / "simple" / "system_prompt.md").read_text()
+ assert content == "You are a math expert.\n"
+
+
+class TestPatchSystemPromptApplier:
+ """Tests for PatchSystemPromptApplier."""
+
+ def test_apply_applies_diff(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ PatchSystemPromptApplier,
+ )
+
+ applier = PatchSystemPromptApplier()
+ ctx = _make_ctx(tmp_path)
+ diff = (
+ "--- a/prompt.md\n"
+ "+++ b/prompt.md\n"
+ "@@ -1,2 +1,2 @@\n"
+ " You are a helpful assistant.\n"
+ "-Be concise.\n"
+ "+Be concise and use math tools.\n"
+ )
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": diff},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ applier.apply(edit, ctx)
+ content = (tmp_path / "agents" / "simple" / "system_prompt.md").read_text()
+ assert "math tools" in content
+
+ def test_validate_fails_for_bad_diff(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ PatchSystemPromptApplier,
+ )
+
+ applier = PatchSystemPromptApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-002",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": "not a valid diff"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ result = applier.validate(edit, ctx)
+ assert not result.ok
+
+
+class TestSetAgentClassApplier:
+ """Tests for SetAgentClassApplier."""
+
+ def test_apply_updates_config(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ SetAgentClassApplier,
+ )
+
+ applier = SetAgentClassApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.SET_AGENT_CLASS,
+ target="agent.simple.class",
+ payload={"agent": "simple", "new_class": "react"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ applier.apply(edit, ctx)
+ content = ctx.config_path.read_text()
+ assert "react" in content
+
+
+class TestSetAgentParamApplier:
+ """Tests for SetAgentParamApplier."""
+
+ def test_apply_updates_param(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ SetAgentParamApplier,
+ )
+
+ applier = SetAgentParamApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.SET_AGENT_PARAM,
+ target="agent.simple.max_turns",
+ payload={"agent": "simple", "param": "max_turns", "value": 10},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ applier.apply(edit, ctx)
+ content = ctx.config_path.read_text()
+ assert "10" in content
+
+
+class TestEditFewShotExemplarsApplier:
+ """Tests for EditFewShotExemplarsApplier."""
+
+ def test_apply_writes_exemplars(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.agent import (
+ EditFewShotExemplarsApplier,
+ )
+
+ applier = EditFewShotExemplarsApplier()
+ ctx = _make_ctx(tmp_path)
+ exemplars = [
+ {"input": "What is 2+2?", "output": "4"},
+ {"input": "Capital of France?", "output": "Paris"},
+ ]
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.EDIT_FEW_SHOT_EXEMPLARS,
+ target="agents.simple.few_shot",
+ payload={"agent": "simple", "exemplars": exemplars},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ applier.apply(edit, ctx)
+ fs_path = tmp_path / "agents" / "simple" / "few_shot.json"
+ assert fs_path.exists()
+ data = json.loads(fs_path.read_text())
+ assert len(data) == 2
+ assert data[0]["input"] == "What is 2+2?"
diff --git a/tests/learning/distillation/test_appliers_intelligence.py b/tests/learning/distillation/test_appliers_intelligence.py
new file mode 100644
index 00000000..181138e2
--- /dev/null
+++ b/tests/learning/distillation/test_appliers_intelligence.py
@@ -0,0 +1,122 @@
+"""Tests for intelligence-pillar appliers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_ctx(tmp_path: Path) -> ApplyContext:
+ (tmp_path / "config.toml").write_text(
+ "[learning.routing]\n"
+ 'policy = "learned"\n'
+ "\n"
+ "[learning.routing.policy_map]\n"
+ 'math = "qwen2.5-coder:3b"\n'
+ 'code = "qwen2.5-coder:7b"\n'
+ )
+ return ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+
+
+def _make_routing_edit(
+ query_class: str = "math",
+ model: str = "qwen2.5-coder:14b",
+) -> Edit:
+ return Edit(
+ id="edit-001",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ target="learning.routing.policy_map.math",
+ payload={"query_class": query_class, "model": model},
+ rationale="Route math to bigger model",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO,
+ )
+
+
+def _make_param_edit(
+ model: str = "qwen2.5-coder:7b",
+ param: str = "temperature",
+ value: float = 0.3,
+) -> Edit:
+ return Edit(
+ id="edit-002",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_PARAM,
+ target=f"models.{model}.{param}",
+ payload={"model": model, "param": param, "value": value},
+ rationale="Lower temperature for code",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO,
+ )
+
+
+class TestSetModelForQueryClassApplier:
+ """Tests for SetModelForQueryClassApplier."""
+
+ def test_validate_ok(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelForQueryClassApplier,
+ )
+
+ applier = SetModelForQueryClassApplier()
+ ctx = _make_ctx(tmp_path)
+ result = applier.validate(_make_routing_edit(), ctx)
+ assert result.ok
+
+ def test_apply_updates_config(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelForQueryClassApplier,
+ )
+
+ applier = SetModelForQueryClassApplier()
+ ctx = _make_ctx(tmp_path)
+ applier.apply(_make_routing_edit(), ctx)
+ content = ctx.config_path.read_text()
+ assert "qwen2.5-coder:14b" in content
+
+ def test_apply_adds_new_query_class(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelForQueryClassApplier,
+ )
+
+ applier = SetModelForQueryClassApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = _make_routing_edit(query_class="science", model="qwen2.5-coder:14b")
+ applier.apply(edit, ctx)
+ content = ctx.config_path.read_text()
+ assert "science" in content
+ assert "qwen2.5-coder:14b" in content
+
+
+class TestSetModelParamApplier:
+ """Tests for SetModelParamApplier."""
+
+ def test_validate_ok(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelParamApplier,
+ )
+
+ applier = SetModelParamApplier()
+ ctx = _make_ctx(tmp_path)
+ result = applier.validate(_make_param_edit(), ctx)
+ assert result.ok
+
+ def test_apply_writes_param(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.intelligence import (
+ SetModelParamApplier,
+ )
+
+ applier = SetModelParamApplier()
+ ctx = _make_ctx(tmp_path)
+ applier.apply(_make_param_edit(), ctx)
+ content = ctx.config_path.read_text()
+ assert "temperature" in content
+ assert "0.3" in content
diff --git a/tests/learning/distillation/test_appliers_tools.py b/tests/learning/distillation/test_appliers_tools.py
new file mode 100644
index 00000000..0035905c
--- /dev/null
+++ b/tests/learning/distillation/test_appliers_tools.py
@@ -0,0 +1,145 @@
+"""Tests for tools-pillar appliers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_ctx(tmp_path: Path) -> ApplyContext:
+ tools_dir = tmp_path / "tools"
+ tools_dir.mkdir(parents=True)
+ (tools_dir / "descriptions.toml").write_text(
+ '[web_search]\ndescription = "Search the web"\n'
+ )
+ (tmp_path / "config.toml").write_text('[agent.simple]\ntools = ["web_search"]\n')
+ return ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+
+
+class TestAddToolToAgentApplier:
+ """Tests for AddToolToAgentApplier."""
+
+ def test_apply_adds_tool(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.tools import (
+ AddToolToAgentApplier,
+ )
+
+ applier = AddToolToAgentApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.TOOLS,
+ op=EditOp.ADD_TOOL_TO_AGENT,
+ target="agent.simple.tools",
+ payload={"agent": "simple", "tool_name": "calculator"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ applier.apply(edit, ctx)
+ content = ctx.config_path.read_text()
+ assert "calculator" in content
+
+ def test_validate_ok(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.tools import (
+ AddToolToAgentApplier,
+ )
+
+ applier = AddToolToAgentApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.TOOLS,
+ op=EditOp.ADD_TOOL_TO_AGENT,
+ target="agent.simple.tools",
+ payload={"agent": "simple", "tool_name": "calculator"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ assert applier.validate(edit, ctx).ok
+
+
+class TestRemoveToolFromAgentApplier:
+ """Tests for RemoveToolFromAgentApplier."""
+
+ def test_apply_removes_tool(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.tools import (
+ RemoveToolFromAgentApplier,
+ )
+
+ applier = RemoveToolFromAgentApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.TOOLS,
+ op=EditOp.REMOVE_TOOL_FROM_AGENT,
+ target="agent.simple.tools",
+ payload={"agent": "simple", "tool_name": "web_search"},
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ applier.apply(edit, ctx)
+ content = ctx.config_path.read_text()
+ assert "web_search" not in content
+
+
+class TestEditToolDescriptionApplier:
+ """Tests for EditToolDescriptionApplier."""
+
+ def test_apply_updates_description(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.tools import (
+ EditToolDescriptionApplier,
+ )
+
+ applier = EditToolDescriptionApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.TOOLS,
+ op=EditOp.EDIT_TOOL_DESCRIPTION,
+ target="tools.web_search.description",
+ payload={
+ "tool_name": "web_search",
+ "new_description": "Search the internet for current information",
+ },
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ applier.apply(edit, ctx)
+ content = (ctx.tools_dir / "descriptions.toml").read_text()
+ assert "Search the internet" in content
+
+ def test_apply_adds_new_tool_section(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.appliers.tools import (
+ EditToolDescriptionApplier,
+ )
+
+ applier = EditToolDescriptionApplier()
+ ctx = _make_ctx(tmp_path)
+ edit = Edit(
+ id="edit-002",
+ pillar=EditPillar.TOOLS,
+ op=EditOp.EDIT_TOOL_DESCRIPTION,
+ target="tools.calculator.description",
+ payload={
+ "tool_name": "calculator",
+ "new_description": "Evaluate mathematical expressions",
+ },
+ rationale="test",
+ expected_improvement="c1",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ applier.apply(edit, ctx)
+ content = (ctx.tools_dir / "descriptions.toml").read_text()
+ assert "[calculator]" in content
+ assert "Evaluate mathematical" in content
diff --git a/tests/learning/distillation/test_benchmark_gate.py b/tests/learning/distillation/test_benchmark_gate.py
new file mode 100644
index 00000000..fd644446
--- /dev/null
+++ b/tests/learning/distillation/test_benchmark_gate.py
@@ -0,0 +1,147 @@
+"""Tests for openjarvis.learning.distillation.gate.benchmark_gate module.
+
+All tests use mock scorers — no live EvalRunner.
+"""
+
+from __future__ import annotations
+
+from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+
+def _make_scorer(scores: dict[str, float], overall: float | None = None):
+ """Return a callable that produces a BenchmarkSnapshot with given scores."""
+
+ def scorer(
+ *, benchmark_version: str, subsample_size: int, seed: int
+ ) -> BenchmarkSnapshot:
+ computed = sum(scores.values()) / max(len(scores), 1)
+ return BenchmarkSnapshot(
+ benchmark_version=benchmark_version,
+ overall_score=overall if overall is not None else computed,
+ cluster_scores=scores,
+ task_count=subsample_size,
+ elapsed_seconds=5.0,
+ )
+
+ return scorer
+
+
+class TestBenchmarkGate:
+ """Tests for BenchmarkGate."""
+
+ def test_accepts_improving_edit(self) -> None:
+ from openjarvis.learning.distillation.gate.benchmark_gate import (
+ BenchmarkGate,
+ )
+
+ before = BenchmarkSnapshot(
+ benchmark_version="v1",
+ overall_score=0.6,
+ cluster_scores={"c1": 0.5, "c2": 0.7},
+ task_count=50,
+ elapsed_seconds=10.0,
+ )
+ gate = BenchmarkGate(
+ scorer=_make_scorer({"c1": 0.6, "c2": 0.75}, overall=0.68),
+ benchmark_version="v1",
+ min_improvement=0.0,
+ max_regression=0.05,
+ subsample_size=50,
+ )
+ result = gate.evaluate(before=before, session_seed=42)
+ assert result.accepted
+ assert result.snapshot.overall_score == 0.68
+ assert result.delta > 0
+
+ def test_rejects_no_improvement(self) -> None:
+ from openjarvis.learning.distillation.gate.benchmark_gate import (
+ BenchmarkGate,
+ )
+
+ before = BenchmarkSnapshot(
+ benchmark_version="v1",
+ overall_score=0.7,
+ cluster_scores={"c1": 0.6, "c2": 0.8},
+ task_count=50,
+ elapsed_seconds=10.0,
+ )
+ gate = BenchmarkGate(
+ scorer=_make_scorer({"c1": 0.6, "c2": 0.8}, overall=0.7),
+ benchmark_version="v1",
+ min_improvement=0.0,
+ max_regression=0.05,
+ subsample_size=50,
+ )
+ result = gate.evaluate(before=before, session_seed=42)
+ assert not result.accepted
+ assert "no improvement" in result.reason.lower()
+
+ def test_rejects_regression(self) -> None:
+ from openjarvis.learning.distillation.gate.benchmark_gate import (
+ BenchmarkGate,
+ )
+
+ before = BenchmarkSnapshot(
+ benchmark_version="v1",
+ overall_score=0.7,
+ cluster_scores={"c1": 0.6, "c2": 0.8},
+ task_count=50,
+ elapsed_seconds=10.0,
+ )
+ # overall improves but c2 regresses badly
+ gate = BenchmarkGate(
+ scorer=_make_scorer({"c1": 0.75, "c2": 0.65}, overall=0.72),
+ benchmark_version="v1",
+ min_improvement=0.0,
+ max_regression=0.05,
+ subsample_size=50,
+ )
+ result = gate.evaluate(before=before, session_seed=42)
+ assert not result.accepted
+ assert "regression" in result.reason.lower()
+
+ def test_min_improvement_threshold(self) -> None:
+ from openjarvis.learning.distillation.gate.benchmark_gate import (
+ BenchmarkGate,
+ )
+
+ before = BenchmarkSnapshot(
+ benchmark_version="v1",
+ overall_score=0.7,
+ cluster_scores={"c1": 0.6, "c2": 0.8},
+ task_count=50,
+ elapsed_seconds=10.0,
+ )
+ # Tiny improvement of 0.01, but min_improvement requires 0.05
+ gate = BenchmarkGate(
+ scorer=_make_scorer({"c1": 0.61, "c2": 0.81}, overall=0.71),
+ benchmark_version="v1",
+ min_improvement=0.05,
+ max_regression=0.05,
+ subsample_size=50,
+ )
+ result = gate.evaluate(before=before, session_seed=42)
+ assert not result.accepted
+
+ def test_result_contains_snapshot(self) -> None:
+ from openjarvis.learning.distillation.gate.benchmark_gate import (
+ BenchmarkGate,
+ )
+
+ before = BenchmarkSnapshot(
+ benchmark_version="v1",
+ overall_score=0.5,
+ cluster_scores={"c1": 0.5},
+ task_count=50,
+ elapsed_seconds=10.0,
+ )
+ gate = BenchmarkGate(
+ scorer=_make_scorer({"c1": 0.7}, overall=0.7),
+ benchmark_version="v1",
+ min_improvement=0.0,
+ max_regression=0.05,
+ subsample_size=50,
+ )
+ result = gate.evaluate(before=before, session_seed=42)
+ assert isinstance(result.snapshot, BenchmarkSnapshot)
+ assert result.snapshot.benchmark_version == "v1"
diff --git a/tests/learning/distillation/test_checkpoint_store.py b/tests/learning/distillation/test_checkpoint_store.py
new file mode 100644
index 00000000..9d19d5dd
--- /dev/null
+++ b/tests/learning/distillation/test_checkpoint_store.py
@@ -0,0 +1,218 @@
+"""Tests for openjarvis.learning.distillation.checkpoint.store module."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+
+def _git(cwd: Path, *args: str) -> str:
+ """Helper to run git commands in tests."""
+ result = subprocess.run(
+ ["git", *args],
+ cwd=cwd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ return result.stdout.strip()
+
+
+def _setup_isolated_repo_root(tmp_path: Path) -> Path:
+ """Create a fake openjarvis-home directory tree with config files for the
+ CheckpointStore to track. Returns the root."""
+ root = tmp_path / "openjarvis_home"
+ (root / "agents" / "simple").mkdir(parents=True)
+ (root / "tools").mkdir(parents=True)
+ (root / "config.toml").write_text("[learning]\nenabled = true\n")
+ (root / "agents" / "simple" / "system_prompt.md").write_text(
+ "You are a helpful assistant.\n"
+ )
+ (root / "tools" / "descriptions.toml").write_text("[web_search]\n")
+ return root
+
+
+class TestCheckpointStoreInit:
+ """Tests for CheckpointStore.init()."""
+
+ def test_creates_repo_with_baseline_commit(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+
+ assert (root / ".git").exists()
+ log = _git(root, "log", "--oneline")
+ assert "baseline" in log
+
+ def test_init_idempotent(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+ first_sha = store.current_sha()
+
+ store.init() # Should not raise or create a second baseline.
+ assert store.current_sha() == first_sha
+
+ def test_init_refuses_inside_source_tree(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+ from openjarvis.learning.distillation.storage import paths
+
+ source_root = paths._find_source_root()
+ assert source_root is not None
+ bad_root = source_root / "fake_openjarvis_home"
+
+ store = CheckpointStore(bad_root)
+ with pytest.raises(paths.ConfigurationError):
+ store.init()
+
+
+class TestStageCommitDiscard:
+ """Tests for begin_stage / commit_stage / discard_stage."""
+
+ def test_commit_stage_creates_commit_with_trailers(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+
+ handle = store.begin_stage("edit-001")
+ # Mutate a tracked file in the working tree.
+ (root / "agents" / "simple" / "system_prompt.md").write_text(
+ "You are a helpful, math-aware assistant.\n"
+ )
+
+ new_sha = store.commit_stage(
+ handle,
+ message="learning: edit-001 add math hint",
+ session_id="session-001",
+ risk_tier="review",
+ )
+
+ # New commit exists.
+ assert new_sha != handle.pre_stage_sha
+ # Commit message contains structured trailers.
+ body = _git(root, "log", "-1", "--format=%B", new_sha)
+ assert "Edit-ID: edit-001" in body
+ assert "Session-ID: session-001" in body
+ assert "Risk-Tier: review" in body
+
+ def test_discard_stage_restores_working_tree(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+ original = (root / "agents" / "simple" / "system_prompt.md").read_text()
+
+ handle = store.begin_stage("edit-002")
+ (root / "agents" / "simple" / "system_prompt.md").write_text(
+ "Mutated content that should be discarded.\n"
+ )
+
+ store.discard_stage(handle)
+
+ restored = (root / "agents" / "simple" / "system_prompt.md").read_text()
+ assert restored == original
+ # HEAD must equal the pre-stage sha.
+ assert store.current_sha() == handle.pre_stage_sha
+
+ def test_begin_stage_refuses_dirty_working_tree(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ DirtyWorkingTreeError,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+
+ # Create an untracked, uncommitted change.
+ (root / "agents" / "simple" / "system_prompt.md").write_text(
+ "Pre-existing manual edit.\n"
+ )
+
+ with pytest.raises(DirtyWorkingTreeError):
+ store.begin_stage("edit-003")
+
+
+class TestRevertSession:
+ """Tests for revert_session."""
+
+ def test_revert_creates_new_commits_and_does_not_rewrite(
+ self, tmp_path: Path
+ ) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+
+ # Apply two commits tagged with the same session id.
+ handle1 = store.begin_stage("edit-001")
+ (root / "agents" / "simple" / "system_prompt.md").write_text("version A\n")
+ store.commit_stage(
+ handle1,
+ message="learning: edit-001 v A",
+ session_id="session-XYZ",
+ risk_tier="auto",
+ )
+
+ handle2 = store.begin_stage("edit-002")
+ (root / "tools" / "descriptions.toml").write_text(
+ "[web_search]\nupdated = true\n"
+ )
+ store.commit_stage(
+ handle2,
+ message="learning: edit-002 update tool",
+ session_id="session-XYZ",
+ risk_tier="auto",
+ )
+
+ before_revert_log_count = len(_git(root, "log", "--oneline").splitlines())
+
+ revert_shas = store.revert_session("session-XYZ")
+ assert len(revert_shas) == 2
+
+ # Two new commits added (the reverts), no history rewriting.
+ after_revert_log_count = len(_git(root, "log", "--oneline").splitlines())
+ assert after_revert_log_count == before_revert_log_count + 2
+
+ # Files restored to baseline.
+ assert (
+ root / "agents" / "simple" / "system_prompt.md"
+ ).read_text() == "You are a helpful assistant.\n"
+ assert (root / "tools" / "descriptions.toml").read_text() == "[web_search]\n"
+
+ def test_revert_session_with_no_commits_returns_empty(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = _setup_isolated_repo_root(tmp_path)
+ store = CheckpointStore(root)
+ store.init()
+
+ result = store.revert_session("session-with-no-commits")
+ assert result == []
diff --git a/tests/learning/distillation/test_cli.py b/tests/learning/distillation/test_cli.py
new file mode 100644
index 00000000..52a2c994
--- /dev/null
+++ b/tests/learning/distillation/test_cli.py
@@ -0,0 +1,44 @@
+"""Tests for the jarvis learning CLI subcommand group."""
+
+from __future__ import annotations
+
+from click.testing import CliRunner
+
+
+class TestLearningCLI:
+ def test_learning_group_exists(self) -> None:
+ from openjarvis.learning.distillation.cli import learning_group
+
+ runner = CliRunner()
+ result = runner.invoke(learning_group, ["--help"])
+ assert result.exit_code == 0
+ out = result.output.lower()
+ assert "learning" in out or "distillation" in out
+
+ def test_init_subcommand(self) -> None:
+ from openjarvis.learning.distillation.cli import learning_group
+
+ runner = CliRunner()
+ result = runner.invoke(learning_group, ["init", "--help"])
+ assert result.exit_code == 0
+
+ def test_run_subcommand(self) -> None:
+ from openjarvis.learning.distillation.cli import learning_group
+
+ runner = CliRunner()
+ result = runner.invoke(learning_group, ["run", "--help"])
+ assert result.exit_code == 0
+
+ def test_history_subcommand(self) -> None:
+ from openjarvis.learning.distillation.cli import learning_group
+
+ runner = CliRunner()
+ result = runner.invoke(learning_group, ["history", "--help"])
+ assert result.exit_code == 0
+
+ def test_rollback_subcommand(self) -> None:
+ from openjarvis.learning.distillation.cli import learning_group
+
+ runner = CliRunner()
+ result = runner.invoke(learning_group, ["rollback", "--help"])
+ assert result.exit_code == 0
diff --git a/tests/learning/distillation/test_cold_start.py b/tests/learning/distillation/test_cold_start.py
new file mode 100644
index 00000000..5714c6f9
--- /dev/null
+++ b/tests/learning/distillation/test_cold_start.py
@@ -0,0 +1,101 @@
+"""Tests for openjarvis.learning.distillation.gate.cold_start module."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+from unittest.mock import MagicMock
+
+
+@dataclass
+class _StubTrace:
+ trace_id: str = "t1"
+ query: str = "test"
+ feedback: Optional[float] = 0.8
+ agent: str = "simple"
+ model: str = "qwen"
+ outcome: Optional[str] = "success"
+ result: str = "answer"
+ started_at: float = 1712534400.0
+ ended_at: float = 1712534401.0
+ steps: list = field(default_factory=list)
+ messages: list = field(default_factory=list)
+ total_tokens: int = 100
+ total_latency_seconds: float = 1.0
+ metadata: Dict[str, Any] = field(default_factory=dict)
+ engine: str = "ollama"
+
+
+def _make_trace_store(count: int = 0, high_feedback_count: int = 0) -> MagicMock:
+ store = MagicMock()
+ store.count.return_value = count
+ # list_traces returns traces with high feedback
+ high_traces = [
+ _StubTrace(trace_id=f"t{i}", feedback=0.9) for i in range(high_feedback_count)
+ ]
+ store.list_traces.return_value = high_traces
+ return store
+
+
+class TestCheckReadiness:
+ """Tests for check_readiness()."""
+
+ def test_not_ready_with_no_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_readiness,
+ )
+
+ store = _make_trace_store(count=0)
+ result = check_readiness(store, min_traces=20)
+ assert not result.ready
+ assert "not enough traces" in result.message.lower()
+
+ def test_not_ready_with_few_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_readiness,
+ )
+
+ store = _make_trace_store(count=10)
+ result = check_readiness(store, min_traces=20)
+ assert not result.ready
+
+ def test_ready_with_enough_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_readiness,
+ )
+
+ store = _make_trace_store(count=25)
+ result = check_readiness(store, min_traces=20)
+ assert result.ready
+
+
+class TestCheckBenchmarkReady:
+ """Tests for check_benchmark_ready()."""
+
+ def test_not_ready_with_no_high_feedback_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_benchmark_ready,
+ )
+
+ store = _make_trace_store(count=30, high_feedback_count=0)
+ result = check_benchmark_ready(store, min_feedback=0.7, min_samples=10)
+ assert not result.ready
+ assert "benchmark" in result.message.lower()
+
+ def test_not_ready_with_few_high_feedback_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_benchmark_ready,
+ )
+
+ store = _make_trace_store(count=30, high_feedback_count=5)
+ result = check_benchmark_ready(store, min_feedback=0.7, min_samples=10)
+ assert not result.ready
+
+ def test_ready_with_enough_high_feedback_traces(self) -> None:
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_benchmark_ready,
+ )
+
+ store = _make_trace_store(count=30, high_feedback_count=15)
+ result = check_benchmark_ready(store, min_feedback=0.7, min_samples=10)
+ assert result.ready
diff --git a/tests/learning/distillation/test_diagnose_tools.py b/tests/learning/distillation/test_diagnose_tools.py
new file mode 100644
index 00000000..bad05e89
--- /dev/null
+++ b/tests/learning/distillation/test_diagnose_tools.py
@@ -0,0 +1,282 @@
+"""Tests for openjarvis.learning.distillation.diagnose.tools module.
+
+All tests use fixture stubs — no live TraceStore, CloudEngine, or ToolRegistry.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Optional
+from unittest.mock import MagicMock
+
+# ---------------------------------------------------------------------------
+# Stubs for dependencies
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _StubTrace:
+ """Minimal stub matching the Trace fields that tools access."""
+
+ trace_id: str = "trace-001"
+ query: str = "What is 2+2?"
+ agent: str = "simple"
+ model: str = "qwen2.5-coder:7b"
+ outcome: Optional[str] = "success"
+ feedback: Optional[float] = 0.8
+ started_at: float = 1712534400.0
+ result: str = "4"
+ steps: list = field(default_factory=list)
+ messages: list = field(default_factory=list)
+ total_tokens: int = 100
+ total_latency_seconds: float = 1.0
+ metadata: Dict[str, Any] = field(default_factory=dict)
+ ended_at: float = 1712534401.0
+ engine: str = "ollama"
+
+
+def _make_stub_trace_store(traces: list[_StubTrace] | None = None) -> MagicMock:
+ """Create a mock TraceStore with canned responses."""
+ store = MagicMock()
+ traces = traces or [_StubTrace()]
+ store.list_traces.return_value = traces
+ store.get.side_effect = lambda tid: next(
+ (t for t in traces if t.trace_id == tid), None
+ )
+ store.search.return_value = [
+ {"trace_id": t.trace_id, "query": t.query, "score": 1.0} for t in traces
+ ]
+ return store
+
+
+def _make_stub_benchmark_samples() -> list:
+ """Return a list of stub PersonalBenchmarkSample objects."""
+ sample = MagicMock()
+ sample.trace_id = "task-001"
+ sample.query = "What is quantum computing?"
+ sample.reference_answer = "Quantum computing uses qubits..."
+ sample.category = "reasoning"
+ sample.feedback_score = 0.9
+ return [sample]
+
+
+def _make_stub_config(tmp_path: Path) -> dict:
+ """Create a minimal config dict and on-disk files."""
+ agents_dir = tmp_path / "agents" / "simple"
+ agents_dir.mkdir(parents=True)
+ (agents_dir / "system_prompt.md").write_text("You are a helpful assistant.\n")
+ tools_dir = tmp_path / "tools"
+ tools_dir.mkdir(parents=True)
+ (tools_dir / "descriptions.toml").write_text(
+ '[web_search]\ndescription = "Search the web"\n'
+ )
+ config_path = tmp_path / "config.toml"
+ config_path.write_text("[learning]\nenabled = true\n")
+ return {
+ "config_path": config_path,
+ "openjarvis_home": tmp_path,
+ }
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestBuildDiagnosticTools:
+ """Tests for the build_diagnostic_tools factory."""
+
+ def test_returns_expected_tool_names(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=_make_stub_benchmark_samples(),
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ names = {t.name for t in tools}
+ assert "list_traces" in names
+ assert "get_trace" in names
+ assert "search_traces" in names
+ assert "get_current_config" in names
+ assert "get_agent_prompt" in names
+ assert "get_tool_description" in names
+ assert "list_available_tools" in names
+ assert "list_personal_benchmark" in names
+ assert "run_student_on_task" in names
+ assert "run_self_on_task" in names
+ assert "compare_outputs" in names
+
+ def test_all_tools_have_openai_format(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=_make_stub_benchmark_samples(),
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ for tool in tools:
+ spec = tool.to_openai_function()
+ assert spec["type"] == "function"
+ assert "name" in spec["function"]
+ assert "parameters" in spec["function"]
+
+
+class TestListTraces:
+ """Tests for the list_traces diagnostic tool."""
+
+ def test_returns_trace_metas(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ traces = [
+ _StubTrace(trace_id="t1", feedback=0.3),
+ _StubTrace(trace_id="t2", feedback=0.9),
+ ]
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(traces),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ list_traces = next(t for t in tools if t.name == "list_traces")
+ result = list_traces.fn(limit=10)
+ parsed = json.loads(result)
+ assert len(parsed) == 2
+ assert parsed[0]["trace_id"] == "t1"
+
+
+class TestGetTrace:
+ """Tests for the get_trace diagnostic tool."""
+
+ def test_returns_trace_details(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ get_trace = next(t for t in tools if t.name == "get_trace")
+ result = get_trace.fn(trace_id="trace-001")
+ parsed = json.loads(result)
+ assert parsed["trace_id"] == "trace-001"
+ assert parsed["query"] == "What is 2+2?"
+
+ def test_returns_error_for_unknown_trace(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ get_trace = next(t for t in tools if t.name == "get_trace")
+ result = get_trace.fn(trace_id="nonexistent")
+ assert "not found" in result.lower()
+
+
+class TestGetCurrentConfig:
+ """Tests for the get_current_config diagnostic tool."""
+
+ def test_returns_config_content(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ get_config = next(t for t in tools if t.name == "get_current_config")
+ result = get_config.fn()
+ assert "learning" in result
+
+
+class TestGetAgentPrompt:
+ """Tests for the get_agent_prompt diagnostic tool."""
+
+ def test_returns_prompt_content(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ get_prompt = next(t for t in tools if t.name == "get_agent_prompt")
+ result = get_prompt.fn(agent_name="simple")
+ assert "helpful assistant" in result
+
+
+class TestListPersonalBenchmark:
+ """Tests for the list_personal_benchmark diagnostic tool."""
+
+ def test_returns_benchmark_tasks(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.tools import (
+ build_diagnostic_tools,
+ )
+
+ tools = build_diagnostic_tools(
+ trace_store=_make_stub_trace_store(),
+ config=_make_stub_config(tmp_path),
+ benchmark_samples=_make_stub_benchmark_samples(),
+ student_runner=MagicMock(),
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ judge=MagicMock(),
+ session_id="session-001",
+ )
+ list_bench = next(t for t in tools if t.name == "list_personal_benchmark")
+ result = list_bench.fn(limit=10)
+ parsed = json.loads(result)
+ assert len(parsed) == 1
+ assert parsed[0]["task_id"] == "task-001"
diff --git a/tests/learning/distillation/test_diagnose_types.py b/tests/learning/distillation/test_diagnose_types.py
new file mode 100644
index 00000000..36046b61
--- /dev/null
+++ b/tests/learning/distillation/test_diagnose_types.py
@@ -0,0 +1,173 @@
+"""Tests for openjarvis.learning.distillation.diagnose.types module."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+
+class TestTraceMeta:
+ """Tests for TraceMeta dataclass."""
+
+ def test_constructs_with_required_fields(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import TraceMeta
+
+ meta = TraceMeta(
+ trace_id="trace-001",
+ query="What is 2+2?",
+ agent="simple",
+ model="qwen2.5-coder:7b",
+ outcome="success",
+ feedback=0.8,
+ started_at=1712534400.0,
+ )
+ assert meta.trace_id == "trace-001"
+ assert meta.feedback == 0.8
+
+ def test_feedback_can_be_none(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import TraceMeta
+
+ meta = TraceMeta(
+ trace_id="trace-002",
+ query="test",
+ agent="simple",
+ model="qwen2.5-coder:7b",
+ outcome=None,
+ feedback=None,
+ started_at=1712534400.0,
+ )
+ assert meta.feedback is None
+
+
+class TestBenchmarkTask:
+ """Tests for BenchmarkTask dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import BenchmarkTask
+
+ task = BenchmarkTask(
+ task_id="task-001",
+ query="Explain quantum computing",
+ reference_answer="Quantum computing uses qubits...",
+ category="reasoning",
+ )
+ assert task.task_id == "task-001"
+ assert task.category == "reasoning"
+
+
+class TestStudentRun:
+ """Tests for StudentRun dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import StudentRun
+
+ run = StudentRun(
+ task_id="task-001",
+ output="The answer is 4.",
+ score=0.9,
+ trace_id="trace-new-001",
+ latency_seconds=2.5,
+ tokens_used=150,
+ )
+ assert run.score == 0.9
+ assert run.trace_id == "trace-new-001"
+
+
+class TestTeacherRun:
+ """Tests for TeacherRun dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import TeacherRun
+
+ run = TeacherRun(
+ task_id="task-001",
+ output="Quantum computing is...",
+ reasoning="I approached this by...",
+ cost_usd=0.05,
+ tokens_used=500,
+ )
+ assert run.cost_usd == 0.05
+
+
+class TestComparisonResult:
+ """Tests for ComparisonResult dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import ComparisonResult
+
+ result = ComparisonResult(
+ task_id="task-001",
+ student_score=0.3,
+ teacher_score=0.9,
+ judge_reasoning="The student missed the key concept...",
+ )
+ assert result.student_score == 0.3
+ assert result.teacher_score == 0.9
+
+
+class TestToolMeta:
+ """Tests for ToolMeta dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import ToolMeta
+
+ meta = ToolMeta(
+ name="calculator",
+ description="Evaluate math expressions",
+ category="math",
+ agents=["simple", "react"],
+ )
+ assert meta.name == "calculator"
+ assert len(meta.agents) == 2
+
+
+class TestDiagnosticTool:
+ """Tests for DiagnosticTool dataclass."""
+
+ def test_constructs_with_callable(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import DiagnosticTool
+
+ def my_func(**kwargs: object) -> str:
+ return "result"
+
+ tool = DiagnosticTool(
+ name="test_tool",
+ description="A test tool",
+ parameters={"type": "object", "properties": {}},
+ fn=my_func,
+ )
+ assert tool.name == "test_tool"
+ assert tool.fn(foo="bar") == "result"
+
+
+class TestToolCallRecord:
+ """Tests for ToolCallRecord dataclass."""
+
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import ToolCallRecord
+
+ record = ToolCallRecord(
+ timestamp=datetime(2026, 4, 9, 3, 0, 0, tzinfo=timezone.utc),
+ tool="list_traces",
+ args={"limit": 10},
+ result="[...]",
+ latency_ms=42.5,
+ cost_usd=0.0,
+ )
+ assert record.tool == "list_traces"
+ assert record.latency_ms == 42.5
+
+ def test_to_jsonl_dict(self) -> None:
+ from openjarvis.learning.distillation.diagnose.types import ToolCallRecord
+
+ record = ToolCallRecord(
+ timestamp=datetime(2026, 4, 9, 3, 0, 0, tzinfo=timezone.utc),
+ tool="get_trace",
+ args={"trace_id": "t1"},
+ result="trace data",
+ latency_ms=10.0,
+ cost_usd=0.01,
+ )
+ d = record.to_jsonl_dict()
+ assert d["tool"] == "get_trace"
+ assert d["timestamp"] == "2026-04-09T03:00:00+00:00"
+ assert d["cost_usd"] == 0.01
diff --git a/tests/learning/distillation/test_diagnosis_runner.py b/tests/learning/distillation/test_diagnosis_runner.py
new file mode 100644
index 00000000..baedbc98
--- /dev/null
+++ b/tests/learning/distillation/test_diagnosis_runner.py
@@ -0,0 +1,235 @@
+"""Tests for openjarvis.learning.distillation.diagnose.runner module.
+
+All tests use mocked dependencies — no live API calls.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock
+
+
+def _make_engine_response(content: str, tool_calls: list | None = None) -> dict:
+ """Create a mock engine.generate() response."""
+ resp = {
+ "content": content,
+ "finish_reason": "stop" if not tool_calls else "tool_calls",
+ "usage": {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300},
+ "cost_usd": 0.02,
+ }
+ if tool_calls:
+ resp["tool_calls"] = tool_calls
+ return resp
+
+
+def _make_diagnosis_content() -> str:
+ """A plausible teacher diagnosis output with embedded cluster JSON."""
+ return (
+ "## Diagnosis\n\n"
+ "The student has two main failure patterns:\n\n"
+ "### Cluster 1: Math Routing\n"
+ "Math queries are being routed to qwen-3b which lacks chain-of-thought.\n\n"
+ "### Cluster 2: Tool Selection\n"
+ "The student frequently fails to use the calculator tool for arithmetic.\n\n"
+ "```json\n"
+ + json.dumps(
+ [
+ {
+ "id": "cluster-001",
+ "description": "Math queries routed to qwen-3b",
+ "sample_trace_ids": ["t1", "t2", "t3"],
+ "student_failure_rate": 0.8,
+ "teacher_success_rate": 0.95,
+ "skill_gap": "Student lacks chain-of-thought on multi-step math",
+ },
+ {
+ "id": "cluster-002",
+ "description": "Calculator tool not used for arithmetic",
+ "sample_trace_ids": ["t4", "t5", "t6"],
+ "student_failure_rate": 0.6,
+ "teacher_success_rate": 0.9,
+ "skill_gap": "Student does not invoke calculator tool",
+ },
+ ]
+ )
+ + "\n```\n"
+ )
+
+
+class TestDiagnosisRunner:
+ """Tests for DiagnosisRunner."""
+
+ def test_produces_diagnosis_artifact(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content=_make_diagnosis_content()
+ )
+
+ runner = DiagnosisRunner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_dir=tmp_path / "session-001",
+ session_id="session-001",
+ config={
+ "config_path": tmp_path / "config.toml",
+ "openjarvis_home": tmp_path,
+ },
+ )
+ # Create minimal config file
+ (tmp_path / "config.toml").write_text("[learning]\n")
+
+ runner.run()
+
+ # Diagnosis artifact written
+ diagnosis_path = tmp_path / "session-001" / "diagnosis.md"
+ assert diagnosis_path.exists()
+ assert "Math" in diagnosis_path.read_text()
+
+ def test_returns_failure_clusters(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content=_make_diagnosis_content()
+ )
+
+ runner = DiagnosisRunner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_dir=tmp_path / "session-001",
+ session_id="session-001",
+ config={
+ "config_path": tmp_path / "config.toml",
+ "openjarvis_home": tmp_path,
+ },
+ )
+ (tmp_path / "config.toml").write_text("[learning]\n")
+
+ result = runner.run()
+
+ assert len(result.clusters) == 2
+ assert result.clusters[0].id == "cluster-001"
+ assert result.clusters[0].student_failure_rate == 0.8
+ assert result.clusters[1].id == "cluster-002"
+
+ def test_persists_teacher_traces_jsonl(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ engine = MagicMock()
+ # Turn 1: tool call
+ engine.generate.side_effect = [
+ _make_engine_response(
+ content="",
+ tool_calls=[
+ {
+ "id": "c1",
+ "name": "get_current_config",
+ "arguments": "{}",
+ }
+ ],
+ ),
+ _make_engine_response(content=_make_diagnosis_content()),
+ ]
+
+ runner = DiagnosisRunner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_dir=tmp_path / "session-001",
+ session_id="session-001",
+ config={
+ "config_path": tmp_path / "config.toml",
+ "openjarvis_home": tmp_path,
+ },
+ )
+ (tmp_path / "config.toml").write_text("[learning]\n")
+
+ runner.run()
+
+ jsonl_path = tmp_path / "session-001" / "teacher_traces" / "diagnose.jsonl"
+ assert jsonl_path.exists()
+ lines = jsonl_path.read_text().strip().splitlines()
+ assert len(lines) >= 1
+ record = json.loads(lines[0])
+ assert "tool" in record
+
+ def test_returns_cost(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content=_make_diagnosis_content()
+ )
+
+ runner = DiagnosisRunner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_dir=tmp_path / "session-001",
+ session_id="session-001",
+ config={
+ "config_path": tmp_path / "config.toml",
+ "openjarvis_home": tmp_path,
+ },
+ )
+ (tmp_path / "config.toml").write_text("[learning]\n")
+
+ result = runner.run()
+ assert result.cost_usd >= 0.0
+
+ def test_handles_no_clusters_in_output(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content="## Diagnosis\nNo clear failure patterns found."
+ )
+
+ runner = DiagnosisRunner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_dir=tmp_path / "session-001",
+ session_id="session-001",
+ config={
+ "config_path": tmp_path / "config.toml",
+ "openjarvis_home": tmp_path,
+ },
+ )
+ (tmp_path / "config.toml").write_text("[learning]\n")
+
+ result = runner.run()
+
+ assert result.clusters == []
+ # Diagnosis artifact is still written
+ assert (tmp_path / "session-001" / "diagnosis.md").exists()
diff --git a/tests/learning/distillation/test_execution_loop.py b/tests/learning/distillation/test_execution_loop.py
new file mode 100644
index 00000000..ff9e27c2
--- /dev/null
+++ b/tests/learning/distillation/test_execution_loop.py
@@ -0,0 +1,150 @@
+"""Tests for the per-edit execution loop."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from openjarvis.learning.distillation.execute.base import ApplyContext
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_ctx(tmp_path: Path) -> ApplyContext:
+ (tmp_path / "config.toml").write_text(
+ '[learning.routing.policy_map]\nmath = "qwen2.5-coder:3b"\n'
+ )
+ agents_dir = tmp_path / "agents" / "simple"
+ agents_dir.mkdir(parents=True)
+ (agents_dir / "system_prompt.md").write_text("You are helpful.\n")
+ tools_dir = tmp_path / "tools"
+ tools_dir.mkdir(parents=True)
+ (tools_dir / "descriptions.toml").write_text(
+ '[web_search]\ndescription = "Search"\n'
+ )
+ return ApplyContext(openjarvis_home=tmp_path, session_id="s1")
+
+
+def _make_auto_edit(edit_id: str = "edit-001") -> Edit:
+ return Edit(
+ id=edit_id,
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ target="routing.math",
+ payload={"query_class": "math", "model": "qwen2.5-coder:14b"},
+ rationale="Route math to bigger model",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO,
+ )
+
+
+def _make_review_edit(edit_id: str = "edit-002") -> Edit:
+ return Edit(
+ id=edit_id,
+ pillar=EditPillar.AGENT,
+ op=EditOp.REPLACE_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"new_content": "New prompt.\n"},
+ rationale="Better prompt",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+
+
+def _make_lora_edit(edit_id: str = "edit-lora") -> Edit:
+ return Edit(
+ id=edit_id,
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.LORA_FINETUNE,
+ target="models.qwen",
+ payload={"target_model": "qwen", "data_source": "all"},
+ rationale="Fine tune",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.MANUAL,
+ )
+
+
+class TestExecuteEdits:
+ """Tests for execute_edits()."""
+
+ def test_applies_auto_tier_edit(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[_make_auto_edit()],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.TIERED,
+ )
+ assert len(outcomes) == 1
+ assert outcomes[0].status == "applied"
+
+ def test_review_edit_goes_to_pending_in_tiered_mode(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[_make_review_edit()],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.TIERED,
+ )
+ assert len(outcomes) == 1
+ assert outcomes[0].status == "pending_review"
+
+ def test_review_edit_applied_in_auto_mode(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[_make_review_edit()],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.AUTO,
+ )
+ assert len(outcomes) == 1
+ assert outcomes[0].status == "applied"
+
+ def test_manual_tier_skipped(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[_make_lora_edit()],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.TIERED,
+ )
+ assert len(outcomes) == 1
+ assert outcomes[0].status == "skipped"
+
+ def test_all_edits_pending_in_manual_mode(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[_make_auto_edit()],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.MANUAL,
+ )
+ assert len(outcomes) == 1
+ assert outcomes[0].status == "pending_review"
+
+ def test_multiple_edits_processed(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.execute.loop import execute_edits
+
+ ctx = _make_ctx(tmp_path)
+ outcomes = execute_edits(
+ edits=[
+ _make_auto_edit("e1"),
+ _make_review_edit("e2"),
+ _make_lora_edit("e3"),
+ ],
+ ctx=ctx,
+ autonomy_mode=AutonomyMode.TIERED,
+ )
+ assert len(outcomes) == 3
+ assert outcomes[0].status == "applied"
+ assert outcomes[1].status == "pending_review"
+ assert outcomes[2].status == "skipped"
diff --git a/tests/learning/distillation/test_live_integration.py b/tests/learning/distillation/test_live_integration.py
new file mode 100644
index 00000000..e8ffd1ec
--- /dev/null
+++ b/tests/learning/distillation/test_live_integration.py
@@ -0,0 +1,177 @@
+"""Live integration tests for the distillation subsystem.
+
+These tests use REAL API calls (CloudEngine with Anthropic) and real
+TraceStore data. They are gated on the ``cloud`` marker — skip them
+with ``pytest -m "not cloud"``.
+
+Requires:
+- ANTHROPIC_API_KEY environment variable set
+- TraceStore at ~/.openjarvis/traces.db with some traces
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+# Skip entire module if no API key
+pytestmark = pytest.mark.cloud
+
+
+@pytest.fixture
+def anthropic_key():
+ key = os.environ.get("ANTHROPIC_API_KEY")
+ if not key:
+ pytest.skip("ANTHROPIC_API_KEY not set")
+ return key
+
+
+@pytest.fixture
+def cloud_engine(anthropic_key):
+ from openjarvis.engine.cloud import CloudEngine
+
+ return CloudEngine()
+
+
+@pytest.fixture
+def real_trace_store():
+ from openjarvis.traces.store import TraceStore
+
+ db_path = Path.home() / ".openjarvis" / "traces.db"
+ if not db_path.exists():
+ pytest.skip("No traces.db found at ~/.openjarvis/")
+ store = TraceStore(db_path)
+ if store.count() < 5:
+ pytest.skip("Need at least 5 traces for live test")
+ return store
+
+
+class TestCloudEngineDirectCall:
+ """Verify CloudEngine works with real API."""
+
+ def test_generate_produces_content(self, cloud_engine) -> None:
+ from openjarvis.core.types import Message, Role
+
+ result = cloud_engine.generate(
+ messages=[
+ Message(
+ role=Role.USER,
+ content="What is 2+2? Answer with just the number.",
+ )
+ ],
+ model="claude-sonnet-4-6",
+ max_tokens=10,
+ )
+ assert "content" in result
+ assert "4" in result["content"]
+ assert result.get("cost_usd", 0) > 0
+
+
+class TestTeacherAgentLive:
+ """Test TeacherAgent with a real CloudEngine."""
+
+ def test_teacher_agent_single_turn(self, cloud_engine) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ agent = TeacherAgent(
+ engine=cloud_engine,
+ model="claude-sonnet-4-6",
+ tools=[],
+ max_turns=2,
+ max_cost_usd=0.50,
+ )
+ result = agent.run(
+ "You are being tested. Simply respond with: 'TeacherAgent works.'",
+ system_prompt="You are a test assistant. Follow instructions exactly.",
+ )
+ assert result.content
+ assert result.turns >= 1
+ assert result.total_cost_usd > 0
+ print(f" Teacher response: {result.content[:100]}")
+ print(f" Cost: ${result.total_cost_usd:.4f}, Turns: {result.turns}")
+
+
+class TestDiagnosisRunnerLive:
+ """Test DiagnosisRunner with real CloudEngine + real traces."""
+
+ def test_diagnosis_produces_output(
+ self, cloud_engine, real_trace_store, tmp_path
+ ) -> None:
+ from openjarvis.learning.distillation.diagnose.runner import (
+ DiagnosisRunner,
+ )
+
+ # Create minimal config
+ config_dir = tmp_path / "oj_home"
+ config_dir.mkdir()
+ (config_dir / "config.toml").write_text("[learning]\nenabled = true\n")
+
+ session_dir = tmp_path / "session"
+ runner = DiagnosisRunner(
+ teacher_engine=cloud_engine,
+ teacher_model="claude-sonnet-4-6",
+ trace_store=real_trace_store,
+ benchmark_samples=[],
+ student_runner=lambda q, **kw: type(
+ "R", (), {"content": "mock", "score": 0.5}
+ )(),
+ judge=type("J", (), {"score_trace": lambda self, t: (0.5, "mock")})(),
+ session_dir=session_dir,
+ session_id="live-test-001",
+ config={
+ "config_path": config_dir / "config.toml",
+ "openjarvis_home": config_dir,
+ },
+ max_turns=5, # Keep it cheap
+ max_cost_usd=1.0,
+ )
+
+ result = runner.run()
+
+ # Diagnosis should produce output
+ assert result.diagnosis_md, "Diagnosis produced no markdown"
+ assert len(result.diagnosis_md) > 50, "Diagnosis too short"
+ print(f" Diagnosis length: {len(result.diagnosis_md)} chars")
+ print(f" Clusters found: {len(result.clusters)}")
+ print(f" Cost: ${result.cost_usd:.4f}")
+ print(f" Tool calls: {len(result.tool_call_records)}")
+
+ # Artifacts should exist
+ assert (session_dir / "diagnosis.md").exists()
+ assert (session_dir / "teacher_traces" / "diagnose.jsonl").exists()
+
+ # Print first 200 chars of diagnosis
+ print(f" Diagnosis preview: {result.diagnosis_md[:200]}...")
+
+
+class TestColdStartLive:
+ """Test cold start behavior with real trace store."""
+
+ def test_orchestrator_cold_start_with_no_feedback(
+ self, cloud_engine, real_trace_store, tmp_path
+ ) -> None:
+ """With 373 traces but 0 feedback, the orchestrator should handle
+ this gracefully — either by running (traces > 20) or by giving
+ a clear message about what's missing."""
+ from openjarvis.learning.distillation.gate.cold_start import (
+ check_benchmark_ready,
+ check_readiness,
+ )
+
+ # Check trace readiness
+ trace_ready = check_readiness(real_trace_store, min_traces=20)
+ print(f" Trace readiness: {trace_ready.ready} ({trace_ready.message})")
+
+ # Check benchmark readiness
+ bench_ready = check_benchmark_ready(
+ real_trace_store, min_feedback=0.7, min_samples=10
+ )
+ print(f" Benchmark readiness: {bench_ready.ready} ({bench_ready.message})")
+
+ # With 373 traces but 0 feedback: traces ready, benchmark not ready
+ assert trace_ready.ready, "Should have enough traces"
+ assert not bench_ready.ready, "Should not have enough high-feedback traces"
diff --git a/tests/learning/distillation/test_models.py b/tests/learning/distillation/test_models.py
new file mode 100644
index 00000000..c04f663e
--- /dev/null
+++ b/tests/learning/distillation/test_models.py
@@ -0,0 +1,586 @@
+"""Tests for openjarvis.learning.distillation.models module."""
+
+from __future__ import annotations
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class TestEditPillar:
+ """Tests for EditPillar enum."""
+
+ def test_has_four_pillars(self) -> None:
+ from openjarvis.learning.distillation.models import EditPillar
+
+ assert EditPillar.INTELLIGENCE.value == "intelligence"
+ assert EditPillar.AGENT.value == "agent"
+ assert EditPillar.TOOLS.value == "tools"
+ assert EditPillar.ENGINE.value == "engine"
+
+ def test_is_string_enum(self) -> None:
+ from openjarvis.learning.distillation.models import EditPillar
+
+ assert isinstance(EditPillar.AGENT, str)
+ assert EditPillar("agent") is EditPillar.AGENT
+
+
+class TestEditRiskTier:
+ """Tests for EditRiskTier enum."""
+
+ def test_has_three_tiers(self) -> None:
+ from openjarvis.learning.distillation.models import EditRiskTier
+
+ assert EditRiskTier.AUTO.value == "auto"
+ assert EditRiskTier.REVIEW.value == "review"
+ assert EditRiskTier.MANUAL.value == "manual"
+
+
+class TestEditOp:
+ """Tests for EditOp enum — must contain all v1 ops plus v2 placeholders."""
+
+ def test_intelligence_ops(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+
+ assert EditOp.SET_MODEL_FOR_QUERY_CLASS.value == "set_model_for_query_class"
+ assert EditOp.SET_MODEL_PARAM.value == "set_model_param"
+
+ def test_agent_ops(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+
+ assert EditOp.PATCH_SYSTEM_PROMPT.value == "patch_system_prompt"
+ assert EditOp.REPLACE_SYSTEM_PROMPT.value == "replace_system_prompt"
+ assert EditOp.SET_AGENT_CLASS.value == "set_agent_class"
+ assert EditOp.SET_AGENT_PARAM.value == "set_agent_param"
+ assert EditOp.EDIT_FEW_SHOT_EXEMPLARS.value == "edit_few_shot_exemplars"
+
+ def test_tools_ops(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+
+ assert EditOp.ADD_TOOL_TO_AGENT.value == "add_tool_to_agent"
+ assert EditOp.REMOVE_TOOL_FROM_AGENT.value == "remove_tool_from_agent"
+ assert EditOp.EDIT_TOOL_DESCRIPTION.value == "edit_tool_description"
+
+ def test_v2_placeholder_ops(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+
+ assert EditOp.LORA_FINETUNE.value == "lora_finetune"
+
+
+class TestTriggerKind:
+ """Tests for TriggerKind enum."""
+
+ def test_four_trigger_kinds(self) -> None:
+ from openjarvis.learning.distillation.models import TriggerKind
+
+ assert TriggerKind.SCHEDULED.value == "scheduled"
+ assert TriggerKind.CLUSTER.value == "cluster"
+ assert TriggerKind.USER_FLAG.value == "user_flag"
+ assert TriggerKind.ON_DEMAND.value == "on_demand"
+
+
+class TestAutonomyMode:
+ """Tests for AutonomyMode enum."""
+
+ def test_three_modes(self) -> None:
+ from openjarvis.learning.distillation.models import AutonomyMode
+
+ assert AutonomyMode.AUTO.value == "auto"
+ assert AutonomyMode.TIERED.value == "tiered"
+ assert AutonomyMode.MANUAL.value == "manual"
+
+
+class TestSessionStatus:
+ """Tests for SessionStatus enum."""
+
+ def test_all_statuses(self) -> None:
+ from openjarvis.learning.distillation.models import SessionStatus
+
+ assert SessionStatus.INITIATED.value == "initiated"
+ assert SessionStatus.DIAGNOSING.value == "diagnosing"
+ assert SessionStatus.PLANNING.value == "planning"
+ assert SessionStatus.EXECUTING.value == "executing"
+ assert SessionStatus.AWAITING_REVIEW.value == "awaiting_review"
+ assert SessionStatus.COMPLETED.value == "completed"
+ assert SessionStatus.FAILED.value == "failed"
+ assert SessionStatus.ROLLED_BACK.value == "rolled_back"
+
+
+# ---------------------------------------------------------------------------
+# Edit
+# ---------------------------------------------------------------------------
+
+
+class TestEdit:
+ """Tests for Edit pydantic model."""
+
+ def _valid_edit_kwargs(self) -> dict:
+ from openjarvis.learning.distillation.models import (
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+
+ return {
+ "id": "11111111-2222-3333-4444-555555555555",
+ "pillar": EditPillar.INTELLIGENCE,
+ "op": EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ "target": "learning.routing.policy_map.math",
+ "payload": {"query_class": "math", "model": "qwen2.5-coder:14b"},
+ "rationale": "Math queries are misrouted to qwen-3b",
+ "expected_improvement": "math_failures cluster",
+ "risk_tier": EditRiskTier.AUTO,
+ "references": ["trace-001", "trace-002"],
+ }
+
+ def test_constructs_with_valid_fields(self) -> None:
+ from openjarvis.learning.distillation.models import Edit
+
+ edit = Edit(**self._valid_edit_kwargs())
+
+ assert edit.id == "11111111-2222-3333-4444-555555555555"
+ assert edit.target == "learning.routing.policy_map.math"
+ assert edit.payload == {"query_class": "math", "model": "qwen2.5-coder:14b"}
+ assert edit.references == ["trace-001", "trace-002"]
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import Edit
+
+ edit = Edit(**self._valid_edit_kwargs())
+ as_json = edit.model_dump_json()
+ restored = Edit.model_validate_json(as_json)
+
+ assert restored == edit
+
+ def test_pillar_must_be_valid_enum(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import Edit
+
+ kwargs = self._valid_edit_kwargs()
+ kwargs["pillar"] = "not_a_pillar"
+
+ with pytest.raises(ValidationError):
+ Edit(**kwargs)
+
+ def test_op_must_be_valid_enum(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import Edit
+
+ kwargs = self._valid_edit_kwargs()
+ kwargs["op"] = "not_an_op"
+
+ with pytest.raises(ValidationError):
+ Edit(**kwargs)
+
+ def test_payload_can_be_empty_dict(self) -> None:
+ from openjarvis.learning.distillation.models import Edit
+
+ kwargs = self._valid_edit_kwargs()
+ kwargs["payload"] = {}
+
+ edit = Edit(**kwargs)
+ assert edit.payload == {}
+
+ def test_references_default_empty_list(self) -> None:
+ from openjarvis.learning.distillation.models import Edit
+
+ kwargs = self._valid_edit_kwargs()
+ del kwargs["references"]
+
+ edit = Edit(**kwargs)
+ assert edit.references == []
+
+
+# ---------------------------------------------------------------------------
+# FailureCluster
+# ---------------------------------------------------------------------------
+
+
+class TestFailureCluster:
+ """Tests for FailureCluster pydantic model."""
+
+ def _valid_cluster_kwargs(self) -> dict:
+ return {
+ "id": "cluster-001",
+ "description": "Math word problems routed to qwen-3b",
+ "sample_trace_ids": ["trace-001", "trace-002", "trace-003"],
+ "student_failure_rate": 0.85,
+ "teacher_success_rate": 0.95,
+ "skill_gap": (
+ "Student lacks chain-of-thought reasoning on multi-step arithmetic."
+ ),
+ "addressed_by_edit_ids": ["edit-001", "edit-002"],
+ }
+
+ def test_constructs_with_valid_fields(self) -> None:
+ from openjarvis.learning.distillation.models import FailureCluster
+
+ cluster = FailureCluster(**self._valid_cluster_kwargs())
+
+ assert cluster.id == "cluster-001"
+ assert cluster.student_failure_rate == 0.85
+ assert cluster.teacher_success_rate == 0.95
+ assert len(cluster.sample_trace_ids) == 3
+ assert len(cluster.addressed_by_edit_ids) == 2
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import FailureCluster
+
+ cluster = FailureCluster(**self._valid_cluster_kwargs())
+ as_json = cluster.model_dump_json()
+ restored = FailureCluster.model_validate_json(as_json)
+
+ assert restored == cluster
+
+ def test_addressed_by_edit_ids_defaults_empty(self) -> None:
+ from openjarvis.learning.distillation.models import FailureCluster
+
+ kwargs = self._valid_cluster_kwargs()
+ del kwargs["addressed_by_edit_ids"]
+
+ cluster = FailureCluster(**kwargs)
+ assert cluster.addressed_by_edit_ids == []
+
+ def test_failure_rate_must_be_between_zero_and_one(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import FailureCluster
+
+ kwargs = self._valid_cluster_kwargs()
+ kwargs["student_failure_rate"] = 1.5
+
+ with pytest.raises(ValidationError):
+ FailureCluster(**kwargs)
+
+ def test_success_rate_must_be_between_zero_and_one(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import FailureCluster
+
+ kwargs = self._valid_cluster_kwargs()
+ kwargs["teacher_success_rate"] = -0.1
+
+ with pytest.raises(ValidationError):
+ FailureCluster(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# LearningPlan
+# ---------------------------------------------------------------------------
+
+
+class TestLearningPlan:
+ """Tests for LearningPlan pydantic model."""
+
+ def _valid_plan_kwargs(self) -> dict:
+ from datetime import datetime, timezone
+
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ FailureCluster,
+ )
+
+ cluster = FailureCluster(
+ id="cluster-001",
+ description="Math routed to qwen-3b",
+ sample_trace_ids=["t1", "t2", "t3"],
+ student_failure_rate=0.8,
+ teacher_success_rate=0.9,
+ skill_gap="needs CoT",
+ addressed_by_edit_ids=["edit-001"],
+ )
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ target="learning.routing.policy_map.math",
+ payload={"query_class": "math", "model": "qwen2.5-coder:14b"},
+ rationale="Math fails on small model",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO,
+ references=["t1"],
+ )
+ return {
+ "session_id": "session-001",
+ "diagnosis_summary": "## Diagnosis\nThe student misroutes math.",
+ "failure_clusters": [cluster],
+ "edits": [edit],
+ "teacher_model": "claude-opus-4-6",
+ "estimated_cost_usd": 1.42,
+ "created_at": datetime(2026, 4, 8, 14, 22, 1, tzinfo=timezone.utc),
+ }
+
+ def test_constructs_with_valid_fields(self) -> None:
+ from openjarvis.learning.distillation.models import LearningPlan
+
+ plan = LearningPlan(**self._valid_plan_kwargs())
+
+ assert plan.session_id == "session-001"
+ assert len(plan.failure_clusters) == 1
+ assert len(plan.edits) == 1
+ assert plan.teacher_model == "claude-opus-4-6"
+ assert plan.estimated_cost_usd == 1.42
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import LearningPlan
+
+ plan = LearningPlan(**self._valid_plan_kwargs())
+ as_json = plan.model_dump_json()
+ restored = LearningPlan.model_validate_json(as_json)
+
+ assert restored == plan
+
+ def test_empty_clusters_and_edits_allowed(self) -> None:
+ # An aborted session may produce a plan with no clusters and no edits.
+ from openjarvis.learning.distillation.models import LearningPlan
+
+ kwargs = self._valid_plan_kwargs()
+ kwargs["failure_clusters"] = []
+ kwargs["edits"] = []
+
+ plan = LearningPlan(**kwargs)
+ assert plan.failure_clusters == []
+ assert plan.edits == []
+
+ def test_estimated_cost_must_be_non_negative(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import LearningPlan
+
+ kwargs = self._valid_plan_kwargs()
+ kwargs["estimated_cost_usd"] = -1.0
+
+ with pytest.raises(ValidationError):
+ LearningPlan(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# BenchmarkSnapshot
+# ---------------------------------------------------------------------------
+
+
+class TestBenchmarkSnapshot:
+ """Tests for BenchmarkSnapshot pydantic model."""
+
+ def _valid_snapshot_kwargs(self) -> dict:
+ return {
+ "benchmark_version": "personal_v3",
+ "overall_score": 0.72,
+ "cluster_scores": {"cluster-001": 0.65, "cluster-002": 0.80},
+ "task_count": 50,
+ "elapsed_seconds": 184.3,
+ }
+
+ def test_constructs_with_valid_fields(self) -> None:
+ from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+ snap = BenchmarkSnapshot(**self._valid_snapshot_kwargs())
+
+ assert snap.benchmark_version == "personal_v3"
+ assert snap.overall_score == 0.72
+ assert snap.cluster_scores["cluster-001"] == 0.65
+ assert snap.task_count == 50
+ assert snap.elapsed_seconds == 184.3
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+ snap = BenchmarkSnapshot(**self._valid_snapshot_kwargs())
+ restored = BenchmarkSnapshot.model_validate_json(snap.model_dump_json())
+
+ assert restored == snap
+
+ def test_score_bounds(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+ kwargs = self._valid_snapshot_kwargs()
+ kwargs["overall_score"] = 1.5
+
+ with pytest.raises(ValidationError):
+ BenchmarkSnapshot(**kwargs)
+
+ def test_task_count_must_be_non_negative(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+ kwargs = self._valid_snapshot_kwargs()
+ kwargs["task_count"] = -1
+
+ with pytest.raises(ValidationError):
+ BenchmarkSnapshot(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# EditOutcome
+# ---------------------------------------------------------------------------
+
+
+class TestEditOutcome:
+ """Tests for EditOutcome pydantic model."""
+
+ def _valid_outcome_kwargs(self) -> dict:
+ from datetime import datetime, timezone
+
+ return {
+ "edit_id": "edit-001",
+ "status": "applied",
+ "benchmark_delta": 0.04,
+ "cluster_deltas": {"cluster-001": 0.10, "cluster-002": 0.0},
+ "error": None,
+ "applied_at": datetime(2026, 4, 8, 14, 25, 0, tzinfo=timezone.utc),
+ }
+
+ def test_applied_outcome(self) -> None:
+ from openjarvis.learning.distillation.models import EditOutcome
+
+ outcome = EditOutcome(**self._valid_outcome_kwargs())
+ assert outcome.status == "applied"
+ assert outcome.benchmark_delta == 0.04
+
+ def test_rejected_outcome_has_no_applied_at(self) -> None:
+ from openjarvis.learning.distillation.models import EditOutcome
+
+ outcome = EditOutcome(
+ edit_id="edit-002",
+ status="rejected_by_gate",
+ benchmark_delta=-0.02,
+ cluster_deltas={},
+ error="regression: cluster-001 dropped 0.06",
+ applied_at=None,
+ )
+ assert outcome.status == "rejected_by_gate"
+ assert outcome.applied_at is None
+ assert outcome.error is not None
+
+ def test_status_must_be_valid_literal(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import EditOutcome
+
+ kwargs = self._valid_outcome_kwargs()
+ kwargs["status"] = "totally_made_up"
+
+ with pytest.raises(ValidationError):
+ EditOutcome(**kwargs)
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import EditOutcome
+
+ outcome = EditOutcome(**self._valid_outcome_kwargs())
+ restored = EditOutcome.model_validate_json(outcome.model_dump_json())
+
+ assert restored == outcome
+
+
+# ---------------------------------------------------------------------------
+# LearningSession
+# ---------------------------------------------------------------------------
+
+
+class TestLearningSession:
+ """Tests for LearningSession pydantic model."""
+
+ def _valid_session_kwargs(self) -> dict:
+ from datetime import datetime, timezone
+ from pathlib import Path
+
+ from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ BenchmarkSnapshot,
+ SessionStatus,
+ TriggerKind,
+ )
+
+ snap = BenchmarkSnapshot(
+ benchmark_version="personal_v1",
+ overall_score=0.65,
+ cluster_scores={"cluster-001": 0.50},
+ task_count=30,
+ elapsed_seconds=92.0,
+ )
+ return {
+ "id": "session-001",
+ "parent_session_id": None,
+ "trigger": TriggerKind.SCHEDULED,
+ "trigger_metadata": {"cron": "0 3 * * *"},
+ "status": SessionStatus.INITIATED,
+ "autonomy_mode": AutonomyMode.TIERED,
+ "started_at": datetime(2026, 4, 8, 3, 0, 0, tzinfo=timezone.utc),
+ "ended_at": None,
+ "diagnosis_path": Path("/tmp/sessions/session-001/diagnosis.md"),
+ "plan_path": Path("/tmp/sessions/session-001/plan.json"),
+ "benchmark_before": snap,
+ "benchmark_after": None,
+ "edit_outcomes": [],
+ "git_checkpoint_pre": "abc1234",
+ "git_checkpoint_post": None,
+ "teacher_cost_usd": 0.0,
+ "error": None,
+ }
+
+ def test_constructs_with_valid_fields(self) -> None:
+ from openjarvis.learning.distillation.models import LearningSession
+
+ session = LearningSession(**self._valid_session_kwargs())
+ assert session.id == "session-001"
+ assert session.parent_session_id is None
+ assert session.git_checkpoint_pre == "abc1234"
+ assert session.benchmark_after is None
+
+ def test_round_trip_via_json(self) -> None:
+ from openjarvis.learning.distillation.models import LearningSession
+
+ session = LearningSession(**self._valid_session_kwargs())
+ as_json = session.model_dump_json()
+ restored = LearningSession.model_validate_json(as_json)
+
+ assert restored == session
+
+ def test_supports_parent_session_chain(self) -> None:
+ from openjarvis.learning.distillation.models import LearningSession
+
+ kwargs = self._valid_session_kwargs()
+ kwargs["parent_session_id"] = "session-000"
+
+ session = LearningSession(**kwargs)
+ assert session.parent_session_id == "session-000"
+
+ def test_status_must_be_valid_enum(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import LearningSession
+
+ kwargs = self._valid_session_kwargs()
+ kwargs["status"] = "not_a_status"
+
+ with pytest.raises(ValidationError):
+ LearningSession(**kwargs)
+
+ def test_teacher_cost_must_be_non_negative(self) -> None:
+ import pytest
+ from pydantic import ValidationError
+
+ from openjarvis.learning.distillation.models import LearningSession
+
+ kwargs = self._valid_session_kwargs()
+ kwargs["teacher_cost_usd"] = -0.01
+
+ with pytest.raises(ValidationError):
+ LearningSession(**kwargs)
diff --git a/tests/learning/distillation/test_orchestrator.py b/tests/learning/distillation/test_orchestrator.py
new file mode 100644
index 00000000..f489d03d
--- /dev/null
+++ b/tests/learning/distillation/test_orchestrator.py
@@ -0,0 +1,174 @@
+"""Tests for DistillationOrchestrator — full session with mocks."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ BenchmarkSnapshot,
+ FailureCluster,
+ SessionStatus,
+)
+from openjarvis.learning.distillation.triggers import OnDemandTrigger
+
+
+def _make_snapshot(overall: float = 0.6) -> BenchmarkSnapshot:
+ return BenchmarkSnapshot(
+ benchmark_version="personal_v1",
+ overall_score=overall,
+ cluster_scores={"c1": overall},
+ task_count=10,
+ elapsed_seconds=5.0,
+ )
+
+
+def _make_diagnosis_result():
+ from openjarvis.learning.distillation.diagnose.runner import DiagnosisResult
+
+ return DiagnosisResult(
+ diagnosis_md="## Diagnosis\nMath routing is broken.",
+ clusters=[
+ FailureCluster(
+ id="c1",
+ description="Math routing",
+ sample_trace_ids=["t1", "t2", "t3"],
+ student_failure_rate=0.8,
+ teacher_success_rate=0.95,
+ skill_gap="needs CoT",
+ )
+ ],
+ cost_usd=0.05,
+ tool_call_records=[],
+ )
+
+
+def _make_mock_engine():
+ engine = MagicMock()
+ engine.generate.return_value = {
+ "content": json.dumps(
+ {
+ "edits": [
+ {
+ "id": "edit-001",
+ "pillar": "intelligence",
+ "op": "set_model_for_query_class",
+ "target": "routing.math",
+ "payload": {
+ "query_class": "math",
+ "model": "qwen2.5-coder:14b",
+ },
+ "rationale": "Route math to bigger model",
+ "expected_improvement": "c1",
+ "risk_tier": "auto",
+ "references": ["t1"],
+ }
+ ]
+ }
+ ),
+ "usage": {"total_tokens": 500},
+ "cost_usd": 0.03,
+ "finish_reason": "stop",
+ }
+ return engine
+
+
+class TestDistillationOrchestrator:
+ def test_full_session_completes(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.orchestrator import (
+ DistillationOrchestrator,
+ )
+
+ orch = DistillationOrchestrator(
+ teacher_engine=_make_mock_engine(),
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(count=MagicMock(return_value=30)),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_store=MagicMock(),
+ checkpoint_store=MagicMock(
+ current_sha=MagicMock(return_value="abc123"),
+ begin_stage=MagicMock(return_value=MagicMock(pre_stage_sha="abc123")),
+ ),
+ openjarvis_home=tmp_path,
+ autonomy_mode=AutonomyMode.AUTO,
+ scorer=lambda **kw: _make_snapshot(0.65),
+ benchmark_version="personal_v1",
+ )
+
+ with patch(
+ "openjarvis.learning.distillation.orchestrator.DiagnosisRunner"
+ ) as MockDiag:
+ MockDiag.return_value.run.return_value = _make_diagnosis_result()
+ session = orch.run(OnDemandTrigger())
+
+ assert session.status in (
+ SessionStatus.COMPLETED,
+ SessionStatus.AWAITING_REVIEW,
+ )
+ assert session.teacher_cost_usd >= 0
+
+ def test_cold_start_returns_failed(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.orchestrator import (
+ DistillationOrchestrator,
+ )
+
+ trace_store = MagicMock()
+ trace_store.count.return_value = 5 # Not enough
+
+ orch = DistillationOrchestrator(
+ teacher_engine=MagicMock(),
+ teacher_model="claude-opus-4-6",
+ trace_store=trace_store,
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_store=MagicMock(),
+ checkpoint_store=MagicMock(
+ current_sha=MagicMock(return_value="abc123"),
+ ),
+ openjarvis_home=tmp_path,
+ autonomy_mode=AutonomyMode.TIERED,
+ scorer=lambda **kw: _make_snapshot(),
+ benchmark_version="personal_v1",
+ )
+
+ session = orch.run(OnDemandTrigger())
+ assert session.status == SessionStatus.FAILED
+ assert "not enough traces" in (session.error or "").lower()
+
+ def test_session_persisted_to_store(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.orchestrator import (
+ DistillationOrchestrator,
+ )
+
+ session_store = MagicMock()
+
+ orch = DistillationOrchestrator(
+ teacher_engine=_make_mock_engine(),
+ teacher_model="claude-opus-4-6",
+ trace_store=MagicMock(count=MagicMock(return_value=30)),
+ benchmark_samples=[],
+ student_runner=MagicMock(),
+ judge=MagicMock(),
+ session_store=session_store,
+ checkpoint_store=MagicMock(
+ current_sha=MagicMock(return_value="abc123"),
+ begin_stage=MagicMock(return_value=MagicMock(pre_stage_sha="abc123")),
+ ),
+ openjarvis_home=tmp_path,
+ autonomy_mode=AutonomyMode.AUTO,
+ scorer=lambda **kw: _make_snapshot(0.65),
+ benchmark_version="personal_v1",
+ )
+
+ with patch(
+ "openjarvis.learning.distillation.orchestrator.DiagnosisRunner"
+ ) as MockDiag:
+ MockDiag.return_value.run.return_value = _make_diagnosis_result()
+ orch.run(OnDemandTrigger())
+
+ assert session_store.save_session.called
diff --git a/tests/learning/distillation/test_paths.py b/tests/learning/distillation/test_paths.py
new file mode 100644
index 00000000..b123a3b0
--- /dev/null
+++ b/tests/learning/distillation/test_paths.py
@@ -0,0 +1,107 @@
+"""Tests for openjarvis.learning.distillation.storage.paths module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# resolve_distillation_root
+# ---------------------------------------------------------------------------
+
+
+class TestResolveDistillationRoot:
+ """Tests for resolve_distillation_root()."""
+
+ def test_default_is_under_home(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ monkeypatch.delenv("OPENJARVIS_HOME", raising=False)
+ result = paths.resolve_distillation_root()
+ assert result == Path.home() / ".openjarvis" / "learning"
+
+ def test_respects_openjarvis_home_env_var(
+ self,
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ custom = tmp_path / "custom_oj"
+ monkeypatch.setenv("OPENJARVIS_HOME", str(custom))
+ result = paths.resolve_distillation_root()
+ assert result == custom / "learning"
+
+ def test_returns_absolute_path(
+ self,
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ monkeypatch.setenv("OPENJARVIS_HOME", str(tmp_path / "rel"))
+ result = paths.resolve_distillation_root()
+ assert result.is_absolute()
+
+ def test_rejects_path_inside_source_tree(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ # Find the OpenJarvis source root by walking up from the paths module.
+ source_root = paths._find_source_root()
+ assert source_root is not None # We must be running inside the repo.
+
+ # Force OPENJARVIS_HOME to point inside the source tree.
+ monkeypatch.setenv("OPENJARVIS_HOME", str(source_root / "junk_dir"))
+
+ with pytest.raises(paths.ConfigurationError, match="inside the source tree"):
+ paths.resolve_distillation_root()
+
+ def test_find_source_root_returns_repo_root(self) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ result = paths._find_source_root()
+ assert result is not None
+ assert (result / "pyproject.toml").exists()
+
+
+# ---------------------------------------------------------------------------
+# ensure_distillation_dirs
+# ---------------------------------------------------------------------------
+
+
+class TestEnsureDistillationDirs:
+ """Tests for ensure_distillation_dirs()."""
+
+ def test_creates_subdirs(
+ self,
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ monkeypatch.setenv("OPENJARVIS_HOME", str(tmp_path / "oj"))
+ root = paths.ensure_distillation_dirs()
+
+ assert root.exists()
+ assert (root / "sessions").exists()
+ assert (root / "benchmarks").exists()
+ assert (root / "benchmarks" / "reference_outputs").exists()
+ assert (root / "pending_review").exists()
+
+ def test_idempotent(
+ self,
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ ) -> None:
+ from openjarvis.learning.distillation.storage import paths
+
+ monkeypatch.setenv("OPENJARVIS_HOME", str(tmp_path / "oj"))
+ first = paths.ensure_distillation_dirs()
+ second = paths.ensure_distillation_dirs()
+
+ assert first == second
+ assert first.exists()
diff --git a/tests/learning/distillation/test_pending_queue.py b/tests/learning/distillation/test_pending_queue.py
new file mode 100644
index 00000000..ba67d171
--- /dev/null
+++ b/tests/learning/distillation/test_pending_queue.py
@@ -0,0 +1,73 @@
+"""Tests for openjarvis.learning.distillation.pending_queue module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+)
+
+
+def _make_edit(edit_id: str = "edit-001") -> Edit:
+ return Edit(
+ id=edit_id,
+ pillar=EditPillar.AGENT,
+ op=EditOp.REPLACE_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"new_content": "New prompt.\n"},
+ rationale="Better prompt",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+
+
+class TestPendingQueue:
+ def test_enqueue_creates_file(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.pending_queue import PendingQueue
+
+ queue = PendingQueue(tmp_path / "pending_review")
+ queue.enqueue("session-001", _make_edit())
+ files = list((tmp_path / "pending_review").glob("*.json"))
+ assert len(files) == 1
+ assert "session-001" in files[0].name
+ assert "edit-001" in files[0].name
+
+ def test_list_pending(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.pending_queue import PendingQueue
+
+ queue = PendingQueue(tmp_path / "pending_review")
+ queue.enqueue("session-001", _make_edit("e1"))
+ queue.enqueue("session-001", _make_edit("e2"))
+ pending = queue.list_pending()
+ assert len(pending) == 2
+ ids = {p["edit"]["id"] for p in pending}
+ assert ids == {"e1", "e2"}
+
+ def test_resolve_removes_file(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.pending_queue import PendingQueue
+
+ queue = PendingQueue(tmp_path / "pending_review")
+ queue.enqueue("session-001", _make_edit())
+ assert len(queue.list_pending()) == 1
+ queue.resolve("session-001", "edit-001")
+ assert len(queue.list_pending()) == 0
+
+ def test_list_empty_queue(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.pending_queue import PendingQueue
+
+ queue = PendingQueue(tmp_path / "pending_review")
+ assert queue.list_pending() == []
+
+ def test_get_pending_edit(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.pending_queue import PendingQueue
+
+ queue = PendingQueue(tmp_path / "pending_review")
+ queue.enqueue("session-001", _make_edit())
+ edit_data = queue.get("session-001", "edit-001")
+ assert edit_data is not None
+ assert edit_data["edit"]["id"] == "edit-001"
+ assert edit_data["session_id"] == "session-001"
diff --git a/tests/learning/distillation/test_planner.py b/tests/learning/distillation/test_planner.py
new file mode 100644
index 00000000..be7602fb
--- /dev/null
+++ b/tests/learning/distillation/test_planner.py
@@ -0,0 +1,272 @@
+"""Tests for openjarvis.learning.distillation.plan.planner module.
+
+All tests use mocked CloudEngine — no live API calls.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock
+
+from openjarvis.learning.distillation.models import (
+ EditOp,
+ EditRiskTier,
+ FailureCluster,
+)
+
+
+def _make_clusters() -> list[FailureCluster]:
+ return [
+ FailureCluster(
+ id="cluster-001",
+ description="Math queries routed to qwen-3b",
+ sample_trace_ids=["t1", "t2", "t3"],
+ student_failure_rate=0.8,
+ teacher_success_rate=0.95,
+ skill_gap="Student lacks chain-of-thought on multi-step math",
+ ),
+ FailureCluster(
+ id="cluster-002",
+ description="Calculator tool not used",
+ sample_trace_ids=["t4", "t5", "t6"],
+ student_failure_rate=0.6,
+ teacher_success_rate=0.9,
+ skill_gap="Student does not invoke calculator",
+ ),
+ ]
+
+
+def _make_teacher_response(edits_json: list[dict]) -> dict:
+ """Create a mock engine.generate() response with edit list."""
+ return {
+ "content": json.dumps({"edits": edits_json}),
+ "usage": {"prompt_tokens": 500, "completion_tokens": 300, "total_tokens": 800},
+ "cost_usd": 0.03,
+ "finish_reason": "stop",
+ }
+
+
+def _make_edit_dict(
+ edit_id: str = "edit-001",
+ pillar: str = "intelligence",
+ op: str = "set_model_for_query_class",
+ target: str = "learning.routing.policy_map.math",
+ payload: dict | None = None,
+ expected_improvement: str = "cluster-001",
+) -> dict:
+ return {
+ "id": edit_id,
+ "pillar": pillar,
+ "op": op,
+ "target": target,
+ "payload": payload or {"query_class": "math", "model": "qwen2.5-coder:14b"},
+ "rationale": "Route math queries to a bigger model",
+ "expected_improvement": expected_improvement,
+ "risk_tier": "auto",
+ "references": ["t1", "t2"],
+ }
+
+
+class TestLearningPlanner:
+ """Tests for LearningPlanner."""
+
+ def test_produces_learning_plan(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_teacher_response([_make_edit_dict()])
+
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=tmp_path / "session-001",
+ prompt_reader=lambda t: "",
+ )
+ plan = planner.run(
+ diagnosis_md="## Diagnosis\nMath routing is broken.",
+ clusters=_make_clusters(),
+ )
+
+ assert plan.session_id == "session-001"
+ assert len(plan.edits) == 1
+ assert plan.edits[0].op == EditOp.SET_MODEL_FOR_QUERY_CLASS
+ assert plan.teacher_model == "claude-opus-4-6"
+
+ def test_assigns_risk_tiers(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ # Teacher incorrectly sets MANUAL for an auto-tier op
+ edit_dict = _make_edit_dict()
+ edit_dict["risk_tier"] = "manual"
+ engine.generate.return_value = _make_teacher_response([edit_dict])
+
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=tmp_path / "session-001",
+ prompt_reader=lambda t: "",
+ )
+ plan = planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=_make_clusters(),
+ )
+
+ # Should be overwritten to AUTO
+ assert plan.edits[0].risk_tier == EditRiskTier.AUTO
+
+ def test_persists_plan_json(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_teacher_response([_make_edit_dict()])
+
+ session_dir = tmp_path / "session-001"
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=session_dir,
+ prompt_reader=lambda t: "",
+ )
+ planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=_make_clusters(),
+ )
+
+ plan_path = session_dir / "plan.json"
+ assert plan_path.exists()
+ data = json.loads(plan_path.read_text())
+ assert data["session_id"] == "session-001"
+
+ def test_drops_cluster_with_zero_rates(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_teacher_response([_make_edit_dict()])
+
+ clusters = [
+ FailureCluster(
+ id="cluster-bad",
+ description="No evidence",
+ sample_trace_ids=[],
+ student_failure_rate=0.0,
+ teacher_success_rate=0.0,
+ skill_gap="Speculative",
+ ),
+ FailureCluster(
+ id="cluster-good",
+ description="Real evidence",
+ sample_trace_ids=["t1", "t2", "t3"],
+ student_failure_rate=0.8,
+ teacher_success_rate=0.9,
+ skill_gap="Verified gap",
+ ),
+ ]
+
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=tmp_path / "session-001",
+ prompt_reader=lambda t: "",
+ )
+ plan = planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=clusters,
+ )
+
+ # cluster-bad should be dropped (marked in skill_gap)
+ bad = next(c for c in plan.failure_clusters if c.id == "cluster-bad")
+ assert "dropped" in bad.skill_gap.lower()
+ assert bad.addressed_by_edit_ids == []
+ # cluster-good should survive
+ good = next(c for c in plan.failure_clusters if c.id == "cluster-good")
+ assert "dropped" not in good.skill_gap.lower()
+
+ def test_all_clusters_dropped_returns_empty_edits(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_teacher_response([])
+
+ clusters = [
+ FailureCluster(
+ id="cluster-bad",
+ description="No evidence",
+ sample_trace_ids=[],
+ student_failure_rate=0.0,
+ teacher_success_rate=0.0,
+ skill_gap="Speculative",
+ ),
+ ]
+
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=tmp_path / "session-001",
+ prompt_reader=lambda t: "",
+ )
+ plan = planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=clusters,
+ )
+
+ assert plan.edits == []
+ assert plan.failure_clusters[0].addressed_by_edit_ids == []
+
+ def test_persists_teacher_trace_jsonl(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_teacher_response([_make_edit_dict()])
+
+ session_dir = tmp_path / "session-001"
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=session_dir,
+ prompt_reader=lambda t: "",
+ )
+ planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=_make_clusters(),
+ )
+
+ jsonl_path = session_dir / "teacher_traces" / "plan.jsonl"
+ assert jsonl_path.exists()
+ lines = jsonl_path.read_text().strip().splitlines()
+ assert len(lines) >= 1
+ record = json.loads(lines[0])
+ assert "cost_usd" in record
+
+ def test_handles_malformed_teacher_output(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.plan.planner import LearningPlanner
+
+ engine = MagicMock()
+ engine.generate.return_value = {
+ "content": "This is not valid JSON at all",
+ "usage": {"total_tokens": 100},
+ "cost_usd": 0.01,
+ "finish_reason": "stop",
+ }
+
+ planner = LearningPlanner(
+ teacher_engine=engine,
+ teacher_model="claude-opus-4-6",
+ session_id="session-001",
+ session_dir=tmp_path / "session-001",
+ prompt_reader=lambda t: "",
+ )
+ plan = planner.run(
+ diagnosis_md="## Diagnosis",
+ clusters=_make_clusters(),
+ )
+
+ # Should return a plan with no edits rather than crashing
+ assert plan.edits == []
diff --git a/tests/learning/distillation/test_prompt_diff.py b/tests/learning/distillation/test_prompt_diff.py
new file mode 100644
index 00000000..9ba38a0b
--- /dev/null
+++ b/tests/learning/distillation/test_prompt_diff.py
@@ -0,0 +1,212 @@
+"""Tests for openjarvis.learning.distillation.plan.prompt_diff module."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+class TestChangedLineRatio:
+ """Tests for changed_line_ratio()."""
+
+ def test_identical_strings(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ changed_line_ratio,
+ )
+
+ assert changed_line_ratio("hello\nworld\n", "hello\nworld\n") == 0.0
+
+ def test_completely_different(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ changed_line_ratio,
+ )
+
+ ratio = changed_line_ratio("aaa\nbbb\nccc\n", "xxx\nyyy\nzzz\n")
+ assert ratio == 1.0
+
+ def test_partial_change(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ changed_line_ratio,
+ )
+
+ original = "line1\nline2\nline3\nline4\n"
+ modified = "line1\nchanged\nline3\nline4\n"
+ ratio = changed_line_ratio(original, modified)
+ # 1 out of 4 lines changed
+ assert 0.2 <= ratio <= 0.35
+
+ def test_empty_original(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ changed_line_ratio,
+ )
+
+ # Adding to empty = 100% change
+ ratio = changed_line_ratio("", "new content\n")
+ assert ratio == 1.0
+
+ def test_empty_both(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ changed_line_ratio,
+ )
+
+ assert changed_line_ratio("", "") == 0.0
+
+
+class TestApplyUnifiedDiff:
+ """Tests for apply_unified_diff()."""
+
+ def test_applies_simple_patch(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ apply_unified_diff,
+ )
+
+ original = "line1\nline2\nline3\n"
+ diff = (
+ "--- a/prompt.md\n"
+ "+++ b/prompt.md\n"
+ "@@ -1,3 +1,3 @@\n"
+ " line1\n"
+ "-line2\n"
+ "+changed_line2\n"
+ " line3\n"
+ )
+ result = apply_unified_diff(original, diff)
+ assert result == "line1\nchanged_line2\nline3\n"
+
+ def test_returns_none_on_bad_diff(self) -> None:
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ apply_unified_diff,
+ )
+
+ result = apply_unified_diff("hello\n", "not a valid diff")
+ assert result is None
+
+
+class TestMaybeDowngradeToReplace:
+ """Tests for maybe_downgrade_to_replace()."""
+
+ def test_non_patch_op_passes_through(self) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ maybe_downgrade_to_replace,
+ )
+
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ target="routing.math",
+ payload={"model": "qwen2.5-coder:14b"},
+ rationale="test",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO,
+ )
+ result = maybe_downgrade_to_replace(edit, prompt_reader=lambda t: "")
+ assert result.op == EditOp.SET_MODEL_FOR_QUERY_CLASS
+
+ def test_small_diff_stays_patch(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ maybe_downgrade_to_replace,
+ )
+
+ original = (
+ "line1\nline2\nline3\nline4\nline5\nline6\nline7\nline8\nline9\nline10\n"
+ )
+ diff = (
+ "--- a/prompt.md\n"
+ "+++ b/prompt.md\n"
+ "@@ -2,1 +2,1 @@\n"
+ " line1\n"
+ "-line2\n"
+ "+changed\n"
+ " line3\n"
+ )
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": diff},
+ rationale="Small fix",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ result = maybe_downgrade_to_replace(edit, prompt_reader=lambda t: original)
+ assert result.op == EditOp.PATCH_SYSTEM_PROMPT
+
+ def test_large_diff_downgrades_to_replace(self) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ maybe_downgrade_to_replace,
+ )
+
+ original = "old1\nold2\nold3\nold4\n"
+ diff = (
+ "--- a/prompt.md\n"
+ "+++ b/prompt.md\n"
+ "@@ -1,4 +1,4 @@\n"
+ "-old1\n"
+ "-old2\n"
+ "-old3\n"
+ "-old4\n"
+ "+new1\n"
+ "+new2\n"
+ "+new3\n"
+ "+new4\n"
+ )
+ edit = Edit(
+ id="edit-002",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": diff},
+ rationale="Major rewrite",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ result = maybe_downgrade_to_replace(edit, prompt_reader=lambda t: original)
+ assert result.op == EditOp.REPLACE_SYSTEM_PROMPT
+ assert "new_content" in result.payload
+ assert "new1" in result.payload["new_content"]
+
+ def test_bad_diff_downgrades_to_replace_with_raw(self) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.prompt_diff import (
+ maybe_downgrade_to_replace,
+ )
+
+ edit = Edit(
+ id="edit-003",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": "this is not a valid unified diff"},
+ rationale="Bad diff",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.REVIEW,
+ )
+ result = maybe_downgrade_to_replace(
+ edit, prompt_reader=lambda t: "original content\n"
+ )
+ # Can't apply the diff, so it should downgrade
+ assert result.op == EditOp.REPLACE_SYSTEM_PROMPT
diff --git a/tests/learning/distillation/test_regression.py b/tests/learning/distillation/test_regression.py
new file mode 100644
index 00000000..7f1c1231
--- /dev/null
+++ b/tests/learning/distillation/test_regression.py
@@ -0,0 +1,89 @@
+"""Tests for openjarvis.learning.distillation.gate.regression module."""
+
+from __future__ import annotations
+
+from openjarvis.learning.distillation.models import BenchmarkSnapshot
+
+
+def _make_snapshot(
+ overall: float = 0.7,
+ clusters: dict[str, float] | None = None,
+) -> BenchmarkSnapshot:
+ return BenchmarkSnapshot(
+ benchmark_version="personal_v1",
+ overall_score=overall,
+ cluster_scores=clusters or {"c1": 0.6, "c2": 0.8},
+ task_count=50,
+ elapsed_seconds=60.0,
+ )
+
+
+class TestRegressionCheck:
+ """Tests for regression_check()."""
+
+ def test_no_regression_when_all_improve(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.6, clusters={"c1": 0.5, "c2": 0.7})
+ after = _make_snapshot(overall=0.7, clusters={"c1": 0.6, "c2": 0.8})
+ result = regression_check(before, after, max_regression=0.05)
+ assert not result.has_regression
+
+ def test_detects_cluster_regression(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.7, clusters={"c1": 0.6, "c2": 0.8})
+ after = _make_snapshot(overall=0.72, clusters={"c1": 0.65, "c2": 0.70})
+ result = regression_check(before, after, max_regression=0.05)
+ assert result.has_regression
+ assert "c2" in result.regressed_clusters
+
+ def test_small_drop_within_threshold(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.7, clusters={"c1": 0.6, "c2": 0.8})
+ after = _make_snapshot(overall=0.72, clusters={"c1": 0.63, "c2": 0.76})
+ result = regression_check(before, after, max_regression=0.05)
+ assert not result.has_regression
+
+ def test_new_cluster_in_after_not_flagged(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.7, clusters={"c1": 0.6})
+ after = _make_snapshot(overall=0.75, clusters={"c1": 0.65, "c2": 0.8})
+ result = regression_check(before, after, max_regression=0.05)
+ assert not result.has_regression
+
+ def test_missing_cluster_in_after_flagged(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.7, clusters={"c1": 0.6, "c2": 0.8})
+ after = _make_snapshot(overall=0.72, clusters={"c1": 0.65})
+ # c2 disappeared — treat as regression (score went from 0.8 to 0.0)
+ result = regression_check(before, after, max_regression=0.05)
+ assert result.has_regression
+ assert "c2" in result.regressed_clusters
+
+ def test_result_has_details(self) -> None:
+ from openjarvis.learning.distillation.gate.regression import (
+ regression_check,
+ )
+
+ before = _make_snapshot(overall=0.7, clusters={"c1": 0.6, "c2": 0.8})
+ after = _make_snapshot(overall=0.68, clusters={"c1": 0.55, "c2": 0.75})
+ result = regression_check(before, after, max_regression=0.03)
+ assert result.has_regression
+ assert len(result.regressed_clusters) >= 1
+ # Check that deltas are provided
+ for cluster_id, delta in result.regressed_clusters.items():
+ assert delta < 0
diff --git a/tests/learning/distillation/test_risk_tier.py b/tests/learning/distillation/test_risk_tier.py
new file mode 100644
index 00000000..d7510e6d
--- /dev/null
+++ b/tests/learning/distillation/test_risk_tier.py
@@ -0,0 +1,120 @@
+"""Tests for openjarvis.learning.distillation.plan.risk_tier module."""
+
+from __future__ import annotations
+
+
+class TestTierTable:
+ """Tests for TIER_TABLE completeness."""
+
+ def test_every_edit_op_has_a_tier(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+ from openjarvis.learning.distillation.plan.risk_tier import TIER_TABLE
+
+ for op in EditOp:
+ assert op in TIER_TABLE, f"Missing tier for {op}"
+
+ def test_no_extra_keys(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp
+ from openjarvis.learning.distillation.plan.risk_tier import TIER_TABLE
+
+ for key in TIER_TABLE:
+ assert key in EditOp, f"Extra key in TIER_TABLE: {key}"
+
+
+class TestAssignTier:
+ """Tests for assign_tier() function."""
+
+ def test_intelligence_ops_are_auto(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.SET_MODEL_FOR_QUERY_CLASS) == EditRiskTier.AUTO
+ assert assign_tier(EditOp.SET_MODEL_PARAM) == EditRiskTier.AUTO
+
+ def test_tool_ops_are_auto(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.ADD_TOOL_TO_AGENT) == EditRiskTier.AUTO
+ assert assign_tier(EditOp.REMOVE_TOOL_FROM_AGENT) == EditRiskTier.AUTO
+ assert assign_tier(EditOp.EDIT_TOOL_DESCRIPTION) == EditRiskTier.AUTO
+
+ def test_agent_param_is_auto(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.SET_AGENT_PARAM) == EditRiskTier.AUTO
+
+ def test_prompt_ops_are_review(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.PATCH_SYSTEM_PROMPT) == EditRiskTier.REVIEW
+ assert assign_tier(EditOp.REPLACE_SYSTEM_PROMPT) == EditRiskTier.REVIEW
+
+ def test_agent_class_is_review(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.SET_AGENT_CLASS) == EditRiskTier.REVIEW
+
+ def test_few_shot_is_review(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.EDIT_FEW_SHOT_EXEMPLARS) == EditRiskTier.REVIEW
+
+ def test_lora_is_manual(self) -> None:
+ from openjarvis.learning.distillation.models import EditOp, EditRiskTier
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tier
+
+ assert assign_tier(EditOp.LORA_FINETUNE) == EditRiskTier.MANUAL
+
+
+class TestAssignTiers:
+ """Tests for assign_tiers() batch function."""
+
+ def test_overwrites_teacher_tier(self) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tiers
+
+ # Teacher incorrectly sets AUTO for a prompt edit
+ edit = Edit(
+ id="edit-001",
+ pillar=EditPillar.AGENT,
+ op=EditOp.PATCH_SYSTEM_PROMPT,
+ target="agents.simple.system_prompt",
+ payload={"diff": "--- a\n+++ b\n@@ -1 +1 @@\n-old\n+new"},
+ rationale="Improve prompt",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO, # Wrong — should be REVIEW
+ )
+ result = assign_tiers([edit])
+ assert result[0].risk_tier == EditRiskTier.REVIEW
+
+ def test_preserves_correct_tier(self) -> None:
+ from openjarvis.learning.distillation.models import (
+ Edit,
+ EditOp,
+ EditPillar,
+ EditRiskTier,
+ )
+ from openjarvis.learning.distillation.plan.risk_tier import assign_tiers
+
+ edit = Edit(
+ id="edit-002",
+ pillar=EditPillar.INTELLIGENCE,
+ op=EditOp.SET_MODEL_FOR_QUERY_CLASS,
+ target="learning.routing.policy_map.math",
+ payload={"query_class": "math", "model": "qwen2.5-coder:14b"},
+ rationale="Route math to bigger model",
+ expected_improvement="cluster-001",
+ risk_tier=EditRiskTier.AUTO, # Correct
+ )
+ result = assign_tiers([edit])
+ assert result[0].risk_tier == EditRiskTier.AUTO
diff --git a/tests/learning/distillation/test_rollback.py b/tests/learning/distillation/test_rollback.py
new file mode 100644
index 00000000..4219159a
--- /dev/null
+++ b/tests/learning/distillation/test_rollback.py
@@ -0,0 +1,87 @@
+"""Integration test: apply session then rollback via CheckpointStore."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+
+def _git(cwd: Path, *args: str) -> str:
+ result = subprocess.run(
+ ["git", *args], cwd=cwd, capture_output=True, text=True, check=True
+ )
+ return result.stdout.strip()
+
+
+def _setup_config_tree(root: Path) -> None:
+ (root / "agents" / "simple").mkdir(parents=True)
+ (root / "tools").mkdir(parents=True)
+ (root / "config.toml").write_text("[learning]\nenabled = true\n")
+ (root / "agents" / "simple" / "system_prompt.md").write_text(
+ "You are a helpful assistant.\n"
+ )
+ (root / "tools" / "descriptions.toml").write_text("[web_search]\n")
+
+
+class TestRollbackIntegration:
+ def test_rollback_restores_files(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = tmp_path / "oj_home"
+ _setup_config_tree(root)
+
+ store = CheckpointStore(root)
+ store.init()
+ original_prompt = (root / "agents" / "simple" / "system_prompt.md").read_text()
+
+ # Simulate two edits in a session
+ h1 = store.begin_stage("edit-001")
+ prompt_file = root / "agents" / "simple" / "system_prompt.md"
+ prompt_file.write_text("Modified prompt v1.\n")
+ store.commit_stage(
+ h1,
+ message="learning: edit-001",
+ session_id="session-A",
+ risk_tier="auto",
+ )
+
+ h2 = store.begin_stage("edit-002")
+ tools_file = root / "tools" / "descriptions.toml"
+ tools_file.write_text("[web_search]\nupdated = true\n")
+ store.commit_stage(
+ h2,
+ message="learning: edit-002",
+ session_id="session-A",
+ risk_tier="auto",
+ )
+
+ # Verify files changed
+ assert prompt_file.read_text() == "Modified prompt v1.\n"
+
+ # Rollback the session
+ revert_shas = store.revert_session("session-A")
+ assert len(revert_shas) == 2
+
+ # Files restored
+ assert prompt_file.read_text() == original_prompt
+ assert tools_file.read_text() == "[web_search]\n"
+
+ # History preserved (reverts are new commits)
+ log_count = len(_git(root, "log", "--oneline").splitlines())
+ assert log_count >= 5 # baseline + 2 edits + 2 reverts
+
+ def test_rollback_nonexistent_session(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.checkpoint.store import (
+ CheckpointStore,
+ )
+
+ root = tmp_path / "oj_home"
+ _setup_config_tree(root)
+
+ store = CheckpointStore(root)
+ store.init()
+
+ result = store.revert_session("nonexistent-session")
+ assert result == []
diff --git a/tests/learning/distillation/test_session_store.py b/tests/learning/distillation/test_session_store.py
new file mode 100644
index 00000000..9243a97b
--- /dev/null
+++ b/tests/learning/distillation/test_session_store.py
@@ -0,0 +1,285 @@
+"""Tests for openjarvis.learning.distillation.storage.session_store module."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from pathlib import Path
+
+from openjarvis.learning.distillation.models import (
+ AutonomyMode,
+ BenchmarkSnapshot,
+ EditOutcome,
+ LearningSession,
+ SessionStatus,
+ TriggerKind,
+)
+
+
+def _make_session(
+ session_id: str = "session-001",
+ parent_session_id: str | None = None,
+ status: SessionStatus = SessionStatus.INITIATED,
+ teacher_cost_usd: float = 0.0,
+) -> LearningSession:
+ snap = BenchmarkSnapshot(
+ benchmark_version="personal_v1",
+ overall_score=0.6,
+ cluster_scores={"cluster-001": 0.5},
+ task_count=20,
+ elapsed_seconds=60.0,
+ )
+ return LearningSession(
+ id=session_id,
+ parent_session_id=parent_session_id,
+ trigger=TriggerKind.SCHEDULED,
+ trigger_metadata={},
+ status=status,
+ autonomy_mode=AutonomyMode.TIERED,
+ started_at=datetime(2026, 4, 8, 3, 0, 0, tzinfo=timezone.utc),
+ ended_at=None,
+ diagnosis_path=Path(f"/tmp/{session_id}/diagnosis.md"),
+ plan_path=Path(f"/tmp/{session_id}/plan.json"),
+ benchmark_before=snap,
+ benchmark_after=None,
+ edit_outcomes=[],
+ git_checkpoint_pre="abc1234",
+ git_checkpoint_post=None,
+ teacher_cost_usd=teacher_cost_usd,
+ error=None,
+ )
+
+
+def _make_outcome(
+ edit_id: str = "edit-001",
+ pillar: str = "intelligence",
+ op: str = "set_model_for_query_class",
+ target: str = "learning.routing.policy_map.math",
+ risk_tier: str = "auto",
+ status: str = "applied",
+ benchmark_delta: float | None = 0.04,
+) -> tuple[EditOutcome, dict]:
+ """Return (EditOutcome, extra_columns) — the SessionStore stores some
+ metadata that isn't on the EditOutcome model itself (pillar/op/target/
+ risk_tier/rationale come from the parent Edit)."""
+ outcome = EditOutcome(
+ edit_id=edit_id,
+ status=status, # type: ignore[arg-type]
+ benchmark_delta=benchmark_delta,
+ cluster_deltas={"cluster-001": 0.1} if benchmark_delta else {},
+ error=None if status == "applied" else "test failure",
+ applied_at=(
+ datetime(2026, 4, 8, 3, 5, 0, tzinfo=timezone.utc)
+ if status == "applied"
+ else None
+ ),
+ )
+ extras = {
+ "pillar": pillar,
+ "op": op,
+ "target": target,
+ "risk_tier": risk_tier,
+ "rationale": "test rationale",
+ }
+ return outcome, extras
+
+
+class TestSessionStoreInit:
+ """Tests for SessionStore initialization."""
+
+ def test_creates_tables(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ db = tmp_path / "learning.db"
+ store = SessionStore(db)
+
+ assert store.list_sessions() == []
+ assert store.get_session("nonexistent") is None
+ assert store.list_outcomes("nonexistent") == []
+ store.close()
+
+ def test_idempotent_init(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ db = tmp_path / "learning.db"
+ SessionStore(db).close()
+ store = SessionStore(db) # Should not raise.
+ assert store.list_sessions() == []
+ store.close()
+
+
+class TestSessionStoreSaveAndGet:
+ """Tests for save_session/get_session round-trip."""
+
+ def test_round_trip(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ original = _make_session()
+ store.save_session(original)
+
+ restored = store.get_session(original.id)
+ assert restored is not None
+ assert restored == original
+ store.close()
+
+ def test_update_existing_session(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ session = _make_session(status=SessionStatus.INITIATED)
+ store.save_session(session)
+
+ # Update status and re-save.
+ session = session.model_copy(update={"status": SessionStatus.COMPLETED})
+ store.save_session(session)
+
+ restored = store.get_session(session.id)
+ assert restored is not None
+ assert restored.status == SessionStatus.COMPLETED
+ store.close()
+
+ def test_get_returns_none_for_unknown(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ assert store.get_session("does-not-exist") is None
+ store.close()
+
+
+class TestSessionStoreList:
+ """Tests for list_sessions ordering and filters."""
+
+ def test_lists_in_started_at_desc_order(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ s1 = _make_session(session_id="s1")
+ s2 = _make_session(session_id="s2")
+ s2 = s2.model_copy(
+ update={"started_at": datetime(2026, 4, 9, 3, 0, 0, tzinfo=timezone.utc)}
+ )
+
+ store.save_session(s1)
+ store.save_session(s2)
+
+ listed = store.list_sessions()
+ assert [s.id for s in listed] == ["s2", "s1"]
+ store.close()
+
+ def test_filter_by_status(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ s1 = _make_session(session_id="s1", status=SessionStatus.COMPLETED)
+ s2 = _make_session(session_id="s2", status=SessionStatus.FAILED)
+ store.save_session(s1)
+ store.save_session(s2)
+
+ completed = store.list_sessions(status=SessionStatus.COMPLETED)
+ failed = store.list_sessions(status=SessionStatus.FAILED)
+ assert [s.id for s in completed] == ["s1"]
+ assert [s.id for s in failed] == ["s2"]
+ store.close()
+
+ def test_limit(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ for i in range(5):
+ session = _make_session(session_id=f"s{i}")
+ session = session.model_copy(
+ update={
+ "started_at": datetime(2026, 4, 8, 3 + i, 0, 0, tzinfo=timezone.utc)
+ }
+ )
+ store.save_session(session)
+
+ listed = store.list_sessions(limit=3)
+ assert len(listed) == 3
+ store.close()
+
+
+class TestEditOutcomes:
+ """Tests for save_outcome / list_outcomes."""
+
+ def test_round_trip(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ session = _make_session()
+ store.save_session(session)
+
+ outcome, extras = _make_outcome()
+ store.save_outcome(session.id, outcome, **extras)
+
+ listed = store.list_outcomes(session.id)
+ assert len(listed) == 1
+ assert listed[0].edit_id == "edit-001"
+ assert listed[0].status == "applied"
+ assert listed[0].benchmark_delta == 0.04
+ store.close()
+
+ def test_multiple_outcomes_per_session(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ session = _make_session()
+ store.save_session(session)
+
+ for i in range(3):
+ outcome, extras = _make_outcome(edit_id=f"edit-{i}")
+ store.save_outcome(session.id, outcome, **extras)
+
+ listed = store.list_outcomes(session.id)
+ assert len(listed) == 3
+ assert {o.edit_id for o in listed} == {"edit-0", "edit-1", "edit-2"}
+ store.close()
+
+ def test_outcomes_for_unknown_session_returns_empty(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ assert store.list_outcomes("nonexistent") == []
+ store.close()
+
+
+class TestParentSessionChain:
+ """Tests for parent_session_id foreign key."""
+
+ def test_parent_id_round_trips(self, tmp_path: Path) -> None:
+ from openjarvis.learning.distillation.storage.session_store import (
+ SessionStore,
+ )
+
+ store = SessionStore(tmp_path / "learning.db")
+ parent = _make_session(session_id="parent")
+ child = _make_session(session_id="child", parent_session_id="parent")
+ store.save_session(parent)
+ store.save_session(child)
+
+ restored = store.get_session("child")
+ assert restored is not None
+ assert restored.parent_session_id == "parent"
+ store.close()
diff --git a/tests/learning/distillation/test_teacher_agent.py b/tests/learning/distillation/test_teacher_agent.py
new file mode 100644
index 00000000..0d595797
--- /dev/null
+++ b/tests/learning/distillation/test_teacher_agent.py
@@ -0,0 +1,252 @@
+"""Tests for openjarvis.learning.distillation.diagnose.teacher_agent module.
+
+All tests use a mocked CloudEngine — no live API calls.
+"""
+
+from __future__ import annotations
+
+import json
+from unittest.mock import MagicMock
+
+from openjarvis.learning.distillation.diagnose.types import DiagnosticTool
+
+
+def _make_tool(
+ name: str = "test_tool", return_value: str = "tool result"
+) -> DiagnosticTool:
+ """Create a minimal diagnostic tool for testing."""
+ return DiagnosticTool(
+ name=name,
+ description=f"A test tool named {name}",
+ parameters={"type": "object", "properties": {"arg": {"type": "string"}}},
+ fn=lambda **kwargs: return_value,
+ )
+
+
+def _make_engine_response(
+ content: str = "",
+ tool_calls: list | None = None,
+ cost_usd: float = 0.01,
+) -> dict:
+ """Create a mock engine.generate() response."""
+ resp = {
+ "content": content,
+ "finish_reason": "stop" if not tool_calls else "tool_calls",
+ "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
+ "cost_usd": cost_usd,
+ }
+ if tool_calls:
+ resp["tool_calls"] = tool_calls
+ return resp
+
+
+class TestTeacherAgentNoTools:
+ """Teacher responds without using any tools."""
+
+ def test_returns_content_from_single_turn(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content="## Diagnosis\nThe student struggles with math routing."
+ )
+
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[],
+ max_turns=5,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Analyze student failures.")
+
+ assert "math routing" in result.content
+ assert result.turns == 1
+ assert result.total_cost_usd > 0
+
+ def test_tracks_cost(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content="Done", cost_usd=0.05
+ )
+
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[],
+ max_turns=5,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Analyze.")
+ assert result.total_cost_usd == 0.05
+
+
+class TestTeacherAgentWithTools:
+ """Teacher uses tools in a multi-turn loop."""
+
+ def test_executes_tool_call_and_continues(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ # Turn 1: teacher calls a tool
+ tool_call_response = _make_engine_response(
+ content="Let me look at the traces.",
+ tool_calls=[
+ {
+ "id": "call_1",
+ "name": "test_tool",
+ "arguments": json.dumps({"arg": "value"}),
+ }
+ ],
+ )
+ # Turn 2: teacher produces final answer
+ final_response = _make_engine_response(
+ content="## Diagnosis\nBased on the traces, I found...",
+ )
+ engine.generate.side_effect = [tool_call_response, final_response]
+
+ tool = _make_tool("test_tool", return_value="trace data here")
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[tool],
+ max_turns=5,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Diagnose failures.")
+
+ assert result.turns == 2
+ assert "Based on the traces" in result.content
+ assert len(result.tool_call_records) == 1
+ assert result.tool_call_records[0].tool == "test_tool"
+
+ def test_stops_at_max_turns(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ # Every turn calls a tool — should stop after max_turns
+ engine.generate.return_value = _make_engine_response(
+ content="",
+ tool_calls=[
+ {
+ "id": "call_loop",
+ "name": "test_tool",
+ "arguments": "{}",
+ }
+ ],
+ )
+
+ tool = _make_tool("test_tool")
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[tool],
+ max_turns=3,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Diagnose.")
+
+ assert result.turns == 3
+
+ def test_stops_at_max_cost(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ # Each call costs 2.0 — should stop after exceeding max_cost
+ engine.generate.return_value = _make_engine_response(
+ content="",
+ tool_calls=[
+ {
+ "id": "call_cost",
+ "name": "test_tool",
+ "arguments": "{}",
+ }
+ ],
+ cost_usd=2.0,
+ )
+
+ tool = _make_tool("test_tool")
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[tool],
+ max_turns=100,
+ max_cost_usd=3.0,
+ )
+ result = agent.run("Diagnose.")
+
+ assert result.total_cost_usd >= 2.0
+ # Should have stopped before exhausting all 100 turns
+ assert result.turns < 100
+
+ def test_multiple_tool_calls_in_one_turn(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ # Turn 1: two tool calls
+ multi_call = _make_engine_response(
+ content="Checking multiple things.",
+ tool_calls=[
+ {"id": "call_a", "name": "tool_a", "arguments": "{}"},
+ {"id": "call_b", "name": "tool_b", "arguments": "{}"},
+ ],
+ )
+ final = _make_engine_response(content="All done.")
+ engine.generate.side_effect = [multi_call, final]
+
+ tool_a = _make_tool("tool_a", "result_a")
+ tool_b = _make_tool("tool_b", "result_b")
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[tool_a, tool_b],
+ max_turns=5,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Diagnose.")
+
+ assert result.turns == 2
+ assert len(result.tool_call_records) == 2
+
+
+class TestTeacherAgentResult:
+ """Tests for TeacherAgentResult structure."""
+
+ def test_result_has_all_fields(self) -> None:
+ from openjarvis.learning.distillation.diagnose.teacher_agent import (
+ TeacherAgent,
+ )
+
+ engine = MagicMock()
+ engine.generate.return_value = _make_engine_response(
+ content="Diagnosis here.", cost_usd=0.1
+ )
+
+ agent = TeacherAgent(
+ engine=engine,
+ model="claude-opus-4-6",
+ tools=[],
+ max_turns=5,
+ max_cost_usd=5.0,
+ )
+ result = agent.run("Go.", system_prompt="You are a meta-engineer.")
+
+ assert hasattr(result, "content")
+ assert hasattr(result, "turns")
+ assert hasattr(result, "total_cost_usd")
+ assert hasattr(result, "tool_call_records")
+ assert hasattr(result, "total_tokens")
diff --git a/tests/learning/distillation/test_triggers.py b/tests/learning/distillation/test_triggers.py
new file mode 100644
index 00000000..e51c4697
--- /dev/null
+++ b/tests/learning/distillation/test_triggers.py
@@ -0,0 +1,52 @@
+"""Tests for openjarvis.learning.distillation.triggers module."""
+
+from __future__ import annotations
+
+
+class TestOnDemandTrigger:
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.triggers import OnDemandTrigger
+
+ t = OnDemandTrigger()
+ assert t.kind.value == "on_demand"
+ assert t.metadata == {}
+
+ def test_metadata(self) -> None:
+ from openjarvis.learning.distillation.triggers import OnDemandTrigger
+
+ t = OnDemandTrigger(metadata={"source": "cli"})
+ assert t.metadata["source"] == "cli"
+
+
+class TestUserFlagTrigger:
+ def test_constructs_with_trace_id(self) -> None:
+ from openjarvis.learning.distillation.triggers import UserFlagTrigger
+
+ t = UserFlagTrigger(trace_id="trace-001")
+ assert t.kind.value == "user_flag"
+ assert t.trace_id == "trace-001"
+ assert "trace_id" in t.metadata
+
+
+class TestScheduledTrigger:
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.triggers import ScheduledTrigger
+
+ t = ScheduledTrigger(cron="0 3 * * *", new_trace_count=25)
+ assert t.kind.value == "scheduled"
+ assert t.cron == "0 3 * * *"
+ assert t.new_trace_count == 25
+
+
+class TestClusterTrigger:
+ def test_constructs(self) -> None:
+ from openjarvis.learning.distillation.triggers import ClusterTrigger
+
+ t = ClusterTrigger(
+ cluster_description="math failures",
+ trace_ids=["t1", "t2", "t3"],
+ failure_rate=0.8,
+ )
+ assert t.kind.value == "cluster"
+ assert len(t.trace_ids) == 3
+ assert t.failure_rate == 0.8