Compare commits

...
Author SHA1 Message Date
Andrew Park 1a843d1ece orchestrator: SFT training + eval harness
Tokenize the cleaned routing data (all-turns masking for ChatML),
full-parameter FSDP SFT (fsdp4/8 configs for Lambda 8xH100), and an eval
backend that drives the fine-tuned orchestrator as the scored model over
GAIA/MMLU-Pro/SuperGPQA. Includes split-making, Braintrust upload, eval
sample rendering, scorer fixes, and a boxed-tool-call salvage for SFT'd
checkpoints that emit their action inside \boxed{}.
2026-07-15 13:47:06 -07:00
Andrew Park cef5eb749a orchestrator: SFT data generation + verification pipeline
Rejection-sampling data gen for the router: run anonymized ToolOrchestra
rollouts over a task pool, verify correctness (judge with backoff, math
normalization, essay/code gates), and keep only correct, well-formed,
actually-routing trajectories. Includes the clean-gate that strips
self-answering and control-token garble, cost-aware reward, dataset
loaders, unified trajectory serialization, and naming (single source of
truth for dataset names). ~76% of rollouts are dropped by design.
2026-07-15 13:46:44 -07:00
Andrew Park bc7ac3e14f toolorchestra: split the monolith into a package + anonymized expert registry
Break the 1.8k-line toolorchestra.py into a focused package (agent,
rollout, workers, experts, unified, parsing, prompts, clients, sandbox,
tracing) and add expert_registry.py, which serves the orchestrator an
anonymized expert catalog (opaque model_xxxx labels, hidden costs) so a
router can't route on brand prior. retry.py adds backoff to the rollout's
orchestrator + expert cloud calls. Prices corrected to OpenRouter list.
2026-07-15 13:45:07 -07:00
51 changed files with 8826 additions and 1907 deletions
+10
View File
@@ -130,3 +130,13 @@ learning.db
minion_logs/
*.oj-debug.json
oj-debug.*.json
# Generated orchestrator SFT/eval data artifacts
data/
# Training run artifacts (model checkpoints + W&B run dirs — regenerated, large)
checkpoints/
wandb/
# Local LiteLLM spend-tracking proxy (machine-specific, not part of the package)
litellm/
+1
View File
@@ -28,6 +28,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"braintrust>=0.0.150",
"click>=8",
"datasets>=4.5.0",
"ddgs>=9.11.4",
@@ -0,0 +1,399 @@
#!/usr/bin/env python
"""Build orchestrator SFT cold-start data by **base Qwen3-8B self-sampling** (v1).
The orchestrator *is* the local Qwen3-8B. We roll the base (un-trained) model out
over the reasoning-SFT task pool (``load_sft_tasks``: GeneralThought + OpenThoughts
code/math/science), verify each trajectory's final answer against the gold
(``verify.make_verifier`` — math/code checkers + Gemini judge fallback), and keep
the cheapest-correct trajectory per task. Those passing trajectories become the
``conversations`` JSONL the SFT trainer consumes — i.e. the model learns from its
own successful rollouts (rejection-sampling cold-start, STaR-style).
This is **v1** (base self-sampling): point ``--orchestrator-endpoint`` at the vLLM
server hosting the *base* ``Qwen/Qwen3-8B``.
For **v2**, train on the v1 data, serve the v1 checkpoint with vLLM, then re-run
this exact driver with ``--orchestrator-endpoint`` pointed at the v1 checkpoint's
endpoint, ``--orchestrator-model`` set to the checkpoint name, and
``--samples-per-task 8`` — the (now stronger) policy self-samples a v2 dataset.
The math/coder specialist tools are only wired when ``--math-endpoint`` /
``--coder-endpoint`` are passed (those are separately-served vLLM specialists);
otherwise they're omitted so the orchestrator only sees the tools it can actually
call.
Example (v1, base self-sampling):
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
--orchestrator-label qwen \
--orchestrator-endpoint http://localhost:8001/v1 \
--orchestrator-model qwen3-8b \
--samples-per-task 8 --max-keep-per-task 1
Example (v2, re-point at the v1 checkpoint):
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
--orchestrator-label qwen \
--orchestrator-endpoint http://localhost:8010/v1 \
--orchestrator-model orchestrator-sft-v1 \
--samples-per-task 8
"""
from __future__ import annotations
import argparse
import json
import logging
import os
import time
from pathlib import Path
from typing import Optional
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
from openjarvis.agents.hybrid.toolorchestra.unified import (
make_call_orchestrator,
make_dispatch,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
load_sft_tasks,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.naming import raw_dir_name
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import make_verifier
def main(argv: Optional[list[str]] = None) -> int:
p = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
# --out is a TAG, not a path. The run dir is always
# <data-root>/raw/{label}-{month}-{day}-{year}-{hhmm}{am|pm}[-{tag}]/ and the
# file inside is always data.jsonl — see the run_dir block below. Use --out
# only to distinguish two runs of the same orchestrator (e.g. --out balanced50).
p.add_argument(
"--out",
default=None,
help="Optional tag appended to the run dir name (not a path).",
)
# The orchestrator == the local Qwen3-8B self-sampling over its own rollouts.
p.add_argument(
"--orchestrator-endpoint",
default="http://localhost:8001/v1",
help="OpenAI-compatible vLLM base URL serving the orchestrator.",
)
p.add_argument("--orchestrator-model", default="qwen3-8b")
p.add_argument("--orchestrator-api-key", default="EMPTY")
# Provenance stamps: when set, every written record carries these fields so
# the JSONL is self-identifying once pooled across model families (gemma vs
# qwen). Default None -> no stamping (existing qwen runs unaffected).
p.add_argument(
"--gen-model",
default=None,
help="Full HF id of the generating orchestrator, stamped as "
"record['gen_model'] (e.g. google/gemma-4-26B-A4B-it).",
)
p.add_argument(
"--orchestrator-label",
default=None,
help="Short label stamped as record['orchestrator_model'] (e.g. gemma-4-26b).",
)
# Local OSS model endpoints; repeatable "model_id=base_url" (e.g.
# "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are
# still listed but served unconfigured (base_url=None).
p.add_argument(
"--local-endpoint",
action="append",
default=[],
metavar="MODEL_ID=URL",
help="Local model id -> vLLM base URL (repeatable).",
)
p.add_argument(
"--max-tasks",
type=int,
default=None,
help="Cap on tasks (default: all of load_sft_tasks()).",
)
p.add_argument(
"--skip-task-ids-from",
action="append",
default=[],
metavar="GLOB",
help="Glob of prior data.jsonl files; skip task_ids already "
"generated there so this run only does unseen prompts "
"(resume). Repeatable; applied before sharding.",
)
p.add_argument("--samples-per-task", type=int, default=8)
p.add_argument("--max-keep-per-task", type=int, default=1)
p.add_argument("--max-turns", type=int, default=8)
p.add_argument("--temperature", type=float, default=0.7)
# Orchestrator completion cap per turn. The library default (4096) intermittently
# truncates the final answer mid-sentence on longer reasoning turns; bump it so
# the trace ends cleanly. Raise further (e.g. 16384) if traces still cut off.
p.add_argument(
"--max-tokens",
type=int,
default=8192,
help="Per-turn completion cap for the orchestrator (default 8192).",
)
# Number of tasks rolled out concurrently against vLLM. Each in-flight task
# issues its samples sequentially, so ~concurrency requests hit the server at
# once. ~30-50 is comfortable on 1xL40S (prefix caching + continuous batching).
p.add_argument(
"--concurrency",
type=int,
default=32,
help="Tasks rolled out in parallel (default 32).",
)
# Balanced-smoke pull: cap//4 from each of GeneralThought + OpenThoughts
# code/math/science instead of the GeneralThought-only fast cap path. Use for a
# representative smoke; the real run omits --max-tasks for the full 8K balanced set.
p.add_argument(
"--balanced",
action=argparse.BooleanOptionalAction,
default=True,
help="Draw an EVEN cross-domain sample (GeneralThought + "
"OpenThoughts code/math/science). Default ON — pass "
"--no-balanced for the old GeneralThought-only skew.",
)
# Data-parallel sharding: split the (deterministic, seed-42) task list across
# N independent driver processes, each pointed at its own orchestrator vLLM
# replica. Shard i takes tasks[i::N] (a strided slice, so every shard stays
# domain-balanced). Each writes its own --out file; concatenate the shard
# JSONLs afterward. Lets the GPU-bound orchestrator scale across idle GPUs.
p.add_argument(
"--shard-index",
type=int,
default=0,
help="This shard's index in [0, shard-count).",
)
p.add_argument(
"--shard-count",
type=int,
default=1,
help="Total number of shards (default 1 = no sharding).",
)
# Throughput knob: stop sampling a task as soon as max-keep-per-task passing
# trajectories are found, instead of always running all --samples-per-task.
# ~3-4x faster when most tasks solve early, but keeps the *first* passers
# rather than the *cheapest* of N (drops the cost-optimisation signal).
p.add_argument(
"--stop-at-keep",
action="store_true",
help="Short-circuit a task once max-keep passing samples found.",
)
# Anonymize model experts (opaque random labels, uniform description, no cost
# line, shuffled order) so the policy can't route on a model's name/position/
# cost. The anon->real map is saved per record (metrics.anon_map) for analysis.
p.add_argument(
"--anonymize-experts",
action="store_true",
help="Hide expert identity (random labels + shuffle) when routing.",
)
# By default we keep EVERY rolled-out trajectory (correct + incorrect), each
# tagged with ``correct`` (verifier verdict) and ``kept`` (the cheapest-correct
# sample the rejection sampler would pick). Lets you compute accuracy / inspect
# failures from the JSONL directly. --rejection-only restores the old
# drop-the-failures behaviour (only cheapest-correct written).
p.add_argument(
"--rejection-only",
action="store_true",
help="Drop incorrect rollouts; write only the cheapest-correct "
"sample per task (the original behaviour).",
)
# Rollouts call shell_exec / file_write with model-chosen relative paths
# (e.g. ``solution.py``), which otherwise land in the repo root. Run them
# from a throwaway scratch dir so generated files never dirty the tree.
# gitignored via the existing ``scratch/`` rule.
p.add_argument(
"--scratch-dir",
default="scratch/sft-rollouts",
help="CWD for rollouts; stray tool-written files go here.",
)
args = p.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(message)s")
# Quiet the per-request HTTP / dataset-stream spam so the run log stays
# readable (rollout calls and dataset shards otherwise flood it).
for _noisy in (
"httpx",
"httpcore",
"urllib3",
"datasets",
"fsspec",
"huggingface_hub",
"openai",
):
logging.getLogger(_noisy).setLevel(logging.WARNING)
# Every run lands in its OWN folder under <data-root>/raw/, sharing its stamp
# with every file later carved from it — only the stage word differs:
#
# raw/qwen-july-7-2026-0553pm/data.jsonl <- every rollout, incl. failures
# sft/qwen-clean-july-7-2026-0553pm.jsonl <- reject-sampled from it
#
# so any curated file traces back to its generation run by eye. raw/ sits ABOVE
# the sft/rl fork on purpose: SFT keeps only correct+clean rows, but GRPO needs
# the failures, and both read the same pool. Naming lives in sft_data.naming.
# --out is an optional extra TAG, not a path; the file inside is always data.jsonl.
#
# OJ_DATA_ROOT keeps the data OUT of the git checkout (repo-relative default so
# a fresh clone still works); this workspace points it at the experiments tree.
data_root = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
prefix = args.orchestrator_label or "orch"
tag = Path(args.out).stem if args.out else ""
base = data_root / "raw" / raw_dir_name(prefix, tag=tag)
run_dir = base.resolve()
n = 1
while run_dir.exists(): # same name+minute (parallel shards) -> disambiguate
n += 1
run_dir = base.with_name(f"{base.name}-{n}").resolve()
label = run_dir.name
out_p = run_dir / "data.jsonl"
lock_p = out_p.with_suffix(out_p.suffix + ".lock")
out_p.parent.mkdir(parents=True, exist_ok=True)
lock_p.write_text(f"{time.strftime('%m-%d-%I%M%p').lower()} pid={os.getpid()}")
args.out = str(out_p)
import atexit
atexit.register(lambda: lock_p.exists() and lock_p.unlink())
logging.info("Run dir: %s", run_dir)
# Enrich Braintrust rollout traces with run-level provenance (no-op-safe;
# tracing.run_context() reads these). setdefault so an explicit env wins.
os.environ.setdefault("OJ_RUN_LABEL", label)
os.environ.setdefault("OJ_GEN_MODEL", args.gen_model or args.orchestrator_model)
os.environ.setdefault("OJ_RUN_STAGE", "smoke" if args.max_tasks else "prod")
os.environ["OJ_CFG_TEMPERATURE"] = str(args.temperature)
os.environ["OJ_CFG_MAX_TURNS"] = str(args.max_turns)
os.environ["OJ_CFG_ANONYMIZE"] = str(bool(args.anonymize_experts))
os.environ["OJ_CFG_REJECTION_ONLY"] = str(bool(args.rejection_only))
# Sandbox the rollouts: chdir into a scratch dir so any file a rollout writes
# (shell_exec redirects, file_write with a relative path) lands there instead
# of the repo root. (--out is already absolute via run_dir.resolve() above.)
scratch = Path(args.scratch_dir).resolve()
scratch.mkdir(parents=True, exist_ok=True)
os.chdir(scratch)
logging.info("Rollout scratch dir (cwd): %s", scratch)
local_endpoints = {}
for item in args.local_endpoint:
model_id, _, url = item.partition("=")
if model_id and url:
local_endpoints[model_id] = url
tools = orchestrator_catalog(local_endpoints=local_endpoints or None)
# Drop the no-op ``think`` scratchpad tool: it adds turns/cost and is a
# harness-leakage vector (the model narrates the routing rule inside a think
# call, which the reasoning-scrub can't reach). The model still reasons in
# its own <think> blocks — it just can't emit reasoning AS a tool call.
tools = [t for t in tools if t.name != "think"]
logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools])
call_orch = make_call_orchestrator(
args.orchestrator_model,
base_url=args.orchestrator_endpoint,
api_key=args.orchestrator_api_key,
temperature=args.temperature,
max_tokens=args.max_tokens,
)
dispatch = make_dispatch({})
# Build the provenance stamp once (None if neither flag given).
record_extra = {}
if args.gen_model:
record_extra["gen_model"] = args.gen_model
if args.orchestrator_label:
record_extra["orchestrator_model"] = args.orchestrator_label
record_extra = record_extra or None
def rollout_fn(task):
try:
return run_unified_rollout(
task.instruction,
tools,
call_orchestrator=call_orch,
dispatch=dispatch,
max_turns=args.max_turns,
anonymize=args.anonymize_experts,
)
except Exception as exc: # network/key failures shouldn't kill the run
logging.warning("rollout failed for %s: %s", task.task_id, exc)
return None
# cap= caps the task count for smoke runs; --balanced makes that small sample
# cross-domain (GeneralThought + OpenThoughts code/math/science) instead of the
# GeneralThought-only fast path. The real run omits --max-tasks (full 8K balanced).
tasks = load_sft_tasks(cap=args.max_tasks, balanced=args.balanced)
# Resume on unseen prompts: drop task_ids already generated in prior runs
# (per-track seen set via --skip-task-ids-from). Filter BEFORE sharding so the
# remaining unseen tasks stride evenly and never re-cover finished prompts.
if args.skip_task_ids_from:
import glob as _glob
skip_ids = set()
for pattern in args.skip_task_ids_from:
for fp in _glob.glob(pattern):
try:
with open(fp) as fh:
for line in fh:
line = line.strip()
if not line:
continue
try:
tid = json.loads(line).get("task_id")
except Exception:
continue
if tid:
skip_ids.add(tid)
except OSError:
continue
if skip_ids:
n_before = len(tasks)
tasks = [t for t in tasks if t.task_id not in skip_ids]
logging.info(
"Resume: %d done task_ids -> skipping; %d of %d tasks remain",
len(skip_ids),
len(tasks),
n_before,
)
if args.shard_count > 1:
n_all = len(tasks)
tasks = tasks[args.shard_index :: args.shard_count]
logging.info(
"Shard %d/%d -> %d of %d tasks",
args.shard_index,
args.shard_count,
len(tasks),
n_all,
)
from collections import Counter as _Counter
_dom = _Counter(getattr(t, "domain", "unknown") for t in tasks)
logging.info(
"Loaded %d SFT tasks | balanced=%s | domain split: %s",
len(tasks),
args.balanced,
", ".join(f"{d}={n}" for d, n in sorted(_dom.items())),
)
stats = generate_sft_dataset(
args.out,
tasks=tasks,
tools=tools,
rollout_fn=rollout_fn,
verify_fn=make_verifier(),
samples_per_task=args.samples_per_task,
max_keep_per_task=args.max_keep_per_task,
reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets highest reward
concurrency=args.concurrency,
stop_at_keep=args.stop_at_keep,
keep_all=not args.rejection_only,
record_extra=record_extra,
)
print(json.dumps(stats, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
+358
View File
@@ -0,0 +1,358 @@
#!/usr/bin/env python
"""Run the ToolOrchestra orchestrator over N-sample subsets of several
benchmarks and score them with the existing eval infra.
Reuses, verbatim:
* ``openjarvis.evals.cli._build_dataset`` / ``_build_scorer`` /
``_build_judge_backend`` (loaders + scorers + judge wiring),
* ``openjarvis.evals.core.runner.EvalRunner`` (parallel run + scoring),
* ``openjarvis.evals.core.types.RunConfig`` (run config),
and plugs in our orchestrator as the "model" via
``openjarvis.learning.intelligence.orchestrator.eval_backend.OrchestratorBackend``.
For each benchmark it runs ``EvalRunner(config, dataset, backend, scorer).run()``,
collects accuracy / cost / latency, writes a combined ``summary.json``, and
prints a table.
NOTE on the judge: the OpenAI and Gemini keys are both dead — Anthropic is the
only live provider, so the judge is ``claude-haiku-4-5-20251001`` (the same model
the generation-side verifier uses, see ``sft_data/verify.py:JUDGE_MODEL``). The
judge backend is still built via the ``cloud`` engine (``_build_judge_backend``)
— the model string selects the provider/route. Override with ``--judge-model`` /
``--judge-engine``.
Example
-------
.venv/bin/python scripts/orchestrator/eval_orchestrator.py \\
--benchmarks gaia,mmlu_pro --n 100 \\
--orchestrator-endpoint http://localhost:8001/v1 \\
--orchestrator-model qwen3-8b \\
--output-dir ~/.openjarvis/experiments/hybrid/runs/orch-eval
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional
# Friendly aliases (underscored / shorthand) -> the EXACT registry keys used by
# openjarvis.evals.cli.BENCHMARKS / _build_dataset / _build_scorer.
BENCHMARK_ALIASES: Dict[str, str] = {
"terminalbench_v2_1": "terminalbench-v2.1",
"terminalbench-v2_1": "terminalbench-v2.1",
"terminalbench_v21": "terminalbench-v2.1",
"terminalbench-v2.1": "terminalbench-v2.1",
"mmlu_pro": "mmlu-pro",
"mmlu-pro": "mmlu-pro",
"gaia": "gaia",
"taubench": "taubench",
"supergpqa": "supergpqa",
}
DEFAULT_BENCHMARKS = "gaia,terminalbench_v2_1,taubench,mmlu_pro,supergpqa"
def _normalize_benchmark(name: str) -> str:
key = name.strip()
return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key))
def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
p = argparse.ArgumentParser(
description="Run the orchestrator over benchmark subsets and score them.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
p.add_argument(
"--benchmarks",
default=DEFAULT_BENCHMARKS,
help="Comma-separated benchmark names (aliases normalized to registry keys).",
)
p.add_argument("--n", type=int, default=100, help="Samples per benchmark.")
p.add_argument("--seed", type=int, default=42, help="Subset seed.")
p.add_argument(
"--orchestrator-endpoint",
default="http://localhost:8001/v1",
help="OpenAI-compatible base URL for the served orchestrator.",
)
p.add_argument(
"--orchestrator-model",
default="qwen3-8b",
help="Served orchestrator model id.",
)
p.add_argument(
"--orchestrator-api-key",
default="EMPTY",
help="API key for the orchestrator endpoint (EMPTY for local vLLM).",
)
p.add_argument(
"--local-endpoint",
action="append",
default=[],
metavar="MODEL_ID=URL",
help=(
"Wire a local OSS tool: map a catalog model id to its vLLM base "
"URL, e.g. 'Qwen/Qwen3.5-9B=http://localhost:8003/v1'. Repeatable."
),
)
p.add_argument(
"--max-turns",
type=int,
default=8,
help="Max orchestrator turns per sample.",
)
p.add_argument(
"--temperature",
type=float,
default=1.0,
help="Orchestrator sampling temperature.",
)
p.add_argument(
"--max-workers",
type=int,
default=4,
help="Parallel samples per benchmark.",
)
p.add_argument(
"--output-dir",
default="results/orchestrator-eval",
help="Directory for per-benchmark JSONL + combined summary.json.",
)
# OpenAI AND Gemini keys are both dead — Anthropic is the only live provider,
# and this matches the generation-side verifier (verify.JUDGE_MODEL).
p.add_argument(
"--base-model",
action="store_true",
help="Orchestrator is an UNTRAINED base model: serve with native tools= "
"instead of the baked fine-tuned JSON prompt.",
)
p.add_argument(
"--judge-model",
default="claude-haiku-4-5-20251001",
help="LLM-judge model.",
)
p.add_argument(
"--judge-engine",
default="cloud",
help="Engine key for the judge backend.",
)
return p.parse_args(argv)
def _run_one(
benchmark: str,
*,
n: int,
seed: int,
backend,
judge_model: str,
judge_engine: str,
max_workers: int,
orchestrator_model: str,
output_dir: Path,
) -> Dict[str, Any]:
"""Run + score one benchmark, returning a row dict (or an error row)."""
# Imports here so --help works without importing the (heavy) eval stack.
from openjarvis.evals.cli import (
_build_dataset,
_build_judge_backend,
_build_scorer,
)
from openjarvis.evals.core.runner import EvalRunner
from openjarvis.evals.core.types import RunConfig
output_path = output_dir / f"{benchmark.replace('.', '_')}_orchestrator.jsonl"
dataset = _build_dataset(benchmark)
# The runner reloads internally; we also load here per spec so callers can
# inspect/size the subset up front. load() is idempotent for these sets.
dataset.load(max_samples=n, seed=seed)
judge_backend = _build_judge_backend(judge_model, engine_key=judge_engine)
scorer = _build_scorer(benchmark, judge_backend, judge_model)
config = RunConfig(
benchmark=benchmark,
backend="orchestrator",
model=orchestrator_model,
max_samples=n,
max_workers=max_workers,
seed=seed,
judge_model=judge_model,
judge_engine=judge_engine,
output_path=str(output_path),
)
runner = EvalRunner(config, dataset, backend, scorer)
started = time.time()
try:
summary = runner.run()
finally:
if judge_backend is not None:
judge_backend.close()
elapsed = time.time() - started
return {
"benchmark": benchmark,
"accuracy": summary.accuracy,
"scored_samples": summary.scored_samples,
"correct": summary.correct,
"errors": summary.errors,
"total_samples": summary.total_samples,
"mean_latency_seconds": summary.mean_latency_seconds,
"total_cost_usd": summary.total_cost_usd,
"mean_continuous_score": summary.mean_continuous_score,
"wall_seconds": round(elapsed, 2),
"output_path": str(output_path),
}
def _print_table(rows: List[Dict[str, Any]]) -> None:
cols = [
("benchmark", 22, "{}"),
("accuracy", 9, "{:.4f}"),
("correct", 8, "{}"),
("scored", 7, "{}"),
("errors", 7, "{}"),
("cost($)", 10, "{:.4f}"),
("lat(s)", 9, "{:.2f}"),
("wall(s)", 9, "{:.1f}"),
]
header = " ".join(f"{name:<{w}}" for name, w, _ in cols)
print("\n" + header)
print("-" * len(header))
for r in rows:
if r.get("error"):
print(f"{r['benchmark']:<22} ERROR: {r['error']}")
continue
vals = {
"benchmark": r["benchmark"],
"accuracy": r["accuracy"],
"correct": r["correct"],
"scored": r["scored_samples"],
"errors": r["errors"],
"cost($)": r["total_cost_usd"],
"lat(s)": r["mean_latency_seconds"],
"wall(s)": r["wall_seconds"],
}
cells = []
for name, w, fmt in cols:
v = vals[name]
try:
s = fmt.format(v)
except (ValueError, TypeError):
s = str(v)
cells.append(f"{s:<{w}}")
print(" ".join(cells))
def main(argv: Optional[List[str]] = None) -> int:
args = _parse_args(argv)
benchmarks = [
_normalize_benchmark(b) for b in args.benchmarks.split(",") if b.strip()
]
output_dir = Path(args.output_dir).expanduser()
output_dir.mkdir(parents=True, exist_ok=True)
# Build the orchestrator backend once (stateless across benchmarks).
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
OrchestratorBackend,
)
local_endpoints: Dict[str, str] = {}
for pair in args.local_endpoint:
if "=" not in pair:
raise SystemExit(f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}")
model_id, url = pair.split("=", 1)
local_endpoints[model_id.strip()] = url.strip()
backend = OrchestratorBackend(
orchestrator_endpoint=args.orchestrator_endpoint,
orchestrator_model=args.orchestrator_model,
api_key=args.orchestrator_api_key,
local_endpoints=local_endpoints,
max_turns=args.max_turns,
temperature=args.temperature,
# BASE models need native tools=; FINE-TUNED models need the baked
# JSON prompt they were trained on. Serving either in the other mode
# handicaps it (base-in-baked scored 0.184 on a rerun; the June 0.40
# baseline used native mode).
finetuned=not args.base_model,
)
rows: List[Dict[str, Any]] = []
try:
for benchmark in benchmarks:
print(f"\n=== {benchmark} (n={args.n}, seed={args.seed}) ===")
try:
row = _run_one(
benchmark,
n=args.n,
seed=args.seed,
backend=backend,
judge_model=args.judge_model,
judge_engine=args.judge_engine,
max_workers=args.max_workers,
orchestrator_model=args.orchestrator_model,
output_dir=output_dir,
)
print(
f" accuracy={row['accuracy']:.4f} "
f"({row['correct']}/{row['scored_samples']}) "
f"errors={row['errors']} cost=${row['total_cost_usd']:.4f}"
)
except Exception as exc: # noqa: BLE001 - one bad bench shouldn't kill all
print(f" FAILED: {type(exc).__name__}: {exc}", file=sys.stderr)
row = {"benchmark": benchmark, "error": f"{type(exc).__name__}: {exc}"}
rows.append(row)
finally:
backend.close()
# Resolve the EXACT model behind the served alias (e.g. "gemma-sft" ->
# the real id + checkpoint path vLLM reports), so results are self-identifying
# instead of an opaque label. Best-effort: never fail the summary over it.
served_id, served_path = args.orchestrator_model, None
try:
import urllib.request
req = urllib.request.Request(
args.orchestrator_endpoint.rstrip("/") + "/models",
headers={"Authorization": f"Bearer {args.orchestrator_api_key}"},
)
with urllib.request.urlopen(req, timeout=10) as resp:
data = json.load(resp).get("data", [])
if data:
served_id = data[0].get("id", served_id)
# vLLM reports the loaded checkpoint path in `root` (or `parent`).
served_path = data[0].get("root") or data[0].get("parent")
except Exception:
pass
combined = {
"orchestrator_model": args.orchestrator_model,
"orchestrator_served_id": served_id,
"orchestrator_model_path": served_path,
"orchestrator_endpoint": args.orchestrator_endpoint,
"judge_model": args.judge_model,
"judge_engine": args.judge_engine,
"n": args.n,
"seed": args.seed,
"max_turns": args.max_turns,
"benchmarks": rows,
}
summary_path = output_dir / "summary.json"
with open(summary_path, "w") as f:
json.dump(combined, f, indent=2, default=str)
_print_table(rows)
print(f"\nCombined summary written to {summary_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+176
View File
@@ -0,0 +1,176 @@
#!/usr/bin/env python
"""Carve train / holdout / overfit100 splits from a clean orchestrator-SFT pool.
Replaces make_tranches.py + make_gemma_tranches.py, which carved 1k/2k/4k/8k
data-scaling tranches from the June pool. That lineage is dead: the tranches it
produced trained checkpoints that evaluated at base ~= 1k ~= 2k (the pool taught
the orchestrator to re-derive answers itself instead of routing), and the pool
was deleted. The lesson is in data/orchestrator/README.md — it was a data-QUALITY
ceiling, so re-running the old scaling ladder buys nothing.
Reads a clean pool (reject-sampled from raw/) and writes the SFT splits to sft/.
Naming is uniform — ``{name}-{split}-{stamp}`` with name in {qwen, gemma, pooled};
pass several pools to merge-and-restratify into `pooled`. The splits INHERIT the
pool's stamp, so they stay tied to the run that generated the data.
The holdout is domain-stratified so val-loss is leak-free, and `overfit100` is a
strict prefix of train (it is a memorisation sanity check — the model must be
able to fit 100 rows it *has* seen, else the format/masking is broken).
Deterministic (seed 42).
# one pool -> qwen-{train,holdout,overfit100}-july-7-2026-0553pm.jsonl
python scripts/orchestrator/make_splits.py --name qwen \
--pool .../sft/qwen-clean-july-7-2026-0553pm.jsonl
# merge both -> pooled-{train,holdout,overfit100}-july-7-2026-0553pm.jsonl
python scripts/orchestrator/make_splits.py --name pooled \
--pool .../sft/qwen-clean-july-7-2026-0553pm.jsonl \
--pool .../sft/gemma-clean-july-7-2026-0553pm.jsonl
"""
import argparse
import json
import os
import random
import sys
from collections import defaultdict
from pathlib import Path
from openjarvis.learning.intelligence.orchestrator.sft_data.naming import (
dataset_name,
run_stamp,
stamp_from,
)
# Where the orchestrator data tree lives. Repo-relative by default so a fresh
# clone works out of the box; set OJ_DATA_ROOT to keep the data OUT of the git
# checkout (this workspace points it at ~/experiments/orchestrator/data, so a
# stray `git reset` can't touch hundreds of GB of generations).
DATA_ROOT = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
OUT = DATA_ROOT / "sft"
SEED = 42
OVERFIT_N = 100
# Orchestrator that generated each pool — stamped onto every row so provenance
# survives a merge (the filename encodes it too, but the field is cheap insurance
# and is what the `pooled` splits rely on to stay attributable).
ORCH_MODEL = {"qwen": "qwen3.5-9b", "gemma": "gemma-4-26b"}
def main() -> int:
ap = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
ap.add_argument(
"--name",
required=True,
help="split family: qwen | gemma | pooled (drives the filename)",
)
ap.add_argument(
"--pool",
action="append",
required=True,
metavar="PATH",
help="clean pool jsonl; repeat to merge (use with --name pooled)",
)
ap.add_argument(
"--holdout-frac",
type=float,
default=0.15,
help="fraction held out, domain-stratified (default 0.15)",
)
ap.add_argument(
"--stamp",
default=None,
help="Stamp for the output names, e.g. july-7-2026-0553pm. Defaults to the "
"POOL's stamp (read off its filename) so a split stays tied to the run "
"that generated it, NOT to the day it happened to be carved. Falls back "
"to now only if no pool name carries a stamp.",
)
ap.add_argument("--out", type=Path, default=OUT)
args = ap.parse_args()
# A split belongs to the generation run, not to the day it was carved — so
# inherit the pool's stamp unless told otherwise.
stamp = args.stamp or next(
(s for p in args.pool if (s := stamp_from(Path(p).name))), None
) or run_stamp()
rows = []
for p in args.pool:
path = Path(p)
if not path.exists():
print(f"!! pool not found: {path}", file=sys.stderr)
return 1
# Infer the source orchestrator from the filename so merged pools stay
# attributable; don't clobber a stamp the pool already carries.
stem = path.name.split("-", 1)[0]
loaded = [json.loads(line) for line in path.open() if line.strip()]
for r in loaded:
r.setdefault("orchestrator_model", ORCH_MODEL.get(stem, stem))
rows.extend(loaded)
print(f"loaded {len(loaded):>5} rows from {path}")
print(f"pool total: {len(rows)}")
# GROUPED stratified holdout: split by task_id, never by row.
#
# The pool keeps SEVERAL rollouts per question (different attempts at the same
# task, same gold answer). Splitting by row therefore leaks: attempt A of a
# question lands in train and attempt B of the SAME question lands in holdout,
# so val-loss is measured on questions the model trained on. That is exactly
# what happened to the 0707 splits — 80% of qwen's holdout tasks were also in
# train, making every val-loss number optimistically biased.
#
# So: group rows by task_id, stratify the TASKS by domain, and assign whole
# tasks to one side. A question is never split across train and holdout.
by_task = defaultdict(list)
for i, r in enumerate(rows):
by_task[r.get("task_id")].append(i)
task_dom = {t: rows[idxs[0]].get("domain", "misc") for t, idxs in by_task.items()}
by_dom = defaultdict(list)
for t, dom in task_dom.items():
by_dom[dom].append(t)
rng = random.Random(SEED)
holdout_tasks: set = set()
for dom, tasks in sorted(by_dom.items()):
tasks = sorted(tasks) # deterministic before sampling
k = round(args.holdout_frac * len(tasks))
pick = rng.sample(tasks, min(k, len(tasks)))
holdout_tasks.update(pick)
print(f" {dom:<9} tasks={len(tasks):>4} holdout={len(pick)}")
holdout_idx = {i for t in holdout_tasks for i in by_task[t]}
holdout = [rows[i] for i in sorted(holdout_idx)]
train = [r for i, r in enumerate(rows) if i not in holdout_idx] # order preserved
overfit = train[:OVERFIT_N]
# Assert the leak is actually gone — this is the whole point of the grouping.
tr_tasks = {r.get("task_id") for r in train}
ho_tasks = {r.get("task_id") for r in holdout}
overlap = tr_tasks & ho_tasks
if overlap:
raise SystemExit(f"!! LEAK: {len(overlap)} task_ids in BOTH train and holdout")
print(
f"train={len(train)} rows / {len(tr_tasks)} tasks "
f"holdout={len(holdout)} rows / {len(ho_tasks)} tasks "
f"overfit={len(overfit)} (0 tasks shared)"
)
args.out.mkdir(parents=True, exist_ok=True)
written = {}
for split, data in (
("train", train),
("holdout", holdout),
(f"overfit{OVERFIT_N}", overfit),
):
f = args.out / f"{dataset_name(args.name, split, stamp)}.jsonl"
f.write_text("".join(json.dumps(r) + "\n" for r in data))
written[split] = f
print(f"wrote {f} ({len(data)})")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+407
View File
@@ -0,0 +1,407 @@
#!/usr/bin/env python
"""Full-parameter SFT for the orchestrator policy via FSDP (multi-GPU).
Launch with accelerate so the model + optimizer shard across GPUs (an 8-9B full
fine-tune won't fit one L40S otherwise):
accelerate launch --config_file <fsdp.yaml> \
scripts/orchestrator/run_sft_fsdp.py \
--data data/orchestrator/sft/qwen_train_0707.jsonl [--data <more> ...] \
--variant correct --model Qwen/Qwen3.5-9B \
--out checkpoints/sft_correct --epochs 3 --batch-size 8 --grad-accum 8 \
--wandb-project orchestrator-sft
--variant selects which trajectories to train on:
all -> every record
correct -> only records the verifier marked correct
correct_routed -> only correct records that delegated to a model expert
By default (--require-clean) records explicitly marked clean=False by the clean
gate (bloated / garbled / unrouted) are also dropped; rows missing the flag
(legacy data) are kept. Pass --no-require-clean to skip this filter.
Reuses the conversation->tokens + assistant-only masking from sft_tokenize.py.
On these no-NVLink L40S, the accelerate/FSDP NCCL env (NCCL_P2P_DISABLE=1 etc.)
must be set by the launcher.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import re
import sys
import time
from pathlib import Path
# reuse the data builder from the LoRA launcher
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from sft_tokenize import build_example # noqa: E402
MODEL_EXPERTS = {
"gpt_5_5",
"claude_opus_4_8",
"qwen3_5_9b",
"qwen3_6_27b_fp8",
"qwen3_5_122b_a10b_fp8",
"qwen3_5_397b_a17b_fp8",
}
def record_is_correct(r: dict) -> bool:
c = r.get("correct")
if c is None:
c = r.get("metrics", {}).get("correct")
return bool(c)
def record_routed_to_expert(r: dict) -> bool:
"""True if the trajectory called a model expert (resolving anon labels)."""
amap = r.get("metrics", {}).get("anon_map", {})
names = re.findall(
r'"name"\s*:\s*"([a-z0-9_]+)"', json.dumps(r.get("conversations", []))
)
for n in names:
real = amap.get(n, n)
if real in MODEL_EXPERTS:
return True
return False
def record_clean_flag(r: dict):
"""The record's `clean` flag: True/False, or None if absent (legacy data).
None means the clean gate never ran on this row, so we can't judge it -> we
keep it (legacy behavior) rather than silently dropping older datasets.
"""
c = r.get("clean")
if c is None:
c = r.get("metrics", {}).get("clean")
return c
def select(records, variant: str, require_clean: bool = True):
"""Filter by --variant, then (default) drop records explicitly marked unclean.
require_clean drops rows whose `clean` flag is False. Rows missing the flag
(legacy) are kept regardless, matching pre-clean-gate behavior.
"""
if variant == "all":
base = list(records)
elif variant == "correct":
base = [r for r in records if record_is_correct(r)]
elif variant == "correct_routed":
base = [
r for r in records if record_is_correct(r) and record_routed_to_expert(r)
]
else:
raise ValueError(f"unknown variant {variant}")
if not require_clean:
return base
return [r for r in base if record_clean_flag(r) is not False]
def main() -> int:
p = argparse.ArgumentParser()
p.add_argument(
"--data", action="append", required=True, help="JSONL(s) (repeatable)"
)
p.add_argument(
"--val-data",
default=None,
help="Held-out JSONL for val-loss (excluded from --data). Eval'd per epoch.",
)
p.add_argument(
"--variant", choices=["all", "correct", "correct_routed"], default="correct"
)
p.add_argument(
"--require-clean",
action=argparse.BooleanOptionalAction,
default=True,
help="Drop records whose `clean` flag is False (bloated / garbled / "
"unrouted). Rows missing the flag (legacy data) are kept. "
"Use --no-require-clean to disable.",
)
p.add_argument("--model", default="Qwen/Qwen3.5-9B")
p.add_argument("--out", required=True)
p.add_argument("--epochs", type=float, default=3.0)
p.add_argument("--batch-size", type=int, default=8)
p.add_argument("--grad-accum", type=int, default=8)
p.add_argument("--lr", type=float, default=1e-5)
p.add_argument("--max-seq", type=int, default=8192)
p.add_argument(
"--supervise-last-only",
action="store_true",
help="Legacy: supervise ONLY the final assistant turn. Default "
"(off) supervises every assistant turn incl. routing.",
)
p.add_argument("--warmup-ratio", type=float, default=0.03)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--wandb-project", default="orchestrator-sft")
p.add_argument("--wandb-name", default="")
args = p.parse_args()
import random
from datetime import timedelta
import torch
from accelerate import Accelerator
from accelerate.utils import InitProcessGroupKwargs, set_seed
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
set_seed(args.seed)
# 30-min collective timeout: a slow first-step shard gather / graph build on a
# contended node was tripping the default 10-min NCCL watchdog (job 3892 died
# at step 1). This only delays a *real* hang's crash; it fixes spurious ones.
pg_timeout_min = int(os.environ.get("SFT_PG_TIMEOUT_MIN", "30"))
accelerator = Accelerator(
gradient_accumulation_steps=args.grad_accum,
kwargs_handlers=[
InitProcessGroupKwargs(timeout=timedelta(minutes=pg_timeout_min))
],
)
is_main = accelerator.is_main_process
def log(m):
if is_main:
print(f"[fsdp-sft] {time.strftime('%H:%M:%S')} {m}", flush=True)
tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
if tok.pad_token_id is None:
tok.pad_token = tok.eos_token
# ---- data ----
records = []
for f in args.data:
records += [json.loads(l) for l in open(f) if l.strip()]
sel = select(records, args.variant, require_clean=args.require_clean)
log(f"variant={args.variant}: {len(sel)}/{len(records)} records")
if args.require_clean:
pre_clean = select(records, args.variant, require_clean=False)
n_unclean = len(pre_clean) - len(sel)
log(
f"require_clean=True: dropped {n_unclean} unclean rows "
f"({len(pre_clean)}->{len(sel)}); use --no-require-clean to keep them"
)
else:
log("require_clean=False: NOT filtering on `clean` flag")
examples = []
for r in sel:
ex = build_example(
tok,
r.get("conversations", []),
args.max_seq,
supervise_all_turns=not args.supervise_last_only,
)
if ex:
examples.append(ex)
log(f"built {len(examples)} training examples")
if not examples:
log("no examples; abort")
return 1
random.Random(args.seed).shuffle(examples)
def collate(batch):
maxlen = max(len(b["input_ids"]) for b in batch)
ii, lab, am = [], [], []
for b in batch:
pad = maxlen - len(b["input_ids"])
ii.append(b["input_ids"] + [tok.pad_token_id] * pad)
lab.append(b["labels"] + [-100] * pad)
am.append([1] * len(b["input_ids"]) + [0] * pad)
return (torch.tensor(ii), torch.tensor(lab), torch.tensor(am))
dl = DataLoader(
examples, batch_size=args.batch_size, shuffle=True, collate_fn=collate
)
# ---- held-out val set (leak-free; excluded from --data upstream) ----
val_dl = None
if args.val_data:
vrecs = [json.loads(l) for l in open(args.val_data) if l.strip()]
vex = []
for r in vrecs:
ex = build_example(
tok,
r.get("conversations", []),
args.max_seq,
supervise_all_turns=not args.supervise_last_only,
)
if ex:
vex.append(ex)
log(f"val: {len(vex)} examples from {args.val_data}")
if vex:
val_dl = DataLoader(
vex, batch_size=args.batch_size, shuffle=False, collate_fn=collate
)
model = AutoModelForCausalLM.from_pretrained(
args.model,
dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="sdpa",
)
if os.environ.get("SFT_NO_GRAD_CKPT") == "1":
log("gradient checkpointing DISABLED (SFT_NO_GRAD_CKPT=1)")
else:
model.gradient_checkpointing_enable(
gradient_checkpointing_kwargs={"use_reentrant": False}
)
model.config.use_cache = False
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.0)
model, opt, dl = accelerator.prepare(model, opt, dl)
if val_dl is not None:
val_dl = accelerator.prepare(val_dl)
@torch.no_grad()
def evaluate():
# Corpus-level mean token loss over the holdout. FSDP needs every rank to
# run the forward collectively, so each rank evals its shard and we reduce
# token-weighted sums (not a mean-of-batch-means, which would mis-weight).
model.eval()
tot_loss = torch.zeros((), device=accelerator.device)
tot_tok = torch.zeros((), device=accelerator.device)
for ii, lab, am in val_dl:
out = model(input_ids=ii, attention_mask=am, labels=lab)
ntok = (lab != -100).sum()
tot_loss += out.loss.detach() * ntok
tot_tok += ntok
tot_loss = accelerator.reduce(tot_loss, reduction="sum")
tot_tok = accelerator.reduce(tot_tok, reduction="sum")
model.train()
return (tot_loss / tot_tok.clamp(min=1)).item()
steps_per_epoch = math.ceil(len(dl) / args.grad_accum)
total_steps = max(1, int(steps_per_epoch * args.epochs))
warmup = max(1, int(total_steps * args.warmup_ratio))
def lr_at(s):
if s < warmup:
return s / warmup
return 0.5 * (
1 + math.cos(math.pi * (s - warmup) / max(1, total_steps - warmup))
)
use_wandb = False
if is_main:
try:
import wandb
base_name = (
args.wandb_name or f"{Path(args.model).name}-{args.variant}-fsdp"
)
run_name = f"{base_name}-{time.strftime('%m%d-%H%M', time.gmtime())}"
wandb.init(
project=args.wandb_project,
name=run_name,
config=vars(args)
| {"n_examples": len(examples), "total_steps": total_steps},
)
use_wandb = True
except Exception as e:
log(f"wandb off ({e})")
log(f"total_steps={total_steps} steps/epoch={steps_per_epoch} warmup={warmup}")
def _save_ckpt(out_path):
# Full-model checkpoint (gathers FSDP shards to rank0). Called per-epoch
# so a long run leaves usable partial models + lets us eval intermediate
# checkpoints (e.g. epoch1/epoch2) for the data-scaling sweep.
accelerator.wait_for_everyone()
unwrapped = accelerator.unwrap_model(model)
p = Path(out_path)
if is_main:
p.mkdir(parents=True, exist_ok=True)
tok.save_pretrained(str(p))
unwrapped.save_pretrained(
str(p),
is_main_process=is_main,
save_function=accelerator.save,
state_dict=accelerator.get_state_dict(model),
)
if is_main:
log(f"checkpoint saved -> {p}")
model.train()
gstep = 0
t0 = time.time()
if val_dl is not None:
vloss = evaluate()
log(f"val/loss (baseline, step 0) = {vloss:.4f}")
if use_wandb:
wandb.log({"val/loss": vloss}, step=0)
for epoch in range(math.ceil(args.epochs)):
for ii, lab, am in dl:
with accelerator.accumulate(model):
out = model(input_ids=ii, attention_mask=am, labels=lab)
accelerator.backward(out.loss)
if accelerator.sync_gradients:
lr = args.lr * lr_at(gstep)
for g in opt.param_groups:
g["lr"] = lr
accelerator.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
opt.zero_grad()
if accelerator.sync_gradients:
gstep += 1
# accelerator.gather() is a COLLECTIVE — EVERY rank must call it.
# Guarding it behind `is_main` made only rank0 run the [1]-float
# all-gather while ranks 1-3 marched on to the next step's FSDP
# param all-gather, desyncing the process group and hanging the job
# at the step1->step2 boundary (NCCL collective-shape mismatch).
loss = accelerator.gather(out.loss.detach()).mean().item()
if is_main:
log(f"step {gstep}/{total_steps} loss={loss:.4f} lr={lr:.2e}")
if use_wandb:
wandb.log({"train/loss": loss, "train/lr": lr}, step=gstep)
if gstep >= total_steps:
break
_save_ckpt(Path(args.out) / f"epoch{epoch + 1}")
if val_dl is not None:
vloss = evaluate()
log(f"val/loss (epoch {epoch + 1}) = {vloss:.4f}")
if use_wandb:
wandb.log({"val/loss": vloss, "epoch": epoch + 1}, step=gstep)
if gstep >= total_steps:
break
accelerator.wait_for_everyone()
log("saving full model...")
unwrapped = accelerator.unwrap_model(model)
out_dir = Path(args.out)
if is_main:
out_dir.mkdir(parents=True, exist_ok=True)
tok.save_pretrained(str(out_dir))
unwrapped.save_pretrained(
str(out_dir),
is_main_process=is_main,
save_function=accelerator.save,
state_dict=accelerator.get_state_dict(model),
)
if is_main:
log(f"saved -> {out_dir}")
if use_wandb:
try:
wandb.finish()
except Exception:
pass
print(
"FSDP_SFT_DONE "
+ json.dumps(
{
"variant": args.variant,
"steps": gstep,
"wall_s": round(time.time() - t0, 1),
"out": str(out_dir),
}
)
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+151
View File
@@ -0,0 +1,151 @@
#!/usr/bin/env python
"""Conversation->tokens + assistant-only masking shared by the SFT trainers.
Extracted from the old LoRA trainer so ``run_sft_fsdp.py`` (full-parameter FSDP,
the path we actually use) doesn't depend on it. ``build_example`` tokenizes one
``conversations`` record and returns ``input_ids`` + ``labels`` with everything
but the supervised assistant turns masked to -100.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
def normalize_messages(conversations: List[Dict[str, Any]]) -> List[Dict[str, str]]:
"""Map raw conversation turns to {role, content} with role in
system/user/assistant. tool turns fold into user (matches the native
OrchestratorSFTDataset fallback)."""
msgs: List[Dict[str, str]] = []
for turn in conversations:
role = turn.get("role") or turn.get("from", "")
content = turn.get("content") or turn.get("value", "") or ""
if role in ("human", "user"):
role = "user"
elif role in ("gpt", "assistant"):
role = "assistant"
elif role == "tool":
content = f"[Tool '{turn.get('name', 'tool')}' returned]: {content}"
role = "user"
if role in ("system", "user", "assistant"):
msgs.append({"role": role, "content": str(content)})
return msgs
def _lcp_len(a: List[int], b: List[int]) -> int:
"""Length of the longest common prefix of two token-id lists."""
n = 0
for x, y in zip(a, b):
if x != y:
break
n += 1
return n
def _chatml_assistant_spans(ids: List[int], tok) -> Optional[List[tuple]]:
"""``(start, end)`` token spans of each assistant turn's content in a
ChatML-rendered sequence, located by ``<|im_start|>assistant ... <|im_end|>``
markers (the ``<|im_end|>`` is included so the model learns to stop).
Operates on the FINAL token stream, so it is correct even when the chat
template strips prior-turn ``<think>`` reasoning from history (Qwen3.x does),
which breaks any incremental re-rendering / prefix approach. Returns None for
non-ChatML templates (e.g. gemma) so the caller can fall back."""
im_start = tok.convert_tokens_to_ids("<|im_start|>")
im_end = tok.convert_tokens_to_ids("<|im_end|>")
asst = tok.convert_tokens_to_ids("assistant")
unk = tok.unk_token_id
if im_start is None or im_end is None or im_start == unk or im_end == unk:
return None
nl = tok.convert_tokens_to_ids("Ċ") # the '\n' after the role header
spans: List[tuple] = []
i, n = 0, len(ids)
while i < n:
if ids[i] == im_start and i + 1 < n and ids[i + 1] == asst:
j = i + 2
if j < n and ids[j] == nl:
j += 1
k = j
while k < n and ids[k] != im_end:
k += 1
end = min(k + 1, n)
if end > j:
spans.append((j, end))
i = end
else:
i += 1
return spans
def build_example(
tok,
conversations,
max_seq: int,
supervise_all_turns: bool = True,
) -> Optional[Dict[str, List[int]]]:
"""Tokenize one record -> input_ids + labels.
``supervise_all_turns=True`` (default): supervise EVERY assistant turn — the
intermediate routing ``<tool_call>`` turns AND the final answer; only
system/user/tool tokens are masked. This teaches the model to actually route,
not just synthesize the last answer given someone else's tool calls.
``supervise_all_turns=False`` (legacy): supervise only the final assistant
turn; everything before it is masked to -100.
Assistant spans are found by the ChatML turn markers on the rendered stream
(robust to Qwen3's prior-turn think-stripping). For non-ChatML templates
(gemma) we fall back to the proven last-turn longest-common-prefix boundary;
all-turns supervision there needs per-template markers and is not yet wired,
so it degrades to last-turn. Records whose last turn isn't assistant, that
supervise nothing, or whose final turn alone overflows max_seq are skipped."""
msgs = normalize_messages(conversations)
if len(msgs) < 2 or msgs[-1]["role"] != "assistant":
return None
try:
full_text = tok.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=False
)
full = tok(full_text, add_special_tokens=False)["input_ids"]
except Exception:
return None
if not full:
return None
spans = _chatml_assistant_spans(full, tok)
if spans:
selected = spans if supervise_all_turns else spans[-1:]
labels = [-100] * len(full)
for s, e in selected:
for k in range(s, e):
labels[k] = full[k]
last_start = selected[-1][0]
else:
# Non-ChatML template: proven last-turn boundary via longest-common-prefix
# (robust to templates whose add_generation_prompt emits extra preamble).
try:
prompt_text = tok.apply_chat_template(
msgs[:-1], tokenize=False, add_generation_prompt=True
)
prompt = tok(prompt_text, add_special_tokens=False)["input_ids"]
except Exception:
return None
boundary = _lcp_len(prompt, full)
if boundary >= len(full):
return None
labels = list(full)
for i in range(min(boundary, len(labels))):
labels[i] = -100
last_start = boundary
if all(v == -100 for v in labels): # nothing left to supervise
return None
# Keep the final assistant turn if too long: left-truncate from the front.
if len(full) > max_seq:
if len(full) - last_start >= max_seq: # final answer alone overflows -> drop
return None
cut = len(full) - max_seq
full = full[cut:]
labels = labels[cut:]
return {"input_ids": full, "labels": labels}
+25
View File
@@ -0,0 +1,25 @@
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false
fsdp_config:
fsdp_version: 1
fsdp_auto_wrap_policy: SIZE_BASED_WRAP
fsdp_min_num_params: 100000000
fsdp_sharding_strategy: FULL_SHARD
fsdp_state_dict_type: FULL_STATE_DICT
fsdp_offload_params: false
fsdp_backward_prefetch: BACKWARD_PRE
fsdp_forward_prefetch: false
fsdp_use_orig_params: true
fsdp_cpu_ram_efficient_loading: true
fsdp_sync_module_states: true
+25
View File
@@ -0,0 +1,25 @@
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
downcast_bf16: 'no'
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
use_cpu: false
fsdp_config:
fsdp_version: 1
fsdp_auto_wrap_policy: SIZE_BASED_WRAP
fsdp_min_num_params: 100000000
fsdp_sharding_strategy: FULL_SHARD
fsdp_state_dict_type: FULL_STATE_DICT
fsdp_offload_params: false
fsdp_backward_prefetch: BACKWARD_PRE
fsdp_forward_prefetch: false
fsdp_use_orig_params: true
fsdp_cpu_ram_efficient_loading: true
fsdp_sync_module_states: true
+4 -4
View File
@@ -445,8 +445,8 @@ class LocalCloudAgent(BaseAgent):
tools: Optional[list] = None,
tool_choice: Optional[dict] = None,
output_config: Optional[dict] = None,
timeout: float = 600.0,
max_retries: int = 12,
timeout: float = 60.0,
max_retries: int = 2,
trace_role: str = "cloud",
) -> Tuple[str, int, int, int]:
"""Single Anthropic call. Returns (text, p_tok, c_tok, n_web_searches).
@@ -536,7 +536,7 @@ class LocalCloudAgent(BaseAgent):
response_format: Optional[dict] = None,
tools: Optional[list] = None,
tool_choice: Optional[Any] = None,
timeout: float = 600.0,
timeout: float = 60.0,
trace_role: str = "cloud",
) -> Tuple[str, int, int]:
"""Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured;
@@ -604,7 +604,7 @@ class LocalCloudAgent(BaseAgent):
system: Optional[str] = None,
max_tokens: int = 4096,
temperature: float = 0.0,
timeout: float = 600.0,
timeout: float = 60.0,
trace_role: str = "cloud",
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, int, int]:
+16
View File
@@ -11,6 +11,7 @@ from __future__ import annotations
# USD per million tokens, (input, output). Local models = 0.
PRICES: dict[str, tuple[float, float]] = {
"claude-opus-4-8": (5.00, 25.0),
"claude-opus-4-7": (5.00, 25.0),
"claude-sonnet-4-6": (3.00, 15.0),
"claude-haiku-4-5": (1.00, 5.00),
@@ -33,6 +34,17 @@ PRICES: dict[str, tuple[float, float]] = {
"qwen/qwen-2.5-coder-32b-instruct": (0.08, 0.18),
"qwen/qwen3-32b": (0.10, 0.30),
"meta-llama/llama-3.3-70b-instruct": (0.13, 0.39),
# OpenRouter slugs for the orchestrator's local-OSS class when routed via
# OpenRouter instead of self-hosted vLLM. OpenRouter list price, checked
# 2026-07-11. These replace earlier active-param-scaled guesses that were
# 2-5x low on output; the cost-aware reward reads these directly, so the
# guesses were systematically undercharging the mid/large Qwen experts.
"qwen/qwen3.5-9b": (0.10, 0.15),
# qwen3.6-27b is not listed on OpenRouter; qwen3.5-27b is the nearest real
# quote and is what we charge for it.
"qwen/qwen3.6-27b": (0.195, 1.56),
"qwen/qwen3.5-122b-a10b": (0.26, 2.08),
"qwen/qwen3.5-397b-a17b": (0.385, 2.45),
}
# Models whose API rejects an explicit `temperature` param — callers should
@@ -41,6 +53,10 @@ NO_TEMP_PREFIXES: tuple[str, ...] = (
"claude-opus-4-7",
"claude-sonnet-4-7",
"claude-haiku-4-7",
# 4-8 family also rejects `temperature` ("deprecated for this model").
"claude-opus-4-8",
"claude-sonnet-4-8",
"claude-haiku-4-8",
)
@@ -0,0 +1,751 @@
"""Faithful ToolOrchestra "unified tool calling" registry (arXiv:2511.21689 §3.1).
The paper exposes **every tool AND every model through a single flat tool
interface** — each is its own named function with a description and a typed
parameter schema, and for each training instance a *random subset* of tools is
sampled with *randomized pricing* (§3.3, "General tool configuration"). This is
unlike the eval-port shortcut in ``toolorchestra.py``, which collapses the whole
catalog into three meta-tools (``search``/``enhance_reasoning``/``answer``) with
a ``model`` slot. This module restores the faithful design.
Each :class:`ExpertTool` knows:
* the orchestrator-visible ``name`` / ``description`` / param schema (what goes
into the tools JSON the policy conditions on), and
* the concrete backend (``backend_type`` + ``model`` + ``base_url``) so a caller
can turn it into the worker dict that ``toolorchestra._call_worker`` dispatches.
Everything here is pure data + deterministic transforms (no network, no model
calls), so the spec building, sampling, and pricing logic is offline-testable.
Dispatch stays in ``toolorchestra.py`` (via :func:`to_worker_dict`) to avoid a
circular import.
"""
from __future__ import annotations
import os
import re
from dataclasses import dataclass
from typing import Dict, List, Optional
from openjarvis.agents.hybrid._prices import PRICES
# Kinds of tool in the unified interface.
KIND_MODEL = "model" # an LLM exposed as a tool (the paper's "models as tools")
KIND_WEB_SEARCH = "web_search"
KIND_CODE = "code_interpreter"
KIND_TOOL = "tool" # a bridged real OpenJarvis tool (custom param schema)
VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_CODE, KIND_TOOL)
# Flat-catalog category label, surfaced in each tool's spec so the orchestrator
# can tell tool types apart without us imposing any hierarchy (the menu stays
# flat — this is just a tag).
CATEGORY_BASIC = "basic_tool"
# Two-model-class taxonomy for the orchestrator catalog — the only model tiers.
CATEGORY_CLOUD_FRONTIER = "cloud_frontier"
CATEGORY_LOCAL_OSS = "local_open_source"
# Backend dispatch types understood by ``toolorchestra._call_worker`` (plus the
# ``openjarvis-tool`` bridge, dispatched in ``unified.make_dispatch`` via the
# OpenJarvis ToolExecutor rather than ``_call_worker``).
VALID_BACKENDS = (
"vllm",
"openai",
"anthropic",
"gemini",
"openrouter",
"anthropic-web-search",
"tavily-search",
"modal-python",
"openjarvis-tool",
)
@dataclass(frozen=True)
class ExpertTool:
"""One entry in the unified tool catalog.
``price_in`` / ``price_out`` are USD per 1M tokens (0.0 for local / non-LLM
tools). ``latency_s`` is a rough average used only to populate the
description's cost/latency line — the orchestrator was trained to read that
table, so we surface it verbatim in the spec.
"""
name: str
kind: str
backend_type: str
summary: str
model: Optional[str] = None
base_url: Optional[str] = None
price_in: float = 0.0
price_out: float = 0.0
latency_s: float = 5.0
category: str = "" # cloud_frontier | local_open_source | basic_tool
# Optional explicit JSON-schema for the tool's arguments. Set for bridged
# real OpenJarvis tools (``openjarvis-tool`` backend) whose params don't fit
# the fixed kind-based schemas; takes precedence over the kind default.
param_schema: Optional[dict] = None
# When True, ``description()`` omits the price/latency line — used by the
# anonymized catalog so the orchestrator can't route on cost/identity.
hide_cost: bool = False
def __post_init__(self) -> None:
if self.kind not in VALID_KINDS:
raise ValueError(f"{self.name}: invalid kind {self.kind!r}")
if self.backend_type not in VALID_BACKENDS:
raise ValueError(f"{self.name}: invalid backend {self.backend_type!r}")
if self.kind == KIND_MODEL and not self.model:
raise ValueError(f"{self.name}: model-kind tool needs a concrete model")
# ---- orchestrator-visible spec -------------------------------------
def _param_schema(self) -> Dict[str, object]:
"""JSON-schema for the tool's arguments (one typed param per kind).
An explicit ``param_schema`` (set by :func:`openjarvis_tool` for bridged
real tools) overrides the kind-based default.
"""
if self.param_schema is not None:
return self.param_schema
if self.kind == KIND_WEB_SEARCH:
return {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query string.",
}
},
"required": ["query"],
}
if self.kind == KIND_CODE:
return {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python code to execute. Print results.",
}
},
"required": ["code"],
}
# model tool
return {
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "The sub-question or instruction for this model.",
}
},
"required": ["input"],
}
def description(self) -> str:
"""Full description incl. the price/latency line (paper bakes this in)."""
if self.hide_cost:
return self.summary
if self.kind == KIND_MODEL:
cost_line = (
f" Pricing: ${self.price_in:.2f}/1M input, "
f"${self.price_out:.2f}/1M output; avg latency ~{self.latency_s:.0f}s."
)
else:
cost_line = f" Avg latency ~{self.latency_s:.0f}s."
return self.summary.rstrip(".") + "." + cost_line
def to_spec(self) -> Dict[str, object]:
"""OpenAI-style tool spec the orchestrator conditions on.
Flat list, but each function carries a ``category`` tag so the policy can
distinguish generalist vs specialized models vs basic tools.
"""
fn: Dict[str, object] = {
"name": self.name,
"description": self.description(),
"parameters": self._param_schema(),
}
if self.category:
fn["category"] = self.category
return {"type": "function", "function": fn}
def _price(model: str) -> tuple[float, float]:
return PRICES.get(model, (0.0, 0.0))
def _tool_name(model: str) -> str:
"""Tool-safe function name derived from a model id (``qwen3-8b`` -> ``qwen3_8b``).
Strips any provider prefix and replaces non-alphanumerics with underscores so
the catalog exposes one named tool per concrete model.
"""
base = model.split("/")[-1].lower()
safe = re.sub(r"[^a-z0-9]+", "_", base).strip("_")
return safe or "local_model"
_ANON_ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789"
# Rough size/scale hint per model — a capability signal WITHOUT the brand name,
# so the orchestrator can route big-vs-small on task difficulty but can't route on
# a proprietary name/tier prior. Closed models have no public param count, so they
# get a scale tier; open models get their param count (else parsed from the name).
_MODEL_SIZE = {
# The real orchestrator_catalog: 2 cloud frontier + 4 local OSS Qwen.
"gpt-5.5": "frontier-scale",
"claude-opus-4-8": "frontier-scale",
"Qwen/Qwen3.5-9B": "~9B params",
"Qwen/Qwen3.6-27B-FP8": "~27B params",
"Qwen/Qwen3.5-122B-A10B-FP8": "~122B total, ~10B active (MoE)",
"Qwen/Qwen3.5-397B-A17B-FP8": "~397B total, ~17B active (MoE)",
# OpenRouter slugs (the tool.model id when these route through OpenRouter).
"qwen/qwen3.5-122b-a10b": "~122B total, ~10B active (MoE)",
"qwen/qwen3.5-397b-a17b": "~397B total, ~17B active (MoE)",
}
def _size_hint(model: str) -> str:
"""Anonymized scale descriptor for an expert (no brand). Falls back to a
param count parsed from the model id (``...-9b`` -> ``~9B params``)."""
if model in _MODEL_SIZE:
return _MODEL_SIZE[model]
m = re.search(r"(\d+(?:\.\d+)?)\s*b\b", model.lower())
return f"~{m.group(1)}B params" if m else "unspecified scale"
def _class_hint(tool) -> str:
"""Which of the two model classes this expert belongs to (anonymized)."""
cat = getattr(tool, "category", "")
if cat == CATEGORY_CLOUD_FRONTIER:
return "cloud frontier"
if cat == CATEGORY_LOCAL_OSS:
return "local open-source"
return ""
def _cost_tier_hint(tool) -> str:
"""Coarse cost tier (not exact pricing — that's the bias we anonymize away).
Shown so the policy learns the strong-but-expensive vs cheap tradeoff instead
of always delegating to the biggest model. Driven by class: cloud frontier is
expensive, self-hosted open-source is cheap."""
cat = getattr(tool, "category", "")
if cat == CATEGORY_CLOUD_FRONTIER:
return "expensive"
if cat == CATEGORY_LOCAL_OSS:
return "cheap"
if getattr(tool, "backend_type", "") == "vllm" or (tool.price_out or 0) <= 0:
return "cheap"
po = tool.price_out
return "cheap" if po < 5 else "moderate" if po < 20 else "expensive"
def anonymize_tools(tools, rng):
"""Strip model identity for unbiased routing data.
Each MODEL expert is replaced with an opaque random label
(``expert_<4 rand>``), a uniform description, a uniform category and no
price/latency line — so the orchestrator cannot route on a model's name,
position, cost, or tier (all of which we found dominate the choice). The full
list is shuffled to kill position bias. Basic tools (calculator, web_search,
…) keep their real names — the policy must still know what they do.
Returns ``(anon_tools, anon_to_real)`` where ``anon_to_real`` maps each opaque
label back to the real tool name (for offline analysis only; never shown to
the model). ``.model`` is preserved on each tool so dispatch still reaches the
right backend. Pass a fresh ``rng`` per rollout so labels don't stabilise.
"""
from dataclasses import replace
anon_to_real: Dict[str, str] = {}
experts: List[ExpertTool] = []
basics: List[ExpertTool] = []
for t in tools:
if t.kind == KIND_MODEL:
tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
while tag in anon_to_real:
tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
anon_to_real[tag] = t.name
bits = [
b
for b in (_class_hint(t), _size_hint(t.model), _cost_tier_hint(t))
if b
]
experts.append(
replace(
t,
name=tag,
summary="Another model — "
+ ", ".join(bits)
+ ". Send it a sub-question.",
category="model",
hide_cost=True,
)
)
else:
basics.append(t)
# Shuffle WITHIN the experts to kill per-expert position bias, but keep all
# experts as a block on top and the basic tools underneath — a clean split
# (experts first) rather than experts and utilities interleaved.
rng.shuffle(experts)
out = experts + basics
return out, anon_to_real
def openjarvis_tool(
registered_name: str,
*,
summary: str,
params: dict,
latency_s: float = 5.0,
) -> ExpertTool:
"""Build an :class:`ExpertTool` that bridges a real OpenJarvis tool.
``registered_name`` is the tool's key in ``ToolRegistry`` (e.g. ``calculator``,
``shell_exec``). ``params`` is the JSON-schema *properties*-style dict for the
tool's arguments; it is surfaced verbatim by :meth:`ExpertTool.to_spec`. The
resulting tool dispatches through the OpenJarvis ``ToolExecutor`` (backend
``openjarvis-tool``) rather than ``_call_worker``.
"""
return ExpertTool(
name=registered_name,
kind=KIND_TOOL,
backend_type="openjarvis-tool",
summary=summary,
model=registered_name,
latency_s=latency_s,
category=CATEGORY_BASIC,
param_schema=params,
)
# Real OpenJarvis tools bridged into the orchestrator catalog as basic tools.
# Names must match the ``ToolRegistry`` keys (confirmed present: calculator,
# shell_exec, file_read, file_write, http_request).
def _openjarvis_basic_tools() -> List[ExpertTool]:
def obj(properties: dict, required: List[str]) -> dict:
return {"type": "object", "properties": properties, "required": required}
return [
openjarvis_tool(
"calculator",
summary="Evaluate an arithmetic / math expression and return the result.",
params=obj(
{
"expression": {
"type": "string",
"description": "Math expression to evaluate.",
}
},
["expression"],
),
latency_s=1.0,
),
openjarvis_tool(
"shell_exec",
summary=(
"Run a shell command and return its stdout/stderr. Critical "
"for terminal / TerminalBench-style tasks."
),
params=obj(
{
"command": {
"type": "string",
"description": "Shell command to execute.",
}
},
["command"],
),
latency_s=4.0,
),
openjarvis_tool(
"file_read",
summary="Read the contents of a file at the given path.",
params=obj(
{
"path": {
"type": "string",
"description": "Path of the file to read.",
}
},
["path"],
),
latency_s=1.0,
),
openjarvis_tool(
"file_write",
summary="Write content to a file at the given path.",
params=obj(
{
"path": {
"type": "string",
"description": "Path of the file to write.",
},
"content": {"type": "string", "description": "Content to write."},
},
["path", "content"],
),
latency_s=1.0,
),
openjarvis_tool(
"http_request",
summary="Make an HTTP request to a URL and return the response body.",
params={
"type": "object",
"properties": {
"url": {"type": "string", "description": "Request URL."},
"method": {
"type": "string",
"description": "HTTP method (GET, POST, ...). Default GET.",
},
},
"required": ["url"],
},
latency_s=4.0,
),
openjarvis_tool(
"think",
summary=(
"Record a private reasoning step (scratchpad). No external "
"effect; use to plan before acting on hard reasoning tasks."
),
params=obj(
{
"thought": {
"type": "string",
"description": "Your reasoning or thought process.",
}
},
["thought"],
),
latency_s=0.5,
),
openjarvis_tool(
"apply_patch",
summary=(
"Apply a unified-diff patch to a file. Use to edit code for "
"terminal / SWE-style tasks."
),
params=obj(
{
"patch": {
"type": "string",
"description": "The unified diff patch text to apply.",
},
"path": {
"type": "string",
"description": "Target file path (auto-detected from the "
"patch header if omitted).",
},
},
["patch"],
),
latency_s=2.0,
),
openjarvis_tool(
"pdf_extract",
summary=(
"Extract text from a PDF file. Use for GAIA-style tasks with "
"PDF attachments."
),
params=obj(
{
"file_path": {
"type": "string",
"description": "Path to the PDF file.",
},
"pages": {
"type": "string",
"description": "Page range, e.g. '1-5' or '1,3,5'. "
"Omit for all pages.",
},
},
["file_path"],
),
latency_s=3.0,
),
openjarvis_tool(
"db_query",
summary=(
"Run a SQL query against a SQLite/Postgres database and return "
"rows. Read-only by default."
),
params=obj(
{
"query": {"type": "string", "description": "SQL query to execute."},
"db_path": {
"type": "string",
"description": "Path to a SQLite DB file. Defaults to "
"in-memory.",
},
"read_only": {
"type": "boolean",
"description": "Restrict to SELECT/EXPLAIN/PRAGMA. "
"Default: true.",
},
},
["query"],
),
latency_s=3.0,
),
]
# Local-Cloud Hybrid orchestrator catalog. The menu is two model classes — cloud
# frontier models and open-source models — plus the basic tools (web search, code
# interpreter, and the bridged real OpenJarvis tools). The orchestrator model
# itself is NOT in the catalog.
#
# Routing default is OpenRouter for every model tool, so the catalog works with
# no self-hosted servers. A model is routed to local vLLM instead when its id is
# in ``local_endpoints`` or ``model_backends`` overrides it. The
# orchestrator-visible tool *name* is always derived from the canonical id, so
# routing can change (vLLM <-> OpenRouter <-> native API) without shifting the
# menu the policy was trained on.
_LOCAL_OSS_MODELS = (
# (canonical_id, openrouter_slug)
("Qwen/Qwen3.5-9B", "qwen/qwen3.5-9b"),
("Qwen/Qwen3.6-27B-FP8", "qwen/qwen3.6-27b"),
("Qwen/Qwen3.5-122B-A10B-FP8", "qwen/qwen3.5-122b-a10b"),
("Qwen/Qwen3.5-397B-A17B-FP8", "qwen/qwen3.5-397b-a17b"),
)
_CLOUD_FRONTIER_MODELS = (
# (canonical, native_backend, openrouter_slug, summary, latency_s)
# Neutral, uniform summaries (no capability ranking) so the orchestrator
# doesn't just pick whichever model is labelled "strongest" — routing should
# be learned from the reward, not hand-labelled here.
("gpt-5.5", "openai", "openai/gpt-5.5", "Expert model (GPT-5.5).", 30.0),
(
"claude-opus-4-8",
"anthropic",
"anthropic/claude-opus-4.8",
"Expert model (Claude Opus 4.8).",
26.0,
),
)
def _model_tool(
canonical: str,
*,
native_backend: str,
or_slug: str,
summary: str,
lat: float,
category: str,
local_endpoints: Dict[str, str],
model_backends: Dict[str, str],
openrouter_slugs: Dict[str, str],
) -> ExpertTool:
"""Build one model tool, resolving its backend.
Backend precedence: explicit ``model_backends[canonical]`` > vLLM if the model
has a ``local_endpoints`` entry > the model's NATIVE provider (openai /
anthropic / gemini) when it has one > OpenRouter. Frontier models thus hit
their first-party API by default (OpenRouter's gpt-5.5 was returning 400s);
OSS Qwen experts (native_backend="vllm") fall through to OpenRouter.
"""
backend = (
model_backends.get(canonical)
or ("vllm" if canonical in local_endpoints else None)
or (
native_backend
if native_backend in ("openai", "anthropic", "gemini")
else "openrouter"
)
)
name = _tool_name(canonical)
if backend == "vllm":
# Self-hosted: free per the cost model.
return ExpertTool(
name=name,
kind=KIND_MODEL,
backend_type="vllm",
summary=summary,
model=canonical,
base_url=local_endpoints.get(canonical),
price_in=0.0,
price_out=0.0,
latency_s=lat,
category=category,
)
if backend == "openrouter":
slug = openrouter_slugs.get(canonical, or_slug)
pi, po = _price(slug)
if (pi, po) == (0.0, 0.0): # fall back to the canonical id's price
pi, po = _price(canonical)
return ExpertTool(
name=name,
kind=KIND_MODEL,
backend_type="openrouter",
summary=summary,
model=slug,
base_url=None,
price_in=pi,
price_out=po,
latency_s=lat,
category=category,
)
# native provider API (openai / anthropic / gemini)
pi, po = _price(canonical)
return ExpertTool(
name=name,
kind=KIND_MODEL,
backend_type=backend,
summary=summary,
model=canonical,
price_in=pi,
price_out=po,
latency_s=lat,
category=category,
)
def orchestrator_catalog(
*,
local_endpoints: Optional[Dict[str, str]] = None,
model_backends: Optional[Dict[str, str]] = None,
openrouter_slugs: Optional[Dict[str, str]] = None,
include_tools: bool = True,
) -> List[ExpertTool]:
"""Return the orchestrator's tool catalog: two model classes + basic tools.
Routing for the model tools defaults to **OpenRouter** (so the catalog works
with no self-hosted servers). Overrides, in precedence order:
* ``model_backends`` maps a canonical model id -> ``"vllm" | "openrouter" |
"openai" | "anthropic" | "gemini"`` to force that model's backend.
* ``local_endpoints`` maps a canonical id (e.g. ``"Qwen/Qwen3.5-9B"``) to a
vLLM ``base_url``; a model present here routes to vLLM (free) unless
``model_backends`` says otherwise.
* ``openrouter_slugs`` overrides the per-model OpenRouter slug used when a
model routes through OpenRouter.
``include_tools`` (default True) appends the basic tools — web search, code
interpreter, and the bridged real OpenJarvis tools (calculator, shell_exec,
file_read, file_write, http_request).
"""
local_endpoints = local_endpoints or {}
model_backends = model_backends or {}
openrouter_slugs = openrouter_slugs or {}
cat: List[ExpertTool] = []
# Env-gated expert exclusion: OJ_EXCLUDE_EXPERTS is a comma-separated list of
# case-insensitive substrings matched against a model's canonical id. Any
# match is skipped from the catalog. Used to temporarily drop unreliable
# experts (e.g. OpenRouter giants during a provider outage) without editing
# the registry — unset the var to restore them.
_excl = {
s.strip().lower()
for s in os.environ.get("OJ_EXCLUDE_EXPERTS", "").split(",")
if s.strip()
}
def _excluded(canonical: str) -> bool:
c = canonical.lower()
return any(x in c for x in _excl)
# ---- cloud frontier models ----
for canonical, native_backend, or_slug, summary, lat in _CLOUD_FRONTIER_MODELS:
if _excluded(canonical):
continue
cat.append(
_model_tool(
canonical,
native_backend=native_backend,
or_slug=or_slug,
summary=summary,
lat=lat,
category=CATEGORY_CLOUD_FRONTIER,
local_endpoints=local_endpoints,
model_backends=model_backends,
openrouter_slugs=openrouter_slugs,
)
)
# ---- open-source models (OpenRouter by default; vLLM when an endpoint or
# a model_backends override is supplied) ----
for canonical, or_slug in _LOCAL_OSS_MODELS:
if _excluded(canonical):
continue
cat.append(
_model_tool(
canonical,
native_backend="vllm",
or_slug=or_slug,
summary=f"Expert model ({canonical}).",
lat=4.0,
category=CATEGORY_LOCAL_OSS,
local_endpoints=local_endpoints,
model_backends=model_backends,
openrouter_slugs=openrouter_slugs,
)
)
if include_tools:
# ---- basic tools ----
cat.append(
ExpertTool(
name="web_search",
kind=KIND_WEB_SEARCH,
backend_type="tavily-search",
summary="Web search (Tavily). Use for facts that need a live lookup.",
model="tavily",
latency_s=8.0,
category=CATEGORY_BASIC,
)
)
cat.append(
ExpertTool(
name="code_interpreter",
kind=KIND_CODE,
backend_type="modal-python",
summary="Python sandbox. Execute code and return stdout/stderr.",
model="modal-python",
latency_s=6.0,
category=CATEGORY_BASIC,
)
)
cat.extend(_openjarvis_basic_tools())
return cat
def build_tool_specs(tools: List[ExpertTool]) -> List[Dict[str, object]]:
"""Turn a tool list into the OpenAI-style tools JSON the policy sees."""
return [t.to_spec() for t in tools]
def tools_by_name(tools: List[ExpertTool]) -> Dict[str, ExpertTool]:
return {t.name: t for t in tools}
def to_worker_dict(tool: ExpertTool) -> Dict[str, object]:
"""Convert a tool into the worker dict ``toolorchestra._call_worker`` eats."""
d: Dict[str, object] = {
"name": tool.name,
"type": tool.backend_type,
"model": tool.model,
}
if tool.base_url:
d["base_url"] = tool.base_url
return d
__all__ = [
"CATEGORY_BASIC",
"CATEGORY_CLOUD_FRONTIER",
"CATEGORY_LOCAL_OSS",
"ExpertTool",
"KIND_CODE",
"KIND_MODEL",
"KIND_TOOL",
"KIND_WEB_SEARCH",
"build_tool_specs",
"openjarvis_tool",
"orchestrator_catalog",
"to_worker_dict",
"tools_by_name",
]
+98
View File
@@ -0,0 +1,98 @@
"""Exponential backoff + jitter for the cloud calls a rollout makes.
The judge already had this (``evals/core/scorer.py``) — it was added after a run
where 93/100 GAIA judge calls 429'd and zeroed the whole bench. The rollout's own
two cloud paths did NOT:
* the ORCHESTRATOR call (``toolorchestra/unified.py``), and
* the EXPERT dispatch (``expert_registry.py``).
Without backoff a single transient 429 makes the expert return an empty/error
observation, the clean gate then rejects the whole trajectory ("error or empty
tool observation"), and the rollout is silently wasted. That's tolerable at
concurrency 12; at 100 it would shred the batch. Same policy as the judge:
retry only on transient errors, exponential + jittered, then give up and let the
caller record the failure.
"""
from __future__ import annotations
import logging
import random
import time
from typing import Callable, TypeVar
LOGGER = logging.getLogger(__name__)
MAX_RETRIES = 6
BASE_DELAY_S = 2.0
MAX_DELAY_S = 60.0
# Transient / server-side failures. A 400 (bad request) or 401 (bad key) is NOT
# here on purpose: retrying those just burns time on a deterministic failure.
RETRYABLE_MARKERS = (
"429",
"rate_limit",
"rate limit",
"overloaded",
"timeout",
"timed out",
"503",
"502",
"500",
"connection",
"temporarily unavailable",
"internalservererror",
"apiconnectionerror",
# An expert that returns 200-OK with an EMPTY body. The HTTP call "succeeds",
# so nothing raises and no retry fires — the rollout just gets an empty
# observation and the clean gate then bins the whole trajectory. Seen from the
# OpenRouter-hosted Qwen 122B/397B (audit 2026-07-13: 6 of 72 rollouts). We
# raise EmptyExpertResponse for it so it retries like any other transient.
"empty expert response",
)
class EmptyExpertResponse(RuntimeError):
"""An expert returned 200-OK with no content — transient, worth retrying."""
def __init__(self, model: str) -> None:
super().__init__(f"empty expert response from {model}")
T = TypeVar("T")
def is_retryable(exc: Exception) -> bool:
msg = str(exc).lower()
return any(marker in msg for marker in RETRYABLE_MARKERS)
def with_backoff(fn: Callable[[], T], *, what: str = "cloud call") -> T:
"""Run ``fn``, retrying transient failures with exponential backoff + jitter.
Re-raises the last exception on a non-retryable error or once the retry
budget is exhausted, so the caller still sees the failure.
"""
last_exc: Exception | None = None
for attempt in range(MAX_RETRIES):
try:
return fn()
except Exception as exc: # noqa: BLE001 — re-raised below
last_exc = exc
if attempt == MAX_RETRIES - 1 or not is_retryable(exc):
raise
delay = min(BASE_DELAY_S * (2**attempt), MAX_DELAY_S)
delay += random.uniform(0.0, delay * 0.25) # jitter: de-sync the herd
LOGGER.warning(
"%s failed (attempt %d/%d): %s — retrying in %.1fs",
what,
attempt + 1,
MAX_RETRIES,
exc,
delay,
)
time.sleep(delay)
raise last_exc # type: ignore[misc]
__all__ = ["with_backoff", "is_retryable", "MAX_RETRIES"]
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,20 @@
"""ToolOrchestraAgent package (split from the former toolorchestra.py module).
Importing this package registers the agent and re-exports the public surface,
so ``from openjarvis.agents.hybrid.toolorchestra import ToolOrchestraAgent``
keeps working unchanged. Submodules:
prompts system prompts, RL tool specs / arg schema
experts slot -> backend worker mapping (default + paper-match)
sandbox Tavily search + Modal Python sandbox helpers
clients orchestrator vLLM tool-call client
parsing action / tool-call parsing + user-prompt assembly
workers worker pool resolution + dispatch (_call_worker etc.)
agent ToolOrchestraAgent (the registered agent class)
"""
from __future__ import annotations
from openjarvis.agents.hybrid.toolorchestra.agent import ToolOrchestraAgent
__all__ = ["ToolOrchestraAgent"]
@@ -0,0 +1,848 @@
"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
Two modes, gated by ``method_cfg.orchestrator_mode``:
* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
orchestrator, dispatching to a numbered worker pool via JSON
``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
a prompted upper-bound reference point NOT the paper's setup.
* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
served on a local vLLM is the orchestrator. It emits OpenAI-style
``tool_calls`` (or ``<tool_call>{...}</tool_call>`` text blocks when
vLLM's tool parser doesn't catch them) for three expert tools
``enhance_reasoning``, ``answer``, ``search`` exactly as in the
upstream ``evaluation/tools.json``. Each tool's ``model`` arg
(``answer-1``, ``reasoner-2``, ``search-3``, ) is mapped to a real
backend through ``EXPERT_MODEL_MAPPING`` by default the frontier
Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
for `*-3`. Search routes to the Anthropic server-side web_search.
We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
worker types. Energy-wise, "expert" answers are cloud calls.
Pipeline per task (RL mode):
1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
appropriate tool.` with the three tools declared.
2. It emits one ``tool_call`` per turn ``search`` updates the
context, ``enhance_reasoning`` appends code/exec output (we run the
tool as a plain LLM call, no sandbox the model just gets prose
back), ``answer`` produces the final answer and the loop stops.
3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
back to the strongest expert worker.
Prompted-mode pipeline:
1. Orchestrator (cloud) reads question + numbered worker pool.
2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
"input": str}`` or ``{"action": "final_answer", "answer": str}``.
3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
prompt; fallback to strongest worker on parse failure.
Workers come from ``cfg["workers"]`` or a sensible default pool (local
Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
gpt-5-mini).
"""
from __future__ import annotations
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from openjarvis.agents._stubs import AgentContext
from openjarvis.agents.hybrid._base import LocalCloudAgent
from openjarvis.agents.hybrid._prices import PRICES
from openjarvis.agents.hybrid.mini_swe_agent import (
_clone_repo,
_extract_diff,
)
from openjarvis.agents.hybrid.toolorchestra.clients import (
_call_orchestrator_with_tool_calls,
)
from openjarvis.agents.hybrid.toolorchestra.experts import (
_PAPER_CODER_OPENROUTER,
_expert_for,
_paper_expert_for,
)
from openjarvis.agents.hybrid.toolorchestra.parsing import (
_build_user_prompt,
_extract_final_answer_text,
_parse_action,
_parse_rl_tool_call,
)
from openjarvis.agents.hybrid.toolorchestra.prompts import (
FORCE_FINAL_PROMPT,
ORCHESTRATOR_SYS,
RL_ALL_TOOLS,
RL_ORCHESTRATOR_SYS,
RL_TOOLS_SPEC,
)
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
_call_modal_python,
_extract_first_python_block,
)
from openjarvis.agents.hybrid.toolorchestra.workers import (
_TOOLORCH_SEARCH_TYPES,
_call_worker,
_resolve_worker_pool,
_swe_call_worker,
)
from openjarvis.core.registry import AgentRegistry
@AgentRegistry.register("toolorchestra")
class ToolOrchestraAgent(LocalCloudAgent):
"""Multi-turn dispatcher over a mixed worker pool.
Two modes (see module docstring): ``method_cfg.orchestrator_mode``
is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
(paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
"""
agent_id = "toolorchestra"
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
# Validate `method_cfg.worker_pool` early — surfaces config errors
# at agent construction rather than on the first task. No-op when
# the override is absent.
if self._cfg.get("worker_pool") is not None:
_resolve_worker_pool(
self._cfg,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# Validate `orchestrator_mode` (typo-checked here, not on first task).
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
if mode not in ("prompted", "rl"):
raise ValueError(
f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
f"got {mode!r}"
)
def _run_paradigm(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
if mode == "rl":
return self._run_rl(input, context, **kwargs)
return self._run_prompted(input, context, **kwargs)
# ------------------------------------------------------------------
# Legacy prompted-orchestrator path.
# ------------------------------------------------------------------
def _run_prompted(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
cfg = self._cfg
question = input
# Resolution order (strict replace, no merge):
# 1. `cfg["workers"]` — legacy direct override, used by tests.
# 2. `cfg["worker_pool"]` — cell-config override; validated +
# $local/$cloud substituted.
# 3. `_default_pool(...)` — heterogeneous default.
if cfg.get("workers"):
workers = cfg["workers"]
else:
workers = _resolve_worker_pool(
cfg,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
if not workers:
raise RuntimeError("toolorchestra: empty worker pool")
max_turns = int(cfg.get("max_turns", 6))
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
swe_mode = (
bool(cfg.get("swe_use_agent_loop"))
and bool(task_meta.get("problem_statement"))
and bool(task_meta.get("repo"))
and bool(task_meta.get("base_commit"))
)
shared_workdir: Optional[Path] = None
if swe_mode:
shared_workdir = Path(
tempfile.mkdtemp(
prefix=f"toolorch-swe-{task_meta.get('task_id', 'x')}-"
)
)
try:
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
except Exception:
shutil.rmtree(shared_workdir, ignore_errors=True)
raise
self.record_trace_event(
{
"kind": "toolorchestra_swe_workdir",
"workdir": str(shared_workdir),
"repo": task_meta["repo"],
"base_commit": task_meta["base_commit"],
}
)
# try/finally guards ``shared_workdir`` against exceptions raised
# anywhere in the turn loop, the worker calls, the fallback, or
# the diff-extraction step. Without this, at n=500 SWE-bench an
# exception leaves hundreds of MB of cloned repos in tempdir.
try:
history: List[Dict[str, Any]] = []
tokens_local = 0
tokens_cloud = 0
cost = 0.0
n_web_searches_total = 0
# tool_calls: bash turns from SWE subloops + web_search uses
# from GAIA. Orchestrator dispatch turns are NOT counted (they
# produce text only — calling a worker is one tool call's worth
# of "delegation" but the actual tool action happens inside).
tool_calls = 0
final_answer: Optional[str] = None
forced_final = False
parse_failures = 0
for turn in range(1, max_turns + 1):
sys_prompt = ORCHESTRATOR_SYS
if turn == max_turns and final_answer is None:
sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
forced_final = True
user = _build_user_prompt(question, workers, history)
text, o_in, o_out = self._call_cloud(
user=user,
system=sys_prompt,
max_tokens=orch_max_tokens,
temperature=0.0,
)
tokens_cloud += o_in + o_out
cost += self.cost_usd(self._cloud_model, o_in, o_out)
action = _parse_action(text)
history.append(
{
"role": "orchestrator",
"turn": turn,
"raw": text,
"action": action,
}
)
self.record_trace_event(
{
"kind": "toolorchestra_action",
"turn": turn,
"action": action,
"raw": text,
}
)
if action is None:
parse_failures += 1
if parse_failures >= 2 or forced_final:
final_answer = _extract_final_answer_text(text)
break
continue
kind = action.get("action")
if kind == "final_answer":
final_answer = str(action.get("answer", "")).strip()
break
if kind == "call_worker":
wid = action.get("worker_id")
w_input = action.get("input", "")
if not isinstance(wid, int) or not (0 <= wid < len(workers)):
parse_failures += 1
if parse_failures >= 2 or forced_final:
final_answer = _extract_final_answer_text(text)
break
continue
worker = workers[wid]
if swe_mode and shared_workdir is not None:
(
w_text,
w_in,
w_out,
is_local,
extra_cost,
n_searches,
bash_turns,
) = _swe_call_worker(
worker,
str(w_input),
cfg,
task_meta,
shared_workdir,
turn,
)
tool_calls += bash_turns
else:
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
_call_worker(worker, str(w_input), cfg)
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
n_web_searches_total += n_searches
tool_calls += n_searches
history.append(
{
"role": "worker",
"turn": turn,
"worker_id": wid,
"worker_name": worker["name"],
"worker_model": worker["model"],
"output": w_text,
"tokens_in": w_in,
"tokens_out": w_out,
"n_web_searches": n_searches,
}
)
continue
# Unknown action kind — treat as parse failure.
parse_failures += 1
if final_answer is None:
# Hard fallback: call the strongest non-search worker directly.
# "Strongest" = highest output-token price in `_prices.PRICES`,
# which tracks model capability tier closely enough for this.
# Search workers are excluded — they answer fact-lookup
# questions, not synthesis.
non_search = [
w for w in workers if w.get("type") not in _TOOLORCH_SEARCH_TYPES
] or workers
worker = max(
non_search,
key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
)
if swe_mode and shared_workdir is not None:
(ans, w_in, w_out, is_local, extra_cost, _, bash_turns) = (
_swe_call_worker(
worker,
question,
cfg,
task_meta,
shared_workdir,
max_turns + 1,
)
)
tool_calls += bash_turns
else:
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
worker, question, cfg
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
history.append(
{
"role": "worker",
"turn": max_turns + 1,
"worker_id": worker["id"],
"worker_name": worker["name"],
"worker_model": worker["model"],
"output": ans,
"tokens_in": w_in,
"tokens_out": w_out,
"fallback": True,
}
)
final_answer = ans
# In SWE mode, the authoritative output is the working-tree diff —
# frame it (the runner extracts it via the scorer's ```diff fence).
if swe_mode and shared_workdir is not None:
patch = _extract_diff(shared_workdir)
if patch.strip():
final_answer = (
f"{final_answer}\n\n```diff\n{patch}```"
if final_answer
else f"```diff\n{patch}```"
)
meta = {
"tokens_local": tokens_local,
"tokens_cloud": tokens_cloud,
"cost_usd": cost,
"turns": len([h for h in history if h["role"] == "orchestrator"]),
"web_search_uses": n_web_searches_total,
"tool_calls": int(tool_calls),
"traces": {
"history": history,
"forced_final": forced_final,
"parse_failures": parse_failures,
"workers": workers,
"n_web_searches": n_web_searches_total,
"note": (
"inference-only port; the RL-trained Nemotron-Orchestrator-8B "
"is NOT in the loop. Results are preliminary."
),
},
}
return final_answer, meta
finally:
if shared_workdir is not None:
shutil.rmtree(shared_workdir, ignore_errors=True)
# ------------------------------------------------------------------
# Paper-faithful Orchestrator-8B path.
# ------------------------------------------------------------------
def _run_rl(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
cfg = self._cfg
question = input
# Orchestrator endpoint / model (where the RL'd 8B lives).
orch_endpoint = str(
cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
)
orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
max_turns = int(cfg.get("max_turns", 8))
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
# Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
# replaces `_expert_for` and `enhance_reasoning` is post-processed
# through a Modal Python sandbox. See module docstring + paper-match
# doc at `docs/26.5.19/toolorchestra-papermatch.md`.
paper_mode = str(cfg.get("pool", "")).lower() == "paper"
# SWE-bench detection: same gate as the prompted path. Requires
# `method_cfg.swe_use_agent_loop = true` AND the task carries the
# SWE-bench fields. When active, the `enhance_reasoning` and
# `answer` workers route through `run_swe_agent_loop` on a shared
# workdir; search workers stay one-shot. At end, the working-tree
# diff is appended to final_answer so `_score_swebench` can extract
# it via the ```diff fence.
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
swe_mode = (
bool(cfg.get("swe_use_agent_loop"))
and bool(task_meta.get("problem_statement"))
and bool(task_meta.get("repo"))
and bool(task_meta.get("base_commit"))
)
shared_workdir: Optional[Path] = None
if swe_mode:
shared_workdir = Path(
tempfile.mkdtemp(
prefix=f"toolorch-rl-swe-{task_meta.get('task_id', 'x')}-"
)
)
try:
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
except Exception:
shutil.rmtree(shared_workdir, ignore_errors=True)
raise
self.record_trace_event(
{
"kind": "toolorchestra_rl_swe_workdir",
"workdir": str(shared_workdir),
"repo": task_meta["repo"],
"base_commit": task_meta["base_commit"],
}
)
# ``context_str`` mirrors the upstream's running context — accumulates
# search documents and code/exec snippets across turns. We keep this
# as a single string for prompt simplicity; the upstream uses
# tokenized cutoffs (we cap at ~24k chars instead).
context_str = ""
doc_list: List[str] = []
history: List[Dict[str, Any]] = []
tokens_local = 0
tokens_cloud = 0
cost = 0.0
n_web_searches_total = 0
tool_calls = 0
final_answer: Optional[str] = None
parse_failures = 0
# Single outer try/finally guards `shared_workdir` against any
# exception in the orchestrator loop, the post-loop fallback, or
# the diff-extraction step. Matches the prompted path's pattern.
try:
for turn in range(1, max_turns + 1):
user = (
f"Problem: {question}\n\n{context_str}\n\n"
"Choose an appropriate tool."
)
# Orchestrator-8B served on local vLLM. We pass the three NVlabs
# tools verbatim. In paper-mode we use the local helper so we
# get the SDK-level ``tool_calls`` object back — `_call_vllm`
# returns just text and loses the call when vLLM's parser
# caught it. Orchestrator-8B emits its routing decision in the
# OpenAI-native ``tool_calls`` array with an empty text body,
# so the legacy `_call_vllm` path saw nothing and silently fell
# through to the answer-1 fallback (parse_failures: 2 on every
# non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
# modes now use `_call_orchestrator_with_tool_calls` so the
# parser can read structured tool calls; the text-tag path in
# `_parse_rl_tool_call` is still the fallback when `tool_calls`
# is empty.
text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
orch_model,
orch_endpoint,
user=user,
system=RL_ORCHESTRATOR_SYS,
max_tokens=orch_max_tokens,
temperature=orch_temp,
tools=RL_TOOLS_SPEC,
)
self.record_trace_event(
{
"kind": "vllm",
"role": "orchestrator",
"model": orch_model,
"endpoint": orch_endpoint,
"system": RL_ORCHESTRATOR_SYS,
"user": user,
"response": text,
"tool_calls": [
{
"id": getattr(tc, "id", None),
"type": getattr(tc, "type", None),
"function": {
"name": getattr(
getattr(tc, "function", None), "name", None
),
"arguments": getattr(
getattr(tc, "function", None), "arguments", None
),
},
}
for tc in (sdk_tool_calls or [])
],
"tokens_in": o_in,
"tokens_out": o_out,
}
)
tokens_local += o_in + o_out
action = _parse_rl_tool_call(text, sdk_tool_calls)
history.append(
{
"role": "orchestrator",
"turn": turn,
"raw": text,
"action": action,
}
)
self.record_trace_event(
{
"kind": "toolorchestra_rl_action",
"turn": turn,
"action": action,
"raw": text,
}
)
if action is None:
parse_failures += 1
if parse_failures >= 2:
break
continue
name = action["name"]
args = action.get("arguments", {})
slot = args.get("model", "")
# Validate against the upstream tool/arg schema.
valid = (
name in RL_ALL_TOOLS
and isinstance(slot, str)
and (slot in RL_ALL_TOOLS[name]["model"])
)
if not valid:
parse_failures += 1
if parse_failures >= 2:
break
# Replay with a softer nudge in the context.
context_str += (
f"\n[Orchestrator emitted invalid tool call "
f"name={name!r} slot={slot!r} — try again.]\n"
)
continue
# Paper-match (`method_cfg.pool == "paper"`) routes through
# the Tavily/OpenRouter/Modal pool instead of the default
# Anthropic-web-search-driven mapping. For `search` this
# also forces the worker prompt to a raw query string
# (Tavily takes a single search string, not a chat-style
# framing).
if paper_mode:
worker = _paper_expert_for(
slot,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# In paper mode, `enhance_reasoning` is always the coder
# specialist regardless of the orchestrator's chosen tier.
# The coder is then expected to emit a python block which
# we exec in Modal (below).
if name == "enhance_reasoning":
worker = {
"name": f"coder:{slot}",
"type": "openrouter",
"model": _PAPER_CODER_OPENROUTER,
}
else:
worker = _expert_for(
slot,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# Dispatch — the orchestrator only conveys a tool/model
# choice, NOT a question rewrite; the prompt we send the
# expert is the same context the orchestrator saw, framed
# appropriately for the tool.
if name == "search":
if paper_mode:
# Tavily takes a query string. Orchestrator-8B often
# emits an extra `query` arg (not in the upstream
# schema but useful) — prefer it; else fall back to
# the raw question.
q = args.get("query")
w_input = q if isinstance(q, str) and q.strip() else question
else:
w_input = (
f"Search the web to gather information that helps answer:\n"
f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
)
elif name == "enhance_reasoning":
if paper_mode:
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Write a short Python script that computes intermediate "
"results which help answer the problem. Output ONLY the "
"code inside one ```python ... ``` fenced block. Print "
"any results you derive using `print(...)`. The script "
"must run with the Python stdlib only — no extra pip "
"installs."
)
else:
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Reason carefully. Outline the key intermediate steps and any "
"computations or facts you can derive. Do NOT give a final "
"answer — the orchestrator will collect your reasoning and "
"call the answer tool next."
)
else: # name == "answer"
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Provide the final answer to the user. Respect any "
"answer-format rules in the question (e.g. GAIA's "
"FINAL ANSWER: <value> convention)."
)
# SWE mode: route enhance_reasoning / answer workers through
# the SWE agent loop on the shared workdir so they can read
# files, run tests, and edit the working tree. Search workers
# stay one-shot (no agent loop). The `_swe_call_worker`
# one-shot fallbacks (openai-typed workers, search) return
# bash_turns=0; vllm/anthropic-typed workers run the loop.
bash_turns = 0
if swe_mode and shared_workdir is not None and name != "search":
(
w_text,
w_in,
w_out,
is_local,
extra_cost,
n_searches,
bash_turns,
) = _swe_call_worker(
worker,
w_input,
cfg,
task_meta,
shared_workdir,
turn,
)
else:
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
_call_worker(worker, w_input, cfg)
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
n_web_searches_total += n_searches
# SWE bash turns count as tool calls (each one is a $BASH block
# the agent executed). On non-SWE turns fall back to the
# original "at least one expert call" accounting.
tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches)
# Paper-match: pipe coder output through a Modal sandbox so
# `enhance_reasoning` actually executes the code the coder
# wrote. Append the exec output to the worker's text. No-op
# when no python block is found.
modal_exec_output: Optional[str] = None
modal_exec_rc: Optional[int] = None
if paper_mode and name == "enhance_reasoning" and not swe_mode:
code = _extract_first_python_block(w_text)
if code:
timeout_s = int(cfg.get("modal_python_timeout_s", 60))
modal_exec_output, modal_exec_rc = _call_modal_python(
code,
timeout_s=timeout_s,
)
tool_calls += 1
w_text = (
f"{w_text}\n\n[modal-python stdout/stderr "
f"(rc={modal_exec_rc})]\n{modal_exec_output}"
)
history.append(
{
"role": "worker",
"turn": turn,
"tool": name,
"slot": slot,
"worker_model": worker["model"],
"worker_type": worker["type"],
"output": w_text,
"tokens_in": w_in,
"tokens_out": w_out,
"n_web_searches": n_searches,
"bash_turns": bash_turns,
"modal_exec_rc": modal_exec_rc,
}
)
# Update accumulated context for the next turn.
if name == "search":
# Treat the search worker's response as a document.
doc_list.append(w_text)
ctx_docs = "\n\n".join(
f"Doc {i + 1}: {d}" for i, d in enumerate(doc_list)
)
# Crude char-level cap mirrors the upstream's ~24k token cap.
context_str = ("Documents:\n" + ctx_docs)[-24000:]
elif name == "enhance_reasoning":
snippet = f"\n\nReasoning/exec output:\n{w_text}"
context_str = (context_str + snippet)[-24000:]
else: # answer
final_answer = w_text.strip()
break
if final_answer is None:
# Hard fallback: ask the frontier worker directly. In SWE
# mode route this final call through the agent loop too so
# it can still touch the workdir and emit a diff.
expert_fn = _paper_expert_for if paper_mode else _expert_for
worker = expert_fn(
"answer-1",
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
fb_bash_turns = 0
if swe_mode and shared_workdir is not None:
(ans, w_in, w_out, is_local, extra_cost, _, fb_bash_turns) = (
_swe_call_worker(
worker,
question,
cfg,
task_meta,
shared_workdir,
max_turns + 1,
)
)
tool_calls += fb_bash_turns
else:
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
worker, question, cfg
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
history.append(
{
"role": "worker",
"turn": max_turns + 1,
"tool": "answer",
"slot": "answer-1",
"worker_model": worker["model"],
"worker_type": worker["type"],
"output": ans,
"tokens_in": w_in,
"tokens_out": w_out,
"bash_turns": fb_bash_turns,
"fallback": True,
}
)
final_answer = ans
# In SWE mode, the authoritative output is the working-tree diff —
# frame it so `_score_swebench`'s extract_patch picks it up.
if swe_mode and shared_workdir is not None:
patch = _extract_diff(shared_workdir)
if patch.strip():
final_answer = (
f"{final_answer}\n\n```diff\n{patch}```"
if final_answer
else f"```diff\n{patch}```"
)
meta = {
"tokens_local": tokens_local,
"tokens_cloud": tokens_cloud,
"cost_usd": cost,
"turns": len([h for h in history if h["role"] == "orchestrator"]),
"web_search_uses": n_web_searches_total,
"tool_calls": int(tool_calls),
"traces": {
"history": history,
"parse_failures": parse_failures,
"orchestrator_model": orch_model,
"orchestrator_endpoint": orch_endpoint,
"mode": "rl",
"pool": "paper" if paper_mode else "default",
"swe_mode": swe_mode,
"note": (
"RL-trained nvidia/Orchestrator-8B as orchestrator. "
"Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto "
"our hybrid worker types — see toolorchestra.py docstring."
),
},
}
return final_answer, meta
finally:
if shared_workdir is not None:
shutil.rmtree(shared_workdir, ignore_errors=True)
__all__ = ["ToolOrchestraAgent"]
@@ -0,0 +1,51 @@
"""Orchestrator vLLM tool-call client for ToolOrchestraAgent."""
from __future__ import annotations
from typing import Any, Dict, List, Tuple
def _call_orchestrator_with_tool_calls(
model: str,
endpoint: str,
*,
user: str,
system: str,
max_tokens: int,
temperature: float,
tools: List[Dict[str, Any]],
timeout: float = 600.0,
) -> Tuple[str, int, int, Any]:
"""Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls).
Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level
``tool_calls`` object so the RL-mode parser can match against it
directly. Otherwise vLLM's tool parser silently swallows the tool call
into the SDK field while leaving ``content == ''`` and the text-tag
parser sees nothing, falling through to the answer-1 fallback. (Bug
observed 2026-05-19 on the paper-match smoke; same path was buggy on
the default pool too, just less reproducibly.)
"""
from openai import OpenAI
client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
messages = [
{"role": "system", "content": system},
{"role": "user", "content": user},
]
resp = client.chat.completions.create(
model=model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
tools=tools,
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
choice = resp.choices[0]
message = choice.message
text = message.content or ""
tool_calls = getattr(message, "tool_calls", None)
u = resp.usage
p = getattr(u, "prompt_tokens", 0) if u else 0
c = getattr(u, "completion_tokens", 0) if u else 0
return text, p, c, tool_calls
@@ -0,0 +1,187 @@
"""Slot -> worker expert mapping for ToolOrchestraAgent."""
from __future__ import annotations
from typing import Any, Dict, Optional
# Default model used when an `anthropic-web-search` entry omits `model`.
_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
# tiers since we don't have Qwen-Math served.
#
# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
# so the substitution is deferred until we know the cell's resolved local/cloud
# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
def _expert_for(
slot: str,
local_model: Optional[str],
local_endpoint: Optional[str],
cloud_model: str,
cloud_endpoint: str = "anthropic",
) -> Dict[str, Any]:
"""Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
Routing policy:
- `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off
`cloud_endpoint` ("anthropic"/"openai"/"gemini")
- `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's
cost tier for mid OpenAI calls)
- `*-3` (local tier) -> local vLLM (`local_model`)
- `answer-math-*` -> same tiers as the numeric suffix
- `search-*` -> provider-native web search when the cloud
endpoint supports it; otherwise Anthropic
"""
if slot.startswith("search"):
ep = (cloud_endpoint or "anthropic").lower()
if ep == "openai":
return {
"name": f"search:{slot}",
"type": "openai-web-search",
"model": cloud_model,
}
if ep == "gemini":
return {
"name": f"search:{slot}",
"type": "gemini-web-search",
"model": cloud_model,
}
return {
"name": f"search:{slot}",
"type": "anthropic-web-search",
"model": _DEFAULT_WEB_SEARCH_MODEL,
}
if slot.endswith("-1") or slot.endswith("-math-1"):
ep = (cloud_endpoint or "anthropic").lower()
if ep not in ("anthropic", "openai", "gemini"):
ep = "anthropic"
return {
"name": f"frontier:{slot}",
"type": ep,
"model": cloud_model,
}
if slot.endswith("-2") or slot.endswith("-math-2"):
return {
"name": f"mid:{slot}",
"type": "openai",
"model": "gpt-5-mini",
}
# `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
# we substitute whatever local model the cell wired up).
if local_model and local_endpoint:
return {
"name": f"local:{slot}",
"type": "vllm",
"model": local_model,
"base_url": local_endpoint,
}
# Fallback if no local — gpt-5-mini.
return {
"name": f"mid-fallback:{slot}",
"type": "openai",
"model": "gpt-5-mini",
}
# ============================================================================
# Paper-match expert mapping (2026-05-19).
# ============================================================================
# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
# models and routes `search` through Tavily, while `enhance_reasoning` is
# expected to produce code that the caller pipes through a Modal sandbox
# (handled at dispatch time, not here).
#
# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
#
# reasoner-1 -> GPT-5 (frontier reasoner)
# reasoner-2 -> GPT-5-mini (mid)
# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
# answer-1 -> GPT-5
# answer-2 -> GPT-5-mini
# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
# answer-4 -> local Qwen
# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
# search-* -> Tavily search (paper)
#
# `enhance_reasoning` is dispatched through the coder specialist regardless of
# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
# routes the same way in all three cases, then optionally extracts a python
# code block and execs it in Modal. (We keep the slot-aware routing inside the
# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
# pins the coder regardless. See `_run_rl_paper` dispatch.)
_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
def _paper_expert_for(
slot: str,
local_model: Optional[str],
local_endpoint: Optional[str],
cloud_model: str,
cloud_endpoint: str = "openai",
) -> Dict[str, Any]:
"""Paper-match counterpart of ``_expert_for``.
Differs from ``_expert_for``:
- Search slots go to ``tavily-search`` (not Anthropic web_search).
- Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
OpenRouter rather than collapsing onto the local vLLM.
- Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
- ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
"""
if slot.startswith("search"):
return {
"name": f"tavily:{slot}",
"type": "tavily-search",
"model": "tavily",
}
if slot in ("answer-math-1", "answer-math-2"):
return {
"name": f"math-coder:{slot}",
"type": "openrouter",
"model": _PAPER_CODER_OPENROUTER,
}
if slot == "answer-3":
return {
"name": f"generalist-llama:{slot}",
"type": "openrouter",
"model": _PAPER_GENERALIST_TIER3_OPENROUTER,
}
if slot.endswith("-1"):
# Tier-1 frontier reasoner / answer — paper uses GPT-5.
return {
"name": f"frontier:{slot}",
"type": "openai",
"model": "gpt-5",
}
if slot.endswith("-2"):
return {
"name": f"mid:{slot}",
"type": "openai",
"model": "gpt-5-mini",
}
# `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
# also serves the local Qwen for the rare local-tier slot).
if local_model and local_endpoint:
return {
"name": f"local:{slot}",
"type": "vllm",
"model": local_model,
"base_url": local_endpoint,
}
return {
"name": f"mid-fallback:{slot}",
"type": "openai",
"model": "gpt-5-mini",
}
@@ -0,0 +1,161 @@
"""Action / tool-call parsing + prompt assembly for ToolOrchestraAgent."""
from __future__ import annotations
import json
import re
from typing import Any, Dict, List, Optional
# Regex for ``<tool_call>{...}</tool_call>`` blocks emitted by Orchestrator-8B
# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a
# hermes-style template). Captures the JSON payload.
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
# Fallback for a known failure mode: the SFT'd orchestrator often emits its
# delegation as ``\boxed{expert_ab12: <sub-question>}`` (math-data habit bleeding
# into routing) instead of a real ``<tool_call>`` JSON block. The name before the
# colon is a valid (anonymized) tool label; the rest is the sub-question/args.
#
# Real outputs use several shapes, all handled here:
# \boxed{web_search: <query>} -> name, args
# \boxed{web_search query: <query>} -> the word "query" is dropped
# \boxed{file_read, path: <path>} -> the ", <key>" hint sets the arg key
# We require a ``name <optional , key | query> : <args>`` structure, so a genuine
# answer like ``\boxed{10}``, ``\boxed{b,e}`` or ``\boxed{The answer is: 42}``
# (first token isn't followed by a bare ``:`` / ``, key:`` / ``query:``) is left
# alone. ``group("key")`` is the explicit arg key when the model wrote ``, key:``.
_BOXED_DELEGATION_RE = re.compile(
r"\\boxed\{\s*([A-Za-z_][\w-]*)\s*"
r"(?:,\s*(?P<key>\w+)\s*)?" # optional ", key" hint (e.g. file_read, path:)
r"(?:query\s*)?" # optional literal "query" word before the colon
r":\s*(.+?)\s*\}\s*$",
re.DOTALL,
)
def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
"""Return ``{"name": str, "arguments": dict}`` or None.
Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
back to scraping ``<tool_call>{...}</tool_call>`` tags from the raw
content. We take the first tool call only Orchestrator-8B was trained
to emit exactly one per turn.
"""
# SDK-level path.
if sdk_tool_calls:
first = sdk_tool_calls[0]
name = getattr(getattr(first, "function", None), "name", None)
args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
try:
args = json.loads(args_raw)
except json.JSONDecodeError:
args = {}
if isinstance(name, str) and isinstance(args, dict):
return {"name": name, "arguments": args}
# Text-tag fallback.
if not isinstance(content, str):
return None
m = _TOOL_CALL_TAG_RE.search(content)
if m:
try:
obj = json.loads(m.group(1))
except json.JSONDecodeError:
obj = None
if isinstance(obj, dict):
name = obj.get("name")
args = obj.get("arguments", {})
if isinstance(name, str) and isinstance(args, dict):
return {"name": name, "arguments": args}
# \boxed{name[, key][ query]: <args>} fallback (see _BOXED_DELEGATION_RE).
bm = _BOXED_DELEGATION_RE.search(content.strip())
if bm:
arg_val = bm.groups()[-1].strip()
key = bm.group("key") or "input"
return {"name": bm.group(1), "arguments": {key: arg_val}}
return None
def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
return "\n".join(
f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
)
def _build_user_prompt(
question: str,
workers: List[Dict[str, Any]],
history: List[Dict[str, Any]],
) -> str:
pieces = [
f"Worker pool:\n{_build_pool_block(workers)}",
f"User question:\n{question}",
]
if history:
pieces.append("Conversation so far (orchestrator turns and worker outputs):")
for h in history:
if h["role"] == "orchestrator":
pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
else:
pieces.append(
f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
f"{h['output']}"
)
pieces.append(
"Emit the next JSON action object now — exactly one object, no prose."
)
return "\n\n".join(pieces)
def _strip_fences(s: str) -> str:
s = s.strip()
if s.startswith("```"):
first_nl = s.find("\n")
if first_nl != -1:
s = s[first_nl + 1 :]
if s.endswith("```"):
s = s[:-3]
s = s.strip()
return s
def _parse_action(text: str) -> Optional[Dict[str, Any]]:
s = _strip_fences(text)
# First try direct parse, then balanced-brace extraction.
try:
obj = json.loads(s)
if isinstance(obj, dict) and "action" in obj:
return obj
except json.JSONDecodeError:
pass
start = s.find("{")
if start == -1:
return None
depth = 0
for i in range(start, len(s)):
c = s[i]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
if depth == 0:
try:
obj = json.loads(s[start : i + 1])
if isinstance(obj, dict) and "action" in obj:
return obj
except json.JSONDecodeError:
return None
return None
def _extract_final_answer_text(text: str) -> str:
"""Best-effort: pull the answer string from a malformed action emission.
Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
"""
m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
if m:
return m.group(1).encode("utf-8").decode("unicode_escape")
m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
if m:
return m.group(1).strip()
return text.strip()
@@ -0,0 +1,110 @@
"""Prompt strings + tool specs for ToolOrchestraAgent (split from toolorchestra.py)."""
from __future__ import annotations
from typing import Any, Dict, List
ORCHESTRATOR_SYS = """\
You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
{"action": "call_worker", "worker_id": <int>, "input": "<question or instruction for that worker>"}
{"action": "final_answer", "answer": "<final answer to the user, respecting the question's answer-format rules>"}
Strategy:
- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math).
- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass.
- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase.
- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question.
"""
FORCE_FINAL_PROMPT = (
"Worker-call budget exhausted. Emit `final_answer` now using everything "
"you've learned. Respect the question's answer-format rules."
)
# ============================================================================
# RL-mode constants (Orchestrator-8B, paper-faithful).
# ============================================================================
#
# Verbatim copies of the upstream system prompt / user-prompt template / tools
# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
# edit the description text — Orchestrator-8B was RL-trained against this
# exact wording and pricing/latency table.
RL_ORCHESTRATOR_SYS = "You are good at using tools."
RL_TOOLS_SPEC: List[Dict[str, Any]] = [
{
"type": "function",
"function": {
"name": "enhance_reasoning",
"description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
"parameters": {
"properties": {
"model": {
"description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
{
"type": "function",
"function": {
"name": "answer",
"description": "give the final answer. Not allowed to call if documents is empty.",
"parameters": {
"properties": {
"model": {
"description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
{
"type": "function",
"function": {
"name": "search",
"description": "Search for missing information",
"parameters": {
"properties": {
"model": {
"description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
]
# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
"enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
"answer": {
"model": [
"answer-1",
"answer-2",
"answer-3",
"answer-4",
"answer-math-1",
"answer-math-2",
],
},
"search": {"model": ["search-1", "search-2", "search-3"]},
}
@@ -0,0 +1,417 @@
"""Faithful unified-tool rollout loop for ToolOrchestra (arXiv:2511.21689 §2.2).
One reasoning->action->observation loop where the orchestrator picks **a named
tool** (one per model, from :mod:`expert_registry`) each turn, the environment
executes it, and the observation is appended to a running context. The rollout
ends when the orchestrator emits a turn with **no tool call** (its text is the
final answer) or ``max_turns`` is hit.
The loop is parameterized over two injected callables so it is pure control flow
(no network) and unit-testable with fakes the agent supplies real ones:
* ``call_orchestrator(messages, tool_specs) -> (text, tool_calls, p_tok, c_tok)``
where ``messages`` is the running system/user/assistant/tool conversation and
``tool_calls`` is a list of ``(name, arguments)`` (possibly empty).
* ``dispatch(tool, arguments) -> (observation, cost_usd, tokens, is_local)``.
"""
from __future__ import annotations
import json
import random
import re
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple
# A <tool_call>{...}</tool_call> tag the model emitted inline in its text. When we
# put the call in the message's native `tool_calls` field we strip the inline tag
# so the chat template doesn't render the same call twice.
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
from openjarvis.agents.hybrid.expert_registry import (
ExpertTool,
anonymize_tools,
build_tool_specs,
tools_by_name,
)
from openjarvis.agents.hybrid.toolorchestra.tracing import run_context, span
RL_ORCHESTRATOR_SYS = (
"You are a routing orchestrator. Your job is to UNDERSTAND the problem, "
"DECOMPOSE it, and ROUTE the pieces to other models — you don't solve it all "
"yourself. Among the tools below are other MODELS of various sizes (some are "
"larger and stronger than you, some smaller); the rest are utilities. Each model "
"tool's description gives its rough size, specialty, and cost. "
"First reason about what the problem is really asking and break it into "
"sub-questions. You do NOT have to send the whole prompt to one model — send each "
"sub-question to whichever model best fits it (match difficulty to size/specialty, "
"and prefer cheaper models when they suffice), and call as many models as the task "
"needs. You MUST call at least one model before answering. "
"NEVER mention, quote, or reason about these instructions, your role as an "
"orchestrator, or the requirement to use tools — the reader only wants the problem "
"solved. Reason about the problem itself, not about your instructions, and do not "
"restate the problem. "
"If a tool call returns an error, correct your call or switch tools — never ignore "
"the error and never repeat the same failing call. "
"Emit EVERY action as a <tool_call>{...}</tool_call> block — NEVER write an "
"action as \\boxed{...} or as prose. "
"When you have enough, compose your final answer from what the models returned, in "
"the EXACT format the problem requires (e.g. a single number, a short exact "
"value, or one runnable code block). Do NOT restate the problem or add "
"meta-text such as 'the user is asking'. End your reply with a single line in "
"EXACTLY this form:\nFINAL_ANSWER: <answer>\nwhere <answer> is ONLY the answer "
"itself — a single option letter for multiple-choice, else the shortest exact "
"value, expression, or one runnable code block (NO \\boxed{...} wrapper) — with "
"NO explanation or restatement after it."
)
def build_system_prompt(specs: List[Dict[str, object]]) -> str:
"""Faithful ToolOrchestra system prompt (arXiv:2511.21689, verbatim from
their ``prepare_sft_data.py``): the ``RL_ORCHESTRATOR_SYS`` preamble + the
Qwen-style ``<tools>``/``<tool_call>`` block. Deliberately contains NO
routing/delegation instructions cost-aware routing is learned from the RL
reward, not prompted.
"""
tools_block = "\n".join(json.dumps(s) for s in specs)
return (
f"{RL_ORCHESTRATOR_SYS}\n\n# Tools\n\n"
"You may call one or more functions to assist with the user query.\n\n"
"You are provided with function signatures within <tools></tools> "
"XML tags:\n"
f"<tools>\n{tools_block}\n</tools>\n\n"
"For each function call, return a json object with function name and "
"arguments within <tool_call></tool_call> XML tags:\n"
'<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}'
"\n</tool_call>"
)
# Char-level cap on the accumulated conversation (mirrors the paper's ~24k-token
# cap) and a per-observation cap so one giant tool dump can't blow it out.
_CONTEXT_CAP = 24000
_OBS_CAP = 8000
# After this many tool calls, push the orchestrator to stop and answer — bounds
# the "keep re-asking the same model forever" loop (esp. gemma).
_SOFT_CALL_CAP = 6
def _trim_history(messages: List[Dict[str, str]]) -> None:
"""Keep the message history under ``_CONTEXT_CAP`` chars by dropping the
OLDEST assistant/tool exchange, never the system prompt or the problem."""
def total() -> int:
return sum(len(m.get("content") or "") for m in messages)
# messages[0]=system, messages[1]=user problem — always keep those two.
while total() > _CONTEXT_CAP and len(messages) > 4:
del messages[2:4]
@dataclass
class UnifiedTurn:
"""One orchestrator turn. ``tool_name is None`` marks the final-answer turn."""
reasoning: str
tool_name: Optional[str] = None
arguments: Dict[str, object] = field(default_factory=dict)
observation: Optional[str] = None
@dataclass
class UnifiedRollout:
turns: List[UnifiedTurn]
final_answer: str
cost_usd: float = 0.0
tokens: int = 0
num_tool_calls: int = 0
parse_failures: int = 0
# When experts were anonymized for this rollout, maps the opaque label the
# policy saw (e.g. ``expert_a3f9``) back to the real tool name (``gpt_5_5``).
anon_map: Optional[Dict[str, str]] = None
# The EXACT tool specs the policy saw this rollout (anonymized when
# ``anonymize=True``). Serialization MUST build the saved system prompt from
# these — not from the real registry — or the prompt's ``<tools>`` block ends
# up with real model names+pricing while the assistant ``<tool_call>`` tags
# use the anon labels, which both breaks SFT (calls a tool absent from its
# list) and re-injects the name/cost bias anonymization removed.
tool_specs: Optional[List[Dict[str, object]]] = None
def tool_calls(self) -> List[Tuple[str, Dict[str, object]]]:
return [(t.tool_name, t.arguments) for t in self.turns if t.tool_name]
def _tool_prompt(tool: ExpertTool, arguments: Dict[str, object], question: str) -> str:
"""The text we actually send the dispatched tool, framed by its arg schema."""
for key in ("input", "query", "code"):
val = arguments.get(key)
if isinstance(val, str) and val.strip():
return val
return question
def _run_unified_rollout_inner(
question: str,
tools: List[ExpertTool],
*,
call_orchestrator: Callable[
..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
],
dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
max_turns: int = 50,
system: str = RL_ORCHESTRATOR_SYS,
anonymize: bool = False,
) -> UnifiedRollout:
"""Drive the faithful unified-tool rollout for one task.
``anonymize``: replace each model expert with an opaque random label, a
uniform description and no cost line, shuffled so the policy can't route on
a model's name/position/cost (which we found dominate the choice). The
anon->real mapping is returned on the rollout for offline analysis.
"""
anon_map: Optional[Dict[str, str]] = None
if anonymize:
tools, anon_map = anonymize_tools(tools, random.Random())
specs = build_tool_specs(tools)
by_name = tools_by_name(tools)
# Proper multi-turn conversation (NOT a flattened user blob): the orchestrator
# sees its own calls as `assistant` turns and each observation as a `tool`
# turn, so it can tell a tool result from user input and doesn't re-derive the
# whole problem every turn. This mirrors the serialized SFT format exactly.
messages: List[Dict[str, str]] = [
{"role": "system", "content": system},
{"role": "user", "content": f"Problem: {question}"},
]
turns: List[UnifiedTurn] = []
cost = 0.0
tokens = 0
n_tool_calls = 0
parse_failures = 0
nudges = 0
empty_input_nudges = 0
final_answer = ""
for _ in range(max_turns):
text, tool_calls, p_tok, c_tok = call_orchestrator(messages, specs)
tokens += int(p_tok) + int(c_tok)
if not tool_calls:
# Enforce the "MUST delegate to a model" rule: if the orchestrator
# tries to answer before ANY tool call, nudge it to route instead of
# accepting a solve-it-yourself answer (which reject-sampling drops).
if n_tool_calls == 0 and nudges < 2:
nudges += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": (
"Make progress by delegating a concrete sub-question to one of the "
"models now via a <tool_call>."
),
}
)
continue
# No tool call -> the orchestrator is answering. Terminate.
final_answer = (text or "").strip()
turns.append(UnifiedTurn(reasoning=text or "", tool_name=None))
messages.append({"role": "assistant", "content": text or ""})
break
name, arguments = tool_calls[0]
if name not in by_name:
parse_failures += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": f"[invalid tool {name!r} — choose one from the provided tool list]",
}
)
if parse_failures >= 2:
final_answer = (text or "").strip()
break
continue
tool = by_name[name]
# Empty-input guard: the small orchestrator often emits a <tool_call> with
# a blank/missing input on harder multi-turn tasks. Dispatching it returns
# a "no input provided" error observation that poisons the whole (often
# otherwise-correct) trajectory — the dominant clean-yield killer on hard
# tasks. Instead, nudge the model to resend WITH input and drop the
# malformed turn (don't record it), so it never enters the training data.
if tool.kind == "model" or tool.backend_type != "openjarvis-tool":
_has_input = any(
isinstance(arguments.get(k), str) and arguments.get(k).strip()
for k in ("input", "query", "code")
)
if not _has_input and empty_input_nudges < 3:
empty_input_nudges += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": (
f"Your call to {name} had an empty 'input'. Resend the "
"<tool_call> with a non-empty 'input' field containing the "
"concrete sub-question to delegate."
),
}
)
continue
obs, dcost, dtok, _is_local = dispatch(tool, arguments)
cost += float(dcost)
tokens += int(dtok)
n_tool_calls += 1
turns.append(
UnifiedTurn(
reasoning=text or "",
tool_name=name,
arguments=dict(arguments),
observation=obs,
)
)
# The model's own action as an assistant turn, then the observation as a
# distinct `tool` turn it reads as a tool response.
#
# Use the NATIVE OpenAI tool-call protocol (assistant carries `tool_calls`
# with an id; the tool message references it via `tool_call_id`) rather
# than embedding the <tool_call> tag as plain text. vLLM tolerates the
# text form, but strict OpenAI-compatible providers (Anthropic) reject a
# `tool` message with no `tool_call_id` — 400 "tool_call_id: Field
# required" — which made cloud teachers unusable as the orchestrator.
# The SFT target is unaffected: it's rendered from `roll.turns` by the
# serializer, not from this message list.
call_id = f"call_{n_tool_calls}"
# Strip the tag if the model emitted it inline, so the call isn't rendered
# twice (once from content, once from the template's tool_calls block).
call_content = _TOOL_CALL_TAG_RE.sub("", text or "").strip()
messages.append(
{
"role": "assistant",
"content": call_content,
"tool_calls": [
{
"id": call_id,
"type": "function",
"function": {
"name": name,
"arguments": json.dumps(arguments, ensure_ascii=False),
},
}
],
}
)
obs_text = obs or ""
if len(obs_text) > _OBS_CAP:
obs_text = obs_text[:_OBS_CAP] + "\n…[truncated]"
messages.append(
{
"role": "tool",
"tool_call_id": call_id,
"name": name,
"content": obs_text,
}
)
if n_tool_calls >= _SOFT_CALL_CAP:
messages.append(
{
"role": "user",
"content": (
"You now have enough information from the models. Do NOT call any "
"more tools — reply with your FINAL_ANSWER line only."
),
}
)
_trim_history(messages)
else:
# Hit max_turns with no explicit answer: use the last observation/text.
final_answer = (
(turns[-1].observation or turns[-1].reasoning).strip() if turns else ""
)
return UnifiedRollout(
turns=turns,
final_answer=final_answer,
cost_usd=cost,
tokens=tokens,
num_tool_calls=n_tool_calls,
parse_failures=parse_failures,
anon_map=anon_map,
tool_specs=specs,
)
def run_unified_rollout(
question: str,
tools: List[ExpertTool],
*,
call_orchestrator: Callable[
..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
],
dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
max_turns: int = 50,
system: str = RL_ORCHESTRATOR_SYS,
anonymize: bool = False,
) -> UnifiedRollout:
"""One ToolOrchestra rollout = one Braintrust trace. Opens the root span,
runs the loop (orchestrator turns + expert/tool dispatches nest inside as
child spans), then logs the final answer + cost/tokens. No-op when tracing
is disabled (see :mod:`.tracing`)."""
_run_meta, _run_tags = run_context()
_fields = dict(
input={"question": question},
metadata={
"max_turns": max_turns,
"anonymize": anonymize,
"n_tools": len(tools),
**_run_meta,
},
)
if _run_tags:
_fields["tags"] = _run_tags
with span("toolorchestra.rollout", span_type="task", **_fields) as _s:
roll = _run_unified_rollout_inner(
question,
tools,
call_orchestrator=call_orchestrator,
dispatch=dispatch,
max_turns=max_turns,
system=system,
anonymize=anonymize,
)
_s.log(
output=roll.final_answer,
metrics={
"cost_usd": roll.cost_usd,
"tokens": roll.tokens,
"num_tool_calls": roll.num_tool_calls,
"parse_failures": roll.parse_failures,
},
metadata={
"n_experts_available": len(roll.anon_map or {}),
"answered": bool(roll.final_answer),
# label -> real model, so every anonymized route span in this
# trace can be decoded back to the model that actually ran.
"anon_map": roll.anon_map or {},
},
)
return roll
def tool_call_tag(name: str, arguments: Dict[str, object]) -> str:
"""Render a tool call as the ``<tool_call>{...}</tool_call>`` text the model emits."""
return (
f"<tool_call>{json.dumps({'name': name, 'arguments': arguments})}</tool_call>"
)
__all__ = [
"RL_ORCHESTRATOR_SYS",
"UnifiedRollout",
"UnifiedTurn",
"build_system_prompt",
"run_unified_rollout",
"tool_call_tag",
]
@@ -0,0 +1,82 @@
"""Tavily search + Modal Python sandbox helpers for ToolOrchestraAgent."""
from __future__ import annotations
import re
from typing import Optional, Tuple
# ---- Tavily + Modal helpers -------------------------------------------------
def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]:
"""One-shot Tavily search. Returns (text, p_tok=0, c_tok=0).
Token counts are reported as zero (no LLM was billed); the OpenJarvis
accounting layer separately tallies tool-call counts. Falls back to
DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``).
"""
from openjarvis.tools.web_search import WebSearchTool
tool = WebSearchTool(max_results=max_results)
res = tool.execute(query=query, max_results=max_results)
text = res.content or ""
if not res.success and not text:
text = "(no results)"
return text, 0, 0
_MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox"
def _call_modal_python(code: str, timeout_s: int = 60) -> Tuple[str, int]:
"""Execute a single Python snippet in a fresh Modal Sandbox.
Returns ``(combined_stdout_stderr, returncode)``. Logs are capped at 8 KiB.
Any exception (modal auth, network, sandbox boot failure) is captured into
the returned string with a non-zero rc we never raise back to the
orchestrator loop. The sandbox is torn down at the end via ``terminate()``.
"""
try:
import modal
app = modal.App.lookup(_MODAL_APP_NAME, create_if_missing=True)
# python:3.12-slim is small + boots fast; the paper uses a generic
# Python image too. We rely on stdlib only — no extra pip installs.
image = modal.Image.debian_slim(python_version="3.12")
sb = modal.Sandbox.create(
"python",
"-c",
code,
app=app,
image=image,
timeout=int(timeout_s),
)
sb.wait()
try:
out = sb.stdout.read() or ""
except Exception:
out = ""
try:
err = sb.stderr.read() or ""
except Exception:
err = ""
rc = sb.returncode if sb.returncode is not None else -1
try:
sb.terminate()
except Exception:
pass
combined = out + (("\n" + err) if err else "")
if len(combined) > 8192:
combined = combined[:8192] + "\n... (output truncated)"
return combined, int(rc)
except Exception as exc:
return f"[modal-python error: {type(exc).__name__}: {exc}]", -1
_PY_CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL)
def _extract_first_python_block(text: str) -> Optional[str]:
"""Return the first ```python ... ``` block (or ```...```), or None."""
m = _PY_CODE_RE.search(text or "")
return m.group(1).strip() if m else None
@@ -0,0 +1,170 @@
"""Braintrust telemetry for the ToolOrchestra rollout.
Each ``run_unified_rollout`` becomes ONE Braintrust trace (root span); the
orchestrator turns and every expert/tool dispatch nest inside it, and the
underlying OpenAI/Anthropic calls (wrapped clients) nest one level deeper so
you see the full routing tree with inputs/outputs/tokens/cost per node.
On by default. Degrades to a total no-op (never raises, never changes behavior)
when: ``OJ_BRAINTRUST=0``, the ``braintrust`` package isn't installed, or
``BRAINTRUST_API_KEY`` is unset. Concurrency: rollouts run one-per-thread
(ThreadPoolExecutor); threads don't inherit contextvars, so each rollout opens
its own independent root trace and same-thread child calls nest correctly.
"""
from __future__ import annotations
import contextlib
import logging
import os
from typing import Any
logger = logging.getLogger(__name__)
_STATE: dict = {"resolved": False, "enabled": False, "bt": None}
def _truthy(v: str) -> bool:
return v.strip().lower() not in ("0", "false", "no", "off", "")
def _resolve() -> bool:
"""Lazily decide whether tracing is active and init the logger once."""
if _STATE["resolved"]:
return _STATE["enabled"]
_STATE["resolved"] = True
if not _truthy(os.getenv("OJ_BRAINTRUST", "1")): # on by default
return False
if not os.getenv("BRAINTRUST_API_KEY"):
logger.info(
"braintrust on-by-default but BRAINTRUST_API_KEY unset — tracing disabled"
)
return False
try:
import braintrust as _bt
proj_id = os.getenv("OJ_BRAINTRUST_PROJECT_ID")
if proj_id:
_bt.init_logger(project_id=proj_id)
else:
_bt.init_logger(project=os.getenv("OJ_BRAINTRUST_PROJECT", "toolorchestra"))
_STATE["bt"] = _bt
_STATE["enabled"] = True
logger.info(
"braintrust tracing ENABLED (%s)",
f"project_id={proj_id}"
if proj_id
else f"project={os.getenv('OJ_BRAINTRUST_PROJECT', 'toolorchestra')}",
)
except Exception as exc: # missing pkg / bad key / init failure — never crash
logger.warning("braintrust init failed (%s) — tracing disabled", exc)
return _STATE["enabled"]
def enabled() -> bool:
return _resolve()
def run_context() -> tuple[dict, list]:
"""Run-level metadata + tags for the ROOT rollout span, sourced from env
(set by the generation driver). No-op-safe: returns ({}, []) if nothing is
set, and never raises. Env keys:
- ``OJ_RUN_LABEL`` human run label (also stamped on the uploaded dataset).
- ``OJ_GEN_MODEL`` SPECIFIC gen model id (e.g. ``Qwen/Qwen3.5-9B``).
- ``OJ_RUN_STAGE`` ``prod``|``smoke`` (inferred from the label if unset).
- ``OJ_CFG_*`` config knobs (temperature/max_turns/anonymize/...).
"""
import datetime
meta: dict = {}
tags: list = []
try:
label = os.getenv("OJ_RUN_LABEL")
gen_model = os.getenv("OJ_GEN_MODEL")
stage = os.getenv("OJ_RUN_STAGE")
if not stage and label:
stage = "smoke" if "smoke" in label.lower() else "prod"
if label:
meta["run_label"] = label
if gen_model:
meta["gen_model"] = gen_model
if stage:
meta["stage"] = stage
cfg = {}
for env_key, key in (
("OJ_CFG_TEMPERATURE", "temperature"),
("OJ_CFG_MAX_TURNS", "max_turns"),
("OJ_CFG_ANONYMIZE", "anonymize"),
("OJ_CFG_REJECTION_ONLY", "rejection_only"),
):
v = os.getenv(env_key)
if v not in (None, ""):
cfg[key] = v
if cfg:
meta["config"] = cfg
date = datetime.date.today().isoformat()
tags = [t for t in (gen_model, stage, date) if t]
except Exception: # never let telemetry enrichment break a rollout
return {}, []
return meta, tags
def wrap_client(client: Any) -> Any:
"""Wrap an OpenAI/Anthropic client so its calls auto-log under the current
span. Pass-through (unchanged client) when tracing is off or on any error
so this is always safe to call at client construction."""
if not _resolve():
return client
try:
bt = _STATE["bt"]
mod = type(client).__module__.lower()
if "anthropic" in mod:
return bt.wrap_anthropic(client)
return bt.wrap_openai(client)
except Exception as exc:
logger.warning("braintrust wrap_client failed (%s) — using raw client", exc)
return client
class _NullSpan:
"""No-op span used when tracing is disabled (keeps call sites branch-free)."""
def log(self, **_kw: Any) -> None:
pass
def __enter__(self) -> "_NullSpan":
return self
def __exit__(self, *_a: Any) -> bool:
return False
@contextlib.contextmanager
def span(name: str, *, span_type: str = "task", **fields: Any):
"""Open a Braintrust span (or a no-op). Use as::
with span("toolorchestra.rollout", input=...) as s:
...
s.log(output=..., metrics=..., metadata=...)
``fields`` (input/metadata/...) are forwarded to ``start_span``.
"""
if not _resolve():
yield _NullSpan()
return
# Guard only span CREATION (so a Braintrust hiccup degrades to a no-op). Do
# NOT wrap the yielded body in try/except: an exception from the rollout would
# be thrown into this generator, caught here, and masked by a second yield
# ("generator didn't stop after throw()"). Let the body's own exceptions
# propagate through the `with` untouched — Braintrust records + re-raises.
try:
cm = _STATE["bt"].start_span(name=name, type=span_type, **fields)
except Exception as exc: # span creation failed — run trace-less, never break
logger.warning(
"braintrust start_span(%s) failed (%s) — continuing untraced", name, exc
)
yield _NullSpan()
return
with cm as s:
yield s
@@ -0,0 +1,280 @@
"""Real backends for the unified-tool rollout — the bridge between the pure
:func:`run_unified_rollout` loop and live model/tool calls.
``make_call_orchestrator`` returns the ``call_orchestrator`` callable (a teacher
LLM emitting tool calls over the unified spec); ``make_dispatch`` returns the
``dispatch`` callable (executes a chosen tool via ``_call_worker``). Both are
import-safe the OpenAI SDK is imported lazily so this module loads without
network or keys.
"""
from __future__ import annotations
import json
import threading
from typing import Any, Callable, Dict, List, Optional, Tuple
from openjarvis.agents.hybrid._prices import cost as _model_cost
from openjarvis.agents.hybrid.expert_registry import KIND_MODEL, ExpertTool, to_worker_dict
from openjarvis.agents.hybrid.retry import EmptyExpertResponse, with_backoff
from openjarvis.agents.hybrid.toolorchestra.parsing import _parse_rl_tool_call
from openjarvis.agents.hybrid.toolorchestra.rollout import (
build_system_prompt,
)
from openjarvis.agents.hybrid.toolorchestra.tracing import span
from openjarvis.agents.hybrid.toolorchestra.workers import _call_worker
# Guards the lazy, one-time build of the shared ToolExecutor in
# ``_dispatch_openjarvis_tool``: with concurrent rollouts (parallel rejection
# sampling) several threads can reach the build at once and would each instantiate
# the full tool registry. The lock makes it build-once.
_EXECUTOR_BUILD_LOCK = threading.Lock()
def make_call_orchestrator(
model: str,
*,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
temperature: float = 1.0,
max_tokens: int = 4096,
timeout: float = 600.0,
native_tools: bool = True,
frequency_penalty: float = 0.2,
presence_penalty: float = 0.0,
repetition_penalty: float = 1.15,
) -> Callable[..., Tuple[str, List[Tuple[str, Dict[str, Any]]], int, int]]:
"""Teacher-orchestrator caller. ``base_url=None`` → OpenAI cloud; set it to a
vLLM endpoint (with ``api_key="EMPTY"``) to drive a local served teacher.
``native_tools`` picks the tool-calling convention, and the two modes MUST NOT
be mixed on the same model:
* ``True`` (default DATA GENERATION with a *base* model): pass the OpenAI
``tools=`` param so the server's native tool template + parser (gemma4 /
qwen3_xml / hermes) drive the base model, which only emits reliable tool
calls that way. The rollout captures ``(name, arguments)`` and the serializer
re-renders them into the canonical JSON ``<tool_call>`` form regardless.
* ``False`` (SERVING a *fine-tuned* model): bake the ``<tools>`` block + JSON
``<tool_call>`` format into the system prompt (exactly what the serializer
trained on) and DROP ``tools=``. If ``tools=`` is set here, the served chat
template injects its own XML ``<function=>`` instructions, which conflict
with the JSON format the model learned -> it falls back to ``\\boxed{}`` and
never routes. ``_parse_rl_tool_call`` scrapes the JSON tag from raw text.
"""
# Create the client ONCE and reuse it for every turn. Creating a fresh OpenAI
# client per call (as before) leaks an httpx connection pool each time -> under
# heavy parallel generation, sockets pile up in CLOSE-WAIT, the run degrades and
# has to be restarted. The orchestrator is the highest-frequency call, so reusing
# one client here removes the dominant leak. httpx clients are thread-safe.
from openai import OpenAI
from openjarvis.agents.hybrid.toolorchestra.tracing import wrap_client
client = wrap_client(
OpenAI(base_url=base_url, api_key=api_key or "EMPTY", timeout=timeout)
)
def call_orchestrator(messages: List[Dict[str, Any]], specs: List[Dict[str, Any]]):
# ``messages`` is the full running conversation (system/user/assistant/tool)
# built by run_unified_rollout — pass it through so the model sees its own
# prior calls + tool responses, not a flattened user blob.
if native_tools:
# Base-model generation: native tool template drives the tool calls.
send = messages
kwargs = {"tools": specs} if specs else {}
else:
# Fine-tuned serving: JSON format is baked into the system prompt and
# tools= is dropped (see make_call_orchestrator docstring).
send = list(messages)
if specs and send and send[0].get("role") == "system":
send[0] = {"role": "system", "content": build_system_prompt(specs)}
kwargs = {}
# repetition_penalty is a vLLM extra (not OpenAI-native); only send it to a
# local vLLM endpoint (base_url set), never to the cloud frontier APIs.
if base_url and repetition_penalty and repetition_penalty != 1.0:
kwargs.setdefault("extra_body", {})["repetition_penalty"] = (
repetition_penalty
)
# Retry transient failures (429 / 5xx / connection) with exponential
# backoff + jitter. Without this a single rate-limit kills the rollout —
# fine at concurrency 12, ruinous at 100.
resp = with_backoff(
lambda: client.chat.completions.create(
model=model,
messages=send,
temperature=temperature,
max_tokens=max_tokens,
frequency_penalty=frequency_penalty,
presence_penalty=presence_penalty,
**kwargs,
),
what=f"orchestrator call ({model})",
)
msg = resp.choices[0].message
text = msg.content or ""
sdk_tool_calls = getattr(msg, "tool_calls", None)
u = resp.usage
p = getattr(u, "prompt_tokens", 0) if u else 0
c = getattr(u, "completion_tokens", 0) if u else 0
parsed = _parse_rl_tool_call(text, sdk_tool_calls)
tool_calls = [(parsed["name"], parsed["arguments"])] if parsed else []
return text, tool_calls, int(p), int(c)
return call_orchestrator
def _dispatch_openjarvis_tool(
tool: ExpertTool,
arguments: Dict[str, Any],
executor_holder: Dict[str, Any],
) -> Tuple[str, float, int, bool]:
"""Execute a bridged real OpenJarvis tool via its ``ToolExecutor``.
Builds the executor lazily (cached in ``executor_holder``) and degrades to a
clear error string never a crash if the tool registry / executor can't be
instantiated. Returns ``(content, cost_usd, total_tokens, is_local=True)``.
"""
try:
executor = executor_holder.get("executor")
if executor is None:
with _EXECUTOR_BUILD_LOCK:
# Re-check inside the lock: another thread may have built it while
# we waited.
executor = executor_holder.get("executor")
if executor is None:
from openjarvis.core.registry import ToolRegistry
from openjarvis.tools._stubs import ToolExecutor
instances = []
for name in ToolRegistry.keys():
entry = ToolRegistry.get(name)
try:
instances.append(
entry() if isinstance(entry, type) else entry
)
except Exception:
continue
# Auto-approve confirmation-gated tools (shell_exec,
# git_commit, ...): rollouts/eval run headless in a sandbox, so
# there is no human to confirm. Without this, those tools return
# a "requires confirmation" error instead of executing — which
# would silently break TerminalBench.
executor = ToolExecutor(
instances,
interactive=True,
confirm_callback=lambda _prompt: True,
)
executor_holder["executor"] = executor
from openjarvis.core.types import ToolCall
call = ToolCall(
id=f"orch-{tool.name}",
name=str(tool.model),
arguments=json.dumps(arguments or {}),
)
result = executor.execute(call)
usage = getattr(result, "usage", None) or {}
total_tokens = int(usage.get("total_tokens", 0) or 0)
return (
str(getattr(result, "content", "")),
float(getattr(result, "cost_usd", 0.0) or 0.0),
total_tokens,
True,
)
except Exception as exc: # never crash the rollout on a tool-bridge failure
return (f"[openjarvis-tool error: {tool.model}: {exc}]", 0.0, 0, True)
def make_dispatch(
cfg: Optional[Dict[str, Any]] = None,
) -> Callable[[ExpertTool, Dict[str, Any]], Tuple[str, float, int, bool]]:
"""Tool-execution caller: run the chosen tool and return (obs, cost, tokens, is_local)."""
cfg = cfg or {}
executor_holder: Dict[str, Any] = {}
def _dispatch_inner(tool: ExpertTool, arguments: Dict[str, Any]):
# Bridged real OpenJarvis tools run through the ToolExecutor, not
# the model/worker path.
if tool.backend_type == "openjarvis-tool":
return _dispatch_openjarvis_tool(tool, arguments or {}, executor_holder)
worker = to_worker_dict(tool)
prompt = ""
for key in ("input", "query", "code"):
val = arguments.get(key)
if isinstance(val, str) and val.strip():
prompt = val
break
if not prompt.strip():
# The orchestrator emitted a tool call with no/empty input. Don't hit
# the API (it 400s on empty content) — return a usable error so the
# rollout keeps going instead of dropping.
return (
f"[{tool.name}: no input provided — supply a non-empty "
"'input' to delegate]",
0.0,
0,
True,
)
# Retry transient expert failures (429 / 5xx / connection) AND the sneaky
# one: a 200-OK with an empty body. Nothing raises on that, so no retry
# fires, the rollout gets an empty observation, and the clean gate bins the
# whole trajectory. The OpenRouter-hosted Qwen 122B/397B do this (audit
# 2026-07-13: 6 of 72 rollouts lost).
#
# MODEL EXPERTS ONLY. A *sandbox* tool returning nothing is not a failure —
# `code_interpreter` legitimately prints nothing when the code just defines
# a function. Retrying that burned 18 of 72 retry-chains (~2 min of a
# blocked worker thread each) and then killed the rollout outright. Only an
# expert that answers with silence is broken.
is_model_expert = tool.kind == KIND_MODEL
def _call_expert():
out = _call_worker(worker, prompt, cfg)
if is_model_expert and not (out[0] or "").strip():
raise EmptyExpertResponse(str(tool.model))
return out
text, p, c, is_local, extra_cost, _n = with_backoff(
_call_expert,
what=f"expert call ({tool.model})",
)
usd = (0.0 if is_local else _model_cost(str(tool.model), p, c)) + float(
extra_cost
)
return text, usd, int(p) + int(c), bool(is_local)
def dispatch(tool: ExpertTool, arguments: Dict[str, Any]):
# One span per route so the trace tree shows which expert/tool was called,
# with the delegated input, the returned observation, and cost/tokens. The
# wrapped OpenAI/Anthropic call (if any) nests one level deeper.
# Span is titled with the REAL model that actually ran (``tool.model``) so
# the trace reads clearly even under anonymization; ``anon_label`` in the
# metadata records the opaque label the orchestrator actually saw/chose.
real_model = str(tool.model)
with span(
f"route:{real_model}",
span_type="tool",
input=arguments,
metadata={
"real_model": real_model,
"anon_label": tool.name,
"backend": tool.backend_type,
},
) as s:
obs, usd, toks, is_local = _dispatch_inner(tool, arguments)
s.log(
output=obs,
metrics={"cost_usd": usd, "tokens": toks},
metadata={"is_local": is_local},
)
return obs, usd, toks, is_local
return dispatch
__all__ = ["make_call_orchestrator", "make_dispatch"]
@@ -0,0 +1,477 @@
"""Worker pool resolution + dispatch for ToolOrchestraAgent."""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from openjarvis.agents.hybrid._base import (
ANTHROPIC_WEB_SEARCH_TOOL,
GEMINI_SEARCH_COST_PER_CALL,
OPENAI_WEB_SEARCH_COST_PER_CALL,
WEB_SEARCH_COST_PER_CALL,
LocalCloudAgent,
)
from openjarvis.agents.hybrid._prices import (
PRICES,
is_gpt5_family,
supports_temperature,
)
from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
_call_modal_python,
_call_tavily_search,
)
def _default_pool(
local_model: Optional[str],
local_endpoint: Optional[str],
cloud_model: str = "claude-opus-4-7",
cloud_endpoint: str = "anthropic",
) -> List[Dict[str, Any]]:
"""Default heterogeneous worker pool.
The frontier worker's ``type`` + ``model`` track the cell's resolved
``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5,
gemini-2.5-pro, ) route their frontier slot to the right SDK.
"""
ep = (cloud_endpoint or "anthropic").lower()
if ep not in ("anthropic", "openai", "gemini"):
ep = "anthropic"
pool: List[Dict[str, Any]] = []
if local_model and local_endpoint:
pool.append(
{
"id": len(pool),
"name": "local-qwen",
"type": "vllm",
"model": local_model,
"base_url": local_endpoint,
"description": (
"Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
"concise extraction, formatting, arithmetic on given data."
),
}
)
if ep == "openai":
search_type = "openai-web-search"
search_model = cloud_model
search_desc = "OpenAI hosted web search on the configured frontier model."
elif ep == "gemini":
search_type = "gemini-web-search"
search_model = cloud_model
search_desc = "Gemini Google Search grounding on the configured frontier model."
else:
search_type = "anthropic-web-search"
search_model = _DEFAULT_WEB_SEARCH_MODEL
search_desc = "Anthropic server-side web_search."
pool.append(
{
"id": len(pool),
"name": "web-search",
"type": search_type,
"model": search_model,
"description": (
f"{search_desc} Use for facts that need a lookup "
"(recent events, rare names/dates, niche sources). Returns a digest."
),
}
)
pool.append(
{
"id": len(pool),
"name": f"frontier-{ep}",
"type": ep,
"model": cloud_model,
"description": (
"Frontier reasoning model. Use for hard multi-step reasoning, "
"code review, or a final synthesis pass. Expensive — use sparingly."
),
}
)
pool.append(
{
"id": len(pool),
"name": "frontier-openai-mini",
"type": "openai",
"model": "gpt-5-mini",
"description": (
"Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
"fraction of frontier cost."
),
}
)
return pool
# Worker types toolorchestra's `_call_worker` actually dispatches.
#
# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
# `tavily-search` — Tavily API search (the paper's web tool).
# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
# Used for the code/math specialists and Llama-3.3-70B /
# Qwen3-32B generalists.
# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
# paper's "Python sandbox" inside `enhance_reasoning`).
_TOOLORCH_VALID_TYPES = (
"vllm",
"openai",
"anthropic",
"anthropic-web-search",
"openai-web-search",
"gemini",
"gemini-web-search",
"tavily-search",
"openrouter",
"modal-python",
)
_TOOLORCH_SEARCH_TYPES = (
"anthropic-web-search",
"openai-web-search",
"gemini-web-search",
"tavily-search",
)
# Default model used when an `anthropic-web-search` entry omits `model`.
_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
def _resolve_worker_pool(
cfg: Dict[str, Any],
local_model: Optional[str],
local_endpoint: Optional[str],
cloud_model: str,
cloud_endpoint: str = "anthropic",
) -> List[Dict[str, Any]]:
"""Return the worker pool for this run.
Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
default pool is ignored entirely. Falls back to ``_default_pool`` when
the override is absent.
Each user-supplied entry must be a dict with keys ``id``, ``name``,
``type``, and (for non-search types) ``model``. Search worker types are
``anthropic-web-search``, ``openai-web-search``, ``gemini-web-search``,
and ``tavily-search``. ``anthropic-web-search`` entries may omit
``model`` it defaults to ``claude-haiku-4-5``. OpenAI and Gemini
search workers default to the configured cloud model. Tavily does not
require a model.
Substitution: ``model = "$local"`` (or ``"<local>"``) resolves to
``local_model``; ``model = "$cloud"`` / ``"<cloud>"`` to ``cloud_model``.
On any validation failure, raises ``ValueError`` with the message
``"Invalid worker_pool entry [<id>]: <reason>"``. Fails fast at agent
init rather than mid-task.
"""
override = cfg.get("worker_pool")
if override is None:
return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint)
if not isinstance(override, list) or not override:
raise ValueError(
"Invalid worker_pool entry [-]: worker_pool must be a non-empty list"
)
resolved: List[Dict[str, Any]] = []
seen_ids: set = set()
has_non_search = False
for raw in override:
wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?"
if not isinstance(raw, dict):
raise ValueError(
f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict"
)
entry = dict(raw)
wid = entry.get("id")
if not isinstance(wid, int):
raise ValueError(
f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int"
)
if wid in seen_ids:
raise ValueError(f"Invalid worker_pool entry [{wid}]: duplicate id")
seen_ids.add(wid)
if not entry.get("name") or not isinstance(entry["name"], str):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string"
)
wtype = entry.get("type") or entry.get("endpoint")
if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'type' must be one of "
f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})"
)
wtype = wtype.lower()
entry["type"] = wtype
# Substitute $local / $cloud placeholders (before any model check).
model = entry.get("model")
if isinstance(model, str) and model in ("$local", "<local>"):
if not local_model:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model='{model}' "
"requires a local_model to be configured for this cell"
)
model = local_model
entry["model"] = model
elif isinstance(model, str) and model in ("$cloud", "<cloud>"):
model = cloud_model
entry["model"] = model
if wtype in _TOOLORCH_SEARCH_TYPES:
if model in (None, ""):
if wtype == "anthropic-web-search":
model = _DEFAULT_WEB_SEARCH_MODEL
elif wtype in ("openai-web-search", "gemini-web-search"):
model = cloud_model
else:
model = wtype
entry["model"] = model
elif not isinstance(model, str):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
)
if (
wtype in ("openai-web-search", "gemini-web-search")
and model not in PRICES
):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model {model!r} "
f"is not in PRICES (known: {sorted(PRICES)})"
)
# Search workers don't satisfy the "needs a solver" requirement.
else:
if not isinstance(model, str) or not model:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
)
if wtype == "vllm":
if not entry.get("base_url"):
if not local_endpoint:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: vllm worker needs "
"'base_url' (or a configured local_endpoint to fall back to)"
)
entry["base_url"] = local_endpoint
entry.setdefault("api_key", "EMPTY")
else:
if model not in PRICES:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model {model!r} "
f"is not in PRICES (known: {sorted(PRICES)})"
)
has_non_search = True
entry.setdefault(
"description",
f"User-supplied {wtype} worker ({model}).",
)
resolved.append(entry)
if not has_non_search:
raise ValueError(
"Invalid worker_pool entry [-]: worker_pool must contain at least "
"one non-search worker (vllm / openai / anthropic / gemini)"
)
return resolved
# Anthropic model id -> OpenRouter slug (for OJ_ANTHROPIC_VIA_OPENROUTER=1).
_ANTHROPIC_OPENROUTER_SLUGS = {
"claude-opus-4-8": "anthropic/claude-opus-4.8",
"claude-opus-4-7": "anthropic/claude-opus-4.7",
"claude-haiku-4-5-20251001": "anthropic/claude-haiku-4.5",
}
def _call_worker(
worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
) -> Tuple[str, int, int, bool, float, int]:
"""Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
wtype = worker.get("type", "openai")
max_tok = int(
cfg.get("worker_max_tokens") or os.environ.get("OJ_WORKER_MAX_TOKENS", "4096")
)
temp = float(cfg.get("worker_temperature", 0.2))
if wtype == "vllm":
text, p, c = LocalCloudAgent._call_vllm(
worker["model"],
worker["base_url"],
user=prompt,
max_tokens=max_tok,
temperature=temp,
enable_thinking=False,
)
return text, p, c, True, 0.0, 0
if wtype == "openai":
is_gpt5 = is_gpt5_family(worker["model"])
eff_temp = 1.0 if is_gpt5 else temp
# GPT-5 is a reasoning model: hidden reasoning tokens count against
# `max_completion_tokens`, so a 4096 cap can be fully consumed by
# reasoning and leave 0 visible content (empty answer). Give the
# reasoning headroom on top of the answer budget.
eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
text, p, c = LocalCloudAgent._call_openai(
worker["model"],
user=prompt,
max_tokens=eff_max_tok,
temperature=eff_temp,
)
return text, p, c, False, 0.0, 0
if wtype == "gemini":
text, p, c = LocalCloudAgent._call_gemini(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=temp,
)
return text, p, c, False, 0.0, 0
if wtype == "anthropic":
# Escape hatch for a dead/rotated ANTHROPIC_API_KEY: the same Claude
# models are served BY Anthropic through OpenRouter, billed on the
# OpenRouter key. Flip OJ_ANTHROPIC_VIA_OPENROUTER=1 and every
# anthropic-typed expert keeps working (2026-07-13: key rotated
# mid-eval and zeroed a run — the judge 401'd and every answer
# defaulted to score 0).
if os.environ.get("OJ_ANTHROPIC_VIA_OPENROUTER") == "1":
slug = _ANTHROPIC_OPENROUTER_SLUGS.get(
worker["model"], f"anthropic/{worker['model']}"
)
text, p, c = LocalCloudAgent._call_openrouter(
slug, user=prompt, max_tokens=max_tok, temperature=temp
)
return text, p, c, False, 0.0, 0
eff_temp = temp if supports_temperature(worker["model"]) else 0.0
text, p, c, _ = LocalCloudAgent._call_anthropic(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=eff_temp,
)
return text, p, c, False, 0.0, 0
if wtype == "anthropic-web-search":
eff_temp = temp if supports_temperature(worker["model"]) else 0.0
text, p, c, n_searches = LocalCloudAgent._call_anthropic(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=eff_temp,
tools=[ANTHROPIC_WEB_SEARCH_TOOL],
tool_choice={"type": "any"},
)
extra = n_searches * WEB_SEARCH_COST_PER_CALL
return text, p, c, False, extra, n_searches
if wtype == "openai-web-search":
eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp
text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent(
worker["model"],
user=prompt,
max_tokens=max(max_tok, 16384)
if is_gpt5_family(worker["model"])
else max_tok,
temperature=eff_temp,
)
extra = n_searches * OPENAI_WEB_SEARCH_COST_PER_CALL
return text, p, c, False, extra, n_searches
if wtype == "gemini-web-search":
text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=temp,
)
extra = n_searches * GEMINI_SEARCH_COST_PER_CALL
return text, p, c, False, extra, n_searches
if wtype == "tavily-search":
# Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL`
# for parity with the Anthropic web-search worker. One call = one
# "n_search" for accounting.
max_results = int(cfg.get("tavily_max_results", 5))
text, p, c = _call_tavily_search(str(prompt), max_results=max_results)
return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1
if wtype == "openrouter":
# Pin to providers measured to actually return completions. OpenRouter
# load-balances across upstreams, and several return 200-OK with an EMPTY
# body (measured 2026-07-13: Chutes 2/2 empty, SiliconFlow 1/1,
# DigitalOcean 1/1, Parasail 1/1 — a blocklist was whack-a-mole, new
# broken providers kept appearing). Allowlist verified 12/12 non-empty;
# ``allow_fallbacks: False`` stops OpenRouter from silently routing
# outside it. Retries only ever "worked" by re-rolling this dice.
text, p, c = LocalCloudAgent._call_openrouter(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=temp,
extra_body={
"provider": {
"order": ["Alibaba", "Novita", "GMICloud", "AtlasCloud"],
"allow_fallbacks": False,
}
},
)
return text, p, c, False, 0.0, 0
if wtype == "modal-python":
# `prompt` is the python code string to exec.
timeout_s = int(cfg.get("modal_python_timeout_s", 60))
out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
# No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
# charges per sandbox-second — not tracked here.
return out, 0, 0, False, 0.0, 0
raise ValueError(f"unsupported worker type: {wtype!r}")
def _swe_call_worker(
worker: Dict[str, Any],
prompt: str,
cfg: Dict[str, Any],
task: Dict[str, Any],
workdir: Path,
turn: int,
) -> Tuple[str, int, int, bool, float, int, int]:
"""SWE-bench worker dispatch: route solver workers through
run_swe_agent_loop on a shared workdir. Web-search workers fall back
to the regular one-shot dispatch (search isn't an agent loop).
Trailing ``bash_turns`` (last element) counts agent-loop turns so the
caller can surface ``tool_calls`` per row. Fallbacks to one-shot
workers return 0 bash turns (no agent loop ran)."""
wtype = worker.get("type", "openai")
if wtype in _TOOLORCH_SEARCH_TYPES:
# Search workers stay one-shot.
text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
return text, p, c, is_local, extra, n_searches, 0
if wtype == "vllm":
backbone = "local"
endpoint = worker.get("base_url")
loop_cloud_endpoint = "anthropic" # unused when backbone=local
elif wtype in ("anthropic", "openai", "gemini"):
backbone = "cloud"
endpoint = None
loop_cloud_endpoint = wtype
else:
# Unknown type — one-shot fallback.
text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
return text, p, c, is_local, extra, n_searches, 0
out = run_swe_agent_loop(
task,
backbone=backbone,
backbone_model=worker["model"],
cloud_endpoint=loop_cloud_endpoint,
local_endpoint=endpoint,
initial_prompt=prompt,
max_turns=int(cfg.get("swe_max_turns", 30)),
bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
output_cap=int(cfg.get("swe_output_cap", 10_000)),
turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
trace_prefix=f"toolorch_turn{turn}",
workdir=workdir,
)
is_local = backbone == "local"
return (
out["final_summary"] or out["answer"],
out["tokens_in"],
out["tokens_out"],
is_local,
0.0,
0,
int(out["turns"]),
)
+89 -1
View File
@@ -13,6 +13,7 @@ Supports two modes:
from __future__ import annotations
import concurrent.futures
import json
import re
from typing import Any, List, Optional
@@ -201,6 +202,70 @@ class OrchestratorAgent(ToolUsingAgent):
return result
@staticmethod
def _boxed_tool_call(content, openai_tools, call_id):
r"""Salvage a tool call the model wrote as ``\boxed{tool: query}``.
Some SFT'd checkpoints (Qwen reasoning-mode habit) emit their action
as a boxed string instead of a real ``<tool_call>`` e.g.
``\boxed{web_search: latest news}`` or ``\boxed{web_search query: ...}``.
The engine's tool-call parser sees no structured call, so without this
the agent would grade the half-finished thought as its final answer and
never actually act. Here we recover ``(tool, args)`` and build a real
``ToolCall`` keyed on the tool's primary parameter. Returns ``None`` when
the boxed content isn't a recognisable tool action (leaving genuine
boxed final answers untouched).
"""
if not openai_tools or "\\boxed" not in content:
return None
# tool name -> primary parameter (first required, else first property)
param_of: dict[str, str] = {}
for t in openai_tools:
fn = t.get("function", t)
name = fn.get("name")
if not name:
continue
params = fn.get("parameters") or {}
props = list((params.get("properties") or {}).keys())
req = params.get("required") or []
param_of[name] = req[0] if req else (props[0] if props else "query")
# The action is the LAST boxed expression in the response.
boxed = re.findall(r"\\boxed\{([^{}]*)\}", content)
if not boxed:
return None
inner = boxed[-1].strip()
# Match a leading tool name (longest first to avoid prefix collisions),
# then strip the separator between name and args. Handles
# ``tool: q`` / ``tool query: q`` / ``tool q`` / ``tool(q)``.
for name in sorted(param_of, key=len, reverse=True):
m = re.match(
r"^\s*" + re.escape(name) + r"\b(.*)$",
inner,
re.DOTALL | re.IGNORECASE,
)
if not m:
continue
rest = m.group(1).strip()
if rest.startswith("(") and rest.endswith(")"):
rest = rest[1:-1] # tool(args) wrapper -> drop matching parens
else:
# optional "query" word, then a ":" / "-" separator
rest = re.sub(
r"^\s*(?:query\s*)?[:\-]\s*", "", rest, flags=re.IGNORECASE
)
arg = rest.strip().strip("\"'").strip()
if not arg:
return None
return ToolCall(
id=call_id,
name=name,
arguments=json.dumps({param_of[name]: arg}),
)
return None
# ------------------------------------------------------------------
# Function-calling mode (original behaviour)
# ------------------------------------------------------------------
@@ -245,8 +310,31 @@ class OrchestratorAgent(ToolUsingAgent):
content = result.get("content", "")
raw_tool_calls = result.get("tool_calls", [])
# No tool calls -> check continuation, then final answer
# No tool calls -> try to salvage a \boxed{tool: query} action the
# model wrote instead of a real <tool_call>; else final answer.
if not raw_tool_calls:
boxed_call = self._boxed_tool_call(
content, openai_tools, f"boxed_{turns}"
)
if boxed_call is not None:
messages.append(
Message(
role=Role.ASSISTANT,
content=content,
tool_calls=[boxed_call],
)
)
tool_result = self._executor.execute(boxed_call)
all_tool_results.append(tool_result)
messages.append(
Message(
role=Role.TOOL,
content=tool_result.content,
tool_call_id=boxed_call.id,
name=boxed_call.name,
)
)
continue
content = self._check_continuation(result, messages)
content = self._strip_think_tags(content)
self._emit_turn_end(turns=turns, content_length=len(content))
+67 -8
View File
@@ -2,12 +2,44 @@
from __future__ import annotations
import logging
import random
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Tuple
from openjarvis.evals.core.backend import InferenceBackend
from openjarvis.evals.core.types import EvalRecord
LOGGER = logging.getLogger(__name__)
# Transient failures worth retrying the judge call on. A single rate-limit
# (429) used to fall straight through to `llm_fallback_error` and score the
# sample WRONG — e.g. 93/100 GAIA judge calls 429'd on one run and zeroed the
# whole bench. Retry with exponential backoff so a busy judge endpoint doesn't
# get mis-scored as a wrong answer.
_JUDGE_MAX_RETRIES = 6
_JUDGE_BASE_DELAY_S = 2.0
_JUDGE_MAX_DELAY_S = 60.0
_RETRYABLE_MARKERS = (
"429",
"rate_limit",
"rate limit",
"overloaded",
"timeout",
"timed out",
"503",
"502",
"500",
"connection",
"temporarily unavailable",
)
def _is_retryable_judge_error(exc: Exception) -> bool:
msg = str(exc).lower()
return any(marker in msg for marker in _RETRYABLE_MARKERS)
class Scorer(ABC):
"""Base class for all scorers."""
@@ -42,14 +74,41 @@ class LLMJudgeScorer(Scorer):
temperature: float = 0.0,
max_tokens: int = 2048,
) -> str:
"""Send a prompt to the judge LLM and return the response text."""
return self._judge_backend.generate(
prompt,
model=self._judge_model,
system=system,
temperature=temperature,
max_tokens=max_tokens,
)
"""Send a prompt to the judge LLM and return the response text.
Retries transient failures (429 / rate-limit / 5xx / timeout) with
exponential backoff + jitter so a busy judge endpoint doesn't get
mis-scored as a wrong answer. Non-retryable errors, or exhaustion of
the retry budget, re-raise to the caller (which records the failure).
"""
last_exc: Optional[Exception] = None
for attempt in range(_JUDGE_MAX_RETRIES):
try:
return self._judge_backend.generate(
prompt,
model=self._judge_model,
system=system,
temperature=temperature,
max_tokens=max_tokens,
)
except Exception as exc: # noqa: BLE001 - re-raised below
last_exc = exc
if attempt == _JUDGE_MAX_RETRIES - 1 or not _is_retryable_judge_error(
exc
):
raise
delay = min(_JUDGE_BASE_DELAY_S * (2**attempt), _JUDGE_MAX_DELAY_S)
delay += random.uniform(0.0, delay * 0.25) # jitter
LOGGER.warning(
"judge call failed (attempt %d/%d): %s — retrying in %.1fs",
attempt + 1,
_JUDGE_MAX_RETRIES,
exc,
delay,
)
time.sleep(delay)
# Unreachable (loop either returns or raises), but keeps type-checkers happy.
raise last_exc # type: ignore[misc]
__all__ = ["LLMJudgeScorer", "Scorer"]
+8 -1
View File
@@ -140,7 +140,14 @@ class GAIAScorer(LLMJudgeScorer):
if structured_match:
is_correct = structured_match.group(1).lower() == "yes"
else:
is_correct = "CORRECT" in raw.upper() and "INCORRECT" not in raw.upper()
up = raw.upper()
# "not correct" / "is not correct" contain "CORRECT" and lack
# "INCORRECT" — guard the negations explicitly or they read True.
is_correct = (
"CORRECT" in up
and "INCORRECT" not in up
and "NOT CORRECT" not in up
)
meta: Dict[str, Any] = {
"match_type": "llm_fallback",
+36 -7
View File
@@ -27,6 +27,25 @@ class MMLUProScorer(LLMJudgeScorer):
return "".join(chr(ord("A") + i) for i in range(n))
return "ABCDEFGHIJ"
def _extract_answer_direct(
self,
model_answer: str,
valid_letters: str,
) -> Optional[str]:
"""High-confidence, unambiguous letter straight from the output.
Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex
prose like "the answer is X": the letters "I" and "A" are real English
words, so "FINAL_ANSWER: I need ..." would grab "I" the exact bug that
mislabels answers. Semantic extraction is left to the judge below.
"""
if not model_answer:
return None
for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)):
if cand.upper() in valid_letters:
return cand.upper()
return None
def _extract_answer_with_llm(
self,
problem: str,
@@ -56,10 +75,13 @@ class MMLUProScorer(LLMJudgeScorer):
)
extracted = raw_response.strip().upper()
if "NONE" in extracted:
return None
# Handle "The answer is: A" etc.
# Accept only an ISOLATED letter — never the first cap of a word like
# "It"/"None". Prefer an explicit "the answer is X" phrasing.
answer_match = re.search(
r"(?:THE ANSWER IS:?\s*)?([A-Z])",
r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])",
extracted,
re.IGNORECASE,
)
@@ -86,11 +108,17 @@ class MMLUProScorer(LLMJudgeScorer):
valid_letters = self._valid_letters_from_options(record.metadata)
candidate = self._extract_answer_with_llm(
record.problem,
model_answer,
valid_letters,
)
# Parse straight from the output first (fast, deterministic, no "I" bug);
# only fall back to the judge when no isolated letter is present.
method = "direct"
candidate = self._extract_answer_direct(model_answer, valid_letters)
if not candidate:
method = "llm"
candidate = self._extract_answer_with_llm(
record.problem,
model_answer,
valid_letters,
)
if not candidate:
return None, {"reason": "no_choice_letter_extracted"}
@@ -99,6 +127,7 @@ class MMLUProScorer(LLMJudgeScorer):
"reference_letter": ref,
"candidate_letter": candidate,
"valid_letters": valid_letters,
"extract_method": method,
}
return is_correct, meta
+36 -7
View File
@@ -27,6 +27,25 @@ class SuperGPQAScorer(LLMJudgeScorer):
return "".join(chr(ord("A") + i) for i in range(n))
return "ABCD"
def _extract_answer_direct(
self,
model_answer: str,
valid_letters: str,
) -> Optional[str]:
"""High-confidence, unambiguous letter straight from the output.
Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex
prose like "the answer is X": the letters "I" and "A" are real English
words, so "FINAL_ANSWER: I need ..." would grab "I" the exact bug that
mislabels answers. Semantic extraction is left to the judge below.
"""
if not model_answer:
return None
for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)):
if cand.upper() in valid_letters:
return cand.upper()
return None
def _extract_answer_with_llm(
self,
problem: str,
@@ -56,10 +75,13 @@ class SuperGPQAScorer(LLMJudgeScorer):
)
extracted = raw_response.strip().upper()
if "NONE" in extracted:
return None
# Handle "The answer is: A" etc.
# Accept only an ISOLATED letter — never the first cap of a word like
# "It"/"None". Prefer an explicit "the answer is X" phrasing.
answer_match = re.search(
r"(?:THE ANSWER IS:?\s*)?([A-Z])",
r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])",
extracted,
re.IGNORECASE,
)
@@ -86,11 +108,17 @@ class SuperGPQAScorer(LLMJudgeScorer):
valid_letters = self._valid_letters_from_options(record.metadata)
candidate = self._extract_answer_with_llm(
record.problem,
model_answer,
valid_letters,
)
# Parse straight from the output first (fast, deterministic, no "I" bug);
# only fall back to the judge when no isolated letter is present.
method = "direct"
candidate = self._extract_answer_direct(model_answer, valid_letters)
if not candidate:
method = "llm"
candidate = self._extract_answer_with_llm(
record.problem,
model_answer,
valid_letters,
)
if not candidate:
return None, {"reason": "no_choice_letter_extracted"}
@@ -99,6 +127,7 @@ class SuperGPQAScorer(LLMJudgeScorer):
"reference_letter": ref,
"candidate_letter": candidate,
"valid_letters": valid_letters,
"extract_method": method,
}
return is_correct, meta
@@ -0,0 +1,34 @@
# Orchestrator: Qwen3.5-9B. v1 = base self-sampling (the served orchestrator_model is this same base).
#
# v1 data = base Qwen3.5-9B self-sampling: roll the base model out over
# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the
# dataset first (no GPU, needs the vLLM orchestrator endpoint up):
# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
# --out data/orchestrator_qwen3.5-9b_sft_v1.jsonl --samples-per-task 8
# then train on a rented GPU with regenerate_traces = false.
[model]
model_name = "Qwen/Qwen3.5-9B"
max_seq_length = 4096
[training]
num_epochs = 3
batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch
learning_rate = 2e-5
weight_decay = 0.01
warmup_ratio = 0.1
gradient_checkpointing = true
[data]
trace_cache_path = "data/orchestrator_qwen3.5-9b_sft_v1.jsonl"
regenerate_traces = false
# Base Qwen3.5-9B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
orchestrator_endpoint = "http://localhost:8001/v1"
orchestrator_model = "qwen3.5-9b"
samples_per_task = 8
max_keep_per_task = 1
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_qwen3.5-9b_sft"
save_every_n_epochs = 1
@@ -0,0 +1,288 @@
"""Eval backend that drives the ToolOrchestra orchestrator as the "model".
This adapts our unified-tool orchestrator rollout
(:func:`openjarvis.agents.hybrid.toolorchestra.rollout.run_unified_rollout`)
to the eval framework's :class:`~openjarvis.evals.core.backend.InferenceBackend`
interface so it can be scored by the existing
:class:`~openjarvis.evals.core.runner.EvalRunner` over any registered benchmark.
The orchestrator is an OpenAI-compatible chat model (served locally via vLLM,
default ``http://localhost:8001/v1``) that emits tool calls over the fixed
8-tool catalog from :func:`expert_registry.orchestrator_catalog`. For each eval
sample we run one full rollout and return ``rollout.final_answer`` as the model
answer, plus token/cost telemetry in the ``generate_full`` payload.
Construction is import-safe and network-free; the OpenAI SDK is only touched
inside ``call_orchestrator`` at generate time (lazy import in ``unified.py``).
"""
from __future__ import annotations
import re
import time
from typing import Any, Dict, Optional
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
from openjarvis.agents.hybrid.toolorchestra import rollout as _rollout_mod
from openjarvis.agents.hybrid.toolorchestra.unified import (
make_call_orchestrator,
make_dispatch,
)
from openjarvis.evals.core.backend import InferenceBackend
# The served fine-tuned model sometimes over-emits the answer marker, e.g.
# "FINAL_ANSWER: FINAL_ANSWER: 42" — a doubled prefix that breaks answer
# extraction and auto-scores the sample 0. Collapse any run of FINAL_ANSWER
# markers to a single one, keeping the text after the LAST marker as the answer.
_FA_MARK = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
def _clean_final_answer(text: str) -> str:
t = (text or "").strip()
marks = list(_FA_MARK.finditer(t))
if not marks:
return t
answer = t[marks[-1].end() :].strip()
return f"FINAL_ANSWER: {answer}" if answer else t
_OBS_CAP = 4000 # cap persisted observations so the eval JSONL stays readable
DEFAULT_ENDPOINT = "http://localhost:8001/v1"
DEFAULT_MODEL = "qwen3-8b"
class OrchestratorBackend(InferenceBackend):
"""Run the ToolOrchestra orchestrator rollout as an eval "model".
Parameters
----------
orchestrator_endpoint:
OpenAI-compatible base URL for the served orchestrator (vLLM). Defaults
to ``http://localhost:8001/v1``.
orchestrator_model:
Served model id for the orchestrator. Defaults to ``"qwen3-8b"`` (will
be the Qwen3.5-9B checkpoint later).
api_key:
API key for the endpoint. ``"EMPTY"`` for a local vLLM server.
local_endpoints:
Optional dict mapping a local OSS model id (e.g. ``"Qwen/Qwen3.5-9B"``)
to its vLLM base URL. Unmapped local models are still listed in the
catalog but served as unconfigured (``base_url=None``).
max_turns:
Maximum orchestrator reasoning->action->observation turns per sample.
temperature:
Orchestrator sampling temperature.
dispatch_cfg:
Optional config dict passed to :func:`make_dispatch` (worker dispatch).
"""
backend_id = "orchestrator"
framework_name = "openjarvis-orchestrator"
def _dump_trace(self, prompt, rollout, orch_model) -> None:
"""Append one trajectory_to_record row to $OJ_EVAL_TRACE_DIR/traces.jsonl."""
import hashlib
import json
import os
import threading
from pathlib import Path
trace_dir = os.environ.get("OJ_EVAL_TRACE_DIR")
if not trace_dir:
return
try:
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
rec = trajectory_to_record(
hashlib.md5(prompt.encode()).hexdigest()[:12],
prompt,
self._tools,
rollout,
)
rec["orchestrator_model"] = orch_model
Path(trace_dir).mkdir(parents=True, exist_ok=True)
if not hasattr(self, "_trace_lock"):
self._trace_lock = threading.Lock()
with self._trace_lock:
with open(Path(trace_dir) / "traces.jsonl", "a") as fh:
fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
except Exception: # noqa: BLE001 — tracing must never fail an eval sample
pass
def __init__(
self,
*,
orchestrator_endpoint: str = DEFAULT_ENDPOINT,
orchestrator_model: str = DEFAULT_MODEL,
api_key: str = "EMPTY",
local_endpoints: Optional[Dict[str, str]] = None,
max_turns: int = 8,
temperature: float = 1.0,
finetuned: bool = True,
dispatch_cfg: Optional[Dict[str, Any]] = None,
) -> None:
self.orchestrator_endpoint = orchestrator_endpoint
self.orchestrator_model = orchestrator_model
self.api_key = api_key
self.local_endpoints = dict(local_endpoints or {})
self.max_turns = int(max_turns)
self.temperature = float(temperature)
# False -> serve with native tools= (a BASE model needs the server's tool
# template to emit reliable calls; the baked-JSON prompt is what a
# FINE-TUNED model was trained on). Getting this wrong handicaps the model:
# the first fixed-harness baseline rerun served the base model in baked
# mode and scored 0.184 — not comparable to anything.
self.finetuned = bool(finetuned)
self._dispatch_cfg = dict(dispatch_cfg or {})
# Build the catalog once; it is stateless. ``local_endpoints`` now maps
# full local model ids (e.g. ``"Qwen/Qwen3.5-9B"``) -> vLLM base_url.
self._tools = orchestrator_catalog(local_endpoints=self.local_endpoints)
# ------------------------------------------------------------------
# InferenceBackend abstract methods
# ------------------------------------------------------------------
def generate(
self,
prompt: str,
*,
model: str,
system: str = "",
temperature: float = 0.0,
max_tokens: int = 2048,
) -> str:
"""Run a rollout and return just the final-answer text."""
full = self.generate_full(
prompt,
model=model,
system=system,
temperature=temperature,
max_tokens=max_tokens,
)
return full.get("content", "") or ""
def generate_full(
self,
prompt: str,
*,
model: str,
system: str = "",
temperature: float = 0.0,
max_tokens: int = 2048,
) -> Dict[str, Any]:
"""Run one orchestrator rollout over ``prompt`` and return full details.
Returns a dict with the keys the runner reads: ``content``, ``usage``,
``model``, ``latency_seconds``, ``cost_usd``, plus ``tool_calls`` /
``turn_count`` / ``framework``. On error, returns ``content=""`` and a
non-empty ``error`` field so the runner records it as a failed sample
rather than crashing the whole run.
"""
# The caller-supplied ``model`` is the RunConfig.model (the orchestrator
# model name). Fall back to the configured one if blank.
orch_model = model or self.orchestrator_model
# Prefer the explicit per-call temperature only when the caller set a
# non-default value; otherwise use the backend's configured temperature.
temp = temperature if temperature else self.temperature
started = time.time()
try:
call_orch = make_call_orchestrator(
orch_model,
base_url=self.orchestrator_endpoint,
api_key=self.api_key,
temperature=temp,
max_tokens=max_tokens,
# Fine-tuned model: baked JSON <tool_call> prompt, no tools=.
# Base model (finetuned=False): native tools= — it can't emit
# reliable calls otherwise.
native_tools=not self.finetuned,
)
dispatch = make_dispatch(self._dispatch_cfg)
rollout = _rollout_mod.run_unified_rollout(
prompt,
self._tools,
call_orchestrator=call_orch,
dispatch=dispatch,
max_turns=self.max_turns,
# Match TRAINING: the model was fine-tuned on anonymized expert
# labels (expert_xxxx), so it only picks valid labels / routes
# correctly when it sees the same anonymized names at inference.
anonymize=True,
)
# Persist the FULL trajectory when asked (OJ_EVAL_TRACE_DIR). The
# scored JSONL keeps only the final answer — without this every eval
# conversation is discarded and failures can't be audited (e.g.
# "is GAIA 0.20 turn-exhaustion or wrong answers?" was unanswerable
# for the first sft879 run). Rows are the same `conversations`
# records the generation pipeline writes, so render_sft_data.py
# renders them as per-sample markdown.
self._dump_trace(prompt, rollout, orch_model)
except Exception as exc: # noqa: BLE001 - surface as a failed sample
return {
"content": "",
"error": f"{type(exc).__name__}: {exc}",
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
"model": orch_model,
"latency_seconds": time.time() - started,
"cost_usd": 0.0,
"tool_calls": 0,
"turn_count": 0,
"framework": self.framework_name,
}
latency = time.time() - started
total_tokens = int(getattr(rollout, "tokens", 0) or 0)
anon_map = getattr(rollout, "anon_map", None) or {}
return {
"content": _clean_final_answer(rollout.final_answer or ""),
# The rollout reports a single combined token count; we surface it
# as completion_tokens so it still flows into total-token telemetry.
"usage": {
"prompt_tokens": 0,
"completion_tokens": total_tokens,
},
"model": orch_model,
"latency_seconds": latency,
"cost_usd": float(getattr(rollout, "cost_usd", 0.0) or 0.0),
"tool_calls": int(getattr(rollout, "num_tool_calls", 0) or 0),
"turn_count": len(getattr(rollout, "turns", []) or []),
"framework": self.framework_name,
"trace_data": {
"parse_failures": int(getattr(rollout, "parse_failures", 0) or 0),
"tool_calls": [
{"name": n, "arguments": a} for n, a in rollout.tool_calls()
],
# Full step-by-step trace so eval samples are inspectable
# (format_eval_sample.py renders this). Observations capped.
"anon_map": anon_map,
"turns": [
{
"reasoning": t.reasoning or "",
"tool_name": t.tool_name,
"real_model": anon_map.get(t.tool_name)
if t.tool_name
else None,
"arguments": t.arguments,
"observation": (
(t.observation or "")[:_OBS_CAP]
+ (
"…[truncated]"
if t.observation and len(t.observation) > _OBS_CAP
else ""
)
),
}
for t in (getattr(rollout, "turns", []) or [])
],
"final_answer": _clean_final_answer(rollout.final_answer or ""),
},
}
__all__ = ["OrchestratorBackend", "DEFAULT_ENDPOINT", "DEFAULT_MODEL"]
@@ -0,0 +1,29 @@
"""SFT-data generation for the orchestrator cold-start.
Execution-grounded rejection sampling: for each reasoning task (GeneralThought +
OpenThoughts, via :func:`~...sft_data.datasets.load_sft_tasks`), roll out a
teacher orchestrator over the unified tool catalog N times, verify each
trajectory, keep the cheapest passing one(s), and serialize them into the
``<tool_call>`` ``conversations`` JSONL that
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
loads directly.
Pipeline::
reasoning task -> N teacher rollouts -> verify -> keep cheapest passing
-> conversations JSONL
"""
from __future__ import annotations
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
__all__ = [
"generate_sft_dataset",
"trajectory_to_record",
]
@@ -0,0 +1,419 @@
"""Loaders for the orchestrator reasoning-SFT task sources.
Two HuggingFace reasoning datasets give us verifiable question/answer tasks for
the small orchestrator's cold-start SFT and the GRPO prompt pool:
- ``natolambert/GeneralThought-430K-filtered`` open reasoning traces scraped
from gr.inc. Each row carries a ``question``, a short ``reference_answer``
(the gold), a long ``model_answer`` (R1's full solution), and provenance in
``question_source`` / ``task``. There is **no** clean single "category" field;
we derive a coarse ``domain`` (math / medical / code / chat / misc) from
``question_source`` so the verifier can pick the right checker.
- ``open-thoughts/OpenThoughts3-1.2M`` OpenThoughts3 distillation set. Each
row has a ``domain`` in {code, math, science}, a ``source``, a ``difficulty``,
and a ``conversations`` list ``[{from: human, value}, {from: gpt, value}]``.
The human turn is the question; the gpt turn is a ``<think></think>`` trace
followed by the final solution, from which we extract the gold answer.
Mirrors the style of the sibling ``hotpotqa.py`` / ``toolscale.py`` loaders: a
plain dataclass plus loaders that accept a ``source=`` iterable override so the
normalization path is exercised offline with no network. Network imports of
``datasets`` are kept lazy (inside the function).
"""
from __future__ import annotations
import random
import re
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Iterator, List, Optional
GENERALTHOUGHT_ID = "natolambert/GeneralThought-430K-filtered"
OPENTHOUGHTS_ID = "open-thoughts/OpenThoughts3-1.2M"
@dataclass
class Task:
task_id: str
question: str
answer: str
domain: str # coarse area: math / code / science / medical / chat / misc
difficulty: str = "" # OpenThoughts difficulty tier (GeneralThought has none)
dataset: str = "" # source dataset: GeneralThought / OpenThoughts3
subsector: str = "" # fine source: NuminaMath / NHSQA / TACO / glaive / ...
@property
def instruction(self) -> str:
return self.question
# --- GeneralThought -----------------------------------------------------------
# question_source -> coarse domain. GeneralThought has no literal category field;
# the source string is the most reliable signal (NuminaMath is math, NHSQA is
# medical, TACO/glaive are code, oasst1/lmsys are open chat, everything else misc).
_SOURCE_DOMAIN = (
("numina", "math"),
("math", "math"),
("nhsqa", "medical"),
("medical", "medical"),
("medicine", "medical"),
("taco", "code"),
("code", "code"),
("glaive", "code"),
("oasst", "chat"),
("lmsys", "chat"),
("chat", "chat"),
)
def _generalthought_domain(row: Dict[str, Any]) -> str:
src = str(row.get("question_source") or "").lower()
task = str(row.get("task") or "").lower()
hay = f"{src} {task}"
for needle, dom in _SOURCE_DOMAIN:
if needle in hay:
return dom
return "misc"
def _normalize_generalthought(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
question = str(row.get("question") or "").strip()
# Prefer the short reference answer (exact-match-able); fall back to the
# long model_answer when no reference is present.
answer = str(row.get("reference_answer") or "").strip()
if not answer:
answer = str(row.get("model_answer") or "").strip()
if not question or not answer:
return None
domain = _generalthought_domain(row)
# Same unwinnable-gold guards as the OpenThoughts path — GeneralThought also
# carries NuminaMath proof problems whose "answer" is just a QED marker, and
# the math verifier has no judge to fall back on. (Audit 2026-07-12: these
# were slipping through because the filter only lived on the other loader.)
if _gold_is_unusable(answer, domain):
return None
task_id = str(row.get("question_id") or f"generalthought-{index}")
return Task(
task_id=task_id,
question=question,
answer=answer,
domain=domain,
difficulty=str(row.get("difficulty") or "").strip(), # usually absent for GT
dataset="GeneralThought",
subsector=str(row.get("question_source") or "").strip(),
)
# Domains whose questions are *lookup*-shaped, not *reasoning*-shaped ("What is
# idiopathic anaphylaxis?"). The orchestrator correctly answers these with a
# single web_search and never delegates to a model expert — so the trajectory is
# rejected by the clean gate ("NEVER routed to a model expert") and the rollout is
# wasted. Audit (2026-07-12) on the 213-row July train split:
#
# medical 41/57 never routed (72%) | math 0/48 (0%)
# chat 2/3 never routed (67%) | code 0/9 (0%)
#
# i.e. ~25% of every generation batch was spent on tasks that cannot teach routing.
# Excluded by default; pass ``exclude_domains=()`` to restore the old mix.
LOOKUP_DOMAINS = frozenset({"medical", "chat"})
def load_generalthought(
*,
n: int = 2000,
seed: int = 42,
source: Optional[Iterable[Dict[str, Any]]] = None,
buffer: Optional[int] = None,
exclude_domains: Iterable[str] = LOOKUP_DOMAINS,
) -> Iterator[Task]:
"""Yield up to ``n`` GeneralThought tasks, randomly mixed across categories.
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
We over-stream a buffer and shuffle it so the yielded tasks are a random mix
of the source's categories rather than a single leading shard. ``buffer``
overrides the shuffle-buffer size; pass a small value (e.g. for smoke runs)
to avoid streaming the default 6000-row floor.
``exclude_domains`` drops lookup-shaped questions (see ``LOOKUP_DOMAINS``);
the buffer is grown to compensate so we still yield ``n`` tasks.
"""
drop = {d.lower() for d in exclude_domains}
if source is None:
from datasets import load_dataset # lazy: optional dep / network
source = load_dataset(GENERALTHOUGHT_ID, split="train", streaming=True)
rng = random.Random(seed)
# Buffer a generous multiple of n so the shuffle actually mixes categories,
# but stay bounded for the multi-hundred-K streams.
buf_cap = buffer if buffer is not None else max(n * 6, 6000)
buf: List[Task] = []
seen = 0
for i, row in enumerate(source):
task = _normalize_generalthought(dict(row), index=i)
if task is None:
continue
seen += 1
if task.domain in drop:
continue
buf.append(task)
if len(buf) >= buf_cap:
break
# Excluded domains are a large slice of GeneralThought, so cap how far we
# read rather than streaming the whole shard hunting for survivors.
if seen >= buf_cap * 6:
break
rng.shuffle(buf)
for task in buf[:n]:
yield task
# --- OpenThoughts3 ------------------------------------------------------------
_BOXED_RE = re.compile(r"\\boxed\{")
def _extract_boxed(text: str) -> Optional[str]:
"""Return the contents of the last ``\\boxed{...}`` in ``text`` (brace-balanced)."""
last = None
for m in _BOXED_RE.finditer(text):
start = m.end()
depth = 1
i = start
while i < len(text) and depth > 0:
c = text[i]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
i += 1
if depth == 0:
last = text[start : i - 1].strip()
return last
def _conversation_qa(conversations: Any) -> Optional[tuple[str, str]]:
"""Pull (question, gold_answer) from an OpenThoughts ``conversations`` list."""
if not isinstance(conversations, list):
return None
question = ""
full = ""
for turn in conversations:
if not isinstance(turn, dict):
continue
who = str(turn.get("from") or "").lower()
val = str(turn.get("value") or "")
if who in ("human", "user") and not question:
question = val.strip()
elif who in ("gpt", "assistant"):
full = val
if not question or not full:
return None
# Drop the <think>…</think> trace; keep the post-reasoning solution.
visible = re.sub(r"(?s)<think>.*?</think>", "", full).strip()
# …and an UNCLOSED <think> (the source row was truncated mid-thought). The
# regex above needs a closing tag, so without this the strip silently no-ops
# and — when there's no \boxed{} either — `answer` falls through to the whole
# reasoning dump. That gold is unusable: math is verified by string/numeric
# match with no judge, so `_math_equal("19", "<think> Okay, so...")` is always
# False and the task becomes unwinnable no matter what the model answers.
if "<think>" in visible:
visible = visible.split("<think>", 1)[0].strip()
body = visible or full
answer = _extract_boxed(body) or _extract_boxed(full) or body.strip()
if not answer:
return None
return question, answer
# Sources that cannot produce a usable training example, dropped at load time so
# we never spend a rollout on them (audit 2026-07-12, 100-task Haiku batch):
#
# stackexchange_codegolf 0/13 correct (0%)
#
# Code-golf asks for the SHORTEST program, which is adversarial to both the model
# and any verifier — there's no canonical answer to match and no judge can score
# "is this minimal". By contrast normal code is fine (nvidia/OpenCodeReasoning
# scored 7/12 = 58%), so this drops the pathological source, not the domain.
UNUSABLE_SOURCES = ("codegolf",)
# A gold that is just a QED / proof marker — there's no value to match against.
_PROOF_GOLD_RE = re.compile(
r"^\s*\\?\(?\s*(\\blacksquare|\\qed|\\square|QED)\s*\\?\)?\s*$", re.IGNORECASE
)
def _gold_is_unusable(answer: str, domain: str) -> bool:
"""Would this task be unwinnable no matter what the model answers?
Skip it at LOAD time a rollout spent here is guaranteed waste.
* ``<think>`` in the gold means answer-extraction *failed* and we kept the raw
reasoning trace. Unwinnable in EVERY domain (audit 2026-07-13: 12/90 rollouts,
correct rate 2/12 = 17% vs 65-68% on a good gold). Note this is about
extraction failing, NOT about length: a *long* gold is fine, because the LLM
judge copes with it (code 67%, science 67%).
* A proof marker (``\\(\\blacksquare\\)``) is a QED symbol, not an answer.
* A long prose gold is unusable for MATH specifically, because math is verified
by string/numeric match with no judge fallback.
"""
if "<think>" in answer:
return True
if _PROOF_GOLD_RE.search(answer):
return True
if domain == "math" and len(answer) > 300:
return True
return False
def _normalize_openthoughts(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
source = str(row.get("source") or "").lower()
if any(s in source for s in UNUSABLE_SOURCES):
return None
qa = _conversation_qa(row.get("conversations"))
if qa is None:
return None
question, answer = qa
if not question or not answer:
return None
domain = str(row.get("domain") or "unknown").strip().lower() or "unknown"
if _gold_is_unusable(answer, domain):
return None
# PROOF tasks: the "answer" is a QED marker, not a value ("\(\blacksquare\)").
# There is nothing to match against — the real answer is the proof itself, and
# the math path has no judge. Unwinnable by construction; skip.
if _PROOF_GOLD_RE.search(answer):
return None
task_id = (
str(row.get("id") or row.get("source") or f"openthoughts-{index}") + f"-{index}"
)
return Task(
task_id=task_id,
question=question,
answer=answer,
domain=domain,
difficulty=str(row.get("difficulty") or "").strip(),
dataset="OpenThoughts3",
subsector=str(row.get("source") or "").strip(),
)
# The 1.2M set ships as 120 uniform parquet shards with no domain in the file
# names, but the rows are front-ordered by domain. Probing one row per shard puts
# code in shards ~0-30, math ~32-104, science ~105-119. Streaming the single
# combined split therefore has to read the entire code (+ math) prefix before it
# reaches math/science — minutes of wasted I/O. Instead we stream each domain from
# shards *inside* its own region. A few shards (~10K rows each) cover any quota.
_OT_SHARD_FMT = "data/train-{:05d}-of-00120.parquet"
_OT_DOMAIN_SHARDS: Dict[str, List[int]] = {
"code": [0, 1, 2, 3],
"math": [45, 46, 47, 48],
"science": [112, 113, 114, 115, 116, 117, 118, 119],
}
def load_openthoughts(
*,
n_code: int = 2000,
n_math: int = 2000,
n_science: int = 2000,
seed: int = 42,
source: Optional[Iterable[Dict[str, Any]]] = None,
) -> Iterator[Task]:
"""Yield OpenThoughts3 tasks balanced across code / math / science.
``source`` overrides the stream with raw row dicts (tests): rows are routed
into per-domain buffers, stopping once every quota is met. With no override we
stream each domain from its own shard region (see ``_OT_DOMAIN_SHARDS``) so a
balanced pull never has to read the entire front-ordered code/math prefix.
"""
quotas = {"code": n_code, "math": n_math, "science": n_science}
bufs: Dict[str, List[Task]] = {k: [] for k in quotas}
if source is not None:
for i, row in enumerate(source):
task = _normalize_openthoughts(dict(row), index=i)
if task is None:
continue
dom = task.domain
if dom in bufs and len(bufs[dom]) < quotas[dom]:
bufs[dom].append(task)
if all(len(bufs[k]) >= quotas[k] for k in quotas):
break
else:
from datasets import load_dataset # lazy: optional dep / network
idx = 0
for dom, quota in quotas.items():
if quota <= 0:
continue
files = [_OT_SHARD_FMT.format(s) for s in _OT_DOMAIN_SHARDS[dom]]
ds = load_dataset(
OPENTHOUGHTS_ID, data_files=files, split="train", streaming=True
)
for row in ds:
task = _normalize_openthoughts(dict(row), index=idx)
idx += 1
if task is None or task.domain != dom:
continue
bufs[dom].append(task)
if len(bufs[dom]) >= quota:
break
rng = random.Random(seed)
out: List[Task] = []
for k in quotas:
out.extend(bufs[k])
rng.shuffle(out)
for task in out:
yield task
# --- combined SFT / GRPO sets -------------------------------------------------
def load_sft_tasks(
*, seed: int = 42, cap: Optional[int] = None, balanced: bool = True
) -> List[Task]:
"""The 8K cold-start SFT set: 2K GeneralThought + 2K code + 2K math + 2K science.
``cap`` (smoke runs): when set, draw ~``cap`` tasks total.
- default (``balanced=False``): GeneralThought only with a small stream
buffer fastest, but the domain mix is whatever GeneralThought yields
(math/code/medical/chat), so a given sample can skew (e.g. all-medical).
- ``balanced=True``: ~cap/4 from each of GeneralThought + OpenThoughts
code/math/science, so the smoke is representative of the real run. Now
cheap because OpenThoughts streams from per-domain shards (no front-prefix).
"""
if cap is not None and cap > 0:
if balanced:
per = max(cap // 4, 1)
tasks = list(load_generalthought(n=per, seed=seed, buffer=max(per * 6, 64)))
tasks.extend(
load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed)
)
rng = random.Random(seed)
rng.shuffle(tasks)
return tasks[:cap]
buf = max(cap * 4, 64)
tasks = list(load_generalthought(n=cap, seed=seed, buffer=buf))
random.Random(seed).shuffle(tasks)
return tasks[:cap]
tasks: List[Task] = []
tasks.extend(load_generalthought(n=2000, seed=seed))
tasks.extend(load_openthoughts(n_code=2000, n_math=2000, n_science=2000, seed=seed))
rng = random.Random(seed)
rng.shuffle(tasks)
return tasks
__all__ = [
"GENERALTHOUGHT_ID",
"OPENTHOUGHTS_ID",
"Task",
"load_generalthought",
"load_openthoughts",
"load_sft_tasks",
]
@@ -0,0 +1,58 @@
"""Single source of truth for orchestrator dataset names.
raw/{name}-{stamp}/data.jsonl raw/qwen-july-7-2026-0553pm/data.jsonl
sft/{name}-{split}-{stamp}.jsonl sft/qwen-train-july-7-2026-0553pm.jsonl
The raw dir and every file carved from it share the same ``stamp``, so a curated
split always traces back to the generation run that produced it by eye.
The stamp is written out (``july-7-2026-0553pm``) rather than numeric because
these names are read by humans far more often than they are parsed. The cost is
that month names sort alphabetically, so ``ls`` is NOT chronological use
``ls -t``.
"""
from __future__ import annotations
import re
import time
from typing import Optional
__all__ = ["run_stamp", "dataset_name", "raw_dir_name", "stamp_from"]
# {month}-{day}-{year}-{hhmm}{am|pm}, e.g. july-7-2026-0553pm
_STAMP_RE = re.compile(r"[a-z]+-\d{1,2}-\d{4}-\d{4}(?:am|pm)", re.I)
def stamp_from(filename: str) -> Optional[str]:
"""Pull the stamp out of a raw dir / dataset name, or None if it carries none.
``qwen-clean-july-7-2026-0553pm.jsonl`` -> ``july-7-2026-0553pm``. Lets a split
inherit its POOL's stamp instead of stamping itself with today's date the
split belongs to the run that generated the data, not to the day it was carved.
"""
m = _STAMP_RE.search(filename)
return m.group(0).lower() if m else None
def run_stamp(when: Optional[time.struct_time] = None) -> str:
"""``july-7-2026-0553pm`` — month-in-words, day, year, 12h clock.
Local time, matching the wall-clock the run was launched at.
"""
t = when or time.localtime()
month = time.strftime("%B", t).lower() # july
day = t.tm_mday # no zero-pad: 7, not 07
clock = time.strftime("%I%M%p", t).lower() # 0553pm (zero-padded hour)
return f"{month}-{day}-{t.tm_year}-{clock}"
def raw_dir_name(name: str, stamp: Optional[str] = None, tag: str = "") -> str:
"""``qwen-july-7-2026-0553pm`` (+ ``-{tag}`` when disambiguating a variant)."""
base = f"{name}-{stamp or run_stamp()}"
return f"{base}-{tag}" if tag else base
def dataset_name(name: str, split: str, stamp: Optional[str] = None) -> str:
"""``qwen-train-july-7-2026-0553pm`` (no extension)."""
return f"{name}-{split}-{stamp or run_stamp()}"
@@ -0,0 +1,518 @@
"""Rejection-sampling SFT-data generator (the ToolOrchestra cold-start).
For each reasoning task: roll out a teacher orchestrator N times, verify each
trajectory, keep the passing ones (optionally just the cheapest), and serialize
them into the unified-tool ``conversations`` JSONL the SFT trainer consumes.
The expensive/network parts are injected so the orchestration is pure and
offline-testable:
* ``rollout_fn(task) -> UnifiedRollout`` one teacher rollout (temperature>0).
* ``verify_fn(task, rollout) -> bool`` did the trajectory solve the task?
(e.g. ``verify.make_verifier()``, an LLM judge over the final answer.)
"""
from __future__ import annotations
import json
import logging
import re
from collections import Counter
from pathlib import Path
from typing import Any, Callable, Iterable, List, Optional
from openjarvis.agents.hybrid.expert_registry import ExpertTool
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
_CONTROL_TOKEN_RE,
_strip_control_tokens,
trajectory_to_record,
)
logger = logging.getLogger(__name__)
# The sampler only touches ``task.task_id`` / ``task.instruction`` / ``task.domain``
# (via ``trajectory_to_record``), so any task dataclass works — the reasoning
# ``datasets.Task`` is what production uses, paired with its own ``verify_fn``
# (e.g. ``verify.make_verifier()``).
TaskLike = Any
RolloutFn = Callable[[TaskLike], Optional[UnifiedRollout]]
VerifyFn = Callable[[TaskLike, UnifiedRollout], bool]
# Human-readable blurb for each source dataset, stamped onto every record as
# ``dataset_description`` so a row is self-describing (what corpus the question
# came from) without a lookup. Keyed on the record's ``dataset`` value.
DATASET_DESCRIPTIONS = {
"GeneralThought": (
"natolambert/GeneralThought-430K-filtered — filtered reasoning Q&A from "
"gr.inc; each item has a question and a short reference (gold) answer "
"distilled from DeepSeek-R1 traces."
),
"OpenThoughts3": (
"open-thoughts/OpenThoughts3-1.2M — code/math/science reasoning tasks; "
"the gold answer is extracted from the boxed final solution."
),
}
def dataset_description(dataset: str) -> str:
"""Blurb for a source dataset name (empty string if unknown)."""
return DATASET_DESCRIPTIONS.get((dataset or "").strip(), "")
# Markers of a broken expert/observation that must never become a training
# target (empty turns, tracebacks, truncated tool errors, dead-provider strings).
_ERR_MARKERS = (
"Traceback (most recent call last)",
"insufficient_quota",
"rate limit",
"RateLimitError",
"[invalid tool",
"Error code: 4",
"Error code: 5",
"InternalServerError",
"ConnectionError",
"timed out",
# A tool call that reached dispatch with no/empty input, or a bridged-tool
# crash — the call was malformed. The trajectory routed garbage; drop it.
"no input provided",
"[openjarvis-tool error",
)
# A final answer cut off mid-expression ends on a dangling connective
# (``a=1, b=``, ``result = ``, ``compute (``) — the decode hit the token cap
# before finishing. A real distilled answer never ends this way.
_TRUNCATED_TAIL_RE = re.compile(r"[=,(\[{+\-*/]\s*$")
# A line that only source code starts with. Used to spot an UNFENCED program so the
# markdown/essay guards don't read its `#` comments as headers.
_CODE_SIGNAL_RE = re.compile(
r"(?m)^\s*(?:#!|def |class |import |from \s*\w+\s+import|function |const |let |"
r"var |public |private |#include|package |using |fn |func )"
)
def clean_reason(roll: UnifiedRollout) -> Optional[str]:
"""Why this trajectory is unfit to be an SFT *target* — ``None`` if it's fine.
Returns the FIRST failing check as a human-readable string so rejections are
auditable (persisted as ``clean_reason`` on every record). Previously this
returned a bare bool, which meant a rejected record gave no clue why and
the reason is genuinely not recoverable from the saved JSONL, because the
gate runs on the RAW rollout while the serializer scrubs the stored copy.
The judge decides semantic correctness; this catches the ~10-12% of records
whose targets are empty / a Traceback / a truncated tool error / unrouted
garbage the model must never be trained to imitate. Requirements:
* non-empty final answer, no error marker, balanced ``<think>`` (not
truncated mid-thought);
* at least one model-expert tool call (it actually *routed*);
* no error-marker / empty observation in the kept trajectory.
"""
fa = (roll.final_answer or "").strip()
if not fa:
return "empty final answer"
if any(m in fa for m in _ERR_MARKERS):
return "error marker in final answer"
# A thought that OPENS and never closes = the decode hit the cap mid-thought.
# A stray CLOSING </think> with no opener is NOT truncation — Qwen3.x's chat
# template pre-fills the opening `<think>` tag, so the model's completion
# begins inside the reasoning block and only ever emits the closing tag. The
# old check (`!=`) treated that template artifact as truncation and rejected
# essentially every thinking rollout; the serializer strips the stray tag
# anyway, so the stored target was fine. Only unclosed thoughts are bad.
if fa.count("<think>") > fa.count("</think>"):
return "unclosed <think> (truncated mid-thought)"
if _TRUNCATED_TAIL_RE.search(fa): # truncated mid-expression (e.g. "a=1, b=")
return "final truncated mid-expression"
# Malformed final: a proper final answer is plain text, NOT a stray/broken
# <tool_call> tag (the model's answer leaking inside a tool call that failed
# to parse). (\boxed{} is NOT rejected here — the serializer de-boxes it, so
# an otherwise-good boxed answer is salvaged rather than thrown away.)
if "<tool_call>" in fa or "</tool_call>" in fa:
return "tool_call tag leaked into final answer"
# Control-token backstop (defense in depth). The serializer strips leaked
# control/special tokens (<|im_end|>, <|tool_call>, <start_of_turn>, …) to
# salvage good answers, so we mirror that strip and gate on the RESULT:
# * strips to empty -> the answer WAS nothing but control tokens -> drop.
# * a token survives the strip -> something malformed the stripper can't
# safely delete mid-answer -> drop rather than train on it.
# A clean answer (or one the stripper fully salvages) sails through.
fa_stripped = _strip_control_tokens(fa)
if not fa_stripped:
return "final answer was only control tokens"
if _CONTROL_TOKEN_RE.search(fa_stripped):
return "control tokens in final answer"
# Degenerate repetition: the same substantive line emitted many times (the
# small-model decode loop). Reject rather than train on it.
_lines = [ln.strip() for ln in fa.split("\n") if len(ln.strip()) > 30]
if _lines and Counter(_lines).most_common(1)[0][1] >= 4:
return "degenerate repetition"
# Word-salad run-on: a giant unbroken line (no newline) is the other decode
# collapse — the model spraying novel tokens instead of repeating. Reject.
if any(len(ln) > 2000 for ln in fa.split("\n")):
return "word-salad run-on line"
# Essay-style final: the format wants a distilled answer, not a multi-section
# writeup. Check the answer AFTER the FINAL_ANSWER marker (reasoning before it
# is fine). NOTE: length/bold limits relaxed for the BALANCED/harder mix —
# code & multi-step math answers are legitimately longer and use **bold** for
# the key result, so the old 700-char + any-bold rejects were dropping ~half
# the correct hard-task answers. Keep the STRUCTURAL essay signals (many
# numbered sections, markdown headers, tables) which still catch real essays.
_marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", fa))
_ans = fa[_marks[-1].end() :].strip() if _marks else fa
# Every structural essay guard below runs on the PROSE only — a fenced code
# block is a legitimate answer, not an essay, and its contents are not
# markdown. Without this:
# * a Python comment ("# Use download_as_text()") matches the markdown-header
# regex `^\s*#{1,6}\s` and the whole answer is rejected as an "essay";
# * a numbered list in a docstring trips the numbered-sections guard;
# * an ASCII table in code output trips the markdown-table guard;
# * a real program blows the 2000-char limit on its own.
# Audit 2026-07-13: this was the difference between code correct=12/22 and code
# USABLE=3/22 — we were binning ~3 of every 4 correct code answers.
_prose = re.sub(r"```.*?```", "", _ans, flags=re.DOTALL)
_prose = re.sub(r"(?m)^ {4,}\S.*$", "", _prose) # indented code blocks too
# …and UNFENCED raw code. Haiku often answers a code task with a bare program
# and no ``` fence at all, and then its shebang (`#!/usr/bin/env python3`) and
# its comments (`# Read input`) trip the markdown-header regex exactly like the
# fenced case did. If what's left after stripping fences still reads as code,
# there is no prose to run essay guards against. (Audit 2026-07-13 — the fenced
# fix caught only half of this bug.)
if _CODE_SIGNAL_RE.search(_prose):
_prose = ""
# Markdown-table dump: the model wrote a formatted table instead of the short
# exact value the format demands. A ``|---|`` separator row is the tell.
if re.search(r"\|\s*:?-{2,}", _prose):
return "markdown table dump"
if len(_prose) > 2000: # prose-only: a long program is fine, an essay is not
return "final answer >2000 chars (essay)"
# Tool-status echo as the final answer: on code tasks the orchestrator
# sometimes ends with a file_write/shell status line ("Successfully wrote to
# /tmp/x.py") instead of the actual answer — a non-answer that slips the
# length/format checks. Reject. (Audit: ~code/math trajectories where the
# trace collapsed but the row was still scored correct.)
if re.search(
r"(?i)\b(successfully (wrote|created|saved|executed|ran)|written to /|"
r"file (written|saved|created)|no further actions?)\b",
_ans,
):
return "tool-status echo as the answer"
if (
len(re.findall(r"(?m)^\s*\d+\.\s", _prose)) >= 6
): # many numbered sections = essay (was 4)
return "many numbered sections (essay)"
if re.search(r"(?m)^\s*#{1,6}\s", _prose): # markdown headers = essay
return "markdown headers (essay)"
# Garbled / shouty final: a multi-word answer that's mostly UPPERCASE, or has
# an absurdly long merged all-letter token, is decode garble — the audit found
# e.g. "RABES PEPTETANUS BOOSTERS" (for "rabies PEP, tetanus boosters") passing.
if len(_ans) > 20:
# Measure the caps ratio over PROSE tokens only. A chemistry answer is
# legitimately uppercase-heavy — "HC≡CH → NaNH₂ → CH₃CH₂CH₂Br" is a
# reaction scheme, not shouting — and the old whole-string ratio rejected
# those outright (audit 2026-07-12), which is expensive because organic
# chemistry is one of our highest-yield domains. So ignore any token that
# carries a digit, a subscript/superscript, or a non-ASCII symbol (arrows,
# bond glyphs): that's formula notation, not decode collapse. What's left
# is real words, which is what the guard was built for ("RABES PEPTETANUS
# BOOSTERS").
_words = [
w
for w in _ans.split()
if w.isalpha() and len(w) >= 3 and w.isascii()
]
_alpha = [c for w in _words for c in w]
if (
len(_words) >= 3
and _alpha
and sum(c.isupper() for c in _alpha) / len(_alpha) > 0.7
):
return "ALL-CAPS garble"
if any(len(w) > 25 and w.isalpha() for w in _ans.split()):
return "merged mega-token garble"
# Reject runaway final answers — a distilled answer, not a multi-KB essay or
# a verbatim dump of an expert observation. (Audit: 279 finals >4k chars, the
# worst a 409k-char wordlist; ~363 were prefix-identical to a tool obs.)
if len(fa) > 8000:
return "runaway final (>8k chars)"
fa_head = re.sub(r"\s+", " ", fa[:200]).strip()
for t in roll.turns:
if t.observation:
obs_head = re.sub(r"\s+", " ", t.observation[:200]).strip()
if (
fa_head and fa_head == obs_head and len(fa) > 200
): # long final = copied tool dump (short relays are legit)
return "final answer copied verbatim from a tool result"
# "routed" = delegated to a model EXPERT (not just a utility like web_search).
# When anonymized, expert calls are the anon labels in anon_map; otherwise
# fall back to "made any tool call".
expert_names = set(roll.anon_map or {})
called = {name for name, _ in roll.tool_calls()}
if expert_names:
if not (called & expert_names):
return "NEVER routed to a model expert"
elif roll.num_tool_calls < 1:
return "no tool calls at all"
# A broken EXPERT observation (rate limit / 5xx / dead provider) means the
# trajectory routed into a hole and whatever follows is built on nothing —
# reject it.
#
# A sandbox tool is different. `code_interpreter` returning a Python
# "Traceback (most recent call last)" is the tool WORKING: the model ran code,
# it raised, and the model reads the error and fixes it. That error-recovery
# loop is precisely what we want to teach. Treating it as a broken tool threw
# out 24% of rollouts (audit 2026-07-12: 21 of 23 flagged observations were
# code_interpreter tracebacks the model then recovered from). Only an EMPTY
# observation is fatal for a sandbox tool — that means the tool itself died.
for t in roll.turns:
if t.tool_name is None:
continue
obs = (t.observation or "").strip()
if not obs:
return "empty tool observation"
is_expert = t.tool_name in expert_names if expert_names else False
if is_expert and any(m in obs for m in _ERR_MARKERS):
return "error observation from a model expert"
# Reasoning-side decode collapse: an intermediate turn whose reasoning is a
# giant unbroken line, or is dominated by exotic unicode / emoji spam, is
# garbage even when the final answer looks fine (the answer-only guards above
# miss it). The audit found a trajectory with a thousands-char symbol blob in
# its reasoning that scored correct and slipped through.
for t in roll.turns:
r = t.reasoning or ""
if any(len(ln) > 2000 for ln in r.split("\n")):
return "reasoning decode collapse (giant line)"
if len(r) > 200:
weird = sum(1 for ch in r if ord(ch) > 0x2100 and not ch.isalnum())
if weird / len(r) > 0.10: # >10% exotic-unicode/emoji = decode spam
return "reasoning unicode/emoji spam"
return None
def _target_is_clean(roll: UnifiedRollout) -> bool:
"""Back-compat bool wrapper around :func:`clean_reason`."""
return clean_reason(roll) is None
def _process_task(
task: TaskLike,
*,
rollout_fn: RolloutFn,
verify_fn: VerifyFn,
samples_per_task: int,
max_keep_per_task: int,
stop_at_keep: bool = False,
keep_all: bool = False,
) -> List[tuple[UnifiedRollout, bool]]:
"""One task's rejection-sampling unit of work: roll out ``samples_per_task``
times and verify each. Returns ``(rollout, is_correct)`` for **every** sample
(the caller decides what to write). Runs in a worker thread, so the only
shared state it touches is the injected fns (each rollout builds its own
OpenAI client; the tool executor is built under a lock).
``stop_at_keep``: short-circuit as soon as ``max_keep_per_task`` passing
trajectories are found, instead of always exhausting ``samples_per_task``.
Big throughput win (no wasted rollouts on already-solved tasks), but it
trades away the cheapest-of-N cost optimisation (we keep the first passers,
not the cheapest). Ignored when ``keep_all`` is set (we want every sample).
Default off preserves the original semantics."""
results: List[tuple[UnifiedRollout, bool]] = []
n_correct = 0
for _ in range(samples_per_task):
roll = rollout_fn(task)
if roll is None:
continue
ok = bool(verify_fn(task, roll))
results.append((roll, ok))
if ok:
n_correct += 1
if stop_at_keep and not keep_all and n_correct >= max_keep_per_task:
break
return results
def generate_sft_dataset(
out_path: str,
*,
tasks: Iterable[TaskLike],
tools: List[ExpertTool],
rollout_fn: RolloutFn,
verify_fn: VerifyFn,
samples_per_task: int = 4,
max_keep_per_task: int = 1,
reward_fn: Optional[Callable[[UnifiedRollout], float]] = None,
concurrency: int = 1,
stop_at_keep: bool = False,
keep_all: bool = False,
record_extra: Optional[dict] = None,
) -> dict:
"""Run rejection sampling over ``tasks`` and write the SFT JSONL.
``max_keep_per_task`` caps records kept per task; when >1 the cheapest
passing trajectories are kept first. ``concurrency`` rolls out that many tasks
in parallel (each task issues its samples sequentially, so ~``concurrency``
requests hit the served model at once) set 1 for the original sequential
behaviour. Records are written as each task finishes, so peak memory is bounded
by the in-flight tasks, not the whole dataset. Returns stats + writes a
``.stats.json``.
``keep_all``: write **every** rolled-out trajectory (correct and incorrect)
rather than dropping the failures. Each record gets ``correct`` (verifier
verdict) and ``kept`` (True on the cheapest-correct sample the one the
rejection sampler would have selected). This preserves the full sample set so
you can compute accuracy / inspect failures; the SFT trainer should filter on
``correct`` (or ``kept``). Default off = original drop-the-failures behaviour.
"""
out = Path(out_path)
out.parent.mkdir(parents=True, exist_ok=True)
tasks = list(tasks)
seen = len(tasks)
written = 0
correct_written = 0
incorrect_written = 0
dropped = 0 # tasks that produced no kept record
tasks_solved = 0 # tasks with >=1 correct sample
domain_counts: Counter[str] = Counter()
def _work(task: TaskLike) -> tuple[TaskLike, List[tuple[UnifiedRollout, bool]]]:
return task, _process_task(
task,
rollout_fn=rollout_fn,
verify_fn=verify_fn,
samples_per_task=samples_per_task,
max_keep_per_task=max_keep_per_task,
stop_at_keep=stop_at_keep,
keep_all=keep_all,
)
def _emit(fh, task: TaskLike, roll: UnifiedRollout, ok: bool, kept: bool) -> None:
nonlocal written, correct_written, incorrect_written
reward = reward_fn(roll) if reward_fn else 0.0
record = trajectory_to_record(
task.task_id,
task.instruction,
tools,
roll,
reward=reward,
domain=task.domain,
)
record["correct"] = ok
record["kept"] = kept
# Persist WHY a trajectory was rejected. The gate runs on the raw rollout
# while the serializer scrubs the stored conversations, so the reason is
# not recoverable from the saved record afterwards — capture it here.
why = clean_reason(roll)
record["clean"] = why is None
record["clean_reason"] = why or ""
# Task provenance so every record is self-describing: area (domain),
# difficulty tier, source dataset, and fine subsector. Lets us slice
# train/holdout/analysis by any of these later.
record["area"] = task.domain
record["difficulty"] = getattr(task, "difficulty", "") or ""
record["dataset"] = getattr(task, "dataset", "") or ""
record["dataset_description"] = dataset_description(record["dataset"])
record["subsector"] = getattr(task, "subsector", "") or ""
# Gold reference answer the verifier graded against. Persisted so every
# record supports gold-vs-model inspection downstream (Braintrust, error
# analysis) instead of only a bare `correct` boolean.
record["gold_answer"] = getattr(task, "answer", "") or ""
# Stamp provenance (which model/orchestrator generated this trajectory)
# so records are self-identifying when pooled across model families.
if record_extra:
record.update(record_extra)
fh.write(json.dumps(record) + "\n")
fh.flush()
written += 1
correct_written += int(ok)
incorrect_written += int(not ok)
domain_counts[task.domain] += 1
def _write(fh, task: TaskLike, results: List[tuple[UnifiedRollout, bool]]) -> None:
nonlocal dropped, tasks_solved
# A trajectory is only eligible to be *kept* as a training target if the
# judge passed it AND it's structurally clean (non-error, routed, not
# truncated). Unclean-but-correct rollouts are still written in keep_all
# mode (for analysis) but never marked kept.
correct = [r for r in results if r[1] and _target_is_clean(r[0])]
cheapest = min((r for r, _ in correct), key=lambda r: r.cost_usd, default=None)
if correct:
tasks_solved += 1
if keep_all:
# Write every sample; mark the cheapest-correct one as kept.
if not results:
dropped += 1
return
for roll, ok in results:
_emit(fh, task, roll, ok, kept=(roll is cheapest))
else:
# Original behaviour: keep only cheapest-correct, capped.
if not correct:
dropped += 1
return
for roll in sorted((r for r, _ in correct), key=lambda r: r.cost_usd)[
:max_keep_per_task
]:
_emit(fh, task, roll, True, kept=(roll is cheapest))
with out.open("w") as fh:
if concurrency <= 1:
for task in tasks:
_, results = _work(task)
_write(fh, task, results)
else:
from concurrent.futures import ThreadPoolExecutor, as_completed
done = 0
with ThreadPoolExecutor(max_workers=concurrency) as pool:
futures = {pool.submit(_work, t): t for t in tasks}
for fut in as_completed(futures):
task = futures[fut]
try:
_, results = fut.result()
except Exception as exc: # a task's worker died; count + skip
logger.warning(
"task %s failed: %s", getattr(task, "task_id", "?"), exc
)
results = []
_write(fh, task, results)
done += 1
if done % 50 == 0 or done == seen:
logger.info(
"rejection-sampling: %d/%d tasks done "
"(%d written, %d solved, %d dropped)",
done,
seen,
written,
tasks_solved,
dropped,
)
stats = {
"out_path": str(out),
"tasks_seen": seen,
"records_written": written,
"records_correct": correct_written,
"records_incorrect": incorrect_written,
"tasks_solved": tasks_solved,
"tasks_dropped": dropped,
"task_accuracy": round(tasks_solved / seen, 4) if seen else 0.0,
"samples_per_task": samples_per_task,
"keep_all": keep_all,
"concurrency": concurrency,
"domain_distribution": dict(domain_counts),
}
out.with_suffix(out.suffix + ".stats.json").write_text(json.dumps(stats, indent=2))
logger.info(
"Wrote %d SFT records to %s (%d correct, %d incorrect)",
written,
out,
correct_written,
incorrect_written,
)
return stats
__all__ = ["generate_sft_dataset"]
@@ -0,0 +1,279 @@
"""Serialize a verified unified-tool rollout into an SFT ``conversations`` record.
Output matches what ``OrchestratorSFTDataset`` consumes, and trains the model to
emit the ``<tool_call>{...}</tool_call>`` text form that
``toolorchestra.parsing._parse_rl_tool_call`` already reads back. One record =
one passing trajectory.
Roles: ``system`` (the unified tool catalog), ``user`` (the running ``Problem``
prompt), ``assistant`` (reasoning + a ``<tool_call>`` tag, or the final answer),
``tool`` (the executed observation).
"""
from __future__ import annotations
import re
from typing import Any, Dict, List
from openjarvis.agents.hybrid.expert_registry import ExpertTool, build_tool_specs
from openjarvis.agents.hybrid.toolorchestra.rollout import (
_OBS_CAP as _ROLLOUT_OBS_CAP,
UnifiedRollout,
build_system_prompt,
tool_call_tag,
)
def _system_prompt(tools: List[ExpertTool], rollout: UnifiedRollout) -> str:
# Faithful ToolOrchestra system prompt (paper's prepare_sft_data.py format).
# Prefer the EXACT specs the policy saw this rollout (anonymized labels when
# anonymize=True) so the saved <tools> block matches the assistant's
# <tool_call> tags. Falling back to the real registry here is the bug that
# leaked real model names+pricing into anonymized records.
specs = rollout.tool_specs if rollout.tool_specs else build_tool_specs(tools)
return build_system_prompt(specs)
def _normalize_think(text: str) -> str:
"""Ensure a reasoning block has matched tags. Qwen rollouts routinely yield a
dangling ``</think>`` with NO opening ``<think>`` (the chat template consumes
the opener in the prompt, so only the closer survives in the completion). Add
the opener back so the trained turn is well-formed."""
t = (text or "").lstrip()
if "</think>" in t and "<think>" not in t:
t = "<think>\n" + t
return t
def _debox(text: str) -> str:
"""Unwrap every ``\\boxed{X}`` -> ``X`` (balanced braces). The format kills
\\boxed everywhere, but small models keep emitting it in otherwise-correct
answers de-box to salvage them rather than reject the whole trajectory."""
out, i = [], 0
marker = r"\boxed{"
while i < len(text):
j = text.find(marker, i)
if j == -1:
out.append(text[i:])
break
out.append(text[i:j])
k = j + len(marker)
depth, inner = 1, []
while k < len(text) and depth:
c = text[k]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
if depth == 0:
break
inner.append(c)
k += 1
out.append("".join(inner))
i = k + 1
return "".join(out)
# Reasoning that PAROTS the training constraint back ("I must call a model per
# my instructions", "this is testing whether I need to route") — the model
# reasoning about its own harness. Never a good target: drop the whole line.
_LEAK_RE = re.compile(
r"(?im)^.*\b(?:"
r"the user (?:is|was|wants|has|had|needs|'s|would|will|asked"
r"|'ve asked|has asked)\b|"
r"asked me to\b|" # echoes of an injected nudge turn
r"make progress by delegat|"
r"delegat\w+[^\n]*sub-?question|"
r"(?:to |ask |send .{0,20}to )?one of the models\b|"
r"instructions?|" # any self-reference to the rules
r"requirement|"
r"at least one (?:model|tool|expert)|" # "(call|invoke) at least one model"
r"(?:call|invoke|route to|delegate to|use)\s+(?:a|one|at least one|the)\s+model\b|"
r"route\s+(?:this|it|that)\b[^\n]{0,40}\bthrough\s+(?:a|one|the)\s+model\b|"
r"without mentioning (?:any )?(?:tools?|reasoning|meta|steps)|"
r"in the required format\b|"
r"based on the (?:model|response|analysis|expert)|" # response-acknowledgment
r"the model (?:provided|confirmed|gave|returned|correctly|analy\w+"
r"|respon\w+|said|indicated|identified)|"
r"i (?:got|received|have) (?:a |the )?"
r"(?:comprehensive |clear |good |detailed )?answer\b|"
r"now i can (?:confidently )?(?:give|provide|state|answer)|"
r"perfect[!,]|great[!,]|excellent[!,]|" # narration exclamations
r"as an orchestrator|"
r"whose job is to (?:route|delegate|orchestrate)|"
r"testing whether i (?:need|have|should|must)\b|"
r"i(?:'m| am)? (?:required|instructed|supposed|expected) to\b"
r")[^\n]*(?:\n|$)"
)
# A leading task-restatement opener ("The user is asking…") — the format
# explicitly bans this meta-text; drop it when it opens the reasoning.
_META_OPEN_RE = re.compile(
r"(?is)^\s*(?:the user (?:is|was|wants|has|needs|would|'s)\b[^.\n]*[.\n]+\s*)+"
)
def _scrub_meta(text: str) -> str:
"""Remove harness-leakage lines and a leading 'The user is asking…' meta
opener from a reasoning block. Both are disallowed by the system prompt, so
they must not survive into the supervised target. Preserves a leading
``<think>`` marker so the block stays well-formed."""
if not text:
return text
think = ""
body = text
m = re.match(r"(?is)^\s*(<think>)\s*", body)
if m:
think = "<think>\n"
body = body[m.end() :]
body = _LEAK_RE.sub("", body)
body = _META_OPEN_RE.sub("", body).lstrip()
return (think + body).strip() if think else body.strip()
# Control / special tokens that must never survive into a training target's
# final answer. Three shapes, in order of alternation:
# 1. closed pipe token — ``<|im_end|>``, ``<|im_start|>``, ``<|eot_id|>``,
# ``<|end_of_turn|>``, ``<|"|>`` (the stray-quote leak).
# 2. UNCLOSED pipe token — ``<|tool_call>`` (opened with ``<|`` but closed on a
# bare ``>`` because the decode was cut before the second pipe).
# 3. named angle special tokens — gemma ``<start_of_turn>`` / ``<end_of_turn>``
# and the sentinel set ``<eos> <bos> <pad> <unk> <s> </s>``.
# Deliberately narrow: only the ``<|...|>`` pipe form and this known name list
# match, so legitimate ``<``/``>`` in math/code (``x < 3 and y > 2``) is left
# untouched.
_CONTROL_TOKEN_RE = re.compile(
r"<\|[^>]*?\|>" # closed pipe token
r"|<\|[^|>]*>" # unclosed pipe token
r"|</?(?:start_of_turn|end_of_turn|eos|bos|pad|unk|s)>" # named angle tokens
)
def _strip_control_tokens(text: str) -> str:
"""Delete residual control/special tokens from an answer so a good answer
with a stray token is salvaged rather than dropped. Loops (bounded) because
removing one match can expose a nested/overlapping leftover (e.g.
``<|<|im_end|>|>``). Re-strips whitespace at the end."""
if not text:
return text
out = text
for _ in range(4):
stripped = _CONTROL_TOKEN_RE.sub("", out)
if stripped == out:
break
out = stripped
return out.strip()
def _final_answer_block(text: str) -> str:
"""Render the final-answer assistant message with a single clean
``FINAL_ANSWER: <value>`` line. De-boxes the answer, strips leaked control
tokens, and normalizes whatever spelling the model used (``FINALANSWER``,
``FINAL ANSWER``, no colon, ) to the exact tag. Keeps the think block at
most once."""
text = _debox(_normalize_think((text or "").strip()))
# Normalize any final-answer marker spelling to the canonical form; take the
# LAST one as the real answer (idempotent — never emits two tags).
marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", text))
if marks:
last = marks[-1]
answer = _strip_control_tokens(text[last.end() :].strip())
# Final turn = the bare short answer only. Drop EVERYTHING before
# FINAL_ANSWER — both the visible "Perfect! Based on the model's response…"
# narration AND the final-turn <think>, which is just confirmatory
# "both models agree… I've verified…" fluff. The routing reasoning already
# lives in the earlier tool-call turns, so nothing of value is lost, and
# this kills all final-turn narration deterministically (any wording).
return f"FINAL_ANSWER: {answer}"
if "</think>" in text:
# No explicit marker: the answer is whatever follows the final </think>.
# Drop the think block (same rationale as above) — keep only the answer.
_, _, answer = text.rpartition("</think>")
answer = _strip_control_tokens(answer.strip())
return (
f"FINAL_ANSWER: {answer}"
if answer
else f"FINAL_ANSWER: {_strip_control_tokens(_scrub_meta(text))}"
)
return f"FINAL_ANSWER: {_strip_control_tokens(text)}"
def trajectory_to_record(
task_id: str,
question: str,
tools: List[ExpertTool],
rollout: UnifiedRollout,
*,
reward: float = 0.0,
domain: str = "unknown",
) -> Dict[str, Any]:
"""Convert a passing :class:`UnifiedRollout` into one SFT JSONL record."""
conversations: List[Dict[str, str]] = [
{"role": "system", "content": _system_prompt(tools, rollout)},
{"role": "user", "content": f"Problem: {question}"},
]
for turn in rollout.turns:
if turn.tool_name is None:
# Final-answer turn. turn.reasoning is the model's actual output for
# this turn (which already contains the answer); render it once.
conversations.append(
{
"role": "assistant",
"content": _final_answer_block(
turn.reasoning or rollout.final_answer
),
}
)
continue
tag = tool_call_tag(turn.tool_name, turn.arguments)
reasoning = _scrub_meta(_normalize_think((turn.reasoning or "").rstrip()))
conversations.append(
{
"role": "assistant",
"content": (reasoning + "\n" + tag).strip(),
}
)
# Store the observation EXACTLY as the orchestrator saw it — the rollout
# caps tool output at _OBS_CAP before the model reads it (rollout.py), but
# this serializer used to store the FULL text. That trained the student on
# context the teacher never had, and produced monster rows (a 212k-char
# raw-HTML http_request dump = ~50k tokens of <!DOCTYPE> boilerplate) that
# then got BEHEADED by the trainer's max-seq. Train on what the policy saw.
obs = turn.observation or ""
if len(obs) > _ROLLOUT_OBS_CAP:
obs = obs[:_ROLLOUT_OBS_CAP] + "\n…[truncated]"
conversations.append(
{
"role": "tool",
"name": turn.tool_name,
"content": obs,
}
)
# If the rollout terminated on max_turns (no None turn), append the answer.
if not rollout.turns or rollout.turns[-1].tool_name is not None:
conversations.append(
{
"role": "assistant",
"content": _final_answer_block(rollout.final_answer),
}
)
return {
"conversations": conversations,
"task_id": task_id,
"domain": domain,
"reward": reward,
"metrics": {
"cost_usd": rollout.cost_usd,
"tokens": rollout.tokens,
"num_tool_calls": rollout.num_tool_calls,
"num_turns": len(rollout.turns),
# Present only when experts were anonymized: opaque label -> real
# tool name, so analysis can recover which model was actually picked.
**({"anon_map": rollout.anon_map} if rollout.anon_map else {}),
},
}
__all__ = ["trajectory_to_record"]
@@ -0,0 +1,388 @@
"""Correctness verifier for orchestrator reasoning tasks (:mod:`.datasets`).
The verifier dispatches on :attr:`Task.domain`:
- **math** normalize and compare. Extract ``\\boxed{...}`` from both sides,
try ``sympy`` for symbolic / numeric equality, and fall back to a normalized
string / number match. No network.
- **code** best-effort: if a short expected output / answer exists, do a
normalized substring match; otherwise defer to the LLM judge. (We don't run
arbitrary code here.)
- **science / medical / chat / misc / unknown** Gemini LLM judge given the
question + gold answer + candidate, asked for ``PASS`` / ``FAIL``. When no
``GEMINI_API_KEY`` is set (e.g. offline tests) it falls back to a normalized
string / token-F1 >= 0.6 match.
The OpenAI key is dead, so the LLM judge talks to Gemini through its
OpenAI-compatible endpoint. Normalization helpers are copied (not imported)
from ``hotpotqa.py`` to keep this module self-contained.
"""
from __future__ import annotations
import logging
import os
import re
import string
from typing import Any, Callable, Dict, Optional
from .datasets import Task
# LLM judge model. Switched from Gemini (free-tier rate limit stalled parallel gen)
# to Anthropic Haiku: fast (~0.6s), direct PASS/FAIL, high rate limits.
LOGGER = logging.getLogger(__name__)
JUDGE_MODEL = "claude-haiku-4-5-20251001"
# The Anthropic SDK's own retry does exponential backoff and honours Retry-After.
# A judge call that waits a few seconds beats a silently-wrong label.
JUDGE_MAX_RETRIES = 6
# Reused across calls — creating a fresh Anthropic client per judge call leaks an
# httpx connection pool (CLOSE-WAIT pileup under parallel generation). Thread-safe.
_JUDGE_CLIENT = None
# Every swallowed judge error, so a run can never again silently mislabel while
# reporting a clean bill of health.
_JUDGE_FAILURES: list[str] = []
# --- normalization (copied from hotpotqa.py — keep self-contained) -----------
def _normalize_answer(s: str) -> str:
"""Lowercase, strip punctuation / articles / extra whitespace (SQuAD/HotpotQA)."""
s = s.lower()
s = "".join(ch for ch in s if ch not in set(string.punctuation))
s = re.sub(r"\b(a|an|the)\b", " ", s)
return " ".join(s.split())
def _f1(pred: str, gold: str) -> float:
pt, gt = _normalize_answer(pred).split(), _normalize_answer(gold).split()
if not pt or not gt:
return float(pt == gt)
common: Dict[str, int] = {}
for w in pt:
if w in gt:
common[w] = common.get(w, 0) + 1
num_same = sum(common.values())
if num_same == 0:
return 0.0
precision = num_same / len(pt)
recall = num_same / len(gt)
return 2 * precision * recall / (precision + recall)
def _string_or_f1(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool:
"""Normalized whole-word substring OR token-F1 >= threshold."""
np_, ng = _normalize_answer(prediction), _normalize_answer(gold)
if not ng:
return False
if f" {ng} " in f" {np_} ":
return True
return _f1(prediction, gold) >= f1_threshold
# --- math --------------------------------------------------------------------
_BOXED_RE = re.compile(r"\\boxed\{")
_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?")
def _extract_boxed(text: str) -> Optional[str]:
"""Return the contents of the last ``\\boxed{...}`` (brace-balanced)."""
last = None
for m in _BOXED_RE.finditer(text):
start = m.end()
depth = 1
i = start
while i < len(text) and depth > 0:
c = text[i]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
i += 1
if depth == 0:
last = text[start : i - 1].strip()
return last
# \frac{a}{b} / \dfrac{a}{b} / \tfrac{a}{b} -> (a)/(b)
_FRAC_RE = re.compile(r"\\[dt]?frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}")
def _clean_math(s: str) -> str:
s = s.strip()
boxed = _extract_boxed(s)
if boxed is not None:
s = boxed
# Common LaTeX wrappers / delimiters.
s = s.replace("$", "").replace("\\!", "").replace("\\,", "").replace("\\;", "")
s = s.replace("\\left", "").replace("\\right", "")
# Normalize LaTeX fractions to plain division. Without this, `\frac{15}{64}`
# never becomes `15/64`, the exact/numeric comparisons both fail, and we fall
# through to a crude "grab the first number" match — which reads `15` out of
# BOTH sides and returns True. That is a FALSE POSITIVE: it passed
# gold \frac{15}{64} vs model 15/99 (different fraction)
# gold \frac{1}{2} vs model 1/3 (different denominator)
# gold \frac{15}{64} vs model 15 (just the numerator)
# i.e. any wrong answer sharing a numerator with the gold was marked CORRECT
# and fed into training. It also caused the mirror false-negative
# (`-\dfrac{73}{143}` vs `-73/143`), because the minus sits outside the frac
# and the number-grab loses it. (Audit 2026-07-13.)
for _ in range(3): # nested fractions
s, n = _FRAC_RE.subn(r"(\1)/(\2)", s)
if not n:
break
s = s.strip().strip("$ ")
return s
# A pure rational expression: -73/143, (15)/(64), 3.5/2 — safe to evaluate.
_RATIONAL_RE = re.compile(r"^-?\(?\s*-?\d+(?:\.\d+)?\s*\)?\s*/\s*\(?\s*-?\d+(?:\.\d+)?\s*\)?$")
def _to_float(s: str) -> Optional[float]:
s = s.replace(",", "").strip()
try:
return float(s)
except (ValueError, TypeError):
pass
# Evaluate a plain fraction rather than falling through to the number-grab
# below, which would read only the NUMERATOR and happily equate 15/64 with
# 15/99. Restricted by regex to digits and one '/', so nothing else is eval'd.
if _RATIONAL_RE.match(s):
try:
num, den = s.replace("(", "").replace(")", "").split("/")
d = float(den)
if d != 0:
return float(num) / d
except (ValueError, ZeroDivisionError):
pass
m = _NUM_RE.search(s)
if m:
try:
return float(m.group(0))
except ValueError:
return None
return None
def _math_equal(prediction: str, gold: str) -> bool:
pred = _clean_math(prediction)
g = _clean_math(gold)
if not g:
return False
# Exact normalized-string shortcut.
if pred.replace(" ", "") == g.replace(" ", "") and pred:
return True
# Numeric comparison (handles "42" vs "42.0" vs trailing prose).
pf, gf = _to_float(pred), _to_float(g)
if pf is not None and gf is not None:
if abs(pf - gf) <= 1e-6 * max(1.0, abs(gf)):
return True
# NOTE: a sympy/parse_latex symbolic-equality block used to live here, but
# sympy.simplify / antlr parse_latex can hang on pathological \boxed{} answers
# while holding the GIL -> deadlocks the whole threaded rejection sampler
# (all worker threads freeze in futex_wait, server goes idle, 0 records).
# Symbolic equivalence is now deferred to the Gemini judge in verify_answer().
# Last resort: whole-word substring of the gold in the prediction.
return _string_or_f1(prediction, g, f1_threshold=0.9)
# --- LLM judge (Gemini, OpenAI-compatible) -----------------------------------
def _judge_route() -> Optional[tuple]:
"""(base_url, api_key, model) for the judge — OpenRouter if configured.
The judge must NOT share a rate-limit bucket with the orchestrator. They're
the same model (Haiku), but the orchestrator makes 5-10 calls per rollout and
the judge makes 1; when they contend, the judge 429s, returns None, and the
caller falls back to string matching silently marking correct answers WRONG.
Set ``OJ_JUDGE_VIA_OPENROUTER=1`` to give the judge its own quota (~$2.50 per
3000 rollouts).
"""
if os.environ.get("OJ_JUDGE_VIA_OPENROUTER") == "1":
key = os.environ.get("OPENROUTER_API_KEY")
if key:
return ("https://openrouter.ai/api/v1", key, "anthropic/claude-haiku-4.5")
key = os.environ.get("ANTHROPIC_API_KEY")
return ("https://api.anthropic.com/v1/", key, JUDGE_MODEL) if key else None
def _gemini_judge(task: Task, prediction: str) -> Optional[bool]:
"""Ask the LLM judge PASS/FAIL. Returns None if unavailable (caller falls back)."""
route = _judge_route()
if route is None:
return None
base_url, api_key, model = route
try:
from openai import OpenAI # OpenAI-compatible: works for both routes
except Exception:
return None
try:
# Judge = Haiku (fast ~0.6s, direct PASS/FAIL), reached over whichever route
# _judge_route() picked.
#
# This used to be fail-fast (max_retries=0, timeout=15) so a slow judge
# couldn't stall the threaded sampler. That was actively destroying data:
# under parallel generation the judge 429s, this returns None, the caller
# falls back to string/f1, and a CORRECT answer gets marked WRONG. The
# failure is invisible — the exception is swallowed and never logged, so
# the run reports "0 rate limits" while silently mislabelling.
# (Audit 2026-07-12: 3 consecutive judge calls at concurrency 100 -> 2x429.)
#
# Retries alone were NOT enough: with the judge on the same bucket as the
# orchestrator, the retries just parked every worker thread in backoff and
# stalled the run. The real fix is giving the judge its own quota
# (OJ_JUDGE_VIA_OPENROUTER=1); retries are the belt to that suspenders.
global _JUDGE_CLIENT
if _JUDGE_CLIENT is None:
_JUDGE_CLIENT = OpenAI(
base_url=base_url,
api_key=api_key,
max_retries=JUDGE_MAX_RETRIES,
timeout=30,
)
client = _JUDGE_CLIENT
# Grade on SEMANTIC equivalence, not surface form. The bare "correct and
# consistent with the gold" instruction was far too strict: an audit
# (2026-07-12) found ~half of all FAIL verdicts were correct answers the
# judge rejected purely for phrasing —
# gold "0" vs "Order w.r.t. A = 1; Order w.r.t. B = 0"
# gold "use plt.figtext()" vs "Use fig.text()" (same matplotlib call)
# gold "<terse fragment>" vs the same fact in a full sentence
# Those FAILs silently threw away good training data. The guardrails at the
# bottom keep it from swinging lenient: a different VALUE still fails.
prompt = (
"You are grading a candidate answer against a gold reference.\n\n"
"The gold is often TERSE — a single word ('Yes'), a bare value ('0'), a "
"compact expression ('4 > 3 > 1 > 2'). The candidate is typically a full "
"sentence that states the same thing and adds correct supporting detail. "
"That is a PASS.\n\n"
"Your ONLY question is: does the candidate AGREE with the gold on the "
"point the gold makes?\n\n"
"PASS when:\n"
" * the candidate says the same thing more verbosely, or explains it;\n"
" * it uses different notation, exactly-converting units, or equivalent "
"mathematical form (15/64 vs \\frac{15}{64});\n"
" * the gold is a fragment and the candidate states the same fact in "
"prose, or answers additional parts of the question as well;\n"
" * for code: a different but equivalent implementation, or a different "
"API call with the same effect.\n\n"
"Do NOT fail the candidate merely for saying MORE than the gold, for "
"being longer, for adding a derivation, or for answering other parts of "
"the question too. Extra correct detail is never a reason to fail.\n\n"
"FAIL only when the candidate gives a DIFFERENT VALUE, contradicts the "
"gold, or does not actually answer the question.\n\n"
f"Question:\n{task.question}\n\n"
f"Gold answer:\n{task.answer}\n\n"
f"Candidate answer:\n{prediction}\n\n"
"Reply with exactly one word — PASS or FAIL.\nVerdict:"
)
# temperature=0: a grader must be DETERMINISTIC. Without it the SDK default
# (1.0) sampled the PASS/FAIL token, so the same (question, gold, candidate)
# could be graded differently on two runs — the label became a coin flip on
# any borderline answer. (Audit 2026-07-13: gold "0" vs "order w.r.t. B is 0"
# graded PASS in one run and FAIL in the next.)
resp = client.chat.completions.create(
model=model,
max_tokens=8,
temperature=0,
messages=[{"role": "user", "content": prompt}],
)
verdict = (resp.choices[0].message.content or "").strip().upper()
if "PASS" in verdict:
return True
if "FAIL" in verdict:
return False
return None
except Exception as exc: # noqa: BLE001
# LOUDLY. Swallowing this silently is how a rate-limited judge quietly
# mislabelled correct answers as wrong for a whole run while the monitor
# cheerfully reported "0 rate limits" — the exception never reached a log.
_JUDGE_FAILURES.append(type(exc).__name__)
LOGGER.warning(
"JUDGE CALL FAILED (%d so far this run) — falling back to string match, "
"which will likely mark a correct answer WRONG: %s",
len(_JUDGE_FAILURES),
str(exc)[:160],
)
return None
# --- public API --------------------------------------------------------------
# The orchestrator is instructed to end with "FINAL_ANSWER: <answer>", so the raw
# rollout text is "<reasoning…>\n\nFINAL_ANSWER: A". Everything before the marker
# (and the marker itself) must come off before we compare, or we're matching the
# model's prose against a one-token gold.
_FINAL_ANSWER_RE = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
def _answer_only(prediction: str) -> str:
"""The text AFTER the last FINAL_ANSWER marker (or the whole thing if absent).
Without this, ``_math_equal("FINAL_ANSWER: A", "A")`` is False while
``_math_equal("A", "A")`` is True the literal marker was never stripped, so
every short non-numeric answer (multiple-choice letters, symbolic values) was
auto-marked WRONG. Numeric answers happened to survive via float parsing,
which is why this hid for so long. (audit 2026-07-12)
"""
marks = list(_FINAL_ANSWER_RE.finditer(prediction))
return (prediction[marks[-1].end() :] if marks else prediction).strip()
def verify_answer(task: Task, prediction: str) -> bool:
"""Domain-dispatched correctness check for ``prediction`` vs ``task.answer``."""
pred = _answer_only(prediction or "")
if not pred or not (task.answer or "").strip():
return False
domain = (task.domain or "").lower()
if domain == "math":
# string + numeric + f1 only. sympy removed (GIL-deadlocked the threaded
# sampler); Gemini judge NOT used here either (32 tasks x N samples of math
# judge calls throttle Gemini and stall the whole run). Accept slightly lower
# math yield for a fast, hang-free verifier.
return _math_equal(pred, task.answer)
if domain == "code":
# Best-effort: short gold -> normalized substring; otherwise LLM judge.
gold = task.answer.strip()
if len(gold) <= 200:
if _string_or_f1(pred, gold, f1_threshold=0.8):
return True
judged = _gemini_judge(task, pred)
if judged is not None:
return judged
return _string_or_f1(pred, gold, f1_threshold=0.6)
# science / medical / chat / misc / unknown -> LLM judge, then F1 fallback.
judged = _gemini_judge(task, pred)
if judged is not None:
return judged
return _string_or_f1(pred, task.answer, f1_threshold=0.6)
def make_verifier() -> Callable[[Any, Any], bool]:
"""Return ``verify_fn(task, rollout) -> bool`` for the rejection sampler."""
def verify(task: Any, rollout: Any) -> bool:
ans = getattr(rollout, "final_answer", "") or ""
return verify_answer(task, ans)
return verify
__all__ = [
"make_verifier",
"verify_answer",
]
@@ -52,7 +52,7 @@ class OrchestratorSFTConfig:
"""Configuration for orchestrator SFT training."""
# Model
model_name: str = "Qwen/Qwen3-1.7B"
model_name: str = "Qwen/Qwen3.5-9B"
max_seq_length: int = 4096
# Training
@@ -63,17 +63,28 @@ class OrchestratorSFTConfig:
warmup_ratio: float = 0.1
max_grad_norm: float = 1.0
# Trace generation
teacher_engine_key: str = ""
teacher_model: str = ""
traces_per_query: int = 2
max_attempts_per_trace: int = 3
# Trace generation by base Qwen3-8B self-sampling (v1 cold-start). The
# orchestrator IS the local Qwen3-8B served over an OpenAI-compatible vLLM
# endpoint; we roll it out over load_sft_tasks() and keep correct trajectories.
# For v2, point orchestrator_endpoint/model at the v1 checkpoint's server.
orchestrator_endpoint: str = "http://localhost:8001/v1"
orchestrator_model: str = "qwen3-8b"
orchestrator_api_key: str = "EMPTY"
generation_temperature: float = 0.7
# Data source
trace_cache_path: str = "data/orchestrator_sft_traces.jsonl"
trace_cache_path: str = "data/orchestrator_sft_v1.jsonl"
regenerate_traces: bool = False
# Rejection-sampling cold-start over load_sft_tasks (see sft_data/reject_sample.py).
distill_max_tasks: int = 0 # 0 / None -> all of load_sft_tasks()
samples_per_task: int = 8
max_keep_per_task: int = 1
max_rollout_turns: int = 8
# Optional local OSS model endpoints, mapping full model id (e.g.
# "Qwen/Qwen3.5-9B") -> vLLM base_url; omitted when unset.
local_endpoints: Dict[str, str] = field(default_factory=dict)
# Checkpoint
checkpoint_dir: str = "checkpoints/orchestrator_sft"
save_every_n_epochs: int = 1
@@ -267,11 +278,81 @@ class OrchestratorSFTTrainer:
)
def _generate_traces(self) -> None:
"""Generate SFT traces (placeholder — requires running engine)."""
"""Generate SFT traces by base Qwen3-8B self-sampling (v1 cold-start).
The orchestrator IS the local Qwen3-8B served over an OpenAI-compatible
vLLM endpoint (``orchestrator_endpoint``/``orchestrator_model``). We roll
it out over ``load_sft_tasks()`` on the fixed orchestrator tool catalog,
verify each trajectory's final answer against the gold
(``verify.make_verifier``), keep the cheapest-correct trajectory per task,
and serialize them into the ``<tool_call>`` ``conversations`` JSONL the SFT
dataset consumes. For v2, point the endpoint/model at the v1 checkpoint's
server. See ``sft_data/reject_sample.py`` and the
``scripts/orchestrator/build_orchestrator_sft.py`` driver.
"""
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
from openjarvis.agents.hybrid.toolorchestra.unified import (
make_call_orchestrator,
make_dispatch,
)
from openjarvis.learning.intelligence.orchestrator.sft_data import (
reject_sample as _reject_sample,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
load_sft_tasks,
)
generate_sft_dataset = _reject_sample.generate_sft_dataset
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
make_verifier,
)
trace_path = Path(self.config.trace_cache_path)
trace_path.parent.mkdir(parents=True, exist_ok=True)
if not trace_path.exists():
trace_path.touch()
tools = orchestrator_catalog(
local_endpoints=self.config.local_endpoints or None,
)
call_orch = make_call_orchestrator(
self.config.orchestrator_model,
base_url=self.config.orchestrator_endpoint,
api_key=self.config.orchestrator_api_key,
temperature=self.config.generation_temperature,
)
dispatch = make_dispatch({})
def rollout_fn(task: Any) -> Any:
try:
return run_unified_rollout(
task.instruction,
tools,
call_orchestrator=call_orch,
dispatch=dispatch,
max_turns=self.config.max_rollout_turns,
)
except Exception as exc: # network/key failures shouldn't kill the run
logger.warning("rollout failed for %s: %s", task.task_id, exc)
return None
try:
tasks = load_sft_tasks()
if self.config.distill_max_tasks:
tasks = tasks[: self.config.distill_max_tasks]
stats = generate_sft_dataset(
str(trace_path),
tasks=tasks,
tools=tools,
rollout_fn=rollout_fn,
verify_fn=make_verifier(),
samples_per_task=self.config.samples_per_task,
max_keep_per_task=self.config.max_keep_per_task,
reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets top reward
)
logger.info("Generated SFT traces: %s", stats)
except Exception as exc: # network/datasets optional — fail soft to empty
logger.warning("trace generation failed (%s); writing empty set", exc)
trace_path.parent.mkdir(parents=True, exist_ok=True)
if not trace_path.exists():
trace_path.touch()
def _init_optimizer(self) -> None:
if not HAS_TORCH or self.policy.model is None:
+14
View File
@@ -42,8 +42,22 @@ _MATH_FUNCS = {
"sin": math.sin,
"cos": math.cos,
"tan": math.tan,
"asin": math.asin,
"acos": math.acos,
"atan": math.atan,
"atan2": math.atan2,
"sinh": math.sinh,
"cosh": math.cosh,
"tanh": math.tanh,
"radians": math.radians,
"degrees": math.degrees,
"exp": math.exp,
"fabs": math.fabs,
"hypot": math.hypot,
"factorial": math.factorial,
"pi": math.pi,
"e": math.e,
"tau": math.tau,
"ceil": math.ceil,
"floor": math.floor,
}
+23 -3
View File
@@ -116,11 +116,31 @@ class WebSearchTool(BaseTool):
return text
def _duckduckgo_search(self, query: str, max_results: int) -> str:
"""Search using DuckDuckGo as fallback."""
"""Search using DuckDuckGo as fallback.
ddgs queries several engines (yandex/yahoo/brave/...) that frequently
hang; with no timeout this blocks the whole rollout. Cap each engine
request (DDGS timeout) AND wrap the call in a hard wall-clock deadline so
a flaky search fast-fails instead of stalling data generation.
"""
from concurrent.futures import ThreadPoolExecutor
from ddgs import DDGS
ddgs = DDGS()
raw_results = list(ddgs.text(query, max_results=max_results))
def _run():
ddgs = DDGS(timeout=5)
return list(ddgs.text(query, max_results=max_results))
# Don't use `with` — its shutdown(wait=True) blocks on the hung thread,
# defeating the deadline. submit, wait up to 8s, then abandon the thread
# (it finishes in the background harmlessly) and fast-fail.
_ex = ThreadPoolExecutor(max_workers=1)
try:
raw_results = _ex.submit(_run).result(timeout=8)
except Exception:
_ex.shutdown(wait=False, cancel_futures=True)
return "[web_search: no results (search backend timed out)]"
_ex.shutdown(wait=False)
results = []
for r in raw_results:
title = r.get("title", "Untitled")
+129
View File
@@ -0,0 +1,129 @@
"""Unit coverage for ``anonymize_tools`` — the identity-stripping step that makes
routing data unbiased. See ``expert_registry.anonymize_tools``.
The anonymizer takes the orchestrator catalog and replaces every MODEL expert
with an opaque ``model_xxxx`` label, a uniform brand-free description, no
price/latency line, and shuffles the expert block so the policy can't route on
a model's name, position, cost, or tier. Basic tools keep their real names.
It returns ``(anon_tools, anon_to_real)``.
"""
from __future__ import annotations
import json
import random
from openjarvis.agents.hybrid.expert_registry import (
KIND_MODEL,
anonymize_tools,
build_tool_specs,
orchestrator_catalog,
tools_by_name,
)
# Real brand tokens that must never leak into the anonymized, model-facing specs.
_BRANDS = ("gpt", "claude", "gemini", "qwen")
def _model_names(cat):
return [t.name for t in cat if t.kind == KIND_MODEL]
def _basic_names(cat):
return [t.name for t in cat if t.kind != KIND_MODEL]
def test_no_brand_names_in_anonymized_specs():
cat = orchestrator_catalog()
anon, _ = anonymize_tools(cat, random.Random(0))
specs = build_tool_specs(anon)
# The whole model-facing payload (names + descriptions + categories) must be
# brand-free. `.model` is preserved on the tool for dispatch but is NOT part
# of the spec the orchestrator conditions on, so it's fine that it still holds
# the real id.
blob = json.dumps(specs).lower()
for brand in _BRANDS:
assert brand not in blob, f"brand {brand!r} leaked into anonymized specs"
# And specifically the anonymized expert descriptions carry no brand.
for t in anon:
if t.name.startswith("model_"):
assert not any(b in t.description().lower() for b in _BRANDS)
def test_hide_cost_removes_price_line_from_descriptions():
cat = orchestrator_catalog()
# Sanity: the raw model tools DO surface a price line before anonymizing.
raw_model = next(t for t in cat if t.kind == KIND_MODEL)
assert "Pricing:" in raw_model.description()
anon, anon_to_real = anonymize_tools(cat, random.Random(1))
for t in anon:
if t.name in anon_to_real: # an anonymized model expert
assert t.hide_cost is True
desc = t.description()
assert "Pricing:" not in desc
assert "$" not in desc
assert "/1M" not in desc
assert "latency" not in desc.lower()
def test_each_model_maps_to_opaque_label_and_round_trips():
cat = orchestrator_catalog()
orig_models = _model_names(cat)
anon, anon_to_real = anonymize_tools(cat, random.Random(2))
# One opaque label per real model, all in the model_xxxx namespace.
assert len(anon_to_real) == len(orig_models)
assert all(lbl.startswith("model_") for lbl in anon_to_real)
# No collisions: labels unique, and each real model recovered exactly once.
assert len(set(anon_to_real)) == len(anon_to_real)
assert len(set(anon_to_real.values())) == len(anon_to_real)
assert set(anon_to_real.values()) == set(orig_models)
# Round-trip: every anonymized expert's label resolves back to a real name,
# and `.model` is preserved so dispatch still reaches the right backend.
by_orig = tools_by_name(cat)
for t in anon:
if t.name.startswith("model_"):
real = anon_to_real[t.name]
assert real in by_orig
assert t.model == by_orig[real].model # backend id untouched
# Basic tools keep their real names and are untouched by the label map.
basics = _basic_names(cat)
anon_names = {t.name for t in anon}
assert set(basics) <= anon_names
assert not (set(basics) & set(anon_to_real))
def test_experts_block_on_top_then_basics():
cat = orchestrator_catalog()
anon, anon_to_real = anonymize_tools(cat, random.Random(3))
n_models = len(anon_to_real)
# Experts form a contiguous block at the front, basics underneath.
assert all(t.name.startswith("model_") for t in anon[:n_models])
assert all(not t.name.startswith("model_") for t in anon[n_models:])
def test_labels_and_order_are_shuffled_not_identity():
cat = orchestrator_catalog()
orig_models = _model_names(cat)
orderings = set()
labels_for_first = set()
for seed in range(25):
anon, anon_to_real = anonymize_tools(cat, random.Random(seed))
# Real-model order as it appears in the anonymized expert block.
order = tuple(anon_to_real[t.name] for t in anon if t.name.startswith("model_"))
orderings.add(order)
# Label assigned to the first catalog model varies across rngs.
rev = {real: lbl for lbl, real in anon_to_real.items()}
labels_for_first.add(rev[orig_models[0]])
# Not a fixed identity: across seeds the expert order actually varies...
assert len(orderings) > 1
# ...and at least one ordering differs from the input model order.
assert any(order != tuple(orig_models) for order in orderings)
# Labels are random per-rng, not a stable function of position.
assert len(labels_for_first) > 1
+221
View File
@@ -0,0 +1,221 @@
"""Tests for the faithful ToolOrchestra unified-tool registry."""
from __future__ import annotations
import pytest
from openjarvis.agents.hybrid.expert_registry import (
CATEGORY_BASIC,
CATEGORY_CLOUD_FRONTIER,
CATEGORY_LOCAL_OSS,
KIND_MODEL,
KIND_TOOL,
ExpertTool,
build_tool_specs,
openjarvis_tool,
orchestrator_catalog,
to_worker_dict,
tools_by_name,
)
def test_catalog_names_unique_and_valid():
cat = orchestrator_catalog()
names = [t.name for t in cat]
assert len(names) == len(set(names))
assert all(isinstance(t, ExpertTool) for t in cat)
def test_invalid_tool_rejected():
with pytest.raises(ValueError):
ExpertTool(name="x", kind="bogus", backend_type="openai", summary="", model="m")
with pytest.raises(ValueError):
ExpertTool(
name="x", kind=KIND_MODEL, backend_type="openai", summary="", model=None
)
def test_specs_shape_and_pricing_in_description():
cat = orchestrator_catalog()
specs = build_tool_specs(cat)
by = {s["function"]["name"]: s for s in specs}
gpt = by["gpt_5_5"]
assert gpt["type"] == "function"
assert "input" in gpt["function"]["parameters"]["properties"]
# Price table is surfaced in the description (the policy is trained on it).
assert "/1M input" in gpt["function"]["description"]
# Search tool takes a query, code takes code.
assert "query" in by["web_search"]["function"]["parameters"]["properties"]
assert "code" in by["code_interpreter"]["function"]["parameters"]["properties"]
# Two cloud-frontier + four local-OSS model tools, in catalog order.
_ORCH_MODEL_NAMES = [
"gpt_5_5",
"claude_opus_4_8",
"qwen3_5_9b",
"qwen3_6_27b_fp8",
"qwen3_5_122b_a10b_fp8",
"qwen3_5_397b_a17b_fp8",
]
# Bridged real OpenJarvis tools (basic) appended after web_search/code_interpreter.
_ORCH_BASIC_NAMES = [
"web_search",
"code_interpreter",
"calculator",
"shell_exec",
"file_read",
"file_write",
"http_request",
"think",
"apply_patch",
"pdf_extract",
"db_query",
]
def test_orchestrator_catalog_two_model_classes_plus_basics():
cat = orchestrator_catalog()
names = [t.name for t in cat]
# 6 model tools (2 cloud_frontier + 4 local_oss) come first.
assert names[:6] == _ORCH_MODEL_NAMES
# then the basic tools.
assert set(_ORCH_BASIC_NAMES) <= set(names)
assert len(cat) == 6 + len(_ORCH_BASIC_NAMES)
by = tools_by_name(cat)
assert by["gpt_5_5"].category == CATEGORY_CLOUD_FRONTIER
assert by["claude_opus_4_8"].category == CATEGORY_CLOUD_FRONTIER
# Default routing for every model tool is OpenRouter (no servers required).
for n in (
"qwen3_5_9b",
"qwen3_6_27b_fp8",
"qwen3_5_122b_a10b_fp8",
"qwen3_5_397b_a17b_fp8",
):
assert by[n].category == CATEGORY_LOCAL_OSS
assert by[n].backend_type == "openrouter"
assert by[n].base_url is None
# OpenRouter routing carries the slug + a (estimated) per-token price.
assert "/" in by[n].model and by[n].price_in > 0.0
def test_orchestrator_catalog_categories_present():
cat = orchestrator_catalog()
cats = {t.category for t in cat}
assert cats == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS, CATEGORY_BASIC}
def test_orchestrator_can_drop_tools():
cat = orchestrator_catalog(include_tools=False)
assert {t.category for t in cat} == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS}
assert len(cat) == 6
def test_orchestrator_specs_include_category_field():
specs = build_tool_specs(orchestrator_catalog())
by = {s["function"]["name"]: s for s in specs}
assert by["gpt_5_5"]["function"]["category"] == CATEGORY_CLOUD_FRONTIER
assert by["qwen3_5_9b"]["function"]["category"] == CATEGORY_LOCAL_OSS
assert by["web_search"]["function"]["category"] == CATEGORY_BASIC
assert by["shell_exec"]["function"]["category"] == CATEGORY_BASIC
# Every tool carries a category tag.
assert all("category" in s["function"] for s in specs)
def test_orchestrator_local_models_get_base_url_when_provided():
# An endpoint switches that model from the OpenRouter default to local vLLM.
cat = orchestrator_catalog(
local_endpoints={
"Qwen/Qwen3.5-9B": "http://x/v1",
"Qwen/Qwen3.6-27B-FP8": "http://y/v1",
}
)
by = tools_by_name(cat)
assert by["qwen3_5_9b"].backend_type == "vllm"
assert by["qwen3_5_9b"].base_url == "http://x/v1"
assert by["qwen3_5_9b"].price_in == 0.0 and by["qwen3_5_9b"].price_out == 0.0
assert by["qwen3_6_27b_fp8"].base_url == "http://y/v1"
# Unmapped local model -> stays on the OpenRouter default (base_url None).
assert by["qwen3_5_122b_a10b_fp8"].backend_type == "openrouter"
assert by["qwen3_5_122b_a10b_fp8"].base_url is None
# Cloud frontier carries real pricing.
assert by["claude_opus_4_8"].price_in > 0.0
assert by["gpt_5_5"].price_in > 0.0
def test_orchestrator_model_backends_override():
# Force a cloud model onto OpenRouter and a local model onto vLLM explicitly.
cat = orchestrator_catalog(
model_backends={
"claude-opus-4-8": "openrouter",
"Qwen/Qwen3.5-397B-A17B-FP8": "vllm",
},
local_endpoints={"Qwen/Qwen3.5-397B-A17B-FP8": "http://z/v1"},
)
by = tools_by_name(cat)
assert by["claude_opus_4_8"].backend_type == "openrouter"
assert by["claude_opus_4_8"].model == "anthropic/claude-opus-4.8"
# No override for gpt-5.5 -> it defaults to its NATIVE first-party API
# (openai), not OpenRouter. Frontier models hit their native provider by
# default; OpenRouter is only the fallback for OSS/local models or when
# explicitly requested via model_backends.
assert by["gpt_5_5"].backend_type == "openai" # native default
assert by["gpt_5_5"].model == "gpt-5.5"
assert by["qwen3_5_397b_a17b_fp8"].backend_type == "vllm"
assert by["qwen3_5_397b_a17b_fp8"].base_url == "http://z/v1"
def test_orchestrator_openrouter_slug_override():
cat = orchestrator_catalog(
openrouter_slugs={"Qwen/Qwen3.5-9B": "qwen/qwen3.5-9b-custom"}
)
by = tools_by_name(cat)
assert by["qwen3_5_9b"].model == "qwen/qwen3.5-9b-custom"
def test_openjarvis_tool_bridges_real_tool_with_custom_schema():
params = {
"type": "object",
"properties": {"command": {"type": "string"}},
"required": ["command"],
}
t = openjarvis_tool("shell_exec", summary="Run a shell command.", params=params)
assert t.kind == KIND_TOOL
assert t.backend_type == "openjarvis-tool"
assert t.model == "shell_exec"
assert t.category == CATEGORY_BASIC
spec = t.to_spec()
assert spec["function"]["name"] == "shell_exec"
assert spec["function"]["parameters"] == params
assert spec["function"]["category"] == CATEGORY_BASIC
def test_build_tool_specs_includes_category_for_bridged_tools():
specs = build_tool_specs(
[
openjarvis_tool(
"calculator",
summary="Math.",
params={
"type": "object",
"properties": {"expression": {"type": "string"}},
"required": ["expression"],
},
),
]
)
assert specs[0]["function"]["category"] == CATEGORY_BASIC
assert "expression" in specs[0]["function"]["parameters"]["properties"]
def test_to_worker_dict_maps_backend():
cat = orchestrator_catalog(local_endpoints={"Qwen/Qwen3.5-9B": "http://x/v1"})
by = tools_by_name(cat)
assert to_worker_dict(by["gpt_5_5"]) == {
"name": "gpt_5_5",
"type": "openai",
"model": "gpt-5.5",
}
local = to_worker_dict(by["qwen3_5_9b"])
assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"
@@ -0,0 +1,154 @@
"""Offline tests for the reasoning-SFT dataset loaders.
All loaders take a ``source=`` iterable of raw row dicts so we never touch the
network: we hand them fake rows shaped like the real GeneralThought /
OpenThoughts3 schemas and assert the normalized :class:`Task` fields.
"""
from __future__ import annotations
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
Task,
load_generalthought,
load_openthoughts,
)
# --- fake rows mirroring the real HF schemas ---------------------------------
GENERALTHOUGHT_ROWS = [
{
"question_id": 1,
"question": "What is 2 + 2?",
"reference_answer": "4",
"model_answer": "The answer is 4.",
"task": "High School Math",
"question_source": "Numina/NuminaMath",
},
{
"question_id": 2,
"question": "How do the neural respiratory centers operate?",
"reference_answer": "Via the medulla oblongata.",
"model_answer": "long answer ...",
"task": "NHS QA",
"question_source": "CogStack/NHSQA",
},
{
"question_id": 3,
"question": "Sort a list.",
"reference_answer": "",
"model_answer": "use sorted()",
"task": "Sorting",
"question_source": "BAAI/TACO",
},
# Should be skipped: empty question and answer.
{"question_id": 4, "question": "", "reference_answer": "", "model_answer": ""},
]
OPENTHOUGHTS_ROWS = [
{
"domain": "code",
"source": "nvidia/OpenCodeReasoning",
"difficulty": 7,
"conversations": [
{"from": "human", "value": "Write a program that prints hi."},
{"from": "gpt", "value": "<think>reason</think> Final: print('hi')"},
],
},
{
"domain": "math",
"source": "Numina/NuminaMath",
"difficulty": 5,
"conversations": [
{"from": "human", "value": "What is 6 times 7?"},
{
"from": "gpt",
"value": "<think>6*7=42</think> The answer is \\boxed{42}.",
},
],
},
{
"domain": "science",
"source": "sci",
"difficulty": 6,
"conversations": [
{"from": "human", "value": "What gas do plants release?"},
{"from": "gpt", "value": "Oxygen."},
],
},
# Should be skipped: no human/gpt turns.
{"domain": "math", "conversations": []},
# Should be skipped: codegolf is an UNUSABLE_SOURCE (adversarial, unscorable).
{
"domain": "code",
"source": "stackexchange_codegolf",
"conversations": [
{"from": "human", "value": "Golf: print hi."},
{"from": "gpt", "value": "print('hi')"},
],
},
]
def test_load_generalthought_fields_and_domains():
# exclude_domains=() so the lookup-shaped medical row survives for the
# classification asserts below — by DEFAULT it is dropped (see end of test).
tasks = list(
load_generalthought(n=10, source=GENERALTHOUGHT_ROWS, exclude_domains=())
)
assert all(isinstance(t, Task) for t in tasks)
# Row 4 (empty) is dropped; remaining 3 yielded.
assert len(tasks) == 3
by_q = {t.question: t for t in tasks}
assert by_q["What is 2 + 2?"].answer == "4"
assert by_q["What is 2 + 2?"].domain == "math"
assert by_q["What is 2 + 2?"].instruction == "What is 2 + 2?"
assert by_q["How do the neural respiratory centers operate?"].domain == "medical"
# No reference_answer -> falls back to model_answer; TACO source -> code.
assert by_q["Sort a list."].answer == "use sorted()"
assert by_q["Sort a list."].domain == "code"
# By DEFAULT the lookup-shaped medical row is dropped (LOOKUP_DOMAINS), so the
# same source yields only the math + code reasoning tasks.
default = list(load_generalthought(n=10, source=GENERALTHOUGHT_ROWS))
assert {t.domain for t in default} == {"math", "code"}
def test_load_generalthought_respects_n():
tasks = list(load_generalthought(n=2, source=GENERALTHOUGHT_ROWS))
assert len(tasks) == 2
def test_load_openthoughts_fields_and_domains():
tasks = list(
load_openthoughts(n_code=5, n_math=5, n_science=5, source=OPENTHOUGHTS_ROWS)
)
assert all(isinstance(t, Task) for t in tasks)
assert len(tasks) == 3
by_dom = {t.domain: t for t in tasks}
assert set(by_dom) == {"code", "math", "science"}
# <think> stripped; boxed answer extracted.
assert by_dom["math"].answer == "42"
assert by_dom["math"].question == "What is 6 times 7?"
# Code: think stripped, visible solution kept.
assert "print('hi')" in by_dom["code"].answer
assert by_dom["science"].answer == "Oxygen."
def test_load_openthoughts_quota_caps_per_domain():
rows = [
{
"domain": "code",
"conversations": [
{"from": "human", "value": f"q{i}"},
{"from": "gpt", "value": f"a{i}"},
],
}
for i in range(10)
]
tasks = list(load_openthoughts(n_code=3, n_math=0, n_science=0, source=rows))
assert len(tasks) == 3
assert all(t.domain == "code" for t in tasks)
def test_task_dataclass_instruction_property():
t = Task(task_id="x", question="hi?", answer="yes", domain="misc")
assert t.instruction == "hi?"
@@ -0,0 +1,150 @@
"""Offline tests for the rejection-sampling SFT pipeline (unified tools)."""
from __future__ import annotations
import json
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog, tools_by_name
from openjarvis.agents.hybrid.toolorchestra.rollout import (
UnifiedRollout,
UnifiedTurn,
run_unified_rollout,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
def test_run_unified_rollout_terminates_on_no_tool_call():
tools = orchestrator_catalog()
by = tools_by_name(tools)
name = tools[0].name # any real expert tool from the live catalog
assert name in by
scripted = [
("reason 1\n", [(name, {"input": "do step 1"})], 5, 5),
("here is the answer", [], 3, 3), # no tool call -> terminate
]
calls = iter(scripted)
def call_orch(messages, specs):
# Signature is now (messages, specs): the rollout drives a running
# system/user/assistant/tool conversation, not a flattened (system, user).
return next(calls)
def dispatch(tool, args):
return (f"OBS for {tool.name}", 0.01, 10, False)
roll = run_unified_rollout(
"What is X?",
tools,
call_orchestrator=call_orch,
dispatch=dispatch,
max_turns=5,
)
assert roll.final_answer == "here is the answer"
assert roll.num_tool_calls == 1
assert roll.tool_calls() == [(name, {"input": "do step 1"})]
assert abs(roll.cost_usd - 0.01) < 1e-9
def test_serialize_record_shape_and_tool_call_tags():
tools = orchestrator_catalog()
roll = UnifiedRollout(
turns=[
UnifiedTurn(
reasoning="think",
tool_name="qwen3_32b",
arguments={"input": "q"},
observation="obs",
),
# The final turn's reasoning IS the model's real final output (it
# already contains the answer). The serializer now renders this turn
# via _final_answer_block(turn.reasoning), not rollout.final_answer.
UnifiedTurn(
reasoning="The result is 42.\nFINAL_ANSWER: 42", tool_name=None
),
],
final_answer="42",
cost_usd=0.02,
tokens=30,
num_tool_calls=1,
)
rec = trajectory_to_record("t1", "Q?", tools, roll, reward=0.5, domain="math")
roles = [m["role"] for m in rec["conversations"]]
assert roles[0] == "system" and roles[1] == "user"
assert "tool" in roles and roles[-1] == "assistant"
# Tool call is emitted as a <tool_call> tag (what the parser reads back).
assert any(
"<tool_call>" in m["content"] and "qwen3_32b" in m["content"]
for m in rec["conversations"]
if m["role"] == "assistant"
)
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
assert rec["reward"] == 0.5 and rec["domain"] == "math"
def test_generate_sft_dataset_end_to_end(tmp_path):
tools = orchestrator_catalog()
tasks = [
Task(
task_id="movie-001",
question="Cancel ticket A03 and refund the user.",
answer="refunded $20.90",
domain="entertainment",
),
Task(
task_id="unsolvable",
question="Cancel ticket A03 and refund the user.",
answer="refunded $20.90",
domain="entertainment",
),
]
def rollout_fn(task):
# Solve the first task; the second answers wrongly and never routes.
if task.task_id == "movie-001":
return UnifiedRollout(
turns=[
UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
UnifiedTurn("done", None),
],
# num_tool_calls must reflect the two routed calls: the structural
# _target_is_clean gate drops a trajectory with num_tool_calls < 1.
final_answer="refunded $20.90",
cost_usd=0.03,
num_tool_calls=2,
)
return UnifiedRollout(
turns=[UnifiedTurn("x", "cancel", {}, "ok")],
final_answer="nope",
cost_usd=0.05,
)
# Simple, meaningful verifier: the final answer must mention the refund.
def verify_fn(task, rollout):
return "refund" in (rollout.final_answer or "").lower()
out = tmp_path / "sft.jsonl"
stats = generate_sft_dataset(
str(out),
tasks=tasks,
tools=tools,
rollout_fn=rollout_fn,
verify_fn=verify_fn,
samples_per_task=2,
)
assert stats["tasks_seen"] == 2
assert stats["records_written"] == 1 # only the solvable task
assert stats["tasks_dropped"] == 1
lines = out.read_text().strip().splitlines()
assert len(lines) == 1
rec = json.loads(lines[0])
assert rec["task_id"] == "movie-001"
assert rec["domain"] == "entertainment"
assert (tmp_path / "sft.jsonl.stats.json").exists()
@@ -0,0 +1,111 @@
"""Offline tests for the answer verifier.
Math and code paths run with no network. The Gemini judge is only reachable for
non-math/code domains; we monkeypatch ``_gemini_judge`` (or rely on its no-key
fallback) so nothing hits the network.
"""
from __future__ import annotations
from openjarvis.learning.intelligence.orchestrator.sft_data import verify as V
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
make_verifier,
verify_answer,
)
def _math(answer: str) -> Task:
return Task(task_id="m", question="q", answer=answer, domain="math")
def test_math_boxed_vs_plain_true():
assert verify_answer(_math("42"), "\\boxed{42}") is True
def test_math_boxed_both_sides():
assert verify_answer(_math("\\boxed{42}"), "The answer is \\boxed{42}") is True
def test_math_wrong_number_false():
assert verify_answer(_math("42"), "41") is False
def test_math_numeric_float_match():
assert verify_answer(_math("42"), "42.0") is True
def test_math_symbolic_equivalence_no_longer_verified_locally():
# The sympy symbolic-equality block was removed from _math_equal: sympy.simplify
# / parse_latex could hang on pathological \boxed{} answers while holding the GIL
# and deadlock the threaded rejection sampler. The math domain now uses
# string+numeric+f1 only (no judge call either), so a purely SYMBOLIC
# equivalence that isn't string/numeric-equal is intentionally NOT recognized.
t = _math("(x+1)^2")
assert verify_answer(t, "x^2 + 2*x + 1") is False
def test_empty_prediction_false():
assert verify_answer(_math("42"), "") is False
def test_code_short_gold_substring():
t = Task(task_id="c", question="print hi", answer="hello", domain="code")
assert verify_answer(t, "the output is hello") is True
def test_code_falls_back_to_judge(monkeypatch):
# Long gold + no substring match -> hits judge; monkeypatch to PASS.
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: True)
t = Task(
task_id="c",
question="solve",
answer="x" * 300, # long -> skips substring shortcut
domain="code",
)
assert verify_answer(t, "totally different") is True
def test_judge_domain_uses_gemini(monkeypatch):
calls = {}
def fake_judge(task, pred):
calls["hit"] = (task.domain, pred)
return True
monkeypatch.setattr(V, "_gemini_judge", fake_judge)
t = Task(task_id="s", question="q", answer="Oxygen", domain="science")
assert verify_answer(t, "plants release oxygen") is True
assert calls["hit"][0] == "science"
def test_judge_domain_fallback_f1_when_no_key(monkeypatch):
# Judge unavailable (returns None) -> falls back to string / F1.
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
t = Task(task_id="s", question="q", answer="the capital is paris", domain="misc")
assert verify_answer(t, "I think the capital is paris indeed") is True
assert verify_answer(t, "completely unrelated text here") is False
def test_make_verifier_reads_final_answer(monkeypatch):
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
verify = make_verifier()
class Rollout:
final_answer = "\\boxed{42}"
assert verify(_math("42"), Rollout()) is True
class BadRollout:
final_answer = "999"
assert verify(_math("42"), BadRollout()) is False
def test_make_verifier_missing_final_answer():
verify = make_verifier()
class Empty:
pass
assert verify(_math("42"), Empty()) is False
@@ -0,0 +1,156 @@
"""Offline smoke test for the v1 orchestrator-SFT driver (base self-sampling).
No network / no GPU: we feed fake ``datasets.Task`` objects, a canned
``UnifiedRollout`` (via a fake ``rollout_fn``), and an accept-all verifier, then
assert ``generate_sft_dataset`` writes a JSONL record in the expected
``conversations`` shape that the SFT trainer consumes.
"""
from __future__ import annotations
import json
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
)
def _fake_tasks(**_kwargs) -> list[Task]:
# Accepts **kwargs so it can stand in for ``load_sft_tasks(cap=..., balanced=...)``.
return [
Task(task_id="t-math-1", question="What is 6 * 7?", answer="42", domain="math"),
Task(task_id="t-code-1", question="Reverse 'ab'.", answer="ba", domain="code"),
]
def _canned_rollout(task: Task) -> UnifiedRollout:
# One tool turn + a final-answer turn that echoes the gold answer.
return UnifiedRollout(
turns=[
UnifiedTurn(
reasoning="let me compute",
tool_name="code_interpreter",
arguments={"code": "print(6*7)"},
observation="42",
),
UnifiedTurn(
reasoning=(
f"The tool returned {task.answer}. FINAL_ANSWER: {task.answer}"
),
tool_name=None,
),
],
final_answer=task.answer,
cost_usd=0.01,
tokens=20,
num_tool_calls=1,
)
def test_build_v1_writes_expected_conversations_shape(tmp_path, monkeypatch):
tools = orchestrator_catalog() # specialists unwired (math/coder endpoints None)
# Accept-all verifier (the real make_verifier may hit Gemini / math checkers).
monkeypatch.setattr(
"openjarvis.learning.intelligence.orchestrator.sft_data.verify.make_verifier",
lambda: lambda task, rollout: True,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
make_verifier,
)
out = tmp_path / "orchestrator_sft_v1.jsonl"
stats = generate_sft_dataset(
str(out),
tasks=_fake_tasks(),
tools=tools,
rollout_fn=_canned_rollout,
verify_fn=make_verifier(),
samples_per_task=2,
max_keep_per_task=1,
reward_fn=lambda r: -r.cost_usd,
)
assert stats["tasks_seen"] == 2
assert stats["records_written"] == 2
assert stats["tasks_dropped"] == 0
lines = out.read_text().strip().splitlines()
assert len(lines) == 2
rec = json.loads(lines[0])
assert rec["task_id"] == "t-math-1"
assert rec["domain"] == "math"
roles = [m["role"] for m in rec["conversations"]]
assert roles[0] == "system" and roles[1] == "user"
assert "tool" in roles
assert roles[-1] == "assistant"
# Tool call emitted as a <tool_call> tag the parser reads back.
assert any(
"<tool_call>" in m["content"] and "code_interpreter" in m["content"]
for m in rec["conversations"]
if m["role"] == "assistant"
)
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
assert (tmp_path / "orchestrator_sft_v1.jsonl.stats.json").exists()
def _load_driver():
"""Load the (non-package) CLI driver module by path."""
import importlib.util
from pathlib import Path
root = Path(__file__).resolve().parents[2]
path = root / "scripts" / "orchestrator" / "build_orchestrator_sft.py"
spec = importlib.util.spec_from_file_location("build_orchestrator_sft", path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_driver_runs_end_to_end_with_fakes(tmp_path, monkeypatch):
"""Drive ``main`` with no network: fake task loader, rollout, and verifier."""
drv = _load_driver()
monkeypatch.setattr(drv, "load_sft_tasks", _fake_tasks)
monkeypatch.setattr(
drv,
"run_unified_rollout",
lambda question, tools, **kw: _canned_rollout(_fake_tasks()[0]),
)
monkeypatch.setattr(drv, "make_verifier", lambda: lambda task, rollout: True)
# make_call_orchestrator builds an OpenAI client lazily, so it never connects
# here (run_unified_rollout is stubbed out).
# OJ_DATA_ROOT is an ABSOLUTE path in the real .env (the data tree lives
# outside the repo). If it leaks into the test env, chdir(tmp_path) below no
# longer sandboxes anything and the run writes into the real dataset dir.
# Drop it so the driver falls back to its repo-relative default.
monkeypatch.delenv("OJ_DATA_ROOT", raising=False)
monkeypatch.chdir(tmp_path)
rc = drv.main(
[
"--out",
"v1",
"--samples-per-task",
"1",
"--max-tasks",
"2",
]
)
assert rc == 0
# raw dir is {label}-{month}-{day}-{year}-{hhmm}{am|pm}[-{tag}] — e.g.
# orch-july-7-2026-0553pm-v1 (see sft_data.naming.raw_dir_name).
produced = list(
(tmp_path / "data" / "orchestrator" / "raw").glob("*-v1/data.jsonl")
)
assert len(produced) == 1
lines = produced[0].read_text().strip().splitlines()
assert len(lines) == 2
rec = json.loads(lines[0])
assert rec["conversations"][0]["role"] == "system"
assert "FINAL_ANSWER" in rec["conversations"][-1]["content"]
@@ -0,0 +1,182 @@
"""Offline tests for the orchestrator eval backend + the eval CLI's
benchmark-name normalization.
No network / models: ``run_unified_rollout`` is monkeypatched to return a
canned rollout, and the dataset-key mapping is asserted against the real
``openjarvis.evals.cli`` registry without instantiating any (network-bound)
dataset.
"""
from __future__ import annotations
import pytest
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedTurn
from openjarvis.learning.intelligence.orchestrator import eval_backend as eb
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
OrchestratorBackend,
)
class _CannedRollout:
"""Minimal stand-in for UnifiedRollout."""
def __init__(self, final_answer: str) -> None:
self.final_answer = final_answer
self.cost_usd = 0.123
self.tokens = 42
self.num_tool_calls = 1
self.parse_failures = 0
self.anon_map = {}
# generate_full now serializes a per-turn trace (reasoning / tool_name /
# arguments / observation), so turns must be real UnifiedTurn-shaped
# objects, not bare object() placeholders.
self.turns = [
UnifiedTurn(
reasoning="let me search",
tool_name="web_search",
arguments={"query": "x"},
observation="result",
),
]
def tool_calls(self):
return [("web_search", {"query": "x"})]
def _patch_rollout(monkeypatch, rollout):
"""Patch run_unified_rollout where the backend looks it up."""
monkeypatch.setattr(
eb._rollout_mod,
"run_unified_rollout",
lambda *a, **k: rollout,
)
def test_construct_is_network_free():
"""Constructing the backend builds the catalog but touches no network."""
backend = OrchestratorBackend()
assert backend.backend_id == "orchestrator"
assert backend.orchestrator_endpoint == eb.DEFAULT_ENDPOINT
assert backend.orchestrator_model == eb.DEFAULT_MODEL
# Catalog is non-empty (cloud frontier + local OSS + basic tools).
assert len(backend._tools) > 0
names = {t.name for t in backend._tools}
assert "web_search" in names
def test_local_endpoints_mapping():
"""local_endpoints (model-id -> base_url) wires the local OSS tool."""
backend = OrchestratorBackend(
local_endpoints={"Qwen/Qwen3.5-9B": "http://m:9/v1"},
)
assert backend.local_endpoints["Qwen/Qwen3.5-9B"] == "http://m:9/v1"
# The corresponding local-OSS tool should carry that base_url.
tool = next(t for t in backend._tools if t.model == "Qwen/Qwen3.5-9B")
assert tool.base_url == "http://m:9/v1"
def test_generate_returns_final_answer(monkeypatch):
"""generate() returns the rollout's final_answer string."""
_patch_rollout(monkeypatch, _CannedRollout("42"))
backend = OrchestratorBackend()
out = backend.generate("what is 6*7?", model="qwen3-8b")
assert out == "42"
def test_generate_full_payload(monkeypatch):
"""generate_full() returns the runner-expected payload shape."""
_patch_rollout(monkeypatch, _CannedRollout("42"))
backend = OrchestratorBackend()
full = backend.generate_full("q", model="qwen3-8b")
assert full["content"] == "42"
assert full["cost_usd"] == 0.123
assert full["usage"]["completion_tokens"] == 42
assert full["tool_calls"] == 1
assert full["turn_count"] == 1
assert full["framework"] == "openjarvis-orchestrator"
assert "error" not in full
assert full["latency_seconds"] >= 0.0
def test_generate_full_error_path(monkeypatch):
"""An exception inside the rollout becomes a recorded error, not a crash."""
def _boom(*a, **k):
raise RuntimeError("vllm down")
monkeypatch.setattr(eb._rollout_mod, "run_unified_rollout", _boom)
backend = OrchestratorBackend()
full = backend.generate_full("q", model="qwen3-8b")
assert full["content"] == ""
assert "vllm down" in full["error"]
assert full["usage"]["completion_tokens"] == 0
# ---------------------------------------------------------------------------
# Eval-script benchmark name mapping
# ---------------------------------------------------------------------------
# The 5 benchmarks the eval script targets, by the script's (alias) names.
SCRIPT_BENCHMARKS = [
"gaia",
"terminalbench_v2_1",
"taubench",
"mmlu_pro",
"supergpqa",
]
def _load_script():
import importlib.util
from pathlib import Path
repo_root = Path(__file__).resolve().parents[2]
script_path = repo_root / "scripts" / "orchestrator" / "eval_orchestrator.py"
spec = importlib.util.spec_from_file_location("eval_orchestrator", script_path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
def test_benchmark_aliases_map_to_registry_keys():
"""Each script benchmark normalizes to a real evals.cli BENCHMARKS key."""
from openjarvis.evals.cli import BENCHMARKS
script = _load_script()
for name in SCRIPT_BENCHMARKS:
key = script._normalize_benchmark(name)
assert key in BENCHMARKS, f"{name} -> {key} not in BENCHMARKS registry"
def test_normalize_specific_keys():
"""Spot-check the underscored aliases normalize to the dotted/dashed keys."""
script = _load_script()
assert script._normalize_benchmark("terminalbench_v2_1") == "terminalbench-v2.1"
assert script._normalize_benchmark("mmlu_pro") == "mmlu-pro"
assert script._normalize_benchmark("gaia") == "gaia"
def test_build_dataset_accepts_keys():
"""_build_dataset must recognise the normalized keys.
We only assert it doesn't raise the 'Unknown benchmark' ClickException;
datasets that need network/files at construction are xfail-skipped.
"""
import click
from openjarvis.evals.cli import _build_dataset
script = _load_script()
for name in SCRIPT_BENCHMARKS:
key = script._normalize_benchmark(name)
try:
_build_dataset(key)
except click.ClickException as exc:
if "Unknown benchmark" in str(exc):
pytest.fail(f"{key} not recognised by _build_dataset")
# Other ClickExceptions (missing files/creds) are acceptable here.
except Exception:
# Construction may require network/files — that's fine; the point
# is the key is recognised (didn't hit the unknown-benchmark else).
pytest.skip(f"{key} dataset construction needs resources")
@@ -0,0 +1,120 @@
"""Control-token sanitation in the SFT serializer + clean gate.
Regression guard for the leak audit: leaked control/special tokens
(``<|im_end|>``, ``<|tool_call>``, gemma ``<start_of_turn>``/``<end_of_turn>``,
``<|"|>`` …) used to survive into the supervised final answer because the clean
gate only rejected the bare ``<tool_call>`` form. Defense in depth now:
1. ``_final_answer_block`` STRIPS residual control tokens so a good answer with a
stray token is salvaged (not dropped).
2. ``_target_is_clean`` REJECTS a rollout whose final answer is nothing but
control tokens, or where a token survives the strip.
Both must leave legitimate ``<``/``>`` in math/code untouched.
"""
from __future__ import annotations
import pytest
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
_target_is_clean,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
_CONTROL_TOKEN_RE,
_final_answer_block,
_strip_control_tokens,
)
# --- serializer strips control tokens, leaving a clean FINAL_ANSWER line ------
@pytest.mark.parametrize(
"raw, expected_answer",
[
("FINAL_ANSWER: 42<end_of_turn>", "42"),
("FINAL_ANSWER: 42<|end_of_turn|>", "42"),
("FINAL_ANSWER: 42<|im_end|>", "42"),
("FINAL_ANSWER: 42<|eot_id|>", "42"),
("FINAL_ANSWER: The result<|tool_call>", "The result"),
('FINAL_ANSWER: foo <|"|> bar', "foo bar"),
("<start_of_turn>model", "model"),
],
)
def test_final_answer_block_strips_tokens(raw: str, expected_answer: str) -> None:
out = _final_answer_block(raw)
assert out == f"FINAL_ANSWER: {expected_answer}"
# No control token survives the serializer.
assert not _CONTROL_TOKEN_RE.search(out), out
def test_bare_im_end_only_becomes_empty_answer() -> None:
# An answer that is nothing but a control token strips to empty; the clean
# gate (below) is what rejects such a rollout.
assert _final_answer_block("<|im_end|>") == "FINAL_ANSWER: "
def test_math_and_code_angle_brackets_survive() -> None:
# The narrow regex must NOT eat legitimate comparisons / generics.
for good in [
"FINAL_ANSWER: x < 3 and y > 2",
"FINAL_ANSWER: List<int> and a<b>c",
"FINAL_ANSWER: if a < b: return a > 0",
]:
out = _final_answer_block(good)
assert out == good, out
assert not _CONTROL_TOKEN_RE.search(out)
def test_strip_helper_handles_nested_and_whitespace() -> None:
assert _strip_control_tokens(" hi <|im_end|> ") == "hi"
assert _strip_control_tokens("a<|im_start|>b<|im_end|>c") == "abc"
# --- clean gate: salvage vs reject -------------------------------------------
def _roll(final_answer: str) -> UnifiedRollout:
"""A minimal well-formed rollout (one expert call, clean obs) whose only
variable is the final answer."""
return UnifiedRollout(
turns=[
UnifiedTurn(
reasoning="route it",
tool_name="expert_a",
arguments={"input": "q"},
observation="valid observation",
),
UnifiedTurn(reasoning=final_answer, tool_name=None),
],
final_answer=final_answer,
cost_usd=0.01,
tokens=10,
num_tool_calls=1,
anon_map={"expert_a": "gpt"},
)
@pytest.mark.parametrize(
"final_answer, clean",
[
# Salvageable: a good answer with a trailing stray token stays clean
# (the serializer strips the token from the emitted target).
("42<end_of_turn>", True),
("42<|end_of_turn|>", True),
("The result<|tool_call>", True),
('foo <|"|> bar', True),
("<start_of_turn>model", True),
("x < 3 and y > 2", True),
("a normal plain answer", True),
# Unsalvageable: the answer is nothing but a control token -> rejected.
("<|im_end|>", False),
("<|eot_id|>", False),
("<end_of_turn>", False),
# Bare tool-call tag: already rejected by the existing gate.
("<tool_call>{}</tool_call>", False),
],
)
def test_clean_gate_salvages_or_rejects(final_answer: str, clean: bool) -> None:
assert _target_is_clean(_roll(final_answer)) is clean
@@ -15,7 +15,7 @@ from openjarvis.learning.intelligence.orchestrator.sft_trainer import (
class TestOrchestratorSFTConfig:
def test_defaults(self):
cfg = OrchestratorSFTConfig()
assert cfg.model_name == "Qwen/Qwen3-1.7B"
assert cfg.model_name == "Qwen/Qwen3.5-9B"
assert cfg.num_epochs == 3
assert cfg.batch_size == 8
assert cfg.learning_rate == 2e-5