mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-08-14 08:52:06 +00:00
Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
90dad24a9e | ||
|
|
c9c25895a8 | ||
|
|
3bffa33938 | ||
|
|
82a2d3edde | ||
|
|
9571a158a2 | ||
|
|
4864ea7776 | ||
|
|
e6c47e9532 | ||
|
|
f3d9705961 | ||
|
|
5aa16473ff | ||
|
|
a6afeae721 |
+10
@@ -130,3 +130,13 @@ learning.db
|
||||
minion_logs/
|
||||
*.oj-debug.json
|
||||
oj-debug.*.json
|
||||
|
||||
# Generated orchestrator SFT/eval data artifacts
|
||||
data/
|
||||
|
||||
# Training run artifacts (model checkpoints + W&B run dirs — regenerated, large)
|
||||
checkpoints/
|
||||
wandb/
|
||||
|
||||
# Local LiteLLM spend-tracking proxy (machine-specific, not part of the package)
|
||||
litellm/
|
||||
|
||||
@@ -28,6 +28,7 @@ classifiers = [
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
]
|
||||
dependencies = [
|
||||
"braintrust>=0.0.150",
|
||||
"click>=8",
|
||||
"datasets>=4.5.0",
|
||||
"ddgs>=9.11.4",
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
#!/usr/bin/env python
|
||||
"""Build orchestrator SFT cold-start data by **base Qwen3-8B self-sampling** (v1).
|
||||
|
||||
The orchestrator *is* the local Qwen3-8B. We roll the base (un-trained) model out
|
||||
over the reasoning-SFT task pool (``load_sft_tasks``: GeneralThought + OpenThoughts
|
||||
code/math/science), verify each trajectory's final answer against the gold
|
||||
(``verify.make_verifier`` — math/code checkers + Gemini judge fallback), and keep
|
||||
the cheapest-correct trajectory per task. Those passing trajectories become the
|
||||
``conversations`` JSONL the SFT trainer consumes — i.e. the model learns from its
|
||||
own successful rollouts (rejection-sampling cold-start, STaR-style).
|
||||
|
||||
This is **v1** (base self-sampling): point ``--orchestrator-endpoint`` at the vLLM
|
||||
server hosting the *base* ``Qwen/Qwen3-8B``.
|
||||
|
||||
For **v2**, train on the v1 data, serve the v1 checkpoint with vLLM, then re-run
|
||||
this exact driver with ``--orchestrator-endpoint`` pointed at the v1 checkpoint's
|
||||
endpoint, ``--orchestrator-model`` set to the checkpoint name, and
|
||||
``--samples-per-task 8`` — the (now stronger) policy self-samples a v2 dataset.
|
||||
|
||||
The math/coder specialist tools are only wired when ``--math-endpoint`` /
|
||||
``--coder-endpoint`` are passed (those are separately-served vLLM specialists);
|
||||
otherwise they're omitted so the orchestrator only sees the tools it can actually
|
||||
call.
|
||||
|
||||
Example (v1, base self-sampling):
|
||||
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
--orchestrator-label qwen \
|
||||
--orchestrator-endpoint http://localhost:8001/v1 \
|
||||
--orchestrator-model qwen3-8b \
|
||||
--samples-per-task 8 --max-keep-per-task 1
|
||||
|
||||
Example (v2, re-point at the v1 checkpoint):
|
||||
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
--out data/orchestrator_sft_v2.jsonl \
|
||||
--orchestrator-endpoint http://localhost:8010/v1 \
|
||||
--orchestrator-model orchestrator-sft-v1 \
|
||||
--samples-per-task 8
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
|
||||
from openjarvis.agents.hybrid.toolorchestra.unified import (
|
||||
make_call_orchestrator,
|
||||
make_dispatch,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
|
||||
load_sft_tasks,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import make_verifier
|
||||
|
||||
|
||||
def main(argv: Optional[list[str]] = None) -> int:
|
||||
p = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
# --out is a TAG, not a path. The run dir is always
|
||||
# data/orchestrator/raw/{label}_{MMDD}[_{tag}]/ and the file inside is always
|
||||
# data.jsonl — see the run_dir block below. Use --out only to distinguish two
|
||||
# runs of the same orchestrator on the same day (e.g. --out balanced50).
|
||||
p.add_argument(
|
||||
"--out",
|
||||
default=None,
|
||||
help="Optional tag appended to the run dir name (not a path).",
|
||||
)
|
||||
# The orchestrator == the local Qwen3-8B self-sampling over its own rollouts.
|
||||
p.add_argument(
|
||||
"--orchestrator-endpoint",
|
||||
default="http://localhost:8001/v1",
|
||||
help="OpenAI-compatible vLLM base URL serving the orchestrator.",
|
||||
)
|
||||
p.add_argument("--orchestrator-model", default="qwen3-8b")
|
||||
p.add_argument("--orchestrator-api-key", default="EMPTY")
|
||||
# Provenance stamps: when set, every written record carries these fields so
|
||||
# the JSONL is self-identifying once pooled across model families (gemma vs
|
||||
# qwen). Default None -> no stamping (existing qwen runs unaffected).
|
||||
p.add_argument(
|
||||
"--gen-model",
|
||||
default=None,
|
||||
help="Full HF id of the generating orchestrator, stamped as "
|
||||
"record['gen_model'] (e.g. google/gemma-4-26B-A4B-it).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--orchestrator-label",
|
||||
default=None,
|
||||
help="Short label stamped as record['orchestrator_model'] (e.g. gemma-4-26b).",
|
||||
)
|
||||
# Local OSS model endpoints; repeatable "model_id=base_url" (e.g.
|
||||
# "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are
|
||||
# still listed but served unconfigured (base_url=None).
|
||||
p.add_argument(
|
||||
"--local-endpoint",
|
||||
action="append",
|
||||
default=[],
|
||||
metavar="MODEL_ID=URL",
|
||||
help="Local model id -> vLLM base URL (repeatable).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--max-tasks",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Cap on tasks (default: all of load_sft_tasks()).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--skip-task-ids-from",
|
||||
action="append",
|
||||
default=[],
|
||||
metavar="GLOB",
|
||||
help="Glob of prior data.jsonl files; skip task_ids already "
|
||||
"generated there so this run only does unseen prompts "
|
||||
"(resume). Repeatable; applied before sharding.",
|
||||
)
|
||||
p.add_argument("--samples-per-task", type=int, default=8)
|
||||
p.add_argument("--max-keep-per-task", type=int, default=1)
|
||||
p.add_argument("--max-turns", type=int, default=8)
|
||||
p.add_argument("--temperature", type=float, default=0.7)
|
||||
# Orchestrator completion cap per turn. The library default (4096) intermittently
|
||||
# truncates the final answer mid-sentence on longer reasoning turns; bump it so
|
||||
# the trace ends cleanly. Raise further (e.g. 16384) if traces still cut off.
|
||||
p.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Per-turn completion cap for the orchestrator (default 8192).",
|
||||
)
|
||||
# Number of tasks rolled out concurrently against vLLM. Each in-flight task
|
||||
# issues its samples sequentially, so ~concurrency requests hit the server at
|
||||
# once. ~30-50 is comfortable on 1xL40S (prefix caching + continuous batching).
|
||||
p.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Tasks rolled out in parallel (default 32).",
|
||||
)
|
||||
# Balanced-smoke pull: cap//4 from each of GeneralThought + OpenThoughts
|
||||
# code/math/science instead of the GeneralThought-only fast cap path. Use for a
|
||||
# representative smoke; the real run omits --max-tasks for the full 8K balanced set.
|
||||
p.add_argument(
|
||||
"--balanced",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Draw an EVEN cross-domain sample (GeneralThought + "
|
||||
"OpenThoughts code/math/science). Default ON — pass "
|
||||
"--no-balanced for the old GeneralThought-only skew.",
|
||||
)
|
||||
# Data-parallel sharding: split the (deterministic, seed-42) task list across
|
||||
# N independent driver processes, each pointed at its own orchestrator vLLM
|
||||
# replica. Shard i takes tasks[i::N] (a strided slice, so every shard stays
|
||||
# domain-balanced). Each writes its own --out file; concatenate the shard
|
||||
# JSONLs afterward. Lets the GPU-bound orchestrator scale across idle GPUs.
|
||||
p.add_argument(
|
||||
"--shard-index",
|
||||
type=int,
|
||||
default=0,
|
||||
help="This shard's index in [0, shard-count).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--shard-count",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Total number of shards (default 1 = no sharding).",
|
||||
)
|
||||
# Throughput knob: stop sampling a task as soon as max-keep-per-task passing
|
||||
# trajectories are found, instead of always running all --samples-per-task.
|
||||
# ~3-4x faster when most tasks solve early, but keeps the *first* passers
|
||||
# rather than the *cheapest* of N (drops the cost-optimisation signal).
|
||||
p.add_argument(
|
||||
"--stop-at-keep",
|
||||
action="store_true",
|
||||
help="Short-circuit a task once max-keep passing samples found.",
|
||||
)
|
||||
# Anonymize model experts (opaque random labels, uniform description, no cost
|
||||
# line, shuffled order) so the policy can't route on a model's name/position/
|
||||
# cost. The anon->real map is saved per record (metrics.anon_map) for analysis.
|
||||
p.add_argument(
|
||||
"--anonymize-experts",
|
||||
action="store_true",
|
||||
help="Hide expert identity (random labels + shuffle) when routing.",
|
||||
)
|
||||
# By default we keep EVERY rolled-out trajectory (correct + incorrect), each
|
||||
# tagged with ``correct`` (verifier verdict) and ``kept`` (the cheapest-correct
|
||||
# sample the rejection sampler would pick). Lets you compute accuracy / inspect
|
||||
# failures from the JSONL directly. --rejection-only restores the old
|
||||
# drop-the-failures behaviour (only cheapest-correct written).
|
||||
p.add_argument(
|
||||
"--rejection-only",
|
||||
action="store_true",
|
||||
help="Drop incorrect rollouts; write only the cheapest-correct "
|
||||
"sample per task (the original behaviour).",
|
||||
)
|
||||
# Rollouts call shell_exec / file_write with model-chosen relative paths
|
||||
# (e.g. ``solution.py``), which otherwise land in the repo root. Run them
|
||||
# from a throwaway scratch dir so generated files never dirty the tree.
|
||||
# gitignored via the existing ``scratch/`` rule.
|
||||
p.add_argument(
|
||||
"--scratch-dir",
|
||||
default="scratch/sft-rollouts",
|
||||
help="CWD for rollouts; stray tool-written files go here.",
|
||||
)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
# Quiet the per-request HTTP / dataset-stream spam so the run log stays
|
||||
# readable (rollout calls and dataset shards otherwise flood it).
|
||||
for _noisy in (
|
||||
"httpx",
|
||||
"httpcore",
|
||||
"urllib3",
|
||||
"datasets",
|
||||
"fsspec",
|
||||
"huggingface_hub",
|
||||
"openai",
|
||||
):
|
||||
logging.getLogger(_noisy).setLevel(logging.WARNING)
|
||||
|
||||
# Every run lands in its OWN folder under <data-root>/raw/, named to MATCH the
|
||||
# sft/ file it will eventually produce — only the stage word differs:
|
||||
#
|
||||
# raw/qwen_0707/data.jsonl <- every rollout, incl. failures
|
||||
# sft/qwen_clean_0707.jsonl <- reject-sampled from it
|
||||
#
|
||||
# so any curated file traces back to its generation run by eye. raw/ sits ABOVE
|
||||
# the sft/rl fork on purpose: SFT keeps only correct+clean rows, but GRPO needs
|
||||
# the failures, and both read the same pool. Scheme is {name}_{MMDD}, name =
|
||||
# --orchestrator-label (qwen/gemma/...). Same name twice in a day -> -2, -3
|
||||
# suffix (creation order is preserved, and a MM-DD-HHMMpm stamp is redundant
|
||||
# with mtime). --out is an optional extra tag, NOT a path; the data file inside
|
||||
# is always data.jsonl.
|
||||
#
|
||||
# OJ_DATA_ROOT keeps the data OUT of the git checkout (repo-relative default so
|
||||
# a fresh clone still works); this workspace points it at the experiments tree.
|
||||
data_root = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
|
||||
prefix = args.orchestrator_label or "orch"
|
||||
tag = f"_{Path(args.out).stem}" if args.out else ""
|
||||
base = data_root / "raw" / f"{prefix}_{time.strftime('%m%d')}{tag}"
|
||||
run_dir = base.resolve()
|
||||
n = 1
|
||||
while run_dir.exists(): # same name+day (incl. parallel shards) -> disambiguate
|
||||
n += 1
|
||||
run_dir = base.with_name(f"{base.name}-{n}").resolve()
|
||||
label = run_dir.name
|
||||
out_p = run_dir / "data.jsonl"
|
||||
lock_p = out_p.with_suffix(out_p.suffix + ".lock")
|
||||
out_p.parent.mkdir(parents=True, exist_ok=True)
|
||||
lock_p.write_text(f"{time.strftime('%m-%d-%I%M%p').lower()} pid={os.getpid()}")
|
||||
args.out = str(out_p)
|
||||
import atexit
|
||||
|
||||
atexit.register(lambda: lock_p.exists() and lock_p.unlink())
|
||||
logging.info("Run dir: %s", run_dir)
|
||||
|
||||
# Enrich Braintrust rollout traces with run-level provenance (no-op-safe;
|
||||
# tracing.run_context() reads these). setdefault so an explicit env wins.
|
||||
os.environ.setdefault("OJ_RUN_LABEL", label)
|
||||
os.environ.setdefault("OJ_GEN_MODEL", args.gen_model or args.orchestrator_model)
|
||||
os.environ.setdefault("OJ_RUN_STAGE", "smoke" if args.max_tasks else "prod")
|
||||
os.environ["OJ_CFG_TEMPERATURE"] = str(args.temperature)
|
||||
os.environ["OJ_CFG_MAX_TURNS"] = str(args.max_turns)
|
||||
os.environ["OJ_CFG_ANONYMIZE"] = str(bool(args.anonymize_experts))
|
||||
os.environ["OJ_CFG_REJECTION_ONLY"] = str(bool(args.rejection_only))
|
||||
|
||||
# Sandbox the rollouts: chdir into a scratch dir so any file a rollout writes
|
||||
# (shell_exec redirects, file_write with a relative path) lands there instead
|
||||
# of the repo root. (--out is already absolute via run_dir.resolve() above.)
|
||||
scratch = Path(args.scratch_dir).resolve()
|
||||
scratch.mkdir(parents=True, exist_ok=True)
|
||||
os.chdir(scratch)
|
||||
logging.info("Rollout scratch dir (cwd): %s", scratch)
|
||||
|
||||
local_endpoints = {}
|
||||
for item in args.local_endpoint:
|
||||
model_id, _, url = item.partition("=")
|
||||
if model_id and url:
|
||||
local_endpoints[model_id] = url
|
||||
tools = orchestrator_catalog(local_endpoints=local_endpoints or None)
|
||||
# Drop the no-op ``think`` scratchpad tool: it adds turns/cost and is a
|
||||
# harness-leakage vector (the model narrates the routing rule inside a think
|
||||
# call, which the reasoning-scrub can't reach). The model still reasons in
|
||||
# its own <think> blocks — it just can't emit reasoning AS a tool call.
|
||||
tools = [t for t in tools if t.name != "think"]
|
||||
logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools])
|
||||
|
||||
call_orch = make_call_orchestrator(
|
||||
args.orchestrator_model,
|
||||
base_url=args.orchestrator_endpoint,
|
||||
api_key=args.orchestrator_api_key,
|
||||
temperature=args.temperature,
|
||||
max_tokens=args.max_tokens,
|
||||
)
|
||||
dispatch = make_dispatch({})
|
||||
|
||||
# Build the provenance stamp once (None if neither flag given).
|
||||
record_extra = {}
|
||||
if args.gen_model:
|
||||
record_extra["gen_model"] = args.gen_model
|
||||
if args.orchestrator_label:
|
||||
record_extra["orchestrator_model"] = args.orchestrator_label
|
||||
record_extra = record_extra or None
|
||||
|
||||
def rollout_fn(task):
|
||||
try:
|
||||
return run_unified_rollout(
|
||||
task.instruction,
|
||||
tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=args.max_turns,
|
||||
anonymize=args.anonymize_experts,
|
||||
)
|
||||
except Exception as exc: # network/key failures shouldn't kill the run
|
||||
logging.warning("rollout failed for %s: %s", task.task_id, exc)
|
||||
return None
|
||||
|
||||
# cap= caps the task count for smoke runs; --balanced makes that small sample
|
||||
# cross-domain (GeneralThought + OpenThoughts code/math/science) instead of the
|
||||
# GeneralThought-only fast path. The real run omits --max-tasks (full 8K balanced).
|
||||
tasks = load_sft_tasks(cap=args.max_tasks, balanced=args.balanced)
|
||||
# Resume on unseen prompts: drop task_ids already generated in prior runs
|
||||
# (per-track seen set via --skip-task-ids-from). Filter BEFORE sharding so the
|
||||
# remaining unseen tasks stride evenly and never re-cover finished prompts.
|
||||
if args.skip_task_ids_from:
|
||||
import glob as _glob
|
||||
|
||||
skip_ids = set()
|
||||
for pattern in args.skip_task_ids_from:
|
||||
for fp in _glob.glob(pattern):
|
||||
try:
|
||||
with open(fp) as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
tid = json.loads(line).get("task_id")
|
||||
except Exception:
|
||||
continue
|
||||
if tid:
|
||||
skip_ids.add(tid)
|
||||
except OSError:
|
||||
continue
|
||||
if skip_ids:
|
||||
n_before = len(tasks)
|
||||
tasks = [t for t in tasks if t.task_id not in skip_ids]
|
||||
logging.info(
|
||||
"Resume: %d done task_ids -> skipping; %d of %d tasks remain",
|
||||
len(skip_ids),
|
||||
len(tasks),
|
||||
n_before,
|
||||
)
|
||||
if args.shard_count > 1:
|
||||
n_all = len(tasks)
|
||||
tasks = tasks[args.shard_index :: args.shard_count]
|
||||
logging.info(
|
||||
"Shard %d/%d -> %d of %d tasks",
|
||||
args.shard_index,
|
||||
args.shard_count,
|
||||
len(tasks),
|
||||
n_all,
|
||||
)
|
||||
from collections import Counter as _Counter
|
||||
|
||||
_dom = _Counter(getattr(t, "domain", "unknown") for t in tasks)
|
||||
logging.info(
|
||||
"Loaded %d SFT tasks | balanced=%s | domain split: %s",
|
||||
len(tasks),
|
||||
args.balanced,
|
||||
", ".join(f"{d}={n}" for d, n in sorted(_dom.items())),
|
||||
)
|
||||
|
||||
stats = generate_sft_dataset(
|
||||
args.out,
|
||||
tasks=tasks,
|
||||
tools=tools,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=make_verifier(),
|
||||
samples_per_task=args.samples_per_task,
|
||||
max_keep_per_task=args.max_keep_per_task,
|
||||
reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets highest reward
|
||||
concurrency=args.concurrency,
|
||||
stop_at_keep=args.stop_at_keep,
|
||||
keep_all=not args.rejection_only,
|
||||
record_extra=record_extra,
|
||||
)
|
||||
# Emit human-viewable companions (<out>.pretty.json + <out>.txt) so the
|
||||
# traces / full message turns can be inspected without un-escaping the JSONL.
|
||||
try:
|
||||
from render_sft_data import render # same scripts/orchestrator dir
|
||||
except ImportError:
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from render_sft_data import render
|
||||
if stats.get("records_written", 0) > 0:
|
||||
rinfo = render(args.out)
|
||||
logging.info("Rendered viewable traces: %s + %s", rinfo["pretty"], rinfo["txt"])
|
||||
|
||||
print(json.dumps(stats, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,344 @@
|
||||
#!/usr/bin/env python
|
||||
"""Run the ToolOrchestra orchestrator over N-sample subsets of several
|
||||
benchmarks and score them with the existing eval infra.
|
||||
|
||||
Reuses, verbatim:
|
||||
* ``openjarvis.evals.cli._build_dataset`` / ``_build_scorer`` /
|
||||
``_build_judge_backend`` (loaders + scorers + judge wiring),
|
||||
* ``openjarvis.evals.core.runner.EvalRunner`` (parallel run + scoring),
|
||||
* ``openjarvis.evals.core.types.RunConfig`` (run config),
|
||||
|
||||
and plugs in our orchestrator as the "model" via
|
||||
``openjarvis.learning.intelligence.orchestrator.eval_backend.OrchestratorBackend``.
|
||||
|
||||
For each benchmark it runs ``EvalRunner(config, dataset, backend, scorer).run()``,
|
||||
collects accuracy / cost / latency, writes a combined ``summary.json``, and
|
||||
prints a table.
|
||||
|
||||
NOTE on the judge: the OpenAI key is dead, so the default judge model is a
|
||||
Gemini model (``gemini-2.5-flash``). The judge backend is still built via the
|
||||
``cloud`` engine (``_build_judge_backend``) — the model string selects the
|
||||
provider/route. Override with ``--judge-model`` / ``--judge-engine``.
|
||||
|
||||
Example
|
||||
-------
|
||||
.venv/bin/python scripts/orchestrator/eval_orchestrator.py \\
|
||||
--benchmarks gaia,mmlu_pro --n 100 \\
|
||||
--orchestrator-endpoint http://localhost:8001/v1 \\
|
||||
--orchestrator-model qwen3-8b \\
|
||||
--output-dir ~/.openjarvis/experiments/hybrid/runs/orch-eval
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# Friendly aliases (underscored / shorthand) -> the EXACT registry keys used by
|
||||
# openjarvis.evals.cli.BENCHMARKS / _build_dataset / _build_scorer.
|
||||
BENCHMARK_ALIASES: Dict[str, str] = {
|
||||
"terminalbench_v2_1": "terminalbench-v2.1",
|
||||
"terminalbench-v2_1": "terminalbench-v2.1",
|
||||
"terminalbench_v21": "terminalbench-v2.1",
|
||||
"terminalbench-v2.1": "terminalbench-v2.1",
|
||||
"mmlu_pro": "mmlu-pro",
|
||||
"mmlu-pro": "mmlu-pro",
|
||||
"gaia": "gaia",
|
||||
"taubench": "taubench",
|
||||
"supergpqa": "supergpqa",
|
||||
}
|
||||
|
||||
DEFAULT_BENCHMARKS = "gaia,terminalbench_v2_1,taubench,mmlu_pro,supergpqa"
|
||||
|
||||
|
||||
def _normalize_benchmark(name: str) -> str:
|
||||
key = name.strip()
|
||||
return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key))
|
||||
|
||||
|
||||
def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(
|
||||
description="Run the orchestrator over benchmark subsets and score them.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
p.add_argument(
|
||||
"--benchmarks",
|
||||
default=DEFAULT_BENCHMARKS,
|
||||
help="Comma-separated benchmark names (aliases normalized to registry keys).",
|
||||
)
|
||||
p.add_argument("--n", type=int, default=100, help="Samples per benchmark.")
|
||||
p.add_argument("--seed", type=int, default=42, help="Subset seed.")
|
||||
p.add_argument(
|
||||
"--orchestrator-endpoint",
|
||||
default="http://localhost:8001/v1",
|
||||
help="OpenAI-compatible base URL for the served orchestrator.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--orchestrator-model",
|
||||
default="qwen3-8b",
|
||||
help="Served orchestrator model id.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--orchestrator-api-key",
|
||||
default="EMPTY",
|
||||
help="API key for the orchestrator endpoint (EMPTY for local vLLM).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--local-endpoint",
|
||||
action="append",
|
||||
default=[],
|
||||
metavar="MODEL_ID=URL",
|
||||
help=(
|
||||
"Wire a local OSS tool: map a catalog model id to its vLLM base "
|
||||
"URL, e.g. 'Qwen/Qwen3.5-9B=http://localhost:8003/v1'. Repeatable."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--max-turns",
|
||||
type=int,
|
||||
default=8,
|
||||
help="Max orchestrator turns per sample.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--temperature",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Orchestrator sampling temperature.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--max-workers",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Parallel samples per benchmark.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--output-dir",
|
||||
default="results/orchestrator-eval",
|
||||
help="Directory for per-benchmark JSONL + combined summary.json.",
|
||||
)
|
||||
# NOTE: OpenAI key is dead — default the judge to a Gemini model.
|
||||
p.add_argument(
|
||||
"--judge-model",
|
||||
default="gemini-2.5-flash",
|
||||
help="LLM-judge model (default Gemini — OpenAI key is dead).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--judge-engine",
|
||||
default="cloud",
|
||||
help="Engine key for the judge backend.",
|
||||
)
|
||||
return p.parse_args(argv)
|
||||
|
||||
|
||||
def _run_one(
|
||||
benchmark: str,
|
||||
*,
|
||||
n: int,
|
||||
seed: int,
|
||||
backend,
|
||||
judge_model: str,
|
||||
judge_engine: str,
|
||||
max_workers: int,
|
||||
orchestrator_model: str,
|
||||
output_dir: Path,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run + score one benchmark, returning a row dict (or an error row)."""
|
||||
# Imports here so --help works without importing the (heavy) eval stack.
|
||||
from openjarvis.evals.cli import (
|
||||
_build_dataset,
|
||||
_build_judge_backend,
|
||||
_build_scorer,
|
||||
)
|
||||
from openjarvis.evals.core.runner import EvalRunner
|
||||
from openjarvis.evals.core.types import RunConfig
|
||||
|
||||
output_path = output_dir / f"{benchmark.replace('.', '_')}_orchestrator.jsonl"
|
||||
|
||||
dataset = _build_dataset(benchmark)
|
||||
# The runner reloads internally; we also load here per spec so callers can
|
||||
# inspect/size the subset up front. load() is idempotent for these sets.
|
||||
dataset.load(max_samples=n, seed=seed)
|
||||
|
||||
judge_backend = _build_judge_backend(judge_model, engine_key=judge_engine)
|
||||
scorer = _build_scorer(benchmark, judge_backend, judge_model)
|
||||
|
||||
config = RunConfig(
|
||||
benchmark=benchmark,
|
||||
backend="orchestrator",
|
||||
model=orchestrator_model,
|
||||
max_samples=n,
|
||||
max_workers=max_workers,
|
||||
seed=seed,
|
||||
judge_model=judge_model,
|
||||
judge_engine=judge_engine,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
|
||||
runner = EvalRunner(config, dataset, backend, scorer)
|
||||
started = time.time()
|
||||
try:
|
||||
summary = runner.run()
|
||||
finally:
|
||||
if judge_backend is not None:
|
||||
judge_backend.close()
|
||||
elapsed = time.time() - started
|
||||
|
||||
return {
|
||||
"benchmark": benchmark,
|
||||
"accuracy": summary.accuracy,
|
||||
"scored_samples": summary.scored_samples,
|
||||
"correct": summary.correct,
|
||||
"errors": summary.errors,
|
||||
"total_samples": summary.total_samples,
|
||||
"mean_latency_seconds": summary.mean_latency_seconds,
|
||||
"total_cost_usd": summary.total_cost_usd,
|
||||
"mean_continuous_score": summary.mean_continuous_score,
|
||||
"wall_seconds": round(elapsed, 2),
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
|
||||
def _print_table(rows: List[Dict[str, Any]]) -> None:
|
||||
cols = [
|
||||
("benchmark", 22, "{}"),
|
||||
("accuracy", 9, "{:.4f}"),
|
||||
("correct", 8, "{}"),
|
||||
("scored", 7, "{}"),
|
||||
("errors", 7, "{}"),
|
||||
("cost($)", 10, "{:.4f}"),
|
||||
("lat(s)", 9, "{:.2f}"),
|
||||
("wall(s)", 9, "{:.1f}"),
|
||||
]
|
||||
header = " ".join(f"{name:<{w}}" for name, w, _ in cols)
|
||||
print("\n" + header)
|
||||
print("-" * len(header))
|
||||
for r in rows:
|
||||
if r.get("error"):
|
||||
print(f"{r['benchmark']:<22} ERROR: {r['error']}")
|
||||
continue
|
||||
vals = {
|
||||
"benchmark": r["benchmark"],
|
||||
"accuracy": r["accuracy"],
|
||||
"correct": r["correct"],
|
||||
"scored": r["scored_samples"],
|
||||
"errors": r["errors"],
|
||||
"cost($)": r["total_cost_usd"],
|
||||
"lat(s)": r["mean_latency_seconds"],
|
||||
"wall(s)": r["wall_seconds"],
|
||||
}
|
||||
cells = []
|
||||
for name, w, fmt in cols:
|
||||
v = vals[name]
|
||||
try:
|
||||
s = fmt.format(v)
|
||||
except (ValueError, TypeError):
|
||||
s = str(v)
|
||||
cells.append(f"{s:<{w}}")
|
||||
print(" ".join(cells))
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> int:
|
||||
args = _parse_args(argv)
|
||||
|
||||
benchmarks = [
|
||||
_normalize_benchmark(b) for b in args.benchmarks.split(",") if b.strip()
|
||||
]
|
||||
output_dir = Path(args.output_dir).expanduser()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Build the orchestrator backend once (stateless across benchmarks).
|
||||
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
|
||||
OrchestratorBackend,
|
||||
)
|
||||
|
||||
local_endpoints: Dict[str, str] = {}
|
||||
for pair in args.local_endpoint:
|
||||
if "=" not in pair:
|
||||
raise SystemExit(f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}")
|
||||
model_id, url = pair.split("=", 1)
|
||||
local_endpoints[model_id.strip()] = url.strip()
|
||||
|
||||
backend = OrchestratorBackend(
|
||||
orchestrator_endpoint=args.orchestrator_endpoint,
|
||||
orchestrator_model=args.orchestrator_model,
|
||||
api_key=args.orchestrator_api_key,
|
||||
local_endpoints=local_endpoints,
|
||||
max_turns=args.max_turns,
|
||||
temperature=args.temperature,
|
||||
)
|
||||
|
||||
rows: List[Dict[str, Any]] = []
|
||||
try:
|
||||
for benchmark in benchmarks:
|
||||
print(f"\n=== {benchmark} (n={args.n}, seed={args.seed}) ===")
|
||||
try:
|
||||
row = _run_one(
|
||||
benchmark,
|
||||
n=args.n,
|
||||
seed=args.seed,
|
||||
backend=backend,
|
||||
judge_model=args.judge_model,
|
||||
judge_engine=args.judge_engine,
|
||||
max_workers=args.max_workers,
|
||||
orchestrator_model=args.orchestrator_model,
|
||||
output_dir=output_dir,
|
||||
)
|
||||
print(
|
||||
f" accuracy={row['accuracy']:.4f} "
|
||||
f"({row['correct']}/{row['scored_samples']}) "
|
||||
f"errors={row['errors']} cost=${row['total_cost_usd']:.4f}"
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - one bad bench shouldn't kill all
|
||||
print(f" FAILED: {type(exc).__name__}: {exc}", file=sys.stderr)
|
||||
row = {"benchmark": benchmark, "error": f"{type(exc).__name__}: {exc}"}
|
||||
rows.append(row)
|
||||
finally:
|
||||
backend.close()
|
||||
|
||||
# Resolve the EXACT model behind the served alias (e.g. "gemma-sft" ->
|
||||
# the real id + checkpoint path vLLM reports), so results are self-identifying
|
||||
# instead of an opaque label. Best-effort: never fail the summary over it.
|
||||
served_id, served_path = args.orchestrator_model, None
|
||||
try:
|
||||
import urllib.request
|
||||
|
||||
req = urllib.request.Request(
|
||||
args.orchestrator_endpoint.rstrip("/") + "/models",
|
||||
headers={"Authorization": f"Bearer {args.orchestrator_api_key}"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.load(resp).get("data", [])
|
||||
if data:
|
||||
served_id = data[0].get("id", served_id)
|
||||
# vLLM reports the loaded checkpoint path in `root` (or `parent`).
|
||||
served_path = data[0].get("root") or data[0].get("parent")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
combined = {
|
||||
"orchestrator_model": args.orchestrator_model,
|
||||
"orchestrator_served_id": served_id,
|
||||
"orchestrator_model_path": served_path,
|
||||
"orchestrator_endpoint": args.orchestrator_endpoint,
|
||||
"judge_model": args.judge_model,
|
||||
"judge_engine": args.judge_engine,
|
||||
"n": args.n,
|
||||
"seed": args.seed,
|
||||
"max_turns": args.max_turns,
|
||||
"benchmarks": rows,
|
||||
}
|
||||
summary_path = output_dir / "summary.json"
|
||||
with open(summary_path, "w") as f:
|
||||
json.dump(combined, f, indent=2, default=str)
|
||||
|
||||
_print_table(rows)
|
||||
print(f"\nCombined summary written to {summary_path}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,267 @@
|
||||
#!/usr/bin/env python
|
||||
"""Render graded orchestrator-eval JSONL records into clean, readable .txt files.
|
||||
|
||||
Companion to ``format_sft_sample.py`` (same banner/indent visual style), but
|
||||
for the *scored* eval outputs written by ``eval_orchestrator.py``:
|
||||
|
||||
results/orch-eval-<model>-<tranche>-<suite>/<benchmark>_orchestrator.jsonl
|
||||
|
||||
Each record has: record_id, benchmark, model, model_answer (the model's
|
||||
FINAL_ANSWER string), is_correct, score, latency_seconds, cost_usd, error, and
|
||||
scoring_metadata. The full step-by-step routing trace is NOT saved — only the
|
||||
final answer + gold + score — so this renderer cannot show intermediate steps.
|
||||
|
||||
The QUESTION text is not in the result JSONL; it's loaded from the benchmark
|
||||
dataset by ``record_id`` (via ``openjarvis.evals.cli._build_dataset``). If the
|
||||
dataset can't be loaded / the id doesn't match, we render everything else and
|
||||
show "(question unavailable)".
|
||||
|
||||
Gold answer lives in ``scoring_metadata`` (shape varies by scorer):
|
||||
* MCQ (mmlu-pro, supergpqa): ``reference_letter``.
|
||||
* LLM-judge (gaia): parsed from ``raw_judge_output`` (`gold target "..."`),
|
||||
plus the judge's ``reasoning``.
|
||||
* anything else: dumped verbatim under SCORING (raw).
|
||||
|
||||
Usage:
|
||||
.venv/bin/python scripts/orchestrator/format_eval_sample.py \
|
||||
--input results/orch-eval-gemma-2k-full/gaia_orchestrator.jsonl --n 2
|
||||
|
||||
... --input <file> --lines 1,5,42 # specific 1-indexed lines
|
||||
... --input <file> --all # one .txt per record
|
||||
... --input <file> --all --only-wrong # only incorrect samples
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
WIDTH = 80
|
||||
|
||||
# Friendly aliases -> the registry keys _build_dataset expects (mirrors
|
||||
# eval_orchestrator.BENCHMARK_ALIASES for the ones that write .jsonl here).
|
||||
BENCHMARK_ALIASES = {
|
||||
"mmlu_pro": "mmlu-pro",
|
||||
"mmlu-pro": "mmlu-pro",
|
||||
"gaia": "gaia",
|
||||
"taubench": "taubench",
|
||||
"supergpqa": "supergpqa",
|
||||
"terminalbench_v2_1": "terminalbench-v2.1",
|
||||
"terminalbench-v2.1": "terminalbench-v2.1",
|
||||
}
|
||||
|
||||
|
||||
def _banner(label: str) -> str:
|
||||
"""Full-width headline: ``━━━ LABEL ━━━━━━…``."""
|
||||
prefix = f"━━━ {label} "
|
||||
fill = "━" * max(3, WIDTH - len(prefix))
|
||||
return f"\n{prefix}{fill}\n"
|
||||
|
||||
|
||||
def _indent(text: str, pad: str = " ") -> str:
|
||||
text = (text or "").rstrip()
|
||||
return "\n".join(pad + ln if ln.strip() else ln for ln in text.splitlines())
|
||||
|
||||
|
||||
def _normalize_benchmark(name: str) -> str:
|
||||
key = (name or "").strip()
|
||||
return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Question map: record_id -> problem text, loaded from the benchmark dataset.
|
||||
# ---------------------------------------------------------------------------
|
||||
def build_question_map(benchmark: str, n: int, seed: int = 42) -> dict:
|
||||
"""Best-effort {record_id: problem}. Returns {} if the dataset can't load."""
|
||||
try:
|
||||
from openjarvis.evals.cli import _build_dataset
|
||||
|
||||
ds = _build_dataset(_normalize_benchmark(benchmark))
|
||||
ds.load(max_samples=n, seed=seed)
|
||||
return {r.record_id: r.problem for r in ds.iter_records()}
|
||||
except Exception as exc: # noqa: BLE001 - never fail the render over this
|
||||
print(
|
||||
f" [warn] could not load dataset for {benchmark!r}: "
|
||||
f"{type(exc).__name__}: {exc}"
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gold / judge parsing out of scoring_metadata.
|
||||
# ---------------------------------------------------------------------------
|
||||
_GOLD_TARGET_RE = re.compile(r'gold target\s*"([^"]*)"', re.IGNORECASE)
|
||||
_JUDGE_REASONING_RE = re.compile(
|
||||
r"reasoning:\s*(.*?)(?:\n\s*correct:|\Z)", re.IGNORECASE | re.DOTALL
|
||||
)
|
||||
|
||||
|
||||
def parse_scoring(meta: dict) -> dict:
|
||||
"""Return {gold, judge, kind, clean}. ``clean`` False => dump raw block.
|
||||
|
||||
kind is one of 'mcq' | 'judge' | 'unknown'.
|
||||
"""
|
||||
if not isinstance(meta, dict):
|
||||
return {"gold": None, "judge": None, "kind": "unknown", "clean": False}
|
||||
|
||||
# MCQ scorer (mmlu-pro / supergpqa).
|
||||
if "reference_letter" in meta:
|
||||
return {
|
||||
"gold": meta.get("reference_letter"),
|
||||
"judge": None,
|
||||
"kind": "mcq",
|
||||
"clean": True,
|
||||
}
|
||||
|
||||
# LLM-judge scorer (gaia).
|
||||
raw = meta.get("raw_judge_output")
|
||||
if raw:
|
||||
gold_m = _GOLD_TARGET_RE.search(raw)
|
||||
reason_m = _JUDGE_REASONING_RE.search(raw)
|
||||
return {
|
||||
"gold": gold_m.group(1) if gold_m else None,
|
||||
"judge": (reason_m.group(1).strip() if reason_m else raw.strip()),
|
||||
"kind": "judge",
|
||||
"clean": True,
|
||||
}
|
||||
|
||||
# Unknown shape (e.g. {"reason": "no_choice_letter_extracted"}).
|
||||
return {"gold": None, "judge": None, "kind": "unknown", "clean": False}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Render one record.
|
||||
# ---------------------------------------------------------------------------
|
||||
def format_record(rec: dict, question: str | None) -> str:
|
||||
parsed = parse_scoring(rec.get("scoring_metadata"))
|
||||
is_correct = rec.get("is_correct")
|
||||
verdict = (
|
||||
"CORRECT"
|
||||
if is_correct
|
||||
else ("INCORRECT" if is_correct is False else "UNSCORED")
|
||||
)
|
||||
|
||||
def _fmt(v, fmt):
|
||||
try:
|
||||
return fmt.format(v)
|
||||
except (ValueError, TypeError):
|
||||
return str(v)
|
||||
|
||||
top = " · ".join(
|
||||
[
|
||||
str(rec.get("record_id", "?")),
|
||||
str(rec.get("benchmark", "?")),
|
||||
verdict,
|
||||
f"score {_fmt(rec.get('score'), '{:.3f}')}",
|
||||
f"lat {_fmt(rec.get('latency_seconds'), '{:.1f}')}s",
|
||||
f"cost ${_fmt(rec.get('cost_usd'), '{:.4f}')}",
|
||||
]
|
||||
)
|
||||
parts = [top]
|
||||
|
||||
err = rec.get("error")
|
||||
if err:
|
||||
parts.append(f" error: {err}")
|
||||
|
||||
parts.append(_banner("QUESTION"))
|
||||
parts.append(_indent(question if question else "(question unavailable)"))
|
||||
|
||||
parts.append(_banner("MODEL ANSWER"))
|
||||
parts.append(_indent(str(rec.get("model_answer") or "(empty)")))
|
||||
|
||||
parts.append(_banner("GOLD"))
|
||||
gold = parsed["gold"]
|
||||
parts.append(
|
||||
_indent(
|
||||
str(gold)
|
||||
if gold not in (None, "")
|
||||
else "(gold unavailable — see SCORING below)"
|
||||
)
|
||||
)
|
||||
|
||||
if parsed["kind"] == "judge" and parsed["judge"]:
|
||||
parts.append(_banner("JUDGE"))
|
||||
parts.append(_indent(parsed["judge"]))
|
||||
|
||||
if not parsed["clean"]:
|
||||
parts.append(_banner("SCORING (raw)"))
|
||||
parts.append(
|
||||
_indent(json.dumps(rec.get("scoring_metadata"), indent=2, default=str))
|
||||
)
|
||||
|
||||
return "\n".join(parts).rstrip() + "\n"
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
p = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
p.add_argument("--input", required=True, help="Graded *_orchestrator.jsonl file.")
|
||||
p.add_argument(
|
||||
"--out-dir", default="results/formatted", help="Where the .txt files go."
|
||||
)
|
||||
p.add_argument(
|
||||
"--seed", type=int, default=42, help="Subset seed used at eval time."
|
||||
)
|
||||
p.add_argument(
|
||||
"--only-wrong",
|
||||
action="store_true",
|
||||
help="Render only incorrect/unscored samples.",
|
||||
)
|
||||
g = p.add_mutually_exclusive_group()
|
||||
g.add_argument("--n", type=int, default=1, help="Format the first N records.")
|
||||
g.add_argument("--lines", help="Comma-separated 1-indexed line numbers.")
|
||||
g.add_argument("--all", action="store_true", help="Format every record.")
|
||||
args = p.parse_args(argv)
|
||||
|
||||
src = Path(args.input)
|
||||
out_dir = Path(args.out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
records = [json.loads(l) for l in src.open() if l.strip()]
|
||||
if not records:
|
||||
print("no records in input")
|
||||
return 1
|
||||
|
||||
# Pick indices first (so we only load the dataset once, sized to the file).
|
||||
if args.lines:
|
||||
idxs = [int(x) - 1 for x in args.lines.split(",") if x.strip()]
|
||||
elif args.all:
|
||||
idxs = list(range(len(records)))
|
||||
else:
|
||||
idxs = list(range(min(args.n, len(records))))
|
||||
idxs = [i for i in idxs if 0 <= i < len(records)]
|
||||
|
||||
if args.only_wrong:
|
||||
idxs = [i for i in idxs if not records[i].get("is_correct")]
|
||||
|
||||
benchmark = _normalize_benchmark(
|
||||
records[0].get("benchmark", src.stem.split("_")[0])
|
||||
)
|
||||
# Load the whole subset so any selected id resolves (the eval-time subset is
|
||||
# the first len(records) of seed=42).
|
||||
qmap = build_question_map(benchmark, n=len(records), seed=args.seed)
|
||||
|
||||
written = []
|
||||
for i in idxs:
|
||||
rec = records[i]
|
||||
rid = str(rec.get("record_id", i))
|
||||
bench = rec.get("benchmark", benchmark)
|
||||
fname = re.sub(r"[^A-Za-z0-9_.-]", "_", f"{bench}__{rid}")[:120] + ".txt"
|
||||
out = out_dir / fname
|
||||
out.write_text(format_record(rec, qmap.get(rid)))
|
||||
written.append(out)
|
||||
|
||||
for w in written:
|
||||
print(w)
|
||||
print(
|
||||
f"\nwrote {len(written)} file(s) to {out_dir}/"
|
||||
+ (" (dataset unavailable — questions omitted)" if not qmap else "")
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python
|
||||
"""Carve train / holdout / overfit100 splits from a clean orchestrator-SFT pool.
|
||||
|
||||
Replaces make_tranches.py + make_gemma_tranches.py, which carved 1k/2k/4k/8k
|
||||
data-scaling tranches from the June pool. That lineage is dead: the tranches it
|
||||
produced trained checkpoints that evaluated at base ~= 1k ~= 2k (the pool taught
|
||||
the orchestrator to re-derive answers itself instead of routing), and the pool
|
||||
was deleted. The lesson is in data/orchestrator/README.md — it was a data-QUALITY
|
||||
ceiling, so re-running the old scaling ladder buys nothing.
|
||||
|
||||
Reads a clean pool (reject-sampled from data/orchestrator/raw/) and writes the
|
||||
SFT splits to data/orchestrator/sft/. Naming is uniform — ``{name}_{split}_{date}``
|
||||
with name in {qwen, gemma, pooled}; pass several pools to merge-and-restratify
|
||||
into `pooled`.
|
||||
|
||||
The holdout is domain-stratified so val-loss is leak-free, and `overfit100` is a
|
||||
strict prefix of train (it is a memorisation sanity check — the model must be
|
||||
able to fit 100 rows it *has* seen, else the format/masking is broken).
|
||||
Deterministic (seed 42).
|
||||
|
||||
# one pool -> qwen_{train,holdout,overfit100}_0711.jsonl
|
||||
python scripts/orchestrator/make_splits.py --name qwen \
|
||||
--pool data/orchestrator/sft/qwen_clean_0711.jsonl
|
||||
|
||||
# merge both -> pooled_{train,holdout,overfit100}_0711.jsonl
|
||||
python scripts/orchestrator/make_splits.py --name pooled \
|
||||
--pool data/orchestrator/sft/qwen_clean_0711.jsonl \
|
||||
--pool data/orchestrator/sft/gemma_clean_0711.jsonl
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
# Where the orchestrator data tree lives. Repo-relative by default so a fresh
|
||||
# clone works out of the box; set OJ_DATA_ROOT to keep the data OUT of the git
|
||||
# checkout (this workspace points it at ~/experiments/orchestrator/data, so a
|
||||
# stray `git reset` can't touch hundreds of GB of generations).
|
||||
DATA_ROOT = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
|
||||
OUT = DATA_ROOT / "sft"
|
||||
SEED = 42
|
||||
OVERFIT_N = 100
|
||||
# Orchestrator that generated each pool — stamped onto every row so provenance
|
||||
# survives a merge (the filename encodes it too, but the field is cheap insurance
|
||||
# and is what the `pooled` splits rely on to stay attributable).
|
||||
ORCH_MODEL = {"qwen": "qwen3.5-9b", "gemma": "gemma-4-26b"}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
ap.add_argument(
|
||||
"--name",
|
||||
required=True,
|
||||
help="split family: qwen | gemma | pooled (drives the filename)",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--pool",
|
||||
action="append",
|
||||
required=True,
|
||||
metavar="PATH",
|
||||
help="clean pool jsonl; repeat to merge (use with --name pooled)",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--holdout-frac",
|
||||
type=float,
|
||||
default=0.15,
|
||||
help="fraction held out, domain-stratified (default 0.15)",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--date",
|
||||
default=datetime.now().strftime("%m%d"),
|
||||
help="date tag in the filename (default: today, MMDD)",
|
||||
)
|
||||
ap.add_argument("--out", type=Path, default=OUT)
|
||||
args = ap.parse_args()
|
||||
|
||||
rows = []
|
||||
for p in args.pool:
|
||||
path = Path(p)
|
||||
if not path.exists():
|
||||
print(f"!! pool not found: {path}", file=sys.stderr)
|
||||
return 1
|
||||
# Infer the source orchestrator from the filename so merged pools stay
|
||||
# attributable; don't clobber a stamp the pool already carries.
|
||||
stem = path.name.split("_", 1)[0]
|
||||
loaded = [json.loads(line) for line in path.open() if line.strip()]
|
||||
for r in loaded:
|
||||
r.setdefault("orchestrator_model", ORCH_MODEL.get(stem, stem))
|
||||
rows.extend(loaded)
|
||||
print(f"loaded {len(loaded):>5} rows from {path}")
|
||||
print(f"pool total: {len(rows)}")
|
||||
|
||||
# Stratified holdout: per-domain, proportional to domain size.
|
||||
by_dom = defaultdict(list)
|
||||
for i, r in enumerate(rows):
|
||||
by_dom[r.get("domain", "misc")].append(i)
|
||||
rng = random.Random(SEED)
|
||||
holdout_idx: set = set()
|
||||
for dom, idxs in sorted(by_dom.items()):
|
||||
k = round(args.holdout_frac * len(idxs))
|
||||
pick = rng.sample(idxs, min(k, len(idxs)))
|
||||
holdout_idx.update(pick)
|
||||
print(f" {dom:<9} pool={len(idxs):>4} holdout={len(pick)}")
|
||||
|
||||
holdout = [rows[i] for i in sorted(holdout_idx)]
|
||||
train = [r for i, r in enumerate(rows) if i not in holdout_idx] # order preserved
|
||||
overfit = train[:OVERFIT_N]
|
||||
print(f"train={len(train)} holdout={len(holdout)} overfit={len(overfit)}")
|
||||
|
||||
args.out.mkdir(parents=True, exist_ok=True)
|
||||
written = {}
|
||||
for split, data in (
|
||||
("train", train),
|
||||
("holdout", holdout),
|
||||
(f"overfit{OVERFIT_N}", overfit),
|
||||
):
|
||||
f = args.out / f"{args.name}_{split}_{args.date}.jsonl"
|
||||
f.write_text("".join(json.dumps(r) + "\n" for r in data))
|
||||
written[split] = f
|
||||
print(f"wrote {f} ({len(data)})")
|
||||
|
||||
# ---- auto-upload train + holdout to Braintrust ----
|
||||
# Gated by OJ_BRAINTRUST_AUTOUPLOAD (default ON); a total no-op / never raises
|
||||
# if the key/pkg is missing or the upload errors (see upload_to_braintrust).
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
try:
|
||||
from upload_to_braintrust import autoupload
|
||||
|
||||
autoupload(
|
||||
[
|
||||
f"{written['train']}={args.name}_train_{args.date}",
|
||||
f"{written['holdout']}={args.name}_holdout_{args.date}",
|
||||
],
|
||||
run_label=os.getenv("OJ_RUN_LABEL", f"{args.name}_splits_{args.date}"),
|
||||
description=f"{args.name} orchestrator-SFT splits (make_splits.py)",
|
||||
)
|
||||
except Exception as exc: # telemetry must never break the data pipeline
|
||||
print(f"[braintrust] autoupload hook skipped ({exc})")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,322 @@
|
||||
#!/usr/bin/env python
|
||||
"""Render an orchestrator SFT JSONL into human-viewable companions.
|
||||
|
||||
The trainer-facing ``*.jsonl`` packs every trajectory onto one line with all
|
||||
newlines escaped, which is unreadable by eye. This writes two siblings next to
|
||||
it so you can actually inspect what the orchestrator did:
|
||||
|
||||
* ``<name>.pretty.json`` — the same records as an indented JSON array.
|
||||
* ``<name>.txt`` — a clean plain-text transcript. The shared system
|
||||
prompt + tool catalog (identical across every record) is printed ONCE at the
|
||||
top; each trajectory then shows just its turns: the user question, the
|
||||
assistant's reasoning/answer, tool calls rendered as ``-> name({args})``, and
|
||||
tool results truncated so a single web-search dump doesn't bury the trace.
|
||||
|
||||
Usage:
|
||||
.venv/bin/python scripts/orchestrator/render_sft_data.py <path.jsonl>
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html as _htmllib
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
# Tool observations (web_search dumps, code stdout, file reads) can be huge; cap
|
||||
# them in the transcript so the actual orchestration stays legible. The full
|
||||
# content is always in the .jsonl / .pretty.json.
|
||||
_OBS_CAP = 1200
|
||||
_RULE = "-" * 80
|
||||
_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
|
||||
|
||||
|
||||
def _parse_tool_calls(content: str):
|
||||
"""Parse the model's ``<tool_call>{json}</tool_call>`` tags into a list of
|
||||
parsed ``[(name, arguments_dict), ...]`` calls, de-duplicating the
|
||||
double-emit the template produces (each call shows up twice)."""
|
||||
seen = []
|
||||
for m in _TOOL_CALL_RE.finditer(content):
|
||||
try:
|
||||
call = json.loads(m.group(1))
|
||||
item = (call.get("name"), call.get("arguments", {}))
|
||||
except Exception:
|
||||
item = (None, {"_raw": m.group(1)})
|
||||
key = json.dumps(item, ensure_ascii=False, sort_keys=True, default=str)
|
||||
if not seen or seen[-1][0] != key: # collapse the adjacent duplicate
|
||||
seen.append((key, item))
|
||||
return [it for _, it in seen]
|
||||
|
||||
|
||||
def _call_label(name, anon_map) -> str:
|
||||
"""Display a call name; if it's an anonymized expert label, append the real
|
||||
model it maps to (from the record's ``metrics.anon_map``)."""
|
||||
real = (anon_map or {}).get(name)
|
||||
return f"{name} -> {real}" if real else f"{name}"
|
||||
|
||||
|
||||
def _fmt_tool_calls(content: str, anon_map=None) -> str:
|
||||
"""Plain-text tool calls with real (un-escaped) newlines in each arg value."""
|
||||
blocks = []
|
||||
for name, args in _parse_tool_calls(content):
|
||||
lines = [f"-> {_call_label(name, anon_map)}"]
|
||||
if isinstance(args, dict):
|
||||
for k, v in args.items():
|
||||
val = (
|
||||
v
|
||||
if isinstance(v, str)
|
||||
else json.dumps(v, ensure_ascii=False, indent=2)
|
||||
)
|
||||
indented = "\n".join(" " + ln for ln in val.splitlines())
|
||||
lines.append(f" {k}:\n{indented}")
|
||||
blocks.append("\n".join(lines))
|
||||
return "\n".join(blocks)
|
||||
|
||||
|
||||
def _clip(text: str, cap: int) -> str:
|
||||
text = text or ""
|
||||
if len(text) <= cap:
|
||||
return text
|
||||
return f"{text[:cap]}\n ... [truncated {len(text) - cap} chars]"
|
||||
|
||||
|
||||
def _render_turn(turn: dict, anon_map=None) -> str:
|
||||
role = str(turn.get("role", "?")).lower()
|
||||
content = turn.get("content", "") or ""
|
||||
|
||||
if role == "user":
|
||||
return f"USER:\n{content.strip()}"
|
||||
|
||||
if role in ("tool", "function"):
|
||||
name = turn.get("name", "tool")
|
||||
return f"TOOL [{name}]:\n{_clip(content.strip(), _OBS_CAP)}"
|
||||
|
||||
if role == "assistant":
|
||||
calls = _fmt_tool_calls(content, anon_map)
|
||||
# Strip the raw <tool_call> tags out of the prose so it isn't shown twice.
|
||||
prose = _TOOL_CALL_RE.sub("", content).strip()
|
||||
parts = []
|
||||
if prose:
|
||||
parts.append(f"ASSISTANT:\n{prose}")
|
||||
if calls:
|
||||
parts.append(f"ASSISTANT calls:\n{calls}")
|
||||
return "\n".join(parts) if parts else "ASSISTANT: (empty)"
|
||||
|
||||
# system handled separately (printed once); fall through for anything else
|
||||
return f"{role.upper()}:\n{content.strip()}"
|
||||
|
||||
|
||||
def _transcript(record: dict, index: int) -> str:
|
||||
m = record.get("metrics", {}) or {}
|
||||
correct = record.get("correct")
|
||||
kept = record.get("kept")
|
||||
flags = []
|
||||
if correct is not None:
|
||||
flags.append("correct" if correct else "WRONG")
|
||||
if kept is not None and kept:
|
||||
flags.append("kept")
|
||||
head = (
|
||||
f"### [{index}] task={record.get('task_id')} domain={record.get('domain')}"
|
||||
+ (f" {' '.join(flags)}" if flags else "")
|
||||
+ f"\n cost=${m.get('cost_usd')} tokens={m.get('tokens')} "
|
||||
f"tool_calls={m.get('num_tool_calls')} turns={m.get('num_turns')} "
|
||||
f"reward={record.get('reward')}"
|
||||
)
|
||||
anon_map = m.get("anon_map")
|
||||
body = [
|
||||
_render_turn(t, anon_map)
|
||||
for t in record.get("conversations", [])
|
||||
if str(t.get("role", "")).lower() != "system"
|
||||
]
|
||||
return head + "\n\n" + "\n\n".join(body)
|
||||
|
||||
|
||||
def _system_header(records: List[dict]) -> str:
|
||||
"""The system message is identical across records — render it once."""
|
||||
for r in records:
|
||||
for t in r.get("conversations", []):
|
||||
if str(t.get("role", "")).lower() == "system":
|
||||
return (
|
||||
"=" * 80 + "\nSHARED SYSTEM PROMPT + TOOL CATALOG "
|
||||
"(identical for every record below)\n"
|
||||
+ "=" * 80
|
||||
+ f"\n{t.get('content', '').strip()}\n"
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
_HTML_CSS = """
|
||||
body{font:14px/1.5 system-ui,sans-serif;margin:0;background:#f5f5f7;color:#1d1d1f}
|
||||
header{position:sticky;top:0;background:#1d1d1f;color:#fff;padding:12px 20px;z-index:9}
|
||||
header b{font-size:16px}
|
||||
.rec{background:#fff;margin:14px 20px;border-radius:10px;box-shadow:0 1px 3px #0002;overflow:hidden}
|
||||
.rec>summary{cursor:pointer;padding:10px 16px;font-weight:600;list-style:none;display:flex;gap:14px;align-items:center;flex-wrap:wrap}
|
||||
.rec>summary::-webkit-details-marker{display:none}
|
||||
.rec>summary:hover{background:#fafafa}
|
||||
.body{padding:6px 16px 16px}
|
||||
.turn{margin:10px 0;padding:10px 14px;border-radius:8px;white-space:pre-wrap;word-break:break-word}
|
||||
.user{background:#e8f0fe}
|
||||
.assistant{background:#f1f8e9}
|
||||
.call{background:#fff3e0}
|
||||
.call .cname{font-family:ui-monospace,monospace;font-weight:700;color:#b25b00;margin-bottom:6px}
|
||||
.call .cname .real{background:#b25b00;color:#fff;padding:1px 7px;border-radius:10px;font-size:12px;margin-left:4px}
|
||||
.call .arg{margin:6px 0 0 14px}
|
||||
.call .arg b{font-family:ui-monospace,monospace;color:#7a4500}
|
||||
.call .arg pre{white-space:pre-wrap;word-break:break-word;margin:3px 0 0;padding:8px 10px;background:#fff;border:1px solid #f0d9b8;border-radius:6px;font-size:13px}
|
||||
.tool details{background:#fafafa;border:1px solid #eee;border-radius:8px;padding:6px 10px;margin:10px 0}
|
||||
.tool summary{cursor:pointer;color:#555;font-weight:600}
|
||||
.tool pre{white-space:pre-wrap;word-break:break-word;margin:8px 0 0;font-size:13px}
|
||||
.role{font-size:11px;text-transform:uppercase;letter-spacing:.5px;color:#888;margin-bottom:4px}
|
||||
.pill{font-size:12px;padding:2px 8px;border-radius:10px;font-weight:600}
|
||||
.ok{background:#d7f5dd;color:#1a7f37}.wrong{background:#ffe0e0;color:#c0392b}
|
||||
.kept{background:#fff0c2;color:#9a6700}.dom{background:#eee;color:#444}
|
||||
.meta{font-weight:400;color:#888;font-size:12px}
|
||||
#sys details{margin:14px 20px;background:#fff;border-radius:10px;padding:10px 16px}
|
||||
#sys pre{white-space:pre-wrap;word-break:break-word;font-size:12px;color:#444}
|
||||
"""
|
||||
|
||||
|
||||
def _esc(s: str) -> str:
|
||||
return _htmllib.escape(s or "")
|
||||
|
||||
|
||||
def _html_turns(record: dict) -> str:
|
||||
anon_map = (record.get("metrics", {}) or {}).get("anon_map") or {}
|
||||
out = []
|
||||
for t in record.get("conversations", []):
|
||||
role = str(t.get("role", "")).lower()
|
||||
content = t.get("content", "") or ""
|
||||
if role == "system":
|
||||
continue
|
||||
if role == "user":
|
||||
out.append(
|
||||
f'<div class="turn user"><div class="role">user</div>{_esc(content.strip())}</div>'
|
||||
)
|
||||
elif role in ("tool", "function"):
|
||||
name = _esc(t.get("name", "tool"))
|
||||
out.append(
|
||||
f'<div class="tool"><details><summary>tool result · {name} '
|
||||
f"({len(content)} chars)</summary><pre>{_esc(content.strip())}</pre></details></div>"
|
||||
)
|
||||
elif role == "assistant":
|
||||
prose = _TOOL_CALL_RE.sub("", content).strip()
|
||||
if prose:
|
||||
out.append(
|
||||
f'<div class="turn assistant"><div class="role">assistant</div>{_esc(prose)}</div>'
|
||||
)
|
||||
for name, args in _parse_tool_calls(content):
|
||||
real = anon_map.get(name)
|
||||
label = (
|
||||
f'{_esc(str(name))} <span class="real">→ {_esc(str(real))}</span>'
|
||||
if real
|
||||
else _esc(str(name))
|
||||
)
|
||||
rows = [f'<div class="cname">→ {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> {len(records)} records "
|
||||
f"· 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:]))
|
||||
@@ -0,0 +1,407 @@
|
||||
#!/usr/bin/env python
|
||||
"""Full-parameter SFT for the orchestrator policy via FSDP (multi-GPU).
|
||||
|
||||
Launch with accelerate so the model + optimizer shard across GPUs (an 8-9B full
|
||||
fine-tune won't fit one L40S otherwise):
|
||||
|
||||
accelerate launch --config_file <fsdp.yaml> \
|
||||
scripts/orchestrator/run_sft_fsdp.py \
|
||||
--data data/orchestrator/sft/qwen_train_0707.jsonl [--data <more> ...] \
|
||||
--variant correct --model Qwen/Qwen3.5-9B \
|
||||
--out checkpoints/sft_correct --epochs 3 --batch-size 8 --grad-accum 8 \
|
||||
--wandb-project orchestrator-sft
|
||||
|
||||
--variant selects which trajectories to train on:
|
||||
all -> every record
|
||||
correct -> only records the verifier marked correct
|
||||
correct_routed -> only correct records that delegated to a model expert
|
||||
|
||||
By default (--require-clean) records explicitly marked clean=False by the clean
|
||||
gate (bloated / garbled / unrouted) are also dropped; rows missing the flag
|
||||
(legacy data) are kept. Pass --no-require-clean to skip this filter.
|
||||
|
||||
Reuses the conversation->tokens + assistant-only masking from sft_tokenize.py.
|
||||
On these no-NVLink L40S, the accelerate/FSDP NCCL env (NCCL_P2P_DISABLE=1 etc.)
|
||||
must be set by the launcher.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# reuse the data builder from the LoRA launcher
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from sft_tokenize import build_example # noqa: E402
|
||||
|
||||
MODEL_EXPERTS = {
|
||||
"gpt_5_5",
|
||||
"claude_opus_4_8",
|
||||
"qwen3_5_9b",
|
||||
"qwen3_6_27b_fp8",
|
||||
"qwen3_5_122b_a10b_fp8",
|
||||
"qwen3_5_397b_a17b_fp8",
|
||||
}
|
||||
|
||||
|
||||
def record_is_correct(r: dict) -> bool:
|
||||
c = r.get("correct")
|
||||
if c is None:
|
||||
c = r.get("metrics", {}).get("correct")
|
||||
return bool(c)
|
||||
|
||||
|
||||
def record_routed_to_expert(r: dict) -> bool:
|
||||
"""True if the trajectory called a model expert (resolving anon labels)."""
|
||||
amap = r.get("metrics", {}).get("anon_map", {})
|
||||
names = re.findall(
|
||||
r'"name"\s*:\s*"([a-z0-9_]+)"', json.dumps(r.get("conversations", []))
|
||||
)
|
||||
for n in names:
|
||||
real = amap.get(n, n)
|
||||
if real in MODEL_EXPERTS:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def record_clean_flag(r: dict):
|
||||
"""The record's `clean` flag: True/False, or None if absent (legacy data).
|
||||
|
||||
None means the clean gate never ran on this row, so we can't judge it -> we
|
||||
keep it (legacy behavior) rather than silently dropping older datasets.
|
||||
"""
|
||||
c = r.get("clean")
|
||||
if c is None:
|
||||
c = r.get("metrics", {}).get("clean")
|
||||
return c
|
||||
|
||||
|
||||
def select(records, variant: str, require_clean: bool = True):
|
||||
"""Filter by --variant, then (default) drop records explicitly marked unclean.
|
||||
|
||||
require_clean drops rows whose `clean` flag is False. Rows missing the flag
|
||||
(legacy) are kept regardless, matching pre-clean-gate behavior.
|
||||
"""
|
||||
if variant == "all":
|
||||
base = list(records)
|
||||
elif variant == "correct":
|
||||
base = [r for r in records if record_is_correct(r)]
|
||||
elif variant == "correct_routed":
|
||||
base = [
|
||||
r for r in records if record_is_correct(r) and record_routed_to_expert(r)
|
||||
]
|
||||
else:
|
||||
raise ValueError(f"unknown variant {variant}")
|
||||
if not require_clean:
|
||||
return base
|
||||
return [r for r in base if record_clean_flag(r) is not False]
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument(
|
||||
"--data", action="append", required=True, help="JSONL(s) (repeatable)"
|
||||
)
|
||||
p.add_argument(
|
||||
"--val-data",
|
||||
default=None,
|
||||
help="Held-out JSONL for val-loss (excluded from --data). Eval'd per epoch.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--variant", choices=["all", "correct", "correct_routed"], default="correct"
|
||||
)
|
||||
p.add_argument(
|
||||
"--require-clean",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Drop records whose `clean` flag is False (bloated / garbled / "
|
||||
"unrouted). Rows missing the flag (legacy data) are kept. "
|
||||
"Use --no-require-clean to disable.",
|
||||
)
|
||||
p.add_argument("--model", default="Qwen/Qwen3.5-9B")
|
||||
p.add_argument("--out", required=True)
|
||||
p.add_argument("--epochs", type=float, default=3.0)
|
||||
p.add_argument("--batch-size", type=int, default=8)
|
||||
p.add_argument("--grad-accum", type=int, default=8)
|
||||
p.add_argument("--lr", type=float, default=1e-5)
|
||||
p.add_argument("--max-seq", type=int, default=8192)
|
||||
p.add_argument(
|
||||
"--supervise-last-only",
|
||||
action="store_true",
|
||||
help="Legacy: supervise ONLY the final assistant turn. Default "
|
||||
"(off) supervises every assistant turn incl. routing.",
|
||||
)
|
||||
p.add_argument("--warmup-ratio", type=float, default=0.03)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--wandb-project", default="orchestrator-sft")
|
||||
p.add_argument("--wandb-name", default="")
|
||||
args = p.parse_args()
|
||||
|
||||
import random
|
||||
from datetime import timedelta
|
||||
|
||||
import torch
|
||||
from accelerate import Accelerator
|
||||
from accelerate.utils import InitProcessGroupKwargs, set_seed
|
||||
from torch.utils.data import DataLoader
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
set_seed(args.seed)
|
||||
# 30-min collective timeout: a slow first-step shard gather / graph build on a
|
||||
# contended node was tripping the default 10-min NCCL watchdog (job 3892 died
|
||||
# at step 1). This only delays a *real* hang's crash; it fixes spurious ones.
|
||||
pg_timeout_min = int(os.environ.get("SFT_PG_TIMEOUT_MIN", "30"))
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=args.grad_accum,
|
||||
kwargs_handlers=[
|
||||
InitProcessGroupKwargs(timeout=timedelta(minutes=pg_timeout_min))
|
||||
],
|
||||
)
|
||||
is_main = accelerator.is_main_process
|
||||
|
||||
def log(m):
|
||||
if is_main:
|
||||
print(f"[fsdp-sft] {time.strftime('%H:%M:%S')} {m}", flush=True)
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
if tok.pad_token_id is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
|
||||
# ---- data ----
|
||||
records = []
|
||||
for f in args.data:
|
||||
records += [json.loads(l) for l in open(f) if l.strip()]
|
||||
sel = select(records, args.variant, require_clean=args.require_clean)
|
||||
log(f"variant={args.variant}: {len(sel)}/{len(records)} records")
|
||||
if args.require_clean:
|
||||
pre_clean = select(records, args.variant, require_clean=False)
|
||||
n_unclean = len(pre_clean) - len(sel)
|
||||
log(
|
||||
f"require_clean=True: dropped {n_unclean} unclean rows "
|
||||
f"({len(pre_clean)}->{len(sel)}); use --no-require-clean to keep them"
|
||||
)
|
||||
else:
|
||||
log("require_clean=False: NOT filtering on `clean` flag")
|
||||
examples = []
|
||||
for r in sel:
|
||||
ex = build_example(
|
||||
tok,
|
||||
r.get("conversations", []),
|
||||
args.max_seq,
|
||||
supervise_all_turns=not args.supervise_last_only,
|
||||
)
|
||||
if ex:
|
||||
examples.append(ex)
|
||||
log(f"built {len(examples)} training examples")
|
||||
if not examples:
|
||||
log("no examples; abort")
|
||||
return 1
|
||||
random.Random(args.seed).shuffle(examples)
|
||||
|
||||
def collate(batch):
|
||||
maxlen = max(len(b["input_ids"]) for b in batch)
|
||||
ii, lab, am = [], [], []
|
||||
for b in batch:
|
||||
pad = maxlen - len(b["input_ids"])
|
||||
ii.append(b["input_ids"] + [tok.pad_token_id] * pad)
|
||||
lab.append(b["labels"] + [-100] * pad)
|
||||
am.append([1] * len(b["input_ids"]) + [0] * pad)
|
||||
return (torch.tensor(ii), torch.tensor(lab), torch.tensor(am))
|
||||
|
||||
dl = DataLoader(
|
||||
examples, batch_size=args.batch_size, shuffle=True, collate_fn=collate
|
||||
)
|
||||
|
||||
# ---- held-out val set (leak-free; excluded from --data upstream) ----
|
||||
val_dl = None
|
||||
if args.val_data:
|
||||
vrecs = [json.loads(l) for l in open(args.val_data) if l.strip()]
|
||||
vex = []
|
||||
for r in vrecs:
|
||||
ex = build_example(
|
||||
tok,
|
||||
r.get("conversations", []),
|
||||
args.max_seq,
|
||||
supervise_all_turns=not args.supervise_last_only,
|
||||
)
|
||||
if ex:
|
||||
vex.append(ex)
|
||||
log(f"val: {len(vex)} examples from {args.val_data}")
|
||||
if vex:
|
||||
val_dl = DataLoader(
|
||||
vex, batch_size=args.batch_size, shuffle=False, collate_fn=collate
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model,
|
||||
dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
attn_implementation="sdpa",
|
||||
)
|
||||
if os.environ.get("SFT_NO_GRAD_CKPT") == "1":
|
||||
log("gradient checkpointing DISABLED (SFT_NO_GRAD_CKPT=1)")
|
||||
else:
|
||||
model.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
model.config.use_cache = False
|
||||
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.0)
|
||||
|
||||
model, opt, dl = accelerator.prepare(model, opt, dl)
|
||||
if val_dl is not None:
|
||||
val_dl = accelerator.prepare(val_dl)
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate():
|
||||
# Corpus-level mean token loss over the holdout. FSDP needs every rank to
|
||||
# run the forward collectively, so each rank evals its shard and we reduce
|
||||
# token-weighted sums (not a mean-of-batch-means, which would mis-weight).
|
||||
model.eval()
|
||||
tot_loss = torch.zeros((), device=accelerator.device)
|
||||
tot_tok = torch.zeros((), device=accelerator.device)
|
||||
for ii, lab, am in val_dl:
|
||||
out = model(input_ids=ii, attention_mask=am, labels=lab)
|
||||
ntok = (lab != -100).sum()
|
||||
tot_loss += out.loss.detach() * ntok
|
||||
tot_tok += ntok
|
||||
tot_loss = accelerator.reduce(tot_loss, reduction="sum")
|
||||
tot_tok = accelerator.reduce(tot_tok, reduction="sum")
|
||||
model.train()
|
||||
return (tot_loss / tot_tok.clamp(min=1)).item()
|
||||
|
||||
steps_per_epoch = math.ceil(len(dl) / args.grad_accum)
|
||||
total_steps = max(1, int(steps_per_epoch * args.epochs))
|
||||
warmup = max(1, int(total_steps * args.warmup_ratio))
|
||||
|
||||
def lr_at(s):
|
||||
if s < warmup:
|
||||
return s / warmup
|
||||
return 0.5 * (
|
||||
1 + math.cos(math.pi * (s - warmup) / max(1, total_steps - warmup))
|
||||
)
|
||||
|
||||
use_wandb = False
|
||||
if is_main:
|
||||
try:
|
||||
import wandb
|
||||
|
||||
base_name = (
|
||||
args.wandb_name or f"{Path(args.model).name}-{args.variant}-fsdp"
|
||||
)
|
||||
run_name = f"{base_name}-{time.strftime('%m%d-%H%M', time.gmtime())}"
|
||||
wandb.init(
|
||||
project=args.wandb_project,
|
||||
name=run_name,
|
||||
config=vars(args)
|
||||
| {"n_examples": len(examples), "total_steps": total_steps},
|
||||
)
|
||||
use_wandb = True
|
||||
except Exception as e:
|
||||
log(f"wandb off ({e})")
|
||||
|
||||
log(f"total_steps={total_steps} steps/epoch={steps_per_epoch} warmup={warmup}")
|
||||
|
||||
def _save_ckpt(out_path):
|
||||
# Full-model checkpoint (gathers FSDP shards to rank0). Called per-epoch
|
||||
# so a long run leaves usable partial models + lets us eval intermediate
|
||||
# checkpoints (e.g. epoch1/epoch2) for the data-scaling sweep.
|
||||
accelerator.wait_for_everyone()
|
||||
unwrapped = accelerator.unwrap_model(model)
|
||||
p = Path(out_path)
|
||||
if is_main:
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
tok.save_pretrained(str(p))
|
||||
unwrapped.save_pretrained(
|
||||
str(p),
|
||||
is_main_process=is_main,
|
||||
save_function=accelerator.save,
|
||||
state_dict=accelerator.get_state_dict(model),
|
||||
)
|
||||
if is_main:
|
||||
log(f"checkpoint saved -> {p}")
|
||||
|
||||
model.train()
|
||||
gstep = 0
|
||||
t0 = time.time()
|
||||
if val_dl is not None:
|
||||
vloss = evaluate()
|
||||
log(f"val/loss (baseline, step 0) = {vloss:.4f}")
|
||||
if use_wandb:
|
||||
wandb.log({"val/loss": vloss}, step=0)
|
||||
for epoch in range(math.ceil(args.epochs)):
|
||||
for ii, lab, am in dl:
|
||||
with accelerator.accumulate(model):
|
||||
out = model(input_ids=ii, attention_mask=am, labels=lab)
|
||||
accelerator.backward(out.loss)
|
||||
if accelerator.sync_gradients:
|
||||
lr = args.lr * lr_at(gstep)
|
||||
for g in opt.param_groups:
|
||||
g["lr"] = lr
|
||||
accelerator.clip_grad_norm_(model.parameters(), 1.0)
|
||||
opt.step()
|
||||
opt.zero_grad()
|
||||
if accelerator.sync_gradients:
|
||||
gstep += 1
|
||||
# accelerator.gather() is a COLLECTIVE — EVERY rank must call it.
|
||||
# Guarding it behind `is_main` made only rank0 run the [1]-float
|
||||
# all-gather while ranks 1-3 marched on to the next step's FSDP
|
||||
# param all-gather, desyncing the process group and hanging the job
|
||||
# at the step1->step2 boundary (NCCL collective-shape mismatch).
|
||||
loss = accelerator.gather(out.loss.detach()).mean().item()
|
||||
if is_main:
|
||||
log(f"step {gstep}/{total_steps} loss={loss:.4f} lr={lr:.2e}")
|
||||
if use_wandb:
|
||||
wandb.log({"train/loss": loss, "train/lr": lr}, step=gstep)
|
||||
if gstep >= total_steps:
|
||||
break
|
||||
_save_ckpt(Path(args.out) / f"epoch{epoch + 1}")
|
||||
if val_dl is not None:
|
||||
vloss = evaluate()
|
||||
log(f"val/loss (epoch {epoch + 1}) = {vloss:.4f}")
|
||||
if use_wandb:
|
||||
wandb.log({"val/loss": vloss, "epoch": epoch + 1}, step=gstep)
|
||||
if gstep >= total_steps:
|
||||
break
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
log("saving full model...")
|
||||
unwrapped = accelerator.unwrap_model(model)
|
||||
out_dir = Path(args.out)
|
||||
if is_main:
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
tok.save_pretrained(str(out_dir))
|
||||
unwrapped.save_pretrained(
|
||||
str(out_dir),
|
||||
is_main_process=is_main,
|
||||
save_function=accelerator.save,
|
||||
state_dict=accelerator.get_state_dict(model),
|
||||
)
|
||||
if is_main:
|
||||
log(f"saved -> {out_dir}")
|
||||
if use_wandb:
|
||||
try:
|
||||
wandb.finish()
|
||||
except Exception:
|
||||
pass
|
||||
print(
|
||||
"FSDP_SFT_DONE "
|
||||
+ json.dumps(
|
||||
{
|
||||
"variant": args.variant,
|
||||
"steps": gstep,
|
||||
"wall_s": round(time.time() - t0, 1),
|
||||
"out": str(out_dir),
|
||||
}
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,151 @@
|
||||
#!/usr/bin/env python
|
||||
"""Conversation->tokens + assistant-only masking shared by the SFT trainers.
|
||||
|
||||
Extracted from the old LoRA trainer so ``run_sft_fsdp.py`` (full-parameter FSDP,
|
||||
the path we actually use) doesn't depend on it. ``build_example`` tokenizes one
|
||||
``conversations`` record and returns ``input_ids`` + ``labels`` with everything
|
||||
but the supervised assistant turns masked to -100.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
def normalize_messages(conversations: List[Dict[str, Any]]) -> List[Dict[str, str]]:
|
||||
"""Map raw conversation turns to {role, content} with role in
|
||||
system/user/assistant. tool turns fold into user (matches the native
|
||||
OrchestratorSFTDataset fallback)."""
|
||||
msgs: List[Dict[str, str]] = []
|
||||
for turn in conversations:
|
||||
role = turn.get("role") or turn.get("from", "")
|
||||
content = turn.get("content") or turn.get("value", "") or ""
|
||||
if role in ("human", "user"):
|
||||
role = "user"
|
||||
elif role in ("gpt", "assistant"):
|
||||
role = "assistant"
|
||||
elif role == "tool":
|
||||
content = f"[Tool '{turn.get('name', 'tool')}' returned]: {content}"
|
||||
role = "user"
|
||||
if role in ("system", "user", "assistant"):
|
||||
msgs.append({"role": role, "content": str(content)})
|
||||
return msgs
|
||||
|
||||
|
||||
def _lcp_len(a: List[int], b: List[int]) -> int:
|
||||
"""Length of the longest common prefix of two token-id lists."""
|
||||
n = 0
|
||||
for x, y in zip(a, b):
|
||||
if x != y:
|
||||
break
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def _chatml_assistant_spans(ids: List[int], tok) -> Optional[List[tuple]]:
|
||||
"""``(start, end)`` token spans of each assistant turn's content in a
|
||||
ChatML-rendered sequence, located by ``<|im_start|>assistant ... <|im_end|>``
|
||||
markers (the ``<|im_end|>`` is included so the model learns to stop).
|
||||
|
||||
Operates on the FINAL token stream, so it is correct even when the chat
|
||||
template strips prior-turn ``<think>`` reasoning from history (Qwen3.x does),
|
||||
which breaks any incremental re-rendering / prefix approach. Returns None for
|
||||
non-ChatML templates (e.g. gemma) so the caller can fall back."""
|
||||
im_start = tok.convert_tokens_to_ids("<|im_start|>")
|
||||
im_end = tok.convert_tokens_to_ids("<|im_end|>")
|
||||
asst = tok.convert_tokens_to_ids("assistant")
|
||||
unk = tok.unk_token_id
|
||||
if im_start is None or im_end is None or im_start == unk or im_end == unk:
|
||||
return None
|
||||
nl = tok.convert_tokens_to_ids("Ċ") # the '\n' after the role header
|
||||
spans: List[tuple] = []
|
||||
i, n = 0, len(ids)
|
||||
while i < n:
|
||||
if ids[i] == im_start and i + 1 < n and ids[i + 1] == asst:
|
||||
j = i + 2
|
||||
if j < n and ids[j] == nl:
|
||||
j += 1
|
||||
k = j
|
||||
while k < n and ids[k] != im_end:
|
||||
k += 1
|
||||
end = min(k + 1, n)
|
||||
if end > j:
|
||||
spans.append((j, end))
|
||||
i = end
|
||||
else:
|
||||
i += 1
|
||||
return spans
|
||||
|
||||
|
||||
def build_example(
|
||||
tok,
|
||||
conversations,
|
||||
max_seq: int,
|
||||
supervise_all_turns: bool = True,
|
||||
) -> Optional[Dict[str, List[int]]]:
|
||||
"""Tokenize one record -> input_ids + labels.
|
||||
|
||||
``supervise_all_turns=True`` (default): supervise EVERY assistant turn — the
|
||||
intermediate routing ``<tool_call>`` turns AND the final answer; only
|
||||
system/user/tool tokens are masked. This teaches the model to actually route,
|
||||
not just synthesize the last answer given someone else's tool calls.
|
||||
|
||||
``supervise_all_turns=False`` (legacy): supervise only the final assistant
|
||||
turn; everything before it is masked to -100.
|
||||
|
||||
Assistant spans are found by the ChatML turn markers on the rendered stream
|
||||
(robust to Qwen3's prior-turn think-stripping). For non-ChatML templates
|
||||
(gemma) we fall back to the proven last-turn longest-common-prefix boundary;
|
||||
all-turns supervision there needs per-template markers and is not yet wired,
|
||||
so it degrades to last-turn. Records whose last turn isn't assistant, that
|
||||
supervise nothing, or whose final turn alone overflows max_seq are skipped."""
|
||||
msgs = normalize_messages(conversations)
|
||||
if len(msgs) < 2 or msgs[-1]["role"] != "assistant":
|
||||
return None
|
||||
try:
|
||||
full_text = tok.apply_chat_template(
|
||||
msgs, tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
full = tok(full_text, add_special_tokens=False)["input_ids"]
|
||||
except Exception:
|
||||
return None
|
||||
if not full:
|
||||
return None
|
||||
|
||||
spans = _chatml_assistant_spans(full, tok)
|
||||
if spans:
|
||||
selected = spans if supervise_all_turns else spans[-1:]
|
||||
labels = [-100] * len(full)
|
||||
for s, e in selected:
|
||||
for k in range(s, e):
|
||||
labels[k] = full[k]
|
||||
last_start = selected[-1][0]
|
||||
else:
|
||||
# Non-ChatML template: proven last-turn boundary via longest-common-prefix
|
||||
# (robust to templates whose add_generation_prompt emits extra preamble).
|
||||
try:
|
||||
prompt_text = tok.apply_chat_template(
|
||||
msgs[:-1], tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
prompt = tok(prompt_text, add_special_tokens=False)["input_ids"]
|
||||
except Exception:
|
||||
return None
|
||||
boundary = _lcp_len(prompt, full)
|
||||
if boundary >= len(full):
|
||||
return None
|
||||
labels = list(full)
|
||||
for i in range(min(boundary, len(labels))):
|
||||
labels[i] = -100
|
||||
last_start = boundary
|
||||
|
||||
if all(v == -100 for v in labels): # nothing left to supervise
|
||||
return None
|
||||
|
||||
# Keep the final assistant turn if too long: left-truncate from the front.
|
||||
if len(full) > max_seq:
|
||||
if len(full) - last_start >= max_seq: # final answer alone overflows -> drop
|
||||
return None
|
||||
cut = len(full) - max_seq
|
||||
full = full[cut:]
|
||||
labels = labels[cut:]
|
||||
return {"input_ids": full, "labels": labels}
|
||||
@@ -0,0 +1,439 @@
|
||||
#!/usr/bin/env python
|
||||
"""Upload orchestrator-SFT JSONL datasets to Braintrust so they're browsable/
|
||||
filterable in the UI (by area / difficulty / dataset / routed-model / correct).
|
||||
|
||||
Datasets land in the SAME Braintrust project as the rollout TRACES (the
|
||||
``research`` project) so a run's data + traces sit side by side. The target is
|
||||
resolved from ``OJ_BRAINTRUST_PROJECT_ID`` (default: the research project id);
|
||||
``--project`` overrides by NAME if you want a different project.
|
||||
|
||||
Each record -> a Braintrust dataset row:
|
||||
input = the problem/question
|
||||
expected = the FINAL_ANSWER value
|
||||
tags = [gen_model, domain, correct|incorrect, clean|dirty, kept|dropped]
|
||||
metadata = domain, area, difficulty, dataset, subsector, task_id, correct,
|
||||
clean, kept, reward, gen_model, orchestrator_model, num_tool_calls,
|
||||
n_turns, routed_models (real names)
|
||||
(the full conversation is kept under metadata.conversation for inspection)
|
||||
|
||||
The dataset itself carries run-level metadata (run_label, specific gen_model,
|
||||
git sha, config knobs, counts, domain distribution) so future experiments are
|
||||
distinguishable in the UI.
|
||||
|
||||
Usage (CLI):
|
||||
.venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
|
||||
data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707 \
|
||||
data/orchestrator/sft/qwen_holdout_0707.jsonl=qwen_holdout_0707
|
||||
|
||||
# name is optional; defaults to {model_short}_{split}_{date}
|
||||
.venv/bin/python scripts/orchestrator/upload_to_braintrust.py \
|
||||
data/orchestrator/sft/qwen_holdout_0707.jsonl
|
||||
|
||||
Programmatic (used by the pipeline auto-upload hook in make_splits.py):
|
||||
from upload_to_braintrust import autoupload
|
||||
autoupload(["data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707"],
|
||||
run_label="...")
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
from collections import Counter
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The `research` project — datasets go here, same place the rollout traces land.
|
||||
DEFAULT_PROJECT_ID = "c707124e-ce9d-4187-ad11-f49f19f777ad"
|
||||
|
||||
_FA = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
|
||||
|
||||
|
||||
def _ordered_turns(convs, anon_map, orchestrator):
|
||||
"""Rebuild each turn as role/model/content, in that display order.
|
||||
|
||||
Braintrust sorts object keys alphabetically on write, which would push the
|
||||
(very long) `content` to the top. We prefix the keys (`1_role`, `2_model`,
|
||||
`3_content`) so they sort into role -> model -> content and the content
|
||||
renders last. `model` is the model behind the turn:
|
||||
* assistant turns -> the orchestrator model
|
||||
* tool turns -> the expert model that answered (de-anonymized)
|
||||
* system/user -> None
|
||||
"""
|
||||
out = []
|
||||
for c in convs:
|
||||
role = c.get("role")
|
||||
if role == "tool":
|
||||
model = anon_map.get(c.get("name")) or c.get("name")
|
||||
elif role == "assistant":
|
||||
model = orchestrator
|
||||
else:
|
||||
model = None
|
||||
out.append({"1_role": role, "2_model": model, "3_content": c.get("content")})
|
||||
return out
|
||||
|
||||
|
||||
def _dataset_desc(dataset):
|
||||
"""Fallback dataset blurb for records generated before the field existed."""
|
||||
try:
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
dataset_description,
|
||||
)
|
||||
|
||||
return dataset_description(dataset)
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def _git_sha():
|
||||
try:
|
||||
return (
|
||||
subprocess.check_output(
|
||||
["git", "rev-parse", "--short", "HEAD"],
|
||||
cwd=Path(__file__).resolve().parent,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
.decode()
|
||||
.strip()
|
||||
or None
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _model_short(rows, path):
|
||||
"""Short family tag for the default dataset name (qwen / gemma / ...)."""
|
||||
stem = Path(path).stem.lower()
|
||||
for fam in ("qwen", "gemma"):
|
||||
if stem.startswith(fam) or fam in stem:
|
||||
return fam
|
||||
gm = next((r.get("gen_model") for r in rows if r.get("gen_model")), "") or ""
|
||||
gm = gm.lower()
|
||||
for fam in ("qwen", "gemma"):
|
||||
if fam in gm:
|
||||
return fam
|
||||
return "orch"
|
||||
|
||||
|
||||
def _split_of(path):
|
||||
stem = Path(path).stem.lower()
|
||||
for s in (
|
||||
"holdout",
|
||||
"train",
|
||||
"clean",
|
||||
"overfit",
|
||||
"partial",
|
||||
"8k",
|
||||
"4k",
|
||||
"2k",
|
||||
"1k",
|
||||
):
|
||||
if s in stem:
|
||||
return s
|
||||
return "data"
|
||||
|
||||
|
||||
def _default_name(rows, path):
|
||||
return f"{_model_short(rows, path)}_{_split_of(path)}_{datetime.now():%m%d}"
|
||||
|
||||
|
||||
def _config_from_env():
|
||||
"""Config knobs from the run env (set by build_orchestrator_sft.py). Omit unset."""
|
||||
cfg = {}
|
||||
for env_key, meta_key, cast in [
|
||||
("OJ_CFG_TEMPERATURE", "temperature", float),
|
||||
("OJ_CFG_MAX_TURNS", "max_turns", int),
|
||||
(
|
||||
"OJ_CFG_ANONYMIZE",
|
||||
"anonymize",
|
||||
lambda v: v.strip().lower() in ("1", "true", "yes"),
|
||||
),
|
||||
(
|
||||
"OJ_CFG_REJECTION_ONLY",
|
||||
"rejection_only",
|
||||
lambda v: v.strip().lower() in ("1", "true", "yes"),
|
||||
),
|
||||
]:
|
||||
v = os.getenv(env_key)
|
||||
if v not in (None, ""):
|
||||
try:
|
||||
cfg[meta_key] = cast(v)
|
||||
except Exception:
|
||||
cfg[meta_key] = v
|
||||
return cfg
|
||||
|
||||
|
||||
def _row(rec):
|
||||
convs = rec.get("conversations", [])
|
||||
user = next((c["content"] for c in convs if c["role"] == "user"), "")
|
||||
question = re.sub(r"^\s*Problem:\s*", "", user, count=1).strip()
|
||||
asst = [c["content"] for c in convs if c["role"] == "assistant"]
|
||||
fa = asst[-1] if asst else ""
|
||||
m = list(_FA.finditer(fa))
|
||||
model_answer = fa[m[-1].end() :].strip() if m else fa.strip()
|
||||
# Gold reference the verifier graded against (stamped at generation time as
|
||||
# `gold_answer`). `expected` in Braintrust is the GOLD so the UI shows
|
||||
# gold-vs-model; the model's own answer goes to metadata.model_answer.
|
||||
gold = rec.get("gold_answer") or None
|
||||
has_gold = gold is not None
|
||||
am = (rec.get("metrics", {}) or {}).get("anon_map", {}) or {}
|
||||
routed = []
|
||||
for c in convs:
|
||||
if c["role"] == "assistant":
|
||||
for t in re.findall(
|
||||
r"<tool_call>(\{.*?\})</tool_call>", c["content"], re.DOTALL
|
||||
):
|
||||
try:
|
||||
real = am.get(json.loads(t).get("name"))
|
||||
if real:
|
||||
routed.append(real)
|
||||
except Exception:
|
||||
pass
|
||||
metrics = rec.get("metrics", {}) or {}
|
||||
gen_model = rec.get("gen_model")
|
||||
domain = rec.get("domain")
|
||||
correct = rec.get("correct")
|
||||
clean = rec.get("clean")
|
||||
kept = rec.get("kept")
|
||||
tags = [
|
||||
t
|
||||
for t in [
|
||||
gen_model,
|
||||
f"domain:{domain}" if domain else None,
|
||||
"correct" if correct else "incorrect",
|
||||
"clean" if clean else "dirty",
|
||||
"kept" if kept else "dropped",
|
||||
"gold" if has_gold else "no_gold",
|
||||
]
|
||||
if t
|
||||
]
|
||||
return {
|
||||
"input": question,
|
||||
"expected": gold,
|
||||
"tags": tags,
|
||||
"metadata": {
|
||||
"gold_answer": gold,
|
||||
"model_answer": model_answer,
|
||||
"has_gold": has_gold,
|
||||
"domain": domain,
|
||||
"area": rec.get("area"),
|
||||
"difficulty": rec.get("difficulty") or None,
|
||||
"dataset": rec.get("dataset"),
|
||||
"dataset_description": rec.get("dataset_description")
|
||||
or _dataset_desc(rec.get("dataset")),
|
||||
"subsector": rec.get("subsector"),
|
||||
"task_id": rec.get("task_id"),
|
||||
"correct": correct,
|
||||
"clean": clean,
|
||||
"kept": kept,
|
||||
"reward": rec.get("reward"),
|
||||
"gen_model": gen_model,
|
||||
"orchestrator_model": rec.get("orchestrator_model"),
|
||||
"num_tool_calls": metrics.get("num_tool_calls"),
|
||||
"n_turns": metrics.get("num_turns"),
|
||||
"routed_models": routed,
|
||||
"conversation": _ordered_turns(convs, am, rec.get("orchestrator_model")),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _dist(rows, key):
|
||||
"""{value: count} for a record field, most-common first, None/'' dropped."""
|
||||
c = Counter(r.get(key) for r in rows if (r.get(key) or "") != "")
|
||||
return {str(k): v for k, v in sorted(c.items(), key=lambda x: -x[1])}
|
||||
|
||||
|
||||
def _routed_dist(rows):
|
||||
"""Distribution of real (de-anonymized) expert models routed to."""
|
||||
c = Counter()
|
||||
for r in rows:
|
||||
am = (r.get("metrics", {}) or {}).get("anon_map", {}) or {}
|
||||
for cc in r.get("conversations", []):
|
||||
if cc.get("role") != "assistant":
|
||||
continue
|
||||
for t in re.findall(
|
||||
r"<tool_call>(\{.*?\})</tool_call>", cc.get("content", ""), re.DOTALL
|
||||
):
|
||||
try:
|
||||
real = am.get(json.loads(t).get("name"))
|
||||
if real:
|
||||
c[real] += 1
|
||||
except Exception:
|
||||
pass
|
||||
return {k: v for k, v in sorted(c.items(), key=lambda x: -x[1])}
|
||||
|
||||
|
||||
def _avg(rows, *path):
|
||||
vals = []
|
||||
for r in rows:
|
||||
v = r.get("metrics", {}) or {}
|
||||
for p in path:
|
||||
v = (v or {}).get(p) if isinstance(v, dict) else None
|
||||
if isinstance(v, (int, float)):
|
||||
vals.append(v)
|
||||
return round(sum(vals) / len(vals), 2) if vals else None
|
||||
|
||||
|
||||
def _dataset_metadata(rows, path, run_label, gen_model):
|
||||
n = len(rows)
|
||||
meta = {
|
||||
"model": _model_short(rows, path),
|
||||
"split": _split_of(path),
|
||||
"task": "orchestrator-SFT unified-tool routing traces",
|
||||
"date": "2026-07-07",
|
||||
"run_label": run_label,
|
||||
"gen_model": gen_model,
|
||||
"git_sha": _git_sha(),
|
||||
"source_file": str(path),
|
||||
"source_datasets": ["GeneralThought-430K-filtered", "OpenThoughts3-1.2M"],
|
||||
"task_mix": "balanced (GeneralThought + OpenThoughts code/math/science)",
|
||||
"filter": "correct + clean (rejects file-write echoes, garble, "
|
||||
"reasoning-degeneration, truncated tails, essays/markdown dumps)",
|
||||
"leak_free": "train/holdout disjoint by task_id",
|
||||
"n_total": n,
|
||||
"n_correct": sum(1 for r in rows if r.get("correct")),
|
||||
"n_clean": sum(1 for r in rows if r.get("clean")),
|
||||
"n_kept": sum(1 for r in rows if r.get("kept")),
|
||||
"n_with_gold": sum(1 for r in rows if (r.get("gold_answer") or "").strip()),
|
||||
"avg_tool_calls": _avg(rows, "num_tool_calls"),
|
||||
"avg_turns": _avg(rows, "num_turns"),
|
||||
"area_distribution": _dist(rows, "area"),
|
||||
"difficulty_distribution": _dist(rows, "difficulty"),
|
||||
"dataset_distribution": _dist(rows, "dataset"),
|
||||
"routed_model_distribution": _routed_dist(rows),
|
||||
}
|
||||
cfg = _config_from_env()
|
||||
if cfg:
|
||||
meta["config"] = cfg
|
||||
return {k: v for k, v in meta.items() if v is not None}
|
||||
|
||||
|
||||
def upload_dataset(
|
||||
path,
|
||||
name=None,
|
||||
*,
|
||||
project_id=None,
|
||||
project=None,
|
||||
run_label=None,
|
||||
gen_model=None,
|
||||
description="",
|
||||
):
|
||||
"""Upload one JSONL file as a Braintrust dataset. Returns (name, n_rows, url).
|
||||
|
||||
Target project: ``project`` (name) if given, else ``project_id`` /
|
||||
``OJ_BRAINTRUST_PROJECT_ID`` / the research default.
|
||||
"""
|
||||
import braintrust
|
||||
|
||||
rows = [json.loads(l) for l in open(path) if l.strip()]
|
||||
name = name or _default_name(rows, path)
|
||||
gen_model = (
|
||||
gen_model
|
||||
or os.getenv("OJ_GEN_MODEL")
|
||||
or next((r.get("gen_model") for r in rows if r.get("gen_model")), None)
|
||||
)
|
||||
run_label = run_label or os.getenv("OJ_RUN_LABEL") or name
|
||||
|
||||
init_kwargs = {
|
||||
"name": name,
|
||||
"description": description or None,
|
||||
"metadata": _dataset_metadata(rows, path, run_label, gen_model),
|
||||
}
|
||||
if project:
|
||||
init_kwargs["project"] = project
|
||||
else:
|
||||
init_kwargs["project_id"] = project_id or os.getenv(
|
||||
"OJ_BRAINTRUST_PROJECT_ID", DEFAULT_PROJECT_ID
|
||||
)
|
||||
|
||||
ds = braintrust.init_dataset(**init_kwargs)
|
||||
for r in rows:
|
||||
ds.insert(**_row(r))
|
||||
ds.flush()
|
||||
summ = ds.summarize()
|
||||
url = getattr(summ, "dataset_url", None) or (
|
||||
project or init_kwargs.get("project_id")
|
||||
)
|
||||
return name, len(rows), url
|
||||
|
||||
|
||||
def autoupload(specs, *, run_label=None, gen_model=None, description=""):
|
||||
"""No-op-safe wrapper for pipeline hooks. Honors OJ_BRAINTRUST_AUTOUPLOAD
|
||||
(default ON) and never raises: any failure (missing key/pkg, network) is
|
||||
logged and swallowed so the data pipeline is never broken by telemetry."""
|
||||
if os.getenv("OJ_BRAINTRUST_AUTOUPLOAD", "1").strip().lower() in (
|
||||
"0",
|
||||
"false",
|
||||
"no",
|
||||
"off",
|
||||
"",
|
||||
):
|
||||
logger.info("[braintrust] autoupload disabled (OJ_BRAINTRUST_AUTOUPLOAD)")
|
||||
return
|
||||
if not os.getenv("BRAINTRUST_API_KEY"):
|
||||
logger.info("[braintrust] autoupload skipped — BRAINTRUST_API_KEY unset")
|
||||
return
|
||||
try:
|
||||
import braintrust # noqa: F401
|
||||
except Exception:
|
||||
logger.info("[braintrust] autoupload skipped — braintrust not installed")
|
||||
return
|
||||
for spec in specs:
|
||||
path, _, name = spec.partition("=")
|
||||
try:
|
||||
nm, n, url = upload_dataset(
|
||||
path,
|
||||
name or None,
|
||||
run_label=run_label,
|
||||
gen_model=gen_model,
|
||||
description=description,
|
||||
)
|
||||
logger.info("[braintrust] uploaded %s: %d rows -> %s", nm, n, url)
|
||||
print(f"[braintrust] uploaded {nm}: {n} rows -> {url}")
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[braintrust] autoupload FAILED for %s (%s) — continuing", path, exc
|
||||
)
|
||||
print(
|
||||
f"[braintrust] autoupload FAILED for {path} ({exc}) — pipeline unaffected"
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument(
|
||||
"--project",
|
||||
default=None,
|
||||
help="Target project by NAME (overrides OJ_BRAINTRUST_PROJECT_ID).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--project-id",
|
||||
default=None,
|
||||
help="Target project by id (default: OJ_BRAINTRUST_PROJECT_ID or research).",
|
||||
)
|
||||
p.add_argument("--run-label", default=None)
|
||||
p.add_argument("--gen-model", default=None, help="Specific gen model id override.")
|
||||
p.add_argument("--description", default="")
|
||||
p.add_argument("specs", nargs="+", help="path.jsonl[=dataset_name]")
|
||||
args = p.parse_args()
|
||||
for spec in args.specs:
|
||||
path, _, name = spec.partition("=")
|
||||
nm, n, url = upload_dataset(
|
||||
path,
|
||||
name or None,
|
||||
project_id=args.project_id,
|
||||
project=args.project,
|
||||
run_label=args.run_label,
|
||||
gen_model=args.gen_model,
|
||||
description=args.description,
|
||||
)
|
||||
print(f"[braintrust] {nm}: {n} rows -> {url}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -445,8 +445,8 @@ class LocalCloudAgent(BaseAgent):
|
||||
tools: Optional[list] = None,
|
||||
tool_choice: Optional[dict] = None,
|
||||
output_config: Optional[dict] = None,
|
||||
timeout: float = 600.0,
|
||||
max_retries: int = 12,
|
||||
timeout: float = 60.0,
|
||||
max_retries: int = 2,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int, int]:
|
||||
"""Single Anthropic call. Returns (text, p_tok, c_tok, n_web_searches).
|
||||
@@ -536,7 +536,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
response_format: Optional[dict] = None,
|
||||
tools: Optional[list] = None,
|
||||
tool_choice: Optional[Any] = None,
|
||||
timeout: float = 600.0,
|
||||
timeout: float = 60.0,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int]:
|
||||
"""Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured;
|
||||
@@ -604,7 +604,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
timeout: float = 600.0,
|
||||
timeout: float = 60.0,
|
||||
trace_role: str = "cloud",
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, int, int]:
|
||||
|
||||
@@ -11,6 +11,7 @@ from __future__ import annotations
|
||||
|
||||
# USD per million tokens, (input, output). Local models = 0.
|
||||
PRICES: dict[str, tuple[float, float]] = {
|
||||
"claude-opus-4-8": (5.00, 25.0),
|
||||
"claude-opus-4-7": (5.00, 25.0),
|
||||
"claude-sonnet-4-6": (3.00, 15.0),
|
||||
"claude-haiku-4-5": (1.00, 5.00),
|
||||
@@ -33,6 +34,17 @@ PRICES: dict[str, tuple[float, float]] = {
|
||||
"qwen/qwen-2.5-coder-32b-instruct": (0.08, 0.18),
|
||||
"qwen/qwen3-32b": (0.10, 0.30),
|
||||
"meta-llama/llama-3.3-70b-instruct": (0.13, 0.39),
|
||||
# OpenRouter slugs for the orchestrator's local-OSS class when routed via
|
||||
# OpenRouter instead of self-hosted vLLM. OpenRouter list price, checked
|
||||
# 2026-07-11. These replace earlier active-param-scaled guesses that were
|
||||
# 2-5x low on output; the cost-aware reward reads these directly, so the
|
||||
# guesses were systematically undercharging the mid/large Qwen experts.
|
||||
"qwen/qwen3.5-9b": (0.10, 0.15),
|
||||
# qwen3.6-27b is not listed on OpenRouter; qwen3.5-27b is the nearest real
|
||||
# quote and is what we charge for it.
|
||||
"qwen/qwen3.6-27b": (0.195, 1.56),
|
||||
"qwen/qwen3.5-122b-a10b": (0.26, 2.08),
|
||||
"qwen/qwen3.5-397b-a17b": (0.385, 2.45),
|
||||
}
|
||||
|
||||
# Models whose API rejects an explicit `temperature` param — callers should
|
||||
@@ -41,6 +53,10 @@ NO_TEMP_PREFIXES: tuple[str, ...] = (
|
||||
"claude-opus-4-7",
|
||||
"claude-sonnet-4-7",
|
||||
"claude-haiku-4-7",
|
||||
# 4-8 family also rejects `temperature` ("deprecated for this model").
|
||||
"claude-opus-4-8",
|
||||
"claude-sonnet-4-8",
|
||||
"claude-haiku-4-8",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,979 @@
|
||||
"""Faithful ToolOrchestra "unified tool calling" registry (arXiv:2511.21689 §3.1).
|
||||
|
||||
The paper exposes **every tool AND every model through a single flat tool
|
||||
interface** — each is its own named function with a description and a typed
|
||||
parameter schema, and for each training instance a *random subset* of tools is
|
||||
sampled with *randomized pricing* (§3.3, "General tool configuration"). This is
|
||||
unlike the eval-port shortcut in ``toolorchestra.py``, which collapses the whole
|
||||
catalog into three meta-tools (``search``/``enhance_reasoning``/``answer``) with
|
||||
a ``model`` slot. This module restores the faithful design.
|
||||
|
||||
Each :class:`ExpertTool` knows:
|
||||
|
||||
* the orchestrator-visible ``name`` / ``description`` / param schema (what goes
|
||||
into the tools JSON the policy conditions on), and
|
||||
* the concrete backend (``backend_type`` + ``model`` + ``base_url``) so a caller
|
||||
can turn it into the worker dict that ``toolorchestra._call_worker`` dispatches.
|
||||
|
||||
Everything here is pure data + deterministic transforms (no network, no model
|
||||
calls), so the spec building, sampling, and pricing logic is offline-testable.
|
||||
Dispatch stays in ``toolorchestra.py`` (via :func:`to_worker_dict`) to avoid a
|
||||
circular import.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from openjarvis.agents.hybrid._prices import PRICES
|
||||
|
||||
# Kinds of tool in the unified interface.
|
||||
KIND_MODEL = "model" # an LLM exposed as a tool (the paper's "models as tools")
|
||||
KIND_WEB_SEARCH = "web_search"
|
||||
KIND_LOCAL_SEARCH = "local_search"
|
||||
KIND_CODE = "code_interpreter"
|
||||
KIND_TOOL = "tool" # a bridged real OpenJarvis tool (custom param schema)
|
||||
|
||||
VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_SEARCH, KIND_CODE, KIND_TOOL)
|
||||
|
||||
# Flat-catalog category label, surfaced in each tool's spec so the orchestrator
|
||||
# can tell tool types apart without us imposing any hierarchy (the menu stays
|
||||
# flat — this is just a tag).
|
||||
CATEGORY_BASIC = "basic_tool"
|
||||
# Two-model-class taxonomy for the orchestrator catalog — the only model tiers.
|
||||
# (The legacy default_catalog still uses CATEGORY_GENERALIST/SPECIALIZED below.)
|
||||
CATEGORY_CLOUD_FRONTIER = "cloud_frontier"
|
||||
CATEGORY_LOCAL_OSS = "local_open_source"
|
||||
# Legacy tiers kept only for default_catalog (build_unified_sft.py); not used by
|
||||
# orchestrator_catalog. Superseded by the two-class taxonomy above.
|
||||
CATEGORY_GENERALIST = "generalist_model"
|
||||
CATEGORY_SPECIALIZED = "specialized_model"
|
||||
|
||||
# Backend dispatch types understood by ``toolorchestra._call_worker`` (plus the
|
||||
# ``openjarvis-tool`` bridge, dispatched in ``unified.make_dispatch`` via the
|
||||
# OpenJarvis ToolExecutor rather than ``_call_worker``).
|
||||
VALID_BACKENDS = (
|
||||
"vllm",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"gemini",
|
||||
"openrouter",
|
||||
"anthropic-web-search",
|
||||
"tavily-search",
|
||||
"modal-python",
|
||||
"openjarvis-tool",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExpertTool:
|
||||
"""One entry in the unified tool catalog.
|
||||
|
||||
``price_in`` / ``price_out`` are USD per 1M tokens (0.0 for local / non-LLM
|
||||
tools). ``latency_s`` is a rough average used only to populate the
|
||||
description's cost/latency line — the orchestrator was trained to read that
|
||||
table, so we surface it verbatim in the spec.
|
||||
"""
|
||||
|
||||
name: str
|
||||
kind: str
|
||||
backend_type: str
|
||||
summary: str
|
||||
model: Optional[str] = None
|
||||
base_url: Optional[str] = None
|
||||
price_in: float = 0.0
|
||||
price_out: float = 0.0
|
||||
latency_s: float = 5.0
|
||||
category: str = "" # cloud_frontier | local_open_source | basic_tool
|
||||
# Optional explicit JSON-schema for the tool's arguments. Set for bridged
|
||||
# real OpenJarvis tools (``openjarvis-tool`` backend) whose params don't fit
|
||||
# the fixed kind-based schemas; takes precedence over the kind default.
|
||||
param_schema: Optional[dict] = None
|
||||
# When True, ``description()`` omits the price/latency line — used by the
|
||||
# anonymized catalog so the orchestrator can't route on cost/identity.
|
||||
hide_cost: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.kind not in VALID_KINDS:
|
||||
raise ValueError(f"{self.name}: invalid kind {self.kind!r}")
|
||||
if self.backend_type not in VALID_BACKENDS:
|
||||
raise ValueError(f"{self.name}: invalid backend {self.backend_type!r}")
|
||||
if self.kind == KIND_MODEL and not self.model:
|
||||
raise ValueError(f"{self.name}: model-kind tool needs a concrete model")
|
||||
|
||||
# ---- orchestrator-visible spec -------------------------------------
|
||||
|
||||
def _param_schema(self) -> Dict[str, object]:
|
||||
"""JSON-schema for the tool's arguments (one typed param per kind).
|
||||
|
||||
An explicit ``param_schema`` (set by :func:`openjarvis_tool` for bridged
|
||||
real tools) overrides the kind-based default.
|
||||
"""
|
||||
if self.param_schema is not None:
|
||||
return self.param_schema
|
||||
if self.kind == KIND_WEB_SEARCH or self.kind == KIND_LOCAL_SEARCH:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Search query string.",
|
||||
}
|
||||
},
|
||||
"required": ["query"],
|
||||
}
|
||||
if self.kind == KIND_CODE:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python code to execute. Print results.",
|
||||
}
|
||||
},
|
||||
"required": ["code"],
|
||||
}
|
||||
# model tool
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"input": {
|
||||
"type": "string",
|
||||
"description": "The sub-question or instruction for this model.",
|
||||
}
|
||||
},
|
||||
"required": ["input"],
|
||||
}
|
||||
|
||||
def description(self) -> str:
|
||||
"""Full description incl. the price/latency line (paper bakes this in)."""
|
||||
if self.hide_cost:
|
||||
return self.summary
|
||||
if self.kind == KIND_MODEL:
|
||||
cost_line = (
|
||||
f" Pricing: ${self.price_in:.2f}/1M input, "
|
||||
f"${self.price_out:.2f}/1M output; avg latency ~{self.latency_s:.0f}s."
|
||||
)
|
||||
else:
|
||||
cost_line = f" Avg latency ~{self.latency_s:.0f}s."
|
||||
return self.summary.rstrip(".") + "." + cost_line
|
||||
|
||||
def to_spec(self) -> Dict[str, object]:
|
||||
"""OpenAI-style tool spec the orchestrator conditions on.
|
||||
|
||||
Flat list, but each function carries a ``category`` tag so the policy can
|
||||
distinguish generalist vs specialized models vs basic tools.
|
||||
"""
|
||||
fn: Dict[str, object] = {
|
||||
"name": self.name,
|
||||
"description": self.description(),
|
||||
"parameters": self._param_schema(),
|
||||
}
|
||||
if self.category:
|
||||
fn["category"] = self.category
|
||||
return {"type": "function", "function": fn}
|
||||
|
||||
|
||||
def _price(model: str) -> tuple[float, float]:
|
||||
return PRICES.get(model, (0.0, 0.0))
|
||||
|
||||
|
||||
def _tool_name(model: str) -> str:
|
||||
"""Tool-safe function name derived from a model id (``qwen3-8b`` -> ``qwen3_8b``).
|
||||
|
||||
Strips any provider prefix and replaces non-alphanumerics with underscores so
|
||||
the catalog exposes one named tool per concrete model.
|
||||
"""
|
||||
base = model.split("/")[-1].lower()
|
||||
safe = re.sub(r"[^a-z0-9]+", "_", base).strip("_")
|
||||
return safe or "local_model"
|
||||
|
||||
|
||||
_ANON_ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||
|
||||
# Rough size/scale hint per model — a capability signal WITHOUT the brand name,
|
||||
# so the orchestrator can route big-vs-small on task difficulty but can't route on
|
||||
# a proprietary name/tier prior. Closed models have no public param count, so they
|
||||
# get a scale tier; open models get their param count (else parsed from the name).
|
||||
_MODEL_SIZE = {
|
||||
# The real orchestrator_catalog: 2 cloud frontier + 4 local OSS Qwen.
|
||||
"gpt-5.5": "frontier-scale",
|
||||
"claude-opus-4-8": "frontier-scale",
|
||||
"Qwen/Qwen3.5-9B": "~9B params",
|
||||
"Qwen/Qwen3.6-27B-FP8": "~27B params",
|
||||
"Qwen/Qwen3.5-122B-A10B-FP8": "~122B total, ~10B active (MoE)",
|
||||
"Qwen/Qwen3.5-397B-A17B-FP8": "~397B total, ~17B active (MoE)",
|
||||
# OpenRouter slugs (the tool.model id when these route through OpenRouter).
|
||||
"qwen/qwen3.5-122b-a10b": "~122B total, ~10B active (MoE)",
|
||||
"qwen/qwen3.5-397b-a17b": "~397B total, ~17B active (MoE)",
|
||||
}
|
||||
|
||||
|
||||
def _size_hint(model: str) -> str:
|
||||
"""Anonymized scale descriptor for an expert (no brand). Falls back to a
|
||||
param count parsed from the model id (``...-9b`` -> ``~9B params``)."""
|
||||
if model in _MODEL_SIZE:
|
||||
return _MODEL_SIZE[model]
|
||||
m = re.search(r"(\d+(?:\.\d+)?)\s*b\b", model.lower())
|
||||
return f"~{m.group(1)}B params" if m else "unspecified scale"
|
||||
|
||||
|
||||
def _class_hint(tool) -> str:
|
||||
"""Which of the two model classes this expert belongs to (anonymized)."""
|
||||
cat = getattr(tool, "category", "")
|
||||
if cat == CATEGORY_CLOUD_FRONTIER:
|
||||
return "cloud frontier"
|
||||
if cat == CATEGORY_LOCAL_OSS:
|
||||
return "local open-source"
|
||||
return ""
|
||||
|
||||
|
||||
def _cost_tier_hint(tool) -> str:
|
||||
"""Coarse cost tier (not exact pricing — that's the bias we anonymize away).
|
||||
Shown so the policy learns the strong-but-expensive vs cheap tradeoff instead
|
||||
of always delegating to the biggest model. Driven by class: cloud frontier is
|
||||
expensive, self-hosted open-source is cheap."""
|
||||
cat = getattr(tool, "category", "")
|
||||
if cat == CATEGORY_CLOUD_FRONTIER:
|
||||
return "expensive"
|
||||
if cat == CATEGORY_LOCAL_OSS:
|
||||
return "cheap"
|
||||
if getattr(tool, "backend_type", "") == "vllm" or (tool.price_out or 0) <= 0:
|
||||
return "cheap"
|
||||
po = tool.price_out
|
||||
return "cheap" if po < 5 else "moderate" if po < 20 else "expensive"
|
||||
|
||||
|
||||
def anonymize_tools(tools, rng):
|
||||
"""Strip model identity for unbiased routing data.
|
||||
|
||||
Each MODEL expert is replaced with an opaque random label
|
||||
(``expert_<4 rand>``), a uniform description, a uniform category and no
|
||||
price/latency line — so the orchestrator cannot route on a model's name,
|
||||
position, cost, or tier (all of which we found dominate the choice). The full
|
||||
list is shuffled to kill position bias. Basic tools (calculator, web_search,
|
||||
…) keep their real names — the policy must still know what they do.
|
||||
|
||||
Returns ``(anon_tools, anon_to_real)`` where ``anon_to_real`` maps each opaque
|
||||
label back to the real tool name (for offline analysis only; never shown to
|
||||
the model). ``.model`` is preserved on each tool so dispatch still reaches the
|
||||
right backend. Pass a fresh ``rng`` per rollout so labels don't stabilise.
|
||||
"""
|
||||
from dataclasses import replace
|
||||
|
||||
anon_to_real: Dict[str, str] = {}
|
||||
experts: List[ExpertTool] = []
|
||||
basics: List[ExpertTool] = []
|
||||
for t in tools:
|
||||
if t.kind == KIND_MODEL:
|
||||
tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
|
||||
while tag in anon_to_real:
|
||||
tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
|
||||
anon_to_real[tag] = t.name
|
||||
bits = [
|
||||
b
|
||||
for b in (_class_hint(t), _size_hint(t.model), _cost_tier_hint(t))
|
||||
if b
|
||||
]
|
||||
experts.append(
|
||||
replace(
|
||||
t,
|
||||
name=tag,
|
||||
summary="Another model — "
|
||||
+ ", ".join(bits)
|
||||
+ ". Send it a sub-question.",
|
||||
category="model",
|
||||
hide_cost=True,
|
||||
)
|
||||
)
|
||||
else:
|
||||
basics.append(t)
|
||||
# Shuffle WITHIN the experts to kill per-expert position bias, but keep all
|
||||
# experts as a block on top and the basic tools underneath — a clean split
|
||||
# (experts first) rather than experts and utilities interleaved.
|
||||
rng.shuffle(experts)
|
||||
out = experts + basics
|
||||
return out, anon_to_real
|
||||
|
||||
|
||||
# Default catalog: the paper's tool categories, mapped onto the models/tools
|
||||
# OpenJarvis can actually call. One named tool per model (faithful §3.1).
|
||||
def default_catalog(
|
||||
*,
|
||||
local_model: Optional[str] = None,
|
||||
local_endpoint: Optional[str] = None,
|
||||
) -> List[ExpertTool]:
|
||||
"""Return the full unified tool catalog.
|
||||
|
||||
``local_model`` / ``local_endpoint`` wire the on-device vLLM tool when a
|
||||
local backbone is served; omitted → the local model tool is left out.
|
||||
"""
|
||||
cat: List[ExpertTool] = []
|
||||
|
||||
# ---- generalist / frontier models (one named tool per model VERSION) ----
|
||||
for model, summary, lat in [
|
||||
(
|
||||
"gpt-5",
|
||||
"Frontier generalist (GPT-5). Strongest reasoning across domains.",
|
||||
30.0,
|
||||
),
|
||||
(
|
||||
"gpt-5-mini",
|
||||
"Mid-tier generalist (GPT-5-mini). Solid reasoning, much cheaper.",
|
||||
15.0,
|
||||
),
|
||||
(
|
||||
"gpt-4o",
|
||||
"Fast generalist (GPT-4o). Good for simple steps and formatting.",
|
||||
8.0,
|
||||
),
|
||||
(
|
||||
"claude-opus-4-7",
|
||||
"Frontier generalist (Claude Opus 4.7). Strong long-horizon reasoning.",
|
||||
26.0,
|
||||
),
|
||||
(
|
||||
"claude-sonnet-4-6",
|
||||
"Strong generalist (Claude Sonnet 4.6). Balanced cost/capability.",
|
||||
15.0,
|
||||
),
|
||||
(
|
||||
"gemini-2.5-pro",
|
||||
"Frontier generalist (Gemini 2.5 Pro). Strong multimodal reasoning.",
|
||||
20.0,
|
||||
),
|
||||
("gemini-2.5-flash", "Cheap fast generalist (Gemini 2.5 Flash).", 8.0),
|
||||
(
|
||||
"meta-llama/llama-3.3-70b-instruct",
|
||||
"Open generalist (Llama-3.3-70B). Decent general knowledge, low cost.",
|
||||
10.0,
|
||||
),
|
||||
(
|
||||
"qwen/qwen3-32b",
|
||||
"Open generalist (Qwen3-32B). Strong math/science reasoning, low cost.",
|
||||
9.0,
|
||||
),
|
||||
]:
|
||||
ep = (
|
||||
"openai"
|
||||
if model.startswith("gpt")
|
||||
else "anthropic"
|
||||
if model.startswith("claude")
|
||||
else "gemini"
|
||||
if model.startswith("gemini")
|
||||
else "openrouter"
|
||||
)
|
||||
pi, po = _price(model)
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name=_tool_name(model),
|
||||
kind=KIND_MODEL,
|
||||
backend_type=ep,
|
||||
summary=summary,
|
||||
model=model,
|
||||
price_in=pi,
|
||||
price_out=po,
|
||||
latency_s=lat,
|
||||
category=CATEGORY_GENERALIST,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- specialized: code ----
|
||||
coder = "qwen/qwen-2.5-coder-32b-instruct"
|
||||
pi, po = _price(coder)
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name=_tool_name(coder),
|
||||
kind=KIND_MODEL,
|
||||
backend_type="openrouter",
|
||||
summary="Specialized code model (Qwen2.5-Coder-32B). Writes/debugs code.",
|
||||
model=coder,
|
||||
price_in=pi,
|
||||
price_out=po,
|
||||
latency_s=9.0,
|
||||
category=CATEGORY_SPECIALIZED,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- local backbone as a tool (on-device vLLM), if served ----
|
||||
# Named after the actual served model (faithful "one named tool per model"),
|
||||
# not a generic "local_model" — e.g. "qwen3-8b" -> tool "qwen3_8b".
|
||||
if local_model and local_endpoint:
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name=_tool_name(local_model),
|
||||
kind=KIND_MODEL,
|
||||
backend_type="vllm",
|
||||
summary=(
|
||||
f"On-device open model ({local_model}) served locally. Cheap "
|
||||
"and private; good for extraction, formatting, arithmetic on "
|
||||
"given data."
|
||||
),
|
||||
model=local_model,
|
||||
base_url=local_endpoint,
|
||||
price_in=0.0,
|
||||
price_out=0.0,
|
||||
latency_s=2.0,
|
||||
category=CATEGORY_GENERALIST,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- basic tools ----
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name="web_search",
|
||||
kind=KIND_WEB_SEARCH,
|
||||
backend_type="tavily-search",
|
||||
summary="Web search (Tavily). Use for facts that need a live lookup.",
|
||||
model="tavily",
|
||||
latency_s=8.0,
|
||||
category=CATEGORY_BASIC,
|
||||
)
|
||||
)
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name="code_interpreter",
|
||||
kind=KIND_CODE,
|
||||
backend_type="modal-python",
|
||||
summary="Python sandbox. Execute code and return stdout/stderr.",
|
||||
model="modal-python",
|
||||
latency_s=6.0,
|
||||
category=CATEGORY_BASIC,
|
||||
)
|
||||
)
|
||||
|
||||
return cat
|
||||
|
||||
|
||||
def openjarvis_tool(
|
||||
registered_name: str,
|
||||
*,
|
||||
summary: str,
|
||||
params: dict,
|
||||
latency_s: float = 5.0,
|
||||
) -> ExpertTool:
|
||||
"""Build an :class:`ExpertTool` that bridges a real OpenJarvis tool.
|
||||
|
||||
``registered_name`` is the tool's key in ``ToolRegistry`` (e.g. ``calculator``,
|
||||
``shell_exec``). ``params`` is the JSON-schema *properties*-style dict for the
|
||||
tool's arguments; it is surfaced verbatim by :meth:`ExpertTool.to_spec`. The
|
||||
resulting tool dispatches through the OpenJarvis ``ToolExecutor`` (backend
|
||||
``openjarvis-tool``) rather than ``_call_worker``.
|
||||
"""
|
||||
return ExpertTool(
|
||||
name=registered_name,
|
||||
kind=KIND_TOOL,
|
||||
backend_type="openjarvis-tool",
|
||||
summary=summary,
|
||||
model=registered_name,
|
||||
latency_s=latency_s,
|
||||
category=CATEGORY_BASIC,
|
||||
param_schema=params,
|
||||
)
|
||||
|
||||
|
||||
# Real OpenJarvis tools bridged into the orchestrator catalog as basic tools.
|
||||
# Names must match the ``ToolRegistry`` keys (confirmed present: calculator,
|
||||
# shell_exec, file_read, file_write, http_request).
|
||||
def _openjarvis_basic_tools() -> List[ExpertTool]:
|
||||
def obj(properties: dict, required: List[str]) -> dict:
|
||||
return {"type": "object", "properties": properties, "required": required}
|
||||
|
||||
return [
|
||||
openjarvis_tool(
|
||||
"calculator",
|
||||
summary="Evaluate an arithmetic / math expression and return the result.",
|
||||
params=obj(
|
||||
{
|
||||
"expression": {
|
||||
"type": "string",
|
||||
"description": "Math expression to evaluate.",
|
||||
}
|
||||
},
|
||||
["expression"],
|
||||
),
|
||||
latency_s=1.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"shell_exec",
|
||||
summary=(
|
||||
"Run a shell command and return its stdout/stderr. Critical "
|
||||
"for terminal / TerminalBench-style tasks."
|
||||
),
|
||||
params=obj(
|
||||
{
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "Shell command to execute.",
|
||||
}
|
||||
},
|
||||
["command"],
|
||||
),
|
||||
latency_s=4.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"file_read",
|
||||
summary="Read the contents of a file at the given path.",
|
||||
params=obj(
|
||||
{
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Path of the file to read.",
|
||||
}
|
||||
},
|
||||
["path"],
|
||||
),
|
||||
latency_s=1.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"file_write",
|
||||
summary="Write content to a file at the given path.",
|
||||
params=obj(
|
||||
{
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Path of the file to write.",
|
||||
},
|
||||
"content": {"type": "string", "description": "Content to write."},
|
||||
},
|
||||
["path", "content"],
|
||||
),
|
||||
latency_s=1.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"http_request",
|
||||
summary="Make an HTTP request to a URL and return the response body.",
|
||||
params={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": {"type": "string", "description": "Request URL."},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"description": "HTTP method (GET, POST, ...). Default GET.",
|
||||
},
|
||||
},
|
||||
"required": ["url"],
|
||||
},
|
||||
latency_s=4.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"think",
|
||||
summary=(
|
||||
"Record a private reasoning step (scratchpad). No external "
|
||||
"effect; use to plan before acting on hard reasoning tasks."
|
||||
),
|
||||
params=obj(
|
||||
{
|
||||
"thought": {
|
||||
"type": "string",
|
||||
"description": "Your reasoning or thought process.",
|
||||
}
|
||||
},
|
||||
["thought"],
|
||||
),
|
||||
latency_s=0.5,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"apply_patch",
|
||||
summary=(
|
||||
"Apply a unified-diff patch to a file. Use to edit code for "
|
||||
"terminal / SWE-style tasks."
|
||||
),
|
||||
params=obj(
|
||||
{
|
||||
"patch": {
|
||||
"type": "string",
|
||||
"description": "The unified diff patch text to apply.",
|
||||
},
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Target file path (auto-detected from the "
|
||||
"patch header if omitted).",
|
||||
},
|
||||
},
|
||||
["patch"],
|
||||
),
|
||||
latency_s=2.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"pdf_extract",
|
||||
summary=(
|
||||
"Extract text from a PDF file. Use for GAIA-style tasks with "
|
||||
"PDF attachments."
|
||||
),
|
||||
params=obj(
|
||||
{
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the PDF file.",
|
||||
},
|
||||
"pages": {
|
||||
"type": "string",
|
||||
"description": "Page range, e.g. '1-5' or '1,3,5'. "
|
||||
"Omit for all pages.",
|
||||
},
|
||||
},
|
||||
["file_path"],
|
||||
),
|
||||
latency_s=3.0,
|
||||
),
|
||||
openjarvis_tool(
|
||||
"db_query",
|
||||
summary=(
|
||||
"Run a SQL query against a SQLite/Postgres database and return "
|
||||
"rows. Read-only by default."
|
||||
),
|
||||
params=obj(
|
||||
{
|
||||
"query": {"type": "string", "description": "SQL query to execute."},
|
||||
"db_path": {
|
||||
"type": "string",
|
||||
"description": "Path to a SQLite DB file. Defaults to "
|
||||
"in-memory.",
|
||||
},
|
||||
"read_only": {
|
||||
"type": "boolean",
|
||||
"description": "Restrict to SELECT/EXPLAIN/PRAGMA. "
|
||||
"Default: true.",
|
||||
},
|
||||
},
|
||||
["query"],
|
||||
),
|
||||
latency_s=3.0,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# Local-Cloud Hybrid orchestrator catalog. The menu is two model classes — cloud
|
||||
# frontier models and open-source models — plus the basic tools (web search, code
|
||||
# interpreter, and the bridged real OpenJarvis tools). The orchestrator model
|
||||
# itself is NOT in the catalog.
|
||||
#
|
||||
# Routing default is OpenRouter for every model tool, so the catalog works with
|
||||
# no self-hosted servers. A model is routed to local vLLM instead when its id is
|
||||
# in ``local_endpoints`` or ``model_backends`` overrides it. The
|
||||
# orchestrator-visible tool *name* is always derived from the canonical id, so
|
||||
# routing can change (vLLM <-> OpenRouter <-> native API) without shifting the
|
||||
# menu the policy was trained on.
|
||||
_LOCAL_OSS_MODELS = (
|
||||
# (canonical_id, openrouter_slug)
|
||||
("Qwen/Qwen3.5-9B", "qwen/qwen3.5-9b"),
|
||||
("Qwen/Qwen3.6-27B-FP8", "qwen/qwen3.6-27b"),
|
||||
("Qwen/Qwen3.5-122B-A10B-FP8", "qwen/qwen3.5-122b-a10b"),
|
||||
("Qwen/Qwen3.5-397B-A17B-FP8", "qwen/qwen3.5-397b-a17b"),
|
||||
)
|
||||
|
||||
_CLOUD_FRONTIER_MODELS = (
|
||||
# (canonical, native_backend, openrouter_slug, summary, latency_s)
|
||||
# Neutral, uniform summaries (no capability ranking) so the orchestrator
|
||||
# doesn't just pick whichever model is labelled "strongest" — routing should
|
||||
# be learned from the reward, not hand-labelled here.
|
||||
("gpt-5.5", "openai", "openai/gpt-5.5", "Expert model (GPT-5.5).", 30.0),
|
||||
(
|
||||
"claude-opus-4-8",
|
||||
"anthropic",
|
||||
"anthropic/claude-opus-4.8",
|
||||
"Expert model (Claude Opus 4.8).",
|
||||
26.0,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _model_tool(
|
||||
canonical: str,
|
||||
*,
|
||||
native_backend: str,
|
||||
or_slug: str,
|
||||
summary: str,
|
||||
lat: float,
|
||||
category: str,
|
||||
local_endpoints: Dict[str, str],
|
||||
model_backends: Dict[str, str],
|
||||
openrouter_slugs: Dict[str, str],
|
||||
) -> ExpertTool:
|
||||
"""Build one model tool, resolving its backend.
|
||||
|
||||
Backend precedence: explicit ``model_backends[canonical]`` > vLLM if the model
|
||||
has a ``local_endpoints`` entry > the model's NATIVE provider (openai /
|
||||
anthropic / gemini) when it has one > OpenRouter. Frontier models thus hit
|
||||
their first-party API by default (OpenRouter's gpt-5.5 was returning 400s);
|
||||
OSS Qwen experts (native_backend="vllm") fall through to OpenRouter.
|
||||
"""
|
||||
backend = (
|
||||
model_backends.get(canonical)
|
||||
or ("vllm" if canonical in local_endpoints else None)
|
||||
or (
|
||||
native_backend
|
||||
if native_backend in ("openai", "anthropic", "gemini")
|
||||
else "openrouter"
|
||||
)
|
||||
)
|
||||
name = _tool_name(canonical)
|
||||
if backend == "vllm":
|
||||
# Self-hosted: free per the cost model.
|
||||
return ExpertTool(
|
||||
name=name,
|
||||
kind=KIND_MODEL,
|
||||
backend_type="vllm",
|
||||
summary=summary,
|
||||
model=canonical,
|
||||
base_url=local_endpoints.get(canonical),
|
||||
price_in=0.0,
|
||||
price_out=0.0,
|
||||
latency_s=lat,
|
||||
category=category,
|
||||
)
|
||||
if backend == "openrouter":
|
||||
slug = openrouter_slugs.get(canonical, or_slug)
|
||||
pi, po = _price(slug)
|
||||
if (pi, po) == (0.0, 0.0): # fall back to the canonical id's price
|
||||
pi, po = _price(canonical)
|
||||
return ExpertTool(
|
||||
name=name,
|
||||
kind=KIND_MODEL,
|
||||
backend_type="openrouter",
|
||||
summary=summary,
|
||||
model=slug,
|
||||
base_url=None,
|
||||
price_in=pi,
|
||||
price_out=po,
|
||||
latency_s=lat,
|
||||
category=category,
|
||||
)
|
||||
# native provider API (openai / anthropic / gemini)
|
||||
pi, po = _price(canonical)
|
||||
return ExpertTool(
|
||||
name=name,
|
||||
kind=KIND_MODEL,
|
||||
backend_type=backend,
|
||||
summary=summary,
|
||||
model=canonical,
|
||||
price_in=pi,
|
||||
price_out=po,
|
||||
latency_s=lat,
|
||||
category=category,
|
||||
)
|
||||
|
||||
|
||||
def orchestrator_catalog(
|
||||
*,
|
||||
local_endpoints: Optional[Dict[str, str]] = None,
|
||||
model_backends: Optional[Dict[str, str]] = None,
|
||||
openrouter_slugs: Optional[Dict[str, str]] = None,
|
||||
include_tools: bool = True,
|
||||
) -> List[ExpertTool]:
|
||||
"""Return the orchestrator's tool catalog: two model classes + basic tools.
|
||||
|
||||
Routing for the model tools defaults to **OpenRouter** (so the catalog works
|
||||
with no self-hosted servers). Overrides, in precedence order:
|
||||
|
||||
* ``model_backends`` maps a canonical model id -> ``"vllm" | "openrouter" |
|
||||
"openai" | "anthropic" | "gemini"`` to force that model's backend.
|
||||
* ``local_endpoints`` maps a canonical id (e.g. ``"Qwen/Qwen3.5-9B"``) to a
|
||||
vLLM ``base_url``; a model present here routes to vLLM (free) unless
|
||||
``model_backends`` says otherwise.
|
||||
* ``openrouter_slugs`` overrides the per-model OpenRouter slug used when a
|
||||
model routes through OpenRouter.
|
||||
|
||||
``include_tools`` (default True) appends the basic tools — web search, code
|
||||
interpreter, and the bridged real OpenJarvis tools (calculator, shell_exec,
|
||||
file_read, file_write, http_request).
|
||||
"""
|
||||
local_endpoints = local_endpoints or {}
|
||||
model_backends = model_backends or {}
|
||||
openrouter_slugs = openrouter_slugs or {}
|
||||
cat: List[ExpertTool] = []
|
||||
|
||||
# Env-gated expert exclusion: OJ_EXCLUDE_EXPERTS is a comma-separated list of
|
||||
# case-insensitive substrings matched against a model's canonical id. Any
|
||||
# match is skipped from the catalog. Used to temporarily drop unreliable
|
||||
# experts (e.g. OpenRouter giants during a provider outage) without editing
|
||||
# the registry — unset the var to restore them.
|
||||
_excl = {
|
||||
s.strip().lower()
|
||||
for s in os.environ.get("OJ_EXCLUDE_EXPERTS", "").split(",")
|
||||
if s.strip()
|
||||
}
|
||||
|
||||
def _excluded(canonical: str) -> bool:
|
||||
c = canonical.lower()
|
||||
return any(x in c for x in _excl)
|
||||
|
||||
# ---- cloud frontier models ----
|
||||
for canonical, native_backend, or_slug, summary, lat in _CLOUD_FRONTIER_MODELS:
|
||||
if _excluded(canonical):
|
||||
continue
|
||||
cat.append(
|
||||
_model_tool(
|
||||
canonical,
|
||||
native_backend=native_backend,
|
||||
or_slug=or_slug,
|
||||
summary=summary,
|
||||
lat=lat,
|
||||
category=CATEGORY_CLOUD_FRONTIER,
|
||||
local_endpoints=local_endpoints,
|
||||
model_backends=model_backends,
|
||||
openrouter_slugs=openrouter_slugs,
|
||||
)
|
||||
)
|
||||
|
||||
# ---- open-source models (OpenRouter by default; vLLM when an endpoint or
|
||||
# a model_backends override is supplied) ----
|
||||
for canonical, or_slug in _LOCAL_OSS_MODELS:
|
||||
if _excluded(canonical):
|
||||
continue
|
||||
cat.append(
|
||||
_model_tool(
|
||||
canonical,
|
||||
native_backend="vllm",
|
||||
or_slug=or_slug,
|
||||
summary=f"Expert model ({canonical}).",
|
||||
lat=4.0,
|
||||
category=CATEGORY_LOCAL_OSS,
|
||||
local_endpoints=local_endpoints,
|
||||
model_backends=model_backends,
|
||||
openrouter_slugs=openrouter_slugs,
|
||||
)
|
||||
)
|
||||
|
||||
if include_tools:
|
||||
# ---- basic tools ----
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name="web_search",
|
||||
kind=KIND_WEB_SEARCH,
|
||||
backend_type="tavily-search",
|
||||
summary="Web search (Tavily). Use for facts that need a live lookup.",
|
||||
model="tavily",
|
||||
latency_s=8.0,
|
||||
category=CATEGORY_BASIC,
|
||||
)
|
||||
)
|
||||
cat.append(
|
||||
ExpertTool(
|
||||
name="code_interpreter",
|
||||
kind=KIND_CODE,
|
||||
backend_type="modal-python",
|
||||
summary="Python sandbox. Execute code and return stdout/stderr.",
|
||||
model="modal-python",
|
||||
latency_s=6.0,
|
||||
category=CATEGORY_BASIC,
|
||||
)
|
||||
)
|
||||
cat.extend(_openjarvis_basic_tools())
|
||||
|
||||
return cat
|
||||
|
||||
|
||||
def build_tool_specs(tools: List[ExpertTool]) -> List[Dict[str, object]]:
|
||||
"""Turn a tool list into the OpenAI-style tools JSON the policy sees."""
|
||||
return [t.to_spec() for t in tools]
|
||||
|
||||
|
||||
def tools_by_name(tools: List[ExpertTool]) -> Dict[str, ExpertTool]:
|
||||
return {t.name: t for t in tools}
|
||||
|
||||
|
||||
def sample_tool_config(
|
||||
catalog: List[ExpertTool],
|
||||
*,
|
||||
rng: random.Random,
|
||||
min_tools: int = 4,
|
||||
max_tools: Optional[int] = None,
|
||||
price_jitter: float = 0.0,
|
||||
) -> List[ExpertTool]:
|
||||
"""Sample a random tool subset with optional price randomization (§3.3).
|
||||
|
||||
Guarantees at least one ``model`` tool and at least one non-model (basic)
|
||||
tool so every instance can both reason and act. ``price_jitter`` (e.g. 0.5)
|
||||
multiplies each model's prices by a per-tool factor drawn uniformly from
|
||||
``[1-jitter, 1+jitter]``, modeling heterogeneous pricing across users.
|
||||
Deterministic given ``rng``.
|
||||
"""
|
||||
if not catalog:
|
||||
raise ValueError("empty catalog")
|
||||
models = [t for t in catalog if t.kind == KIND_MODEL]
|
||||
basics = [t for t in catalog if t.kind != KIND_MODEL]
|
||||
if not models:
|
||||
raise ValueError("catalog has no model tools")
|
||||
|
||||
hi = max_tools if max_tools is not None else len(catalog)
|
||||
hi = min(hi, len(catalog))
|
||||
lo = min(max(min_tools, 2), hi)
|
||||
k = rng.randint(lo, hi)
|
||||
|
||||
# Always include >=1 model; include >=1 basic if any exist.
|
||||
chosen: List[ExpertTool] = [rng.choice(models)]
|
||||
if basics:
|
||||
chosen.append(rng.choice(basics))
|
||||
pool = [t for t in catalog if t not in chosen]
|
||||
rng.shuffle(pool)
|
||||
for t in pool:
|
||||
if len(chosen) >= k:
|
||||
break
|
||||
chosen.append(t)
|
||||
|
||||
# Re-order to catalog order for stable specs.
|
||||
order = {t.name: i for i, t in enumerate(catalog)}
|
||||
chosen.sort(key=lambda t: order[t.name])
|
||||
|
||||
if price_jitter > 0.0:
|
||||
jittered: List[ExpertTool] = []
|
||||
for t in chosen:
|
||||
if t.kind == KIND_MODEL and (t.price_in or t.price_out):
|
||||
f = rng.uniform(1.0 - price_jitter, 1.0 + price_jitter)
|
||||
jittered.append(
|
||||
ExpertTool(
|
||||
name=t.name,
|
||||
kind=t.kind,
|
||||
backend_type=t.backend_type,
|
||||
summary=t.summary,
|
||||
model=t.model,
|
||||
base_url=t.base_url,
|
||||
price_in=round(t.price_in * f, 4),
|
||||
price_out=round(t.price_out * f, 4),
|
||||
latency_s=t.latency_s,
|
||||
)
|
||||
)
|
||||
else:
|
||||
jittered.append(t)
|
||||
return jittered
|
||||
return chosen
|
||||
|
||||
|
||||
def to_worker_dict(tool: ExpertTool) -> Dict[str, object]:
|
||||
"""Convert a tool into the worker dict ``toolorchestra._call_worker`` eats."""
|
||||
d: Dict[str, object] = {
|
||||
"name": tool.name,
|
||||
"type": tool.backend_type,
|
||||
"model": tool.model,
|
||||
}
|
||||
if tool.base_url:
|
||||
d["base_url"] = tool.base_url
|
||||
return d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CATEGORY_BASIC",
|
||||
"CATEGORY_CLOUD_FRONTIER",
|
||||
"CATEGORY_GENERALIST",
|
||||
"CATEGORY_LOCAL_OSS",
|
||||
"CATEGORY_SPECIALIZED",
|
||||
"ExpertTool",
|
||||
"KIND_CODE",
|
||||
"KIND_LOCAL_SEARCH",
|
||||
"KIND_MODEL",
|
||||
"KIND_TOOL",
|
||||
"KIND_WEB_SEARCH",
|
||||
"build_tool_specs",
|
||||
"default_catalog",
|
||||
"openjarvis_tool",
|
||||
"orchestrator_catalog",
|
||||
"sample_tool_config",
|
||||
"to_worker_dict",
|
||||
"tools_by_name",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
"""ToolOrchestraAgent package (split from the former toolorchestra.py module).
|
||||
|
||||
Importing this package registers the agent and re-exports the public surface,
|
||||
so ``from openjarvis.agents.hybrid.toolorchestra import ToolOrchestraAgent``
|
||||
keeps working unchanged. Submodules:
|
||||
|
||||
prompts — system prompts, RL tool specs / arg schema
|
||||
experts — slot -> backend worker mapping (default + paper-match)
|
||||
sandbox — Tavily search + Modal Python sandbox helpers
|
||||
clients — orchestrator vLLM tool-call client
|
||||
parsing — action / tool-call parsing + user-prompt assembly
|
||||
workers — worker pool resolution + dispatch (_call_worker etc.)
|
||||
agent — ToolOrchestraAgent (the registered agent class)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.agents.hybrid.toolorchestra.agent import ToolOrchestraAgent
|
||||
|
||||
__all__ = ["ToolOrchestraAgent"]
|
||||
@@ -0,0 +1,848 @@
|
||||
"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
|
||||
|
||||
Two modes, gated by ``method_cfg.orchestrator_mode``:
|
||||
|
||||
* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
|
||||
orchestrator, dispatching to a numbered worker pool via JSON
|
||||
``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
|
||||
a prompted upper-bound reference point — NOT the paper's setup.
|
||||
|
||||
* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
|
||||
served on a local vLLM is the orchestrator. It emits OpenAI-style
|
||||
``tool_calls`` (or ``<tool_call>{...}</tool_call>`` text blocks when
|
||||
vLLM's tool parser doesn't catch them) for three expert tools —
|
||||
``enhance_reasoning``, ``answer``, ``search`` — exactly as in the
|
||||
upstream ``evaluation/tools.json``. Each tool's ``model`` arg
|
||||
(``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real
|
||||
backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier
|
||||
Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
|
||||
for `*-3`. Search routes to the Anthropic server-side web_search.
|
||||
|
||||
We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
|
||||
code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
|
||||
Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
|
||||
worker types. Energy-wise, "expert" answers are cloud calls.
|
||||
|
||||
Pipeline per task (RL mode):
|
||||
|
||||
1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
|
||||
appropriate tool.` with the three tools declared.
|
||||
2. It emits one ``tool_call`` per turn — ``search`` updates the
|
||||
context, ``enhance_reasoning`` appends code/exec output (we run the
|
||||
tool as a plain LLM call, no sandbox — the model just gets prose
|
||||
back), ``answer`` produces the final answer and the loop stops.
|
||||
3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
|
||||
back to the strongest expert worker.
|
||||
|
||||
Prompted-mode pipeline:
|
||||
|
||||
1. Orchestrator (cloud) reads question + numbered worker pool.
|
||||
2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
|
||||
"input": str}`` or ``{"action": "final_answer", "answer": str}``.
|
||||
3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
|
||||
prompt; fallback to strongest worker on parse failure.
|
||||
|
||||
Workers come from ``cfg["workers"]`` or a sensible default pool (local
|
||||
Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
|
||||
gpt-5-mini).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.hybrid._base import LocalCloudAgent
|
||||
from openjarvis.agents.hybrid._prices import PRICES
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import (
|
||||
_clone_repo,
|
||||
_extract_diff,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.clients import (
|
||||
_call_orchestrator_with_tool_calls,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.experts import (
|
||||
_PAPER_CODER_OPENROUTER,
|
||||
_expert_for,
|
||||
_paper_expert_for,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.parsing import (
|
||||
_build_user_prompt,
|
||||
_extract_final_answer_text,
|
||||
_parse_action,
|
||||
_parse_rl_tool_call,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.prompts import (
|
||||
FORCE_FINAL_PROMPT,
|
||||
ORCHESTRATOR_SYS,
|
||||
RL_ALL_TOOLS,
|
||||
RL_ORCHESTRATOR_SYS,
|
||||
RL_TOOLS_SPEC,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
|
||||
_call_modal_python,
|
||||
_extract_first_python_block,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.workers import (
|
||||
_TOOLORCH_SEARCH_TYPES,
|
||||
_call_worker,
|
||||
_resolve_worker_pool,
|
||||
_swe_call_worker,
|
||||
)
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
|
||||
@AgentRegistry.register("toolorchestra")
|
||||
class ToolOrchestraAgent(LocalCloudAgent):
|
||||
"""Multi-turn dispatcher over a mixed worker pool.
|
||||
|
||||
Two modes (see module docstring): ``method_cfg.orchestrator_mode``
|
||||
is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
|
||||
(paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
|
||||
"""
|
||||
|
||||
agent_id = "toolorchestra"
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
# Validate `method_cfg.worker_pool` early — surfaces config errors
|
||||
# at agent construction rather than on the first task. No-op when
|
||||
# the override is absent.
|
||||
if self._cfg.get("worker_pool") is not None:
|
||||
_resolve_worker_pool(
|
||||
self._cfg,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
self._cloud_endpoint,
|
||||
)
|
||||
# Validate `orchestrator_mode` (typo-checked here, not on first task).
|
||||
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
|
||||
if mode not in ("prompted", "rl"):
|
||||
raise ValueError(
|
||||
f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
|
||||
f"got {mode!r}"
|
||||
)
|
||||
|
||||
def _run_paradigm(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
|
||||
if mode == "rl":
|
||||
return self._run_rl(input, context, **kwargs)
|
||||
return self._run_prompted(input, context, **kwargs)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Legacy prompted-orchestrator path.
|
||||
# ------------------------------------------------------------------
|
||||
def _run_prompted(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
cfg = self._cfg
|
||||
question = input
|
||||
# Resolution order (strict replace, no merge):
|
||||
# 1. `cfg["workers"]` — legacy direct override, used by tests.
|
||||
# 2. `cfg["worker_pool"]` — cell-config override; validated +
|
||||
# $local/$cloud substituted.
|
||||
# 3. `_default_pool(...)` — heterogeneous default.
|
||||
if cfg.get("workers"):
|
||||
workers = cfg["workers"]
|
||||
else:
|
||||
workers = _resolve_worker_pool(
|
||||
cfg,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
self._cloud_endpoint,
|
||||
)
|
||||
if not workers:
|
||||
raise RuntimeError("toolorchestra: empty worker pool")
|
||||
|
||||
max_turns = int(cfg.get("max_turns", 6))
|
||||
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
|
||||
|
||||
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
|
||||
swe_mode = (
|
||||
bool(cfg.get("swe_use_agent_loop"))
|
||||
and bool(task_meta.get("problem_statement"))
|
||||
and bool(task_meta.get("repo"))
|
||||
and bool(task_meta.get("base_commit"))
|
||||
)
|
||||
shared_workdir: Optional[Path] = None
|
||||
if swe_mode:
|
||||
shared_workdir = Path(
|
||||
tempfile.mkdtemp(
|
||||
prefix=f"toolorch-swe-{task_meta.get('task_id', 'x')}-"
|
||||
)
|
||||
)
|
||||
try:
|
||||
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
|
||||
except Exception:
|
||||
shutil.rmtree(shared_workdir, ignore_errors=True)
|
||||
raise
|
||||
self.record_trace_event(
|
||||
{
|
||||
"kind": "toolorchestra_swe_workdir",
|
||||
"workdir": str(shared_workdir),
|
||||
"repo": task_meta["repo"],
|
||||
"base_commit": task_meta["base_commit"],
|
||||
}
|
||||
)
|
||||
|
||||
# try/finally guards ``shared_workdir`` against exceptions raised
|
||||
# anywhere in the turn loop, the worker calls, the fallback, or
|
||||
# the diff-extraction step. Without this, at n=500 SWE-bench an
|
||||
# exception leaves hundreds of MB of cloned repos in tempdir.
|
||||
try:
|
||||
history: List[Dict[str, Any]] = []
|
||||
tokens_local = 0
|
||||
tokens_cloud = 0
|
||||
cost = 0.0
|
||||
n_web_searches_total = 0
|
||||
# tool_calls: bash turns from SWE subloops + web_search uses
|
||||
# from GAIA. Orchestrator dispatch turns are NOT counted (they
|
||||
# produce text only — calling a worker is one tool call's worth
|
||||
# of "delegation" but the actual tool action happens inside).
|
||||
tool_calls = 0
|
||||
final_answer: Optional[str] = None
|
||||
forced_final = False
|
||||
parse_failures = 0
|
||||
|
||||
for turn in range(1, max_turns + 1):
|
||||
sys_prompt = ORCHESTRATOR_SYS
|
||||
if turn == max_turns and final_answer is None:
|
||||
sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
|
||||
forced_final = True
|
||||
|
||||
user = _build_user_prompt(question, workers, history)
|
||||
text, o_in, o_out = self._call_cloud(
|
||||
user=user,
|
||||
system=sys_prompt,
|
||||
max_tokens=orch_max_tokens,
|
||||
temperature=0.0,
|
||||
)
|
||||
tokens_cloud += o_in + o_out
|
||||
cost += self.cost_usd(self._cloud_model, o_in, o_out)
|
||||
|
||||
action = _parse_action(text)
|
||||
history.append(
|
||||
{
|
||||
"role": "orchestrator",
|
||||
"turn": turn,
|
||||
"raw": text,
|
||||
"action": action,
|
||||
}
|
||||
)
|
||||
self.record_trace_event(
|
||||
{
|
||||
"kind": "toolorchestra_action",
|
||||
"turn": turn,
|
||||
"action": action,
|
||||
"raw": text,
|
||||
}
|
||||
)
|
||||
|
||||
if action is None:
|
||||
parse_failures += 1
|
||||
if parse_failures >= 2 or forced_final:
|
||||
final_answer = _extract_final_answer_text(text)
|
||||
break
|
||||
continue
|
||||
|
||||
kind = action.get("action")
|
||||
if kind == "final_answer":
|
||||
final_answer = str(action.get("answer", "")).strip()
|
||||
break
|
||||
if kind == "call_worker":
|
||||
wid = action.get("worker_id")
|
||||
w_input = action.get("input", "")
|
||||
if not isinstance(wid, int) or not (0 <= wid < len(workers)):
|
||||
parse_failures += 1
|
||||
if parse_failures >= 2 or forced_final:
|
||||
final_answer = _extract_final_answer_text(text)
|
||||
break
|
||||
continue
|
||||
worker = workers[wid]
|
||||
if swe_mode and shared_workdir is not None:
|
||||
(
|
||||
w_text,
|
||||
w_in,
|
||||
w_out,
|
||||
is_local,
|
||||
extra_cost,
|
||||
n_searches,
|
||||
bash_turns,
|
||||
) = _swe_call_worker(
|
||||
worker,
|
||||
str(w_input),
|
||||
cfg,
|
||||
task_meta,
|
||||
shared_workdir,
|
||||
turn,
|
||||
)
|
||||
tool_calls += bash_turns
|
||||
else:
|
||||
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
|
||||
_call_worker(worker, str(w_input), cfg)
|
||||
)
|
||||
if is_local:
|
||||
tokens_local += w_in + w_out
|
||||
else:
|
||||
tokens_cloud += w_in + w_out
|
||||
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
|
||||
n_web_searches_total += n_searches
|
||||
tool_calls += n_searches
|
||||
history.append(
|
||||
{
|
||||
"role": "worker",
|
||||
"turn": turn,
|
||||
"worker_id": wid,
|
||||
"worker_name": worker["name"],
|
||||
"worker_model": worker["model"],
|
||||
"output": w_text,
|
||||
"tokens_in": w_in,
|
||||
"tokens_out": w_out,
|
||||
"n_web_searches": n_searches,
|
||||
}
|
||||
)
|
||||
continue
|
||||
# Unknown action kind — treat as parse failure.
|
||||
parse_failures += 1
|
||||
|
||||
if final_answer is None:
|
||||
# Hard fallback: call the strongest non-search worker directly.
|
||||
# "Strongest" = highest output-token price in `_prices.PRICES`,
|
||||
# which tracks model capability tier closely enough for this.
|
||||
# Search workers are excluded — they answer fact-lookup
|
||||
# questions, not synthesis.
|
||||
non_search = [
|
||||
w for w in workers if w.get("type") not in _TOOLORCH_SEARCH_TYPES
|
||||
] or workers
|
||||
worker = max(
|
||||
non_search,
|
||||
key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
|
||||
)
|
||||
if swe_mode and shared_workdir is not None:
|
||||
(ans, w_in, w_out, is_local, extra_cost, _, bash_turns) = (
|
||||
_swe_call_worker(
|
||||
worker,
|
||||
question,
|
||||
cfg,
|
||||
task_meta,
|
||||
shared_workdir,
|
||||
max_turns + 1,
|
||||
)
|
||||
)
|
||||
tool_calls += bash_turns
|
||||
else:
|
||||
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
|
||||
worker, question, cfg
|
||||
)
|
||||
if is_local:
|
||||
tokens_local += w_in + w_out
|
||||
else:
|
||||
tokens_cloud += w_in + w_out
|
||||
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
|
||||
history.append(
|
||||
{
|
||||
"role": "worker",
|
||||
"turn": max_turns + 1,
|
||||
"worker_id": worker["id"],
|
||||
"worker_name": worker["name"],
|
||||
"worker_model": worker["model"],
|
||||
"output": ans,
|
||||
"tokens_in": w_in,
|
||||
"tokens_out": w_out,
|
||||
"fallback": True,
|
||||
}
|
||||
)
|
||||
final_answer = ans
|
||||
|
||||
# In SWE mode, the authoritative output is the working-tree diff —
|
||||
# frame it (the runner extracts it via the scorer's ```diff fence).
|
||||
if swe_mode and shared_workdir is not None:
|
||||
patch = _extract_diff(shared_workdir)
|
||||
if patch.strip():
|
||||
final_answer = (
|
||||
f"{final_answer}\n\n```diff\n{patch}```"
|
||||
if final_answer
|
||||
else f"```diff\n{patch}```"
|
||||
)
|
||||
|
||||
meta = {
|
||||
"tokens_local": tokens_local,
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost,
|
||||
"turns": len([h for h in history if h["role"] == "orchestrator"]),
|
||||
"web_search_uses": n_web_searches_total,
|
||||
"tool_calls": int(tool_calls),
|
||||
"traces": {
|
||||
"history": history,
|
||||
"forced_final": forced_final,
|
||||
"parse_failures": parse_failures,
|
||||
"workers": workers,
|
||||
"n_web_searches": n_web_searches_total,
|
||||
"note": (
|
||||
"inference-only port; the RL-trained Nemotron-Orchestrator-8B "
|
||||
"is NOT in the loop. Results are preliminary."
|
||||
),
|
||||
},
|
||||
}
|
||||
return final_answer, meta
|
||||
finally:
|
||||
if shared_workdir is not None:
|
||||
shutil.rmtree(shared_workdir, ignore_errors=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Paper-faithful Orchestrator-8B path.
|
||||
# ------------------------------------------------------------------
|
||||
def _run_rl(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
cfg = self._cfg
|
||||
question = input
|
||||
|
||||
# Orchestrator endpoint / model (where the RL'd 8B lives).
|
||||
orch_endpoint = str(
|
||||
cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
|
||||
)
|
||||
orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
|
||||
max_turns = int(cfg.get("max_turns", 8))
|
||||
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
|
||||
orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
|
||||
|
||||
# Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
|
||||
# replaces `_expert_for` and `enhance_reasoning` is post-processed
|
||||
# through a Modal Python sandbox. See module docstring + paper-match
|
||||
# doc at `docs/26.5.19/toolorchestra-papermatch.md`.
|
||||
paper_mode = str(cfg.get("pool", "")).lower() == "paper"
|
||||
|
||||
# SWE-bench detection: same gate as the prompted path. Requires
|
||||
# `method_cfg.swe_use_agent_loop = true` AND the task carries the
|
||||
# SWE-bench fields. When active, the `enhance_reasoning` and
|
||||
# `answer` workers route through `run_swe_agent_loop` on a shared
|
||||
# workdir; search workers stay one-shot. At end, the working-tree
|
||||
# diff is appended to final_answer so `_score_swebench` can extract
|
||||
# it via the ```diff fence.
|
||||
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
|
||||
swe_mode = (
|
||||
bool(cfg.get("swe_use_agent_loop"))
|
||||
and bool(task_meta.get("problem_statement"))
|
||||
and bool(task_meta.get("repo"))
|
||||
and bool(task_meta.get("base_commit"))
|
||||
)
|
||||
shared_workdir: Optional[Path] = None
|
||||
if swe_mode:
|
||||
shared_workdir = Path(
|
||||
tempfile.mkdtemp(
|
||||
prefix=f"toolorch-rl-swe-{task_meta.get('task_id', 'x')}-"
|
||||
)
|
||||
)
|
||||
try:
|
||||
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
|
||||
except Exception:
|
||||
shutil.rmtree(shared_workdir, ignore_errors=True)
|
||||
raise
|
||||
self.record_trace_event(
|
||||
{
|
||||
"kind": "toolorchestra_rl_swe_workdir",
|
||||
"workdir": str(shared_workdir),
|
||||
"repo": task_meta["repo"],
|
||||
"base_commit": task_meta["base_commit"],
|
||||
}
|
||||
)
|
||||
|
||||
# ``context_str`` mirrors the upstream's running context — accumulates
|
||||
# search documents and code/exec snippets across turns. We keep this
|
||||
# as a single string for prompt simplicity; the upstream uses
|
||||
# tokenized cutoffs (we cap at ~24k chars instead).
|
||||
context_str = ""
|
||||
doc_list: List[str] = []
|
||||
history: List[Dict[str, Any]] = []
|
||||
tokens_local = 0
|
||||
tokens_cloud = 0
|
||||
cost = 0.0
|
||||
n_web_searches_total = 0
|
||||
tool_calls = 0
|
||||
final_answer: Optional[str] = None
|
||||
parse_failures = 0
|
||||
|
||||
# Single outer try/finally guards `shared_workdir` against any
|
||||
# exception in the orchestrator loop, the post-loop fallback, or
|
||||
# the diff-extraction step. Matches the prompted path's pattern.
|
||||
try:
|
||||
for turn in range(1, max_turns + 1):
|
||||
user = (
|
||||
f"Problem: {question}\n\n{context_str}\n\n"
|
||||
"Choose an appropriate tool."
|
||||
)
|
||||
|
||||
# Orchestrator-8B served on local vLLM. We pass the three NVlabs
|
||||
# tools verbatim. In paper-mode we use the local helper so we
|
||||
# get the SDK-level ``tool_calls`` object back — `_call_vllm`
|
||||
# returns just text and loses the call when vLLM's parser
|
||||
# caught it. Orchestrator-8B emits its routing decision in the
|
||||
# OpenAI-native ``tool_calls`` array with an empty text body,
|
||||
# so the legacy `_call_vllm` path saw nothing and silently fell
|
||||
# through to the answer-1 fallback (parse_failures: 2 on every
|
||||
# non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
|
||||
# modes now use `_call_orchestrator_with_tool_calls` so the
|
||||
# parser can read structured tool calls; the text-tag path in
|
||||
# `_parse_rl_tool_call` is still the fallback when `tool_calls`
|
||||
# is empty.
|
||||
text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
|
||||
orch_model,
|
||||
orch_endpoint,
|
||||
user=user,
|
||||
system=RL_ORCHESTRATOR_SYS,
|
||||
max_tokens=orch_max_tokens,
|
||||
temperature=orch_temp,
|
||||
tools=RL_TOOLS_SPEC,
|
||||
)
|
||||
self.record_trace_event(
|
||||
{
|
||||
"kind": "vllm",
|
||||
"role": "orchestrator",
|
||||
"model": orch_model,
|
||||
"endpoint": orch_endpoint,
|
||||
"system": RL_ORCHESTRATOR_SYS,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": getattr(tc, "id", None),
|
||||
"type": getattr(tc, "type", None),
|
||||
"function": {
|
||||
"name": getattr(
|
||||
getattr(tc, "function", None), "name", None
|
||||
),
|
||||
"arguments": getattr(
|
||||
getattr(tc, "function", None), "arguments", None
|
||||
),
|
||||
},
|
||||
}
|
||||
for tc in (sdk_tool_calls or [])
|
||||
],
|
||||
"tokens_in": o_in,
|
||||
"tokens_out": o_out,
|
||||
}
|
||||
)
|
||||
tokens_local += o_in + o_out
|
||||
|
||||
action = _parse_rl_tool_call(text, sdk_tool_calls)
|
||||
history.append(
|
||||
{
|
||||
"role": "orchestrator",
|
||||
"turn": turn,
|
||||
"raw": text,
|
||||
"action": action,
|
||||
}
|
||||
)
|
||||
self.record_trace_event(
|
||||
{
|
||||
"kind": "toolorchestra_rl_action",
|
||||
"turn": turn,
|
||||
"action": action,
|
||||
"raw": text,
|
||||
}
|
||||
)
|
||||
|
||||
if action is None:
|
||||
parse_failures += 1
|
||||
if parse_failures >= 2:
|
||||
break
|
||||
continue
|
||||
|
||||
name = action["name"]
|
||||
args = action.get("arguments", {})
|
||||
slot = args.get("model", "")
|
||||
|
||||
# Validate against the upstream tool/arg schema.
|
||||
valid = (
|
||||
name in RL_ALL_TOOLS
|
||||
and isinstance(slot, str)
|
||||
and (slot in RL_ALL_TOOLS[name]["model"])
|
||||
)
|
||||
if not valid:
|
||||
parse_failures += 1
|
||||
if parse_failures >= 2:
|
||||
break
|
||||
# Replay with a softer nudge in the context.
|
||||
context_str += (
|
||||
f"\n[Orchestrator emitted invalid tool call "
|
||||
f"name={name!r} slot={slot!r} — try again.]\n"
|
||||
)
|
||||
continue
|
||||
|
||||
# Paper-match (`method_cfg.pool == "paper"`) routes through
|
||||
# the Tavily/OpenRouter/Modal pool instead of the default
|
||||
# Anthropic-web-search-driven mapping. For `search` this
|
||||
# also forces the worker prompt to a raw query string
|
||||
# (Tavily takes a single search string, not a chat-style
|
||||
# framing).
|
||||
if paper_mode:
|
||||
worker = _paper_expert_for(
|
||||
slot,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
self._cloud_endpoint,
|
||||
)
|
||||
# In paper mode, `enhance_reasoning` is always the coder
|
||||
# specialist regardless of the orchestrator's chosen tier.
|
||||
# The coder is then expected to emit a python block which
|
||||
# we exec in Modal (below).
|
||||
if name == "enhance_reasoning":
|
||||
worker = {
|
||||
"name": f"coder:{slot}",
|
||||
"type": "openrouter",
|
||||
"model": _PAPER_CODER_OPENROUTER,
|
||||
}
|
||||
else:
|
||||
worker = _expert_for(
|
||||
slot,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
self._cloud_endpoint,
|
||||
)
|
||||
|
||||
# Dispatch — the orchestrator only conveys a tool/model
|
||||
# choice, NOT a question rewrite; the prompt we send the
|
||||
# expert is the same context the orchestrator saw, framed
|
||||
# appropriately for the tool.
|
||||
if name == "search":
|
||||
if paper_mode:
|
||||
# Tavily takes a query string. Orchestrator-8B often
|
||||
# emits an extra `query` arg (not in the upstream
|
||||
# schema but useful) — prefer it; else fall back to
|
||||
# the raw question.
|
||||
q = args.get("query")
|
||||
w_input = q if isinstance(q, str) and q.strip() else question
|
||||
else:
|
||||
w_input = (
|
||||
f"Search the web to gather information that helps answer:\n"
|
||||
f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
|
||||
)
|
||||
elif name == "enhance_reasoning":
|
||||
if paper_mode:
|
||||
w_input = (
|
||||
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
|
||||
"Write a short Python script that computes intermediate "
|
||||
"results which help answer the problem. Output ONLY the "
|
||||
"code inside one ```python ... ``` fenced block. Print "
|
||||
"any results you derive using `print(...)`. The script "
|
||||
"must run with the Python stdlib only — no extra pip "
|
||||
"installs."
|
||||
)
|
||||
else:
|
||||
w_input = (
|
||||
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
|
||||
"Reason carefully. Outline the key intermediate steps and any "
|
||||
"computations or facts you can derive. Do NOT give a final "
|
||||
"answer — the orchestrator will collect your reasoning and "
|
||||
"call the answer tool next."
|
||||
)
|
||||
else: # name == "answer"
|
||||
w_input = (
|
||||
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
|
||||
"Provide the final answer to the user. Respect any "
|
||||
"answer-format rules in the question (e.g. GAIA's "
|
||||
"FINAL ANSWER: <value> convention)."
|
||||
)
|
||||
|
||||
# SWE mode: route enhance_reasoning / answer workers through
|
||||
# the SWE agent loop on the shared workdir so they can read
|
||||
# files, run tests, and edit the working tree. Search workers
|
||||
# stay one-shot (no agent loop). The `_swe_call_worker`
|
||||
# one-shot fallbacks (openai-typed workers, search) return
|
||||
# bash_turns=0; vllm/anthropic-typed workers run the loop.
|
||||
bash_turns = 0
|
||||
if swe_mode and shared_workdir is not None and name != "search":
|
||||
(
|
||||
w_text,
|
||||
w_in,
|
||||
w_out,
|
||||
is_local,
|
||||
extra_cost,
|
||||
n_searches,
|
||||
bash_turns,
|
||||
) = _swe_call_worker(
|
||||
worker,
|
||||
w_input,
|
||||
cfg,
|
||||
task_meta,
|
||||
shared_workdir,
|
||||
turn,
|
||||
)
|
||||
else:
|
||||
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
|
||||
_call_worker(worker, w_input, cfg)
|
||||
)
|
||||
if is_local:
|
||||
tokens_local += w_in + w_out
|
||||
else:
|
||||
tokens_cloud += w_in + w_out
|
||||
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
|
||||
n_web_searches_total += n_searches
|
||||
# SWE bash turns count as tool calls (each one is a $BASH block
|
||||
# the agent executed). On non-SWE turns fall back to the
|
||||
# original "at least one expert call" accounting.
|
||||
tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches)
|
||||
|
||||
# Paper-match: pipe coder output through a Modal sandbox so
|
||||
# `enhance_reasoning` actually executes the code the coder
|
||||
# wrote. Append the exec output to the worker's text. No-op
|
||||
# when no python block is found.
|
||||
modal_exec_output: Optional[str] = None
|
||||
modal_exec_rc: Optional[int] = None
|
||||
if paper_mode and name == "enhance_reasoning" and not swe_mode:
|
||||
code = _extract_first_python_block(w_text)
|
||||
if code:
|
||||
timeout_s = int(cfg.get("modal_python_timeout_s", 60))
|
||||
modal_exec_output, modal_exec_rc = _call_modal_python(
|
||||
code,
|
||||
timeout_s=timeout_s,
|
||||
)
|
||||
tool_calls += 1
|
||||
w_text = (
|
||||
f"{w_text}\n\n[modal-python stdout/stderr "
|
||||
f"(rc={modal_exec_rc})]\n{modal_exec_output}"
|
||||
)
|
||||
|
||||
history.append(
|
||||
{
|
||||
"role": "worker",
|
||||
"turn": turn,
|
||||
"tool": name,
|
||||
"slot": slot,
|
||||
"worker_model": worker["model"],
|
||||
"worker_type": worker["type"],
|
||||
"output": w_text,
|
||||
"tokens_in": w_in,
|
||||
"tokens_out": w_out,
|
||||
"n_web_searches": n_searches,
|
||||
"bash_turns": bash_turns,
|
||||
"modal_exec_rc": modal_exec_rc,
|
||||
}
|
||||
)
|
||||
|
||||
# Update accumulated context for the next turn.
|
||||
if name == "search":
|
||||
# Treat the search worker's response as a document.
|
||||
doc_list.append(w_text)
|
||||
ctx_docs = "\n\n".join(
|
||||
f"Doc {i + 1}: {d}" for i, d in enumerate(doc_list)
|
||||
)
|
||||
# Crude char-level cap mirrors the upstream's ~24k token cap.
|
||||
context_str = ("Documents:\n" + ctx_docs)[-24000:]
|
||||
elif name == "enhance_reasoning":
|
||||
snippet = f"\n\nReasoning/exec output:\n{w_text}"
|
||||
context_str = (context_str + snippet)[-24000:]
|
||||
else: # answer
|
||||
final_answer = w_text.strip()
|
||||
break
|
||||
|
||||
if final_answer is None:
|
||||
# Hard fallback: ask the frontier worker directly. In SWE
|
||||
# mode route this final call through the agent loop too so
|
||||
# it can still touch the workdir and emit a diff.
|
||||
expert_fn = _paper_expert_for if paper_mode else _expert_for
|
||||
worker = expert_fn(
|
||||
"answer-1",
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
self._cloud_endpoint,
|
||||
)
|
||||
fb_bash_turns = 0
|
||||
if swe_mode and shared_workdir is not None:
|
||||
(ans, w_in, w_out, is_local, extra_cost, _, fb_bash_turns) = (
|
||||
_swe_call_worker(
|
||||
worker,
|
||||
question,
|
||||
cfg,
|
||||
task_meta,
|
||||
shared_workdir,
|
||||
max_turns + 1,
|
||||
)
|
||||
)
|
||||
tool_calls += fb_bash_turns
|
||||
else:
|
||||
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
|
||||
worker, question, cfg
|
||||
)
|
||||
if is_local:
|
||||
tokens_local += w_in + w_out
|
||||
else:
|
||||
tokens_cloud += w_in + w_out
|
||||
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
|
||||
history.append(
|
||||
{
|
||||
"role": "worker",
|
||||
"turn": max_turns + 1,
|
||||
"tool": "answer",
|
||||
"slot": "answer-1",
|
||||
"worker_model": worker["model"],
|
||||
"worker_type": worker["type"],
|
||||
"output": ans,
|
||||
"tokens_in": w_in,
|
||||
"tokens_out": w_out,
|
||||
"bash_turns": fb_bash_turns,
|
||||
"fallback": True,
|
||||
}
|
||||
)
|
||||
final_answer = ans
|
||||
|
||||
# In SWE mode, the authoritative output is the working-tree diff —
|
||||
# frame it so `_score_swebench`'s extract_patch picks it up.
|
||||
if swe_mode and shared_workdir is not None:
|
||||
patch = _extract_diff(shared_workdir)
|
||||
if patch.strip():
|
||||
final_answer = (
|
||||
f"{final_answer}\n\n```diff\n{patch}```"
|
||||
if final_answer
|
||||
else f"```diff\n{patch}```"
|
||||
)
|
||||
|
||||
meta = {
|
||||
"tokens_local": tokens_local,
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost,
|
||||
"turns": len([h for h in history if h["role"] == "orchestrator"]),
|
||||
"web_search_uses": n_web_searches_total,
|
||||
"tool_calls": int(tool_calls),
|
||||
"traces": {
|
||||
"history": history,
|
||||
"parse_failures": parse_failures,
|
||||
"orchestrator_model": orch_model,
|
||||
"orchestrator_endpoint": orch_endpoint,
|
||||
"mode": "rl",
|
||||
"pool": "paper" if paper_mode else "default",
|
||||
"swe_mode": swe_mode,
|
||||
"note": (
|
||||
"RL-trained nvidia/Orchestrator-8B as orchestrator. "
|
||||
"Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto "
|
||||
"our hybrid worker types — see toolorchestra.py docstring."
|
||||
),
|
||||
},
|
||||
}
|
||||
return final_answer, meta
|
||||
finally:
|
||||
if shared_workdir is not None:
|
||||
shutil.rmtree(shared_workdir, ignore_errors=True)
|
||||
|
||||
|
||||
__all__ = ["ToolOrchestraAgent"]
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Orchestrator vLLM tool-call client for ToolOrchestraAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
|
||||
def _call_orchestrator_with_tool_calls(
|
||||
model: str,
|
||||
endpoint: str,
|
||||
*,
|
||||
user: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
tools: List[Dict[str, Any]],
|
||||
timeout: float = 600.0,
|
||||
) -> Tuple[str, int, int, Any]:
|
||||
"""Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls).
|
||||
|
||||
Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level
|
||||
``tool_calls`` object so the RL-mode parser can match against it
|
||||
directly. Otherwise vLLM's tool parser silently swallows the tool call
|
||||
into the SDK field while leaving ``content == ''`` — and the text-tag
|
||||
parser sees nothing, falling through to the answer-1 fallback. (Bug
|
||||
observed 2026-05-19 on the paper-match smoke; same path was buggy on
|
||||
the default pool too, just less reproducibly.)
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
|
||||
messages = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
]
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tools,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
text = message.content or ""
|
||||
tool_calls = getattr(message, "tool_calls", None)
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
return text, p, c, tool_calls
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Slot -> worker expert mapping for ToolOrchestraAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
# Default model used when an `anthropic-web-search` entry omits `model`.
|
||||
_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
|
||||
|
||||
# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
|
||||
# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
|
||||
# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
|
||||
# tiers since we don't have Qwen-Math served.
|
||||
#
|
||||
# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
|
||||
# so the substitution is deferred until we know the cell's resolved local/cloud
|
||||
# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
|
||||
|
||||
|
||||
def _expert_for(
|
||||
slot: str,
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str,
|
||||
cloud_endpoint: str = "anthropic",
|
||||
) -> Dict[str, Any]:
|
||||
"""Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
|
||||
|
||||
Routing policy:
|
||||
- `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off
|
||||
`cloud_endpoint` ("anthropic"/"openai"/"gemini")
|
||||
- `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's
|
||||
cost tier for mid OpenAI calls)
|
||||
- `*-3` (local tier) -> local vLLM (`local_model`)
|
||||
- `answer-math-*` -> same tiers as the numeric suffix
|
||||
- `search-*` -> provider-native web search when the cloud
|
||||
endpoint supports it; otherwise Anthropic
|
||||
"""
|
||||
if slot.startswith("search"):
|
||||
ep = (cloud_endpoint or "anthropic").lower()
|
||||
if ep == "openai":
|
||||
return {
|
||||
"name": f"search:{slot}",
|
||||
"type": "openai-web-search",
|
||||
"model": cloud_model,
|
||||
}
|
||||
if ep == "gemini":
|
||||
return {
|
||||
"name": f"search:{slot}",
|
||||
"type": "gemini-web-search",
|
||||
"model": cloud_model,
|
||||
}
|
||||
return {
|
||||
"name": f"search:{slot}",
|
||||
"type": "anthropic-web-search",
|
||||
"model": _DEFAULT_WEB_SEARCH_MODEL,
|
||||
}
|
||||
if slot.endswith("-1") or slot.endswith("-math-1"):
|
||||
ep = (cloud_endpoint or "anthropic").lower()
|
||||
if ep not in ("anthropic", "openai", "gemini"):
|
||||
ep = "anthropic"
|
||||
return {
|
||||
"name": f"frontier:{slot}",
|
||||
"type": ep,
|
||||
"model": cloud_model,
|
||||
}
|
||||
if slot.endswith("-2") or slot.endswith("-math-2"):
|
||||
return {
|
||||
"name": f"mid:{slot}",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
}
|
||||
# `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
|
||||
# we substitute whatever local model the cell wired up).
|
||||
if local_model and local_endpoint:
|
||||
return {
|
||||
"name": f"local:{slot}",
|
||||
"type": "vllm",
|
||||
"model": local_model,
|
||||
"base_url": local_endpoint,
|
||||
}
|
||||
# Fallback if no local — gpt-5-mini.
|
||||
return {
|
||||
"name": f"mid-fallback:{slot}",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Paper-match expert mapping (2026-05-19).
|
||||
# ============================================================================
|
||||
# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
|
||||
# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
|
||||
# models and routes `search` through Tavily, while `enhance_reasoning` is
|
||||
# expected to produce code that the caller pipes through a Modal sandbox
|
||||
# (handled at dispatch time, not here).
|
||||
#
|
||||
# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
|
||||
# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
|
||||
#
|
||||
# reasoner-1 -> GPT-5 (frontier reasoner)
|
||||
# reasoner-2 -> GPT-5-mini (mid)
|
||||
# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
|
||||
# answer-1 -> GPT-5
|
||||
# answer-2 -> GPT-5-mini
|
||||
# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
|
||||
# answer-4 -> local Qwen
|
||||
# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
|
||||
# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
|
||||
# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
|
||||
# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
|
||||
# search-* -> Tavily search (paper)
|
||||
#
|
||||
# `enhance_reasoning` is dispatched through the coder specialist regardless of
|
||||
# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
|
||||
# routes the same way in all three cases, then optionally extracts a python
|
||||
# code block and execs it in Modal. (We keep the slot-aware routing inside the
|
||||
# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
|
||||
# pins the coder regardless. See `_run_rl_paper` dispatch.)
|
||||
|
||||
_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
|
||||
_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
|
||||
|
||||
|
||||
def _paper_expert_for(
|
||||
slot: str,
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str,
|
||||
cloud_endpoint: str = "openai",
|
||||
) -> Dict[str, Any]:
|
||||
"""Paper-match counterpart of ``_expert_for``.
|
||||
|
||||
Differs from ``_expert_for``:
|
||||
- Search slots go to ``tavily-search`` (not Anthropic web_search).
|
||||
- Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
|
||||
OpenRouter rather than collapsing onto the local vLLM.
|
||||
- Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
|
||||
as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
|
||||
- ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
|
||||
"""
|
||||
if slot.startswith("search"):
|
||||
return {
|
||||
"name": f"tavily:{slot}",
|
||||
"type": "tavily-search",
|
||||
"model": "tavily",
|
||||
}
|
||||
if slot in ("answer-math-1", "answer-math-2"):
|
||||
return {
|
||||
"name": f"math-coder:{slot}",
|
||||
"type": "openrouter",
|
||||
"model": _PAPER_CODER_OPENROUTER,
|
||||
}
|
||||
if slot == "answer-3":
|
||||
return {
|
||||
"name": f"generalist-llama:{slot}",
|
||||
"type": "openrouter",
|
||||
"model": _PAPER_GENERALIST_TIER3_OPENROUTER,
|
||||
}
|
||||
if slot.endswith("-1"):
|
||||
# Tier-1 frontier reasoner / answer — paper uses GPT-5.
|
||||
return {
|
||||
"name": f"frontier:{slot}",
|
||||
"type": "openai",
|
||||
"model": "gpt-5",
|
||||
}
|
||||
if slot.endswith("-2"):
|
||||
return {
|
||||
"name": f"mid:{slot}",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
}
|
||||
# `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
|
||||
# also serves the local Qwen for the rare local-tier slot).
|
||||
if local_model and local_endpoint:
|
||||
return {
|
||||
"name": f"local:{slot}",
|
||||
"type": "vllm",
|
||||
"model": local_model,
|
||||
"base_url": local_endpoint,
|
||||
}
|
||||
return {
|
||||
"name": f"mid-fallback:{slot}",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Action / tool-call parsing + prompt assembly for ToolOrchestraAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# Regex for ``<tool_call>{...}</tool_call>`` blocks emitted by Orchestrator-8B
|
||||
# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a
|
||||
# hermes-style template). Captures the JSON payload.
|
||||
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
|
||||
|
||||
# Fallback for a known failure mode: the SFT'd orchestrator often emits its
|
||||
# delegation as ``\boxed{expert_ab12: <sub-question>}`` (math-data habit bleeding
|
||||
# into routing) instead of a real ``<tool_call>`` JSON block. The name before the
|
||||
# colon is a valid (anonymized) tool label; the rest is the sub-question/args.
|
||||
#
|
||||
# Real outputs use several shapes, all handled here:
|
||||
# \boxed{web_search: <query>} -> name, args
|
||||
# \boxed{web_search query: <query>} -> the word "query" is dropped
|
||||
# \boxed{file_read, path: <path>} -> the ", <key>" hint sets the arg key
|
||||
# We require a ``name <optional , key | query> : <args>`` structure, so a genuine
|
||||
# answer like ``\boxed{10}``, ``\boxed{b,e}`` or ``\boxed{The answer is: 42}``
|
||||
# (first token isn't followed by a bare ``:`` / ``, key:`` / ``query:``) is left
|
||||
# alone. ``group("key")`` is the explicit arg key when the model wrote ``, key:``.
|
||||
_BOXED_DELEGATION_RE = re.compile(
|
||||
r"\\boxed\{\s*([A-Za-z_][\w-]*)\s*"
|
||||
r"(?:,\s*(?P<key>\w+)\s*)?" # optional ", key" hint (e.g. file_read, path:)
|
||||
r"(?:query\s*)?" # optional literal "query" word before the colon
|
||||
r":\s*(.+?)\s*\}\s*$",
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
|
||||
"""Return ``{"name": str, "arguments": dict}`` or None.
|
||||
|
||||
Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
|
||||
back to scraping ``<tool_call>{...}</tool_call>`` tags from the raw
|
||||
content. We take the first tool call only — Orchestrator-8B was trained
|
||||
to emit exactly one per turn.
|
||||
"""
|
||||
# SDK-level path.
|
||||
if sdk_tool_calls:
|
||||
first = sdk_tool_calls[0]
|
||||
name = getattr(getattr(first, "function", None), "name", None)
|
||||
args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
|
||||
try:
|
||||
args = json.loads(args_raw)
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
if isinstance(name, str) and isinstance(args, dict):
|
||||
return {"name": name, "arguments": args}
|
||||
# Text-tag fallback.
|
||||
if not isinstance(content, str):
|
||||
return None
|
||||
m = _TOOL_CALL_TAG_RE.search(content)
|
||||
if m:
|
||||
try:
|
||||
obj = json.loads(m.group(1))
|
||||
except json.JSONDecodeError:
|
||||
obj = None
|
||||
if isinstance(obj, dict):
|
||||
name = obj.get("name")
|
||||
args = obj.get("arguments", {})
|
||||
if isinstance(name, str) and isinstance(args, dict):
|
||||
return {"name": name, "arguments": args}
|
||||
# \boxed{name[, key][ query]: <args>} fallback (see _BOXED_DELEGATION_RE).
|
||||
bm = _BOXED_DELEGATION_RE.search(content.strip())
|
||||
if bm:
|
||||
arg_val = bm.groups()[-1].strip()
|
||||
key = bm.group("key") or "input"
|
||||
return {"name": bm.group(1), "arguments": {key: arg_val}}
|
||||
return None
|
||||
|
||||
|
||||
def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
|
||||
return "\n".join(
|
||||
f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
|
||||
)
|
||||
|
||||
|
||||
def _build_user_prompt(
|
||||
question: str,
|
||||
workers: List[Dict[str, Any]],
|
||||
history: List[Dict[str, Any]],
|
||||
) -> str:
|
||||
pieces = [
|
||||
f"Worker pool:\n{_build_pool_block(workers)}",
|
||||
f"User question:\n{question}",
|
||||
]
|
||||
if history:
|
||||
pieces.append("Conversation so far (orchestrator turns and worker outputs):")
|
||||
for h in history:
|
||||
if h["role"] == "orchestrator":
|
||||
pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
|
||||
else:
|
||||
pieces.append(
|
||||
f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
|
||||
f"{h['output']}"
|
||||
)
|
||||
pieces.append(
|
||||
"Emit the next JSON action object now — exactly one object, no prose."
|
||||
)
|
||||
return "\n\n".join(pieces)
|
||||
|
||||
|
||||
def _strip_fences(s: str) -> str:
|
||||
s = s.strip()
|
||||
if s.startswith("```"):
|
||||
first_nl = s.find("\n")
|
||||
if first_nl != -1:
|
||||
s = s[first_nl + 1 :]
|
||||
if s.endswith("```"):
|
||||
s = s[:-3]
|
||||
s = s.strip()
|
||||
return s
|
||||
|
||||
|
||||
def _parse_action(text: str) -> Optional[Dict[str, Any]]:
|
||||
s = _strip_fences(text)
|
||||
# First try direct parse, then balanced-brace extraction.
|
||||
try:
|
||||
obj = json.loads(s)
|
||||
if isinstance(obj, dict) and "action" in obj:
|
||||
return obj
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
start = s.find("{")
|
||||
if start == -1:
|
||||
return None
|
||||
depth = 0
|
||||
for i in range(start, len(s)):
|
||||
c = s[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
try:
|
||||
obj = json.loads(s[start : i + 1])
|
||||
if isinstance(obj, dict) and "action" in obj:
|
||||
return obj
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _extract_final_answer_text(text: str) -> str:
|
||||
"""Best-effort: pull the answer string from a malformed action emission.
|
||||
|
||||
Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
|
||||
"""
|
||||
m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
|
||||
if m:
|
||||
return m.group(1).encode("utf-8").decode("unicode_escape")
|
||||
m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
|
||||
if m:
|
||||
return m.group(1).strip()
|
||||
return text.strip()
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Prompt strings + tool specs for ToolOrchestraAgent (split from toolorchestra.py)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
ORCHESTRATOR_SYS = """\
|
||||
You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
|
||||
|
||||
{"action": "call_worker", "worker_id": <int>, "input": "<question or instruction for that worker>"}
|
||||
|
||||
{"action": "final_answer", "answer": "<final answer to the user, respecting the question's answer-format rules>"}
|
||||
|
||||
Strategy:
|
||||
|
||||
- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math).
|
||||
- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass.
|
||||
- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase.
|
||||
- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question.
|
||||
"""
|
||||
|
||||
FORCE_FINAL_PROMPT = (
|
||||
"Worker-call budget exhausted. Emit `final_answer` now using everything "
|
||||
"you've learned. Respect the question's answer-format rules."
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# RL-mode constants (Orchestrator-8B, paper-faithful).
|
||||
# ============================================================================
|
||||
#
|
||||
# Verbatim copies of the upstream system prompt / user-prompt template / tools
|
||||
# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
|
||||
# edit the description text — Orchestrator-8B was RL-trained against this
|
||||
# exact wording and pricing/latency table.
|
||||
|
||||
RL_ORCHESTRATOR_SYS = "You are good at using tools."
|
||||
|
||||
RL_TOOLS_SPEC: List[Dict[str, Any]] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "enhance_reasoning",
|
||||
"description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "answer",
|
||||
"description": "give the final answer. Not allowed to call if documents is empty.",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"description": "Search for missing information",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
|
||||
RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
|
||||
"enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
|
||||
"answer": {
|
||||
"model": [
|
||||
"answer-1",
|
||||
"answer-2",
|
||||
"answer-3",
|
||||
"answer-4",
|
||||
"answer-math-1",
|
||||
"answer-math-2",
|
||||
],
|
||||
},
|
||||
"search": {"model": ["search-1", "search-2", "search-3"]},
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
"""Faithful unified-tool rollout loop for ToolOrchestra (arXiv:2511.21689 §2.2).
|
||||
|
||||
One reasoning->action->observation loop where the orchestrator picks **a named
|
||||
tool** (one per model, from :mod:`expert_registry`) each turn, the environment
|
||||
executes it, and the observation is appended to a running context. The rollout
|
||||
ends when the orchestrator emits a turn with **no tool call** (its text is the
|
||||
final answer) or ``max_turns`` is hit.
|
||||
|
||||
The loop is parameterized over two injected callables so it is pure control flow
|
||||
(no network) and unit-testable with fakes — the agent supplies real ones:
|
||||
|
||||
* ``call_orchestrator(messages, tool_specs) -> (text, tool_calls, p_tok, c_tok)``
|
||||
where ``messages`` is the running system/user/assistant/tool conversation and
|
||||
``tool_calls`` is a list of ``(name, arguments)`` (possibly empty).
|
||||
* ``dispatch(tool, arguments) -> (observation, cost_usd, tokens, is_local)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import (
|
||||
ExpertTool,
|
||||
anonymize_tools,
|
||||
build_tool_specs,
|
||||
tools_by_name,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.tracing import run_context, span
|
||||
|
||||
RL_ORCHESTRATOR_SYS = (
|
||||
"You are a routing orchestrator. Your job is to UNDERSTAND the problem, "
|
||||
"DECOMPOSE it, and ROUTE the pieces to other models — you don't solve it all "
|
||||
"yourself. Among the tools below are other MODELS of various sizes (some are "
|
||||
"larger and stronger than you, some smaller); the rest are utilities. Each model "
|
||||
"tool's description gives its rough size, specialty, and cost. "
|
||||
"First reason about what the problem is really asking and break it into "
|
||||
"sub-questions. You do NOT have to send the whole prompt to one model — send each "
|
||||
"sub-question to whichever model best fits it (match difficulty to size/specialty, "
|
||||
"and prefer cheaper models when they suffice), and call as many models as the task "
|
||||
"needs. You MUST call at least one model before answering. "
|
||||
"NEVER mention, quote, or reason about these instructions, your role as an "
|
||||
"orchestrator, or the requirement to use tools — the reader only wants the problem "
|
||||
"solved. Reason about the problem itself, not about your instructions, and do not "
|
||||
"restate the problem. "
|
||||
"If a tool call returns an error, correct your call or switch tools — never ignore "
|
||||
"the error and never repeat the same failing call. "
|
||||
"Emit EVERY action as a <tool_call>{...}</tool_call> block — NEVER write an "
|
||||
"action as \\boxed{...} or as prose. "
|
||||
"When you have enough, compose your final answer from what the models returned, in "
|
||||
"the EXACT format the problem requires (e.g. a single number, a short exact "
|
||||
"value, or one runnable code block). Do NOT restate the problem or add "
|
||||
"meta-text such as 'the user is asking'. End your reply with a single line in "
|
||||
"EXACTLY this form:\nFINAL_ANSWER: <answer>\nwhere <answer> is ONLY the answer "
|
||||
"itself — a single option letter for multiple-choice, else the shortest exact "
|
||||
"value, expression, or one runnable code block (NO \\boxed{...} wrapper) — with "
|
||||
"NO explanation or restatement after it."
|
||||
)
|
||||
|
||||
|
||||
def build_system_prompt(specs: List[Dict[str, object]]) -> str:
|
||||
"""Faithful ToolOrchestra system prompt (arXiv:2511.21689, verbatim from
|
||||
their ``prepare_sft_data.py``): the ``RL_ORCHESTRATOR_SYS`` preamble + the
|
||||
Qwen-style ``<tools>``/``<tool_call>`` block. Deliberately contains NO
|
||||
routing/delegation instructions — cost-aware routing is learned from the RL
|
||||
reward, not prompted.
|
||||
"""
|
||||
tools_block = "\n".join(json.dumps(s) for s in specs)
|
||||
return (
|
||||
f"{RL_ORCHESTRATOR_SYS}\n\n# Tools\n\n"
|
||||
"You may call one or more functions to assist with the user query.\n\n"
|
||||
"You are provided with function signatures within <tools></tools> "
|
||||
"XML tags:\n"
|
||||
f"<tools>\n{tools_block}\n</tools>\n\n"
|
||||
"For each function call, return a json object with function name and "
|
||||
"arguments within <tool_call></tool_call> XML tags:\n"
|
||||
'<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}'
|
||||
"\n</tool_call>"
|
||||
)
|
||||
|
||||
|
||||
# Char-level cap on the accumulated conversation (mirrors the paper's ~24k-token
|
||||
# cap) and a per-observation cap so one giant tool dump can't blow it out.
|
||||
_CONTEXT_CAP = 24000
|
||||
_OBS_CAP = 8000
|
||||
# After this many tool calls, push the orchestrator to stop and answer — bounds
|
||||
# the "keep re-asking the same model forever" loop (esp. gemma).
|
||||
_SOFT_CALL_CAP = 6
|
||||
|
||||
|
||||
def _trim_history(messages: List[Dict[str, str]]) -> None:
|
||||
"""Keep the message history under ``_CONTEXT_CAP`` chars by dropping the
|
||||
OLDEST assistant/tool exchange, never the system prompt or the problem."""
|
||||
|
||||
def total() -> int:
|
||||
return sum(len(m.get("content") or "") for m in messages)
|
||||
|
||||
# messages[0]=system, messages[1]=user problem — always keep those two.
|
||||
while total() > _CONTEXT_CAP and len(messages) > 4:
|
||||
del messages[2:4]
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnifiedTurn:
|
||||
"""One orchestrator turn. ``tool_name is None`` marks the final-answer turn."""
|
||||
|
||||
reasoning: str
|
||||
tool_name: Optional[str] = None
|
||||
arguments: Dict[str, object] = field(default_factory=dict)
|
||||
observation: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnifiedRollout:
|
||||
turns: List[UnifiedTurn]
|
||||
final_answer: str
|
||||
cost_usd: float = 0.0
|
||||
tokens: int = 0
|
||||
num_tool_calls: int = 0
|
||||
parse_failures: int = 0
|
||||
# When experts were anonymized for this rollout, maps the opaque label the
|
||||
# policy saw (e.g. ``expert_a3f9``) back to the real tool name (``gpt_5_5``).
|
||||
anon_map: Optional[Dict[str, str]] = None
|
||||
# The EXACT tool specs the policy saw this rollout (anonymized when
|
||||
# ``anonymize=True``). Serialization MUST build the saved system prompt from
|
||||
# these — not from the real registry — or the prompt's ``<tools>`` block ends
|
||||
# up with real model names+pricing while the assistant ``<tool_call>`` tags
|
||||
# use the anon labels, which both breaks SFT (calls a tool absent from its
|
||||
# list) and re-injects the name/cost bias anonymization removed.
|
||||
tool_specs: Optional[List[Dict[str, object]]] = None
|
||||
|
||||
def tool_calls(self) -> List[Tuple[str, Dict[str, object]]]:
|
||||
return [(t.tool_name, t.arguments) for t in self.turns if t.tool_name]
|
||||
|
||||
|
||||
def _tool_prompt(tool: ExpertTool, arguments: Dict[str, object], question: str) -> str:
|
||||
"""The text we actually send the dispatched tool, framed by its arg schema."""
|
||||
for key in ("input", "query", "code"):
|
||||
val = arguments.get(key)
|
||||
if isinstance(val, str) and val.strip():
|
||||
return val
|
||||
return question
|
||||
|
||||
|
||||
def _run_unified_rollout_inner(
|
||||
question: str,
|
||||
tools: List[ExpertTool],
|
||||
*,
|
||||
call_orchestrator: Callable[
|
||||
..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
|
||||
],
|
||||
dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
|
||||
max_turns: int = 50,
|
||||
system: str = RL_ORCHESTRATOR_SYS,
|
||||
anonymize: bool = False,
|
||||
) -> UnifiedRollout:
|
||||
"""Drive the faithful unified-tool rollout for one task.
|
||||
|
||||
``anonymize``: replace each model expert with an opaque random label, a
|
||||
uniform description and no cost line, shuffled — so the policy can't route on
|
||||
a model's name/position/cost (which we found dominate the choice). The
|
||||
anon->real mapping is returned on the rollout for offline analysis.
|
||||
"""
|
||||
anon_map: Optional[Dict[str, str]] = None
|
||||
if anonymize:
|
||||
tools, anon_map = anonymize_tools(tools, random.Random())
|
||||
specs = build_tool_specs(tools)
|
||||
by_name = tools_by_name(tools)
|
||||
|
||||
# Proper multi-turn conversation (NOT a flattened user blob): the orchestrator
|
||||
# sees its own calls as `assistant` turns and each observation as a `tool`
|
||||
# turn, so it can tell a tool result from user input and doesn't re-derive the
|
||||
# whole problem every turn. This mirrors the serialized SFT format exactly.
|
||||
messages: List[Dict[str, str]] = [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": f"Problem: {question}"},
|
||||
]
|
||||
turns: List[UnifiedTurn] = []
|
||||
cost = 0.0
|
||||
tokens = 0
|
||||
n_tool_calls = 0
|
||||
parse_failures = 0
|
||||
nudges = 0
|
||||
empty_input_nudges = 0
|
||||
final_answer = ""
|
||||
|
||||
for _ in range(max_turns):
|
||||
text, tool_calls, p_tok, c_tok = call_orchestrator(messages, specs)
|
||||
tokens += int(p_tok) + int(c_tok)
|
||||
|
||||
if not tool_calls:
|
||||
# Enforce the "MUST delegate to a model" rule: if the orchestrator
|
||||
# tries to answer before ANY tool call, nudge it to route instead of
|
||||
# accepting a solve-it-yourself answer (which reject-sampling drops).
|
||||
if n_tool_calls == 0 and nudges < 2:
|
||||
nudges += 1
|
||||
messages.append({"role": "assistant", "content": text or ""})
|
||||
messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Make progress by delegating a concrete sub-question to one of the "
|
||||
"models now via a <tool_call>."
|
||||
),
|
||||
}
|
||||
)
|
||||
continue
|
||||
# No tool call -> the orchestrator is answering. Terminate.
|
||||
final_answer = (text or "").strip()
|
||||
turns.append(UnifiedTurn(reasoning=text or "", tool_name=None))
|
||||
messages.append({"role": "assistant", "content": text or ""})
|
||||
break
|
||||
|
||||
name, arguments = tool_calls[0]
|
||||
if name not in by_name:
|
||||
parse_failures += 1
|
||||
messages.append({"role": "assistant", "content": text or ""})
|
||||
messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"[invalid tool {name!r} — choose one from the provided tool list]",
|
||||
}
|
||||
)
|
||||
if parse_failures >= 2:
|
||||
final_answer = (text or "").strip()
|
||||
break
|
||||
continue
|
||||
|
||||
tool = by_name[name]
|
||||
# Empty-input guard: the small orchestrator often emits a <tool_call> with
|
||||
# a blank/missing input on harder multi-turn tasks. Dispatching it returns
|
||||
# a "no input provided" error observation that poisons the whole (often
|
||||
# otherwise-correct) trajectory — the dominant clean-yield killer on hard
|
||||
# tasks. Instead, nudge the model to resend WITH input and drop the
|
||||
# malformed turn (don't record it), so it never enters the training data.
|
||||
if tool.kind == "model" or tool.backend_type != "openjarvis-tool":
|
||||
_has_input = any(
|
||||
isinstance(arguments.get(k), str) and arguments.get(k).strip()
|
||||
for k in ("input", "query", "code")
|
||||
)
|
||||
if not _has_input and empty_input_nudges < 3:
|
||||
empty_input_nudges += 1
|
||||
messages.append({"role": "assistant", "content": text or ""})
|
||||
messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"Your call to {name} had an empty 'input'. Resend the "
|
||||
"<tool_call> with a non-empty 'input' field containing the "
|
||||
"concrete sub-question to delegate."
|
||||
),
|
||||
}
|
||||
)
|
||||
continue
|
||||
obs, dcost, dtok, _is_local = dispatch(tool, arguments)
|
||||
cost += float(dcost)
|
||||
tokens += int(dtok)
|
||||
n_tool_calls += 1
|
||||
turns.append(
|
||||
UnifiedTurn(
|
||||
reasoning=text or "",
|
||||
tool_name=name,
|
||||
arguments=dict(arguments),
|
||||
observation=obs,
|
||||
)
|
||||
)
|
||||
# The model's own action as an assistant turn (the <tool_call> tag in
|
||||
# content, matching the SFT serialization), then the observation as a
|
||||
# distinct `tool` turn the model reads as a tool response.
|
||||
call_content = text or ""
|
||||
if "<tool_call>" not in call_content:
|
||||
call_content = (
|
||||
call_content + "\n" + tool_call_tag(name, arguments)
|
||||
).strip()
|
||||
messages.append({"role": "assistant", "content": call_content})
|
||||
obs_text = obs or ""
|
||||
if len(obs_text) > _OBS_CAP:
|
||||
obs_text = obs_text[:_OBS_CAP] + "\n…[truncated]"
|
||||
messages.append({"role": "tool", "name": name, "content": obs_text})
|
||||
if n_tool_calls >= _SOFT_CALL_CAP:
|
||||
messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"You now have enough information from the models. Do NOT call any "
|
||||
"more tools — reply with your FINAL_ANSWER line only."
|
||||
),
|
||||
}
|
||||
)
|
||||
_trim_history(messages)
|
||||
else:
|
||||
# Hit max_turns with no explicit answer: use the last observation/text.
|
||||
final_answer = (
|
||||
(turns[-1].observation or turns[-1].reasoning).strip() if turns else ""
|
||||
)
|
||||
|
||||
return UnifiedRollout(
|
||||
turns=turns,
|
||||
final_answer=final_answer,
|
||||
cost_usd=cost,
|
||||
tokens=tokens,
|
||||
num_tool_calls=n_tool_calls,
|
||||
parse_failures=parse_failures,
|
||||
anon_map=anon_map,
|
||||
tool_specs=specs,
|
||||
)
|
||||
|
||||
|
||||
def run_unified_rollout(
|
||||
question: str,
|
||||
tools: List[ExpertTool],
|
||||
*,
|
||||
call_orchestrator: Callable[
|
||||
..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
|
||||
],
|
||||
dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
|
||||
max_turns: int = 50,
|
||||
system: str = RL_ORCHESTRATOR_SYS,
|
||||
anonymize: bool = False,
|
||||
) -> UnifiedRollout:
|
||||
"""One ToolOrchestra rollout = one Braintrust trace. Opens the root span,
|
||||
runs the loop (orchestrator turns + expert/tool dispatches nest inside as
|
||||
child spans), then logs the final answer + cost/tokens. No-op when tracing
|
||||
is disabled (see :mod:`.tracing`)."""
|
||||
_run_meta, _run_tags = run_context()
|
||||
_fields = dict(
|
||||
input={"question": question},
|
||||
metadata={
|
||||
"max_turns": max_turns,
|
||||
"anonymize": anonymize,
|
||||
"n_tools": len(tools),
|
||||
**_run_meta,
|
||||
},
|
||||
)
|
||||
if _run_tags:
|
||||
_fields["tags"] = _run_tags
|
||||
with span("toolorchestra.rollout", span_type="task", **_fields) as _s:
|
||||
roll = _run_unified_rollout_inner(
|
||||
question,
|
||||
tools,
|
||||
call_orchestrator=call_orchestrator,
|
||||
dispatch=dispatch,
|
||||
max_turns=max_turns,
|
||||
system=system,
|
||||
anonymize=anonymize,
|
||||
)
|
||||
_s.log(
|
||||
output=roll.final_answer,
|
||||
metrics={
|
||||
"cost_usd": roll.cost_usd,
|
||||
"tokens": roll.tokens,
|
||||
"num_tool_calls": roll.num_tool_calls,
|
||||
"parse_failures": roll.parse_failures,
|
||||
},
|
||||
metadata={
|
||||
"n_experts_available": len(roll.anon_map or {}),
|
||||
"answered": bool(roll.final_answer),
|
||||
# label -> real model, so every anonymized route span in this
|
||||
# trace can be decoded back to the model that actually ran.
|
||||
"anon_map": roll.anon_map or {},
|
||||
},
|
||||
)
|
||||
return roll
|
||||
|
||||
|
||||
def tool_call_tag(name: str, arguments: Dict[str, object]) -> str:
|
||||
"""Render a tool call as the ``<tool_call>{...}</tool_call>`` text the model emits."""
|
||||
return (
|
||||
f"<tool_call>{json.dumps({'name': name, 'arguments': arguments})}</tool_call>"
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"RL_ORCHESTRATOR_SYS",
|
||||
"UnifiedRollout",
|
||||
"UnifiedTurn",
|
||||
"build_system_prompt",
|
||||
"run_unified_rollout",
|
||||
"tool_call_tag",
|
||||
]
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Tavily search + Modal Python sandbox helpers for ToolOrchestraAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
|
||||
# ---- Tavily + Modal helpers -------------------------------------------------
|
||||
|
||||
|
||||
def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]:
|
||||
"""One-shot Tavily search. Returns (text, p_tok=0, c_tok=0).
|
||||
|
||||
Token counts are reported as zero (no LLM was billed); the OpenJarvis
|
||||
accounting layer separately tallies tool-call counts. Falls back to
|
||||
DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``).
|
||||
"""
|
||||
from openjarvis.tools.web_search import WebSearchTool
|
||||
|
||||
tool = WebSearchTool(max_results=max_results)
|
||||
res = tool.execute(query=query, max_results=max_results)
|
||||
text = res.content or ""
|
||||
if not res.success and not text:
|
||||
text = "(no results)"
|
||||
return text, 0, 0
|
||||
|
||||
|
||||
_MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox"
|
||||
|
||||
|
||||
def _call_modal_python(code: str, timeout_s: int = 60) -> Tuple[str, int]:
|
||||
"""Execute a single Python snippet in a fresh Modal Sandbox.
|
||||
|
||||
Returns ``(combined_stdout_stderr, returncode)``. Logs are capped at 8 KiB.
|
||||
Any exception (modal auth, network, sandbox boot failure) is captured into
|
||||
the returned string with a non-zero rc — we never raise back to the
|
||||
orchestrator loop. The sandbox is torn down at the end via ``terminate()``.
|
||||
"""
|
||||
try:
|
||||
import modal
|
||||
|
||||
app = modal.App.lookup(_MODAL_APP_NAME, create_if_missing=True)
|
||||
# python:3.12-slim is small + boots fast; the paper uses a generic
|
||||
# Python image too. We rely on stdlib only — no extra pip installs.
|
||||
image = modal.Image.debian_slim(python_version="3.12")
|
||||
sb = modal.Sandbox.create(
|
||||
"python",
|
||||
"-c",
|
||||
code,
|
||||
app=app,
|
||||
image=image,
|
||||
timeout=int(timeout_s),
|
||||
)
|
||||
sb.wait()
|
||||
try:
|
||||
out = sb.stdout.read() or ""
|
||||
except Exception:
|
||||
out = ""
|
||||
try:
|
||||
err = sb.stderr.read() or ""
|
||||
except Exception:
|
||||
err = ""
|
||||
rc = sb.returncode if sb.returncode is not None else -1
|
||||
try:
|
||||
sb.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
combined = out + (("\n" + err) if err else "")
|
||||
if len(combined) > 8192:
|
||||
combined = combined[:8192] + "\n... (output truncated)"
|
||||
return combined, int(rc)
|
||||
except Exception as exc:
|
||||
return f"[modal-python error: {type(exc).__name__}: {exc}]", -1
|
||||
|
||||
|
||||
_PY_CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL)
|
||||
|
||||
|
||||
def _extract_first_python_block(text: str) -> Optional[str]:
|
||||
"""Return the first ```python ... ``` block (or ```...```), or None."""
|
||||
m = _PY_CODE_RE.search(text or "")
|
||||
return m.group(1).strip() if m else None
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Braintrust telemetry for the ToolOrchestra rollout.
|
||||
|
||||
Each ``run_unified_rollout`` becomes ONE Braintrust trace (root span); the
|
||||
orchestrator turns and every expert/tool dispatch nest inside it, and the
|
||||
underlying OpenAI/Anthropic calls (wrapped clients) nest one level deeper — so
|
||||
you see the full routing tree with inputs/outputs/tokens/cost per node.
|
||||
|
||||
On by default. Degrades to a total no-op (never raises, never changes behavior)
|
||||
when: ``OJ_BRAINTRUST=0``, the ``braintrust`` package isn't installed, or
|
||||
``BRAINTRUST_API_KEY`` is unset. Concurrency: rollouts run one-per-thread
|
||||
(ThreadPoolExecutor); threads don't inherit contextvars, so each rollout opens
|
||||
its own independent root trace and same-thread child calls nest correctly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_STATE: dict = {"resolved": False, "enabled": False, "bt": None}
|
||||
|
||||
|
||||
def _truthy(v: str) -> bool:
|
||||
return v.strip().lower() not in ("0", "false", "no", "off", "")
|
||||
|
||||
|
||||
def _resolve() -> bool:
|
||||
"""Lazily decide whether tracing is active and init the logger once."""
|
||||
if _STATE["resolved"]:
|
||||
return _STATE["enabled"]
|
||||
_STATE["resolved"] = True
|
||||
if not _truthy(os.getenv("OJ_BRAINTRUST", "1")): # on by default
|
||||
return False
|
||||
if not os.getenv("BRAINTRUST_API_KEY"):
|
||||
logger.info(
|
||||
"braintrust on-by-default but BRAINTRUST_API_KEY unset — tracing disabled"
|
||||
)
|
||||
return False
|
||||
try:
|
||||
import braintrust as _bt
|
||||
|
||||
proj_id = os.getenv("OJ_BRAINTRUST_PROJECT_ID")
|
||||
if proj_id:
|
||||
_bt.init_logger(project_id=proj_id)
|
||||
else:
|
||||
_bt.init_logger(project=os.getenv("OJ_BRAINTRUST_PROJECT", "toolorchestra"))
|
||||
_STATE["bt"] = _bt
|
||||
_STATE["enabled"] = True
|
||||
logger.info(
|
||||
"braintrust tracing ENABLED (%s)",
|
||||
f"project_id={proj_id}"
|
||||
if proj_id
|
||||
else f"project={os.getenv('OJ_BRAINTRUST_PROJECT', 'toolorchestra')}",
|
||||
)
|
||||
except Exception as exc: # missing pkg / bad key / init failure — never crash
|
||||
logger.warning("braintrust init failed (%s) — tracing disabled", exc)
|
||||
return _STATE["enabled"]
|
||||
|
||||
|
||||
def enabled() -> bool:
|
||||
return _resolve()
|
||||
|
||||
|
||||
def run_context() -> tuple[dict, list]:
|
||||
"""Run-level metadata + tags for the ROOT rollout span, sourced from env
|
||||
(set by the generation driver). No-op-safe: returns ({}, []) if nothing is
|
||||
set, and never raises. Env keys:
|
||||
|
||||
- ``OJ_RUN_LABEL`` — human run label (also stamped on the uploaded dataset).
|
||||
- ``OJ_GEN_MODEL`` — SPECIFIC gen model id (e.g. ``Qwen/Qwen3.5-9B``).
|
||||
- ``OJ_RUN_STAGE`` — ``prod``|``smoke`` (inferred from the label if unset).
|
||||
- ``OJ_CFG_*`` — config knobs (temperature/max_turns/anonymize/...).
|
||||
"""
|
||||
import datetime
|
||||
|
||||
meta: dict = {}
|
||||
tags: list = []
|
||||
try:
|
||||
label = os.getenv("OJ_RUN_LABEL")
|
||||
gen_model = os.getenv("OJ_GEN_MODEL")
|
||||
stage = os.getenv("OJ_RUN_STAGE")
|
||||
if not stage and label:
|
||||
stage = "smoke" if "smoke" in label.lower() else "prod"
|
||||
if label:
|
||||
meta["run_label"] = label
|
||||
if gen_model:
|
||||
meta["gen_model"] = gen_model
|
||||
if stage:
|
||||
meta["stage"] = stage
|
||||
cfg = {}
|
||||
for env_key, key in (
|
||||
("OJ_CFG_TEMPERATURE", "temperature"),
|
||||
("OJ_CFG_MAX_TURNS", "max_turns"),
|
||||
("OJ_CFG_ANONYMIZE", "anonymize"),
|
||||
("OJ_CFG_REJECTION_ONLY", "rejection_only"),
|
||||
):
|
||||
v = os.getenv(env_key)
|
||||
if v not in (None, ""):
|
||||
cfg[key] = v
|
||||
if cfg:
|
||||
meta["config"] = cfg
|
||||
date = datetime.date.today().isoformat()
|
||||
tags = [t for t in (gen_model, stage, date) if t]
|
||||
except Exception: # never let telemetry enrichment break a rollout
|
||||
return {}, []
|
||||
return meta, tags
|
||||
|
||||
|
||||
def wrap_client(client: Any) -> Any:
|
||||
"""Wrap an OpenAI/Anthropic client so its calls auto-log under the current
|
||||
span. Pass-through (unchanged client) when tracing is off or on any error —
|
||||
so this is always safe to call at client construction."""
|
||||
if not _resolve():
|
||||
return client
|
||||
try:
|
||||
bt = _STATE["bt"]
|
||||
mod = type(client).__module__.lower()
|
||||
if "anthropic" in mod:
|
||||
return bt.wrap_anthropic(client)
|
||||
return bt.wrap_openai(client)
|
||||
except Exception as exc:
|
||||
logger.warning("braintrust wrap_client failed (%s) — using raw client", exc)
|
||||
return client
|
||||
|
||||
|
||||
class _NullSpan:
|
||||
"""No-op span used when tracing is disabled (keeps call sites branch-free)."""
|
||||
|
||||
def log(self, **_kw: Any) -> None:
|
||||
pass
|
||||
|
||||
def __enter__(self) -> "_NullSpan":
|
||||
return self
|
||||
|
||||
def __exit__(self, *_a: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def span(name: str, *, span_type: str = "task", **fields: Any):
|
||||
"""Open a Braintrust span (or a no-op). Use as::
|
||||
|
||||
with span("toolorchestra.rollout", input=...) as s:
|
||||
...
|
||||
s.log(output=..., metrics=..., metadata=...)
|
||||
|
||||
``fields`` (input/metadata/...) are forwarded to ``start_span``.
|
||||
"""
|
||||
if not _resolve():
|
||||
yield _NullSpan()
|
||||
return
|
||||
# Guard only span CREATION (so a Braintrust hiccup degrades to a no-op). Do
|
||||
# NOT wrap the yielded body in try/except: an exception from the rollout would
|
||||
# be thrown into this generator, caught here, and masked by a second yield
|
||||
# ("generator didn't stop after throw()"). Let the body's own exceptions
|
||||
# propagate through the `with` untouched — Braintrust records + re-raises.
|
||||
try:
|
||||
cm = _STATE["bt"].start_span(name=name, type=span_type, **fields)
|
||||
except Exception as exc: # span creation failed — run trace-less, never break
|
||||
logger.warning(
|
||||
"braintrust start_span(%s) failed (%s) — continuing untraced", name, exc
|
||||
)
|
||||
yield _NullSpan()
|
||||
return
|
||||
with cm as s:
|
||||
yield s
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Real backends for the unified-tool rollout — the bridge between the pure
|
||||
:func:`run_unified_rollout` loop and live model/tool calls.
|
||||
|
||||
``make_call_orchestrator`` returns the ``call_orchestrator`` callable (a teacher
|
||||
LLM emitting tool calls over the unified spec); ``make_dispatch`` returns the
|
||||
``dispatch`` callable (executes a chosen tool via ``_call_worker``). Both are
|
||||
import-safe — the OpenAI SDK is imported lazily so this module loads without
|
||||
network or keys.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents.hybrid._prices import cost as _model_cost
|
||||
from openjarvis.agents.hybrid.expert_registry import ExpertTool, to_worker_dict
|
||||
from openjarvis.agents.hybrid.toolorchestra.parsing import _parse_rl_tool_call
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import (
|
||||
UnifiedRollout,
|
||||
build_system_prompt,
|
||||
run_unified_rollout,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.tracing import span
|
||||
from openjarvis.agents.hybrid.toolorchestra.workers import _call_worker
|
||||
|
||||
# Guards the lazy, one-time build of the shared ToolExecutor in
|
||||
# ``_dispatch_openjarvis_tool``: with concurrent rollouts (parallel rejection
|
||||
# sampling) several threads can reach the build at once and would each instantiate
|
||||
# the full tool registry. The lock makes it build-once.
|
||||
_EXECUTOR_BUILD_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def make_call_orchestrator(
|
||||
model: str,
|
||||
*,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
temperature: float = 1.0,
|
||||
max_tokens: int = 4096,
|
||||
timeout: float = 600.0,
|
||||
native_tools: bool = True,
|
||||
frequency_penalty: float = 0.2,
|
||||
presence_penalty: float = 0.0,
|
||||
repetition_penalty: float = 1.15,
|
||||
) -> Callable[..., Tuple[str, List[Tuple[str, Dict[str, Any]]], int, int]]:
|
||||
"""Teacher-orchestrator caller. ``base_url=None`` → OpenAI cloud; set it to a
|
||||
vLLM endpoint (with ``api_key="EMPTY"``) to drive a local served teacher.
|
||||
|
||||
``native_tools`` picks the tool-calling convention, and the two modes MUST NOT
|
||||
be mixed on the same model:
|
||||
|
||||
* ``True`` (default — DATA GENERATION with a *base* model): pass the OpenAI
|
||||
``tools=`` param so the server's native tool template + parser (gemma4 /
|
||||
qwen3_xml / hermes) drive the base model, which only emits reliable tool
|
||||
calls that way. The rollout captures ``(name, arguments)`` and the serializer
|
||||
re-renders them into the canonical JSON ``<tool_call>`` form regardless.
|
||||
* ``False`` (SERVING a *fine-tuned* model): bake the ``<tools>`` block + JSON
|
||||
``<tool_call>`` format into the system prompt (exactly what the serializer
|
||||
trained on) and DROP ``tools=``. If ``tools=`` is set here, the served chat
|
||||
template injects its own XML ``<function=>`` instructions, which conflict
|
||||
with the JSON format the model learned -> it falls back to ``\\boxed{}`` and
|
||||
never routes. ``_parse_rl_tool_call`` scrapes the JSON tag from raw text.
|
||||
"""
|
||||
|
||||
# Create the client ONCE and reuse it for every turn. Creating a fresh OpenAI
|
||||
# client per call (as before) leaks an httpx connection pool each time -> under
|
||||
# heavy parallel generation, sockets pile up in CLOSE-WAIT, the run degrades and
|
||||
# has to be restarted. The orchestrator is the highest-frequency call, so reusing
|
||||
# one client here removes the dominant leak. httpx clients are thread-safe.
|
||||
from openai import OpenAI
|
||||
|
||||
from openjarvis.agents.hybrid.toolorchestra.tracing import wrap_client
|
||||
|
||||
client = wrap_client(
|
||||
OpenAI(base_url=base_url, api_key=api_key or "EMPTY", timeout=timeout)
|
||||
)
|
||||
|
||||
def call_orchestrator(messages: List[Dict[str, Any]], specs: List[Dict[str, Any]]):
|
||||
# ``messages`` is the full running conversation (system/user/assistant/tool)
|
||||
# built by run_unified_rollout — pass it through so the model sees its own
|
||||
# prior calls + tool responses, not a flattened user blob.
|
||||
if native_tools:
|
||||
# Base-model generation: native tool template drives the tool calls.
|
||||
send = messages
|
||||
kwargs = {"tools": specs} if specs else {}
|
||||
else:
|
||||
# Fine-tuned serving: JSON format is baked into the system prompt and
|
||||
# tools= is dropped (see make_call_orchestrator docstring).
|
||||
send = list(messages)
|
||||
if specs and send and send[0].get("role") == "system":
|
||||
send[0] = {"role": "system", "content": build_system_prompt(specs)}
|
||||
kwargs = {}
|
||||
# repetition_penalty is a vLLM extra (not OpenAI-native); only send it to a
|
||||
# local vLLM endpoint (base_url set), never to the cloud frontier APIs.
|
||||
if base_url and repetition_penalty and repetition_penalty != 1.0:
|
||||
kwargs.setdefault("extra_body", {})["repetition_penalty"] = (
|
||||
repetition_penalty
|
||||
)
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=send,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
frequency_penalty=frequency_penalty,
|
||||
presence_penalty=presence_penalty,
|
||||
**kwargs,
|
||||
)
|
||||
msg = resp.choices[0].message
|
||||
text = msg.content or ""
|
||||
sdk_tool_calls = getattr(msg, "tool_calls", None)
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
|
||||
parsed = _parse_rl_tool_call(text, sdk_tool_calls)
|
||||
tool_calls = [(parsed["name"], parsed["arguments"])] if parsed else []
|
||||
return text, tool_calls, int(p), int(c)
|
||||
|
||||
return call_orchestrator
|
||||
|
||||
|
||||
def _dispatch_openjarvis_tool(
|
||||
tool: ExpertTool,
|
||||
arguments: Dict[str, Any],
|
||||
executor_holder: Dict[str, Any],
|
||||
) -> Tuple[str, float, int, bool]:
|
||||
"""Execute a bridged real OpenJarvis tool via its ``ToolExecutor``.
|
||||
|
||||
Builds the executor lazily (cached in ``executor_holder``) and degrades to a
|
||||
clear error string — never a crash — if the tool registry / executor can't be
|
||||
instantiated. Returns ``(content, cost_usd, total_tokens, is_local=True)``.
|
||||
"""
|
||||
try:
|
||||
executor = executor_holder.get("executor")
|
||||
if executor is None:
|
||||
with _EXECUTOR_BUILD_LOCK:
|
||||
# Re-check inside the lock: another thread may have built it while
|
||||
# we waited.
|
||||
executor = executor_holder.get("executor")
|
||||
if executor is None:
|
||||
from openjarvis.core.registry import ToolRegistry
|
||||
from openjarvis.tools._stubs import ToolExecutor
|
||||
|
||||
instances = []
|
||||
for name in ToolRegistry.keys():
|
||||
entry = ToolRegistry.get(name)
|
||||
try:
|
||||
instances.append(
|
||||
entry() if isinstance(entry, type) else entry
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
# Auto-approve confirmation-gated tools (shell_exec,
|
||||
# git_commit, ...): rollouts/eval run headless in a sandbox, so
|
||||
# there is no human to confirm. Without this, those tools return
|
||||
# a "requires confirmation" error instead of executing — which
|
||||
# would silently break TerminalBench.
|
||||
executor = ToolExecutor(
|
||||
instances,
|
||||
interactive=True,
|
||||
confirm_callback=lambda _prompt: True,
|
||||
)
|
||||
executor_holder["executor"] = executor
|
||||
|
||||
from openjarvis.core.types import ToolCall
|
||||
|
||||
call = ToolCall(
|
||||
id=f"orch-{tool.name}",
|
||||
name=str(tool.model),
|
||||
arguments=json.dumps(arguments or {}),
|
||||
)
|
||||
result = executor.execute(call)
|
||||
usage = getattr(result, "usage", None) or {}
|
||||
total_tokens = int(usage.get("total_tokens", 0) or 0)
|
||||
return (
|
||||
str(getattr(result, "content", "")),
|
||||
float(getattr(result, "cost_usd", 0.0) or 0.0),
|
||||
total_tokens,
|
||||
True,
|
||||
)
|
||||
except Exception as exc: # never crash the rollout on a tool-bridge failure
|
||||
return (f"[openjarvis-tool error: {tool.model}: {exc}]", 0.0, 0, True)
|
||||
|
||||
|
||||
def make_dispatch(
|
||||
cfg: Optional[Dict[str, Any]] = None,
|
||||
) -> Callable[[ExpertTool, Dict[str, Any]], Tuple[str, float, int, bool]]:
|
||||
"""Tool-execution caller: run the chosen tool and return (obs, cost, tokens, is_local)."""
|
||||
cfg = cfg or {}
|
||||
executor_holder: Dict[str, Any] = {}
|
||||
|
||||
def _dispatch_inner(tool: ExpertTool, arguments: Dict[str, Any]):
|
||||
# Bridged real OpenJarvis tools run through the ToolExecutor, not
|
||||
# the model/worker path.
|
||||
if tool.backend_type == "openjarvis-tool":
|
||||
return _dispatch_openjarvis_tool(tool, arguments or {}, executor_holder)
|
||||
|
||||
worker = to_worker_dict(tool)
|
||||
prompt = ""
|
||||
for key in ("input", "query", "code"):
|
||||
val = arguments.get(key)
|
||||
if isinstance(val, str) and val.strip():
|
||||
prompt = val
|
||||
break
|
||||
if not prompt.strip():
|
||||
# The orchestrator emitted a tool call with no/empty input. Don't hit
|
||||
# the API (it 400s on empty content) — return a usable error so the
|
||||
# rollout keeps going instead of dropping.
|
||||
return (
|
||||
f"[{tool.name}: no input provided — supply a non-empty "
|
||||
"'input' to delegate]",
|
||||
0.0,
|
||||
0,
|
||||
True,
|
||||
)
|
||||
text, p, c, is_local, extra_cost, _n = _call_worker(worker, prompt, cfg)
|
||||
usd = (0.0 if is_local else _model_cost(str(tool.model), p, c)) + float(
|
||||
extra_cost
|
||||
)
|
||||
return text, usd, int(p) + int(c), bool(is_local)
|
||||
|
||||
def dispatch(tool: ExpertTool, arguments: Dict[str, Any]):
|
||||
# One span per route so the trace tree shows which expert/tool was called,
|
||||
# with the delegated input, the returned observation, and cost/tokens. The
|
||||
# wrapped OpenAI/Anthropic call (if any) nests one level deeper.
|
||||
# Span is titled with the REAL model that actually ran (``tool.model``) so
|
||||
# the trace reads clearly even under anonymization; ``anon_label`` in the
|
||||
# metadata records the opaque label the orchestrator actually saw/chose.
|
||||
real_model = str(tool.model)
|
||||
with span(
|
||||
f"route:{real_model}",
|
||||
span_type="tool",
|
||||
input=arguments,
|
||||
metadata={
|
||||
"real_model": real_model,
|
||||
"anon_label": tool.name,
|
||||
"backend": tool.backend_type,
|
||||
},
|
||||
) as s:
|
||||
obs, usd, toks, is_local = _dispatch_inner(tool, arguments)
|
||||
s.log(
|
||||
output=obs,
|
||||
metrics={"cost_usd": usd, "tokens": toks},
|
||||
metadata={"is_local": is_local},
|
||||
)
|
||||
return obs, usd, toks, is_local
|
||||
|
||||
return dispatch
|
||||
|
||||
|
||||
def teacher_rollout(
|
||||
question: str,
|
||||
tools: List[ExpertTool],
|
||||
*,
|
||||
teacher_model: str,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
temperature: float = 1.0,
|
||||
max_turns: int = 50,
|
||||
cfg: Optional[Dict[str, Any]] = None,
|
||||
) -> UnifiedRollout:
|
||||
"""Convenience: one full teacher rollout with real backends."""
|
||||
return run_unified_rollout(
|
||||
question,
|
||||
tools,
|
||||
call_orchestrator=make_call_orchestrator(
|
||||
teacher_model,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
temperature=temperature,
|
||||
),
|
||||
dispatch=make_dispatch(cfg),
|
||||
max_turns=max_turns,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["make_call_orchestrator", "make_dispatch", "teacher_rollout"]
|
||||
@@ -0,0 +1,526 @@
|
||||
"""Worker pool resolution + dispatch for ToolOrchestraAgent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents.hybrid._base import (
|
||||
ANTHROPIC_WEB_SEARCH_TOOL,
|
||||
GEMINI_SEARCH_COST_PER_CALL,
|
||||
OPENAI_WEB_SEARCH_COST_PER_CALL,
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
LocalCloudAgent,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
PRICES,
|
||||
is_gpt5_family,
|
||||
supports_temperature,
|
||||
)
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
|
||||
from openjarvis.agents.hybrid.toolorchestra.experts import (
|
||||
_PAPER_CODER_OPENROUTER,
|
||||
_PAPER_GENERALIST_TIER3_OPENROUTER,
|
||||
)
|
||||
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
|
||||
_call_modal_python,
|
||||
_call_tavily_search,
|
||||
)
|
||||
|
||||
|
||||
def _paper_pool(
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Paper-match worker pool (registered for traces / inspection).
|
||||
|
||||
NOTE: in RL mode the orchestrator dispatches via tool/slot rather than
|
||||
worker_id, so this list is purely informational — `_paper_expert_for`
|
||||
is the actual routing function. We still return a list here so the
|
||||
paradigm's trace metadata has something concrete to log.
|
||||
"""
|
||||
pool: List[Dict[str, Any]] = []
|
||||
if local_model and local_endpoint:
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "local-qwen",
|
||||
"type": "vllm",
|
||||
"model": local_model,
|
||||
"base_url": local_endpoint,
|
||||
"description": "Local Qwen vLLM (paper uses Qwen3-32B).",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "tavily-search",
|
||||
"type": "tavily-search",
|
||||
"model": "tavily",
|
||||
"description": "Tavily web search.",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "modal-python",
|
||||
"type": "modal-python",
|
||||
"model": "modal-python",
|
||||
"description": "Modal Sandbox for one-shot Python exec.",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "code-specialist",
|
||||
"type": "openrouter",
|
||||
"model": _PAPER_CODER_OPENROUTER,
|
||||
"description": "Qwen-2.5-Coder-32B via OpenRouter (paper).",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "generalist-llama",
|
||||
"type": "openrouter",
|
||||
"model": _PAPER_GENERALIST_TIER3_OPENROUTER,
|
||||
"description": "Llama-3.3-70B-Instruct via OpenRouter (paper tier-3).",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "generalist-gpt5",
|
||||
"type": "openai",
|
||||
"model": "gpt-5",
|
||||
"description": "GPT-5 frontier generalist.",
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "generalist-gpt5-mini",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
"description": "GPT-5-mini mid generalist.",
|
||||
}
|
||||
)
|
||||
return pool
|
||||
|
||||
|
||||
def _default_pool(
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str = "claude-opus-4-7",
|
||||
cloud_endpoint: str = "anthropic",
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Default heterogeneous worker pool.
|
||||
|
||||
The frontier worker's ``type`` + ``model`` track the cell's resolved
|
||||
``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5,
|
||||
gemini-2.5-pro, …) route their frontier slot to the right SDK.
|
||||
"""
|
||||
ep = (cloud_endpoint or "anthropic").lower()
|
||||
if ep not in ("anthropic", "openai", "gemini"):
|
||||
ep = "anthropic"
|
||||
pool: List[Dict[str, Any]] = []
|
||||
if local_model and local_endpoint:
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "local-qwen",
|
||||
"type": "vllm",
|
||||
"model": local_model,
|
||||
"base_url": local_endpoint,
|
||||
"description": (
|
||||
"Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
|
||||
"concise extraction, formatting, arithmetic on given data."
|
||||
),
|
||||
}
|
||||
)
|
||||
if ep == "openai":
|
||||
search_type = "openai-web-search"
|
||||
search_model = cloud_model
|
||||
search_desc = "OpenAI hosted web search on the configured frontier model."
|
||||
elif ep == "gemini":
|
||||
search_type = "gemini-web-search"
|
||||
search_model = cloud_model
|
||||
search_desc = "Gemini Google Search grounding on the configured frontier model."
|
||||
else:
|
||||
search_type = "anthropic-web-search"
|
||||
search_model = _DEFAULT_WEB_SEARCH_MODEL
|
||||
search_desc = "Anthropic server-side web_search."
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "web-search",
|
||||
"type": search_type,
|
||||
"model": search_model,
|
||||
"description": (
|
||||
f"{search_desc} Use for facts that need a lookup "
|
||||
"(recent events, rare names/dates, niche sources). Returns a digest."
|
||||
),
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": f"frontier-{ep}",
|
||||
"type": ep,
|
||||
"model": cloud_model,
|
||||
"description": (
|
||||
"Frontier reasoning model. Use for hard multi-step reasoning, "
|
||||
"code review, or a final synthesis pass. Expensive — use sparingly."
|
||||
),
|
||||
}
|
||||
)
|
||||
pool.append(
|
||||
{
|
||||
"id": len(pool),
|
||||
"name": "frontier-openai-mini",
|
||||
"type": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
"description": (
|
||||
"Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
|
||||
"fraction of frontier cost."
|
||||
),
|
||||
}
|
||||
)
|
||||
return pool
|
||||
|
||||
|
||||
# Worker types toolorchestra's `_call_worker` actually dispatches.
|
||||
#
|
||||
# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
|
||||
# `tavily-search` — Tavily API search (the paper's web tool).
|
||||
# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
|
||||
# Used for the code/math specialists and Llama-3.3-70B /
|
||||
# Qwen3-32B generalists.
|
||||
# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
|
||||
# paper's "Python sandbox" inside `enhance_reasoning`).
|
||||
_TOOLORCH_VALID_TYPES = (
|
||||
"vllm",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"anthropic-web-search",
|
||||
"openai-web-search",
|
||||
"gemini",
|
||||
"gemini-web-search",
|
||||
"tavily-search",
|
||||
"openrouter",
|
||||
"modal-python",
|
||||
)
|
||||
_TOOLORCH_SEARCH_TYPES = (
|
||||
"anthropic-web-search",
|
||||
"openai-web-search",
|
||||
"gemini-web-search",
|
||||
"tavily-search",
|
||||
)
|
||||
|
||||
# Default model used when an `anthropic-web-search` entry omits `model`.
|
||||
_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
|
||||
|
||||
|
||||
def _resolve_worker_pool(
|
||||
cfg: Dict[str, Any],
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str,
|
||||
cloud_endpoint: str = "anthropic",
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Return the worker pool for this run.
|
||||
|
||||
Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
|
||||
default pool is ignored entirely. Falls back to ``_default_pool`` when
|
||||
the override is absent.
|
||||
|
||||
Each user-supplied entry must be a dict with keys ``id``, ``name``,
|
||||
``type``, and (for non-search types) ``model``. Search worker types are
|
||||
``anthropic-web-search``, ``openai-web-search``, ``gemini-web-search``,
|
||||
and ``tavily-search``. ``anthropic-web-search`` entries may omit
|
||||
``model`` — it defaults to ``claude-haiku-4-5``. OpenAI and Gemini
|
||||
search workers default to the configured cloud model. Tavily does not
|
||||
require a model.
|
||||
|
||||
Substitution: ``model = "$local"`` (or ``"<local>"``) resolves to
|
||||
``local_model``; ``model = "$cloud"`` / ``"<cloud>"`` to ``cloud_model``.
|
||||
|
||||
On any validation failure, raises ``ValueError`` with the message
|
||||
``"Invalid worker_pool entry [<id>]: <reason>"``. Fails fast at agent
|
||||
init rather than mid-task.
|
||||
"""
|
||||
override = cfg.get("worker_pool")
|
||||
if override is None:
|
||||
return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint)
|
||||
if not isinstance(override, list) or not override:
|
||||
raise ValueError(
|
||||
"Invalid worker_pool entry [-]: worker_pool must be a non-empty list"
|
||||
)
|
||||
|
||||
resolved: List[Dict[str, Any]] = []
|
||||
seen_ids: set = set()
|
||||
has_non_search = False
|
||||
for raw in override:
|
||||
wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?"
|
||||
if not isinstance(raw, dict):
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict"
|
||||
)
|
||||
entry = dict(raw)
|
||||
wid = entry.get("id")
|
||||
if not isinstance(wid, int):
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int"
|
||||
)
|
||||
if wid in seen_ids:
|
||||
raise ValueError(f"Invalid worker_pool entry [{wid}]: duplicate id")
|
||||
seen_ids.add(wid)
|
||||
if not entry.get("name") or not isinstance(entry["name"], str):
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string"
|
||||
)
|
||||
wtype = entry.get("type") or entry.get("endpoint")
|
||||
if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'type' must be one of "
|
||||
f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})"
|
||||
)
|
||||
wtype = wtype.lower()
|
||||
entry["type"] = wtype
|
||||
# Substitute $local / $cloud placeholders (before any model check).
|
||||
model = entry.get("model")
|
||||
if isinstance(model, str) and model in ("$local", "<local>"):
|
||||
if not local_model:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: model='{model}' "
|
||||
"requires a local_model to be configured for this cell"
|
||||
)
|
||||
model = local_model
|
||||
entry["model"] = model
|
||||
elif isinstance(model, str) and model in ("$cloud", "<cloud>"):
|
||||
model = cloud_model
|
||||
entry["model"] = model
|
||||
if wtype in _TOOLORCH_SEARCH_TYPES:
|
||||
if model in (None, ""):
|
||||
if wtype == "anthropic-web-search":
|
||||
model = _DEFAULT_WEB_SEARCH_MODEL
|
||||
elif wtype in ("openai-web-search", "gemini-web-search"):
|
||||
model = cloud_model
|
||||
else:
|
||||
model = wtype
|
||||
entry["model"] = model
|
||||
elif not isinstance(model, str):
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
|
||||
)
|
||||
if (
|
||||
wtype in ("openai-web-search", "gemini-web-search")
|
||||
and model not in PRICES
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: model {model!r} "
|
||||
f"is not in PRICES (known: {sorted(PRICES)})"
|
||||
)
|
||||
# Search workers don't satisfy the "needs a solver" requirement.
|
||||
else:
|
||||
if not isinstance(model, str) or not model:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
|
||||
)
|
||||
if wtype == "vllm":
|
||||
if not entry.get("base_url"):
|
||||
if not local_endpoint:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: vllm worker needs "
|
||||
"'base_url' (or a configured local_endpoint to fall back to)"
|
||||
)
|
||||
entry["base_url"] = local_endpoint
|
||||
entry.setdefault("api_key", "EMPTY")
|
||||
else:
|
||||
if model not in PRICES:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: model {model!r} "
|
||||
f"is not in PRICES (known: {sorted(PRICES)})"
|
||||
)
|
||||
has_non_search = True
|
||||
entry.setdefault(
|
||||
"description",
|
||||
f"User-supplied {wtype} worker ({model}).",
|
||||
)
|
||||
resolved.append(entry)
|
||||
|
||||
if not has_non_search:
|
||||
raise ValueError(
|
||||
"Invalid worker_pool entry [-]: worker_pool must contain at least "
|
||||
"one non-search worker (vllm / openai / anthropic / gemini)"
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def _call_worker(
|
||||
worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
|
||||
) -> Tuple[str, int, int, bool, float, int]:
|
||||
"""Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
|
||||
wtype = worker.get("type", "openai")
|
||||
max_tok = int(
|
||||
cfg.get("worker_max_tokens") or os.environ.get("OJ_WORKER_MAX_TOKENS", "4096")
|
||||
)
|
||||
temp = float(cfg.get("worker_temperature", 0.2))
|
||||
|
||||
if wtype == "vllm":
|
||||
text, p, c = LocalCloudAgent._call_vllm(
|
||||
worker["model"],
|
||||
worker["base_url"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return text, p, c, True, 0.0, 0
|
||||
if wtype == "openai":
|
||||
is_gpt5 = is_gpt5_family(worker["model"])
|
||||
eff_temp = 1.0 if is_gpt5 else temp
|
||||
# GPT-5 is a reasoning model: hidden reasoning tokens count against
|
||||
# `max_completion_tokens`, so a 4096 cap can be fully consumed by
|
||||
# reasoning and leave 0 visible content (empty answer). Give the
|
||||
# reasoning headroom on top of the answer budget.
|
||||
eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
|
||||
text, p, c = LocalCloudAgent._call_openai(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=eff_max_tok,
|
||||
temperature=eff_temp,
|
||||
)
|
||||
return text, p, c, False, 0.0, 0
|
||||
if wtype == "gemini":
|
||||
text, p, c = LocalCloudAgent._call_gemini(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
)
|
||||
return text, p, c, False, 0.0, 0
|
||||
if wtype == "anthropic":
|
||||
eff_temp = temp if supports_temperature(worker["model"]) else 0.0
|
||||
text, p, c, _ = LocalCloudAgent._call_anthropic(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=eff_temp,
|
||||
)
|
||||
return text, p, c, False, 0.0, 0
|
||||
if wtype == "anthropic-web-search":
|
||||
eff_temp = temp if supports_temperature(worker["model"]) else 0.0
|
||||
text, p, c, n_searches = LocalCloudAgent._call_anthropic(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=eff_temp,
|
||||
tools=[ANTHROPIC_WEB_SEARCH_TOOL],
|
||||
tool_choice={"type": "any"},
|
||||
)
|
||||
extra = n_searches * WEB_SEARCH_COST_PER_CALL
|
||||
return text, p, c, False, extra, n_searches
|
||||
if wtype == "openai-web-search":
|
||||
eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp
|
||||
text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max(max_tok, 16384)
|
||||
if is_gpt5_family(worker["model"])
|
||||
else max_tok,
|
||||
temperature=eff_temp,
|
||||
)
|
||||
extra = n_searches * OPENAI_WEB_SEARCH_COST_PER_CALL
|
||||
return text, p, c, False, extra, n_searches
|
||||
if wtype == "gemini-web-search":
|
||||
text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
)
|
||||
extra = n_searches * GEMINI_SEARCH_COST_PER_CALL
|
||||
return text, p, c, False, extra, n_searches
|
||||
if wtype == "tavily-search":
|
||||
# Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL`
|
||||
# for parity with the Anthropic web-search worker. One call = one
|
||||
# "n_search" for accounting.
|
||||
max_results = int(cfg.get("tavily_max_results", 5))
|
||||
text, p, c = _call_tavily_search(str(prompt), max_results=max_results)
|
||||
return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1
|
||||
if wtype == "openrouter":
|
||||
text, p, c = LocalCloudAgent._call_openrouter(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
)
|
||||
return text, p, c, False, 0.0, 0
|
||||
if wtype == "modal-python":
|
||||
# `prompt` is the python code string to exec.
|
||||
timeout_s = int(cfg.get("modal_python_timeout_s", 60))
|
||||
out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
|
||||
# No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
|
||||
# charges per sandbox-second — not tracked here.
|
||||
return out, 0, 0, False, 0.0, 0
|
||||
raise ValueError(f"unsupported worker type: {wtype!r}")
|
||||
|
||||
|
||||
def _swe_call_worker(
|
||||
worker: Dict[str, Any],
|
||||
prompt: str,
|
||||
cfg: Dict[str, Any],
|
||||
task: Dict[str, Any],
|
||||
workdir: Path,
|
||||
turn: int,
|
||||
) -> Tuple[str, int, int, bool, float, int, int]:
|
||||
"""SWE-bench worker dispatch: route solver workers through
|
||||
run_swe_agent_loop on a shared workdir. Web-search workers fall back
|
||||
to the regular one-shot dispatch (search isn't an agent loop).
|
||||
|
||||
Trailing ``bash_turns`` (last element) counts agent-loop turns so the
|
||||
caller can surface ``tool_calls`` per row. Fallbacks to one-shot
|
||||
workers return 0 bash turns (no agent loop ran)."""
|
||||
wtype = worker.get("type", "openai")
|
||||
if wtype in _TOOLORCH_SEARCH_TYPES:
|
||||
# Search workers stay one-shot.
|
||||
text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
|
||||
return text, p, c, is_local, extra, n_searches, 0
|
||||
if wtype == "vllm":
|
||||
backbone = "local"
|
||||
endpoint = worker.get("base_url")
|
||||
loop_cloud_endpoint = "anthropic" # unused when backbone=local
|
||||
elif wtype in ("anthropic", "openai", "gemini"):
|
||||
backbone = "cloud"
|
||||
endpoint = None
|
||||
loop_cloud_endpoint = wtype
|
||||
else:
|
||||
# Unknown type — one-shot fallback.
|
||||
text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
|
||||
return text, p, c, is_local, extra, n_searches, 0
|
||||
out = run_swe_agent_loop(
|
||||
task,
|
||||
backbone=backbone,
|
||||
backbone_model=worker["model"],
|
||||
cloud_endpoint=loop_cloud_endpoint,
|
||||
local_endpoint=endpoint,
|
||||
initial_prompt=prompt,
|
||||
max_turns=int(cfg.get("swe_max_turns", 30)),
|
||||
bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
|
||||
output_cap=int(cfg.get("swe_output_cap", 10_000)),
|
||||
turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
|
||||
trace_prefix=f"toolorch_turn{turn}",
|
||||
workdir=workdir,
|
||||
)
|
||||
is_local = backbone == "local"
|
||||
return (
|
||||
out["final_summary"] or out["answer"],
|
||||
out["tokens_in"],
|
||||
out["tokens_out"],
|
||||
is_local,
|
||||
0.0,
|
||||
0,
|
||||
int(out["turns"]),
|
||||
)
|
||||
@@ -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))
|
||||
|
||||
@@ -2,12 +2,44 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from openjarvis.evals.core.backend import InferenceBackend
|
||||
from openjarvis.evals.core.types import EvalRecord
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
# Transient failures worth retrying the judge call on. A single rate-limit
|
||||
# (429) used to fall straight through to `llm_fallback_error` and score the
|
||||
# sample WRONG — e.g. 93/100 GAIA judge calls 429'd on one run and zeroed the
|
||||
# whole bench. Retry with exponential backoff so a busy judge endpoint doesn't
|
||||
# get mis-scored as a wrong answer.
|
||||
_JUDGE_MAX_RETRIES = 6
|
||||
_JUDGE_BASE_DELAY_S = 2.0
|
||||
_JUDGE_MAX_DELAY_S = 60.0
|
||||
_RETRYABLE_MARKERS = (
|
||||
"429",
|
||||
"rate_limit",
|
||||
"rate limit",
|
||||
"overloaded",
|
||||
"timeout",
|
||||
"timed out",
|
||||
"503",
|
||||
"502",
|
||||
"500",
|
||||
"connection",
|
||||
"temporarily unavailable",
|
||||
)
|
||||
|
||||
|
||||
def _is_retryable_judge_error(exc: Exception) -> bool:
|
||||
msg = str(exc).lower()
|
||||
return any(marker in msg for marker in _RETRYABLE_MARKERS)
|
||||
|
||||
|
||||
class Scorer(ABC):
|
||||
"""Base class for all scorers."""
|
||||
@@ -42,14 +74,41 @@ class LLMJudgeScorer(Scorer):
|
||||
temperature: float = 0.0,
|
||||
max_tokens: int = 2048,
|
||||
) -> str:
|
||||
"""Send a prompt to the judge LLM and return the response text."""
|
||||
return self._judge_backend.generate(
|
||||
prompt,
|
||||
model=self._judge_model,
|
||||
system=system,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
"""Send a prompt to the judge LLM and return the response text.
|
||||
|
||||
Retries transient failures (429 / rate-limit / 5xx / timeout) with
|
||||
exponential backoff + jitter so a busy judge endpoint doesn't get
|
||||
mis-scored as a wrong answer. Non-retryable errors, or exhaustion of
|
||||
the retry budget, re-raise to the caller (which records the failure).
|
||||
"""
|
||||
last_exc: Optional[Exception] = None
|
||||
for attempt in range(_JUDGE_MAX_RETRIES):
|
||||
try:
|
||||
return self._judge_backend.generate(
|
||||
prompt,
|
||||
model=self._judge_model,
|
||||
system=system,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - re-raised below
|
||||
last_exc = exc
|
||||
if attempt == _JUDGE_MAX_RETRIES - 1 or not _is_retryable_judge_error(
|
||||
exc
|
||||
):
|
||||
raise
|
||||
delay = min(_JUDGE_BASE_DELAY_S * (2**attempt), _JUDGE_MAX_DELAY_S)
|
||||
delay += random.uniform(0.0, delay * 0.25) # jitter
|
||||
LOGGER.warning(
|
||||
"judge call failed (attempt %d/%d): %s — retrying in %.1fs",
|
||||
attempt + 1,
|
||||
_JUDGE_MAX_RETRIES,
|
||||
exc,
|
||||
delay,
|
||||
)
|
||||
time.sleep(delay)
|
||||
# Unreachable (loop either returns or raises), but keeps type-checkers happy.
|
||||
raise last_exc # type: ignore[misc]
|
||||
|
||||
|
||||
__all__ = ["LLMJudgeScorer", "Scorer"]
|
||||
|
||||
@@ -140,7 +140,14 @@ class GAIAScorer(LLMJudgeScorer):
|
||||
if structured_match:
|
||||
is_correct = structured_match.group(1).lower() == "yes"
|
||||
else:
|
||||
is_correct = "CORRECT" in raw.upper() and "INCORRECT" not in raw.upper()
|
||||
up = raw.upper()
|
||||
# "not correct" / "is not correct" contain "CORRECT" and lack
|
||||
# "INCORRECT" — guard the negations explicitly or they read True.
|
||||
is_correct = (
|
||||
"CORRECT" in up
|
||||
and "INCORRECT" not in up
|
||||
and "NOT CORRECT" not in up
|
||||
)
|
||||
|
||||
meta: Dict[str, Any] = {
|
||||
"match_type": "llm_fallback",
|
||||
|
||||
@@ -27,6 +27,25 @@ class MMLUProScorer(LLMJudgeScorer):
|
||||
return "".join(chr(ord("A") + i) for i in range(n))
|
||||
return "ABCDEFGHIJ"
|
||||
|
||||
def _extract_answer_direct(
|
||||
self,
|
||||
model_answer: str,
|
||||
valid_letters: str,
|
||||
) -> Optional[str]:
|
||||
"""High-confidence, unambiguous letter straight from the output.
|
||||
|
||||
Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex
|
||||
prose like "the answer is X": the letters "I" and "A" are real English
|
||||
words, so "FINAL_ANSWER: I need ..." would grab "I" — the exact bug that
|
||||
mislabels answers. Semantic extraction is left to the judge below.
|
||||
"""
|
||||
if not model_answer:
|
||||
return None
|
||||
for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)):
|
||||
if cand.upper() in valid_letters:
|
||||
return cand.upper()
|
||||
return None
|
||||
|
||||
def _extract_answer_with_llm(
|
||||
self,
|
||||
problem: str,
|
||||
@@ -56,10 +75,13 @@ class MMLUProScorer(LLMJudgeScorer):
|
||||
)
|
||||
|
||||
extracted = raw_response.strip().upper()
|
||||
if "NONE" in extracted:
|
||||
return None
|
||||
|
||||
# Handle "The answer is: A" etc.
|
||||
# Accept only an ISOLATED letter — never the first cap of a word like
|
||||
# "It"/"None". Prefer an explicit "the answer is X" phrasing.
|
||||
answer_match = re.search(
|
||||
r"(?:THE ANSWER IS:?\s*)?([A-Z])",
|
||||
r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])",
|
||||
extracted,
|
||||
re.IGNORECASE,
|
||||
)
|
||||
@@ -86,11 +108,17 @@ class MMLUProScorer(LLMJudgeScorer):
|
||||
|
||||
valid_letters = self._valid_letters_from_options(record.metadata)
|
||||
|
||||
candidate = self._extract_answer_with_llm(
|
||||
record.problem,
|
||||
model_answer,
|
||||
valid_letters,
|
||||
)
|
||||
# Parse straight from the output first (fast, deterministic, no "I" bug);
|
||||
# only fall back to the judge when no isolated letter is present.
|
||||
method = "direct"
|
||||
candidate = self._extract_answer_direct(model_answer, valid_letters)
|
||||
if not candidate:
|
||||
method = "llm"
|
||||
candidate = self._extract_answer_with_llm(
|
||||
record.problem,
|
||||
model_answer,
|
||||
valid_letters,
|
||||
)
|
||||
if not candidate:
|
||||
return None, {"reason": "no_choice_letter_extracted"}
|
||||
|
||||
@@ -99,6 +127,7 @@ class MMLUProScorer(LLMJudgeScorer):
|
||||
"reference_letter": ref,
|
||||
"candidate_letter": candidate,
|
||||
"valid_letters": valid_letters,
|
||||
"extract_method": method,
|
||||
}
|
||||
return is_correct, meta
|
||||
|
||||
|
||||
@@ -27,6 +27,25 @@ class SuperGPQAScorer(LLMJudgeScorer):
|
||||
return "".join(chr(ord("A") + i) for i in range(n))
|
||||
return "ABCD"
|
||||
|
||||
def _extract_answer_direct(
|
||||
self,
|
||||
model_answer: str,
|
||||
valid_letters: str,
|
||||
) -> Optional[str]:
|
||||
"""High-confidence, unambiguous letter straight from the output.
|
||||
|
||||
Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex
|
||||
prose like "the answer is X": the letters "I" and "A" are real English
|
||||
words, so "FINAL_ANSWER: I need ..." would grab "I" — the exact bug that
|
||||
mislabels answers. Semantic extraction is left to the judge below.
|
||||
"""
|
||||
if not model_answer:
|
||||
return None
|
||||
for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)):
|
||||
if cand.upper() in valid_letters:
|
||||
return cand.upper()
|
||||
return None
|
||||
|
||||
def _extract_answer_with_llm(
|
||||
self,
|
||||
problem: str,
|
||||
@@ -56,10 +75,13 @@ class SuperGPQAScorer(LLMJudgeScorer):
|
||||
)
|
||||
|
||||
extracted = raw_response.strip().upper()
|
||||
if "NONE" in extracted:
|
||||
return None
|
||||
|
||||
# Handle "The answer is: A" etc.
|
||||
# Accept only an ISOLATED letter — never the first cap of a word like
|
||||
# "It"/"None". Prefer an explicit "the answer is X" phrasing.
|
||||
answer_match = re.search(
|
||||
r"(?:THE ANSWER IS:?\s*)?([A-Z])",
|
||||
r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])",
|
||||
extracted,
|
||||
re.IGNORECASE,
|
||||
)
|
||||
@@ -86,11 +108,17 @@ class SuperGPQAScorer(LLMJudgeScorer):
|
||||
|
||||
valid_letters = self._valid_letters_from_options(record.metadata)
|
||||
|
||||
candidate = self._extract_answer_with_llm(
|
||||
record.problem,
|
||||
model_answer,
|
||||
valid_letters,
|
||||
)
|
||||
# Parse straight from the output first (fast, deterministic, no "I" bug);
|
||||
# only fall back to the judge when no isolated letter is present.
|
||||
method = "direct"
|
||||
candidate = self._extract_answer_direct(model_answer, valid_letters)
|
||||
if not candidate:
|
||||
method = "llm"
|
||||
candidate = self._extract_answer_with_llm(
|
||||
record.problem,
|
||||
model_answer,
|
||||
valid_letters,
|
||||
)
|
||||
if not candidate:
|
||||
return None, {"reason": "no_choice_letter_extracted"}
|
||||
|
||||
@@ -99,6 +127,7 @@ class SuperGPQAScorer(LLMJudgeScorer):
|
||||
"reference_letter": ref,
|
||||
"candidate_letter": candidate,
|
||||
"valid_letters": valid_letters,
|
||||
"extract_method": method,
|
||||
}
|
||||
return is_correct, meta
|
||||
|
||||
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
# Orchestrator: gemma-4-12B-it. v1 = base self-sampling (the served orchestrator_model is this same base).
|
||||
#
|
||||
# v1 data = base gemma-4-12B-it self-sampling: roll the base model out over
|
||||
# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the
|
||||
# dataset first (no GPU, needs the vLLM orchestrator endpoint up):
|
||||
# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
# --out data/orchestrator_gemma-4-12b_sft_v1.jsonl --samples-per-task 8
|
||||
# then train on a rented GPU with regenerate_traces = false.
|
||||
|
||||
[model]
|
||||
model_name = "google/gemma-4-12B-it"
|
||||
max_seq_length = 4096
|
||||
|
||||
[training]
|
||||
num_epochs = 3
|
||||
batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch
|
||||
learning_rate = 2e-5
|
||||
weight_decay = 0.01
|
||||
warmup_ratio = 0.1
|
||||
gradient_checkpointing = true
|
||||
|
||||
[data]
|
||||
trace_cache_path = "data/orchestrator_gemma-4-12b_sft_v1.jsonl"
|
||||
regenerate_traces = false
|
||||
# Base gemma-4-12B-it self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
|
||||
orchestrator_endpoint = "http://localhost:8001/v1"
|
||||
orchestrator_model = "gemma-4-12b"
|
||||
samples_per_task = 8
|
||||
max_keep_per_task = 1
|
||||
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
|
||||
|
||||
[checkpoint]
|
||||
checkpoint_dir = "checkpoints/orchestrator_gemma-4-12b_sft"
|
||||
save_every_n_epochs = 1
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
# Orchestrator SFT cold-start — Qwen3-8B (~9B, "easier to train").
|
||||
#
|
||||
# v1 data = base Qwen3-8B self-sampling: roll the base model out over
|
||||
# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the
|
||||
# dataset first (no GPU, needs the vLLM orchestrator endpoint up):
|
||||
# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
# --out data/orchestrator_sft_v1.jsonl --samples-per-task 8
|
||||
# then train on a rented GPU with regenerate_traces = false.
|
||||
|
||||
[model]
|
||||
model_name = "Qwen/Qwen3-8B"
|
||||
max_seq_length = 4096
|
||||
|
||||
[training]
|
||||
num_epochs = 3
|
||||
batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch
|
||||
learning_rate = 2e-5
|
||||
weight_decay = 0.01
|
||||
warmup_ratio = 0.1
|
||||
gradient_checkpointing = true
|
||||
|
||||
[data]
|
||||
trace_cache_path = "data/orchestrator_sft_v1.jsonl"
|
||||
regenerate_traces = false
|
||||
# Base Qwen3-8B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
|
||||
orchestrator_endpoint = "http://localhost:8001/v1"
|
||||
orchestrator_model = "qwen3-8b"
|
||||
samples_per_task = 8
|
||||
max_keep_per_task = 1
|
||||
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
|
||||
|
||||
[checkpoint]
|
||||
checkpoint_dir = "checkpoints/orchestrator_sft"
|
||||
save_every_n_epochs = 1
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
# Orchestrator: Qwen3.5-9B. v1 = base self-sampling (the served orchestrator_model is this same base).
|
||||
#
|
||||
# v1 data = base Qwen3.5-9B self-sampling: roll the base model out over
|
||||
# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the
|
||||
# dataset first (no GPU, needs the vLLM orchestrator endpoint up):
|
||||
# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
|
||||
# --out data/orchestrator_qwen3.5-9b_sft_v1.jsonl --samples-per-task 8
|
||||
# then train on a rented GPU with regenerate_traces = false.
|
||||
|
||||
[model]
|
||||
model_name = "Qwen/Qwen3.5-9B"
|
||||
max_seq_length = 4096
|
||||
|
||||
[training]
|
||||
num_epochs = 3
|
||||
batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch
|
||||
learning_rate = 2e-5
|
||||
weight_decay = 0.01
|
||||
warmup_ratio = 0.1
|
||||
gradient_checkpointing = true
|
||||
|
||||
[data]
|
||||
trace_cache_path = "data/orchestrator_qwen3.5-9b_sft_v1.jsonl"
|
||||
regenerate_traces = false
|
||||
# Base Qwen3.5-9B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py).
|
||||
orchestrator_endpoint = "http://localhost:8001/v1"
|
||||
orchestrator_model = "qwen3.5-9b"
|
||||
samples_per_task = 8
|
||||
max_keep_per_task = 1
|
||||
distill_max_tasks = 0 # 0 -> all of load_sft_tasks()
|
||||
|
||||
[checkpoint]
|
||||
checkpoint_dir = "checkpoints/orchestrator_qwen3.5-9b_sft"
|
||||
save_every_n_epochs = 1
|
||||
@@ -0,0 +1,238 @@
|
||||
"""Eval backend that drives the ToolOrchestra orchestrator as the "model".
|
||||
|
||||
This adapts our unified-tool orchestrator rollout
|
||||
(:func:`openjarvis.agents.hybrid.toolorchestra.rollout.run_unified_rollout`)
|
||||
to the eval framework's :class:`~openjarvis.evals.core.backend.InferenceBackend`
|
||||
interface so it can be scored by the existing
|
||||
:class:`~openjarvis.evals.core.runner.EvalRunner` over any registered benchmark.
|
||||
|
||||
The orchestrator is an OpenAI-compatible chat model (served locally via vLLM,
|
||||
default ``http://localhost:8001/v1``) that emits tool calls over the fixed
|
||||
8-tool catalog from :func:`expert_registry.orchestrator_catalog`. For each eval
|
||||
sample we run one full rollout and return ``rollout.final_answer`` as the model
|
||||
answer, plus token/cost telemetry in the ``generate_full`` payload.
|
||||
|
||||
Construction is import-safe and network-free; the OpenAI SDK is only touched
|
||||
inside ``call_orchestrator`` at generate time (lazy import in ``unified.py``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra import rollout as _rollout_mod
|
||||
from openjarvis.agents.hybrid.toolorchestra.unified import (
|
||||
make_call_orchestrator,
|
||||
make_dispatch,
|
||||
)
|
||||
from openjarvis.evals.core.backend import InferenceBackend
|
||||
|
||||
# The served fine-tuned model sometimes over-emits the answer marker, e.g.
|
||||
# "FINAL_ANSWER: FINAL_ANSWER: 42" — a doubled prefix that breaks answer
|
||||
# extraction and auto-scores the sample 0. Collapse any run of FINAL_ANSWER
|
||||
# markers to a single one, keeping the text after the LAST marker as the answer.
|
||||
_FA_MARK = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?")
|
||||
|
||||
|
||||
def _clean_final_answer(text: str) -> str:
|
||||
t = (text or "").strip()
|
||||
marks = list(_FA_MARK.finditer(t))
|
||||
if not marks:
|
||||
return t
|
||||
answer = t[marks[-1].end() :].strip()
|
||||
return f"FINAL_ANSWER: {answer}" if answer else t
|
||||
|
||||
|
||||
_OBS_CAP = 4000 # cap persisted observations so the eval JSONL stays readable
|
||||
|
||||
DEFAULT_ENDPOINT = "http://localhost:8001/v1"
|
||||
DEFAULT_MODEL = "qwen3-8b"
|
||||
|
||||
|
||||
class OrchestratorBackend(InferenceBackend):
|
||||
"""Run the ToolOrchestra orchestrator rollout as an eval "model".
|
||||
|
||||
Parameters
|
||||
----------
|
||||
orchestrator_endpoint:
|
||||
OpenAI-compatible base URL for the served orchestrator (vLLM). Defaults
|
||||
to ``http://localhost:8001/v1``.
|
||||
orchestrator_model:
|
||||
Served model id for the orchestrator. Defaults to ``"qwen3-8b"`` (will
|
||||
be the Qwen3.5-9B checkpoint later).
|
||||
api_key:
|
||||
API key for the endpoint. ``"EMPTY"`` for a local vLLM server.
|
||||
local_endpoints:
|
||||
Optional dict mapping a local OSS model id (e.g. ``"Qwen/Qwen3.5-9B"``)
|
||||
to its vLLM base URL. Unmapped local models are still listed in the
|
||||
catalog but served as unconfigured (``base_url=None``).
|
||||
max_turns:
|
||||
Maximum orchestrator reasoning->action->observation turns per sample.
|
||||
temperature:
|
||||
Orchestrator sampling temperature.
|
||||
dispatch_cfg:
|
||||
Optional config dict passed to :func:`make_dispatch` (worker dispatch).
|
||||
"""
|
||||
|
||||
backend_id = "orchestrator"
|
||||
framework_name = "openjarvis-orchestrator"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
orchestrator_endpoint: str = DEFAULT_ENDPOINT,
|
||||
orchestrator_model: str = DEFAULT_MODEL,
|
||||
api_key: str = "EMPTY",
|
||||
local_endpoints: Optional[Dict[str, str]] = None,
|
||||
max_turns: int = 8,
|
||||
temperature: float = 1.0,
|
||||
dispatch_cfg: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
self.orchestrator_endpoint = orchestrator_endpoint
|
||||
self.orchestrator_model = orchestrator_model
|
||||
self.api_key = api_key
|
||||
self.local_endpoints = dict(local_endpoints or {})
|
||||
self.max_turns = int(max_turns)
|
||||
self.temperature = float(temperature)
|
||||
self._dispatch_cfg = dict(dispatch_cfg or {})
|
||||
|
||||
# Build the catalog once; it is stateless. ``local_endpoints`` now maps
|
||||
# full local model ids (e.g. ``"Qwen/Qwen3.5-9B"``) -> vLLM base_url.
|
||||
self._tools = orchestrator_catalog(local_endpoints=self.local_endpoints)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# InferenceBackend abstract methods
|
||||
# ------------------------------------------------------------------
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
model: str,
|
||||
system: str = "",
|
||||
temperature: float = 0.0,
|
||||
max_tokens: int = 2048,
|
||||
) -> str:
|
||||
"""Run a rollout and return just the final-answer text."""
|
||||
full = self.generate_full(
|
||||
prompt,
|
||||
model=model,
|
||||
system=system,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
return full.get("content", "") or ""
|
||||
|
||||
def generate_full(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
model: str,
|
||||
system: str = "",
|
||||
temperature: float = 0.0,
|
||||
max_tokens: int = 2048,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run one orchestrator rollout over ``prompt`` and return full details.
|
||||
|
||||
Returns a dict with the keys the runner reads: ``content``, ``usage``,
|
||||
``model``, ``latency_seconds``, ``cost_usd``, plus ``tool_calls`` /
|
||||
``turn_count`` / ``framework``. On error, returns ``content=""`` and a
|
||||
non-empty ``error`` field so the runner records it as a failed sample
|
||||
rather than crashing the whole run.
|
||||
"""
|
||||
# The caller-supplied ``model`` is the RunConfig.model (the orchestrator
|
||||
# model name). Fall back to the configured one if blank.
|
||||
orch_model = model or self.orchestrator_model
|
||||
# Prefer the explicit per-call temperature only when the caller set a
|
||||
# non-default value; otherwise use the backend's configured temperature.
|
||||
temp = temperature if temperature else self.temperature
|
||||
|
||||
started = time.time()
|
||||
try:
|
||||
call_orch = make_call_orchestrator(
|
||||
orch_model,
|
||||
base_url=self.orchestrator_endpoint,
|
||||
api_key=self.api_key,
|
||||
temperature=temp,
|
||||
max_tokens=max_tokens,
|
||||
# Serving a fine-tuned model: JSON <tool_call> format in the system
|
||||
# prompt, no tools= param (matches how it was trained/serialized).
|
||||
native_tools=False,
|
||||
)
|
||||
dispatch = make_dispatch(self._dispatch_cfg)
|
||||
rollout = _rollout_mod.run_unified_rollout(
|
||||
prompt,
|
||||
self._tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=self.max_turns,
|
||||
# Match TRAINING: the model was fine-tuned on anonymized expert
|
||||
# labels (expert_xxxx), so it only picks valid labels / routes
|
||||
# correctly when it sees the same anonymized names at inference.
|
||||
anonymize=True,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - surface as a failed sample
|
||||
return {
|
||||
"content": "",
|
||||
"error": f"{type(exc).__name__}: {exc}",
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0},
|
||||
"model": orch_model,
|
||||
"latency_seconds": time.time() - started,
|
||||
"cost_usd": 0.0,
|
||||
"tool_calls": 0,
|
||||
"turn_count": 0,
|
||||
"framework": self.framework_name,
|
||||
}
|
||||
|
||||
latency = time.time() - started
|
||||
total_tokens = int(getattr(rollout, "tokens", 0) or 0)
|
||||
anon_map = getattr(rollout, "anon_map", None) or {}
|
||||
return {
|
||||
"content": _clean_final_answer(rollout.final_answer or ""),
|
||||
# The rollout reports a single combined token count; we surface it
|
||||
# as completion_tokens so it still flows into total-token telemetry.
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": total_tokens,
|
||||
},
|
||||
"model": orch_model,
|
||||
"latency_seconds": latency,
|
||||
"cost_usd": float(getattr(rollout, "cost_usd", 0.0) or 0.0),
|
||||
"tool_calls": int(getattr(rollout, "num_tool_calls", 0) or 0),
|
||||
"turn_count": len(getattr(rollout, "turns", []) or []),
|
||||
"framework": self.framework_name,
|
||||
"trace_data": {
|
||||
"parse_failures": int(getattr(rollout, "parse_failures", 0) or 0),
|
||||
"tool_calls": [
|
||||
{"name": n, "arguments": a} for n, a in rollout.tool_calls()
|
||||
],
|
||||
# Full step-by-step trace so eval samples are inspectable
|
||||
# (format_eval_sample.py renders this). Observations capped.
|
||||
"anon_map": anon_map,
|
||||
"turns": [
|
||||
{
|
||||
"reasoning": t.reasoning or "",
|
||||
"tool_name": t.tool_name,
|
||||
"real_model": anon_map.get(t.tool_name)
|
||||
if t.tool_name
|
||||
else None,
|
||||
"arguments": t.arguments,
|
||||
"observation": (
|
||||
(t.observation or "")[:_OBS_CAP]
|
||||
+ (
|
||||
"…[truncated]"
|
||||
if t.observation and len(t.observation) > _OBS_CAP
|
||||
else ""
|
||||
)
|
||||
),
|
||||
}
|
||||
for t in (getattr(rollout, "turns", []) or [])
|
||||
],
|
||||
"final_answer": _clean_final_answer(rollout.final_answer or ""),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["OrchestratorBackend", "DEFAULT_ENDPOINT", "DEFAULT_MODEL"]
|
||||
@@ -143,6 +143,61 @@ class MultiObjectiveReward:
|
||||
return [self.compute(ep) for ep in episodes]
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostAwareReward:
|
||||
"""Cost-aware GRPO reward: ``r = +1/-1`` for correctness, minus a
|
||||
cost penalty.
|
||||
|
||||
This is the headline objective from the plan::
|
||||
|
||||
base = +1.0 if episode.correct else -1.0
|
||||
reward = base - lambda * (total_cost_usd / cost_max)
|
||||
|
||||
``lambda`` (``lam``) is the cost-penalty weight and is *swept* across a
|
||||
set of values (see :func:`lambda_sweep`) to trace the accuracy/cost
|
||||
Pareto frontier. ``cost_max`` normalises USD cost so the penalty is
|
||||
O(1) when an episode hits the per-task cost budget.
|
||||
"""
|
||||
|
||||
lam: float = 0.0
|
||||
"""Cost-penalty weight (swept, e.g. [0.0, 0.05, 0.1, 0.2, 0.4])."""
|
||||
|
||||
cost_max: float = 0.10
|
||||
"""USD cost normalizer (per-task cost budget)."""
|
||||
|
||||
def compute(self, episode: Episode) -> float:
|
||||
"""Scalar reward: ``+1/-1`` for correctness minus a cost term."""
|
||||
base = 1.0 if episode.correct else -1.0
|
||||
return base - self.lam * (episode.total_cost_usd / self.cost_max)
|
||||
|
||||
def compute_with_breakdown(self, episode: Episode) -> Dict[str, float]:
|
||||
"""Reward with per-component breakdown."""
|
||||
correct = 1.0 if episode.correct else 0.0
|
||||
base = 1.0 if episode.correct else -1.0
|
||||
cost_term = -self.lam * (episode.total_cost_usd / self.cost_max)
|
||||
return {
|
||||
"correct": correct,
|
||||
"base": base,
|
||||
"cost_term": cost_term,
|
||||
"reward": base + cost_term,
|
||||
}
|
||||
|
||||
def compute_batch(self, episodes: List[Episode]) -> List[float]:
|
||||
"""Compute rewards for a batch of episodes."""
|
||||
return [self.compute(ep) for ep in episodes]
|
||||
|
||||
|
||||
def lambda_sweep(
|
||||
values: List[float], cost_max: float = 0.10
|
||||
) -> List["CostAwareReward"]:
|
||||
"""Build one :class:`CostAwareReward` per lambda value.
|
||||
|
||||
Used to sweep the cost-penalty weight and trace the accuracy/cost
|
||||
Pareto frontier (the ``R -= lambda * cost / cost_max`` objective).
|
||||
"""
|
||||
return [CostAwareReward(lam=lam, cost_max=cost_max) for lam in values]
|
||||
|
||||
|
||||
class AdaptiveRewardWeights:
|
||||
"""Adaptive reward weights that shift during training.
|
||||
|
||||
@@ -211,7 +266,9 @@ class AdaptiveRewardWeights:
|
||||
|
||||
__all__ = [
|
||||
"AdaptiveRewardWeights",
|
||||
"CostAwareReward",
|
||||
"MultiObjectiveReward",
|
||||
"Normalizers",
|
||||
"RewardWeights",
|
||||
"lambda_sweep",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
"""SFT-data generation for the orchestrator cold-start.
|
||||
|
||||
Execution-grounded rejection sampling in the faithful ToolOrchestra action
|
||||
space (arXiv:2511.21689): for each ``nvidia/ToolScale`` task, roll out a teacher
|
||||
orchestrator over the unified tool catalog N times, verify each trajectory, keep
|
||||
the cheapest passing one(s), and serialize them into the ``<tool_call>``
|
||||
``conversations`` JSONL that
|
||||
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
|
||||
loads directly.
|
||||
|
||||
Pipeline::
|
||||
|
||||
ToolScale task -> N teacher rollouts -> verify -> keep cheapest passing
|
||||
-> conversations JSONL
|
||||
|
||||
(The earlier ADP-relabel cold-start was a heuristic re-tiering of demonstrated
|
||||
traces; it was removed in favour of this grounded pipeline.)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
gold_coverage_verify,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
|
||||
GoldAction,
|
||||
ToolScaleTask,
|
||||
load_toolscale,
|
||||
normalize_row,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"GoldAction",
|
||||
"ToolScaleTask",
|
||||
"generate_sft_dataset",
|
||||
"gold_coverage_verify",
|
||||
"load_toolscale",
|
||||
"normalize_row",
|
||||
"trajectory_to_record",
|
||||
]
|
||||
@@ -0,0 +1,355 @@
|
||||
"""Loaders for the orchestrator reasoning-SFT task sources.
|
||||
|
||||
Two HuggingFace reasoning datasets give us verifiable question/answer tasks for
|
||||
the small orchestrator's cold-start SFT and the GRPO prompt pool:
|
||||
|
||||
- ``natolambert/GeneralThought-430K-filtered`` — open reasoning traces scraped
|
||||
from gr.inc. Each row carries a ``question``, a short ``reference_answer``
|
||||
(the gold), a long ``model_answer`` (R1's full solution), and provenance in
|
||||
``question_source`` / ``task``. There is **no** clean single "category" field;
|
||||
we derive a coarse ``domain`` (math / medical / code / chat / misc) from
|
||||
``question_source`` so the verifier can pick the right checker.
|
||||
|
||||
- ``open-thoughts/OpenThoughts3-1.2M`` — OpenThoughts3 distillation set. Each
|
||||
row has a ``domain`` in {code, math, science}, a ``source``, a ``difficulty``,
|
||||
and a ``conversations`` list ``[{from: human, value}, {from: gpt, value}]``.
|
||||
The human turn is the question; the gpt turn is a ``<think>…</think>`` trace
|
||||
followed by the final solution, from which we extract the gold answer.
|
||||
|
||||
Mirrors the style of the sibling ``hotpotqa.py`` / ``toolscale.py`` loaders: a
|
||||
plain dataclass plus loaders that accept a ``source=`` iterable override so the
|
||||
normalization path is exercised offline with no network. Network imports of
|
||||
``datasets`` are kept lazy (inside the function).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Iterable, Iterator, List, Optional
|
||||
|
||||
GENERALTHOUGHT_ID = "natolambert/GeneralThought-430K-filtered"
|
||||
OPENTHOUGHTS_ID = "open-thoughts/OpenThoughts3-1.2M"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
task_id: str
|
||||
question: str
|
||||
answer: str
|
||||
domain: str # coarse area: math / code / science / medical / chat / misc
|
||||
difficulty: str = "" # OpenThoughts difficulty tier (GeneralThought has none)
|
||||
dataset: str = "" # source dataset: GeneralThought / OpenThoughts3
|
||||
subsector: str = "" # fine source: NuminaMath / NHSQA / TACO / glaive / ...
|
||||
|
||||
@property
|
||||
def instruction(self) -> str:
|
||||
return self.question
|
||||
|
||||
|
||||
# --- GeneralThought -----------------------------------------------------------
|
||||
|
||||
# question_source -> coarse domain. GeneralThought has no literal category field;
|
||||
# the source string is the most reliable signal (NuminaMath is math, NHSQA is
|
||||
# medical, TACO/glaive are code, oasst1/lmsys are open chat, everything else misc).
|
||||
_SOURCE_DOMAIN = (
|
||||
("numina", "math"),
|
||||
("math", "math"),
|
||||
("nhsqa", "medical"),
|
||||
("medical", "medical"),
|
||||
("medicine", "medical"),
|
||||
("taco", "code"),
|
||||
("code", "code"),
|
||||
("glaive", "code"),
|
||||
("oasst", "chat"),
|
||||
("lmsys", "chat"),
|
||||
("chat", "chat"),
|
||||
)
|
||||
|
||||
|
||||
def _generalthought_domain(row: Dict[str, Any]) -> str:
|
||||
src = str(row.get("question_source") or "").lower()
|
||||
task = str(row.get("task") or "").lower()
|
||||
hay = f"{src} {task}"
|
||||
for needle, dom in _SOURCE_DOMAIN:
|
||||
if needle in hay:
|
||||
return dom
|
||||
return "misc"
|
||||
|
||||
|
||||
def _normalize_generalthought(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
|
||||
question = str(row.get("question") or "").strip()
|
||||
# Prefer the short reference answer (exact-match-able); fall back to the
|
||||
# long model_answer when no reference is present.
|
||||
answer = str(row.get("reference_answer") or "").strip()
|
||||
if not answer:
|
||||
answer = str(row.get("model_answer") or "").strip()
|
||||
if not question or not answer:
|
||||
return None
|
||||
task_id = str(row.get("question_id") or f"generalthought-{index}")
|
||||
return Task(
|
||||
task_id=task_id,
|
||||
question=question,
|
||||
answer=answer,
|
||||
domain=_generalthought_domain(row),
|
||||
difficulty=str(row.get("difficulty") or "").strip(), # usually absent for GT
|
||||
dataset="GeneralThought",
|
||||
subsector=str(row.get("question_source") or "").strip(),
|
||||
)
|
||||
|
||||
|
||||
def load_generalthought(
|
||||
*,
|
||||
n: int = 2000,
|
||||
seed: int = 42,
|
||||
source: Optional[Iterable[Dict[str, Any]]] = None,
|
||||
buffer: Optional[int] = None,
|
||||
) -> Iterator[Task]:
|
||||
"""Yield up to ``n`` GeneralThought tasks, randomly mixed across categories.
|
||||
|
||||
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
|
||||
We over-stream a buffer and shuffle it so the yielded tasks are a random mix
|
||||
of the source's categories rather than a single leading shard. ``buffer``
|
||||
overrides the shuffle-buffer size; pass a small value (e.g. for smoke runs)
|
||||
to avoid streaming the default 6000-row floor.
|
||||
"""
|
||||
if source is None:
|
||||
from datasets import load_dataset # lazy: optional dep / network
|
||||
|
||||
source = load_dataset(GENERALTHOUGHT_ID, split="train", streaming=True)
|
||||
|
||||
rng = random.Random(seed)
|
||||
# Buffer a generous multiple of n so the shuffle actually mixes categories,
|
||||
# but stay bounded for the multi-hundred-K streams.
|
||||
buf_cap = buffer if buffer is not None else max(n * 6, 6000)
|
||||
buf: List[Task] = []
|
||||
for i, row in enumerate(source):
|
||||
task = _normalize_generalthought(dict(row), index=i)
|
||||
if task is None:
|
||||
continue
|
||||
buf.append(task)
|
||||
if len(buf) >= buf_cap:
|
||||
break
|
||||
rng.shuffle(buf)
|
||||
for task in buf[:n]:
|
||||
yield task
|
||||
|
||||
|
||||
# --- OpenThoughts3 ------------------------------------------------------------
|
||||
|
||||
_BOXED_RE = re.compile(r"\\boxed\{")
|
||||
|
||||
|
||||
def _extract_boxed(text: str) -> Optional[str]:
|
||||
"""Return the contents of the last ``\\boxed{...}`` in ``text`` (brace-balanced)."""
|
||||
last = None
|
||||
for m in _BOXED_RE.finditer(text):
|
||||
start = m.end()
|
||||
depth = 1
|
||||
i = start
|
||||
while i < len(text) and depth > 0:
|
||||
c = text[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
i += 1
|
||||
if depth == 0:
|
||||
last = text[start : i - 1].strip()
|
||||
return last
|
||||
|
||||
|
||||
def _conversation_qa(conversations: Any) -> Optional[tuple[str, str]]:
|
||||
"""Pull (question, gold_answer) from an OpenThoughts ``conversations`` list."""
|
||||
if not isinstance(conversations, list):
|
||||
return None
|
||||
question = ""
|
||||
full = ""
|
||||
for turn in conversations:
|
||||
if not isinstance(turn, dict):
|
||||
continue
|
||||
who = str(turn.get("from") or "").lower()
|
||||
val = str(turn.get("value") or "")
|
||||
if who in ("human", "user") and not question:
|
||||
question = val.strip()
|
||||
elif who in ("gpt", "assistant"):
|
||||
full = val
|
||||
if not question or not full:
|
||||
return None
|
||||
# Drop the <think>…</think> trace; keep the post-reasoning solution.
|
||||
visible = re.sub(r"(?s)<think>.*?</think>", "", full).strip()
|
||||
body = visible or full
|
||||
answer = _extract_boxed(body) or _extract_boxed(full) or body.strip()
|
||||
if not answer:
|
||||
return None
|
||||
return question, answer
|
||||
|
||||
|
||||
def _normalize_openthoughts(row: Dict[str, Any], *, index: int = 0) -> Optional[Task]:
|
||||
qa = _conversation_qa(row.get("conversations"))
|
||||
if qa is None:
|
||||
return None
|
||||
question, answer = qa
|
||||
if not question or not answer:
|
||||
return None
|
||||
domain = str(row.get("domain") or "unknown").strip().lower() or "unknown"
|
||||
task_id = (
|
||||
str(row.get("id") or row.get("source") or f"openthoughts-{index}") + f"-{index}"
|
||||
)
|
||||
return Task(
|
||||
task_id=task_id,
|
||||
question=question,
|
||||
answer=answer,
|
||||
domain=domain,
|
||||
difficulty=str(row.get("difficulty") or "").strip(),
|
||||
dataset="OpenThoughts3",
|
||||
subsector=str(row.get("source") or "").strip(),
|
||||
)
|
||||
|
||||
|
||||
# The 1.2M set ships as 120 uniform parquet shards with no domain in the file
|
||||
# names, but the rows are front-ordered by domain. Probing one row per shard puts
|
||||
# code in shards ~0-30, math ~32-104, science ~105-119. Streaming the single
|
||||
# combined split therefore has to read the entire code (+ math) prefix before it
|
||||
# reaches math/science — minutes of wasted I/O. Instead we stream each domain from
|
||||
# shards *inside* its own region. A few shards (~10K rows each) cover any quota.
|
||||
_OT_SHARD_FMT = "data/train-{:05d}-of-00120.parquet"
|
||||
_OT_DOMAIN_SHARDS: Dict[str, List[int]] = {
|
||||
"code": [0, 1, 2, 3],
|
||||
"math": [45, 46, 47, 48],
|
||||
"science": [112, 113, 114, 115, 116, 117, 118, 119],
|
||||
}
|
||||
|
||||
|
||||
def load_openthoughts(
|
||||
*,
|
||||
n_code: int = 2000,
|
||||
n_math: int = 2000,
|
||||
n_science: int = 2000,
|
||||
seed: int = 42,
|
||||
source: Optional[Iterable[Dict[str, Any]]] = None,
|
||||
) -> Iterator[Task]:
|
||||
"""Yield OpenThoughts3 tasks balanced across code / math / science.
|
||||
|
||||
``source`` overrides the stream with raw row dicts (tests): rows are routed
|
||||
into per-domain buffers, stopping once every quota is met. With no override we
|
||||
stream each domain from its own shard region (see ``_OT_DOMAIN_SHARDS``) so a
|
||||
balanced pull never has to read the entire front-ordered code/math prefix.
|
||||
"""
|
||||
quotas = {"code": n_code, "math": n_math, "science": n_science}
|
||||
bufs: Dict[str, List[Task]] = {k: [] for k in quotas}
|
||||
|
||||
if source is not None:
|
||||
for i, row in enumerate(source):
|
||||
task = _normalize_openthoughts(dict(row), index=i)
|
||||
if task is None:
|
||||
continue
|
||||
dom = task.domain
|
||||
if dom in bufs and len(bufs[dom]) < quotas[dom]:
|
||||
bufs[dom].append(task)
|
||||
if all(len(bufs[k]) >= quotas[k] for k in quotas):
|
||||
break
|
||||
else:
|
||||
from datasets import load_dataset # lazy: optional dep / network
|
||||
|
||||
idx = 0
|
||||
for dom, quota in quotas.items():
|
||||
if quota <= 0:
|
||||
continue
|
||||
files = [_OT_SHARD_FMT.format(s) for s in _OT_DOMAIN_SHARDS[dom]]
|
||||
ds = load_dataset(
|
||||
OPENTHOUGHTS_ID, data_files=files, split="train", streaming=True
|
||||
)
|
||||
for row in ds:
|
||||
task = _normalize_openthoughts(dict(row), index=idx)
|
||||
idx += 1
|
||||
if task is None or task.domain != dom:
|
||||
continue
|
||||
bufs[dom].append(task)
|
||||
if len(bufs[dom]) >= quota:
|
||||
break
|
||||
|
||||
rng = random.Random(seed)
|
||||
out: List[Task] = []
|
||||
for k in quotas:
|
||||
out.extend(bufs[k])
|
||||
rng.shuffle(out)
|
||||
for task in out:
|
||||
yield task
|
||||
|
||||
|
||||
# --- combined SFT / GRPO sets -------------------------------------------------
|
||||
|
||||
|
||||
def load_sft_tasks(
|
||||
*, seed: int = 42, cap: Optional[int] = None, balanced: bool = True
|
||||
) -> List[Task]:
|
||||
"""The 8K cold-start SFT set: 2K GeneralThought + 2K code + 2K math + 2K science.
|
||||
|
||||
``cap`` (smoke runs): when set, draw ~``cap`` tasks total.
|
||||
- default (``balanced=False``): GeneralThought only with a small stream
|
||||
buffer — fastest, but the domain mix is whatever GeneralThought yields
|
||||
(math/code/medical/chat), so a given sample can skew (e.g. all-medical).
|
||||
- ``balanced=True``: ~cap/4 from each of GeneralThought + OpenThoughts
|
||||
code/math/science, so the smoke is representative of the real run. Now
|
||||
cheap because OpenThoughts streams from per-domain shards (no front-prefix).
|
||||
"""
|
||||
if cap is not None and cap > 0:
|
||||
if balanced:
|
||||
per = max(cap // 4, 1)
|
||||
tasks = list(load_generalthought(n=per, seed=seed, buffer=max(per * 6, 64)))
|
||||
tasks.extend(
|
||||
load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed)
|
||||
)
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(tasks)
|
||||
return tasks[:cap]
|
||||
buf = max(cap * 4, 64)
|
||||
tasks = list(load_generalthought(n=cap, seed=seed, buffer=buf))
|
||||
random.Random(seed).shuffle(tasks)
|
||||
return tasks[:cap]
|
||||
tasks: List[Task] = []
|
||||
tasks.extend(load_generalthought(n=2000, seed=seed))
|
||||
tasks.extend(load_openthoughts(n_code=2000, n_math=2000, n_science=2000, seed=seed))
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(tasks)
|
||||
return tasks
|
||||
|
||||
|
||||
def load_grpo_prompts(*, n: int = 30000, seed: int = 42) -> List[Task]:
|
||||
"""Pool ``n`` unique prompts from the same datasets, deduped by question.
|
||||
|
||||
We draw a roughly even split from GeneralThought and OpenThoughts3 (code /
|
||||
math / science), dedupe on the question text, and cap at ``n``.
|
||||
"""
|
||||
per = max(n // 4 + 1, 1)
|
||||
pool: List[Task] = []
|
||||
pool.extend(load_generalthought(n=per, seed=seed))
|
||||
pool.extend(load_openthoughts(n_code=per, n_math=per, n_science=per, seed=seed))
|
||||
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(pool)
|
||||
|
||||
seen: set[str] = set()
|
||||
out: List[Task] = []
|
||||
for task in pool:
|
||||
key = task.question.strip()
|
||||
if not key or key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
out.append(task)
|
||||
if len(out) >= n:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GENERALTHOUGHT_ID",
|
||||
"OPENTHOUGHTS_ID",
|
||||
"Task",
|
||||
"load_generalthought",
|
||||
"load_grpo_prompts",
|
||||
"load_openthoughts",
|
||||
"load_sft_tasks",
|
||||
]
|
||||
@@ -0,0 +1,460 @@
|
||||
"""Rejection-sampling SFT-data generator (the ToolOrchestra cold-start).
|
||||
|
||||
For each ToolScale task: roll out a teacher orchestrator N times, verify each
|
||||
trajectory, keep the passing ones (optionally just the cheapest), and serialize
|
||||
them into the unified-tool ``conversations`` JSONL the SFT trainer consumes.
|
||||
|
||||
The expensive/network parts are injected so the orchestration is pure and
|
||||
offline-testable:
|
||||
|
||||
* ``rollout_fn(task) -> UnifiedRollout`` — one teacher rollout (temperature>0).
|
||||
* ``verify_fn(task, rollout) -> bool`` — did the trajectory solve the task?
|
||||
|
||||
:func:`gold_coverage_verify` is a dependency-free default verifier (checks the
|
||||
trajectory's tool calls cover the task's golden action names); a real run should
|
||||
compose it with an LLM judge on the final answer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, List, Optional
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import ExpertTool
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
_CONTROL_TOKEN_RE,
|
||||
_strip_control_tokens,
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The sampler only touches ``task.task_id`` / ``task.instruction`` / ``task.domain``
|
||||
# (via ``trajectory_to_record``), so any task dataclass works — ToolScaleTask or the
|
||||
# reasoning ``datasets.Task``. ``gold_coverage_verify`` additionally needs
|
||||
# ``gold_action_names()`` (ToolScale-only), but callers using ``datasets.Task`` pass
|
||||
# their own ``verify_fn`` (e.g. ``verify.make_verifier()``) instead.
|
||||
TaskLike = Any
|
||||
RolloutFn = Callable[[TaskLike], Optional[UnifiedRollout]]
|
||||
VerifyFn = Callable[[TaskLike, UnifiedRollout], bool]
|
||||
|
||||
# Human-readable blurb for each source dataset, stamped onto every record as
|
||||
# ``dataset_description`` so a row is self-describing (what corpus the question
|
||||
# came from) without a lookup. Keyed on the record's ``dataset`` value.
|
||||
DATASET_DESCRIPTIONS = {
|
||||
"GeneralThought": (
|
||||
"natolambert/GeneralThought-430K-filtered — filtered reasoning Q&A from "
|
||||
"gr.inc; each item has a question and a short reference (gold) answer "
|
||||
"distilled from DeepSeek-R1 traces."
|
||||
),
|
||||
"OpenThoughts3": (
|
||||
"open-thoughts/OpenThoughts3-1.2M — code/math/science reasoning tasks; "
|
||||
"the gold answer is extracted from the boxed final solution."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def dataset_description(dataset: str) -> str:
|
||||
"""Blurb for a source dataset name (empty string if unknown)."""
|
||||
return DATASET_DESCRIPTIONS.get((dataset or "").strip(), "")
|
||||
|
||||
|
||||
# Markers of a broken expert/observation that must never become a training
|
||||
# target (empty turns, tracebacks, truncated tool errors, dead-provider strings).
|
||||
_ERR_MARKERS = (
|
||||
"Traceback (most recent call last)",
|
||||
"insufficient_quota",
|
||||
"rate limit",
|
||||
"RateLimitError",
|
||||
"[invalid tool",
|
||||
"Error code: 4",
|
||||
"Error code: 5",
|
||||
"InternalServerError",
|
||||
"ConnectionError",
|
||||
"timed out",
|
||||
# A tool call that reached dispatch with no/empty input, or a bridged-tool
|
||||
# crash — the call was malformed. The trajectory routed garbage; drop it.
|
||||
"no input provided",
|
||||
"[openjarvis-tool error",
|
||||
)
|
||||
|
||||
# A final answer cut off mid-expression ends on a dangling connective
|
||||
# (``a=1, b=``, ``result = ``, ``compute (``) — the decode hit the token cap
|
||||
# before finishing. A real distilled answer never ends this way.
|
||||
_TRUNCATED_TAIL_RE = re.compile(r"[=,(\[{+\-*/]\s*$")
|
||||
|
||||
|
||||
def _target_is_clean(roll: UnifiedRollout) -> bool:
|
||||
"""Structural gate on whether a trajectory is fit to be an SFT *target*.
|
||||
|
||||
The judge decides semantic correctness; this catches the ~10-12% of records
|
||||
whose targets are empty / a Traceback / a truncated tool error / unrouted —
|
||||
garbage the model must never be trained to imitate. Requirements:
|
||||
* non-empty final answer, no error marker, balanced ``<think>`` (not
|
||||
truncated mid-thought);
|
||||
* at least one model-expert tool call (it actually *routed*);
|
||||
* no error-marker / empty observation in the kept trajectory.
|
||||
"""
|
||||
fa = (roll.final_answer or "").strip()
|
||||
if not fa:
|
||||
return False
|
||||
if any(m in fa for m in _ERR_MARKERS):
|
||||
return False
|
||||
if fa.count("<think>") != fa.count(
|
||||
"</think>"
|
||||
): # unbalanced think tags (truncated OR stray </think>)
|
||||
return False
|
||||
if _TRUNCATED_TAIL_RE.search(fa): # truncated mid-expression (e.g. "a=1, b=")
|
||||
return False
|
||||
# Malformed final: a proper final answer is plain text, NOT a stray/broken
|
||||
# <tool_call> tag (the model's answer leaking inside a tool call that failed
|
||||
# to parse). (\boxed{} is NOT rejected here — the serializer de-boxes it, so
|
||||
# an otherwise-good boxed answer is salvaged rather than thrown away.)
|
||||
if "<tool_call>" in fa or "</tool_call>" in fa:
|
||||
return False
|
||||
# Control-token backstop (defense in depth). The serializer strips leaked
|
||||
# control/special tokens (<|im_end|>, <|tool_call>, <start_of_turn>, …) to
|
||||
# salvage good answers, so we mirror that strip and gate on the RESULT:
|
||||
# * strips to empty -> the answer WAS nothing but control tokens -> drop.
|
||||
# * a token survives the strip -> something malformed the stripper can't
|
||||
# safely delete mid-answer -> drop rather than train on it.
|
||||
# A clean answer (or one the stripper fully salvages) sails through.
|
||||
fa_stripped = _strip_control_tokens(fa)
|
||||
if not fa_stripped:
|
||||
return False
|
||||
if _CONTROL_TOKEN_RE.search(fa_stripped):
|
||||
return False
|
||||
# Degenerate repetition: the same substantive line emitted many times (the
|
||||
# small-model decode loop). Reject rather than train on it.
|
||||
_lines = [ln.strip() for ln in fa.split("\n") if len(ln.strip()) > 30]
|
||||
if _lines and Counter(_lines).most_common(1)[0][1] >= 4:
|
||||
return False
|
||||
# Word-salad run-on: a giant unbroken line (no newline) is the other decode
|
||||
# collapse — the model spraying novel tokens instead of repeating. Reject.
|
||||
if any(len(ln) > 2000 for ln in fa.split("\n")):
|
||||
return False
|
||||
# Markdown-table dump: the model wrote a formatted table/essay instead of the
|
||||
# short exact value the format demands. A ``|---|`` separator row is the tell.
|
||||
if re.search(r"\|\s*:?-{2,}", fa):
|
||||
return False
|
||||
# Essay-style final: the format wants a distilled answer, not a multi-section
|
||||
# writeup. Check the answer AFTER the FINAL_ANSWER marker (reasoning before it
|
||||
# is fine). NOTE: length/bold limits relaxed for the BALANCED/harder mix —
|
||||
# code & multi-step math answers are legitimately longer and use **bold** for
|
||||
# the key result, so the old 700-char + any-bold rejects were dropping ~half
|
||||
# the correct hard-task answers. Keep the STRUCTURAL essay signals (many
|
||||
# numbered sections, markdown headers, tables) which still catch real essays.
|
||||
_marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", fa))
|
||||
_ans = fa[_marks[-1].end() :].strip() if _marks else fa
|
||||
if len(_ans) > 2000: # was 700 — too tight for code/math
|
||||
return False
|
||||
# Tool-status echo as the final answer: on code tasks the orchestrator
|
||||
# sometimes ends with a file_write/shell status line ("Successfully wrote to
|
||||
# /tmp/x.py") instead of the actual answer — a non-answer that slips the
|
||||
# length/format checks. Reject. (Audit: ~code/math trajectories where the
|
||||
# trace collapsed but the row was still scored correct.)
|
||||
if re.search(
|
||||
r"(?i)\b(successfully (wrote|created|saved|executed|ran)|written to /|"
|
||||
r"file (written|saved|created)|no further actions?)\b",
|
||||
_ans,
|
||||
):
|
||||
return False
|
||||
if (
|
||||
len(re.findall(r"(?m)^\s*\d+\.\s", _ans)) >= 6
|
||||
): # many numbered sections = essay (was 4)
|
||||
return False
|
||||
if re.search(r"(?m)^\s*#{1,6}\s", _ans): # markdown headers = essay
|
||||
return False
|
||||
# Garbled / shouty final: a multi-word answer that's mostly UPPERCASE, or has
|
||||
# an absurdly long merged all-letter token, is decode garble — the audit found
|
||||
# e.g. "RABES PEPTETANUS BOOSTERS" (for "rabies PEP, tetanus boosters") passing.
|
||||
if len(_ans) > 20:
|
||||
_alpha = [c for c in _ans if c.isalpha()]
|
||||
if (
|
||||
_alpha
|
||||
and sum(c.isupper() for c in _alpha) / len(_alpha) > 0.7
|
||||
and len(_ans.split()) >= 3
|
||||
):
|
||||
return False
|
||||
if any(len(w) > 25 and w.isalpha() for w in _ans.split()):
|
||||
return False
|
||||
# Reject runaway final answers — a distilled answer, not a multi-KB essay or
|
||||
# a verbatim dump of an expert observation. (Audit: 279 finals >4k chars, the
|
||||
# worst a 409k-char wordlist; ~363 were prefix-identical to a tool obs.)
|
||||
if len(fa) > 8000:
|
||||
return False
|
||||
fa_head = re.sub(r"\s+", " ", fa[:200]).strip()
|
||||
for t in roll.turns:
|
||||
if t.observation:
|
||||
obs_head = re.sub(r"\s+", " ", t.observation[:200]).strip()
|
||||
if (
|
||||
fa_head and fa_head == obs_head and len(fa) > 200
|
||||
): # long final = copied tool dump (short relays are legit)
|
||||
return False
|
||||
# "routed" = delegated to a model EXPERT (not just a utility like web_search).
|
||||
# When anonymized, expert calls are the anon labels in anon_map; otherwise
|
||||
# fall back to "made any tool call".
|
||||
expert_names = set(roll.anon_map or {})
|
||||
called = {name for name, _ in roll.tool_calls()}
|
||||
if expert_names:
|
||||
if not (called & expert_names):
|
||||
return False
|
||||
elif roll.num_tool_calls < 1:
|
||||
return False
|
||||
for t in roll.turns:
|
||||
if t.tool_name is not None:
|
||||
obs = (t.observation or "").strip()
|
||||
if not obs or any(m in obs for m in _ERR_MARKERS):
|
||||
return False
|
||||
# Reasoning-side decode collapse: an intermediate turn whose reasoning is a
|
||||
# giant unbroken line, or is dominated by exotic unicode / emoji spam, is
|
||||
# garbage even when the final answer looks fine (the answer-only guards above
|
||||
# miss it). The audit found a trajectory with a thousands-char symbol blob in
|
||||
# its reasoning that scored correct and slipped through.
|
||||
for t in roll.turns:
|
||||
r = t.reasoning or ""
|
||||
if any(len(ln) > 2000 for ln in r.split("\n")):
|
||||
return False
|
||||
if len(r) > 200:
|
||||
weird = sum(1 for ch in r if ord(ch) > 0x2100 and not ch.isalnum())
|
||||
if weird / len(r) > 0.10: # >10% exotic-unicode/emoji = decode spam
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def gold_coverage_verify(task: TaskLike, rollout: UnifiedRollout) -> bool:
|
||||
"""Dependency-free proxy verifier: trajectory must (a) produce a non-empty
|
||||
answer and (b) call tools covering every golden action name.
|
||||
|
||||
This is the offline stand-in for ToolScale's execution-correctness checker
|
||||
(which needs the DB simulator). Compose with an LLM judge for real runs.
|
||||
"""
|
||||
if not rollout.final_answer.strip():
|
||||
return False
|
||||
gold = set(task.gold_action_names())
|
||||
if not gold:
|
||||
return True
|
||||
called = {name for name, _ in rollout.tool_calls()}
|
||||
return gold.issubset(called)
|
||||
|
||||
|
||||
def _process_task(
|
||||
task: TaskLike,
|
||||
*,
|
||||
rollout_fn: RolloutFn,
|
||||
verify_fn: VerifyFn,
|
||||
samples_per_task: int,
|
||||
max_keep_per_task: int,
|
||||
stop_at_keep: bool = False,
|
||||
keep_all: bool = False,
|
||||
) -> List[tuple[UnifiedRollout, bool]]:
|
||||
"""One task's rejection-sampling unit of work: roll out ``samples_per_task``
|
||||
times and verify each. Returns ``(rollout, is_correct)`` for **every** sample
|
||||
(the caller decides what to write). Runs in a worker thread, so the only
|
||||
shared state it touches is the injected fns (each rollout builds its own
|
||||
OpenAI client; the tool executor is built under a lock).
|
||||
|
||||
``stop_at_keep``: short-circuit as soon as ``max_keep_per_task`` passing
|
||||
trajectories are found, instead of always exhausting ``samples_per_task``.
|
||||
Big throughput win (no wasted rollouts on already-solved tasks), but it
|
||||
trades away the cheapest-of-N cost optimisation (we keep the first passers,
|
||||
not the cheapest). Ignored when ``keep_all`` is set (we want every sample).
|
||||
Default off preserves the original semantics."""
|
||||
results: List[tuple[UnifiedRollout, bool]] = []
|
||||
n_correct = 0
|
||||
for _ in range(samples_per_task):
|
||||
roll = rollout_fn(task)
|
||||
if roll is None:
|
||||
continue
|
||||
ok = bool(verify_fn(task, roll))
|
||||
results.append((roll, ok))
|
||||
if ok:
|
||||
n_correct += 1
|
||||
if stop_at_keep and not keep_all and n_correct >= max_keep_per_task:
|
||||
break
|
||||
return results
|
||||
|
||||
|
||||
def generate_sft_dataset(
|
||||
out_path: str,
|
||||
*,
|
||||
tasks: Iterable[TaskLike],
|
||||
tools: List[ExpertTool],
|
||||
rollout_fn: RolloutFn,
|
||||
verify_fn: VerifyFn = gold_coverage_verify,
|
||||
samples_per_task: int = 4,
|
||||
max_keep_per_task: int = 1,
|
||||
reward_fn: Optional[Callable[[UnifiedRollout], float]] = None,
|
||||
concurrency: int = 1,
|
||||
stop_at_keep: bool = False,
|
||||
keep_all: bool = False,
|
||||
record_extra: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""Run rejection sampling over ``tasks`` and write the SFT JSONL.
|
||||
|
||||
``max_keep_per_task`` caps records kept per task; when >1 the cheapest
|
||||
passing trajectories are kept first. ``concurrency`` rolls out that many tasks
|
||||
in parallel (each task issues its samples sequentially, so ~``concurrency``
|
||||
requests hit the served model at once) — set 1 for the original sequential
|
||||
behaviour. Records are written as each task finishes, so peak memory is bounded
|
||||
by the in-flight tasks, not the whole dataset. Returns stats + writes a
|
||||
``.stats.json``.
|
||||
|
||||
``keep_all``: write **every** rolled-out trajectory (correct and incorrect)
|
||||
rather than dropping the failures. Each record gets ``correct`` (verifier
|
||||
verdict) and ``kept`` (True on the cheapest-correct sample — the one the
|
||||
rejection sampler would have selected). This preserves the full sample set so
|
||||
you can compute accuracy / inspect failures; the SFT trainer should filter on
|
||||
``correct`` (or ``kept``). Default off = original drop-the-failures behaviour.
|
||||
"""
|
||||
out = Path(out_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
tasks = list(tasks)
|
||||
seen = len(tasks)
|
||||
written = 0
|
||||
correct_written = 0
|
||||
incorrect_written = 0
|
||||
dropped = 0 # tasks that produced no kept record
|
||||
tasks_solved = 0 # tasks with >=1 correct sample
|
||||
domain_counts: Counter[str] = Counter()
|
||||
|
||||
def _work(task: TaskLike) -> tuple[TaskLike, List[tuple[UnifiedRollout, bool]]]:
|
||||
return task, _process_task(
|
||||
task,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=verify_fn,
|
||||
samples_per_task=samples_per_task,
|
||||
max_keep_per_task=max_keep_per_task,
|
||||
stop_at_keep=stop_at_keep,
|
||||
keep_all=keep_all,
|
||||
)
|
||||
|
||||
def _emit(fh, task: TaskLike, roll: UnifiedRollout, ok: bool, kept: bool) -> None:
|
||||
nonlocal written, correct_written, incorrect_written
|
||||
reward = reward_fn(roll) if reward_fn else 0.0
|
||||
record = trajectory_to_record(
|
||||
task.task_id,
|
||||
task.instruction,
|
||||
tools,
|
||||
roll,
|
||||
reward=reward,
|
||||
domain=task.domain,
|
||||
)
|
||||
record["correct"] = ok
|
||||
record["kept"] = kept
|
||||
record["clean"] = _target_is_clean(roll)
|
||||
# Task provenance so every record is self-describing: area (domain),
|
||||
# difficulty tier, source dataset, and fine subsector. Lets us slice
|
||||
# train/holdout/analysis by any of these later.
|
||||
record["area"] = task.domain
|
||||
record["difficulty"] = getattr(task, "difficulty", "") or ""
|
||||
record["dataset"] = getattr(task, "dataset", "") or ""
|
||||
record["dataset_description"] = dataset_description(record["dataset"])
|
||||
record["subsector"] = getattr(task, "subsector", "") or ""
|
||||
# Gold reference answer the verifier graded against. Persisted so every
|
||||
# record supports gold-vs-model inspection downstream (Braintrust, error
|
||||
# analysis) instead of only a bare `correct` boolean.
|
||||
record["gold_answer"] = getattr(task, "answer", "") or ""
|
||||
# Stamp provenance (which model/orchestrator generated this trajectory)
|
||||
# so records are self-identifying when pooled across model families.
|
||||
if record_extra:
|
||||
record.update(record_extra)
|
||||
fh.write(json.dumps(record) + "\n")
|
||||
fh.flush()
|
||||
written += 1
|
||||
correct_written += int(ok)
|
||||
incorrect_written += int(not ok)
|
||||
domain_counts[task.domain] += 1
|
||||
|
||||
def _write(fh, task: TaskLike, results: List[tuple[UnifiedRollout, bool]]) -> None:
|
||||
nonlocal dropped, tasks_solved
|
||||
# A trajectory is only eligible to be *kept* as a training target if the
|
||||
# judge passed it AND it's structurally clean (non-error, routed, not
|
||||
# truncated). Unclean-but-correct rollouts are still written in keep_all
|
||||
# mode (for analysis) but never marked kept.
|
||||
correct = [r for r in results if r[1] and _target_is_clean(r[0])]
|
||||
cheapest = min((r for r, _ in correct), key=lambda r: r.cost_usd, default=None)
|
||||
if correct:
|
||||
tasks_solved += 1
|
||||
if keep_all:
|
||||
# Write every sample; mark the cheapest-correct one as kept.
|
||||
if not results:
|
||||
dropped += 1
|
||||
return
|
||||
for roll, ok in results:
|
||||
_emit(fh, task, roll, ok, kept=(roll is cheapest))
|
||||
else:
|
||||
# Original behaviour: keep only cheapest-correct, capped.
|
||||
if not correct:
|
||||
dropped += 1
|
||||
return
|
||||
for roll in sorted((r for r, _ in correct), key=lambda r: r.cost_usd)[
|
||||
:max_keep_per_task
|
||||
]:
|
||||
_emit(fh, task, roll, True, kept=(roll is cheapest))
|
||||
|
||||
with out.open("w") as fh:
|
||||
if concurrency <= 1:
|
||||
for task in tasks:
|
||||
_, results = _work(task)
|
||||
_write(fh, task, results)
|
||||
else:
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
done = 0
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
||||
futures = {pool.submit(_work, t): t for t in tasks}
|
||||
for fut in as_completed(futures):
|
||||
task = futures[fut]
|
||||
try:
|
||||
_, results = fut.result()
|
||||
except Exception as exc: # a task's worker died; count + skip
|
||||
logger.warning(
|
||||
"task %s failed: %s", getattr(task, "task_id", "?"), exc
|
||||
)
|
||||
results = []
|
||||
_write(fh, task, results)
|
||||
done += 1
|
||||
if done % 50 == 0 or done == seen:
|
||||
logger.info(
|
||||
"rejection-sampling: %d/%d tasks done "
|
||||
"(%d written, %d solved, %d dropped)",
|
||||
done,
|
||||
seen,
|
||||
written,
|
||||
tasks_solved,
|
||||
dropped,
|
||||
)
|
||||
|
||||
stats = {
|
||||
"out_path": str(out),
|
||||
"tasks_seen": seen,
|
||||
"records_written": written,
|
||||
"records_correct": correct_written,
|
||||
"records_incorrect": incorrect_written,
|
||||
"tasks_solved": tasks_solved,
|
||||
"tasks_dropped": dropped,
|
||||
"task_accuracy": round(tasks_solved / seen, 4) if seen else 0.0,
|
||||
"samples_per_task": samples_per_task,
|
||||
"keep_all": keep_all,
|
||||
"concurrency": concurrency,
|
||||
"domain_distribution": dict(domain_counts),
|
||||
}
|
||||
out.with_suffix(out.suffix + ".stats.json").write_text(json.dumps(stats, indent=2))
|
||||
logger.info(
|
||||
"Wrote %d SFT records to %s (%d correct, %d incorrect)",
|
||||
written,
|
||||
out,
|
||||
correct_written,
|
||||
incorrect_written,
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
__all__ = ["generate_sft_dataset", "gold_coverage_verify"]
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Loader for NVIDIA ToolScale (``nvidia/ToolScale``) — the ToolOrchestra
|
||||
RL/SFT task source (arXiv:2511.21689 §3.3).
|
||||
|
||||
Each row is a synthetic user-agent-tool task: an instruction ``I``, golden
|
||||
function calls ``A`` (the ground-truth tool sequence), and short info ``o`` that
|
||||
must be communicated. We normalize the raw HF row into :class:`ToolScaleTask`.
|
||||
|
||||
``load_toolscale`` streams via the HuggingFace ``datasets`` library; tests pass
|
||||
``source=`` an iterable of raw row dicts so normalization is exercised offline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Iterable, Iterator, List, Optional
|
||||
|
||||
DATASET_ID = "nvidia/ToolScale"
|
||||
|
||||
|
||||
@dataclass
|
||||
class GoldAction:
|
||||
name: str
|
||||
arguments: Dict[str, Any] = field(default_factory=dict)
|
||||
action_id: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolScaleTask:
|
||||
task_id: str
|
||||
domain: str
|
||||
instruction: str
|
||||
gold_actions: List[GoldAction] = field(default_factory=list)
|
||||
required_info: List[str] = field(default_factory=list)
|
||||
nl_assertions: List[str] = field(default_factory=list)
|
||||
|
||||
def gold_action_names(self) -> List[str]:
|
||||
return [a.name for a in self.gold_actions]
|
||||
|
||||
|
||||
def _as_list(v: Any) -> List[Any]:
|
||||
if v is None:
|
||||
return []
|
||||
if isinstance(v, list):
|
||||
return v
|
||||
return [v]
|
||||
|
||||
|
||||
def _str_list(v: Any) -> List[str]:
|
||||
out: List[str] = []
|
||||
for item in _as_list(v):
|
||||
if isinstance(item, str):
|
||||
out.append(item)
|
||||
elif isinstance(item, dict):
|
||||
# communicate_info entries are sometimes {"info": "..."} dicts.
|
||||
for key in ("info", "content", "text", "value"):
|
||||
if isinstance(item.get(key), str):
|
||||
out.append(item[key])
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def normalize_row(row: Dict[str, Any], *, index: int = 0) -> ToolScaleTask:
|
||||
"""Turn one raw ToolScale row into a :class:`ToolScaleTask` (pure)."""
|
||||
scenario = row.get("user_scenario") or {}
|
||||
instructions = scenario.get("instructions") or {}
|
||||
instruction = (
|
||||
instructions.get("task_instructions")
|
||||
or instructions.get("reason_for_call")
|
||||
or row.get("task")
|
||||
or row.get("instruction")
|
||||
or ""
|
||||
)
|
||||
domain = scenario.get("domain") or row.get("domain") or "unknown"
|
||||
|
||||
crit = row.get("evaluation_criteria") or {}
|
||||
gold: List[GoldAction] = []
|
||||
for a in _as_list(crit.get("actions")):
|
||||
if isinstance(a, dict) and a.get("name"):
|
||||
gold.append(
|
||||
GoldAction(
|
||||
name=str(a["name"]),
|
||||
arguments=a.get("arguments") or a.get("args") or {},
|
||||
action_id=a.get("action_id"),
|
||||
)
|
||||
)
|
||||
|
||||
required = _str_list(crit.get("communicate_info"))
|
||||
nl = _str_list(crit.get("nl_assertions"))
|
||||
|
||||
task_id = str(row.get("id") or row.get("task_id") or f"toolscale-{index}")
|
||||
return ToolScaleTask(
|
||||
task_id=task_id,
|
||||
domain=str(domain),
|
||||
instruction=str(instruction),
|
||||
gold_actions=gold,
|
||||
required_info=required,
|
||||
nl_assertions=nl,
|
||||
)
|
||||
|
||||
|
||||
def load_toolscale(
|
||||
*,
|
||||
max_tasks: Optional[int] = None,
|
||||
split: str = "train",
|
||||
source: Optional[Iterable[Dict[str, Any]]] = None,
|
||||
) -> Iterator[ToolScaleTask]:
|
||||
"""Yield normalized ToolScale tasks.
|
||||
|
||||
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
|
||||
When ``source`` is None, streams ``nvidia/ToolScale`` via ``datasets``.
|
||||
"""
|
||||
if source is None:
|
||||
from datasets import load_dataset # lazy: optional dep / network
|
||||
|
||||
source = load_dataset(DATASET_ID, split=split, streaming=True)
|
||||
|
||||
n = 0
|
||||
for i, row in enumerate(source):
|
||||
if max_tasks is not None and n >= max_tasks:
|
||||
break
|
||||
task = normalize_row(dict(row), index=i)
|
||||
if not task.instruction.strip():
|
||||
continue
|
||||
yield task
|
||||
n += 1
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DATASET_ID",
|
||||
"GoldAction",
|
||||
"ToolScaleTask",
|
||||
"load_toolscale",
|
||||
"normalize_row",
|
||||
]
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Serialize a verified unified-tool rollout into an SFT ``conversations`` record.
|
||||
|
||||
Output matches what ``OrchestratorSFTDataset`` consumes, and trains the model to
|
||||
emit the ``<tool_call>{...}</tool_call>`` text form that
|
||||
``toolorchestra.parsing._parse_rl_tool_call`` already reads back. One record =
|
||||
one passing trajectory.
|
||||
|
||||
Roles: ``system`` (the unified tool catalog), ``user`` (the running ``Problem``
|
||||
prompt), ``assistant`` (reasoning + a ``<tool_call>`` tag, or the final answer),
|
||||
``tool`` (the executed observation).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import ExpertTool, build_tool_specs
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import (
|
||||
UnifiedRollout,
|
||||
build_system_prompt,
|
||||
tool_call_tag,
|
||||
)
|
||||
|
||||
|
||||
def _system_prompt(tools: List[ExpertTool], rollout: UnifiedRollout) -> str:
|
||||
# Faithful ToolOrchestra system prompt (paper's prepare_sft_data.py format).
|
||||
# Prefer the EXACT specs the policy saw this rollout (anonymized labels when
|
||||
# anonymize=True) so the saved <tools> block matches the assistant's
|
||||
# <tool_call> tags. Falling back to the real registry here is the bug that
|
||||
# leaked real model names+pricing into anonymized records.
|
||||
specs = rollout.tool_specs if rollout.tool_specs else build_tool_specs(tools)
|
||||
return build_system_prompt(specs)
|
||||
|
||||
|
||||
def _normalize_think(text: str) -> str:
|
||||
"""Ensure a reasoning block has matched tags. Qwen rollouts routinely yield a
|
||||
dangling ``</think>`` with NO opening ``<think>`` (the chat template consumes
|
||||
the opener in the prompt, so only the closer survives in the completion). Add
|
||||
the opener back so the trained turn is well-formed."""
|
||||
t = (text or "").lstrip()
|
||||
if "</think>" in t and "<think>" not in t:
|
||||
t = "<think>\n" + t
|
||||
return t
|
||||
|
||||
|
||||
def _debox(text: str) -> str:
|
||||
"""Unwrap every ``\\boxed{X}`` -> ``X`` (balanced braces). The format kills
|
||||
\\boxed everywhere, but small models keep emitting it in otherwise-correct
|
||||
answers — de-box to salvage them rather than reject the whole trajectory."""
|
||||
out, i = [], 0
|
||||
marker = r"\boxed{"
|
||||
while i < len(text):
|
||||
j = text.find(marker, i)
|
||||
if j == -1:
|
||||
out.append(text[i:])
|
||||
break
|
||||
out.append(text[i:j])
|
||||
k = j + len(marker)
|
||||
depth, inner = 1, []
|
||||
while k < len(text) and depth:
|
||||
c = text[k]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
break
|
||||
inner.append(c)
|
||||
k += 1
|
||||
out.append("".join(inner))
|
||||
i = k + 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
# Reasoning that PAROTS the training constraint back ("I must call a model per
|
||||
# my instructions", "this is testing whether I need to route") — the model
|
||||
# reasoning about its own harness. Never a good target: drop the whole line.
|
||||
_LEAK_RE = re.compile(
|
||||
r"(?im)^.*\b(?:"
|
||||
r"the user (?:is|was|wants|has|had|needs|'s|would|will|asked"
|
||||
r"|'ve asked|has asked)\b|"
|
||||
r"asked me to\b|" # echoes of an injected nudge turn
|
||||
r"make progress by delegat|"
|
||||
r"delegat\w+[^\n]*sub-?question|"
|
||||
r"(?:to |ask |send .{0,20}to )?one of the models\b|"
|
||||
r"instructions?|" # any self-reference to the rules
|
||||
r"requirement|"
|
||||
r"at least one (?:model|tool|expert)|" # "(call|invoke) at least one model"
|
||||
r"(?:call|invoke|route to|delegate to|use)\s+(?:a|one|at least one|the)\s+model\b|"
|
||||
r"route\s+(?:this|it|that)\b[^\n]{0,40}\bthrough\s+(?:a|one|the)\s+model\b|"
|
||||
r"without mentioning (?:any )?(?:tools?|reasoning|meta|steps)|"
|
||||
r"in the required format\b|"
|
||||
r"based on the (?:model|response|analysis|expert)|" # response-acknowledgment
|
||||
r"the model (?:provided|confirmed|gave|returned|correctly|analy\w+"
|
||||
r"|respon\w+|said|indicated|identified)|"
|
||||
r"i (?:got|received|have) (?:a |the )?"
|
||||
r"(?:comprehensive |clear |good |detailed )?answer\b|"
|
||||
r"now i can (?:confidently )?(?:give|provide|state|answer)|"
|
||||
r"perfect[!,]|great[!,]|excellent[!,]|" # narration exclamations
|
||||
r"as an orchestrator|"
|
||||
r"whose job is to (?:route|delegate|orchestrate)|"
|
||||
r"testing whether i (?:need|have|should|must)\b|"
|
||||
r"i(?:'m| am)? (?:required|instructed|supposed|expected) to\b"
|
||||
r")[^\n]*(?:\n|$)"
|
||||
)
|
||||
# A leading task-restatement opener ("The user is asking…") — the format
|
||||
# explicitly bans this meta-text; drop it when it opens the reasoning.
|
||||
_META_OPEN_RE = re.compile(
|
||||
r"(?is)^\s*(?:the user (?:is|was|wants|has|needs|would|'s)\b[^.\n]*[.\n]+\s*)+"
|
||||
)
|
||||
|
||||
|
||||
def _scrub_meta(text: str) -> str:
|
||||
"""Remove harness-leakage lines and a leading 'The user is asking…' meta
|
||||
opener from a reasoning block. Both are disallowed by the system prompt, so
|
||||
they must not survive into the supervised target. Preserves a leading
|
||||
``<think>`` marker so the block stays well-formed."""
|
||||
if not text:
|
||||
return text
|
||||
think = ""
|
||||
body = text
|
||||
m = re.match(r"(?is)^\s*(<think>)\s*", body)
|
||||
if m:
|
||||
think = "<think>\n"
|
||||
body = body[m.end() :]
|
||||
body = _LEAK_RE.sub("", body)
|
||||
body = _META_OPEN_RE.sub("", body).lstrip()
|
||||
return (think + body).strip() if think else body.strip()
|
||||
|
||||
|
||||
# Control / special tokens that must never survive into a training target's
|
||||
# final answer. Three shapes, in order of alternation:
|
||||
# 1. closed pipe token — ``<|im_end|>``, ``<|im_start|>``, ``<|eot_id|>``,
|
||||
# ``<|end_of_turn|>``, ``<|"|>`` (the stray-quote leak).
|
||||
# 2. UNCLOSED pipe token — ``<|tool_call>`` (opened with ``<|`` but closed on a
|
||||
# bare ``>`` because the decode was cut before the second pipe).
|
||||
# 3. named angle special tokens — gemma ``<start_of_turn>`` / ``<end_of_turn>``
|
||||
# and the sentinel set ``<eos> <bos> <pad> <unk> <s> </s>``.
|
||||
# Deliberately narrow: only the ``<|...|>`` pipe form and this known name list
|
||||
# match, so legitimate ``<``/``>`` in math/code (``x < 3 and y > 2``) is left
|
||||
# untouched.
|
||||
_CONTROL_TOKEN_RE = re.compile(
|
||||
r"<\|[^>]*?\|>" # closed pipe token
|
||||
r"|<\|[^|>]*>" # unclosed pipe token
|
||||
r"|</?(?:start_of_turn|end_of_turn|eos|bos|pad|unk|s)>" # named angle tokens
|
||||
)
|
||||
|
||||
|
||||
def _strip_control_tokens(text: str) -> str:
|
||||
"""Delete residual control/special tokens from an answer so a good answer
|
||||
with a stray token is salvaged rather than dropped. Loops (bounded) because
|
||||
removing one match can expose a nested/overlapping leftover (e.g.
|
||||
``<|<|im_end|>|>``). Re-strips whitespace at the end."""
|
||||
if not text:
|
||||
return text
|
||||
out = text
|
||||
for _ in range(4):
|
||||
stripped = _CONTROL_TOKEN_RE.sub("", out)
|
||||
if stripped == out:
|
||||
break
|
||||
out = stripped
|
||||
return out.strip()
|
||||
|
||||
|
||||
def _final_answer_block(text: str) -> str:
|
||||
"""Render the final-answer assistant message with a single clean
|
||||
``FINAL_ANSWER: <value>`` line. De-boxes the answer, strips leaked control
|
||||
tokens, and normalizes whatever spelling the model used (``FINALANSWER``,
|
||||
``FINAL ANSWER``, no colon, …) to the exact tag. Keeps the think block at
|
||||
most once."""
|
||||
text = _debox(_normalize_think((text or "").strip()))
|
||||
# Normalize any final-answer marker spelling to the canonical form; take the
|
||||
# LAST one as the real answer (idempotent — never emits two tags).
|
||||
marks = list(re.finditer(r"(?im)FINAL[_\s]?ANSWER\s*:?", text))
|
||||
if marks:
|
||||
last = marks[-1]
|
||||
answer = _strip_control_tokens(text[last.end() :].strip())
|
||||
# Final turn = the bare short answer only. Drop EVERYTHING before
|
||||
# FINAL_ANSWER — both the visible "Perfect! Based on the model's response…"
|
||||
# narration AND the final-turn <think>, which is just confirmatory
|
||||
# "both models agree… I've verified…" fluff. The routing reasoning already
|
||||
# lives in the earlier tool-call turns, so nothing of value is lost, and
|
||||
# this kills all final-turn narration deterministically (any wording).
|
||||
return f"FINAL_ANSWER: {answer}"
|
||||
if "</think>" in text:
|
||||
# No explicit marker: the answer is whatever follows the final </think>.
|
||||
# Drop the think block (same rationale as above) — keep only the answer.
|
||||
_, _, answer = text.rpartition("</think>")
|
||||
answer = _strip_control_tokens(answer.strip())
|
||||
return (
|
||||
f"FINAL_ANSWER: {answer}"
|
||||
if answer
|
||||
else f"FINAL_ANSWER: {_strip_control_tokens(_scrub_meta(text))}"
|
||||
)
|
||||
return f"FINAL_ANSWER: {_strip_control_tokens(text)}"
|
||||
|
||||
|
||||
def trajectory_to_record(
|
||||
task_id: str,
|
||||
question: str,
|
||||
tools: List[ExpertTool],
|
||||
rollout: UnifiedRollout,
|
||||
*,
|
||||
reward: float = 0.0,
|
||||
domain: str = "unknown",
|
||||
) -> Dict[str, Any]:
|
||||
"""Convert a passing :class:`UnifiedRollout` into one SFT JSONL record."""
|
||||
conversations: List[Dict[str, str]] = [
|
||||
{"role": "system", "content": _system_prompt(tools, rollout)},
|
||||
{"role": "user", "content": f"Problem: {question}"},
|
||||
]
|
||||
|
||||
for turn in rollout.turns:
|
||||
if turn.tool_name is None:
|
||||
# Final-answer turn. turn.reasoning is the model's actual output for
|
||||
# this turn (which already contains the answer); render it once.
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": _final_answer_block(
|
||||
turn.reasoning or rollout.final_answer
|
||||
),
|
||||
}
|
||||
)
|
||||
continue
|
||||
tag = tool_call_tag(turn.tool_name, turn.arguments)
|
||||
reasoning = _scrub_meta(_normalize_think((turn.reasoning or "").rstrip()))
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (reasoning + "\n" + tag).strip(),
|
||||
}
|
||||
)
|
||||
conversations.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"name": turn.tool_name,
|
||||
"content": turn.observation or "",
|
||||
}
|
||||
)
|
||||
|
||||
# If the rollout terminated on max_turns (no None turn), append the answer.
|
||||
if not rollout.turns or rollout.turns[-1].tool_name is not None:
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": _final_answer_block(rollout.final_answer),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"conversations": conversations,
|
||||
"task_id": task_id,
|
||||
"domain": domain,
|
||||
"reward": reward,
|
||||
"metrics": {
|
||||
"cost_usd": rollout.cost_usd,
|
||||
"tokens": rollout.tokens,
|
||||
"num_tool_calls": rollout.num_tool_calls,
|
||||
"num_turns": len(rollout.turns),
|
||||
# Present only when experts were anonymized: opaque label -> real
|
||||
# tool name, so analysis can recover which model was actually picked.
|
||||
**({"anon_map": rollout.anon_map} if rollout.anon_map else {}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["trajectory_to_record"]
|
||||
@@ -0,0 +1,255 @@
|
||||
"""Correctness verifier for orchestrator reasoning tasks (:mod:`.datasets`).
|
||||
|
||||
The verifier dispatches on :attr:`Task.domain`:
|
||||
|
||||
- **math** — normalize and compare. Extract ``\\boxed{...}`` from both sides,
|
||||
try ``sympy`` for symbolic / numeric equality, and fall back to a normalized
|
||||
string / number match. No network.
|
||||
- **code** — best-effort: if a short expected output / answer exists, do a
|
||||
normalized substring match; otherwise defer to the LLM judge. (We don't run
|
||||
arbitrary code here.)
|
||||
- **science / medical / chat / misc / unknown** — Gemini LLM judge given the
|
||||
question + gold answer + candidate, asked for ``PASS`` / ``FAIL``. When no
|
||||
``GEMINI_API_KEY`` is set (e.g. offline tests) it falls back to a normalized
|
||||
string / token-F1 >= 0.6 match.
|
||||
|
||||
The OpenAI key is dead, so the LLM judge talks to Gemini through its
|
||||
OpenAI-compatible endpoint. Normalization helpers are copied (not imported)
|
||||
from ``hotpotqa.py`` to keep this module self-contained.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import string
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
from .datasets import Task
|
||||
|
||||
GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
|
||||
GEMINI_MODEL = "gemini-2.5-flash"
|
||||
# LLM judge model. Switched from Gemini (free-tier rate limit stalled parallel gen)
|
||||
# to Anthropic Haiku: fast (~0.6s), direct PASS/FAIL, high rate limits.
|
||||
JUDGE_MODEL = "claude-haiku-4-5-20251001"
|
||||
# Reused across calls — creating a fresh Anthropic client per judge call leaks an
|
||||
# httpx connection pool (CLOSE-WAIT pileup under parallel generation). Thread-safe.
|
||||
_JUDGE_CLIENT = None
|
||||
|
||||
|
||||
# --- normalization (copied from hotpotqa.py — keep self-contained) -----------
|
||||
|
||||
|
||||
def _normalize_answer(s: str) -> str:
|
||||
"""Lowercase, strip punctuation / articles / extra whitespace (SQuAD/HotpotQA)."""
|
||||
s = s.lower()
|
||||
s = "".join(ch for ch in s if ch not in set(string.punctuation))
|
||||
s = re.sub(r"\b(a|an|the)\b", " ", s)
|
||||
return " ".join(s.split())
|
||||
|
||||
|
||||
def _f1(pred: str, gold: str) -> float:
|
||||
pt, gt = _normalize_answer(pred).split(), _normalize_answer(gold).split()
|
||||
if not pt or not gt:
|
||||
return float(pt == gt)
|
||||
common: Dict[str, int] = {}
|
||||
for w in pt:
|
||||
if w in gt:
|
||||
common[w] = common.get(w, 0) + 1
|
||||
num_same = sum(common.values())
|
||||
if num_same == 0:
|
||||
return 0.0
|
||||
precision = num_same / len(pt)
|
||||
recall = num_same / len(gt)
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
def _string_or_f1(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool:
|
||||
"""Normalized whole-word substring OR token-F1 >= threshold."""
|
||||
np_, ng = _normalize_answer(prediction), _normalize_answer(gold)
|
||||
if not ng:
|
||||
return False
|
||||
if f" {ng} " in f" {np_} ":
|
||||
return True
|
||||
return _f1(prediction, gold) >= f1_threshold
|
||||
|
||||
|
||||
# --- math --------------------------------------------------------------------
|
||||
|
||||
_BOXED_RE = re.compile(r"\\boxed\{")
|
||||
_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?")
|
||||
|
||||
|
||||
def _extract_boxed(text: str) -> Optional[str]:
|
||||
"""Return the contents of the last ``\\boxed{...}`` (brace-balanced)."""
|
||||
last = None
|
||||
for m in _BOXED_RE.finditer(text):
|
||||
start = m.end()
|
||||
depth = 1
|
||||
i = start
|
||||
while i < len(text) and depth > 0:
|
||||
c = text[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
i += 1
|
||||
if depth == 0:
|
||||
last = text[start : i - 1].strip()
|
||||
return last
|
||||
|
||||
|
||||
def _clean_math(s: str) -> str:
|
||||
s = s.strip()
|
||||
boxed = _extract_boxed(s)
|
||||
if boxed is not None:
|
||||
s = boxed
|
||||
# Common LaTeX wrappers / delimiters.
|
||||
s = s.replace("$", "").replace("\\!", "").replace("\\,", "").replace("\\;", "")
|
||||
s = s.replace("\\left", "").replace("\\right", "")
|
||||
s = s.strip().strip("$ ")
|
||||
return s
|
||||
|
||||
|
||||
def _to_float(s: str) -> Optional[float]:
|
||||
s = s.replace(",", "").strip()
|
||||
try:
|
||||
return float(s)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
m = _NUM_RE.search(s)
|
||||
if m:
|
||||
try:
|
||||
return float(m.group(0))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _math_equal(prediction: str, gold: str) -> bool:
|
||||
pred = _clean_math(prediction)
|
||||
g = _clean_math(gold)
|
||||
if not g:
|
||||
return False
|
||||
|
||||
# Exact normalized-string shortcut.
|
||||
if pred.replace(" ", "") == g.replace(" ", "") and pred:
|
||||
return True
|
||||
|
||||
# Numeric comparison (handles "42" vs "42.0" vs trailing prose).
|
||||
pf, gf = _to_float(pred), _to_float(g)
|
||||
if pf is not None and gf is not None:
|
||||
if abs(pf - gf) <= 1e-6 * max(1.0, abs(gf)):
|
||||
return True
|
||||
|
||||
# NOTE: a sympy/parse_latex symbolic-equality block used to live here, but
|
||||
# sympy.simplify / antlr parse_latex can hang on pathological \boxed{} answers
|
||||
# while holding the GIL -> deadlocks the whole threaded rejection sampler
|
||||
# (all worker threads freeze in futex_wait, server goes idle, 0 records).
|
||||
# Symbolic equivalence is now deferred to the Gemini judge in verify_answer().
|
||||
|
||||
# Last resort: whole-word substring of the gold in the prediction.
|
||||
return _string_or_f1(prediction, g, f1_threshold=0.9)
|
||||
|
||||
|
||||
# --- LLM judge (Gemini, OpenAI-compatible) -----------------------------------
|
||||
|
||||
|
||||
def _gemini_judge(task: Task, prediction: str) -> Optional[bool]:
|
||||
"""Ask Gemini PASS/FAIL. Returns None if unavailable (caller falls back)."""
|
||||
api_key = os.environ.get("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
return None
|
||||
try:
|
||||
import anthropic
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Judge = Anthropic Haiku (fast ~0.6s, direct PASS/FAIL, high rate limits).
|
||||
# Switched off the Gemini judge: its free-tier rate limit + SDK retry/backoff
|
||||
# blocked worker threads for minutes under parallel generation and stalled
|
||||
# the run. Fail-fast (no retries, short timeout): on error return None and the
|
||||
# caller falls back to string/f1.
|
||||
global _JUDGE_CLIENT
|
||||
if _JUDGE_CLIENT is None:
|
||||
_JUDGE_CLIENT = anthropic.Anthropic(
|
||||
api_key=api_key, max_retries=0, timeout=15
|
||||
)
|
||||
client = _JUDGE_CLIENT
|
||||
prompt = (
|
||||
"You are grading a candidate answer against a gold reference.\n"
|
||||
"Reply with exactly one word: PASS if the candidate is correct and "
|
||||
"consistent with the gold answer, otherwise FAIL.\n\n"
|
||||
f"Question:\n{task.question}\n\n"
|
||||
f"Gold answer:\n{task.answer}\n\n"
|
||||
f"Candidate answer:\n{prediction}\n\n"
|
||||
"Verdict (PASS or FAIL):"
|
||||
)
|
||||
resp = client.messages.create(
|
||||
model=JUDGE_MODEL,
|
||||
max_tokens=8,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
)
|
||||
verdict = (resp.content[0].text or "").strip().upper()
|
||||
if "PASS" in verdict:
|
||||
return True
|
||||
if "FAIL" in verdict:
|
||||
return False
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
# --- public API --------------------------------------------------------------
|
||||
|
||||
|
||||
def verify_answer(task: Task, prediction: str) -> bool:
|
||||
"""Domain-dispatched correctness check for ``prediction`` vs ``task.answer``."""
|
||||
pred = (prediction or "").strip()
|
||||
if not pred or not (task.answer or "").strip():
|
||||
return False
|
||||
|
||||
domain = (task.domain or "").lower()
|
||||
|
||||
if domain == "math":
|
||||
# string + numeric + f1 only. sympy removed (GIL-deadlocked the threaded
|
||||
# sampler); Gemini judge NOT used here either (32 tasks x N samples of math
|
||||
# judge calls throttle Gemini and stall the whole run). Accept slightly lower
|
||||
# math yield for a fast, hang-free verifier.
|
||||
return _math_equal(pred, task.answer)
|
||||
|
||||
if domain == "code":
|
||||
# Best-effort: short gold -> normalized substring; otherwise LLM judge.
|
||||
gold = task.answer.strip()
|
||||
if len(gold) <= 200:
|
||||
if _string_or_f1(pred, gold, f1_threshold=0.8):
|
||||
return True
|
||||
judged = _gemini_judge(task, pred)
|
||||
if judged is not None:
|
||||
return judged
|
||||
return _string_or_f1(pred, gold, f1_threshold=0.6)
|
||||
|
||||
# science / medical / chat / misc / unknown -> LLM judge, then F1 fallback.
|
||||
judged = _gemini_judge(task, pred)
|
||||
if judged is not None:
|
||||
return judged
|
||||
return _string_or_f1(pred, task.answer, f1_threshold=0.6)
|
||||
|
||||
|
||||
def make_verifier() -> Callable[[Any, Any], bool]:
|
||||
"""Return ``verify_fn(task, rollout) -> bool`` for the rejection sampler."""
|
||||
|
||||
def verify(task: Any, rollout: Any) -> bool:
|
||||
ans = getattr(rollout, "final_answer", "") or ""
|
||||
return verify_answer(task, ans)
|
||||
|
||||
return verify
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GEMINI_BASE_URL",
|
||||
"GEMINI_MODEL",
|
||||
"make_verifier",
|
||||
"verify_answer",
|
||||
]
|
||||
@@ -52,7 +52,7 @@ class OrchestratorSFTConfig:
|
||||
"""Configuration for orchestrator SFT training."""
|
||||
|
||||
# Model
|
||||
model_name: str = "Qwen/Qwen3-1.7B"
|
||||
model_name: str = "Qwen/Qwen3.5-9B"
|
||||
max_seq_length: int = 4096
|
||||
|
||||
# Training
|
||||
@@ -63,17 +63,28 @@ class OrchestratorSFTConfig:
|
||||
warmup_ratio: float = 0.1
|
||||
max_grad_norm: float = 1.0
|
||||
|
||||
# Trace generation
|
||||
teacher_engine_key: str = ""
|
||||
teacher_model: str = ""
|
||||
traces_per_query: int = 2
|
||||
max_attempts_per_trace: int = 3
|
||||
# Trace generation by base Qwen3-8B self-sampling (v1 cold-start). The
|
||||
# orchestrator IS the local Qwen3-8B served over an OpenAI-compatible vLLM
|
||||
# endpoint; we roll it out over load_sft_tasks() and keep correct trajectories.
|
||||
# For v2, point orchestrator_endpoint/model at the v1 checkpoint's server.
|
||||
orchestrator_endpoint: str = "http://localhost:8001/v1"
|
||||
orchestrator_model: str = "qwen3-8b"
|
||||
orchestrator_api_key: str = "EMPTY"
|
||||
generation_temperature: float = 0.7
|
||||
|
||||
# Data source
|
||||
trace_cache_path: str = "data/orchestrator_sft_traces.jsonl"
|
||||
trace_cache_path: str = "data/orchestrator_sft_v1.jsonl"
|
||||
regenerate_traces: bool = False
|
||||
|
||||
# Rejection-sampling cold-start over load_sft_tasks (see sft_data/reject_sample.py).
|
||||
distill_max_tasks: int = 0 # 0 / None -> all of load_sft_tasks()
|
||||
samples_per_task: int = 8
|
||||
max_keep_per_task: int = 1
|
||||
max_rollout_turns: int = 8
|
||||
# Optional local OSS model endpoints, mapping full model id (e.g.
|
||||
# "Qwen/Qwen3.5-9B") -> vLLM base_url; omitted when unset.
|
||||
local_endpoints: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# Checkpoint
|
||||
checkpoint_dir: str = "checkpoints/orchestrator_sft"
|
||||
save_every_n_epochs: int = 1
|
||||
@@ -267,11 +278,81 @@ class OrchestratorSFTTrainer:
|
||||
)
|
||||
|
||||
def _generate_traces(self) -> None:
|
||||
"""Generate SFT traces (placeholder — requires running engine)."""
|
||||
"""Generate SFT traces by base Qwen3-8B self-sampling (v1 cold-start).
|
||||
|
||||
The orchestrator IS the local Qwen3-8B served over an OpenAI-compatible
|
||||
vLLM endpoint (``orchestrator_endpoint``/``orchestrator_model``). We roll
|
||||
it out over ``load_sft_tasks()`` on the fixed orchestrator tool catalog,
|
||||
verify each trajectory's final answer against the gold
|
||||
(``verify.make_verifier``), keep the cheapest-correct trajectory per task,
|
||||
and serialize them into the ``<tool_call>`` ``conversations`` JSONL the SFT
|
||||
dataset consumes. For v2, point the endpoint/model at the v1 checkpoint's
|
||||
server. See ``sft_data/reject_sample.py`` and the
|
||||
``scripts/orchestrator/build_orchestrator_sft.py`` driver.
|
||||
"""
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
|
||||
from openjarvis.agents.hybrid.toolorchestra.unified import (
|
||||
make_call_orchestrator,
|
||||
make_dispatch,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data import (
|
||||
reject_sample as _reject_sample,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
|
||||
load_sft_tasks,
|
||||
)
|
||||
|
||||
generate_sft_dataset = _reject_sample.generate_sft_dataset
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
|
||||
make_verifier,
|
||||
)
|
||||
|
||||
trace_path = Path(self.config.trace_cache_path)
|
||||
trace_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not trace_path.exists():
|
||||
trace_path.touch()
|
||||
tools = orchestrator_catalog(
|
||||
local_endpoints=self.config.local_endpoints or None,
|
||||
)
|
||||
call_orch = make_call_orchestrator(
|
||||
self.config.orchestrator_model,
|
||||
base_url=self.config.orchestrator_endpoint,
|
||||
api_key=self.config.orchestrator_api_key,
|
||||
temperature=self.config.generation_temperature,
|
||||
)
|
||||
dispatch = make_dispatch({})
|
||||
|
||||
def rollout_fn(task: Any) -> Any:
|
||||
try:
|
||||
return run_unified_rollout(
|
||||
task.instruction,
|
||||
tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=self.config.max_rollout_turns,
|
||||
)
|
||||
except Exception as exc: # network/key failures shouldn't kill the run
|
||||
logger.warning("rollout failed for %s: %s", task.task_id, exc)
|
||||
return None
|
||||
|
||||
try:
|
||||
tasks = load_sft_tasks()
|
||||
if self.config.distill_max_tasks:
|
||||
tasks = tasks[: self.config.distill_max_tasks]
|
||||
stats = generate_sft_dataset(
|
||||
str(trace_path),
|
||||
tasks=tasks,
|
||||
tools=tools,
|
||||
rollout_fn=rollout_fn,
|
||||
verify_fn=make_verifier(),
|
||||
samples_per_task=self.config.samples_per_task,
|
||||
max_keep_per_task=self.config.max_keep_per_task,
|
||||
reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets top reward
|
||||
)
|
||||
logger.info("Generated SFT traces: %s", stats)
|
||||
except Exception as exc: # network/datasets optional — fail soft to empty
|
||||
logger.warning("trace generation failed (%s); writing empty set", exc)
|
||||
trace_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not trace_path.exists():
|
||||
trace_path.touch()
|
||||
|
||||
def _init_optimizer(self) -> None:
|
||||
if not HAS_TORCH or self.policy.model is None:
|
||||
|
||||
@@ -42,8 +42,22 @@ _MATH_FUNCS = {
|
||||
"sin": math.sin,
|
||||
"cos": math.cos,
|
||||
"tan": math.tan,
|
||||
"asin": math.asin,
|
||||
"acos": math.acos,
|
||||
"atan": math.atan,
|
||||
"atan2": math.atan2,
|
||||
"sinh": math.sinh,
|
||||
"cosh": math.cosh,
|
||||
"tanh": math.tanh,
|
||||
"radians": math.radians,
|
||||
"degrees": math.degrees,
|
||||
"exp": math.exp,
|
||||
"fabs": math.fabs,
|
||||
"hypot": math.hypot,
|
||||
"factorial": math.factorial,
|
||||
"pi": math.pi,
|
||||
"e": math.e,
|
||||
"tau": math.tau,
|
||||
"ceil": math.ceil,
|
||||
"floor": math.floor,
|
||||
}
|
||||
|
||||
@@ -116,11 +116,31 @@ class WebSearchTool(BaseTool):
|
||||
return text
|
||||
|
||||
def _duckduckgo_search(self, query: str, max_results: int) -> str:
|
||||
"""Search using DuckDuckGo as fallback."""
|
||||
"""Search using DuckDuckGo as fallback.
|
||||
|
||||
ddgs queries several engines (yandex/yahoo/brave/...) that frequently
|
||||
hang; with no timeout this blocks the whole rollout. Cap each engine
|
||||
request (DDGS timeout) AND wrap the call in a hard wall-clock deadline so
|
||||
a flaky search fast-fails instead of stalling data generation.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from ddgs import DDGS
|
||||
|
||||
ddgs = DDGS()
|
||||
raw_results = list(ddgs.text(query, max_results=max_results))
|
||||
def _run():
|
||||
ddgs = DDGS(timeout=5)
|
||||
return list(ddgs.text(query, max_results=max_results))
|
||||
|
||||
# Don't use `with` — its shutdown(wait=True) blocks on the hung thread,
|
||||
# defeating the deadline. submit, wait up to 8s, then abandon the thread
|
||||
# (it finishes in the background harmlessly) and fast-fail.
|
||||
_ex = ThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
raw_results = _ex.submit(_run).result(timeout=8)
|
||||
except Exception:
|
||||
_ex.shutdown(wait=False, cancel_futures=True)
|
||||
return "[web_search: no results (search backend timed out)]"
|
||||
_ex.shutdown(wait=False)
|
||||
results = []
|
||||
for r in raw_results:
|
||||
title = r.get("title", "Untitled")
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Unit coverage for ``anonymize_tools`` — the identity-stripping step that makes
|
||||
routing data unbiased. See ``expert_registry.anonymize_tools``.
|
||||
|
||||
The anonymizer takes the orchestrator catalog and replaces every MODEL expert
|
||||
with an opaque ``model_xxxx`` label, a uniform brand-free description, no
|
||||
price/latency line, and shuffles the expert block — so the policy can't route on
|
||||
a model's name, position, cost, or tier. Basic tools keep their real names.
|
||||
It returns ``(anon_tools, anon_to_real)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import (
|
||||
KIND_MODEL,
|
||||
anonymize_tools,
|
||||
build_tool_specs,
|
||||
orchestrator_catalog,
|
||||
tools_by_name,
|
||||
)
|
||||
|
||||
# Real brand tokens that must never leak into the anonymized, model-facing specs.
|
||||
_BRANDS = ("gpt", "claude", "gemini", "qwen")
|
||||
|
||||
|
||||
def _model_names(cat):
|
||||
return [t.name for t in cat if t.kind == KIND_MODEL]
|
||||
|
||||
|
||||
def _basic_names(cat):
|
||||
return [t.name for t in cat if t.kind != KIND_MODEL]
|
||||
|
||||
|
||||
def test_no_brand_names_in_anonymized_specs():
|
||||
cat = orchestrator_catalog()
|
||||
anon, _ = anonymize_tools(cat, random.Random(0))
|
||||
specs = build_tool_specs(anon)
|
||||
# The whole model-facing payload (names + descriptions + categories) must be
|
||||
# brand-free. `.model` is preserved on the tool for dispatch but is NOT part
|
||||
# of the spec the orchestrator conditions on, so it's fine that it still holds
|
||||
# the real id.
|
||||
blob = json.dumps(specs).lower()
|
||||
for brand in _BRANDS:
|
||||
assert brand not in blob, f"brand {brand!r} leaked into anonymized specs"
|
||||
# And specifically the anonymized expert descriptions carry no brand.
|
||||
for t in anon:
|
||||
if t.name.startswith("model_"):
|
||||
assert not any(b in t.description().lower() for b in _BRANDS)
|
||||
|
||||
|
||||
def test_hide_cost_removes_price_line_from_descriptions():
|
||||
cat = orchestrator_catalog()
|
||||
# Sanity: the raw model tools DO surface a price line before anonymizing.
|
||||
raw_model = next(t for t in cat if t.kind == KIND_MODEL)
|
||||
assert "Pricing:" in raw_model.description()
|
||||
|
||||
anon, anon_to_real = anonymize_tools(cat, random.Random(1))
|
||||
for t in anon:
|
||||
if t.name in anon_to_real: # an anonymized model expert
|
||||
assert t.hide_cost is True
|
||||
desc = t.description()
|
||||
assert "Pricing:" not in desc
|
||||
assert "$" not in desc
|
||||
assert "/1M" not in desc
|
||||
assert "latency" not in desc.lower()
|
||||
|
||||
|
||||
def test_each_model_maps_to_opaque_label_and_round_trips():
|
||||
cat = orchestrator_catalog()
|
||||
orig_models = _model_names(cat)
|
||||
anon, anon_to_real = anonymize_tools(cat, random.Random(2))
|
||||
|
||||
# One opaque label per real model, all in the model_xxxx namespace.
|
||||
assert len(anon_to_real) == len(orig_models)
|
||||
assert all(lbl.startswith("model_") for lbl in anon_to_real)
|
||||
|
||||
# No collisions: labels unique, and each real model recovered exactly once.
|
||||
assert len(set(anon_to_real)) == len(anon_to_real)
|
||||
assert len(set(anon_to_real.values())) == len(anon_to_real)
|
||||
assert set(anon_to_real.values()) == set(orig_models)
|
||||
|
||||
# Round-trip: every anonymized expert's label resolves back to a real name,
|
||||
# and `.model` is preserved so dispatch still reaches the right backend.
|
||||
by_orig = tools_by_name(cat)
|
||||
for t in anon:
|
||||
if t.name.startswith("model_"):
|
||||
real = anon_to_real[t.name]
|
||||
assert real in by_orig
|
||||
assert t.model == by_orig[real].model # backend id untouched
|
||||
|
||||
# Basic tools keep their real names and are untouched by the label map.
|
||||
basics = _basic_names(cat)
|
||||
anon_names = {t.name for t in anon}
|
||||
assert set(basics) <= anon_names
|
||||
assert not (set(basics) & set(anon_to_real))
|
||||
|
||||
|
||||
def test_experts_block_on_top_then_basics():
|
||||
cat = orchestrator_catalog()
|
||||
anon, anon_to_real = anonymize_tools(cat, random.Random(3))
|
||||
n_models = len(anon_to_real)
|
||||
# Experts form a contiguous block at the front, basics underneath.
|
||||
assert all(t.name.startswith("model_") for t in anon[:n_models])
|
||||
assert all(not t.name.startswith("model_") for t in anon[n_models:])
|
||||
|
||||
|
||||
def test_labels_and_order_are_shuffled_not_identity():
|
||||
cat = orchestrator_catalog()
|
||||
orig_models = _model_names(cat)
|
||||
|
||||
orderings = set()
|
||||
labels_for_first = set()
|
||||
for seed in range(25):
|
||||
anon, anon_to_real = anonymize_tools(cat, random.Random(seed))
|
||||
# Real-model order as it appears in the anonymized expert block.
|
||||
order = tuple(anon_to_real[t.name] for t in anon if t.name.startswith("model_"))
|
||||
orderings.add(order)
|
||||
# Label assigned to the first catalog model varies across rngs.
|
||||
rev = {real: lbl for lbl, real in anon_to_real.items()}
|
||||
labels_for_first.add(rev[orig_models[0]])
|
||||
|
||||
# Not a fixed identity: across seeds the expert order actually varies...
|
||||
assert len(orderings) > 1
|
||||
# ...and at least one ordering differs from the input model order.
|
||||
assert any(order != tuple(orig_models) for order in orderings)
|
||||
# Labels are random per-rng, not a stable function of position.
|
||||
assert len(labels_for_first) > 1
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Tests for the faithful ToolOrchestra unified-tool registry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import (
|
||||
CATEGORY_BASIC,
|
||||
CATEGORY_CLOUD_FRONTIER,
|
||||
CATEGORY_LOCAL_OSS,
|
||||
KIND_MODEL,
|
||||
KIND_TOOL,
|
||||
ExpertTool,
|
||||
build_tool_specs,
|
||||
default_catalog,
|
||||
openjarvis_tool,
|
||||
orchestrator_catalog,
|
||||
sample_tool_config,
|
||||
to_worker_dict,
|
||||
tools_by_name,
|
||||
)
|
||||
|
||||
|
||||
def test_each_model_is_its_own_tool():
|
||||
"""Faithful §3.1: one named tool per model, not a meta-tool + slot."""
|
||||
cat = default_catalog()
|
||||
names = {t.name for t in cat}
|
||||
# Distinct model tools, each with its own name.
|
||||
for n in (
|
||||
"gpt_5",
|
||||
"gpt_5_mini",
|
||||
"qwen3_32b",
|
||||
"qwen_2_5_coder_32b_instruct",
|
||||
"llama_3_3_70b_instruct",
|
||||
"claude_opus_4_7",
|
||||
):
|
||||
assert n in names, f"missing model tool {n}"
|
||||
# No meta-tool / slot vocabulary leaks in.
|
||||
assert "answer" not in names and "enhance_reasoning" not in names
|
||||
|
||||
|
||||
def test_catalog_names_unique_and_valid():
|
||||
cat = default_catalog()
|
||||
names = [t.name for t in cat]
|
||||
assert len(names) == len(set(names))
|
||||
assert all(isinstance(t, ExpertTool) for t in cat)
|
||||
|
||||
|
||||
def test_local_model_included_only_when_served():
|
||||
# Local tool is named after the served model ("qwen3:8b" -> "qwen3_8b").
|
||||
assert "qwen3_8b" not in {t.name for t in default_catalog()}
|
||||
cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
|
||||
local = tools_by_name(cat)["qwen3_8b"]
|
||||
assert local.backend_type == "vllm"
|
||||
assert local.base_url == "http://x/v1"
|
||||
assert local.price_in == 0.0 and local.price_out == 0.0
|
||||
|
||||
|
||||
def test_invalid_tool_rejected():
|
||||
with pytest.raises(ValueError):
|
||||
ExpertTool(name="x", kind="bogus", backend_type="openai", summary="", model="m")
|
||||
with pytest.raises(ValueError):
|
||||
ExpertTool(
|
||||
name="x", kind=KIND_MODEL, backend_type="openai", summary="", model=None
|
||||
)
|
||||
|
||||
|
||||
def test_specs_shape_and_pricing_in_description():
|
||||
cat = default_catalog()
|
||||
specs = build_tool_specs(cat)
|
||||
by = {s["function"]["name"]: s for s in specs}
|
||||
gpt5 = by["gpt_5"]
|
||||
assert gpt5["type"] == "function"
|
||||
assert "input" in gpt5["function"]["parameters"]["properties"]
|
||||
# Price table is surfaced in the description (the policy is trained on it).
|
||||
assert "$1.25/1M input" in gpt5["function"]["description"]
|
||||
# Search tool takes a query, code takes code.
|
||||
assert "query" in by["web_search"]["function"]["parameters"]["properties"]
|
||||
assert "code" in by["code_interpreter"]["function"]["parameters"]["properties"]
|
||||
|
||||
|
||||
def test_sample_is_deterministic_and_well_formed():
|
||||
cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
|
||||
a = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
|
||||
b = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
|
||||
assert [t.name for t in a] == [t.name for t in b] # deterministic
|
||||
assert len(a) >= 4
|
||||
assert any(t.kind == KIND_MODEL for t in a) # can reason
|
||||
assert any(t.kind != KIND_MODEL for t in a) # can act
|
||||
assert {t.name for t in a} <= {t.name for t in cat} # subset
|
||||
|
||||
|
||||
def test_price_jitter_changes_prices_reproducibly():
|
||||
cat = default_catalog()
|
||||
base = {
|
||||
t.name: t for t in sample_tool_config(cat, rng=random.Random(3), min_tools=8)
|
||||
}
|
||||
jit = {
|
||||
t.name: t
|
||||
for t in sample_tool_config(
|
||||
cat, rng=random.Random(3), min_tools=8, price_jitter=0.5
|
||||
)
|
||||
}
|
||||
# Same subset (same seed/sequence up to jitter draws), but model prices move.
|
||||
moved = [
|
||||
n
|
||||
for n in base
|
||||
if base[n].kind == KIND_MODEL
|
||||
and base[n].price_in
|
||||
and n in jit
|
||||
and jit[n].price_in != base[n].price_in
|
||||
]
|
||||
assert moved, "expected jitter to change at least one model price"
|
||||
for n in moved:
|
||||
f = jit[n].price_in / base[n].price_in
|
||||
assert 0.5 <= f <= 1.5
|
||||
|
||||
|
||||
# Two cloud-frontier + four local-OSS model tools, in catalog order.
|
||||
_ORCH_MODEL_NAMES = [
|
||||
"gpt_5_5",
|
||||
"claude_opus_4_8",
|
||||
"qwen3_5_9b",
|
||||
"qwen3_6_27b_fp8",
|
||||
"qwen3_5_122b_a10b_fp8",
|
||||
"qwen3_5_397b_a17b_fp8",
|
||||
]
|
||||
|
||||
# Bridged real OpenJarvis tools (basic) appended after web_search/code_interpreter.
|
||||
_ORCH_BASIC_NAMES = [
|
||||
"web_search",
|
||||
"code_interpreter",
|
||||
"calculator",
|
||||
"shell_exec",
|
||||
"file_read",
|
||||
"file_write",
|
||||
"http_request",
|
||||
"think",
|
||||
"apply_patch",
|
||||
"pdf_extract",
|
||||
"db_query",
|
||||
]
|
||||
|
||||
|
||||
def test_orchestrator_catalog_two_model_classes_plus_basics():
|
||||
cat = orchestrator_catalog()
|
||||
names = [t.name for t in cat]
|
||||
# 6 model tools (2 cloud_frontier + 4 local_oss) come first.
|
||||
assert names[:6] == _ORCH_MODEL_NAMES
|
||||
# then the basic tools.
|
||||
assert set(_ORCH_BASIC_NAMES) <= set(names)
|
||||
assert len(cat) == 6 + len(_ORCH_BASIC_NAMES)
|
||||
by = tools_by_name(cat)
|
||||
assert by["gpt_5_5"].category == CATEGORY_CLOUD_FRONTIER
|
||||
assert by["claude_opus_4_8"].category == CATEGORY_CLOUD_FRONTIER
|
||||
# Default routing for every model tool is OpenRouter (no servers required).
|
||||
for n in (
|
||||
"qwen3_5_9b",
|
||||
"qwen3_6_27b_fp8",
|
||||
"qwen3_5_122b_a10b_fp8",
|
||||
"qwen3_5_397b_a17b_fp8",
|
||||
):
|
||||
assert by[n].category == CATEGORY_LOCAL_OSS
|
||||
assert by[n].backend_type == "openrouter"
|
||||
assert by[n].base_url is None
|
||||
# OpenRouter routing carries the slug + a (estimated) per-token price.
|
||||
assert "/" in by[n].model and by[n].price_in > 0.0
|
||||
|
||||
|
||||
def test_orchestrator_catalog_categories_present():
|
||||
cat = orchestrator_catalog()
|
||||
cats = {t.category for t in cat}
|
||||
assert cats == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS, CATEGORY_BASIC}
|
||||
|
||||
|
||||
def test_orchestrator_can_drop_tools():
|
||||
cat = orchestrator_catalog(include_tools=False)
|
||||
assert {t.category for t in cat} == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS}
|
||||
assert len(cat) == 6
|
||||
|
||||
|
||||
def test_orchestrator_specs_include_category_field():
|
||||
specs = build_tool_specs(orchestrator_catalog())
|
||||
by = {s["function"]["name"]: s for s in specs}
|
||||
assert by["gpt_5_5"]["function"]["category"] == CATEGORY_CLOUD_FRONTIER
|
||||
assert by["qwen3_5_9b"]["function"]["category"] == CATEGORY_LOCAL_OSS
|
||||
assert by["web_search"]["function"]["category"] == CATEGORY_BASIC
|
||||
assert by["shell_exec"]["function"]["category"] == CATEGORY_BASIC
|
||||
# Every tool carries a category tag.
|
||||
assert all("category" in s["function"] for s in specs)
|
||||
|
||||
|
||||
def test_orchestrator_local_models_get_base_url_when_provided():
|
||||
# An endpoint switches that model from the OpenRouter default to local vLLM.
|
||||
cat = orchestrator_catalog(
|
||||
local_endpoints={
|
||||
"Qwen/Qwen3.5-9B": "http://x/v1",
|
||||
"Qwen/Qwen3.6-27B-FP8": "http://y/v1",
|
||||
}
|
||||
)
|
||||
by = tools_by_name(cat)
|
||||
assert by["qwen3_5_9b"].backend_type == "vllm"
|
||||
assert by["qwen3_5_9b"].base_url == "http://x/v1"
|
||||
assert by["qwen3_5_9b"].price_in == 0.0 and by["qwen3_5_9b"].price_out == 0.0
|
||||
assert by["qwen3_6_27b_fp8"].base_url == "http://y/v1"
|
||||
# Unmapped local model -> stays on the OpenRouter default (base_url None).
|
||||
assert by["qwen3_5_122b_a10b_fp8"].backend_type == "openrouter"
|
||||
assert by["qwen3_5_122b_a10b_fp8"].base_url is None
|
||||
# Cloud frontier carries real pricing.
|
||||
assert by["claude_opus_4_8"].price_in > 0.0
|
||||
assert by["gpt_5_5"].price_in > 0.0
|
||||
|
||||
|
||||
def test_orchestrator_model_backends_override():
|
||||
# Force a cloud model onto OpenRouter and a local model onto vLLM explicitly.
|
||||
cat = orchestrator_catalog(
|
||||
model_backends={
|
||||
"claude-opus-4-8": "openrouter",
|
||||
"Qwen/Qwen3.5-397B-A17B-FP8": "vllm",
|
||||
},
|
||||
local_endpoints={"Qwen/Qwen3.5-397B-A17B-FP8": "http://z/v1"},
|
||||
)
|
||||
by = tools_by_name(cat)
|
||||
assert by["claude_opus_4_8"].backend_type == "openrouter"
|
||||
assert by["claude_opus_4_8"].model == "anthropic/claude-opus-4.8"
|
||||
# No override for gpt-5.5 -> it defaults to its NATIVE first-party API
|
||||
# (openai), not OpenRouter. Frontier models hit their native provider by
|
||||
# default; OpenRouter is only the fallback for OSS/local models or when
|
||||
# explicitly requested via model_backends.
|
||||
assert by["gpt_5_5"].backend_type == "openai" # native default
|
||||
assert by["gpt_5_5"].model == "gpt-5.5"
|
||||
assert by["qwen3_5_397b_a17b_fp8"].backend_type == "vllm"
|
||||
assert by["qwen3_5_397b_a17b_fp8"].base_url == "http://z/v1"
|
||||
|
||||
|
||||
def test_orchestrator_openrouter_slug_override():
|
||||
cat = orchestrator_catalog(
|
||||
openrouter_slugs={"Qwen/Qwen3.5-9B": "qwen/qwen3.5-9b-custom"}
|
||||
)
|
||||
by = tools_by_name(cat)
|
||||
assert by["qwen3_5_9b"].model == "qwen/qwen3.5-9b-custom"
|
||||
|
||||
|
||||
def test_openjarvis_tool_bridges_real_tool_with_custom_schema():
|
||||
params = {
|
||||
"type": "object",
|
||||
"properties": {"command": {"type": "string"}},
|
||||
"required": ["command"],
|
||||
}
|
||||
t = openjarvis_tool("shell_exec", summary="Run a shell command.", params=params)
|
||||
assert t.kind == KIND_TOOL
|
||||
assert t.backend_type == "openjarvis-tool"
|
||||
assert t.model == "shell_exec"
|
||||
assert t.category == CATEGORY_BASIC
|
||||
spec = t.to_spec()
|
||||
assert spec["function"]["name"] == "shell_exec"
|
||||
assert spec["function"]["parameters"] == params
|
||||
assert spec["function"]["category"] == CATEGORY_BASIC
|
||||
|
||||
|
||||
def test_build_tool_specs_includes_category_for_bridged_tools():
|
||||
specs = build_tool_specs(
|
||||
[
|
||||
openjarvis_tool(
|
||||
"calculator",
|
||||
summary="Math.",
|
||||
params={
|
||||
"type": "object",
|
||||
"properties": {"expression": {"type": "string"}},
|
||||
"required": ["expression"],
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
assert specs[0]["function"]["category"] == CATEGORY_BASIC
|
||||
assert "expression" in specs[0]["function"]["parameters"]["properties"]
|
||||
|
||||
|
||||
def test_to_worker_dict_maps_backend():
|
||||
cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
|
||||
by = tools_by_name(cat)
|
||||
assert to_worker_dict(by["gpt_5"]) == {
|
||||
"name": "gpt_5",
|
||||
"type": "openai",
|
||||
"model": "gpt-5",
|
||||
}
|
||||
local = to_worker_dict(by["qwen3_8b"])
|
||||
assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Offline tests for the reasoning-SFT dataset loaders.
|
||||
|
||||
All loaders take a ``source=`` iterable of raw row dicts so we never touch the
|
||||
network: we hand them fake rows shaped like the real GeneralThought /
|
||||
OpenThoughts3 schemas and assert the normalized :class:`Task` fields.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import (
|
||||
Task,
|
||||
load_generalthought,
|
||||
load_openthoughts,
|
||||
)
|
||||
|
||||
# --- fake rows mirroring the real HF schemas ---------------------------------
|
||||
|
||||
GENERALTHOUGHT_ROWS = [
|
||||
{
|
||||
"question_id": 1,
|
||||
"question": "What is 2 + 2?",
|
||||
"reference_answer": "4",
|
||||
"model_answer": "The answer is 4.",
|
||||
"task": "High School Math",
|
||||
"question_source": "Numina/NuminaMath",
|
||||
},
|
||||
{
|
||||
"question_id": 2,
|
||||
"question": "How do the neural respiratory centers operate?",
|
||||
"reference_answer": "Via the medulla oblongata.",
|
||||
"model_answer": "long answer ...",
|
||||
"task": "NHS QA",
|
||||
"question_source": "CogStack/NHSQA",
|
||||
},
|
||||
{
|
||||
"question_id": 3,
|
||||
"question": "Sort a list.",
|
||||
"reference_answer": "",
|
||||
"model_answer": "use sorted()",
|
||||
"task": "Sorting",
|
||||
"question_source": "BAAI/TACO",
|
||||
},
|
||||
# Should be skipped: empty question and answer.
|
||||
{"question_id": 4, "question": "", "reference_answer": "", "model_answer": ""},
|
||||
]
|
||||
|
||||
OPENTHOUGHTS_ROWS = [
|
||||
{
|
||||
"domain": "code",
|
||||
"source": "stackexchange_codegolf",
|
||||
"difficulty": 7,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "Write a program that prints hi."},
|
||||
{"from": "gpt", "value": "<think>reason</think> Final: print('hi')"},
|
||||
],
|
||||
},
|
||||
{
|
||||
"domain": "math",
|
||||
"source": "Numina/NuminaMath",
|
||||
"difficulty": 5,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What is 6 times 7?"},
|
||||
{
|
||||
"from": "gpt",
|
||||
"value": "<think>6*7=42</think> The answer is \\boxed{42}.",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"domain": "science",
|
||||
"source": "sci",
|
||||
"difficulty": 6,
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What gas do plants release?"},
|
||||
{"from": "gpt", "value": "Oxygen."},
|
||||
],
|
||||
},
|
||||
# Should be skipped: no human/gpt turns.
|
||||
{"domain": "math", "conversations": []},
|
||||
]
|
||||
|
||||
|
||||
def test_load_generalthought_fields_and_domains():
|
||||
tasks = list(load_generalthought(n=10, source=GENERALTHOUGHT_ROWS))
|
||||
assert all(isinstance(t, Task) for t in tasks)
|
||||
# Row 4 (empty) is dropped; remaining 3 yielded.
|
||||
assert len(tasks) == 3
|
||||
by_q = {t.question: t for t in tasks}
|
||||
assert by_q["What is 2 + 2?"].answer == "4"
|
||||
assert by_q["What is 2 + 2?"].domain == "math"
|
||||
assert by_q["What is 2 + 2?"].instruction == "What is 2 + 2?"
|
||||
assert by_q["How do the neural respiratory centers operate?"].domain == "medical"
|
||||
# No reference_answer -> falls back to model_answer; TACO source -> code.
|
||||
assert by_q["Sort a list."].answer == "use sorted()"
|
||||
assert by_q["Sort a list."].domain == "code"
|
||||
|
||||
|
||||
def test_load_generalthought_respects_n():
|
||||
tasks = list(load_generalthought(n=2, source=GENERALTHOUGHT_ROWS))
|
||||
assert len(tasks) == 2
|
||||
|
||||
|
||||
def test_load_openthoughts_fields_and_domains():
|
||||
tasks = list(
|
||||
load_openthoughts(n_code=5, n_math=5, n_science=5, source=OPENTHOUGHTS_ROWS)
|
||||
)
|
||||
assert all(isinstance(t, Task) for t in tasks)
|
||||
assert len(tasks) == 3
|
||||
by_dom = {t.domain: t for t in tasks}
|
||||
assert set(by_dom) == {"code", "math", "science"}
|
||||
# <think> stripped; boxed answer extracted.
|
||||
assert by_dom["math"].answer == "42"
|
||||
assert by_dom["math"].question == "What is 6 times 7?"
|
||||
# Code: think stripped, visible solution kept.
|
||||
assert "print('hi')" in by_dom["code"].answer
|
||||
assert by_dom["science"].answer == "Oxygen."
|
||||
|
||||
|
||||
def test_load_openthoughts_quota_caps_per_domain():
|
||||
rows = [
|
||||
{
|
||||
"domain": "code",
|
||||
"conversations": [
|
||||
{"from": "human", "value": f"q{i}"},
|
||||
{"from": "gpt", "value": f"a{i}"},
|
||||
],
|
||||
}
|
||||
for i in range(10)
|
||||
]
|
||||
tasks = list(load_openthoughts(n_code=3, n_math=0, n_science=0, source=rows))
|
||||
assert len(tasks) == 3
|
||||
assert all(t.domain == "code" for t in tasks)
|
||||
|
||||
|
||||
def test_task_dataclass_instruction_property():
|
||||
t = Task(task_id="x", question="hi?", answer="yes", domain="misc")
|
||||
assert t.instruction == "hi?"
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Offline tests for the rejection-sampling SFT pipeline (unified tools)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import default_catalog, tools_by_name
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import (
|
||||
UnifiedRollout,
|
||||
UnifiedTurn,
|
||||
run_unified_rollout,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
gold_coverage_verify,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
|
||||
normalize_row,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
trajectory_to_record,
|
||||
)
|
||||
|
||||
# A representative raw ToolScale row.
|
||||
_RAW_ROW = {
|
||||
"id": "movie-001",
|
||||
"user_scenario": {
|
||||
"domain": "entertainment",
|
||||
"instructions": {"task_instructions": "Cancel ticket A03 and refund the user."},
|
||||
},
|
||||
"evaluation_criteria": {
|
||||
"actions": [
|
||||
{"name": "cancel", "arguments": {"booking": "A03"}, "action_id": "x1"},
|
||||
{"name": "refund", "arguments": {"user": "8612"}, "action_id": "x2"},
|
||||
],
|
||||
"communicate_info": ["refund amount is $20.90"],
|
||||
"nl_assertions": ["the ticket is cancelled"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_normalize_row():
|
||||
t = normalize_row(_RAW_ROW)
|
||||
assert t.task_id == "movie-001"
|
||||
assert t.domain == "entertainment"
|
||||
assert "Cancel ticket A03" in t.instruction
|
||||
assert t.gold_action_names() == ["cancel", "refund"]
|
||||
assert t.required_info == ["refund amount is $20.90"]
|
||||
|
||||
|
||||
def test_run_unified_rollout_terminates_on_no_tool_call():
|
||||
tools = default_catalog()
|
||||
by = tools_by_name(tools)
|
||||
name = "qwen3_32b"
|
||||
assert name in by
|
||||
|
||||
scripted = [
|
||||
("reason 1\n", [(name, {"input": "do step 1"})], 5, 5),
|
||||
("here is the answer", [], 3, 3), # no tool call -> terminate
|
||||
]
|
||||
calls = iter(scripted)
|
||||
|
||||
def call_orch(messages, specs):
|
||||
# Signature is now (messages, specs): the rollout drives a running
|
||||
# system/user/assistant/tool conversation, not a flattened (system, user).
|
||||
return next(calls)
|
||||
|
||||
def dispatch(tool, args):
|
||||
return (f"OBS for {tool.name}", 0.01, 10, False)
|
||||
|
||||
roll = run_unified_rollout(
|
||||
"What is X?",
|
||||
tools,
|
||||
call_orchestrator=call_orch,
|
||||
dispatch=dispatch,
|
||||
max_turns=5,
|
||||
)
|
||||
assert roll.final_answer == "here is the answer"
|
||||
assert roll.num_tool_calls == 1
|
||||
assert roll.tool_calls() == [(name, {"input": "do step 1"})]
|
||||
assert abs(roll.cost_usd - 0.01) < 1e-9
|
||||
|
||||
|
||||
def test_serialize_record_shape_and_tool_call_tags():
|
||||
tools = default_catalog()
|
||||
roll = UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn(
|
||||
reasoning="think",
|
||||
tool_name="qwen3_32b",
|
||||
arguments={"input": "q"},
|
||||
observation="obs",
|
||||
),
|
||||
# The final turn's reasoning IS the model's real final output (it
|
||||
# already contains the answer). The serializer now renders this turn
|
||||
# via _final_answer_block(turn.reasoning), not rollout.final_answer.
|
||||
UnifiedTurn(
|
||||
reasoning="The result is 42.\nFINAL_ANSWER: 42", tool_name=None
|
||||
),
|
||||
],
|
||||
final_answer="42",
|
||||
cost_usd=0.02,
|
||||
tokens=30,
|
||||
num_tool_calls=1,
|
||||
)
|
||||
rec = trajectory_to_record("t1", "Q?", tools, roll, reward=0.5, domain="math")
|
||||
roles = [m["role"] for m in rec["conversations"]]
|
||||
assert roles[0] == "system" and roles[1] == "user"
|
||||
assert "tool" in roles and roles[-1] == "assistant"
|
||||
# Tool call is emitted as a <tool_call> tag (what the parser reads back).
|
||||
assert any(
|
||||
"<tool_call>" in m["content"] and "qwen3_32b" in m["content"]
|
||||
for m in rec["conversations"]
|
||||
if m["role"] == "assistant"
|
||||
)
|
||||
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
|
||||
assert rec["reward"] == 0.5 and rec["domain"] == "math"
|
||||
|
||||
|
||||
def test_gold_coverage_verify():
|
||||
t = normalize_row(_RAW_ROW)
|
||||
good = UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
|
||||
UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
|
||||
],
|
||||
final_answer="done",
|
||||
)
|
||||
missing = UnifiedRollout(
|
||||
turns=[UnifiedTurn("", "cancel", {}, "ok")], final_answer="done"
|
||||
)
|
||||
empty_ans = UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn("", "cancel", {}, "ok"),
|
||||
UnifiedTurn("", "refund", {}, "ok"),
|
||||
],
|
||||
final_answer="",
|
||||
)
|
||||
assert gold_coverage_verify(t, good) is True
|
||||
assert gold_coverage_verify(t, missing) is False
|
||||
assert gold_coverage_verify(t, empty_ans) is False
|
||||
|
||||
|
||||
def test_generate_sft_dataset_end_to_end(tmp_path):
|
||||
tools = default_catalog()
|
||||
tasks = [
|
||||
normalize_row(_RAW_ROW),
|
||||
normalize_row(
|
||||
{
|
||||
**_RAW_ROW,
|
||||
"id": "unsolvable",
|
||||
"evaluation_criteria": {"actions": [{"name": "never_called"}]},
|
||||
}
|
||||
),
|
||||
]
|
||||
|
||||
def rollout_fn(task):
|
||||
# Solve the first task; always miss the gold action of the second.
|
||||
if task.task_id == "movie-001":
|
||||
return UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
|
||||
UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
|
||||
UnifiedTurn("done", None),
|
||||
],
|
||||
# num_tool_calls must reflect the two routed calls: the structural
|
||||
# _target_is_clean gate drops a trajectory with num_tool_calls < 1.
|
||||
final_answer="refunded $20.90",
|
||||
cost_usd=0.03,
|
||||
num_tool_calls=2,
|
||||
)
|
||||
return UnifiedRollout(
|
||||
turns=[UnifiedTurn("x", "cancel", {}, "ok")],
|
||||
final_answer="nope",
|
||||
cost_usd=0.05,
|
||||
)
|
||||
|
||||
out = tmp_path / "sft.jsonl"
|
||||
stats = generate_sft_dataset(
|
||||
str(out),
|
||||
tasks=tasks,
|
||||
tools=tools,
|
||||
rollout_fn=rollout_fn,
|
||||
samples_per_task=2,
|
||||
)
|
||||
assert stats["tasks_seen"] == 2
|
||||
assert stats["records_written"] == 1 # only the solvable task
|
||||
assert stats["tasks_dropped"] == 1
|
||||
lines = out.read_text().strip().splitlines()
|
||||
assert len(lines) == 1
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["task_id"] == "movie-001"
|
||||
assert rec["domain"] == "entertainment"
|
||||
assert (tmp_path / "sft.jsonl.stats.json").exists()
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Offline tests for the answer verifier.
|
||||
|
||||
Math and code paths run with no network. The Gemini judge is only reachable for
|
||||
non-math/code domains; we monkeypatch ``_gemini_judge`` (or rely on its no-key
|
||||
fallback) so nothing hits the network.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data import verify as V
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
|
||||
make_verifier,
|
||||
verify_answer,
|
||||
)
|
||||
|
||||
|
||||
def _math(answer: str) -> Task:
|
||||
return Task(task_id="m", question="q", answer=answer, domain="math")
|
||||
|
||||
|
||||
def test_math_boxed_vs_plain_true():
|
||||
assert verify_answer(_math("42"), "\\boxed{42}") is True
|
||||
|
||||
|
||||
def test_math_boxed_both_sides():
|
||||
assert verify_answer(_math("\\boxed{42}"), "The answer is \\boxed{42}") is True
|
||||
|
||||
|
||||
def test_math_wrong_number_false():
|
||||
assert verify_answer(_math("42"), "41") is False
|
||||
|
||||
|
||||
def test_math_numeric_float_match():
|
||||
assert verify_answer(_math("42"), "42.0") is True
|
||||
|
||||
|
||||
def test_math_symbolic_equivalence_no_longer_verified_locally():
|
||||
# The sympy symbolic-equality block was removed from _math_equal: sympy.simplify
|
||||
# / parse_latex could hang on pathological \boxed{} answers while holding the GIL
|
||||
# and deadlock the threaded rejection sampler. The math domain now uses
|
||||
# string+numeric+f1 only (no judge call either), so a purely SYMBOLIC
|
||||
# equivalence that isn't string/numeric-equal is intentionally NOT recognized.
|
||||
t = _math("(x+1)^2")
|
||||
assert verify_answer(t, "x^2 + 2*x + 1") is False
|
||||
|
||||
|
||||
def test_empty_prediction_false():
|
||||
assert verify_answer(_math("42"), "") is False
|
||||
|
||||
|
||||
def test_code_short_gold_substring():
|
||||
t = Task(task_id="c", question="print hi", answer="hello", domain="code")
|
||||
assert verify_answer(t, "the output is hello") is True
|
||||
|
||||
|
||||
def test_code_falls_back_to_judge(monkeypatch):
|
||||
# Long gold + no substring match -> hits judge; monkeypatch to PASS.
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: True)
|
||||
t = Task(
|
||||
task_id="c",
|
||||
question="solve",
|
||||
answer="x" * 300, # long -> skips substring shortcut
|
||||
domain="code",
|
||||
)
|
||||
assert verify_answer(t, "totally different") is True
|
||||
|
||||
|
||||
def test_judge_domain_uses_gemini(monkeypatch):
|
||||
calls = {}
|
||||
|
||||
def fake_judge(task, pred):
|
||||
calls["hit"] = (task.domain, pred)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(V, "_gemini_judge", fake_judge)
|
||||
t = Task(task_id="s", question="q", answer="Oxygen", domain="science")
|
||||
assert verify_answer(t, "plants release oxygen") is True
|
||||
assert calls["hit"][0] == "science"
|
||||
|
||||
|
||||
def test_judge_domain_fallback_f1_when_no_key(monkeypatch):
|
||||
# Judge unavailable (returns None) -> falls back to string / F1.
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
|
||||
t = Task(task_id="s", question="q", answer="the capital is paris", domain="misc")
|
||||
assert verify_answer(t, "I think the capital is paris indeed") is True
|
||||
assert verify_answer(t, "completely unrelated text here") is False
|
||||
|
||||
|
||||
def test_make_verifier_reads_final_answer(monkeypatch):
|
||||
monkeypatch.setattr(V, "_gemini_judge", lambda task, pred: None)
|
||||
verify = make_verifier()
|
||||
|
||||
class Rollout:
|
||||
final_answer = "\\boxed{42}"
|
||||
|
||||
assert verify(_math("42"), Rollout()) is True
|
||||
|
||||
class BadRollout:
|
||||
final_answer = "999"
|
||||
|
||||
assert verify(_math("42"), BadRollout()) is False
|
||||
|
||||
|
||||
def test_make_verifier_missing_final_answer():
|
||||
verify = make_verifier()
|
||||
|
||||
class Empty:
|
||||
pass
|
||||
|
||||
assert verify(_math("42"), Empty()) is False
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Offline smoke test for the v1 orchestrator-SFT driver (base self-sampling).
|
||||
|
||||
No network / no GPU: we feed fake ``datasets.Task`` objects, a canned
|
||||
``UnifiedRollout`` (via a fake ``rollout_fn``), and an accept-all verifier, then
|
||||
assert ``generate_sft_dataset`` writes a JSONL record in the expected
|
||||
``conversations`` shape that the SFT trainer consumes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import Task
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
generate_sft_dataset,
|
||||
)
|
||||
|
||||
|
||||
def _fake_tasks(**_kwargs) -> list[Task]:
|
||||
# Accepts **kwargs so it can stand in for ``load_sft_tasks(cap=..., balanced=...)``.
|
||||
return [
|
||||
Task(task_id="t-math-1", question="What is 6 * 7?", answer="42", domain="math"),
|
||||
Task(task_id="t-code-1", question="Reverse 'ab'.", answer="ba", domain="code"),
|
||||
]
|
||||
|
||||
|
||||
def _canned_rollout(task: Task) -> UnifiedRollout:
|
||||
# One tool turn + a final-answer turn that echoes the gold answer.
|
||||
return UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn(
|
||||
reasoning="let me compute",
|
||||
tool_name="code_interpreter",
|
||||
arguments={"code": "print(6*7)"},
|
||||
observation="42",
|
||||
),
|
||||
UnifiedTurn(
|
||||
reasoning=(
|
||||
f"The tool returned {task.answer}. FINAL_ANSWER: {task.answer}"
|
||||
),
|
||||
tool_name=None,
|
||||
),
|
||||
],
|
||||
final_answer=task.answer,
|
||||
cost_usd=0.01,
|
||||
tokens=20,
|
||||
num_tool_calls=1,
|
||||
)
|
||||
|
||||
|
||||
def test_build_v1_writes_expected_conversations_shape(tmp_path, monkeypatch):
|
||||
tools = orchestrator_catalog() # specialists unwired (math/coder endpoints None)
|
||||
|
||||
# Accept-all verifier (the real make_verifier may hit Gemini / math checkers).
|
||||
monkeypatch.setattr(
|
||||
"openjarvis.learning.intelligence.orchestrator.sft_data.verify.make_verifier",
|
||||
lambda: lambda task, rollout: True,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
|
||||
make_verifier,
|
||||
)
|
||||
|
||||
out = tmp_path / "orchestrator_sft_v1.jsonl"
|
||||
stats = generate_sft_dataset(
|
||||
str(out),
|
||||
tasks=_fake_tasks(),
|
||||
tools=tools,
|
||||
rollout_fn=_canned_rollout,
|
||||
verify_fn=make_verifier(),
|
||||
samples_per_task=2,
|
||||
max_keep_per_task=1,
|
||||
reward_fn=lambda r: -r.cost_usd,
|
||||
)
|
||||
|
||||
assert stats["tasks_seen"] == 2
|
||||
assert stats["records_written"] == 2
|
||||
assert stats["tasks_dropped"] == 0
|
||||
|
||||
lines = out.read_text().strip().splitlines()
|
||||
assert len(lines) == 2
|
||||
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["task_id"] == "t-math-1"
|
||||
assert rec["domain"] == "math"
|
||||
|
||||
roles = [m["role"] for m in rec["conversations"]]
|
||||
assert roles[0] == "system" and roles[1] == "user"
|
||||
assert "tool" in roles
|
||||
assert roles[-1] == "assistant"
|
||||
# Tool call emitted as a <tool_call> tag the parser reads back.
|
||||
assert any(
|
||||
"<tool_call>" in m["content"] and "code_interpreter" in m["content"]
|
||||
for m in rec["conversations"]
|
||||
if m["role"] == "assistant"
|
||||
)
|
||||
assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
|
||||
assert (tmp_path / "orchestrator_sft_v1.jsonl.stats.json").exists()
|
||||
|
||||
|
||||
def _load_driver():
|
||||
"""Load the (non-package) CLI driver module by path."""
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
path = root / "scripts" / "orchestrator" / "build_orchestrator_sft.py"
|
||||
spec = importlib.util.spec_from_file_location("build_orchestrator_sft", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_driver_runs_end_to_end_with_fakes(tmp_path, monkeypatch):
|
||||
"""Drive ``main`` with no network: fake task loader, rollout, and verifier."""
|
||||
drv = _load_driver()
|
||||
|
||||
monkeypatch.setattr(drv, "load_sft_tasks", _fake_tasks)
|
||||
monkeypatch.setattr(
|
||||
drv,
|
||||
"run_unified_rollout",
|
||||
lambda question, tools, **kw: _canned_rollout(_fake_tasks()[0]),
|
||||
)
|
||||
monkeypatch.setattr(drv, "make_verifier", lambda: lambda task, rollout: True)
|
||||
# make_call_orchestrator builds an OpenAI client lazily, so it never connects
|
||||
# here (run_unified_rollout is stubbed out).
|
||||
|
||||
# The driver treats --out as a TAG and always writes to
|
||||
# <OJ_DATA_ROOT>/raw/<label>_<MMDD>[_<tag>]/data.jsonl, defaulting to a
|
||||
# cwd-relative data/orchestrator. Drop OJ_DATA_ROOT (the repo .env exports
|
||||
# it) so the run folder lands in the test's temp dir and not the real
|
||||
# experiments tree, then chdir there.
|
||||
monkeypatch.delenv("OJ_DATA_ROOT", raising=False)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
rc = drv.main(
|
||||
[
|
||||
"--out",
|
||||
"v1",
|
||||
"--samples-per-task",
|
||||
"1",
|
||||
"--max-tasks",
|
||||
"2",
|
||||
]
|
||||
)
|
||||
assert rc == 0
|
||||
produced = list(
|
||||
(tmp_path / "data" / "orchestrator" / "raw").glob("*_v1/data.jsonl")
|
||||
)
|
||||
assert len(produced) == 1
|
||||
lines = produced[0].read_text().strip().splitlines()
|
||||
assert len(lines) == 2
|
||||
rec = json.loads(lines[0])
|
||||
assert rec["conversations"][0]["role"] == "system"
|
||||
assert "FINAL_ANSWER" in rec["conversations"][-1]["content"]
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Offline tests for the orchestrator eval backend + the eval CLI's
|
||||
benchmark-name normalization.
|
||||
|
||||
No network / models: ``run_unified_rollout`` is monkeypatched to return a
|
||||
canned rollout, and the dataset-key mapping is asserted against the real
|
||||
``openjarvis.evals.cli`` registry without instantiating any (network-bound)
|
||||
dataset.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedTurn
|
||||
from openjarvis.learning.intelligence.orchestrator import eval_backend as eb
|
||||
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
|
||||
OrchestratorBackend,
|
||||
)
|
||||
|
||||
|
||||
class _CannedRollout:
|
||||
"""Minimal stand-in for UnifiedRollout."""
|
||||
|
||||
def __init__(self, final_answer: str) -> None:
|
||||
self.final_answer = final_answer
|
||||
self.cost_usd = 0.123
|
||||
self.tokens = 42
|
||||
self.num_tool_calls = 1
|
||||
self.parse_failures = 0
|
||||
self.anon_map = {}
|
||||
# generate_full now serializes a per-turn trace (reasoning / tool_name /
|
||||
# arguments / observation), so turns must be real UnifiedTurn-shaped
|
||||
# objects, not bare object() placeholders.
|
||||
self.turns = [
|
||||
UnifiedTurn(
|
||||
reasoning="let me search",
|
||||
tool_name="web_search",
|
||||
arguments={"query": "x"},
|
||||
observation="result",
|
||||
),
|
||||
]
|
||||
|
||||
def tool_calls(self):
|
||||
return [("web_search", {"query": "x"})]
|
||||
|
||||
|
||||
def _patch_rollout(monkeypatch, rollout):
|
||||
"""Patch run_unified_rollout where the backend looks it up."""
|
||||
monkeypatch.setattr(
|
||||
eb._rollout_mod,
|
||||
"run_unified_rollout",
|
||||
lambda *a, **k: rollout,
|
||||
)
|
||||
|
||||
|
||||
def test_construct_is_network_free():
|
||||
"""Constructing the backend builds the catalog but touches no network."""
|
||||
backend = OrchestratorBackend()
|
||||
assert backend.backend_id == "orchestrator"
|
||||
assert backend.orchestrator_endpoint == eb.DEFAULT_ENDPOINT
|
||||
assert backend.orchestrator_model == eb.DEFAULT_MODEL
|
||||
# Catalog is non-empty (cloud frontier + local OSS + basic tools).
|
||||
assert len(backend._tools) > 0
|
||||
names = {t.name for t in backend._tools}
|
||||
assert "web_search" in names
|
||||
|
||||
|
||||
def test_local_endpoints_mapping():
|
||||
"""local_endpoints (model-id -> base_url) wires the local OSS tool."""
|
||||
backend = OrchestratorBackend(
|
||||
local_endpoints={"Qwen/Qwen3.5-9B": "http://m:9/v1"},
|
||||
)
|
||||
assert backend.local_endpoints["Qwen/Qwen3.5-9B"] == "http://m:9/v1"
|
||||
# The corresponding local-OSS tool should carry that base_url.
|
||||
tool = next(t for t in backend._tools if t.model == "Qwen/Qwen3.5-9B")
|
||||
assert tool.base_url == "http://m:9/v1"
|
||||
|
||||
|
||||
def test_generate_returns_final_answer(monkeypatch):
|
||||
"""generate() returns the rollout's final_answer string."""
|
||||
_patch_rollout(monkeypatch, _CannedRollout("42"))
|
||||
backend = OrchestratorBackend()
|
||||
out = backend.generate("what is 6*7?", model="qwen3-8b")
|
||||
assert out == "42"
|
||||
|
||||
|
||||
def test_generate_full_payload(monkeypatch):
|
||||
"""generate_full() returns the runner-expected payload shape."""
|
||||
_patch_rollout(monkeypatch, _CannedRollout("42"))
|
||||
backend = OrchestratorBackend()
|
||||
full = backend.generate_full("q", model="qwen3-8b")
|
||||
assert full["content"] == "42"
|
||||
assert full["cost_usd"] == 0.123
|
||||
assert full["usage"]["completion_tokens"] == 42
|
||||
assert full["tool_calls"] == 1
|
||||
assert full["turn_count"] == 1
|
||||
assert full["framework"] == "openjarvis-orchestrator"
|
||||
assert "error" not in full
|
||||
assert full["latency_seconds"] >= 0.0
|
||||
|
||||
|
||||
def test_generate_full_error_path(monkeypatch):
|
||||
"""An exception inside the rollout becomes a recorded error, not a crash."""
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise RuntimeError("vllm down")
|
||||
|
||||
monkeypatch.setattr(eb._rollout_mod, "run_unified_rollout", _boom)
|
||||
backend = OrchestratorBackend()
|
||||
full = backend.generate_full("q", model="qwen3-8b")
|
||||
assert full["content"] == ""
|
||||
assert "vllm down" in full["error"]
|
||||
assert full["usage"]["completion_tokens"] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Eval-script benchmark name mapping
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# The 5 benchmarks the eval script targets, by the script's (alias) names.
|
||||
SCRIPT_BENCHMARKS = [
|
||||
"gaia",
|
||||
"terminalbench_v2_1",
|
||||
"taubench",
|
||||
"mmlu_pro",
|
||||
"supergpqa",
|
||||
]
|
||||
|
||||
|
||||
def _load_script():
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
script_path = repo_root / "scripts" / "orchestrator" / "eval_orchestrator.py"
|
||||
spec = importlib.util.spec_from_file_location("eval_orchestrator", script_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
def test_benchmark_aliases_map_to_registry_keys():
|
||||
"""Each script benchmark normalizes to a real evals.cli BENCHMARKS key."""
|
||||
from openjarvis.evals.cli import BENCHMARKS
|
||||
|
||||
script = _load_script()
|
||||
for name in SCRIPT_BENCHMARKS:
|
||||
key = script._normalize_benchmark(name)
|
||||
assert key in BENCHMARKS, f"{name} -> {key} not in BENCHMARKS registry"
|
||||
|
||||
|
||||
def test_normalize_specific_keys():
|
||||
"""Spot-check the underscored aliases normalize to the dotted/dashed keys."""
|
||||
script = _load_script()
|
||||
assert script._normalize_benchmark("terminalbench_v2_1") == "terminalbench-v2.1"
|
||||
assert script._normalize_benchmark("mmlu_pro") == "mmlu-pro"
|
||||
assert script._normalize_benchmark("gaia") == "gaia"
|
||||
|
||||
|
||||
def test_build_dataset_accepts_keys():
|
||||
"""_build_dataset must recognise the normalized keys.
|
||||
|
||||
We only assert it doesn't raise the 'Unknown benchmark' ClickException;
|
||||
datasets that need network/files at construction are xfail-skipped.
|
||||
"""
|
||||
import click
|
||||
|
||||
from openjarvis.evals.cli import _build_dataset
|
||||
|
||||
script = _load_script()
|
||||
for name in SCRIPT_BENCHMARKS:
|
||||
key = script._normalize_benchmark(name)
|
||||
try:
|
||||
_build_dataset(key)
|
||||
except click.ClickException as exc:
|
||||
if "Unknown benchmark" in str(exc):
|
||||
pytest.fail(f"{key} not recognised by _build_dataset")
|
||||
# Other ClickExceptions (missing files/creds) are acceptable here.
|
||||
except Exception:
|
||||
# Construction may require network/files — that's fine; the point
|
||||
# is the key is recognised (didn't hit the unknown-benchmark else).
|
||||
pytest.skip(f"{key} dataset construction needs resources")
|
||||
@@ -0,0 +1,120 @@
|
||||
"""Control-token sanitation in the SFT serializer + clean gate.
|
||||
|
||||
Regression guard for the leak audit: leaked control/special tokens
|
||||
(``<|im_end|>``, ``<|tool_call>``, gemma ``<start_of_turn>``/``<end_of_turn>``,
|
||||
``<|"|>`` …) used to survive into the supervised final answer because the clean
|
||||
gate only rejected the bare ``<tool_call>`` form. Defense in depth now:
|
||||
|
||||
1. ``_final_answer_block`` STRIPS residual control tokens so a good answer with a
|
||||
stray token is salvaged (not dropped).
|
||||
2. ``_target_is_clean`` REJECTS a rollout whose final answer is nothing but
|
||||
control tokens, or where a token survives the strip.
|
||||
|
||||
Both must leave legitimate ``<``/``>`` in math/code untouched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
|
||||
_target_is_clean,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
|
||||
_CONTROL_TOKEN_RE,
|
||||
_final_answer_block,
|
||||
_strip_control_tokens,
|
||||
)
|
||||
|
||||
# --- serializer strips control tokens, leaving a clean FINAL_ANSWER line ------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw, expected_answer",
|
||||
[
|
||||
("FINAL_ANSWER: 42<end_of_turn>", "42"),
|
||||
("FINAL_ANSWER: 42<|end_of_turn|>", "42"),
|
||||
("FINAL_ANSWER: 42<|im_end|>", "42"),
|
||||
("FINAL_ANSWER: 42<|eot_id|>", "42"),
|
||||
("FINAL_ANSWER: The result<|tool_call>", "The result"),
|
||||
('FINAL_ANSWER: foo <|"|> bar', "foo bar"),
|
||||
("<start_of_turn>model", "model"),
|
||||
],
|
||||
)
|
||||
def test_final_answer_block_strips_tokens(raw: str, expected_answer: str) -> None:
|
||||
out = _final_answer_block(raw)
|
||||
assert out == f"FINAL_ANSWER: {expected_answer}"
|
||||
# No control token survives the serializer.
|
||||
assert not _CONTROL_TOKEN_RE.search(out), out
|
||||
|
||||
|
||||
def test_bare_im_end_only_becomes_empty_answer() -> None:
|
||||
# An answer that is nothing but a control token strips to empty; the clean
|
||||
# gate (below) is what rejects such a rollout.
|
||||
assert _final_answer_block("<|im_end|>") == "FINAL_ANSWER: "
|
||||
|
||||
|
||||
def test_math_and_code_angle_brackets_survive() -> None:
|
||||
# The narrow regex must NOT eat legitimate comparisons / generics.
|
||||
for good in [
|
||||
"FINAL_ANSWER: x < 3 and y > 2",
|
||||
"FINAL_ANSWER: List<int> and a<b>c",
|
||||
"FINAL_ANSWER: if a < b: return a > 0",
|
||||
]:
|
||||
out = _final_answer_block(good)
|
||||
assert out == good, out
|
||||
assert not _CONTROL_TOKEN_RE.search(out)
|
||||
|
||||
|
||||
def test_strip_helper_handles_nested_and_whitespace() -> None:
|
||||
assert _strip_control_tokens(" hi <|im_end|> ") == "hi"
|
||||
assert _strip_control_tokens("a<|im_start|>b<|im_end|>c") == "abc"
|
||||
|
||||
|
||||
# --- clean gate: salvage vs reject -------------------------------------------
|
||||
|
||||
|
||||
def _roll(final_answer: str) -> UnifiedRollout:
|
||||
"""A minimal well-formed rollout (one expert call, clean obs) whose only
|
||||
variable is the final answer."""
|
||||
return UnifiedRollout(
|
||||
turns=[
|
||||
UnifiedTurn(
|
||||
reasoning="route it",
|
||||
tool_name="expert_a",
|
||||
arguments={"input": "q"},
|
||||
observation="valid observation",
|
||||
),
|
||||
UnifiedTurn(reasoning=final_answer, tool_name=None),
|
||||
],
|
||||
final_answer=final_answer,
|
||||
cost_usd=0.01,
|
||||
tokens=10,
|
||||
num_tool_calls=1,
|
||||
anon_map={"expert_a": "gpt"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"final_answer, clean",
|
||||
[
|
||||
# Salvageable: a good answer with a trailing stray token stays clean
|
||||
# (the serializer strips the token from the emitted target).
|
||||
("42<end_of_turn>", True),
|
||||
("42<|end_of_turn|>", True),
|
||||
("The result<|tool_call>", True),
|
||||
('foo <|"|> bar', True),
|
||||
("<start_of_turn>model", True),
|
||||
("x < 3 and y > 2", True),
|
||||
("a normal plain answer", True),
|
||||
# Unsalvageable: the answer is nothing but a control token -> rejected.
|
||||
("<|im_end|>", False),
|
||||
("<|eot_id|>", False),
|
||||
("<end_of_turn>", False),
|
||||
# Bare tool-call tag: already rejected by the existing gate.
|
||||
("<tool_call>{}</tool_call>", False),
|
||||
],
|
||||
)
|
||||
def test_clean_gate_salvages_or_rejects(final_answer: str, clean: bool) -> None:
|
||||
assert _target_is_clean(_roll(final_answer)) is clean
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Tests for the cost-aware GRPO reward (offline, no GPU)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.reward import (
|
||||
CostAwareReward,
|
||||
lambda_sweep,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeEpisode:
|
||||
"""Minimal stand-in exposing the two fields CostAwareReward reads."""
|
||||
|
||||
correct: bool
|
||||
total_cost_usd: float = 0.0
|
||||
|
||||
|
||||
class TestCostAwareReward:
|
||||
def test_correct_zero_cost_is_plus_one(self):
|
||||
r = CostAwareReward(lam=0.0)
|
||||
assert r.compute(FakeEpisode(correct=True, total_cost_usd=0.0)) == 1.0
|
||||
|
||||
def test_incorrect_zero_cost_is_minus_one(self):
|
||||
r = CostAwareReward(lam=0.0)
|
||||
assert r.compute(FakeEpisode(correct=False, total_cost_usd=0.0)) == -1.0
|
||||
|
||||
def test_correct_with_lam_penalty_when_cost(self):
|
||||
# base=+1, penalty = lam * cost / cost_max = 0.2 * 0.05 / 0.10 = 0.1
|
||||
r = CostAwareReward(lam=0.2, cost_max=0.10)
|
||||
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.05))
|
||||
assert abs(val - (1.0 - 0.1)) < 1e-9
|
||||
|
||||
def test_lam_zero_ignores_cost(self):
|
||||
r = CostAwareReward(lam=0.0, cost_max=0.10)
|
||||
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.5))
|
||||
assert val == 1.0
|
||||
|
||||
def test_cost_at_budget_full_penalty(self):
|
||||
# cost == cost_max → penalty == lam exactly.
|
||||
r = CostAwareReward(lam=0.4, cost_max=0.10)
|
||||
val = r.compute(FakeEpisode(correct=True, total_cost_usd=0.10))
|
||||
assert abs(val - (1.0 - 0.4)) < 1e-9
|
||||
|
||||
def test_incorrect_also_penalised(self):
|
||||
r = CostAwareReward(lam=0.2, cost_max=0.10)
|
||||
val = r.compute(FakeEpisode(correct=False, total_cost_usd=0.05))
|
||||
assert abs(val - (-1.0 - 0.1)) < 1e-9
|
||||
|
||||
def test_breakdown_fields(self):
|
||||
r = CostAwareReward(lam=0.2, cost_max=0.10)
|
||||
b = r.compute_with_breakdown(FakeEpisode(correct=True, total_cost_usd=0.05))
|
||||
assert set(b) == {"correct", "base", "cost_term", "reward"}
|
||||
assert b["correct"] == 1.0
|
||||
assert b["base"] == 1.0
|
||||
assert abs(b["cost_term"] - (-0.1)) < 1e-9
|
||||
assert abs(b["reward"] - 0.9) < 1e-9
|
||||
|
||||
def test_breakdown_reward_matches_compute(self):
|
||||
r = CostAwareReward(lam=0.3, cost_max=0.10)
|
||||
ep = FakeEpisode(correct=False, total_cost_usd=0.07)
|
||||
assert abs(r.compute_with_breakdown(ep)["reward"] - r.compute(ep)) < 1e-12
|
||||
|
||||
def test_compute_batch(self):
|
||||
r = CostAwareReward(lam=0.0)
|
||||
eps = [FakeEpisode(correct=True), FakeEpisode(correct=False)]
|
||||
assert r.compute_batch(eps) == [1.0, -1.0]
|
||||
|
||||
|
||||
class TestLambdaSweep:
|
||||
def test_sweep_length(self):
|
||||
values = [0.0, 0.05, 0.1, 0.2, 0.4]
|
||||
rewards = lambda_sweep(values)
|
||||
assert len(rewards) == len(values)
|
||||
|
||||
def test_sweep_assigns_lambdas(self):
|
||||
values = [0.0, 0.05, 0.1, 0.2, 0.4]
|
||||
rewards = lambda_sweep(values, cost_max=0.25)
|
||||
assert [r.lam for r in rewards] == values
|
||||
assert all(r.cost_max == 0.25 for r in rewards)
|
||||
|
||||
def test_sweep_returns_costaware(self):
|
||||
rewards = lambda_sweep([0.1])
|
||||
assert isinstance(rewards[0], CostAwareReward)
|
||||
@@ -15,7 +15,7 @@ from openjarvis.learning.intelligence.orchestrator.sft_trainer import (
|
||||
class TestOrchestratorSFTConfig:
|
||||
def test_defaults(self):
|
||||
cfg = OrchestratorSFTConfig()
|
||||
assert cfg.model_name == "Qwen/Qwen3-1.7B"
|
||||
assert cfg.model_name == "Qwen/Qwen3.5-9B"
|
||||
assert cfg.num_epochs == 3
|
||||
assert cfg.batch_size == 8
|
||||
assert cfg.learning_rate == 2e-5
|
||||
|
||||
Reference in New Issue
Block a user