Compare commits

...
Author SHA1 Message Date
Andrew Park 08fe3a6c98 evals: switch the default LLM judge to claude-haiku-4-5
Replaces gpt-5-mini-2025-08-07 as the default judge everywhere: the 147
eval configs plus the hardcoded defaults in the evals CLI, recipe composer,
comparison config generator, taubench user-simulator, trial runner and
OptimizeConfig.

Haiku is roughly 3x faster on the judging path and we are out of OpenAI
quota, so gpt-5-mini was failing closed on long eval runs.
2026-07-11 11:59:32 -07:00
Andrew Park f338d64b21 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
  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.

Paths in the sbatch and litellm helpers are now env-driven rather than
hardcoded to one cluster home.
2026-07-11 11:55:59 -07:00
Andrew Park 6862b50da8 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-06-18 20:31:40 -07:00
Andrew Park 08916ec468 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-06-17 11:35:55 -07:00
Andrew Park 28ef7e5cb5 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-06-05 10:27:02 -07:00
228 changed files with 10731 additions and 2018 deletions
+7
View File
@@ -130,3 +130,10 @@ 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/
+74
View File
@@ -0,0 +1,74 @@
# LiteLLM proxy config for OpenJarvis ToolOrchestra expert spend tracking.
# Launch: litellm --config litellm/config.yaml --port 4000 --host 127.0.0.1
# Keys are read from process env (source OpenJarvis/.env first).
# No Postgres on this box (docker socket denied) -> admin UI/virtual keys are
# disabled; spend is written to litellm/logs/spend.jsonl by spend_logger.py.
model_list:
# ---- OpenAI ----
- model_name: gpt-5.5
litellm_params:
model: openai/gpt-5.5
api_key: os.environ/OPENAI_API_KEY
- model_name: gpt-5-mini
litellm_params:
model: openai/gpt-5-mini
api_key: os.environ/OPENAI_API_KEY
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
# ---- Anthropic ----
- model_name: claude-opus-4-8
litellm_params:
model: anthropic/claude-opus-4-8
api_key: os.environ/ANTHROPIC_API_KEY
- model_name: claude-sonnet-4-6
litellm_params:
model: anthropic/claude-sonnet-4-6
api_key: os.environ/ANTHROPIC_API_KEY
# ---- Gemini ----
- model_name: gemini-2.5-pro
litellm_params:
model: gemini/gemini-2.5-pro
api_key: os.environ/GEMINI_API_KEY
- model_name: gemini-2.5-flash
litellm_params:
model: gemini/gemini-2.5-flash
api_key: os.environ/GEMINI_API_KEY
# ---- OpenRouter (Qwen coder etc.) ----
- model_name: qwen-coder
litellm_params:
model: openrouter/qwen/qwen-2.5-coder-32b-instruct
api_key: os.environ/OPENROUTER_API_KEY
- model_name: openrouter/*
litellm_params:
model: openrouter/*
api_key: os.environ/OPENROUTER_API_KEY
# ---- Local vLLM (self-hosted students / experts) ----
- model_name: qwen3.5-9b-local
litellm_params:
model: openai/Qwen/Qwen3.5-9B
api_base: http://localhost:8011/v1
api_key: EMPTY
input_cost_per_token: 0.0
output_cost_per_token: 0.0
- model_name: qwen3.6-27b-local
litellm_params:
model: openai/Qwen/Qwen3.6-27B-FP8
api_base: http://localhost:8002/v1
api_key: EMPTY
input_cost_per_token: 0.0
output_cost_per_token: 0.0
litellm_settings:
# JSONL spend logger (Postgres-free spend tracking).
callbacks: spend_logger.proxy_handler_instance
drop_params: true
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
+87
View File
@@ -0,0 +1,87 @@
"""Custom LiteLLM callback that appends per-request cost/usage to a JSONL file.
This is the spend-tracking fallback used when no Postgres is available for the
LiteLLM admin UI. Every successful (and failed) proxy request writes one line to
$OJ_LITELLM_SPEND_LOG. Aggregate with `litellm/spend_report.py`.
"""
from __future__ import annotations
import datetime
import json
import os
import threading
from litellm.integrations.custom_logger import CustomLogger
LOG_PATH = os.environ.get(
"OJ_LITELLM_SPEND_LOG",
os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs", "spend.jsonl"),
)
os.makedirs(os.path.dirname(LOG_PATH), exist_ok=True)
_lock = threading.Lock()
def _write(kwargs, response_obj, start_time, end_time, status: str) -> None:
try:
cost = kwargs.get("response_cost")
litellm_params = kwargs.get("litellm_params") or {}
metadata = litellm_params.get("metadata") or {}
usage = None
try:
usage = getattr(response_obj, "usage", None) or (
response_obj.get("usage") if isinstance(response_obj, dict) else None
)
except Exception:
usage = None
def _u(field):
if usage is None:
return None
return getattr(usage, field, None) if not isinstance(usage, dict) else usage.get(field)
dur = None
try:
if start_time and end_time:
dur = (end_time - start_time).total_seconds()
except Exception:
dur = None
rec = {
"ts": datetime.datetime.utcnow().isoformat() + "Z",
"status": status,
"model": kwargs.get("model"),
"model_group": metadata.get("model_group"),
"call_type": kwargs.get("call_type"),
"cost_usd": cost,
"prompt_tokens": _u("prompt_tokens"),
"completion_tokens": _u("completion_tokens"),
"total_tokens": _u("total_tokens"),
"duration_s": dur,
"key_alias": metadata.get("user_api_key_alias"),
"exception": str(kwargs.get("exception")) if status == "failure" else None,
}
line = json.dumps(rec)
with _lock:
with open(LOG_PATH, "a") as f:
f.write(line + "\n")
except Exception as exc: # never let logging break a request
print(f"[spend_logger] error: {exc}")
class SpendLogger(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
_write(kwargs, response_obj, start_time, end_time, "success")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
_write(kwargs, response_obj, start_time, end_time, "success")
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
_write(kwargs, response_obj, start_time, end_time, "failure")
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
_write(kwargs, response_obj, start_time, end_time, "failure")
proxy_handler_instance = SpendLogger()
+43
View File
@@ -0,0 +1,43 @@
#!/usr/bin/env python3
"""Aggregate litellm/logs/spend.jsonl into a per-model spend summary.
Postgres-free stand-in for the LiteLLM admin UI. Usage:
.venv/bin/python litellm/spend_report.py [path-to-spend.jsonl]
"""
import json
import os
import sys
from collections import defaultdict
path = sys.argv[1] if len(sys.argv) > 1 else os.environ.get(
"OJ_LITELLM_SPEND_LOG",
os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs", "spend.jsonl"),
)
agg = defaultdict(lambda: {"calls": 0, "ok": 0, "fail": 0, "cost": 0.0,
"ptok": 0, "ctok": 0})
total_cost = 0.0
with open(path) as f:
for line in f:
line = line.strip()
if not line:
continue
r = json.loads(line)
g = r.get("model_group") or r.get("model") or "?"
a = agg[g]
a["calls"] += 1
a["ok"] += 1 if r.get("status") == "success" else 0
a["fail"] += 1 if r.get("status") == "failure" else 0
a["cost"] += float(r.get("cost_usd") or 0)
a["ptok"] += int(r.get("prompt_tokens") or 0)
a["ctok"] += int(r.get("completion_tokens") or 0)
total_cost += float(r.get("cost_usd") or 0)
hdr = f"{'model':<24}{'calls':>7}{'ok':>5}{'fail':>6}{'p_tok':>9}{'c_tok':>9}{'cost_usd':>12}"
print(hdr)
print("-" * len(hdr))
for g, a in sorted(agg.items(), key=lambda kv: -kv[1]["cost"]):
print(f"{g:<24}{a['calls']:>7}{a['ok']:>5}{a['fail']:>6}"
f"{a['ptok']:>9}{a['ctok']:>9}{a['cost']:>12.6f}")
print("-" * len(hdr))
print(f"{'TOTAL':<24}{'':>7}{'':>5}{'':>6}{'':>9}{'':>9}{total_cost:>12.6f}")
+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,319 @@
#!/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 \
--out data/orchestrator_sft_v1.jsonl \
--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
)
# Default output is timestamped (e.g. data/orchestrator_sft_06-22-1230am.jsonl)
# so every run lands in its own file and concurrent runs never collide. Pass
# --out explicitly to override (sharded runs do this, one file per shard).
p.add_argument("--out", default=None,
help="Output JSONL. Default: data/orchestrator_sft_<MM-DD-HHMMam>.jsonl")
# 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 datetimed folder, all collected flat under
# data/runs/ — so runs sort by time, never collide, and never litter data/.
# --out is treated as an optional LABEL for the folder, not a path; the data
# file inside is always data.jsonl (+ .stats.json/.pretty.json/.txt/.html).
# Folder name: <MM-DD-HHMMam/pm>_<label> (local PT) e.g. 06-24-0349pm_gemma_20_anon
stamp = time.strftime("%m-%d-%I%M%p").lower()
label = Path(args.out).stem if args.out else "orchestrator_sft"
run_dir = (Path("data/runs") / f"{stamp}_{label}").resolve()
if run_dir.exists(): # parallel runs, same label+minute -> disambiguate
run_dir = run_dir.with_name(f"{run_dir.name}-{os.getpid()}")
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"{stamp} 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())
+109
View File
@@ -0,0 +1,109 @@
#!/usr/bin/env python
"""Interactive REPL to chat with a trained orchestrator checkpoint.
Reuses the SAME routing loop the eval uses (OrchestratorBackend.generate_full),
so what you see here is exactly how it behaves on the benchmarks: it reads your
question, decides whether to answer itself / call a tool / delegate to an expert
model, runs that rollout, and prints the final answer + a routing summary.
Prereqs:
1. Serve the checkpoint you want on a vLLM port (see --endpoint). e.g. 8k:
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True CUDA_VISIBLE_DEVICES=0 \
.venv/bin/python -m vllm.entrypoints.openai.api_server \
--model "$OJ_WORK/qwen_eval_ckpts/8k_served" \
--served-model-name sft-qwen-8k --port 8020 \
--gpu-memory-utilization 0.88 --trust-remote-code --max-model-len 32768 \
--enforce-eager --enable-auto-tool-choice --tool-call-parser qwen3_xml
2. source .env (cloud experts GPT-5.5 / Opus need the API keys)
Usage:
.venv/bin/python scripts/orchestrator/chat_orchestrator.py \
--endpoint http://localhost:8020/v1 --model sft-qwen-8k
# point local expert tiers at their own vLLMs if you have them up (optional):
# --local-endpoint Qwen/Qwen3.6-27B=http://localhost:8002/v1
"""
import argparse
import sys
import time
def _parse_args(argv=None):
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--endpoint", default="http://localhost:8020/v1",
help="vLLM endpoint serving the orchestrator checkpoint.")
p.add_argument("--model", default="sft-qwen-8k",
help="Served-model-name of the checkpoint.")
p.add_argument("--api-key", default="EMPTY", help="API key for the endpoint (local vLLM = EMPTY).")
p.add_argument("--temperature", type=float, default=0.0)
p.add_argument("--max-turns", type=int, default=8, help="Same default as the eval.")
p.add_argument("--max-tokens", type=int, default=3000,
help="Keep high: the model reasons a lot before it emits the delegation.")
p.add_argument("--local-endpoint", action="append", default=[],
help="MODEL_ID=URL for a local expert tier. Repeatable.")
return p.parse_args(argv)
def main(argv=None):
args = _parse_args(argv)
local_endpoints = {}
for pair in args.local_endpoint:
if "=" not in pair:
raise SystemExit(f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}")
mid, url = pair.split("=", 1)
local_endpoints[mid.strip()] = url.strip()
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
OrchestratorBackend,
)
backend = OrchestratorBackend(
orchestrator_endpoint=args.endpoint,
orchestrator_model=args.model,
api_key=args.api_key,
local_endpoints=local_endpoints,
max_turns=args.max_turns,
temperature=args.temperature,
)
print(f"orchestrator: {args.model} @ {args.endpoint} "
f"(temp={args.temperature}, max_turns={args.max_turns})")
print("type a question and hit enter. Ctrl-D or 'quit' to exit.\n")
try:
while True:
try:
prompt = input("you> ").strip()
except EOFError:
print()
break
if not prompt:
continue
if prompt.lower() in {"quit", "exit"}:
break
started = time.time()
full = backend.generate_full(
prompt, model=args.model, temperature=args.temperature,
max_tokens=args.max_tokens,
)
elapsed = time.time() - started
if full.get("error"):
print(f"\n[ERROR] {full['error']}\n", file=sys.stderr)
continue
print(f"\norch> {full.get('content', '').strip()}\n")
print(f" [turns={full.get('turn_count', '?')} "
f"tool_calls={full.get('tool_calls', '?')} "
f"cost=${full.get('cost_usd', 0.0):.4f} "
f"{elapsed:.1f}s]\n")
finally:
backend.close()
return 0
if __name__ == "__main__":
raise SystemExit(main())
+345
View File
@@ -0,0 +1,345 @@
#!/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())
+239
View File
@@ -0,0 +1,239 @@
#!/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())
+120
View File
@@ -0,0 +1,120 @@
#!/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. See data/DELETED_ERA1_MANIFEST.md.
Naming is uniform: ``{name}_orch_{split}_{date}.jsonl`` 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_orch_{train,holdout,overfit100}_0711.jsonl
python scripts/orchestrator/make_splits.py --name qwen \
--pool data/sft_variants/qwen_orch_clean_0711.jsonl
# merge both -> pooled_orch_{train,holdout,overfit100}_0711.jsonl
python scripts/orchestrator/make_splits.py --name pooled \
--pool data/sft_variants/qwen_orch_clean_0711.jsonl \
--pool data/sft_variants/gemma_orch_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
OUT = Path("data/sft_variants")
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}_orch_{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())
+290
View File
@@ -0,0 +1,290 @@
#!/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:]))
+338
View File
@@ -0,0 +1,338 @@
#!/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/runs/<...>/data.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 torch.utils.data import DataLoader
from accelerate import Accelerator
from accelerate.utils import set_seed, InitProcessGroupKwargs
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())
@@ -0,0 +1,13 @@
compute_environment: LOCAL_MACHINE
distributed_type: FSDP
mixed_precision: bf16
num_machines: 1
num_processes: 4
fsdp_config:
fsdp_sharding_strategy: FULL_SHARD
fsdp_auto_wrap_policy: SIZE_BASED_WRAP
fsdp_min_num_params: 100000000
fsdp_state_dict_type: FULL_STATE_DICT
fsdp_use_orig_params: true
fsdp_cpu_ram_efficient_loading: false
fsdp_sync_module_states: false
+148
View File
@@ -0,0 +1,148 @@
#!/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,361 @@
#!/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/sft_variants/qwen_orch_train_0707.jsonl=qwen_train_0707 \
data/sft_variants/qwen_orch_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/sft_variants/qwen_orch_holdout_0707.jsonl
Programmatic (used by the pipeline auto-upload hook in make_splits.py):
from upload_to_braintrust import autoupload
autoupload(["data/sft_variants/qwen_orch_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
@@ -394,8 +394,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).
@@ -471,7 +471,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;
@@ -537,7 +537,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]:
+14
View File
@@ -11,10 +11,12 @@ 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),
"claude-haiku-4-5-20251001": (1.00, 5.00),
"gpt-5.5": (2.00, 15.0),
"gpt-5": (1.25, 10.0),
"gpt-5-mini": (0.25, 2.00),
"gpt-5-mini-2025-08-07": (0.25, 2.00),
@@ -30,6 +32,14 @@ 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. ESTIMATES — no public list price
# exists yet for these Qwen3.5/3.6 builds; scaled by active-param size.
# VERIFY against openrouter.ai before trusting the cost-aware reward numbers.
"qwen/qwen3.5-9b": (0.05, 0.10),
"qwen/qwen3.6-27b": (0.10, 0.30),
"qwen/qwen3.5-122b-a10b": (0.20, 0.60),
"qwen/qwen3.5-397b-a17b": (0.40, 1.20),
}
# Models whose API rejects an explicit `temperature` param — callers should
@@ -38,6 +48,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,783 @@
"""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, field
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,771 @@
"""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 json
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,
run_swe_agent_loop,
)
from openjarvis.core.registry import AgentRegistry
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.experts import (
_PAPER_CODER_OPENROUTER,
_expert_for,
_paper_expert_for,
)
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
_call_modal_python,
_extract_first_python_block,
)
from openjarvis.agents.hybrid.toolorchestra.clients import (
_call_orchestrator_with_tool_calls,
)
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.workers import (
_call_worker,
_default_pool,
_resolve_worker_pool,
_swe_call_worker,
)
@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") != "anthropic-web-search"
] 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,50 @@
"""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,170 @@
"""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-*` -> always the Anthropic web_search tool (the
upstream uses Tavily; we have web_search)
"""
if slot.startswith("search"):
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,163 @@
"""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,106 @@
"""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,324 @@
"""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,79 @@
"""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,161 @@
"""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, Optional
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,255 @@
"""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
# 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()
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.workers import _call_worker
from openjarvis.agents.hybrid.toolorchestra.tracing import span
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,422 @@
"""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,
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."
),
})
pool.append({
"id": len(pool),
"name": "web-search",
"type": "anthropic-web-search",
"model": "claude-haiku-4-5",
"description": (
"Anthropic server-side web_search. 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", "gemini",
"tavily-search", "openrouter", "modal-python",
)
# 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``. ``type`` must be one
of ``vllm`` / ``openai`` / ``anthropic`` / ``anthropic-web-search``.
``anthropic-web-search`` entries may omit ``model`` it defaults to
``claude-haiku-4-5``.
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 == "anthropic-web-search":
if model in (None, ""):
model = _DEFAULT_WEB_SEARCH_MODEL
entry["model"] = model
elif not isinstance(model, str):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
)
# 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)"
)
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 == "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 == "anthropic-web-search":
# 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))
+1 -1
View File
@@ -1408,7 +1408,7 @@ class OptimizeConfig:
optimizer_provider: str = "anthropic"
benchmark: str = ""
max_samples: int = 50
judge_model: str = "gpt-5-mini-2025-08-07"
judge_model: str = "claude-haiku-4-5-20251001"
db_path: str = str(DEFAULT_CONFIG_DIR / "optimize.db")
+2 -2
View File
@@ -1141,7 +1141,7 @@ def main():
"-n", "--max-samples", type=int, default=None, help="Maximum samples to evaluate"
)
@click.option("-w", "--max-workers", type=int, default=4, help="Parallel workers")
@click.option("--judge-model", default="gpt-5-mini-2025-08-07", help="LLM judge model")
@click.option("--judge-model", default="claude-haiku-4-5-20251001", help="LLM judge model")
@click.option("-o", "--output", "output_path", default=None, help="Output JSONL path")
@click.option("--seed", type=int, default=42, help="Random seed")
@click.option("--split", "dataset_split", default=None, help="Dataset split override")
@@ -1387,7 +1387,7 @@ def run(
"-n", "--max-samples", type=int, default=None, help="Max samples per benchmark"
)
@click.option("-w", "--max-workers", type=int, default=4, help="Parallel workers")
@click.option("--judge-model", default="gpt-5-mini-2025-08-07", help="LLM judge model")
@click.option("--judge-model", default="claude-haiku-4-5-20251001", help="LLM judge model")
@click.option("--output-dir", default="results/", help="Output directory for results")
@click.option("--seed", type=int, default=42, help="Random seed")
@click.option("-v", "--verbose", is_flag=True, help="Verbose logging")
@@ -98,7 +98,7 @@ BENCHMARKS: Dict[str, Dict[str, object]] = {
},
}
DEFAULT_JUDGE_MODEL = "gpt-5-mini-2025-08-07"
DEFAULT_JUDGE_MODEL = "claude-haiku-4-5-20251001"
def _template_path() -> Path:
@@ -11,7 +11,7 @@ temperature = 0.6
max_tokens = 32768
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 1024
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
+1 -1
View File
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 1.0
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 8192
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 2048
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 2048
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 2048
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 1024
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 2048
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 1024
@@ -12,7 +12,7 @@ temperature = 0.6
max_tokens = 32768
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -12,7 +12,7 @@ temperature = 1.0 # Kimi-K2.5 thinking mode recommended
max_tokens = 32768
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -11,7 +11,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 1024
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -17,7 +17,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -14,7 +14,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -10,7 +10,7 @@ temperature = 0.0
max_tokens = 4096
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"
@@ -8,7 +8,7 @@ temperature = 0.6
max_tokens = 16384
[judge]
model = "gpt-5-mini-2025-08-07"
model = "claude-haiku-4-5-20251001"
temperature = 0.0
max_tokens = 4096
engine = "cloud"

Some files were not shown because too many files have changed in this diff Show More