Compare commits

...
Author SHA1 Message Date
Andrew Park 90dad24a9e orchestrator: drop the chat REPL
Nothing imports it and it isn't part of the pipeline. Recoverable from
history if the manual routing-inspection loop is wanted again.
2026-07-12 14:12:43 -07:00
Andrew Park c9c25895a8 orchestrator: drop the Modal-era FSDP accelerate config
Nothing references it, and its consumer (modal_sft.py) is already gone.
Training runs on Lambda H100 via scripts/train/fsdp{4,8}.yaml.
2026-07-12 14:11:39 -07:00
Andrew Park 3bffa33938 orchestrator: drop the unused hotpotqa loader
Nothing imports it. verify.py copied the normalization out of it to stay
self-contained, and datasets.py only mentions it as a style reference.
2026-07-12 14:01:35 -07:00
Andrew Park 82a2d3edde hybrid: replace guessed OpenRouter Qwen prices with list prices
The Qwen3.5/3.6 entries were active-param-scaled guesses with a comment saying
to verify them before trusting the cost-aware reward. Verified against
openrouter.ai: every one of them was low on output, by 1.5x (9B) to 5.2x
(27B), with the 122B off by 3.5x. The reward reads PRICES directly, so it has
been undercharging the mid and large Qwen experts the whole time.

qwen3.6-27b is not listed on OpenRouter; charge it at qwen3.5-27b's quote,
which is the nearest real one.
2026-07-11 14:12:45 -07:00
Andrew Park 9571a158a2 toolorchestra: carry provider-aware search through the package split
The monolith->package split landed while #558 was adding provider-aware web
search to toolorchestra.py upstream. Rebasing resolved that as delete-vs-modify,
which would have silently dropped the feature: the package always routed search
to Anthropic. Port it -- search now follows cloud_endpoint (openai-web-search /
gemini-web-search / anthropic-web-search), with the matching worker types,
per-search cost surcharges, and pool validation.

_TOOLORCH_SEARCH_TYPES replaces the scattered == "anthropic-web-search" checks,
so the one-shot guard and the "needs a solver" rule now cover every search type
rather than just the Anthropic one -- previously a search worker could have been
picked as the synthesis fallback.

Also, to get back under the ruff gate CI now enforces (it landed in the 95
commits we were behind): formatting, import order, and one real bug -- calculator
had "abs" twice in the same dict literal, and eval_backend had four imports
stranded below a function body.

The build_orchestrator_sft driver test read OJ_DATA_ROOT from the ambient
environment, so it wrote to the real experiments tree instead of tmp_path for
anyone who had sourced .env. Drop the var in the test.
2026-07-11 12:53:52 -07:00
Andrew Park 4864ea7776 orchestrator: give every generation run its own raw/ dir
--out was being used as a path, so runs overwrote each other and a curated
file couldn't be traced back to the rollout that produced it. It is now a
tag: every run lands in <OJ_DATA_ROOT>/raw/{label}_{MMDD}[_{tag}]/data.jsonl
and keeps the file name fixed.

raw/ sits above the sft/rl fork deliberately. SFT keeps only correct and
well-formed rows, but GRPO needs the failures, so the unfiltered rollouts
have to survive reject-sampling. Naming matches across stages -- only the
stage word differs -- so raw/qwen_0707/ pairs with sft/qwen_clean_0707.jsonl
by eye.

OJ_DATA_ROOT defaults to a repo-relative data/orchestrator so a fresh clone
works, and points outside the checkout in a real workspace.
2026-07-11 12:53:31 -07:00
Andrew Park e6c47e9532 orchestrator: track the SFT training pipeline and fix ToolOrchestra tracing
The orchestrator SFT pipeline was living entirely in the working tree and
never got committed. One script in that lineage (clean_sft_data.py) has
already been lost with no way to recover it, so track the rest before the
same thing happens again:

  run_sft_fsdp.py      full-parameter FSDP trainer (launched via accelerate)
  sft_tokenize.py      conversation -> tokens + assistant-only masking
  render_sft_data.py   record -> chat-template rendering
  make_splits.py       train / holdout / overfit splits
  chat_orchestrator.py REPL for eyeballing a trained checkpoint's routing
  format_eval_sample.py, upload_to_braintrust.py
  scripts/train/       fsdp4.yaml, fsdp8.yaml, eval sbatch

toolorchestra/tracing.py was also untracked despite being imported by
unified.py and rollout.py, which left the branch broken for a fresh
checkout. Add it along with the rest of the ToolOrchestra package split.

Also here: mmlu_pro / supergpqa scorers, expert pricing for the OpenRouter
Qwen builds (estimates, flagged in-file), an expanded calculator/web_search
tool surface, and the reject-sampling clean gate that run_sft_fsdp reads
via --require-clean.

GRPO is left at its upstream state -- the RL stage hasn't been run yet, so
it stays out of this change. Paths in the sbatch are env-driven rather than
hardcoded to one cluster home.
2026-07-11 12:23:47 -07:00
Andrew Park f3d9705961 orchestrator: expand tool catalog, fix shell_exec confirmation gate, wire SFT data pipeline
Tools:
- Bridge 4 more real OpenJarvis tools into the orchestrator catalog
  (think, apply_patch, pdf_extract, db_query); catalog is now 6 models
  + 11 basic tools.
- Fix the dispatch path so confirmation-gated tools actually run: the
  ToolExecutor was built with no confirm callback, so shell_exec (and
  any requires_confirmation tool) returned a "requires confirmation"
  error instead of executing -- which silently broke TerminalBench.
  Headless eval/rollouts now auto-approve.
- Drop dead-end candidates: git_* (need the unbuilt openjarvis_rust
  extension and just duplicate shell_exec), browser_* (needs Playwright),
  and knowledge_search/sql/retrieval (personal-data RAG, empty on the
  academic benchmarks).

SFT data pipeline:
- Reasoning-task loaders (GeneralThought-430K + OpenThoughts3) with an
  8K cold-start set and a 30K GRPO prompt pool.
- Domain-dispatched verifier (math/code checkers + Gemini judge fallback,
  since the OpenAI key is dead).
- Rejection-sampling generator + unified <tool_call> serializer matching
  the GRPO rollout format; base self-sampling driver.
- Drop the old paradigm/tier scaffolding (adp_loader, build, paradigms,
  select, serialize, tiers) replaced by the unified path.

Training/eval:
- SFT + GRPO configs for Qwen3.5-9B and gemma-4-12B-it.
- OrchestratorBackend + eval harness over GAIA/TerminalBench/TauBench/
  MMLU-Pro/SuperGPQA.
- Cost-aware GRPO reward.

Ignore generated data/ artifacts.
2026-07-11 12:23:46 -07:00
Andrew Park 5aa16473ff feat(orchestrator): faithful ToolOrchestra unified-tool + rejection-sampling SFT pipeline
Replace the ADP-relabel cold-start with a ToolScale-backed, execution-grounded
rejection-sampling pipeline in the paper's real action space (arXiv:2511.21689).

- expert_registry: faithful unified tool calling — one named tool per model
  (gpt_5, qwen3_32b, ...) + basic tools, per-instance subset/price sampling (§3.1/3.3)
- toolorchestra/: split the 1663-line module into a package (prompts/experts/
  sandbox/clients/parsing/workers/agent), behavior preserved
- toolorchestra/rollout.py: faithful reason->act->observe loop (pure, injectable)
- toolorchestra/unified.py: real teacher caller + tool dispatch
- sft_data/toolscale.py: nvidia/ToolScale loader
- sft_data/unified_serialize.py: trajectory -> conversations JSONL (<tool_call> tags)
- sft_data/reject_sample.py: generate_sft_dataset (sample N -> verify -> keep)
- scripts/orchestrator/build_unified_sft.py: CLI
- 25 new offline tests; 127 orchestrator-learning tests green

Verifier is a gold-action-coverage proxy (full ToolScale DB checker + LLM judge
are follow-ups); deployed _run_rl still uses the 3-meta-tool eval port.
2026-07-11 12:23:00 -07:00
Andrew Park a6afeae721 feat(orchestrator): ADP-derived SFT cold-start data pipeline
Train a single ~9B orchestrator (Qwen3-8B) for local<->cloud routing via SFT
on synthetic traces derived from NeuLab ADP (agent-data-collection).

- sft_data/: full ADP trajectory loader, tier model + pricing, 4 paradigm
  renderers, reward-ranked selection, conversations-JSONL serializer, CLI
- wire sft_trainer._generate_traces() to the pipeline (no GPU / no API keys)
- Qwen3-8B SFT config + scripts/orchestrator CLI wrapper
- offline fixture tests (12)

Heuristic competence labels for the cold-start; real paradigm execution +
GRPO are the documented v2.
2026-07-11 12:22:44 -07:00
58 changed files with 10143 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,413 @@
#!/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 \
--out data/orchestrator_sft_v2.jsonl \
--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 sys
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.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/orchestrator/raw/{label}_{MMDD}[_{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 on the same day (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/, named to MATCH the
# sft/ file it will eventually produce — only the stage word differs:
#
# raw/qwen_0707/data.jsonl <- every rollout, incl. failures
# sft/qwen_clean_0707.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. Scheme is {name}_{MMDD}, name =
# --orchestrator-label (qwen/gemma/...). Same name twice in a day -> -2, -3
# suffix (creation order is preserved, and a MM-DD-HHMMpm stamp is redundant
# with mtime). --out is an optional extra tag, NOT a path; the data 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 = f"_{Path(args.out).stem}" if args.out else ""
base = data_root / "raw" / f"{prefix}_{time.strftime('%m%d')}{tag}"
run_dir = base.resolve()
n = 1
while run_dir.exists(): # same name+day (incl. 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,
)
# Emit human-viewable companions (<out>.pretty.json + <out>.txt) so the
# traces / full message turns can be inspected without un-escaping the JSONL.
try:
from render_sft_data import render # same scripts/orchestrator dir
except ImportError:
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from render_sft_data import render
if stats.get("records_written", 0) > 0:
rinfo = render(args.out)
logging.info("Rendered viewable traces: %s + %s", rinfo["pretty"], rinfo["txt"])
print(json.dumps(stats, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
+344
View File
@@ -0,0 +1,344 @@
#!/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 key is dead, so the default judge model is a
Gemini model (``gemini-2.5-flash``). 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.",
)
# NOTE: OpenAI key is dead — default the judge to a Gemini model.
p.add_argument(
"--judge-model",
default="gemini-2.5-flash",
help="LLM-judge model (default Gemini — OpenAI key is dead).",
)
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,
)
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())
+267
View File
@@ -0,0 +1,267 @@
#!/usr/bin/env python
"""Render graded orchestrator-eval JSONL records into clean, readable .txt files.
Companion to ``format_sft_sample.py`` (same banner/indent visual style), but
for the *scored* eval outputs written by ``eval_orchestrator.py``:
results/orch-eval-<model>-<tranche>-<suite>/<benchmark>_orchestrator.jsonl
Each record has: record_id, benchmark, model, model_answer (the model's
FINAL_ANSWER string), is_correct, score, latency_seconds, cost_usd, error, and
scoring_metadata. The full step-by-step routing trace is NOT saved — only the
final answer + gold + score — so this renderer cannot show intermediate steps.
The QUESTION text is not in the result JSONL; it's loaded from the benchmark
dataset by ``record_id`` (via ``openjarvis.evals.cli._build_dataset``). If the
dataset can't be loaded / the id doesn't match, we render everything else and
show "(question unavailable)".
Gold answer lives in ``scoring_metadata`` (shape varies by scorer):
* MCQ (mmlu-pro, supergpqa): ``reference_letter``.
* LLM-judge (gaia): parsed from ``raw_judge_output`` (`gold target "..."`),
plus the judge's ``reasoning``.
* anything else: dumped verbatim under SCORING (raw).
Usage:
.venv/bin/python scripts/orchestrator/format_eval_sample.py \
--input results/orch-eval-gemma-2k-full/gaia_orchestrator.jsonl --n 2
... --input <file> --lines 1,5,42 # specific 1-indexed lines
... --input <file> --all # one .txt per record
... --input <file> --all --only-wrong # only incorrect samples
"""
from __future__ import annotations
import argparse
import json
import re
from pathlib import Path
WIDTH = 80
# Friendly aliases -> the registry keys _build_dataset expects (mirrors
# eval_orchestrator.BENCHMARK_ALIASES for the ones that write .jsonl here).
BENCHMARK_ALIASES = {
"mmlu_pro": "mmlu-pro",
"mmlu-pro": "mmlu-pro",
"gaia": "gaia",
"taubench": "taubench",
"supergpqa": "supergpqa",
"terminalbench_v2_1": "terminalbench-v2.1",
"terminalbench-v2.1": "terminalbench-v2.1",
}
def _banner(label: str) -> str:
"""Full-width headline: ``━━━ LABEL ━━━━━━…``."""
prefix = f"━━━ {label} "
fill = "" * max(3, WIDTH - len(prefix))
return f"\n{prefix}{fill}\n"
def _indent(text: str, pad: str = " ") -> str:
text = (text or "").rstrip()
return "\n".join(pad + ln if ln.strip() else ln for ln in text.splitlines())
def _normalize_benchmark(name: str) -> str:
key = (name or "").strip()
return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key))
# ---------------------------------------------------------------------------
# Question map: record_id -> problem text, loaded from the benchmark dataset.
# ---------------------------------------------------------------------------
def build_question_map(benchmark: str, n: int, seed: int = 42) -> dict:
"""Best-effort {record_id: problem}. Returns {} if the dataset can't load."""
try:
from openjarvis.evals.cli import _build_dataset
ds = _build_dataset(_normalize_benchmark(benchmark))
ds.load(max_samples=n, seed=seed)
return {r.record_id: r.problem for r in ds.iter_records()}
except Exception as exc: # noqa: BLE001 - never fail the render over this
print(
f" [warn] could not load dataset for {benchmark!r}: "
f"{type(exc).__name__}: {exc}"
)
return {}
# ---------------------------------------------------------------------------
# Gold / judge parsing out of scoring_metadata.
# ---------------------------------------------------------------------------
_GOLD_TARGET_RE = re.compile(r'gold target\s*"([^"]*)"', re.IGNORECASE)
_JUDGE_REASONING_RE = re.compile(
r"reasoning:\s*(.*?)(?:\n\s*correct:|\Z)", re.IGNORECASE | re.DOTALL
)
def parse_scoring(meta: dict) -> dict:
"""Return {gold, judge, kind, clean}. ``clean`` False => dump raw block.
kind is one of 'mcq' | 'judge' | 'unknown'.
"""
if not isinstance(meta, dict):
return {"gold": None, "judge": None, "kind": "unknown", "clean": False}
# MCQ scorer (mmlu-pro / supergpqa).
if "reference_letter" in meta:
return {
"gold": meta.get("reference_letter"),
"judge": None,
"kind": "mcq",
"clean": True,
}
# LLM-judge scorer (gaia).
raw = meta.get("raw_judge_output")
if raw:
gold_m = _GOLD_TARGET_RE.search(raw)
reason_m = _JUDGE_REASONING_RE.search(raw)
return {
"gold": gold_m.group(1) if gold_m else None,
"judge": (reason_m.group(1).strip() if reason_m else raw.strip()),
"kind": "judge",
"clean": True,
}
# Unknown shape (e.g. {"reason": "no_choice_letter_extracted"}).
return {"gold": None, "judge": None, "kind": "unknown", "clean": False}
# ---------------------------------------------------------------------------
# Render one record.
# ---------------------------------------------------------------------------
def format_record(rec: dict, question: str | None) -> str:
parsed = parse_scoring(rec.get("scoring_metadata"))
is_correct = rec.get("is_correct")
verdict = (
"CORRECT"
if is_correct
else ("INCORRECT" if is_correct is False else "UNSCORED")
)
def _fmt(v, fmt):
try:
return fmt.format(v)
except (ValueError, TypeError):
return str(v)
top = " · ".join(
[
str(rec.get("record_id", "?")),
str(rec.get("benchmark", "?")),
verdict,
f"score {_fmt(rec.get('score'), '{:.3f}')}",
f"lat {_fmt(rec.get('latency_seconds'), '{:.1f}')}s",
f"cost ${_fmt(rec.get('cost_usd'), '{:.4f}')}",
]
)
parts = [top]
err = rec.get("error")
if err:
parts.append(f" error: {err}")
parts.append(_banner("QUESTION"))
parts.append(_indent(question if question else "(question unavailable)"))
parts.append(_banner("MODEL ANSWER"))
parts.append(_indent(str(rec.get("model_answer") or "(empty)")))
parts.append(_banner("GOLD"))
gold = parsed["gold"]
parts.append(
_indent(
str(gold)
if gold not in (None, "")
else "(gold unavailable — see SCORING below)"
)
)
if parsed["kind"] == "judge" and parsed["judge"]:
parts.append(_banner("JUDGE"))
parts.append(_indent(parsed["judge"]))
if not parsed["clean"]:
parts.append(_banner("SCORING (raw)"))
parts.append(
_indent(json.dumps(rec.get("scoring_metadata"), indent=2, default=str))
)
return "\n".join(parts).rstrip() + "\n"
def main(argv=None):
p = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
p.add_argument("--input", required=True, help="Graded *_orchestrator.jsonl file.")
p.add_argument(
"--out-dir", default="results/formatted", help="Where the .txt files go."
)
p.add_argument(
"--seed", type=int, default=42, help="Subset seed used at eval time."
)
p.add_argument(
"--only-wrong",
action="store_true",
help="Render only incorrect/unscored samples.",
)
g = p.add_mutually_exclusive_group()
g.add_argument("--n", type=int, default=1, help="Format the first N records.")
g.add_argument("--lines", help="Comma-separated 1-indexed line numbers.")
g.add_argument("--all", action="store_true", help="Format every record.")
args = p.parse_args(argv)
src = Path(args.input)
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
records = [json.loads(l) for l in src.open() if l.strip()]
if not records:
print("no records in input")
return 1
# Pick indices first (so we only load the dataset once, sized to the file).
if args.lines:
idxs = [int(x) - 1 for x in args.lines.split(",") if x.strip()]
elif args.all:
idxs = list(range(len(records)))
else:
idxs = list(range(min(args.n, len(records))))
idxs = [i for i in idxs if 0 <= i < len(records)]
if args.only_wrong:
idxs = [i for i in idxs if not records[i].get("is_correct")]
benchmark = _normalize_benchmark(
records[0].get("benchmark", src.stem.split("_")[0])
)
# Load the whole subset so any selected id resolves (the eval-time subset is
# the first len(records) of seed=42).
qmap = build_question_map(benchmark, n=len(records), seed=args.seed)
written = []
for i in idxs:
rec = records[i]
rid = str(rec.get("record_id", i))
bench = rec.get("benchmark", benchmark)
fname = re.sub(r"[^A-Za-z0-9_.-]", "_", f"{bench}__{rid}")[:120] + ".txt"
out = out_dir / fname
out.write_text(format_record(rec, qmap.get(rid)))
written.append(out)
for w in written:
print(w)
print(
f"\nwrote {len(written)} file(s) to {out_dir}/"
+ (" (dataset unavailable — questions omitted)" if not qmap else "")
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+150
View File
@@ -0,0 +1,150 @@
#!/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 data/orchestrator/raw/) and writes the
SFT splits to data/orchestrator/sft/. Naming is uniform — ``{name}_{split}_{date}``
with name in {qwen, gemma, pooled}; pass several pools to merge-and-restratify
into `pooled`.
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}_0711.jsonl
python scripts/orchestrator/make_splits.py --name qwen \
--pool data/orchestrator/sft/qwen_clean_0711.jsonl
# merge both -> pooled_{train,holdout,overfit100}_0711.jsonl
python scripts/orchestrator/make_splits.py --name pooled \
--pool data/orchestrator/sft/qwen_clean_0711.jsonl \
--pool data/orchestrator/sft/gemma_clean_0711.jsonl
"""
import argparse
import json
import os
import random
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path
# 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(
"--date",
default=datetime.now().strftime("%m%d"),
help="date tag in the filename (default: today, MMDD)",
)
ap.add_argument("--out", type=Path, default=OUT)
args = ap.parse_args()
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)}")
# Stratified holdout: per-domain, proportional to domain size.
by_dom = defaultdict(list)
for i, r in enumerate(rows):
by_dom[r.get("domain", "misc")].append(i)
rng = random.Random(SEED)
holdout_idx: set = set()
for dom, idxs in sorted(by_dom.items()):
k = round(args.holdout_frac * len(idxs))
pick = rng.sample(idxs, min(k, len(idxs)))
holdout_idx.update(pick)
print(f" {dom:<9} pool={len(idxs):>4} holdout={len(pick)}")
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]
print(f"train={len(train)} holdout={len(holdout)} overfit={len(overfit)}")
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"{args.name}_{split}_{args.date}.jsonl"
f.write_text("".join(json.dumps(r) + "\n" for r in data))
written[split] = f
print(f"wrote {f} ({len(data)})")
# ---- auto-upload train + holdout to Braintrust ----
# Gated by OJ_BRAINTRUST_AUTOUPLOAD (default ON); a total no-op / never raises
# if the key/pkg is missing or the upload errors (see upload_to_braintrust).
sys.path.insert(0, str(Path(__file__).resolve().parent))
try:
from upload_to_braintrust import autoupload
autoupload(
[
f"{written['train']}={args.name}_train_{args.date}",
f"{written['holdout']}={args.name}_holdout_{args.date}",
],
run_label=os.getenv("OJ_RUN_LABEL", f"{args.name}_splits_{args.date}"),
description=f"{args.name} orchestrator-SFT splits (make_splits.py)",
)
except Exception as exc: # telemetry must never break the data pipeline
print(f"[braintrust] autoupload hook skipped ({exc})")
return 0
if __name__ == "__main__":
raise SystemExit(main())
+322
View File
@@ -0,0 +1,322 @@
#!/usr/bin/env python
"""Render an orchestrator SFT JSONL into human-viewable companions.
The trainer-facing ``*.jsonl`` packs every trajectory onto one line with all
newlines escaped, which is unreadable by eye. This writes two siblings next to
it so you can actually inspect what the orchestrator did:
* ``<name>.pretty.json`` — the same records as an indented JSON array.
* ``<name>.txt`` — a clean plain-text transcript. The shared system
prompt + tool catalog (identical across every record) is printed ONCE at the
top; each trajectory then shows just its turns: the user question, the
assistant's reasoning/answer, tool calls rendered as ``-> name({args})``, and
tool results truncated so a single web-search dump doesn't bury the trace.
Usage:
.venv/bin/python scripts/orchestrator/render_sft_data.py <path.jsonl>
"""
from __future__ import annotations
import html as _htmllib
import json
import re
import sys
from pathlib import Path
from typing import List
# Tool observations (web_search dumps, code stdout, file reads) can be huge; cap
# them in the transcript so the actual orchestration stays legible. The full
# content is always in the .jsonl / .pretty.json.
_OBS_CAP = 1200
_RULE = "-" * 80
_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
def _parse_tool_calls(content: str):
"""Parse the model's ``<tool_call>{json}</tool_call>`` tags into a list of
parsed ``[(name, arguments_dict), ...]`` calls, de-duplicating the
double-emit the template produces (each call shows up twice)."""
seen = []
for m in _TOOL_CALL_RE.finditer(content):
try:
call = json.loads(m.group(1))
item = (call.get("name"), call.get("arguments", {}))
except Exception:
item = (None, {"_raw": m.group(1)})
key = json.dumps(item, ensure_ascii=False, sort_keys=True, default=str)
if not seen or seen[-1][0] != key: # collapse the adjacent duplicate
seen.append((key, item))
return [it for _, it in seen]
def _call_label(name, anon_map) -> str:
"""Display a call name; if it's an anonymized expert label, append the real
model it maps to (from the record's ``metrics.anon_map``)."""
real = (anon_map or {}).get(name)
return f"{name} -> {real}" if real else f"{name}"
def _fmt_tool_calls(content: str, anon_map=None) -> str:
"""Plain-text tool calls with real (un-escaped) newlines in each arg value."""
blocks = []
for name, args in _parse_tool_calls(content):
lines = [f"-> {_call_label(name, anon_map)}"]
if isinstance(args, dict):
for k, v in args.items():
val = (
v
if isinstance(v, str)
else json.dumps(v, ensure_ascii=False, indent=2)
)
indented = "\n".join(" " + ln for ln in val.splitlines())
lines.append(f" {k}:\n{indented}")
blocks.append("\n".join(lines))
return "\n".join(blocks)
def _clip(text: str, cap: int) -> str:
text = text or ""
if len(text) <= cap:
return text
return f"{text[:cap]}\n ... [truncated {len(text) - cap} chars]"
def _render_turn(turn: dict, anon_map=None) -> str:
role = str(turn.get("role", "?")).lower()
content = turn.get("content", "") or ""
if role == "user":
return f"USER:\n{content.strip()}"
if role in ("tool", "function"):
name = turn.get("name", "tool")
return f"TOOL [{name}]:\n{_clip(content.strip(), _OBS_CAP)}"
if role == "assistant":
calls = _fmt_tool_calls(content, anon_map)
# Strip the raw <tool_call> tags out of the prose so it isn't shown twice.
prose = _TOOL_CALL_RE.sub("", content).strip()
parts = []
if prose:
parts.append(f"ASSISTANT:\n{prose}")
if calls:
parts.append(f"ASSISTANT calls:\n{calls}")
return "\n".join(parts) if parts else "ASSISTANT: (empty)"
# system handled separately (printed once); fall through for anything else
return f"{role.upper()}:\n{content.strip()}"
def _transcript(record: dict, index: int) -> str:
m = record.get("metrics", {}) or {}
correct = record.get("correct")
kept = record.get("kept")
flags = []
if correct is not None:
flags.append("correct" if correct else "WRONG")
if kept is not None and kept:
flags.append("kept")
head = (
f"### [{index}] task={record.get('task_id')} domain={record.get('domain')}"
+ (f" {' '.join(flags)}" if flags else "")
+ f"\n cost=${m.get('cost_usd')} tokens={m.get('tokens')} "
f"tool_calls={m.get('num_tool_calls')} turns={m.get('num_turns')} "
f"reward={record.get('reward')}"
)
anon_map = m.get("anon_map")
body = [
_render_turn(t, anon_map)
for t in record.get("conversations", [])
if str(t.get("role", "")).lower() != "system"
]
return head + "\n\n" + "\n\n".join(body)
def _system_header(records: List[dict]) -> str:
"""The system message is identical across records — render it once."""
for r in records:
for t in r.get("conversations", []):
if str(t.get("role", "")).lower() == "system":
return (
"=" * 80 + "\nSHARED SYSTEM PROMPT + TOOL CATALOG "
"(identical for every record below)\n"
+ "=" * 80
+ f"\n{t.get('content', '').strip()}\n"
)
return ""
_HTML_CSS = """
body{font:14px/1.5 system-ui,sans-serif;margin:0;background:#f5f5f7;color:#1d1d1f}
header{position:sticky;top:0;background:#1d1d1f;color:#fff;padding:12px 20px;z-index:9}
header b{font-size:16px}
.rec{background:#fff;margin:14px 20px;border-radius:10px;box-shadow:0 1px 3px #0002;overflow:hidden}
.rec>summary{cursor:pointer;padding:10px 16px;font-weight:600;list-style:none;display:flex;gap:14px;align-items:center;flex-wrap:wrap}
.rec>summary::-webkit-details-marker{display:none}
.rec>summary:hover{background:#fafafa}
.body{padding:6px 16px 16px}
.turn{margin:10px 0;padding:10px 14px;border-radius:8px;white-space:pre-wrap;word-break:break-word}
.user{background:#e8f0fe}
.assistant{background:#f1f8e9}
.call{background:#fff3e0}
.call .cname{font-family:ui-monospace,monospace;font-weight:700;color:#b25b00;margin-bottom:6px}
.call .cname .real{background:#b25b00;color:#fff;padding:1px 7px;border-radius:10px;font-size:12px;margin-left:4px}
.call .arg{margin:6px 0 0 14px}
.call .arg b{font-family:ui-monospace,monospace;color:#7a4500}
.call .arg pre{white-space:pre-wrap;word-break:break-word;margin:3px 0 0;padding:8px 10px;background:#fff;border:1px solid #f0d9b8;border-radius:6px;font-size:13px}
.tool details{background:#fafafa;border:1px solid #eee;border-radius:8px;padding:6px 10px;margin:10px 0}
.tool summary{cursor:pointer;color:#555;font-weight:600}
.tool pre{white-space:pre-wrap;word-break:break-word;margin:8px 0 0;font-size:13px}
.role{font-size:11px;text-transform:uppercase;letter-spacing:.5px;color:#888;margin-bottom:4px}
.pill{font-size:12px;padding:2px 8px;border-radius:10px;font-weight:600}
.ok{background:#d7f5dd;color:#1a7f37}.wrong{background:#ffe0e0;color:#c0392b}
.kept{background:#fff0c2;color:#9a6700}.dom{background:#eee;color:#444}
.meta{font-weight:400;color:#888;font-size:12px}
#sys details{margin:14px 20px;background:#fff;border-radius:10px;padding:10px 16px}
#sys pre{white-space:pre-wrap;word-break:break-word;font-size:12px;color:#444}
"""
def _esc(s: str) -> str:
return _htmllib.escape(s or "")
def _html_turns(record: dict) -> str:
anon_map = (record.get("metrics", {}) or {}).get("anon_map") or {}
out = []
for t in record.get("conversations", []):
role = str(t.get("role", "")).lower()
content = t.get("content", "") or ""
if role == "system":
continue
if role == "user":
out.append(
f'<div class="turn user"><div class="role">user</div>{_esc(content.strip())}</div>'
)
elif role in ("tool", "function"):
name = _esc(t.get("name", "tool"))
out.append(
f'<div class="tool"><details><summary>tool result · {name} '
f"({len(content)} chars)</summary><pre>{_esc(content.strip())}</pre></details></div>"
)
elif role == "assistant":
prose = _TOOL_CALL_RE.sub("", content).strip()
if prose:
out.append(
f'<div class="turn assistant"><div class="role">assistant</div>{_esc(prose)}</div>'
)
for name, args in _parse_tool_calls(content):
real = anon_map.get(name)
label = (
f'{_esc(str(name))} <span class="real">&rarr; {_esc(str(real))}</span>'
if real
else _esc(str(name))
)
rows = [f'<div class="cname">&rarr; {label}</div>']
if isinstance(args, dict):
for k, v in args.items():
val = (
v
if isinstance(v, str)
else json.dumps(v, ensure_ascii=False, indent=2)
)
rows.append(
f'<div class="arg"><b>{_esc(str(k))}</b><pre>{_esc(val)}</pre></div>'
)
out.append(
f'<div class="turn call"><div class="role">tool call</div>{"".join(rows)}</div>'
)
return "".join(out)
def _html_record(record: dict, index: int) -> str:
m = record.get("metrics", {}) or {}
pills = [f'<span class="pill dom">{_esc(str(record.get("domain")))}</span>']
correct = record.get("correct")
if correct is not None:
pills.append(
'<span class="pill ok">correct</span>'
if correct
else '<span class="pill wrong">wrong</span>'
)
if record.get("kept"):
pills.append('<span class="pill kept">kept</span>')
meta = (
f'<span class="meta">cost ${m.get("cost_usd")} · {m.get("tokens")} tok · '
f"{m.get('num_tool_calls')} calls · {m.get('num_turns')} turns</span>"
)
summary = (
f"<span>#{index}</span><span>{_esc(str(record.get('task_id')))}</span>"
+ "".join(pills)
+ meta
)
return (
f'<details class="rec"><summary>{summary}</summary>'
f'<div class="body">{_html_turns(record)}</div></details>'
)
def _html_doc(records: List[dict], title: str) -> str:
sys_txt = ""
for r in records:
for t in r.get("conversations", []):
if str(t.get("role", "")).lower() == "system":
sys_txt = t.get("content", "")
break
if sys_txt:
break
recs_html = "".join(_html_record(r, i + 1) for i, r in enumerate(records))
return (
f"<!doctype html><html><head><meta charset='utf-8'><title>{_esc(title)}</title>"
f"<style>{_HTML_CSS}</style></head><body>"
f"<header><b>{_esc(title)}</b> &nbsp; {len(records)} records "
f"&middot; click a row to expand</header>"
f"<div id='sys'><details><summary style='cursor:pointer;font-weight:600;padding:4px'>"
f"Shared system prompt + tool catalog (same for all)</summary>"
f"<pre>{_esc(sys_txt.strip())}</pre></details></div>"
f"{recs_html}</body></html>"
)
def render(jsonl_path: str) -> dict:
src = Path(jsonl_path)
records: List[dict] = []
with src.open() as fh:
for line in fh:
line = line.strip()
if line:
records.append(json.loads(line))
pretty = src.with_suffix(".pretty.json")
pretty.write_text(json.dumps(records, indent=2, ensure_ascii=False))
txt = src.with_suffix(".txt")
banner = f"ORCHESTRATOR SFT TRANSCRIPT — {src.name}\n{len(records)} records\n\n"
sep = f"\n\n{_RULE}\n\n"
parts = [banner + _system_header(records)]
parts.append(sep.join(_transcript(r, i + 1) for i, r in enumerate(records)))
txt.write_text("\n\n".join(parts))
htmlp = src.with_suffix(".html")
htmlp.write_text(_html_doc(records, src.name))
return {
"records": len(records),
"pretty": str(pretty),
"txt": str(txt),
"html": str(htmlp),
}
def main(argv: List[str]) -> int:
if not argv:
print(__doc__)
return 2
info = render(argv[0])
print(json.dumps(info, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main(sys.argv[1:]))
+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}
@@ -0,0 +1,439 @@
#!/usr/bin/env python
"""Upload orchestrator-SFT JSONL datasets to Braintrust so they're browsable/
filterable in the UI (by area / difficulty / dataset / routed-model / correct).
Datasets land in the SAME Braintrust project as the rollout TRACES (the
``research`` project) so a run's data + traces sit side by side. The target is
resolved from ``OJ_BRAINTRUST_PROJECT_ID`` (default: the research project id);
``--project`` overrides by NAME if you want a different project.
Each record -> a Braintrust dataset row:
input = the problem/question
expected = the FINAL_ANSWER value
tags = [gen_model, domain, correct|incorrect, clean|dirty, kept|dropped]
metadata = domain, area, difficulty, dataset, subsector, task_id, correct,
clean, kept, reward, gen_model, orchestrator_model, num_tool_calls,
n_turns, routed_models (real names)
(the full conversation is kept under metadata.conversation for inspection)
The dataset itself carries run-level metadata (run_label, specific gen_model,
git sha, config knobs, counts, domain distribution) so future experiments are
distinguishable in the UI.
Usage (CLI):
.venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707 \
data/orchestrator/sft/qwen_holdout_0707.jsonl=qwen_holdout_0707
# name is optional; defaults to {model_short}_{split}_{date}
.venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
data/orchestrator/sft/qwen_holdout_0707.jsonl
Programmatic (used by the pipeline auto-upload hook in make_splits.py):
from upload_to_braintrust import autoupload
autoupload(["data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707"],
run_label="...")
"""
import argparse
import json
import logging
import os
import re
import subprocess
from collections import Counter
from datetime import datetime
from pathlib import Path
logger = logging.getLogger(__name__)
# The `research` project — datasets go here, same place the rollout traces land.
DEFAULT_PROJECT_ID = "c707124e-ce9d-4187-ad11-f49f19f777ad"
_FA = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
def _ordered_turns(convs, anon_map, orchestrator):
"""Rebuild each turn as role/model/content, in that display order.
Braintrust sorts object keys alphabetically on write, which would push the
(very long) `content` to the top. We prefix the keys (`1_role`, `2_model`,
`3_content`) so they sort into role -> model -> content and the content
renders last. `model` is the model behind the turn:
* assistant turns -> the orchestrator model
* tool turns -> the expert model that answered (de-anonymized)
* system/user -> None
"""
out = []
for c in convs:
role = c.get("role")
if role == "tool":
model = anon_map.get(c.get("name")) or c.get("name")
elif role == "assistant":
model = orchestrator
else:
model = None
out.append({"1_role": role, "2_model": model, "3_content": c.get("content")})
return out
def _dataset_desc(dataset):
"""Fallback dataset blurb for records generated before the field existed."""
try:
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
dataset_description,
)
return dataset_description(dataset)
except Exception:
return ""
def _git_sha():
try:
return (
subprocess.check_output(
["git", "rev-parse", "--short", "HEAD"],
cwd=Path(__file__).resolve().parent,
stderr=subprocess.DEVNULL,
)
.decode()
.strip()
or None
)
except Exception:
return None
def _model_short(rows, path):
"""Short family tag for the default dataset name (qwen / gemma / ...)."""
stem = Path(path).stem.lower()
for fam in ("qwen", "gemma"):
if stem.startswith(fam) or fam in stem:
return fam
gm = next((r.get("gen_model") for r in rows if r.get("gen_model")), "") or ""
gm = gm.lower()
for fam in ("qwen", "gemma"):
if fam in gm:
return fam
return "orch"
def _split_of(path):
stem = Path(path).stem.lower()
for s in (
"holdout",
"train",
"clean",
"overfit",
"partial",
"8k",
"4k",
"2k",
"1k",
):
if s in stem:
return s
return "data"
def _default_name(rows, path):
return f"{_model_short(rows, path)}_{_split_of(path)}_{datetime.now():%m%d}"
def _config_from_env():
"""Config knobs from the run env (set by build_orchestrator_sft.py). Omit unset."""
cfg = {}
for env_key, meta_key, cast in [
("OJ_CFG_TEMPERATURE", "temperature", float),
("OJ_CFG_MAX_TURNS", "max_turns", int),
(
"OJ_CFG_ANONYMIZE",
"anonymize",
lambda v: v.strip().lower() in ("1", "true", "yes"),
),
(
"OJ_CFG_REJECTION_ONLY",
"rejection_only",
lambda v: v.strip().lower() in ("1", "true", "yes"),
),
]:
v = os.getenv(env_key)
if v not in (None, ""):
try:
cfg[meta_key] = cast(v)
except Exception:
cfg[meta_key] = v
return cfg
def _row(rec):
convs = rec.get("conversations", [])
user = next((c["content"] for c in convs if c["role"] == "user"), "")
question = re.sub(r"^\s*Problem:\s*", "", user, count=1).strip()
asst = [c["content"] for c in convs if c["role"] == "assistant"]
fa = asst[-1] if asst else ""
m = list(_FA.finditer(fa))
model_answer = fa[m[-1].end() :].strip() if m else fa.strip()
# Gold reference the verifier graded against (stamped at generation time as
# `gold_answer`). `expected` in Braintrust is the GOLD so the UI shows
# gold-vs-model; the model's own answer goes to metadata.model_answer.
gold = rec.get("gold_answer") or None
has_gold = gold is not None
am = (rec.get("metrics", {}) or {}).get("anon_map", {}) or {}
routed = []
for c in convs:
if c["role"] == "assistant":
for t in re.findall(
r"<tool_call>(\{.*?\})</tool_call>", c["content"], re.DOTALL
):
try:
real = am.get(json.loads(t).get("name"))
if real:
routed.append(real)
except Exception:
pass
metrics = rec.get("metrics", {}) or {}
gen_model = rec.get("gen_model")
domain = rec.get("domain")
correct = rec.get("correct")
clean = rec.get("clean")
kept = rec.get("kept")
tags = [
t
for t in [
gen_model,
f"domain:{domain}" if domain else None,
"correct" if correct else "incorrect",
"clean" if clean else "dirty",
"kept" if kept else "dropped",
"gold" if has_gold else "no_gold",
]
if t
]
return {
"input": question,
"expected": gold,
"tags": tags,
"metadata": {
"gold_answer": gold,
"model_answer": model_answer,
"has_gold": has_gold,
"domain": domain,
"area": rec.get("area"),
"difficulty": rec.get("difficulty") or None,
"dataset": rec.get("dataset"),
"dataset_description": rec.get("dataset_description")
or _dataset_desc(rec.get("dataset")),
"subsector": rec.get("subsector"),
"task_id": rec.get("task_id"),
"correct": correct,
"clean": clean,
"kept": kept,
"reward": rec.get("reward"),
"gen_model": gen_model,
"orchestrator_model": rec.get("orchestrator_model"),
"num_tool_calls": metrics.get("num_tool_calls"),
"n_turns": metrics.get("num_turns"),
"routed_models": routed,
"conversation": _ordered_turns(convs, am, rec.get("orchestrator_model")),
},
}
def _dist(rows, key):
"""{value: count} for a record field, most-common first, None/'' dropped."""
c = Counter(r.get(key) for r in rows if (r.get(key) or "") != "")
return {str(k): v for k, v in sorted(c.items(), key=lambda x: -x[1])}
def _routed_dist(rows):
"""Distribution of real (de-anonymized) expert models routed to."""
c = Counter()
for r in rows:
am = (r.get("metrics", {}) or {}).get("anon_map", {}) or {}
for cc in r.get("conversations", []):
if cc.get("role") != "assistant":
continue
for t in re.findall(
r"<tool_call>(\{.*?\})</tool_call>", cc.get("content", ""), re.DOTALL
):
try:
real = am.get(json.loads(t).get("name"))
if real:
c[real] += 1
except Exception:
pass
return {k: v for k, v in sorted(c.items(), key=lambda x: -x[1])}
def _avg(rows, *path):
vals = []
for r in rows:
v = r.get("metrics", {}) or {}
for p in path:
v = (v or {}).get(p) if isinstance(v, dict) else None
if isinstance(v, (int, float)):
vals.append(v)
return round(sum(vals) / len(vals), 2) if vals else None
def _dataset_metadata(rows, path, run_label, gen_model):
n = len(rows)
meta = {
"model": _model_short(rows, path),
"split": _split_of(path),
"task": "orchestrator-SFT unified-tool routing traces",
"date": "2026-07-07",
"run_label": run_label,
"gen_model": gen_model,
"git_sha": _git_sha(),
"source_file": str(path),
"source_datasets": ["GeneralThought-430K-filtered", "OpenThoughts3-1.2M"],
"task_mix": "balanced (GeneralThought + OpenThoughts code/math/science)",
"filter": "correct + clean (rejects file-write echoes, garble, "
"reasoning-degeneration, truncated tails, essays/markdown dumps)",
"leak_free": "train/holdout disjoint by task_id",
"n_total": n,
"n_correct": sum(1 for r in rows if r.get("correct")),
"n_clean": sum(1 for r in rows if r.get("clean")),
"n_kept": sum(1 for r in rows if r.get("kept")),
"n_with_gold": sum(1 for r in rows if (r.get("gold_answer") or "").strip()),
"avg_tool_calls": _avg(rows, "num_tool_calls"),
"avg_turns": _avg(rows, "num_turns"),
"area_distribution": _dist(rows, "area"),
"difficulty_distribution": _dist(rows, "difficulty"),
"dataset_distribution": _dist(rows, "dataset"),
"routed_model_distribution": _routed_dist(rows),
}
cfg = _config_from_env()
if cfg:
meta["config"] = cfg
return {k: v for k, v in meta.items() if v is not None}
def upload_dataset(
path,
name=None,
*,
project_id=None,
project=None,
run_label=None,
gen_model=None,
description="",
):
"""Upload one JSONL file as a Braintrust dataset. Returns (name, n_rows, url).
Target project: ``project`` (name) if given, else ``project_id`` /
``OJ_BRAINTRUST_PROJECT_ID`` / the research default.
"""
import braintrust
rows = [json.loads(l) for l in open(path) if l.strip()]
name = name or _default_name(rows, path)
gen_model = (
gen_model
or os.getenv("OJ_GEN_MODEL")
or next((r.get("gen_model") for r in rows if r.get("gen_model")), None)
)
run_label = run_label or os.getenv("OJ_RUN_LABEL") or name
init_kwargs = {
"name": name,
"description": description or None,
"metadata": _dataset_metadata(rows, path, run_label, gen_model),
}
if project:
init_kwargs["project"] = project
else:
init_kwargs["project_id"] = project_id or os.getenv(
"OJ_BRAINTRUST_PROJECT_ID", DEFAULT_PROJECT_ID
)
ds = braintrust.init_dataset(**init_kwargs)
for r in rows:
ds.insert(**_row(r))
ds.flush()
summ = ds.summarize()
url = getattr(summ, "dataset_url", None) or (
project or init_kwargs.get("project_id")
)
return name, len(rows), url
def autoupload(specs, *, run_label=None, gen_model=None, description=""):
"""No-op-safe wrapper for pipeline hooks. Honors OJ_BRAINTRUST_AUTOUPLOAD
(default ON) and never raises: any failure (missing key/pkg, network) is
logged and swallowed so the data pipeline is never broken by telemetry."""
if os.getenv("OJ_BRAINTRUST_AUTOUPLOAD", "1").strip().lower() in (
"0",
"false",
"no",
"off",
"",
):
logger.info("[braintrust] autoupload disabled (OJ_BRAINTRUST_AUTOUPLOAD)")
return
if not os.getenv("BRAINTRUST_API_KEY"):
logger.info("[braintrust] autoupload skipped — BRAINTRUST_API_KEY unset")
return
try:
import braintrust # noqa: F401
except Exception:
logger.info("[braintrust] autoupload skipped — braintrust not installed")
return
for spec in specs:
path, _, name = spec.partition("=")
try:
nm, n, url = upload_dataset(
path,
name or None,
run_label=run_label,
gen_model=gen_model,
description=description,
)
logger.info("[braintrust] uploaded %s: %d rows -> %s", nm, n, url)
print(f"[braintrust] uploaded {nm}: {n} rows -> {url}")
except Exception as exc:
logger.warning(
"[braintrust] autoupload FAILED for %s (%s) — continuing", path, exc
)
print(
f"[braintrust] autoupload FAILED for {path} ({exc}) — pipeline unaffected"
)
def main():
logging.basicConfig(level=logging.INFO, format="%(message)s")
p = argparse.ArgumentParser()
p.add_argument(
"--project",
default=None,
help="Target project by NAME (overrides OJ_BRAINTRUST_PROJECT_ID).",
)
p.add_argument(
"--project-id",
default=None,
help="Target project by id (default: OJ_BRAINTRUST_PROJECT_ID or research).",
)
p.add_argument("--run-label", default=None)
p.add_argument("--gen-model", default=None, help="Specific gen model id override.")
p.add_argument("--description", default="")
p.add_argument("specs", nargs="+", help="path.jsonl[=dataset_name]")
args = p.parse_args()
for spec in args.specs:
path, _, name = spec.partition("=")
nm, n, url = upload_dataset(
path,
name or None,
project_id=args.project_id,
project=args.project,
run_label=args.run_label,
gen_model=args.gen_model,
description=args.description,
)
print(f"[braintrust] {nm}: {n} rows -> {url}")
if __name__ == "__main__":
raise SystemExit(main())
+90
View File
@@ -0,0 +1,90 @@
#!/bin/bash
# Serve a trained gemma-4-26B-A4B SFT checkpoint (orchestrator) on 2xL40S (mkt1),
# then run eval_orchestrator over benchmarks. Self-contained: serve -> health -> eval -> teardown.
# Usage: sbatch eval_orch_gemma_matx.sbatch <1k|2k> <benchmarks> <N> <tag>
#SBATCH --job-name=eval_gemma
#SBATCH --account=mkt
#SBATCH --partition=mkt
#SBATCH -w mkt1
#SBATCH --gres=gpu:2
#SBATCH --cpus-per-gpu=10
#SBATCH --mem=200G
#SBATCH -t 24:00:00
#SBATCH --output=logs/eval_gemma_matx_%j.log
set -uo pipefail
SIZE=${1:?usage: <1k|2k> <benchmarks> <N> <tag>}
BENCHES=${2:-gaia}
N=${3:-2}
TAG=${4:-smoke}
PORT=8123
# Paths are resolved relative to the repo, or overridden by env:
# OJ_ROOT repo checkout (default: two levels up from this script)
# OJ_WORK working-state dir (default: $HOME/.openjarvis)
# OJ_CKPT_ROOT trained SFT checkpoints (default: $OJ_WORK/lambda_ckpts)
OJ_ROOT=${OJ_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}
OJ_WORK=${OJ_WORK:-$HOME/.openjarvis}
OJ_CKPT_ROOT=${OJ_CKPT_ROOT:-$OJ_WORK/lambda_ckpts}
cd "$OJ_ROOT"
set -a
source .env
[ -f "$OJ_WORK/eval_key_override.sh" ] && source "$OJ_WORK/eval_key_override.sh"
set +a
export HF_HOME=${HF_HOME:?set HF_HOME to a cache dir with room for the weights}
export PYTHONPATH="$PWD/src"
# judge/expert cloud engine needs a config with an [engine.cloud] section
export OPENJARVIS_CONFIG=${OPENJARVIS_CONFIG:?set OPENJARVIS_CONFIG to a cell config.toml}
source .venv/bin/activate
if [ "$SIZE" = "base" ]; then
CKPT=google/gemma-4-26B-A4B-it # un-fine-tuned base orchestrator (from HF cache)
else
CKPT=$OJ_CKPT_ROOT/sft_gemma_${SIZE}_bf16/epoch3
fi
OUT=results/orch-eval-gemma-${SIZE}-${TAG}
mkdir -p "$OUT" logs
echo "=== [$(date)] eval_gemma SIZE=$SIZE BENCHES=$BENCHES N=$N TAG=$TAG on $(hostname) ==="
nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader
# ---- serve: TP=2 bf16 + MANDATORY no-NVLink L40S flags (else NCCL hangs) ----
echo "=== [$(date)] launching vLLM (TP=2) serve of $CKPT ==="
NCCL_P2P_DISABLE=1 NCCL_CUMEM_HOST_ENABLE=0 VLLM_WORKER_MULTIPROC_METHOD=spawn \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python -m vllm.entrypoints.openai.api_server \
--model "$CKPT" --served-model-name gemma-sft \
--port $PORT --tensor-parallel-size 2 --dtype bfloat16 \
--gpu-memory-utilization 0.90 --trust-remote-code \
--max-model-len 32768 --enforce-eager --disable-custom-all-reduce \
--enable-auto-tool-choice --tool-call-parser gemma4 \
--max-num-batched-tokens 8192 --safetensors-load-strategy eager \
--enable-prefix-caching > logs/serve_gemma_${SIZE}_${SLURM_JOB_ID}.log 2>&1 &
SERVE_PID=$!
echo "serve pid=$SERVE_PID log=logs/serve_gemma_${SIZE}_${SLURM_JOB_ID}.log"
cleanup(){ echo "=== [$(date)] teardown kill $SERVE_PID ==="; kill $SERVE_PID 2>/dev/null; sleep 5; kill -9 $SERVE_PID 2>/dev/null; }
trap cleanup EXIT
# ---- health poll (up to 25 min; TP + 52GB load is slow on L40S) ----
echo "=== [$(date)] waiting for /v1/models ==="
READY=0
for i in $(seq 1 150); do
if ! kill -0 $SERVE_PID 2>/dev/null; then echo "!! serve died during startup"; tail -40 logs/serve_gemma_${SIZE}_${SLURM_JOB_ID}.log; exit 3; fi
if curl -sf http://localhost:$PORT/v1/models 2>/dev/null | grep -q gemma-sft; then READY=1; echo "=== [$(date)] READY after ${i}0s ==="; break; fi
sleep 10
done
[ $READY -eq 1 ] || { echo "!! not ready in 25min"; tail -40 logs/serve_gemma_${SIZE}_${SLURM_JOB_ID}.log; exit 4; }
# ---- eval ----
echo "=== [$(date)] running eval_orchestrator benches=$BENCHES n=$N ==="
python scripts/orchestrator/eval_orchestrator.py \
--benchmarks "$BENCHES" --n "$N" --seed 42 \
--orchestrator-endpoint http://localhost:$PORT/v1 \
--orchestrator-model gemma-sft --orchestrator-api-key EMPTY \
--judge-model gemini-2.5-flash --judge-engine cloud \
--max-workers 4 --max-turns 8 \
--output-dir "$OUT"
RC=$?
echo "=== [$(date)] eval exit=$RC summary=$OUT/summary.json ==="
echo "EVAL_GEMMA_DONE size=$SIZE tag=$TAG rc=$RC"
exit $RC
+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,979 @@
"""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 random
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_LOCAL_SEARCH = "local_search"
KIND_CODE = "code_interpreter"
KIND_TOOL = "tool" # a bridged real OpenJarvis tool (custom param schema)
VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_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.
# (The legacy default_catalog still uses CATEGORY_GENERALIST/SPECIALIZED below.)
CATEGORY_CLOUD_FRONTIER = "cloud_frontier"
CATEGORY_LOCAL_OSS = "local_open_source"
# Legacy tiers kept only for default_catalog (build_unified_sft.py); not used by
# orchestrator_catalog. Superseded by the two-class taxonomy above.
CATEGORY_GENERALIST = "generalist_model"
CATEGORY_SPECIALIZED = "specialized_model"
# 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 or self.kind == KIND_LOCAL_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
# Default catalog: the paper's tool categories, mapped onto the models/tools
# OpenJarvis can actually call. One named tool per model (faithful §3.1).
def default_catalog(
*,
local_model: Optional[str] = None,
local_endpoint: Optional[str] = None,
) -> List[ExpertTool]:
"""Return the full unified tool catalog.
``local_model`` / ``local_endpoint`` wire the on-device vLLM tool when a
local backbone is served; omitted → the local model tool is left out.
"""
cat: List[ExpertTool] = []
# ---- generalist / frontier models (one named tool per model VERSION) ----
for model, summary, lat in [
(
"gpt-5",
"Frontier generalist (GPT-5). Strongest reasoning across domains.",
30.0,
),
(
"gpt-5-mini",
"Mid-tier generalist (GPT-5-mini). Solid reasoning, much cheaper.",
15.0,
),
(
"gpt-4o",
"Fast generalist (GPT-4o). Good for simple steps and formatting.",
8.0,
),
(
"claude-opus-4-7",
"Frontier generalist (Claude Opus 4.7). Strong long-horizon reasoning.",
26.0,
),
(
"claude-sonnet-4-6",
"Strong generalist (Claude Sonnet 4.6). Balanced cost/capability.",
15.0,
),
(
"gemini-2.5-pro",
"Frontier generalist (Gemini 2.5 Pro). Strong multimodal reasoning.",
20.0,
),
("gemini-2.5-flash", "Cheap fast generalist (Gemini 2.5 Flash).", 8.0),
(
"meta-llama/llama-3.3-70b-instruct",
"Open generalist (Llama-3.3-70B). Decent general knowledge, low cost.",
10.0,
),
(
"qwen/qwen3-32b",
"Open generalist (Qwen3-32B). Strong math/science reasoning, low cost.",
9.0,
),
]:
ep = (
"openai"
if model.startswith("gpt")
else "anthropic"
if model.startswith("claude")
else "gemini"
if model.startswith("gemini")
else "openrouter"
)
pi, po = _price(model)
cat.append(
ExpertTool(
name=_tool_name(model),
kind=KIND_MODEL,
backend_type=ep,
summary=summary,
model=model,
price_in=pi,
price_out=po,
latency_s=lat,
category=CATEGORY_GENERALIST,
)
)
# ---- specialized: code ----
coder = "qwen/qwen-2.5-coder-32b-instruct"
pi, po = _price(coder)
cat.append(
ExpertTool(
name=_tool_name(coder),
kind=KIND_MODEL,
backend_type="openrouter",
summary="Specialized code model (Qwen2.5-Coder-32B). Writes/debugs code.",
model=coder,
price_in=pi,
price_out=po,
latency_s=9.0,
category=CATEGORY_SPECIALIZED,
)
)
# ---- local backbone as a tool (on-device vLLM), if served ----
# Named after the actual served model (faithful "one named tool per model"),
# not a generic "local_model" — e.g. "qwen3-8b" -> tool "qwen3_8b".
if local_model and local_endpoint:
cat.append(
ExpertTool(
name=_tool_name(local_model),
kind=KIND_MODEL,
backend_type="vllm",
summary=(
f"On-device open model ({local_model}) served locally. Cheap "
"and private; good for extraction, formatting, arithmetic on "
"given data."
),
model=local_model,
base_url=local_endpoint,
price_in=0.0,
price_out=0.0,
latency_s=2.0,
category=CATEGORY_GENERALIST,
)
)
# ---- 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,
)
)
return cat
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 sample_tool_config(
catalog: List[ExpertTool],
*,
rng: random.Random,
min_tools: int = 4,
max_tools: Optional[int] = None,
price_jitter: float = 0.0,
) -> List[ExpertTool]:
"""Sample a random tool subset with optional price randomization (§3.3).
Guarantees at least one ``model`` tool and at least one non-model (basic)
tool so every instance can both reason and act. ``price_jitter`` (e.g. 0.5)
multiplies each model's prices by a per-tool factor drawn uniformly from
``[1-jitter, 1+jitter]``, modeling heterogeneous pricing across users.
Deterministic given ``rng``.
"""
if not catalog:
raise ValueError("empty catalog")
models = [t for t in catalog if t.kind == KIND_MODEL]
basics = [t for t in catalog if t.kind != KIND_MODEL]
if not models:
raise ValueError("catalog has no model tools")
hi = max_tools if max_tools is not None else len(catalog)
hi = min(hi, len(catalog))
lo = min(max(min_tools, 2), hi)
k = rng.randint(lo, hi)
# Always include >=1 model; include >=1 basic if any exist.
chosen: List[ExpertTool] = [rng.choice(models)]
if basics:
chosen.append(rng.choice(basics))
pool = [t for t in catalog if t not in chosen]
rng.shuffle(pool)
for t in pool:
if len(chosen) >= k:
break
chosen.append(t)
# Re-order to catalog order for stable specs.
order = {t.name: i for i, t in enumerate(catalog)}
chosen.sort(key=lambda t: order[t.name])
if price_jitter > 0.0:
jittered: List[ExpertTool] = []
for t in chosen:
if t.kind == KIND_MODEL and (t.price_in or t.price_out):
f = rng.uniform(1.0 - price_jitter, 1.0 + price_jitter)
jittered.append(
ExpertTool(
name=t.name,
kind=t.kind,
backend_type=t.backend_type,
summary=t.summary,
model=t.model,
base_url=t.base_url,
price_in=round(t.price_in * f, 4),
price_out=round(t.price_out * f, 4),
latency_s=t.latency_s,
)
)
else:
jittered.append(t)
return jittered
return chosen
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_GENERALIST",
"CATEGORY_LOCAL_OSS",
"CATEGORY_SPECIALIZED",
"ExpertTool",
"KIND_CODE",
"KIND_LOCAL_SEARCH",
"KIND_MODEL",
"KIND_TOOL",
"KIND_WEB_SEARCH",
"build_tool_specs",
"default_catalog",
"openjarvis_tool",
"orchestrator_catalog",
"sample_tool_config",
"to_worker_dict",
"tools_by_name",
]
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,382 @@
"""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
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple
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 (the <tool_call> tag in
# content, matching the SFT serialization), then the observation as a
# distinct `tool` turn the model reads as a tool response.
call_content = text or ""
if "<tool_call>" not in call_content:
call_content = (
call_content + "\n" + tool_call_tag(name, arguments)
).strip()
messages.append({"role": "assistant", "content": call_content})
obs_text = obs or ""
if len(obs_text) > _OBS_CAP:
obs_text = obs_text[:_OBS_CAP] + "\n…[truncated]"
messages.append({"role": "tool", "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,279 @@
"""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 ExpertTool, to_worker_dict
from openjarvis.agents.hybrid.toolorchestra.parsing import _parse_rl_tool_call
from openjarvis.agents.hybrid.toolorchestra.rollout import (
UnifiedRollout,
build_system_prompt,
run_unified_rollout,
)
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
)
resp = client.chat.completions.create(
model=model,
messages=send,
temperature=temperature,
max_tokens=max_tokens,
frequency_penalty=frequency_penalty,
presence_penalty=presence_penalty,
**kwargs,
)
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,
)
text, p, c, is_local, extra_cost, _n = _call_worker(worker, prompt, cfg)
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
def teacher_rollout(
question: str,
tools: List[ExpertTool],
*,
teacher_model: str,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
temperature: float = 1.0,
max_turns: int = 50,
cfg: Optional[Dict[str, Any]] = None,
) -> UnifiedRollout:
"""Convenience: one full teacher rollout with real backends."""
return run_unified_rollout(
question,
tools,
call_orchestrator=make_call_orchestrator(
teacher_model,
base_url=base_url,
api_key=api_key,
temperature=temperature,
),
dispatch=make_dispatch(cfg),
max_turns=max_turns,
)
__all__ = ["make_call_orchestrator", "make_dispatch", "teacher_rollout"]
@@ -0,0 +1,526 @@
"""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.experts import (
_PAPER_CODER_OPENROUTER,
_PAPER_GENERALIST_TIER3_OPENROUTER,
)
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
_call_modal_python,
_call_tavily_search,
)
def _paper_pool(
local_model: Optional[str],
local_endpoint: Optional[str],
) -> List[Dict[str, Any]]:
"""Paper-match worker pool (registered for traces / inspection).
NOTE: in RL mode the orchestrator dispatches via tool/slot rather than
worker_id, so this list is purely informational `_paper_expert_for`
is the actual routing function. We still return a list here so the
paradigm's trace metadata has something concrete to log.
"""
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": "Local Qwen vLLM (paper uses Qwen3-32B).",
}
)
pool.append(
{
"id": len(pool),
"name": "tavily-search",
"type": "tavily-search",
"model": "tavily",
"description": "Tavily web search.",
}
)
pool.append(
{
"id": len(pool),
"name": "modal-python",
"type": "modal-python",
"model": "modal-python",
"description": "Modal Sandbox for one-shot Python exec.",
}
)
pool.append(
{
"id": len(pool),
"name": "code-specialist",
"type": "openrouter",
"model": _PAPER_CODER_OPENROUTER,
"description": "Qwen-2.5-Coder-32B via OpenRouter (paper).",
}
)
pool.append(
{
"id": len(pool),
"name": "generalist-llama",
"type": "openrouter",
"model": _PAPER_GENERALIST_TIER3_OPENROUTER,
"description": "Llama-3.3-70B-Instruct via OpenRouter (paper tier-3).",
}
)
pool.append(
{
"id": len(pool),
"name": "generalist-gpt5",
"type": "openai",
"model": "gpt-5",
"description": "GPT-5 frontier generalist.",
}
)
pool.append(
{
"id": len(pool),
"name": "generalist-gpt5-mini",
"type": "openai",
"model": "gpt-5-mini",
"description": "GPT-5-mini mid generalist.",
}
)
return pool
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
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":
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":
text, p, c = LocalCloudAgent._call_openrouter(
worker["model"],
user=prompt,
max_tokens=max_tok,
temperature=temp,
)
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: gemma-4-12B-it. v1 = base self-sampling (the served orchestrator_model is this same base).
#
# v1 data = base gemma-4-12B-it 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_gemma-4-12b_sft_v1.jsonl --samples-per-task 8
# then train on a rented GPU with regenerate_traces = false.
[model]
model_name = "google/gemma-4-12B-it"
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_gemma-4-12b_sft_v1.jsonl"
regenerate_traces = false
# Base gemma-4-12B-it self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
orchestrator_endpoint = "http://localhost:8001/v1"
orchestrator_model = "gemma-4-12b"
samples_per_task = 8
max_keep_per_task = 1
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_gemma-4-12b_sft"
save_every_n_epochs = 1
@@ -0,0 +1,34 @@
# Orchestrator SFT cold-start — Qwen3-8B (~9B, "easier to train").
#
# v1 data = base Qwen3-8B 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_sft_v1.jsonl --samples-per-task 8
# then train on a rented GPU with regenerate_traces = false.
[model]
model_name = "Qwen/Qwen3-8B"
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_sft_v1.jsonl"
regenerate_traces = false
# Base Qwen3-8B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
orchestrator_endpoint = "http://localhost:8001/v1"
orchestrator_model = "qwen3-8b"
samples_per_task = 8
max_keep_per_task = 1
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_sft"
save_every_n_epochs = 1
@@ -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,238 @@
"""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 __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,
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)
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,
# Serving a fine-tuned model: JSON <tool_call> format in the system
# prompt, no tools= param (matches how it was trained/serialized).
native_tools=False,
)
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,
)
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"]
@@ -143,6 +143,61 @@ class MultiObjectiveReward:
return [self.compute(ep) for ep in episodes]
@dataclass
class CostAwareReward:
"""Cost-aware GRPO reward: ``r = +1/-1`` for correctness, minus a
cost penalty.
This is the headline objective from the plan::
base = +1.0 if episode.correct else -1.0
reward = base - lambda * (total_cost_usd / cost_max)
``lambda`` (``lam``) is the cost-penalty weight and is *swept* across a
set of values (see :func:`lambda_sweep`) to trace the accuracy/cost
Pareto frontier. ``cost_max`` normalises USD cost so the penalty is
O(1) when an episode hits the per-task cost budget.
"""
lam: float = 0.0
"""Cost-penalty weight (swept, e.g. [0.0, 0.05, 0.1, 0.2, 0.4])."""
cost_max: float = 0.10
"""USD cost normalizer (per-task cost budget)."""
def compute(self, episode: Episode) -> float:
"""Scalar reward: ``+1/-1`` for correctness minus a cost term."""
base = 1.0 if episode.correct else -1.0
return base - self.lam * (episode.total_cost_usd / self.cost_max)
def compute_with_breakdown(self, episode: Episode) -> Dict[str, float]:
"""Reward with per-component breakdown."""
correct = 1.0 if episode.correct else 0.0
base = 1.0 if episode.correct else -1.0
cost_term = -self.lam * (episode.total_cost_usd / self.cost_max)
return {
"correct": correct,
"base": base,
"cost_term": cost_term,
"reward": base + cost_term,
}
def compute_batch(self, episodes: List[Episode]) -> List[float]:
"""Compute rewards for a batch of episodes."""
return [self.compute(ep) for ep in episodes]
def lambda_sweep(
values: List[float], cost_max: float = 0.10
) -> List["CostAwareReward"]:
"""Build one :class:`CostAwareReward` per lambda value.
Used to sweep the cost-penalty weight and trace the accuracy/cost
Pareto frontier (the ``R -= lambda * cost / cost_max`` objective).
"""
return [CostAwareReward(lam=lam, cost_max=cost_max) for lam in values]
class AdaptiveRewardWeights:
"""Adaptive reward weights that shift during training.
@@ -211,7 +266,9 @@ class AdaptiveRewardWeights:
__all__ = [
"AdaptiveRewardWeights",
"CostAwareReward",
"MultiObjectiveReward",
"Normalizers",
"RewardWeights",
"lambda_sweep",
]
@@ -0,0 +1,44 @@
"""SFT-data generation for the orchestrator cold-start.
Execution-grounded rejection sampling in the faithful ToolOrchestra action
space (arXiv:2511.21689): for each ``nvidia/ToolScale`` task, 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::
ToolScale task -> N teacher rollouts -> verify -> keep cheapest passing
-> conversations JSONL
(The earlier ADP-relabel cold-start was a heuristic re-tiering of demonstrated
traces; it was removed in favour of this grounded pipeline.)
"""
from __future__ import annotations
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
gold_coverage_verify,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
GoldAction,
ToolScaleTask,
load_toolscale,
normalize_row,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
__all__ = [
"GoldAction",
"ToolScaleTask",
"generate_sft_dataset",
"gold_coverage_verify",
"load_toolscale",
"normalize_row",
"trajectory_to_record",
]
@@ -0,0 +1,355 @@
"""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
task_id = str(row.get("question_id") or f"generalthought-{index}")
return Task(
task_id=task_id,
question=question,
answer=answer,
domain=_generalthought_domain(row),
difficulty=str(row.get("difficulty") or "").strip(), # usually absent for GT
dataset="GeneralThought",
subsector=str(row.get("question_source") or "").strip(),
)
def load_generalthought(
*,
n: int = 2000,
seed: int = 42,
source: Optional[Iterable[Dict[str, Any]]] = None,
buffer: Optional[int] = None,
) -> 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.
"""
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] = []
for i, row in enumerate(source):
task = _normalize_generalthought(dict(row), index=i)
if task is None:
continue
buf.append(task)
if len(buf) >= buf_cap:
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()
body = visible or full
answer = _extract_boxed(body) or _extract_boxed(full) or body.strip()
if not answer:
return None
return question, answer
def _normalize_openthoughts(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
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"
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
def load_grpo_prompts(*, n: int = 30000, seed: int = 42) -> List[Task]:
"""Pool ``n`` unique prompts from the same datasets, deduped by question.
We draw a roughly even split from GeneralThought and OpenThoughts3 (code /
math / science), dedupe on the question text, and cap at ``n``.
"""
per = max(n // 4 + 1, 1)
pool: List[Task] = []
pool.extend(load_generalthought(n=per, seed=seed))
pool.extend(load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed))
rng = random.Random(seed)
rng.shuffle(pool)
seen: set[str] = set()
out: List[Task] = []
for task in pool:
key = task.question.strip()
if not key or key in seen:
continue
seen.add(key)
out.append(task)
if len(out) >= n:
break
return out
__all__ = [
"GENERALTHOUGHT_ID",
"OPENTHOUGHTS_ID",
"Task",
"load_generalthought",
"load_grpo_prompts",
"load_openthoughts",
"load_sft_tasks",
]
@@ -0,0 +1,460 @@
"""Rejection-sampling SFT-data generator (the ToolOrchestra cold-start).
For each ToolScale 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?
:func:`gold_coverage_verify` is a dependency-free default verifier (checks the
trajectory's tool calls cover the task's golden action names); a real run should
compose it with an LLM judge on 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 — ToolScaleTask or the
# reasoning ``datasets.Task``. ``gold_coverage_verify`` additionally needs
# ``gold_action_names()`` (ToolScale-only), but callers using ``datasets.Task`` pass
# their own ``verify_fn`` (e.g. ``verify.make_verifier()``) instead.
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*$")
def _target_is_clean(roll: UnifiedRollout) -> bool:
"""Structural gate on whether a trajectory is fit to be an SFT *target*.
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 False
if any(m in fa for m in _ERR_MARKERS):
return False
if fa.count("<think>") != fa.count(
"</think>"
): # unbalanced think tags (truncated OR stray </think>)
return False
if _TRUNCATED_TAIL_RE.search(fa): # truncated mid-expression (e.g. "a=1, b=")
return False
# 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 False
# 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 False
if _CONTROL_TOKEN_RE.search(fa_stripped):
return False
# 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 False
# 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 False
# Markdown-table dump: the model wrote a formatted table/essay instead of the
# short exact value the format demands. A ``|---|`` separator row is the tell.
if re.search(r"\|\s*:?-{2,}", fa):
return False
# 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
if len(_ans) > 2000: # was 700 — too tight for code/math
return False
# 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 False
if (
len(re.findall(r"(?m)^\s*\d+\.\s", _ans)) >= 6
): # many numbered sections = essay (was 4)
return False
if re.search(r"(?m)^\s*#{1,6}\s", _ans): # markdown headers = essay
return False
# 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:
_alpha = [c for c in _ans if c.isalpha()]
if (
_alpha
and sum(c.isupper() for c in _alpha) / len(_alpha) > 0.7
and len(_ans.split()) >= 3
):
return False
if any(len(w) > 25 and w.isalpha() for w in _ans.split()):
return False
# 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 False
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 False
# "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 False
elif roll.num_tool_calls < 1:
return False
for t in roll.turns:
if t.tool_name is not None:
obs = (t.observation or "").strip()
if not obs or any(m in obs for m in _ERR_MARKERS):
return False
# 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 False
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 False
return True
def gold_coverage_verify(task: TaskLike, rollout: UnifiedRollout) -> bool:
"""Dependency-free proxy verifier: trajectory must (a) produce a non-empty
answer and (b) call tools covering every golden action name.
This is the offline stand-in for ToolScale's execution-correctness checker
(which needs the DB simulator). Compose with an LLM judge for real runs.
"""
if not rollout.final_answer.strip():
return False
gold = set(task.gold_action_names())
if not gold:
return True
called = {name for name, _ in rollout.tool_calls()}
return gold.issubset(called)
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 = gold_coverage_verify,
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
record["clean"] = _target_is_clean(roll)
# 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", "gold_coverage_verify"]
@@ -0,0 +1,134 @@
"""Loader for NVIDIA ToolScale (``nvidia/ToolScale``) — the ToolOrchestra
RL/SFT task source (arXiv:2511.21689 §3.3).
Each row is a synthetic user-agent-tool task: an instruction ``I``, golden
function calls ``A`` (the ground-truth tool sequence), and short info ``o`` that
must be communicated. We normalize the raw HF row into :class:`ToolScaleTask`.
``load_toolscale`` streams via the HuggingFace ``datasets`` library; tests pass
``source=`` an iterable of raw row dicts so normalization is exercised offline.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, Iterator, List, Optional
DATASET_ID = "nvidia/ToolScale"
@dataclass
class GoldAction:
name: str
arguments: Dict[str, Any] = field(default_factory=dict)
action_id: Optional[str] = None
@dataclass
class ToolScaleTask:
task_id: str
domain: str
instruction: str
gold_actions: List[GoldAction] = field(default_factory=list)
required_info: List[str] = field(default_factory=list)
nl_assertions: List[str] = field(default_factory=list)
def gold_action_names(self) -> List[str]:
return [a.name for a in self.gold_actions]
def _as_list(v: Any) -> List[Any]:
if v is None:
return []
if isinstance(v, list):
return v
return [v]
def _str_list(v: Any) -> List[str]:
out: List[str] = []
for item in _as_list(v):
if isinstance(item, str):
out.append(item)
elif isinstance(item, dict):
# communicate_info entries are sometimes {"info": "..."} dicts.
for key in ("info", "content", "text", "value"):
if isinstance(item.get(key), str):
out.append(item[key])
break
return out
def normalize_row(row: Dict[str, Any], *, index: int = 0) -> ToolScaleTask:
"""Turn one raw ToolScale row into a :class:`ToolScaleTask` (pure)."""
scenario = row.get("user_scenario") or {}
instructions = scenario.get("instructions") or {}
instruction = (
instructions.get("task_instructions")
or instructions.get("reason_for_call")
or row.get("task")
or row.get("instruction")
or ""
)
domain = scenario.get("domain") or row.get("domain") or "unknown"
crit = row.get("evaluation_criteria") or {}
gold: List[GoldAction] = []
for a in _as_list(crit.get("actions")):
if isinstance(a, dict) and a.get("name"):
gold.append(
GoldAction(
name=str(a["name"]),
arguments=a.get("arguments") or a.get("args") or {},
action_id=a.get("action_id"),
)
)
required = _str_list(crit.get("communicate_info"))
nl = _str_list(crit.get("nl_assertions"))
task_id = str(row.get("id") or row.get("task_id") or f"toolscale-{index}")
return ToolScaleTask(
task_id=task_id,
domain=str(domain),
instruction=str(instruction),
gold_actions=gold,
required_info=required,
nl_assertions=nl,
)
def load_toolscale(
*,
max_tasks: Optional[int] = None,
split: str = "train",
source: Optional[Iterable[Dict[str, Any]]] = None,
) -> Iterator[ToolScaleTask]:
"""Yield normalized ToolScale tasks.
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
When ``source`` is None, streams ``nvidia/ToolScale`` via ``datasets``.
"""
if source is None:
from datasets import load_dataset # lazy: optional dep / network
source = load_dataset(DATASET_ID, split=split, streaming=True)
n = 0
for i, row in enumerate(source):
if max_tasks is not None and n >= max_tasks:
break
task = normalize_row(dict(row), index=i)
if not task.instruction.strip():
continue
yield task
n += 1
__all__ = [
"DATASET_ID",
"GoldAction",
"ToolScaleTask",
"load_toolscale",
"normalize_row",
]
@@ -0,0 +1,269 @@
"""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 (
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(),
}
)
conversations.append(
{
"role": "tool",
"name": turn.tool_name,
"content": turn.observation or "",
}
)
# 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,255 @@
"""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 os
import re
import string
from typing import Any, Callable, Dict, Optional
from .datasets import Task
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
GEMINI_MODEL = "gemini-2.5-flash"
# 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.
JUDGE_MODEL = "claude-haiku-4-5-20251001"
# 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
# --- 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
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", "")
s = s.strip().strip("$ ")
return s
def _to_float(s: str) -> Optional[float]:
s = s.replace(",", "").strip()
try:
return float(s)
except (ValueError, TypeError):
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 _gemini_judge(task: Task, prediction: str) -> Optional[bool]:
"""Ask Gemini PASS/FAIL. Returns None if unavailable (caller falls back)."""
api_key = os.environ.get("ANTHROPIC_API_KEY")
if not api_key:
return None
try:
import anthropic
except Exception:
return None
try:
# Judge = Anthropic Haiku (fast ~0.6s, direct PASS/FAIL, high rate limits).
# Switched off the Gemini judge: its free-tier rate limit + SDK retry/backoff
# blocked worker threads for minutes under parallel generation and stalled
# the run. Fail-fast (no retries, short timeout): on error return None and the
# caller falls back to string/f1.
global _JUDGE_CLIENT
if _JUDGE_CLIENT is None:
_JUDGE_CLIENT = anthropic.Anthropic(
api_key=api_key, max_retries=0, timeout=15
)
client = _JUDGE_CLIENT
prompt = (
"You are grading a candidate answer against a gold reference.\n"
"Reply with exactly one word: PASS if the candidate is correct and "
"consistent with the gold answer, otherwise FAIL.\n\n"
f"Question:\n{task.question}\n\n"
f"Gold answer:\n{task.answer}\n\n"
f"Candidate answer:\n{prediction}\n\n"
"Verdict (PASS or FAIL):"
)
resp = client.messages.create(
model=JUDGE_MODEL,
max_tokens=8,
messages=[{"role": "user", "content": prompt}],
)
verdict = (resp.content[0].text or "").strip().upper()
if "PASS" in verdict:
return True
if "FAIL" in verdict:
return False
return None
except Exception:
return None
# --- public API --------------------------------------------------------------
def verify_answer(task: Task, prediction: str) -> bool:
"""Domain-dispatched correctness check for ``prediction`` vs ``task.answer``."""
pred = (prediction or "").strip()
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__ = [
"GEMINI_BASE_URL",
"GEMINI_MODEL",
"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
+290
View File
@@ -0,0 +1,290 @@
"""Tests for the faithful ToolOrchestra unified-tool registry."""
from __future__ import annotations
import random
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,
default_catalog,
openjarvis_tool,
orchestrator_catalog,
sample_tool_config,
to_worker_dict,
tools_by_name,
)
def test_each_model_is_its_own_tool():
"""Faithful §3.1: one named tool per model, not a meta-tool + slot."""
cat = default_catalog()
names = {t.name for t in cat}
# Distinct model tools, each with its own name.
for n in (
"gpt_5",
"gpt_5_mini",
"qwen3_32b",
"qwen_2_5_coder_32b_instruct",
"llama_3_3_70b_instruct",
"claude_opus_4_7",
):
assert n in names, f"missing model tool {n}"
# No meta-tool / slot vocabulary leaks in.
assert "answer" not in names and "enhance_reasoning" not in names
def test_catalog_names_unique_and_valid():
cat = default_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_local_model_included_only_when_served():
# Local tool is named after the served model ("qwen3:8b" -> "qwen3_8b").
assert "qwen3_8b" not in {t.name for t in default_catalog()}
cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
local = tools_by_name(cat)["qwen3_8b"]
assert local.backend_type == "vllm"
assert local.base_url == "http://x/v1"
assert local.price_in == 0.0 and local.price_out == 0.0
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 = default_catalog()
specs = build_tool_specs(cat)
by = {s["function"]["name"]: s for s in specs}
gpt5 = by["gpt_5"]
assert gpt5["type"] == "function"
assert "input" in gpt5["function"]["parameters"]["properties"]
# Price table is surfaced in the description (the policy is trained on it).
assert "$1.25/1M input" in gpt5["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"]
def test_sample_is_deterministic_and_well_formed():
cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
a = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
b = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
assert [t.name for t in a] == [t.name for t in b] # deterministic
assert len(a) >= 4
assert any(t.kind == KIND_MODEL for t in a) # can reason
assert any(t.kind != KIND_MODEL for t in a) # can act
assert {t.name for t in a} <= {t.name for t in cat} # subset
def test_price_jitter_changes_prices_reproducibly():
cat = default_catalog()
base = {
t.name: t for t in sample_tool_config(cat, rng=random.Random(3), min_tools=8)
}
jit = {
t.name: t
for t in sample_tool_config(
cat, rng=random.Random(3), min_tools=8, price_jitter=0.5
)
}
# Same subset (same seed/sequence up to jitter draws), but model prices move.
moved = [
n
for n in base
if base[n].kind == KIND_MODEL
and base[n].price_in
and n in jit
and jit[n].price_in != base[n].price_in
]
assert moved, "expected jitter to change at least one model price"
for n in moved:
f = jit[n].price_in / base[n].price_in
assert 0.5 <= f <= 1.5
# 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 = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
by = tools_by_name(cat)
assert to_worker_dict(by["gpt_5"]) == {
"name": "gpt_5",
"type": "openai",
"model": "gpt-5",
}
local = to_worker_dict(by["qwen3_8b"])
assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"
@@ -0,0 +1,137 @@
"""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": "stackexchange_codegolf",
"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": []},
]
def test_load_generalthought_fields_and_domains():
tasks = list(load_generalthought(n=10, source=GENERALTHOUGHT_ROWS))
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"
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,194 @@
"""Offline tests for the rejection-sampling SFT pipeline (unified tools)."""
from __future__ import annotations
import json
from openjarvis.agents.hybrid.expert_registry import default_catalog, tools_by_name
from openjarvis.agents.hybrid.toolorchestra.rollout import (
UnifiedRollout,
UnifiedTurn,
run_unified_rollout,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
gold_coverage_verify,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
normalize_row,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
# A representative raw ToolScale row.
_RAW_ROW = {
"id": "movie-001",
"user_scenario": {
"domain": "entertainment",
"instructions": {"task_instructions": "Cancel ticket A03 and refund the user."},
},
"evaluation_criteria": {
"actions": [
{"name": "cancel", "arguments": {"booking": "A03"}, "action_id": "x1"},
{"name": "refund", "arguments": {"user": "8612"}, "action_id": "x2"},
],
"communicate_info": ["refund amount is $20.90"],
"nl_assertions": ["the ticket is cancelled"],
},
}
def test_normalize_row():
t = normalize_row(_RAW_ROW)
assert t.task_id == "movie-001"
assert t.domain == "entertainment"
assert "Cancel ticket A03" in t.instruction
assert t.gold_action_names() == ["cancel", "refund"]
assert t.required_info == ["refund amount is $20.90"]
def test_run_unified_rollout_terminates_on_no_tool_call():
tools = default_catalog()
by = tools_by_name(tools)
name = "qwen3_32b"
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 = default_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_gold_coverage_verify():
t = normalize_row(_RAW_ROW)
good = UnifiedRollout(
turns=[
UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
],
final_answer="done",
)
missing = UnifiedRollout(
turns=[UnifiedTurn("", "cancel", {}, "ok")], final_answer="done"
)
empty_ans = UnifiedRollout(
turns=[
UnifiedTurn("", "cancel", {}, "ok"),
UnifiedTurn("", "refund", {}, "ok"),
],
final_answer="",
)
assert gold_coverage_verify(t, good) is True
assert gold_coverage_verify(t, missing) is False
assert gold_coverage_verify(t, empty_ans) is False
def test_generate_sft_dataset_end_to_end(tmp_path):
tools = default_catalog()
tasks = [
normalize_row(_RAW_ROW),
normalize_row(
{
**_RAW_ROW,
"id": "unsolvable",
"evaluation_criteria": {"actions": [{"name": "never_called"}]},
}
),
]
def rollout_fn(task):
# Solve the first task; always miss the gold action of the second.
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,
)
out = tmp_path / "sft.jsonl"
stats = generate_sft_dataset(
str(out),
tasks=tasks,
tools=tools,
rollout_fn=rollout_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,155 @@
"""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).
# The driver treats --out as a TAG and always writes to
# <OJ_DATA_ROOT>/raw/<label>_<MMDD>[_<tag>]/data.jsonl, defaulting to a
# cwd-relative data/orchestrator. Drop OJ_DATA_ROOT (the repo .env exports
# it) so the run folder lands in the test's temp dir and not the real
# experiments tree, then chdir there.
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
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
@@ -0,0 +1,86 @@
"""Tests for the cost-aware GRPO reward (offline, no GPU)."""
from __future__ import annotations
from dataclasses import dataclass
from openjarvis.learning.intelligence.orchestrator.reward import (
CostAwareReward,
lambda_sweep,
)
@dataclass
class FakeEpisode:
"""Minimal stand-in exposing the two fields CostAwareReward reads."""
correct: bool
total_cost_usd: float = 0.0
class TestCostAwareReward:
def test_correct_zero_cost_is_plus_one(self):
r = CostAwareReward(lam=0.0)
assert r.compute(FakeEpisode(correct=True, total_cost_usd=0.0)) == 1.0
def test_incorrect_zero_cost_is_minus_one(self):
r = CostAwareReward(lam=0.0)
assert r.compute(FakeEpisode(correct=False, total_cost_usd=0.0)) == -1.0
def test_correct_with_lam_penalty_when_cost(self):
# base=+1, penalty = lam * cost / cost_max = 0.2 * 0.05 / 0.10 = 0.1
r = CostAwareReward(lam=0.2, cost_max=0.10)
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.05))
assert abs(val - (1.0 - 0.1)) < 1e-9
def test_lam_zero_ignores_cost(self):
r = CostAwareReward(lam=0.0, cost_max=0.10)
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.5))
assert val == 1.0
def test_cost_at_budget_full_penalty(self):
# cost == cost_max → penalty == lam exactly.
r = CostAwareReward(lam=0.4, cost_max=0.10)
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.10))
assert abs(val - (1.0 - 0.4)) < 1e-9
def test_incorrect_also_penalised(self):
r = CostAwareReward(lam=0.2, cost_max=0.10)
val = r.compute(FakeEpisode(correct=False, total_cost_usd=0.05))
assert abs(val - (-1.0 - 0.1)) < 1e-9
def test_breakdown_fields(self):
r = CostAwareReward(lam=0.2, cost_max=0.10)
b = r.compute_with_breakdown(FakeEpisode(correct=True, total_cost_usd=0.05))
assert set(b) == {"correct", "base", "cost_term", "reward"}
assert b["correct"] == 1.0
assert b["base"] == 1.0
assert abs(b["cost_term"] - (-0.1)) < 1e-9
assert abs(b["reward"] - 0.9) < 1e-9
def test_breakdown_reward_matches_compute(self):
r = CostAwareReward(lam=0.3, cost_max=0.10)
ep = FakeEpisode(correct=False, total_cost_usd=0.07)
assert abs(r.compute_with_breakdown(ep)["reward"] - r.compute(ep)) < 1e-12
def test_compute_batch(self):
r = CostAwareReward(lam=0.0)
eps = [FakeEpisode(correct=True), FakeEpisode(correct=False)]
assert r.compute_batch(eps) == [1.0, -1.0]
class TestLambdaSweep:
def test_sweep_length(self):
values = [0.0, 0.05, 0.1, 0.2, 0.4]
rewards = lambda_sweep(values)
assert len(rewards) == len(values)
def test_sweep_assigns_lambdas(self):
values = [0.0, 0.05, 0.1, 0.2, 0.4]
rewards = lambda_sweep(values, cost_max=0.25)
assert [r.lam for r in rewards] == values
assert all(r.cost_max == 0.25 for r in rewards)
def test_sweep_returns_costaware(self):
rewards = lambda_sweep([0.1])
assert isinstance(rewards[0], CostAwareReward)
@@ -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