admin/OpenJarvis
1
0
Fork 0
mirror of https://github.com/open-jarvis/OpenJarvis.git synced 2026-07-27 21:05:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

orchestrator: expand tool catalog, fix shell_exec confirmation gate, wire SFT data pipeline

Browse Source 
Tools:
- Bridge 4 more real OpenJarvis tools into the orchestrator catalog
  (think, apply_patch, pdf_extract, db_query); catalog is now 6 models
  + 11 basic tools.
- Fix the dispatch path so confirmation-gated tools actually run: the
  ToolExecutor was built with no confirm callback, so shell_exec (and
  any requires_confirmation tool) returned a "requires confirmation"
  error instead of executing -- which silently broke TerminalBench.
  Headless eval/rollouts now auto-approve.
- Drop dead-end candidates: git_* (need the unbuilt openjarvis_rust
  extension and just duplicate shell_exec), browser_* (needs Playwright),
  and knowledge_search/sql/retrieval (personal-data RAG, empty on the
  academic benchmarks).

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

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

Ignore generated data/ artifacts.
This commit is contained in:
Andrew Park
2026-07-11 12:23:46 -07:00
parent 5aa16473ff
commit f3d9705961
40 changed files with 3707 additions and 1250 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -130,3 +130,6 @@ learning.db
minion_logs/
*.oj-debug.json
oj-debug.*.json
# Generated orchestrator SFT/eval data artifacts
data/
scripts/orchestrator/build_orchestrator_sft.py
+135
View File
@@ -0,0 +1,135 @@
#!/usr/bin/env python
"""Build orchestrator SFT cold-start data by **base Qwen3-8B self-sampling** (v1).
The orchestrator *is* the local Qwen3-8B. We roll the base (un-trained) model out
over the reasoning-SFT task pool (``load_sft_tasks``: GeneralThought + OpenThoughts
code/math/science), verify each trajectory's final answer against the gold
(``verify.make_verifier`` — math/code checkers + Gemini judge fallback), and keep
the cheapest-correct trajectory per task. Those passing trajectories become the
``conversations`` JSONL the SFT trainer consumes — i.e. the model learns from its
own successful rollouts (rejection-sampling cold-start, STaR-style).
This is **v1** (base self-sampling): point ``--orchestrator-endpoint`` at the vLLM
server hosting the *base* ``Qwen/Qwen3-8B``.
For **v2**, train on the v1 data, serve the v1 checkpoint with vLLM, then re-run
this exact driver with ``--orchestrator-endpoint`` pointed at the v1 checkpoint's
endpoint, ``--orchestrator-model`` set to the checkpoint name, and
``--samples-per-task 8`` — the (now stronger) policy self-samples a v2 dataset.
The math/coder specialist tools are only wired when ``--math-endpoint`` /
``--coder-endpoint`` are passed (those are separately-served vLLM specialists);
otherwise they're omitted so the orchestrator only sees the tools it can actually
call.
Example (v1, base self-sampling):
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
--out data/orchestrator_sft_v1.jsonl \
--orchestrator-endpoint http://localhost:8001/v1 \
--orchestrator-model qwen3-8b \
--samples-per-task 8 --max-keep-per-task 1
Example (v2, re-point at the v1 checkpoint):
.venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \
--out data/orchestrator_sft_v2.jsonl \
--orchestrator-endpoint http://localhost:8010/v1 \
--orchestrator-model orchestrator-sft-v1 \
--samples-per-task 8
"""
from __future__ import annotations
import argparse
import json
import logging
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
)
p.add_argument("--out", default="data/orchestrator_sft_v1.jsonl")
# The orchestrator == the local Qwen3-8B self-sampling over its own rollouts.
p.add_argument("--orchestrator-endpoint", default="http://localhost:8001/v1",
help="OpenAI-compatible vLLM base URL serving the orchestrator.")
p.add_argument("--orchestrator-model", default="qwen3-8b")
p.add_argument("--orchestrator-api-key", default="EMPTY")
# 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("--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)
args = p.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(message)s")
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)
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,
)
dispatch = make_dispatch({})
def rollout_fn(task):
try:
return run_unified_rollout(
task.instruction, tools,
call_orchestrator=call_orch, dispatch=dispatch,
max_turns=args.max_turns,
)
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
tasks = load_sft_tasks()
if args.max_tasks is not None:
tasks = tasks[: args.max_tasks]
logging.info("Loaded %d SFT tasks", len(tasks))
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
)
print(json.dumps(stats, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
scripts/orchestrator/build_sft_data.py
-19
View File
@@ -1,19 +0,0 @@
#!/usr/bin/env python
"""Thin CLI wrapper: build the orchestrator SFT dataset from NeuLab ADP.
Equivalent to::
python -m openjarvis.learning.intelligence.orchestrator.sft_data.build [...]
Usage:
python scripts/orchestrator/build_sft_data.py \
--out data/orchestrator_sft_traces.jsonl \
--max-tasks 2000 --adp-configs codeactinstruct,code_feedback,openhands
"""
from __future__ import annotations
from openjarvis.learning.intelligence.orchestrator.sft_data.build import _main
if __name__ == "__main__":
raise SystemExit(_main())
scripts/orchestrator/build_unified_sft.py
+11 -5
View File
@@ -40,9 +40,9 @@ from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
def main(argv: Optional[list[str]] = None) -> int:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--out", default="data/orchestrator_unified_sft.jsonl")
p.add_argument("--teacher-model", default="gpt-5")
p.add_argument("--teacher-base-url", default=None,
help="OpenAI-compatible base URL; omit for OpenAI cloud.")
p.add_argument("--teacher-model", default="claude-sonnet-4-6")
p.add_argument("--teacher-base-url", default="https://api.anthropic.com/v1/",
help="OpenAI-compatible base URL; pass '' for OpenAI cloud.")
p.add_argument("--max-tasks", type=int, default=200)
p.add_argument("--samples-per-task", type=int, default=4)
p.add_argument("--max-keep-per-task", type=int, default=1)
@@ -59,10 +59,16 @@ def main(argv: Optional[list[str]] = None) -> int:
)
logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools])
base_url = args.teacher_base_url or None
api_key = (
os.environ.get("ANTHROPIC_API_KEY")
if base_url and "anthropic" in base_url
else os.environ.get("OPENAI_API_KEY")
)
call_orch = make_call_orchestrator(
args.teacher_model,
base_url=args.teacher_base_url,
api_key=os.environ.get("OPENAI_API_KEY"),
base_url=base_url,
api_key=api_key,
temperature=args.temperature,
)
dispatch = make_dispatch({})
scripts/orchestrator/eval_orchestrator.py
+324
View File
@@ -0,0 +1,324 @@
#!/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()
combined = {
"orchestrator_model": args.orchestrator_model,
"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())
src/openjarvis/agents/hybrid/_prices.py
+1
View File
@@ -11,6 +11,7 @@ from __future__ import annotations
# USD per million tokens, (input, output). Local models = 0.
PRICES: dict[str, tuple[float, float]] = {
"claude-opus-4-8": (5.00, 25.0),
"claude-opus-4-7": (5.00, 25.0),
"claude-sonnet-4-6": (3.00, 15.0),
"claude-haiku-4-5": (1.00, 5.00),
src/openjarvis/agents/hybrid/expert_registry.py
+321 -37
View File
@@ -24,6 +24,7 @@ circular import.
from __future__ import annotations
import random
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional
@@ -34,13 +35,30 @@ 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)
VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_SEARCH, KIND_CODE, KIND_TOOL)
# Backend dispatch types understood by ``toolorchestra._call_worker``.
# 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_GENERALIST = "generalist_model"
CATEGORY_SPECIALIZED = "specialized_model"
CATEGORY_BASIC = "basic_tool"
CATEGORY_SMALL_GENERALIST = "small_generalist"
CATEGORY_STRONG_GENERALIST = "strong_generalist"
# Two-model-class taxonomy for the orchestrator catalog (the specialized /
# small-vs-strong generalist tiers are superseded).
CATEGORY_CLOUD_FRONTIER = "cloud_frontier"
CATEGORY_LOCAL_OSS = "local_open_source"
# Backend dispatch types understood by ``toolorchestra._call_worker`` (plus the
# ``openjarvis-tool`` bridge, dispatched in ``unified.make_dispatch`` via the
# OpenJarvis ToolExecutor rather than ``_call_worker``).
VALID_BACKENDS = (
"vllm", "openai", "anthropic", "gemini", "openrouter",
"anthropic-web-search", "tavily-search", "modal-python",
"openjarvis-tool",
)
@@ -63,6 +81,11 @@ class ExpertTool:
price_in: float = 0.0
price_out: float = 0.0
latency_s: float = 5.0
category: str = "" # generalist_model | specialized_model | 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
def __post_init__(self) -> None:
if self.kind not in VALID_KINDS:
@@ -75,7 +98,13 @@ class ExpertTool:
# ---- orchestrator-visible spec -------------------------------------
def _param_schema(self) -> Dict[str, object]:
"""JSON-schema for the tool's arguments (one typed param per kind)."""
"""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",
@@ -122,21 +151,36 @@ class ExpertTool:
return self.summary.rstrip(".") + "." + cost_line
def to_spec(self) -> Dict[str, object]:
"""OpenAI-style tool spec the orchestrator conditions on."""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description(),
"parameters": self._param_schema(),
},
"""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"
# 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(
@@ -151,70 +195,300 @@ def default_catalog(
"""
cat: List[ExpertTool] = []
# ---- generalist / frontier models ----
for name, model, summary, lat in [
("gpt_5", "gpt-5",
# ---- 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", "gpt-5-mini",
("gpt-5-mini",
"Mid-tier generalist (GPT-5-mini). Solid reasoning, much cheaper.", 15.0),
("gpt_4o", "gpt-4o",
("gpt-4o",
"Fast generalist (GPT-4o). Good for simple steps and formatting.", 8.0),
("claude_opus", "claude-opus-4-7",
"Frontier generalist (Claude Opus). Strong long-horizon reasoning.", 26.0),
("claude_sonnet", "claude-sonnet-4-6",
"Strong generalist (Claude Sonnet). Balanced cost/capability.", 15.0),
("gemini_2_5_pro", "gemini-2.5-pro",
("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", "gemini-2.5-flash",
("gemini-2.5-flash",
"Cheap fast generalist (Gemini 2.5 Flash).", 8.0),
("llama_3_3_70b", "meta-llama/llama-3.3-70b-instruct",
("meta-llama/llama-3.3-70b-instruct",
"Open generalist (Llama-3.3-70B). Decent general knowledge, low cost.", 10.0),
("qwen3_32b", "qwen/qwen3-32b",
("qwen/qwen3-32b",
"Open generalist (Qwen3-32B). Strong math/science reasoning, low cost.", 9.0),
]:
ep = "openai" if name.startswith("gpt") else (
"anthropic" if name.startswith("claude") else (
"gemini" if name.startswith("gemini") else "openrouter"))
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=name, kind=KIND_MODEL, backend_type=ep, summary=summary,
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 ----
pi, po = _price("qwen/qwen-2.5-coder-32b-instruct")
coder = "qwen/qwen-2.5-coder-32b-instruct"
pi, po = _price(coder)
cat.append(ExpertTool(
name="qwen2_5_coder_32b", kind=KIND_MODEL, backend_type="openrouter",
name=_tool_name(coder), kind=KIND_MODEL, backend_type="openrouter",
summary="Specialized code model (Qwen2.5-Coder-32B). Writes/debugs code.",
model="qwen/qwen-2.5-coder-32b-instruct",
price_in=pi, price_out=po, latency_s=9.0,
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="local_model", kind=KIND_MODEL, backend_type="vllm",
summary=("On-device open model served locally. Cheap and private; "
"good for extraction, formatting, arithmetic on given data."),
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,
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,
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 locally-served 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.
#
# Local model ids -> vLLM base_url come from ``local_endpoints``; unmapped ids
# get ``base_url=None`` (tool still listed, dispatch wires the endpoint later).
_LOCAL_OSS_MODELS = (
"Qwen/Qwen3.5-9B",
"Qwen/Qwen3.6-27B-FP8",
"Qwen/Qwen3.5-122B-A10B-FP8",
"Qwen/Qwen3.5-397B-A17B-FP8",
)
_CLOUD_FRONTIER_MODELS = (
# (model, backend, summary, latency_s)
("gpt-5.5", "openai",
"Cloud frontier generalist (GPT-5.5). Strongest reasoning across domains.", 30.0),
("claude-opus-4-8", "anthropic",
"Cloud frontier generalist (Claude Opus 4.8). Strong long-horizon reasoning.", 26.0),
)
def orchestrator_catalog(
*,
local_endpoints: Optional[Dict[str, str]] = None,
include_tools: bool = True,
) -> List[ExpertTool]:
"""Return the orchestrator's tool catalog: two model classes + basic tools.
``local_endpoints`` maps a local model id (e.g. ``"Qwen/Qwen3.5-9B"``) to its
vLLM ``base_url``; ids not present get ``base_url=None``. ``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 {}
cat: List[ExpertTool] = []
# ---- cloud frontier models ----
for model, backend, summary, lat in _CLOUD_FRONTIER_MODELS:
pi, po = _price(model)
cat.append(ExpertTool(
name=_tool_name(model), kind=KIND_MODEL, backend_type=backend,
summary=summary, model=model, price_in=pi, price_out=po,
latency_s=lat, category=CATEGORY_CLOUD_FRONTIER,
))
# ---- local open-source models (served via vLLM, price 0/0) ----
for model in _LOCAL_OSS_MODELS:
cat.append(ExpertTool(
name=_tool_name(model), kind=KIND_MODEL, backend_type="vllm",
summary=(f"Locally-served open-source model ({model}). Cheap and "
"private; no per-token cost."),
model=model, base_url=local_endpoints.get(model),
price_in=0.0, price_out=0.0, latency_s=4.0,
category=CATEGORY_LOCAL_OSS,
))
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]
@@ -298,13 +572,23 @@ def to_worker_dict(tool: ExpertTool) -> Dict[str, object]:
__all__ = [
"CATEGORY_BASIC",
"CATEGORY_CLOUD_FRONTIER",
"CATEGORY_GENERALIST",
"CATEGORY_LOCAL_OSS",
"CATEGORY_SMALL_GENERALIST",
"CATEGORY_SPECIALIZED",
"CATEGORY_STRONG_GENERALIST",
"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",
src/openjarvis/agents/hybrid/toolorchestra/rollout.py
+22
View File
@@ -28,6 +28,27 @@ from openjarvis.agents.hybrid.expert_registry import (
RL_ORCHESTRATOR_SYS = "You are good at using tools."
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 context (mirrors the paper's ~24k-token cap).
_CONTEXT_CAP = 24000
@@ -137,6 +158,7 @@ __all__ = [
"RL_ORCHESTRATOR_SYS",
"UnifiedRollout",
"UnifiedTurn",
"build_system_prompt",
"run_unified_rollout",
"tool_call_tag",
]
src/openjarvis/agents/hybrid/toolorchestra/unified.py
+62
View File
@@ -10,6 +10,7 @@ network or keys.
from __future__ import annotations
import json
from typing import Any, Callable, Dict, List, Optional, Tuple
from openjarvis.agents.hybrid._prices import cost as _model_cost
@@ -63,13 +64,74 @@ def make_call_orchestrator(
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:
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(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"):
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-grpo.toml
+54
View File
@@ -0,0 +1,54 @@
# Orchestrator GRPO — Gemma-4-12B (cost-aware, lambda-swept).
#
# Trajectory-level GRPO over the UNIFIED tool-use rollout (the same faithful
# <tool_call> multi-turn path the SFT data is built from). For each prompt we
# sample group_size rollouts with the *policy* as orchestrator
# (run_unified_rollout), verify each rollout's final_answer against the gold
# answer, score with the cost-aware reward (r = +1/-1, R -= lam * cost /
# cost_max), normalise advantages within the group, and PPO-clip the
# trajectory log-prob update (summed over assistant-masked tokens only).
#
# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 rollouts per
# step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches
# per epoch.
[model]
model_name = "google/gemma-4-12B-it"
max_prompt_length = 24000
max_response_length = 8768
[training]
num_epochs = 5
learning_rate = 1e-6
max_grad_norm = 1.0
gradient_checkpointing = true
[grpo]
# group_size = number of trajectories sampled per prompt.
num_samples_per_prompt = 8
temperature = 1.0
kl_coef = 0.0001
clip_ratio = 0.2
# Per-rollout tool-use turn cap (one assistant turn per loop iteration).
max_turns = 10
[data]
grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer)
prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step)
[reward]
# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max).
# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier,
# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run.
lam = 0.0
cost_max = 0.10
# Optional: map local OSS model ids -> served vLLM base_url so the local-model
# tools in the orchestrator catalog dispatch to a real endpoint. Omit to leave
# them unbacked (tool still listed).
[local_endpoints]
# "Qwen/Qwen3.5-9B" = "http://localhost:8001/v1"
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_grpo_gemma_4_12b"
save_every_n_epochs = 1
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-sft.toml
+34
View File
@@ -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
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-grpo.toml
+44
View File
@@ -0,0 +1,44 @@
# Orchestrator GRPO — Qwen3-8B (cost-aware, lambda-swept).
#
# Fields map onto OrchestratorGRPOConfig. GRPO uses group-relative
# advantages (no critic): for each prompt, sample group_size trajectories,
# verify against the gold answer, score with the cost-aware reward
# (r = +1/-1, R -= lam * cost / cost_max), normalise advantages within the
# group, and PPO-clip the policy update.
#
# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 trajectories
# per step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches
# per epoch.
[model]
model_name = "Qwen/Qwen3-8B"
max_prompt_length = 24000
max_response_length = 8768
[training]
num_epochs = 5
learning_rate = 1e-6
max_grad_norm = 1.0
gradient_checkpointing = true
[grpo]
# group_size = number of trajectories sampled per prompt.
num_samples_per_prompt = 8
temperature = 1.0
kl_coef = 0.0001
clip_ratio = 0.2
[data]
grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer)
prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step)
[reward]
# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max).
# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier,
# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run.
lam = 0.0
cost_max = 0.10
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_grpo"
save_every_n_epochs = 1
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml
+13 -8
View File
@@ -1,8 +1,10 @@
# Orchestrator SFT cold-start — Qwen3-8B (~9B, "easier to train").
#
# Fields map onto OrchestratorSFTConfig. Build the dataset first (no GPU):
# python -m openjarvis.learning.intelligence.orchestrator.sft_data.build \
# --out data/orchestrator_sft_traces.jsonl --max-tasks 2000
# 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]
@@ -11,18 +13,21 @@ max_seq_length = 4096
[training]
num_epochs = 3
batch_size = 8
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_traces.jsonl"
trace_cache_path = "data/orchestrator_sft_v1.jsonl"
regenerate_traces = false
# Coding / agentic / tool-use ADP sub-configs (per the meeting steer).
adp_configs = "codeactinstruct,code_feedback,openhands,agenttuning_os,swe-smith"
distill_max_tasks = 2000
# 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"
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-grpo.toml
+54
View File
@@ -0,0 +1,54 @@
# Orchestrator GRPO — Qwen3.5-9B (cost-aware, lambda-swept).
#
# Trajectory-level GRPO over the UNIFIED tool-use rollout (the same faithful
# <tool_call> multi-turn path the SFT data is built from). For each prompt we
# sample group_size rollouts with the *policy* as orchestrator
# (run_unified_rollout), verify each rollout's final_answer against the gold
# answer, score with the cost-aware reward (r = +1/-1, R -= lam * cost /
# cost_max), normalise advantages within the group, and PPO-clip the
# trajectory log-prob update (summed over assistant-masked tokens only).
#
# Batch math: prompts_per_step * group_size = 256 * 8 = 2048 rollouts per
# step. grpo_max_prompts / prompts_per_step = 20000 / 256 ~= 78 batches
# per epoch.
[model]
model_name = "Qwen/Qwen3.5-9B"
max_prompt_length = 24000
max_response_length = 8768
[training]
num_epochs = 5
learning_rate = 1e-6
max_grad_norm = 1.0
gradient_checkpointing = true
[grpo]
# group_size = number of trajectories sampled per prompt.
num_samples_per_prompt = 8
temperature = 1.0
kl_coef = 0.0001
clip_ratio = 0.2
# Per-rollout tool-use turn cap (one assistant turn per loop iteration).
max_turns = 10
[data]
grpo_max_prompts = 20000 # pooled GRPO prompts (cached on the trainer)
prompts_per_step = 256 # prompts per _grpo_step (256 * 8 = 2048 traj/step)
[reward]
# Cost-aware reward: R = base(+1/-1) - lam * (cost_usd / cost_max).
# `lam` is SWEPT across runs to trace the accuracy/cost Pareto frontier,
# e.g. [0.0, 0.05, 0.1, 0.2, 0.4]. The value below is the default run.
lam = 0.0
cost_max = 0.10
# Optional: map local OSS model ids -> served vLLM base_url so the local-model
# tools in the orchestrator catalog dispatch to a real endpoint. Omit to leave
# them unbacked (tool still listed).
[local_endpoints]
# "Qwen/Qwen3.5-9B" = "http://localhost:8001/v1"
[checkpoint]
checkpoint_dir = "checkpoints/orchestrator_grpo_qwen3.5_9b"
save_every_n_epochs = 1
src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml
+34
View File
@@ -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
src/openjarvis/learning/intelligence/orchestrator/eval_backend.py
+188
View File
@@ -0,0 +1,188 @@
"""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 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
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,
)
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,
)
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)
return {
"content": 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()
],
},
}
__all__ = ["OrchestratorBackend", "DEFAULT_ENDPOINT", "DEFAULT_MODEL"]
src/openjarvis/learning/intelligence/orchestrator/grpo_trainer.py
+387 -148
View File
@@ -6,8 +6,23 @@ doesn't require a separate critic model — instead, it uses
trajectories, compute rewards, normalise within the group, and update
the policy to increase the probability of better solutions.
**Trajectory-level GRPO over the unified tool-use rollout.** Generation
flows through :func:`run_unified_rollout` (the SAME faithful
``<tool_call>`` multi-turn loop our SFT data is built from), NOT the old
single-turn ``THOUGHT/TOOL/INPUT`` environment path. The policy *is* the
orchestrator: each turn we tokenize ``system+user`` with the model's chat
template, generate ONE assistant turn, parse any
``<tool_call>{...}</tool_call>``, and record the exact
``(input_ids, generated_ids)`` for that turn. Once a rollout finishes, the
recorded turns are stitched into one trajectory token sequence with an
*assistant mask* (1 only on policy-generated tokens) so the policy gradient
is computed solely over the tokens the model actually authored, under both
the current policy (with grad) and the frozen reference (no grad).
All ``torch``/``transformers`` imports are guarded so the module can be
imported without GPU dependencies.
imported without GPU dependencies; the rollout + log-prob paths are
monkeypatchable so ``_grpo_step``/``_train_epoch`` are CPU-smoke-testable
with fakes (``policy.model is None``).
"""
from __future__ import annotations
@@ -17,7 +32,7 @@ import logging
import shutil
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Dict, List
from typing import Any, Dict, List, Optional, Tuple
# Optional imports -----------------------------------------------------------
try:
@@ -68,11 +83,24 @@ class OrchestratorGRPOConfig:
max_grad_norm: float = 1.0
# GRPO specific
# ``num_samples_per_prompt`` is the GRPO *group size*: N trajectories
# sampled per prompt for group-relative advantages.
num_samples_per_prompt: int = 8
temperature: float = 1.0
kl_coef: float = 0.0001
clip_ratio: float = 0.2
# Data / batching
# How many prompts to pool from the GRPO dataset (cached on the trainer).
grpo_max_prompts: int = 20000
# How many prompts go into one ``_grpo_step`` (a batch / step). With
# group_size N this is prompts_per_step * N trajectories per step.
prompts_per_step: int = 256
# Cost-aware reward (r=+1/-1, R -= lam * cost / cost_max). ``lam`` is swept.
lam: float = 0.0
cost_max: float = 0.10
# Environment
available_tools: List[str] = field(
default_factory=lambda: [
@@ -83,6 +111,11 @@ class OrchestratorGRPOConfig:
]
)
max_turns: int = 10
# Optional local OSS model endpoints, mapping full model id (e.g.
# "Qwen/Qwen3.5-9B") -> vLLM base_url. Passed to ``orchestrator_catalog``
# so the local-model tools dispatch to a served endpoint; omitted when
# unset (tool still listed, just unbacked).
local_endpoints: Dict[str, str] = field(default_factory=dict)
# Checkpoint
checkpoint_dir: str = "checkpoints/orchestrator_grpo"
@@ -95,6 +128,24 @@ class OrchestratorGRPOConfig:
use_8bit_optimizer: bool = False
@dataclass
class RolloutRecord:
"""One sampled unified-rollout trajectory, with the tokens + assistant mask
needed to compute its policy log-prob.
``token_ids`` is the full stitched trajectory (prompt / observation tokens
interleaved with policy-generated tokens over turns). ``assistant_mask``
is parallel: 1 only on policy-generated positions. ``advantage`` is filled
in after group-relative normalisation.
"""
final_answer: str
cost_usd: float
token_ids: List[int] = field(default_factory=list)
assistant_mask: List[int] = field(default_factory=list)
advantage: float = 0.0
# ---------------------------------------------------------------------------
# Trainer
# ---------------------------------------------------------------------------
@@ -111,12 +162,37 @@ class OrchestratorGRPOTrainer:
self.device = None
self.global_step = 0
# Cost-aware reward (r=+1/-1, R -= lam * cost / cost_max). ``lam``
# is swept across runs to trace the accuracy/cost Pareto frontier.
from openjarvis.learning.intelligence.orchestrator.reward import (
CostAwareReward,
)
self._reward = CostAwareReward(lam=config.lam, cost_max=config.cost_max)
# Lazily-loaded, cached GRPO prompt pool (list[Task]).
self._prompts: List[Any] | None = None
# Lazily-built, cached unified tool catalog (list[ExpertTool]).
self._tool_catalog: Optional[List[Any]] = None
if HAS_TORCH and torch is not None:
self.device = _select_torch_device()
self._init_model()
self._init_optimizer()
# -- Data ----------------------------------------------------------------
def _load_prompts(self) -> List[Any]:
"""Load (and cache) the GRPO prompt pool — a ``list[Task]``."""
if self._prompts is None:
from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( # noqa: E501
load_grpo_prompts,
)
self._prompts = load_grpo_prompts(n=self.config.grpo_max_prompts)
return self._prompts
# -- Initialisation ------------------------------------------------------
def _init_model(self) -> None:
@@ -188,82 +264,270 @@ class OrchestratorGRPOTrainer:
self._save_checkpoint(epoch)
def _train_epoch(self, epoch: int) -> Dict[str, float]:
if self.policy.model is None:
return {"epoch": epoch, "loss": 0.0, "reward": 0.0}
"""Iterate the GRPO prompt pool in chunks and run ``_grpo_step``.
Prompts are loaded once via :meth:`_load_prompts` (cached on the
trainer) and chunked into batches of ``config.prompts_per_step``.
Each chunk is one ``_grpo_step`` call (group_size trajectories per
prompt). Real model work happens inside ``_grpo_step`` and is
guarded there; ``_train_epoch`` itself does no torch work, so it
can be smoke-tested on CPU by monkeypatching ``_grpo_step``.
"""
if HAS_TORCH and self.policy.model is not None:
self.policy.model.train()
prompts = self._load_prompts()
step_size = max(int(self.config.prompts_per_step), 1)
self.policy.model.train()
total_loss = 0.0
total_reward = 0.0
total_accuracy = 0.0
n_prompts = 0
num_batches = 0
# In a real implementation, iterate over a task dataset.
# Here we provide the skeleton; actual data loading is
# trainer-specific.
self.global_step += 1
num_batches = max(num_batches, 1)
for start in range(0, len(prompts), step_size):
chunk = prompts[start : start + step_size]
if not chunk:
continue
metrics = self._grpo_step(chunk)
loss = float(metrics.get("loss", 0.0))
reward = float(metrics.get("reward", 0.0))
accuracy = float(metrics.get("accuracy", 0.0))
n = int(metrics.get("n_prompts", len(chunk)))
avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
avg_reward = total_reward / num_batches if num_batches > 0 else 0.0
return {"epoch": epoch, "loss": avg_loss, "reward": avg_reward}
total_loss += loss * n
total_reward += reward * n
total_accuracy += accuracy * n
n_prompts += n
num_batches += 1
self.global_step += 1
def _grpo_step(
self,
prompts: List[str],
ground_truths: List[str],
) -> tuple:
"""Perform one GRPO training step.
denom = n_prompts if n_prompts > 0 else 1
return {
"epoch": epoch,
"loss": total_loss / denom,
"reward": total_reward / denom,
"accuracy": total_accuracy / denom,
"n_prompts": n_prompts,
"n_batches": num_batches,
}
For each prompt:
1. Generate N candidate trajectories.
2. Compute reward for each.
3. Normalise advantages within the group.
4. Compute clipped policy gradient + KL penalty.
# -- Tool catalog --------------------------------------------------------
def _tools(self) -> List[Any]:
"""The orchestrator's unified tool catalog (cached).
Two model classes (cloud frontier + local OSS) plus the basic tools,
exactly what the rollout / SFT data condition on. ``local_endpoints``
wires the local-model tools to their served vLLM endpoints.
"""
if getattr(self, "_tool_catalog", None) is None:
from openjarvis.agents.hybrid.expert_registry import (
orchestrator_catalog,
)
self._tool_catalog = orchestrator_catalog(
local_endpoints=dict(self.config.local_endpoints) or None,
)
return self._tool_catalog
# -- Rollout / trajectory collection -------------------------------------
def _rollout_group(self, task: Any) -> List["RolloutRecord"]:
"""Sample ``group_size`` rollouts for one task through the unified loop.
Each rollout is driven by a *policy* ``call_orchestrator`` (the HF
model being trained generating one assistant turn at a time) and the
real :func:`make_dispatch` tool executor. Per turn we capture the exact
``(input_ids, generated_ids)`` so the trajectory token sequence +
assistant mask can be rebuilt for the log-prob pass.
Returns one :class:`RolloutRecord` per sampled rollout. Overridable /
monkeypatchable in tests so ``_grpo_step`` runs with fakes when
``policy.model is None``.
"""
from openjarvis.agents.hybrid.toolorchestra.rollout import (
build_system_prompt,
run_unified_rollout,
)
from openjarvis.agents.hybrid.toolorchestra.unified import make_dispatch
from openjarvis.agents.hybrid.expert_registry import build_tool_specs
tools = self._tools()
system = build_system_prompt(build_tool_specs(tools))
dispatch = make_dispatch({})
records: List[RolloutRecord] = []
for _ in range(self.config.num_samples_per_prompt):
turn_buffer: List[Tuple[List[int], List[int]]] = []
call_orchestrator = self._make_policy_caller(turn_buffer)
rollout = run_unified_rollout(
task.instruction,
tools,
call_orchestrator=call_orchestrator,
dispatch=dispatch,
max_turns=self.config.max_turns,
system=system,
)
token_ids, assistant_mask = self._stitch_trajectory(turn_buffer)
records.append(
RolloutRecord(
final_answer=rollout.final_answer,
cost_usd=float(rollout.cost_usd),
token_ids=token_ids,
assistant_mask=assistant_mask,
)
)
return records
def _make_policy_caller(self, turn_buffer):
"""Build a policy-driven ``call_orchestrator(system, user, specs)``.
Tokenizes ``system+user`` with the model's chat template, generates ONE
assistant turn with the current policy, parses any
``<tool_call>{...}</tool_call>``, appends the per-turn
``(input_ids, generated_ids)`` to ``turn_buffer`` (for the later
log-prob pass), and returns
``(text, tool_calls, prompt_tokens, completion_tokens)`` as the unified
rollout expects.
"""
from openjarvis.agents.hybrid.toolorchestra.parsing import (
_parse_rl_tool_call,
)
tok = self.policy.tokenizer
def call_orchestrator(system: str, user: str, specs):
messages = [
{"role": "system", "content": system},
{"role": "user", "content": user},
]
input_ids = tok.apply_chat_template(
messages, add_generation_prompt=True, tokenize=True
)
inputs = {
"input_ids": torch.tensor([input_ids], device=self.device),
}
input_len = len(input_ids)
max_new = min(self.config.max_response_length, 32000 - input_len - 100)
max_new = max(min(max_new, 2048), 16)
with torch.no_grad():
out = self.policy.model.generate(
**inputs,
max_new_tokens=max_new,
temperature=self.config.temperature,
do_sample=True,
)
gen_ids = out[0][input_len:].tolist()
text = tok.decode(gen_ids, skip_special_tokens=True)
# Record exact tokens for the trajectory log-prob pass.
turn_buffer.append((list(input_ids), list(gen_ids)))
parsed = _parse_rl_tool_call(text, None)
tool_calls = (
[(parsed["name"], parsed["arguments"])] if parsed else []
)
return text, tool_calls, int(input_len), int(len(gen_ids))
return call_orchestrator
def _stitch_trajectory(
self, turn_buffer: List[Tuple[List[int], List[int]]]
) -> Tuple[List[int], List[int]]:
"""Concatenate per-turn tokens into one trajectory + assistant mask.
For turn ``i`` the buffer holds ``(prompt_ids_i, gen_ids_i)`` where
``prompt_ids_i`` already contains all prior turns' assistant text and
the tool observations re-rendered into the chat by the rollout loop
(the rollout rebuilds ``user`` each turn from the running context). So
the faithful trajectory token stream is, for each turn, the *new*
prefix tokens (mask 0 — prompt / observation / chat-template tokens)
followed by the generated assistant tokens (mask 1).
We reconstruct it by diffing successive prompts: turn ``i`` contributes
the prompt tokens that weren't already emitted, then its generated
tokens. This keeps the assistant mask exact (1 only on policy-generated
positions) without re-tokenizing the whole transcript.
"""
token_ids: List[int] = []
mask: List[int] = []
emitted = 0 # length of token_ids already produced
for prompt_ids, gen_ids in turn_buffer:
# New prompt-side tokens for this turn (everything past what we've
# already emitted). These are observation / chat-template tokens →
# mask 0 (not policy-authored).
new_prompt = prompt_ids[emitted:] if len(prompt_ids) > emitted else []
token_ids.extend(new_prompt)
mask.extend([0] * len(new_prompt))
# Generated assistant tokens → mask 1.
token_ids.extend(gen_ids)
mask.extend([1] * len(gen_ids))
emitted = len(token_ids)
return token_ids, mask
def _grpo_step(self, tasks: List[Any]) -> Dict[str, float]:
"""Perform one GRPO training step over a batch of ``Task`` objects.
For each task:
1. Sample N candidate trajectories (``group_size`` =
``num_samples_per_prompt``) through the unified tool-use rollout —
the SAME ``<tool_call>`` multi-turn path the SFT data uses, driven
by the *policy* as orchestrator (see :meth:`_rollout_group`).
2. Verify each rollout's ``final_answer`` against the gold answer
(:func:`verify_answer`) and score with :class:`CostAwareReward`
over an :class:`Episode` carrying ``.correct`` and
``.total_cost_usd``.
3. Normalise advantages within the group (mean/std).
4. Compute the trajectory-level clipped policy gradient + KL penalty.
The trajectory log-prob is the sum of per-token log-probs over the
*assistant-masked* positions (:meth:`_trajectory_logprob`).
5. Backward + clip + step.
Returns ``(loss_value, avg_reward)``.
Returns a metrics dict with ``loss``, ``reward``, ``accuracy``,
and ``n_prompts``.
"""
if self.policy.model is None or not HAS_TORCH:
raise RuntimeError("Cannot train without PyTorch and model.")
self.policy.model.train()
all_prompts: list[str] = []
all_responses: list[str] = []
all_advantages: list[float] = []
all_rewards: list[float] = []
from openjarvis.learning.intelligence.orchestrator.reward import (
MultiObjectiveReward,
Normalizers,
RewardWeights,
from openjarvis.learning.intelligence.orchestrator.sft_data.verify import (
verify_answer,
)
from openjarvis.learning.intelligence.orchestrator.types import (
Episode,
)
reward_fn = MultiObjectiveReward(RewardWeights(), Normalizers())
reward_fn = self._reward
for prompt, gt in zip(prompts, ground_truths):
# (token_ids, assistant_mask, advantage) per trajectory, plus running
# reward / accuracy stats.
traj_records: list[RolloutRecord] = []
all_advantages: list[float] = []
all_rewards: list[float] = []
all_correct: list[float] = []
for task in tasks:
records = self._rollout_group(task)
group_rewards: list[float] = []
group_responses: list[str] = []
for _ in range(self.config.num_samples_per_prompt):
response, _log_probs = self._generate_with_log_probs(prompt)
# Build a minimal episode for reward computation
for rec in records:
correct = verify_answer(task, rec.final_answer)
episode = Episode(
task_id="grpo",
initial_prompt=prompt,
ground_truth=gt,
final_answer=response,
correct=(response.strip() == gt.strip()),
task_id=getattr(task, "task_id", ""),
initial_prompt=task.instruction,
ground_truth=getattr(task, "answer", ""),
final_answer=rec.final_answer,
correct=bool(correct),
total_cost_usd=rec.cost_usd,
)
reward = reward_fn.compute(episode)
group_rewards.append(reward)
group_responses.append(response)
all_correct.append(1.0 if correct else 0.0)
# Group-relative advantages
# Group-relative advantages (mean/std within the group).
mean_r = sum(group_rewards) / len(group_rewards)
std_r = (
sum((r - mean_r) ** 2 for r in group_rewards) / len(group_rewards)
@@ -273,21 +537,25 @@ class OrchestratorGRPOTrainer:
else:
advantages = [0.0] * len(group_rewards)
for resp, adv, rew in zip(group_responses, advantages, group_rewards):
all_prompts.append(prompt)
all_responses.append(resp)
for rec, adv, rew in zip(records, advantages, group_rewards):
rec.advantage = adv
traj_records.append(rec)
all_advantages.append(adv)
all_rewards.append(rew)
# Policy gradient loss
# Trajectory-level policy gradient loss.
total_loss = torch.tensor(0.0, device=self.device, requires_grad=True)
for prompt, response, advantage in zip(
all_prompts, all_responses, all_advantages
):
current_lp = self._compute_log_probs(prompt, response)
for rec in traj_records:
if not rec.token_ids or sum(rec.assistant_mask) == 0:
continue
current_lp = self._trajectory_logprob(
rec.token_ids, rec.assistant_mask, ref=False
)
with torch.no_grad():
ref_lp = self._compute_log_probs_ref(prompt, response)
ref_lp = self._trajectory_logprob(
rec.token_ids, rec.assistant_mask, ref=True
)
log_ratio = current_lp - ref_lp
ratio = torch.exp(log_ratio)
@@ -296,16 +564,25 @@ class OrchestratorGRPOTrainer:
clip = self.config.clip_ratio
clipped = torch.clamp(ratio, 1 - clip, 1 + clip)
policy_loss = -torch.min(ratio * advantage, clipped * advantage)
policy_loss = -torch.min(ratio * rec.advantage, clipped * rec.advantage)
kl = (ratio - 1) - log_ratio
total_loss = total_loss + policy_loss + self.config.kl_coef * kl
avg_loss = total_loss / max(len(all_prompts), 1)
n_prompts = len(tasks)
avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
avg_acc = sum(all_correct) / len(all_correct) if all_correct else 0.0
n_traj = max(len(traj_records), 1)
avg_loss = total_loss / n_traj
loss_val = avg_loss.item()
if torch.isnan(avg_loss) or torch.isinf(avg_loss):
avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
return 0.0, float(avg_reward)
return {
"loss": 0.0,
"reward": float(avg_reward),
"accuracy": float(avg_acc),
"n_prompts": n_prompts,
}
self.optimizer.zero_grad()
avg_loss.backward()
@@ -314,107 +591,68 @@ class OrchestratorGRPOTrainer:
for param in self.policy.model.parameters():
if param.grad is not None and torch.isnan(param.grad).any():
self.optimizer.zero_grad()
avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
return float(loss_val), float(avg_reward)
return {
"loss": float(loss_val),
"reward": float(avg_reward),
"accuracy": float(avg_acc),
"n_prompts": n_prompts,
}
torch.nn.utils.clip_grad_norm_(
self.policy.model.parameters(), self.config.max_grad_norm
)
self.optimizer.step()
avg_reward = sum(all_rewards) / len(all_rewards) if all_rewards else 0.0
return float(loss_val), float(avg_reward)
return {
"loss": float(loss_val),
"reward": float(avg_reward),
"accuracy": float(avg_acc),
"n_prompts": n_prompts,
}
# -- Generation / log-prob helpers ---------------------------------------
# -- Trajectory log-prob -------------------------------------------------
def _generate_with_log_probs(self, prompt: str) -> "tuple[str, Any]":
"""Generate a response and return ``(text, log_probs)``."""
inputs = self.policy.tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=min(self.config.max_prompt_length, 16000),
).to(self.device)
def _trajectory_logprob(
self,
token_ids: List[int],
assistant_mask: List[int],
*,
ref: bool,
) -> "torch.Tensor":
"""Sum of per-token log-probs over the assistant-masked positions.
input_len = inputs.input_ids.shape[1]
max_new = min(self.config.max_response_length, 32000 - input_len - 100)
max_new = max(min(max_new, 2048), 128)
``token_ids`` is the full trajectory (prompt + observation + assistant
tokens interleaved over turns); ``assistant_mask[i] == 1`` iff token
``i`` was generated by the policy. We run the model once over the whole
sequence and accumulate ``log p(token_{i} | token_{<i})`` only where the
*predicted* token (position ``i``) is masked 1.
with torch.no_grad():
outputs = self.policy.model.generate(
**inputs,
max_new_tokens=max_new,
temperature=self.config.temperature,
do_sample=True,
output_scores=True,
return_dict_in_generate=True,
)
``ref=True`` → frozen reference model, no grad; else current policy with
grad. Mirrors the single-turn log-prob math the old code used, lifted to
the masked multi-turn trajectory.
"""
model = self.ref_policy.model if ref else self.policy.model
ids = torch.tensor([token_ids], device=self.device)
generated_ids = outputs.sequences[0][input_len:]
if len(generated_ids) == 0:
return "", torch.tensor(0.0, device=self.device)
text = self.policy.tokenizer.decode(generated_ids, skip_special_tokens=True)
log_probs = []
for token_id, logits in zip(generated_ids, outputs.scores):
probs = F.softmax(logits[0], dim=-1)
tid = token_id.item()
if 0 <= tid < probs.shape[0]:
log_probs.append(torch.log(probs[tid] + 1e-10))
total_lp = (
torch.stack(log_probs).sum()
if log_probs
else torch.tensor(0.0, device=self.device)
)
return text, total_lp
def _compute_log_probs(self, prompt: str, response: str) -> "torch.Tensor":
"""Log-probs of *response* given *prompt* under current policy."""
full = prompt + response
inputs = self.policy.tokenizer(full, return_tensors="pt", truncation=True).to(
self.device
)
prompt_inputs = self.policy.tokenizer(
prompt, return_tensors="pt", truncation=True
).to(self.device)
with torch.enable_grad():
logits = self.policy.model(**inputs).logits
start = len(prompt_inputs.input_ids[0])
end = len(inputs.input_ids[0])
if ref:
ctx = torch.no_grad()
else:
ctx = torch.enable_grad()
with ctx:
logits = model(input_ids=ids).logits
lps = []
for i in range(start, end - 1):
lp = F.log_softmax(logits[0, i, :], dim=-1)[inputs.input_ids[0, i + 1]]
# Predict token i from logits at position i-1; include it iff token i is
# an assistant-generated token.
for i in range(1, len(token_ids)):
if assistant_mask[i] != 1:
continue
lp = F.log_softmax(logits[0, i - 1, :], dim=-1)[token_ids[i]]
lps.append(lp)
return torch.stack(lps).sum() if lps else torch.tensor(0.0)
def _compute_log_probs_ref(self, prompt: str, response: str) -> "torch.Tensor":
"""Log-probs under the frozen reference policy (no grad)."""
full = prompt + response
inputs = self.ref_policy.tokenizer(
full, return_tensors="pt", truncation=True
).to(self.device)
prompt_inputs = self.ref_policy.tokenizer(
prompt, return_tensors="pt", truncation=True
).to(self.device)
with torch.no_grad():
logits = self.ref_policy.model(**inputs).logits
start = len(prompt_inputs.input_ids[0])
end = len(inputs.input_ids[0])
lps = []
for i in range(start, end - 1):
lp = F.log_softmax(logits[0, i, :], dim=-1)[inputs.input_ids[0, i + 1]]
lps.append(lp)
return torch.stack(lps).sum() if lps else torch.tensor(0.0)
if not lps:
return torch.tensor(0.0, device=self.device)
return torch.stack(lps).sum()
# -- Checkpointing -------------------------------------------------------
@@ -488,4 +726,5 @@ _ensure_registered()
__all__ = [
"OrchestratorGRPOConfig",
"OrchestratorGRPOTrainer",
"RolloutRecord",
]
src/openjarvis/learning/intelligence/orchestrator/reward.py
+57
View File
@@ -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",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py
+27 -38
View File
@@ -1,55 +1,44 @@
"""Synthetic SFT-data generation for the orchestrator cold-start.
"""SFT-data generation for the orchestrator cold-start.
Turns NeuLab ADP (``neulab/agent-data-collection``) agent trajectories into
THOUGHT/TOOL/INPUT ``conversations`` JSONL that
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::
ADP trajectory -> canonical Episode -> per-paradigm tiered renderings
-> reward-ranked best-correct rendering -> conversations JSONL
ToolScale task -> N teacher rollouts -> verify -> keep cheapest passing
-> conversations JSONL
Everything here runs with **no GPU and no API keys** (the cold-start does not
re-execute models; it re-tiers the demonstrated ADP traces).
(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.adp_loader import (
CanonicalStep,
iter_trajectories,
trajectory_rows_to_episode,
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
generate_sft_dataset,
gold_coverage_verify,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.build import (
build_sft_dataset,
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
GoldAction,
ToolScaleTask,
load_toolscale,
normalize_row,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
PARADIGMS,
RenderedEpisode,
render_all,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
Difficulty,
Tier,
step_difficulty,
tier_telemetry,
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
__all__ = [
"CanonicalStep",
"Difficulty",
"PARADIGMS",
"RenderedEpisode",
"Tier",
"build_sft_dataset",
"iter_trajectories",
"render_all",
"select_best",
"step_difficulty",
"tier_telemetry",
"to_record",
"trajectory_rows_to_episode",
"GoldAction",
"ToolScaleTask",
"generate_sft_dataset",
"gold_coverage_verify",
"load_toolscale",
"normalize_row",
"trajectory_to_record",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py
-198
View File
@@ -1,198 +0,0 @@
"""Full-trajectory reader for NeuLab ADP (``neulab/agent-data-collection``).
The eval-side loader (:mod:`openjarvis.evals.datasets.adp`) keeps only
``(problem, reference)`` per row — it drops every intermediate step, which is
exactly the signal the orchestrator needs. This module instead transcribes
**all** turns of a trajectory into a canonical
:class:`~openjarvis.learning.intelligence.orchestrator.types.Episode`.
No GPU / no API keys: we read the demonstrated traces and re-tier them
downstream; we never re-execute a model here.
"""
from __future__ import annotations
import ast
from dataclasses import dataclass
from typing import Iterable, Iterator, List, MutableMapping, Optional
from openjarvis.learning.intelligence.orchestrator.types import (
Episode,
OrchestratorAction,
OrchestratorObservation,
)
HF_DATASET_ID = "neulab/agent-data-collection"
HF_SPLIT = "std"
# Same configs the eval loader concatenates; the meeting steer favours the
# coding / agentic / tool-use ones.
DEFAULT_CONFIGS: tuple[str, ...] = (
"codeactinstruct",
"code_feedback",
"openhands",
"agenttuning_os",
"agenttuning_db",
"swe-smith",
)
# ADP turn ``class_`` values that denote a tool/code action vs. plain message.
_CODE_CLASSES = {"code_action", "ipython_action", "bash_action"}
_SEARCH_CLASSES = {"search_action", "browse_action", "web_action"}
_MESSAGE_CLASSES = {"message_action"}
@dataclass
class CanonicalStep:
"""One transcribed ADP turn: what the demonstration did at this step."""
kind: str
"""Normalised step kind: ``reason`` | ``code`` | ``search`` | ``message``."""
content: str
"""The turn's text (the action the demonstration took)."""
observation: str = ""
"""The environment/tool result that followed, if any."""
is_final: bool = False
"""Whether this is the trajectory's final answer turn."""
def _parse_content(raw: object) -> List[MutableMapping[str, object]]:
"""Parse the ``content`` field (a list, or a string repr of one)."""
if isinstance(raw, list):
return raw # type: ignore[return-value]
if isinstance(raw, str):
try:
parsed = ast.literal_eval(raw)
if isinstance(parsed, list):
return parsed # type: ignore[return-value]
except (ValueError, SyntaxError):
pass
return []
def _classify(turn: MutableMapping[str, object]) -> str:
cls = str(turn.get("class_") or "").lower()
if cls in _CODE_CLASSES:
return "code"
if cls in _SEARCH_CLASSES:
return "search"
if cls in _MESSAGE_CLASSES:
return "message"
return "reason"
def trajectory_rows_to_episode(
record_id: str,
turns: List[MutableMapping[str, object]],
) -> Optional[Episode]:
"""Transcribe one ADP trajectory's turns into a canonical ``Episode``.
The orchestrator-relevant structure is preserved: the first user turn is the
problem; every subsequent agent turn becomes a step (with its following
observation, if the trace recorded one). Returns ``None`` if there is no
usable problem or no agent steps.
"""
problem: Optional[str] = None
steps: List[CanonicalStep] = []
# Group agent turns with the observation/user turn that follows them.
pending: Optional[CanonicalStep] = None
for turn in turns:
source = str(turn.get("source") or "").lower()
text = str(turn.get("content") or "").strip()
if not text:
continue
if source == "user":
if problem is None:
problem = text
elif pending is not None:
# A user/environment turn after an agent action = its observation.
pending.observation = text[:2000]
continue
# Agent-sourced turn -> a new step.
if pending is not None:
steps.append(pending)
pending = CanonicalStep(kind=_classify(turn), content=text)
if pending is not None:
steps.append(pending)
if not problem or not steps:
return None
steps[-1].is_final = True
steps[-1].kind = "message" if steps[-1].kind == "reason" else steps[-1].kind
episode = Episode(
task_id=record_id,
initial_prompt=problem,
ground_truth=steps[-1].content[:2000],
# ADP rows are demonstrated solutions; treat them as correct teachers.
correct=True,
)
for st in steps:
episode.add_step(
OrchestratorAction(
thought="", # filled in by the paradigm renderer
tool_name=st.kind,
tool_input=st.content,
is_final_answer=st.is_final,
),
OrchestratorObservation(content=st.observation or st.content),
)
episode.final_answer = steps[-1].content[:2000]
episode.metadata["source"] = "adp"
return episode
def iter_trajectories(
*,
max_tasks: Optional[int] = None,
configs: Iterable[str] = DEFAULT_CONFIGS,
min_steps: int = 1,
max_steps: int = 24,
) -> Iterator[Episode]:
"""Stream canonical ``Episode`` objects from ADP.
Network is touched lazily inside the loop (``datasets.load_dataset`` with
``streaming=True``) so importing this module stays free. Configs that fail
to load (gated/missing) are skipped.
"""
from datasets import load_dataset
emitted = 0
for cfg in configs:
if max_tasks is not None and emitted >= max_tasks:
break
try:
stream = load_dataset(HF_DATASET_ID, cfg, split=HF_SPLIT, streaming=True)
except Exception:
continue
for i, row in enumerate(stream):
if max_tasks is not None and emitted >= max_tasks:
break
row = dict(row) # type: ignore[arg-type]
turns = _parse_content(row.get("content"))
row_id = row.get("id")
rec_id = str(row_id) if row_id is not None else f"{cfg}-{i}"
episode = trajectory_rows_to_episode(rec_id, turns)
if episode is None:
continue
n = episode.num_turns()
if n < min_steps or n > max_steps:
continue
emitted += 1
yield episode
__all__ = [
"CanonicalStep",
"DEFAULT_CONFIGS",
"iter_trajectories",
"trajectory_rows_to_episode",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py
-112
View File
@@ -1,112 +0,0 @@
"""End-to-end SFT dataset builder + CLI.
ADP trajectory -> canonical Episode -> render_all -> select_best
-> to_record -> JSONL (+ a sidecar ``.stats.json``)
Runs with no GPU and no API keys (the cold-start re-tiers demonstrated traces;
it does not execute models). Network is only touched to stream ADP rows.
"""
from __future__ import annotations
import argparse
import json
import logging
from collections import Counter
from pathlib import Path
from typing import Callable, Iterable, Iterator, Optional
from openjarvis.learning.intelligence.orchestrator.reward import (
MultiObjectiveReward,
Normalizers,
RewardWeights,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import (
DEFAULT_CONFIGS,
iter_trajectories,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import render_all
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
from openjarvis.learning.intelligence.orchestrator.types import Episode
logger = logging.getLogger(__name__)
def build_sft_dataset(
out_path: str,
*,
max_tasks: Optional[int] = 2000,
configs: Iterable[str] = DEFAULT_CONFIGS,
source: Optional[Callable[..., Iterator[Episode]]] = None,
) -> dict:
"""Build the SFT JSONL at ``out_path`` and return stats.
``source`` overrides the ADP stream (used by tests to inject fixtures); it
must be a callable returning an iterator of canonical :class:`Episode`.
"""
reward = MultiObjectiveReward(RewardWeights(), Normalizers())
out = Path(out_path)
out.parent.mkdir(parents=True, exist_ok=True)
episodes = (
source(max_tasks=max_tasks, configs=configs)
if source is not None
else iter_trajectories(max_tasks=max_tasks, configs=configs)
)
seen = 0
written = 0
dropped = 0
paradigm_counts: Counter[str] = Counter()
with out.open("w") as fh:
for episode in episodes:
seen += 1
best = select_best(render_all(episode), reward=reward)
if best is None:
dropped += 1
continue
record = to_record(best, reward=reward.compute(best.episode))
fh.write(json.dumps(record) + "\n")
written += 1
paradigm_counts[best.paradigm] += 1
stats = {
"out_path": str(out),
"tasks_seen": seen,
"records_written": written,
"tasks_dropped": dropped,
"paradigm_distribution": dict(paradigm_counts),
}
stats_path = out.with_suffix(out.suffix + ".stats.json")
stats_path.write_text(json.dumps(stats, indent=2))
logger.info("Wrote %d SFT records to %s (%s)", written, out, dict(paradigm_counts))
return stats
def _main(argv: Optional[list[str]] = None) -> int:
parser = argparse.ArgumentParser(
description="Build orchestrator SFT data from ADP."
)
parser.add_argument("--out", default="data/orchestrator_sft_traces.jsonl")
parser.add_argument("--max-tasks", type=int, default=2000)
parser.add_argument(
"--adp-configs",
default=",".join(DEFAULT_CONFIGS),
help="comma-separated ADP sub-configs to stream",
)
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(message)s")
configs = [c.strip() for c in args.adp_configs.split(",") if c.strip()]
stats = build_sft_dataset(args.out, max_tasks=args.max_tasks, configs=configs)
print(json.dumps(stats, indent=2))
return 0
if __name__ == "__main__": # pragma: no cover
raise SystemExit(_main())
__all__ = ["build_sft_dataset"]
src/openjarvis/learning/intelligence/orchestrator/sft_data/datasets.py
+282
View File
@@ -0,0 +1,282 @@
"""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
@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),
)
def load_generalthought(
*,
n: int = 2000,
seed: int = 42,
source: Optional[Iterable[Dict[str, Any]]] = 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.
"""
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 = 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)
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.
The HF stream is sharded by domain (all-code first), so we route rows into
per-domain buffers and stop once every quota is met. ``source`` overrides the
stream with raw row dicts (tests).
"""
if source is None:
from datasets import load_dataset # lazy: optional dep / network
source = load_dataset(OPENTHOUGHTS_ID, split="train", streaming=True)
quotas = {"code": n_code, "math": n_math, "science": n_science}
bufs: Dict[str, List[Task]] = {k: [] for k in quotas}
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
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) -> List[Task]:
"""The 8K cold-start SFT set: 2K GeneralThought + 2K code + 2K math + 2K science."""
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",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/hotpotqa.py
+155
View File
@@ -0,0 +1,155 @@
"""Loader + answer verifier for HotpotQA (``hotpotqa/hotpot_qa``).
HotpotQA is multi-hop Wikipedia QA: each task is a question with a short
factoid ``answer`` and a difficulty ``level`` (easy/medium/hard). We use it as
the orchestrator SFT cold-start source because (a) answers are exact-match
verifiable, so rejection sampling needs **no LLM judge**, and (b) solving a
question requires ``web_search`` + multi-step reasoning, which is exactly the
local↔cloud routing signal we want to teach.
We load the ``fullwiki`` config (the agent must retrieve from all of Wikipedia
rather than being handed the gold paragraphs); we only keep ``question`` /
``answer`` / ``level`` / ``type`` since the orchestrator does its own retrieval.
``load_hotpotqa`` streams via the HuggingFace ``datasets`` library; tests pass
``source=`` an iterable of raw row dicts so normalization is exercised offline.
"""
from __future__ import annotations
import re
import string
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Iterator, List, Optional
DATASET_ID = "hotpotqa/hotpot_qa"
CONFIG = "fullwiki"
SPLIT = "validation" # train/validation carry gold answers; test does not
@dataclass
class HotpotTask:
task_id: str
question: str
answer: str
level: str = "" # easy | medium | hard
qtype: str = "" # comparison | bridge
# Parity with ToolScaleTask so the rejection-sampling loop is dataset-agnostic.
@property
def instruction(self) -> str:
return self.question
@property
def domain(self) -> str:
return f"hotpotqa/{self.level or 'unknown'}"
def normalize_row(row: Dict[str, Any], *, index: int = 0) -> Optional[HotpotTask]:
"""Turn one raw HotpotQA row into a :class:`HotpotTask` (pure)."""
question = str(row.get("question") or "").strip()
answer = str(row.get("answer") or "").strip()
if not question or not answer:
return None
task_id = str(row.get("id") or row.get("_id") or f"hotpotqa-{index}")
return HotpotTask(
task_id=task_id,
question=question,
answer=answer,
level=str(row.get("level") or ""),
qtype=str(row.get("type") or ""),
)
def load_hotpotqa(
*,
max_tasks: Optional[int] = None,
split: str = SPLIT,
config: str = CONFIG,
source: Optional[Iterable[Dict[str, Any]]] = None,
) -> Iterator[HotpotTask]:
"""Yield normalized HotpotQA tasks.
``source`` overrides the HF stream with an iterable of raw row dicts (tests).
When ``source`` is None, streams ``hotpotqa/hotpot_qa`` via ``datasets``.
"""
if source is None:
from datasets import load_dataset # lazy: optional dep / network
source = load_dataset(DATASET_ID, config, 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 task is None:
continue
yield task
n += 1
# --- answer verification (official HotpotQA-style normalization) -------------
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 answer_matches(prediction: str, gold: str, *, f1_threshold: float = 0.6) -> bool:
"""True if ``prediction`` (free-form) contains / matches the gold answer.
The orchestrator answers in prose, so we accept (a) the normalized gold as a
whole-word substring of the prediction, or (b) token-F1 >= threshold. This is
the dependency-free, no-LLM verifier that makes rejection sampling automatic.
"""
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
def make_verifier(f1_threshold: float = 0.6):
"""Return a ``verify_fn(task, rollout) -> bool`` for the rejection sampler."""
def verify(task: HotpotTask, rollout: Any) -> bool:
ans = getattr(rollout, "final_answer", "") or ""
if not ans.strip():
return False
return answer_matches(ans, task.answer, f1_threshold=f1_threshold)
return verify
__all__ = [
"CONFIG",
"DATASET_ID",
"HotpotTask",
"SPLIT",
"answer_matches",
"load_hotpotqa",
"make_verifier",
"normalize_row",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py
-188
View File
@@ -1,188 +0,0 @@
"""Render a canonical ADP ``Episode`` under each local↔cloud paradigm.
Each renderer assigns a model **tier** to every step according to that
paradigm's policy, attaches the estimated telemetry (cost/energy/latency/power)
to the step's observation, and predicts whether the rendering would still solve
the task. Selection (``select.py``) then keeps the cheapest predicted-correct
rendering per task.
These are the cold-start *tiering* scaffolds of the real
``agents/hybrid/*`` paradigms — same names, same local-first spirit, no live
execution. Swapping in real runs is the v2 item in the design doc.
"""
from __future__ import annotations
import copy
from dataclasses import dataclass, field
from typing import Callable, List
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
Difficulty,
Tier,
covers,
min_covering_tier,
step_difficulty,
tier_name,
tier_telemetry,
)
from openjarvis.learning.intelligence.orchestrator.types import Episode
# Detected in an observation -> the demonstration had to retry -> harder task.
_RETRY_MARKERS = (
"your answer is wrong",
"incorrect",
"traceback",
"error:",
"failed",
"try again",
)
@dataclass
class RenderedEpisode:
"""An ``Episode`` re-tiered under one paradigm, plus its prediction."""
paradigm: str
episode: Episode
"""Deep copy with per-step tier in ``action.tool_name`` and telemetry filled."""
step_tiers: List[str] = field(default_factory=list)
"""Tier key (``local``/``mid``/``frontier``/``search``) chosen per step."""
predicted_correct: bool = False
def _difficulties(episode: Episode) -> List[Difficulty]:
"""Per-step difficulty, with a trajectory-wide retry bump."""
retry = any(
any(m in (step.observation.content or "").lower() for m in _RETRY_MARKERS)
for step in episode.steps
)
out: List[Difficulty] = []
for step in episode.steps:
kind = step.action.tool_name
out.append(step_difficulty(kind, step.action.tool_input, retry_signal=retry))
return out
def _apply(
episode: Episode,
paradigm: str,
tier_keys: List[str],
predicted_correct: bool,
) -> RenderedEpisode:
"""Build a RenderedEpisode: stamp each step with its tier + telemetry."""
ep = copy.deepcopy(episode)
ep.total_cost_usd = 0.0
ep.total_energy_joules = 0.0
ep.total_latency_seconds = 0.0
ep.total_tokens = 0
ep.max_power_watts = 0.0
for step, tier_key in zip(ep.steps, tier_keys):
tel = tier_telemetry(tier_key)
step.observation.cost_usd = tel["cost_usd"]
step.observation.energy_joules = tel["energy_joules"]
step.observation.latency_seconds = tel["latency_seconds"]
step.observation.power_watts = tel["power_watts"]
step.observation.tokens = int(tel["tokens"])
# Record the routing decision on the action for the serializer.
step.action.thought = _thought_for(tier_key, step.action.tool_name)
ep.total_cost_usd += tel["cost_usd"]
ep.total_energy_joules += tel["energy_joules"]
ep.total_latency_seconds += tel["latency_seconds"]
ep.total_tokens += int(tel["tokens"])
ep.max_power_watts = max(ep.max_power_watts, tel["power_watts"])
ep.correct = predicted_correct
ep.metadata["paradigm"] = paradigm
return RenderedEpisode(
paradigm=paradigm,
episode=ep,
step_tiers=tier_keys,
predicted_correct=predicted_correct,
)
def _thought_for(tier_key: str, kind: str) -> str:
if kind == "search" or tier_key == "search":
return "This step needs retrieval; route to web_search."
reason = {
"local": "Cheap and on-device; the local model can handle this step.",
"mid": "Beyond easy; escalate to a cheap mid-tier cloud model.",
"frontier": "Hard step; escalate to the frontier cloud model.",
}[tier_key]
return reason
def _tier_key_for_step(kind: str, tier: Tier) -> str:
return "search" if kind == "search" else tier_name(tier)
# --- paradigm renderers -----------------------------------------------------
def render_baseline_local(episode: Episode) -> RenderedEpisode:
diffs = _difficulties(episode)
tiers = [
_tier_key_for_step(s.action.tool_name, Tier.LOCAL) for s in episode.steps
]
predicted = all(d == Difficulty.EASY for d in diffs)
return _apply(episode, "baseline_local", tiers, predicted)
def render_baseline_cloud(episode: Episode) -> RenderedEpisode:
tiers = [
_tier_key_for_step(s.action.tool_name, Tier.FRONTIER) for s in episode.steps
]
return _apply(episode, "baseline_cloud", tiers, True)
def render_advisor(episode: Episode) -> RenderedEpisode:
"""Local executor, frontier rewrite of the final answer."""
diffs = _difficulties(episode)
tiers: List[str] = []
n = len(episode.steps)
for i, s in enumerate(episode.steps):
tier = Tier.FRONTIER if i == n - 1 else Tier.LOCAL
tiers.append(_tier_key_for_step(s.action.tool_name, tier))
predicted = all(d == Difficulty.EASY for d in diffs[:-1]) if n > 1 else True
return _apply(episode, "advisor", tiers, predicted)
def render_toolorchestra(episode: Episode) -> RenderedEpisode:
"""Local-first per-step: assign each step its minimum covering tier."""
diffs = _difficulties(episode)
tiers = [
_tier_key_for_step(s.action.tool_name, min_covering_tier(d))
for s, d in zip(episode.steps, diffs)
]
predicted = all(
covers(min_covering_tier(d), d) for d in diffs
) # True by construction
return _apply(episode, "toolorchestra", tiers, predicted)
PARADIGMS: dict[str, Callable[[Episode], RenderedEpisode]] = {
"baseline_local": render_baseline_local,
"baseline_cloud": render_baseline_cloud,
"advisor": render_advisor,
"toolorchestra": render_toolorchestra,
}
def render_all(episode: Episode) -> List[RenderedEpisode]:
"""Render ``episode`` under every paradigm."""
return [render(episode) for render in PARADIGMS.values()]
__all__ = [
"PARADIGMS",
"RenderedEpisode",
"render_advisor",
"render_all",
"render_baseline_cloud",
"render_baseline_local",
"render_toolorchestra",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py
+12 -7
View File
@@ -23,22 +23,27 @@ from collections import Counter
from pathlib import Path
from typing import Callable, Iterable, List, Optional
from typing import Any
from openjarvis.agents.hybrid.expert_registry import ExpertTool
from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout
from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
ToolScaleTask,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
trajectory_to_record,
)
logger = logging.getLogger(__name__)
RolloutFn = Callable[[ToolScaleTask], Optional[UnifiedRollout]]
VerifyFn = Callable[[ToolScaleTask, UnifiedRollout], bool]
# 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]
def gold_coverage_verify(task: ToolScaleTask, rollout: UnifiedRollout) -> bool:
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.
@@ -57,7 +62,7 @@ def gold_coverage_verify(task: ToolScaleTask, rollout: UnifiedRollout) -> bool:
def generate_sft_dataset(
out_path: str,
*,
tasks: Iterable[ToolScaleTask],
tasks: Iterable[TaskLike],
tools: List[ExpertTool],
rollout_fn: RolloutFn,
verify_fn: VerifyFn = gold_coverage_verify,
src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py
-42
View File
@@ -1,42 +0,0 @@
"""Pick the best paradigm rendering per task.
The label for SFT is the *cheapest strategy that still solves the task*: among
the predicted-correct renderings, keep the one with the highest multi-objective
reward (accuracy minus cost/energy/latency/power). This is the local-first
thesis expressed as a training target.
"""
from __future__ import annotations
from typing import List, Optional
from openjarvis.learning.intelligence.orchestrator.reward import (
MultiObjectiveReward,
Normalizers,
RewardWeights,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
RenderedEpisode,
)
_DEFAULT_REWARD = MultiObjectiveReward(RewardWeights(), Normalizers())
def select_best(
renderings: List[RenderedEpisode],
*,
reward: Optional[MultiObjectiveReward] = None,
) -> Optional[RenderedEpisode]:
"""Return the highest-reward predicted-correct rendering, or ``None``.
``None`` means no paradigm was predicted to solve the task — the task is
dropped from the SFT set rather than teaching a wrong trajectory.
"""
scorer = reward or _DEFAULT_REWARD
correct = [r for r in renderings if r.predicted_correct]
if not correct:
return None
return max(correct, key=lambda r: scorer.compute(r.episode))
__all__ = ["select_best"]
src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py
-103
View File
@@ -1,103 +0,0 @@
"""Serialize a winning :class:`RenderedEpisode` into a ``conversations`` record.
Output schema matches what
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
already consumes::
{"conversations": [{"role": "system"|"user"|"assistant"|"tool", "content": ...}],
"paradigm": ..., "reward": ..., "metrics": {...}}
The assistant turns use the canonical THOUGHT/TOOL/INPUT format from
``prompt_registry``; the routing tool encodes the tier the orchestrator chose
(``local_model`` / ``mid_model`` / ``frontier_model`` / ``web_search``).
"""
from __future__ import annotations
from typing import Any, Dict
from openjarvis.learning.intelligence.orchestrator.prompt_registry import (
build_system_prompt,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
RenderedEpisode,
)
# Routing vocabulary the orchestrator learns to emit.
TIER_TOOL = {
"local": "local_model",
"mid": "mid_model",
"frontier": "frontier_model",
"search": "web_search",
}
ROUTING_TOOLS = ["local_model", "mid_model", "frontier_model", "web_search"]
_SYSTEM_PROMPT = build_system_prompt(ROUTING_TOOLS)
def to_record(rendered: RenderedEpisode, *, reward: float = 0.0) -> Dict[str, Any]:
"""Convert a winning rendering into one SFT JSONL record."""
ep = rendered.episode
conversations: list[dict[str, str]] = [
{"role": "system", "content": _SYSTEM_PROMPT},
{"role": "user", "content": ep.initial_prompt},
]
n = len(ep.steps)
for i, (step, tier_key) in enumerate(zip(ep.steps, rendered.step_tiers)):
tool = TIER_TOOL.get(tier_key, "local_model")
if step.action.is_final_answer or i == n - 1:
conversations.append(
{
"role": "assistant",
"content": (
f"THOUGHT: {step.action.thought}\n"
f"TOOL: {tool}\n"
f"INPUT: {step.action.tool_input}"
),
}
)
conversations.append(
{"role": "tool", "name": tool, "content": step.observation.content}
)
conversations.append(
{
"role": "assistant",
"content": (
"THOUGHT: The result answers the task.\n"
f"FINAL_ANSWER: {ep.final_answer}"
),
}
)
else:
conversations.append(
{
"role": "assistant",
"content": (
f"THOUGHT: {step.action.thought}\n"
f"TOOL: {tool}\n"
f"INPUT: {step.action.tool_input}"
),
}
)
conversations.append(
{"role": "tool", "name": tool, "content": step.observation.content}
)
return {
"conversations": conversations,
"task_id": ep.task_id,
"paradigm": rendered.paradigm,
"reward": reward,
"metrics": {
"cost_usd": ep.total_cost_usd,
"energy_joules": ep.total_energy_joules,
"latency_seconds": ep.total_latency_seconds,
"tokens": ep.total_tokens,
"max_power_watts": ep.max_power_watts,
"num_steps": n,
},
}
__all__ = ["ROUTING_TOOLS", "TIER_TOOL", "to_record"]
src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py
-119
View File
@@ -1,119 +0,0 @@
"""Model tiers, per-tier telemetry estimates, and step-difficulty heuristics.
This is the *one estimate* of the cold-start (see the design doc): because we
do not execute models, we approximate "what would this step cost on tier X?"
with flat per-tier telemetry, and "how hard is this step?" with a heuristic over
the ADP step kind and the trajectory's retry signals.
Costs reuse the authoritative hybrid pricing table so the orchestrator is
trained against the same dollars the paradigm harness charges.
"""
from __future__ import annotations
from enum import IntEnum
from openjarvis.agents.hybrid import _prices
# Representative model per tier (drives the $ side of telemetry).
TIER_MODEL: dict[str, str] = {
"local": "qwen3:8b", # local vLLM -> unknown to PRICES -> $0
"mid": "gemini-2.5-flash",
"frontier": "claude-sonnet-4-6",
"search": "qwen3:8b", # retrieval is a tool call, not a model tier
}
# Rough token budgets per step kind, for the cost estimate.
_TOKENS_IN = 800
_TOKENS_OUT = {"local": 400, "mid": 400, "frontier": 600, "search": 80}
# Non-$ telemetry per step (joules / seconds / watts). Local work burns some
# on-device energy/power but no dollars; cloud work is ~free locally but costs
# dollars and adds latency. Tuned (against reward.py's default weights and
# normalizers) so the local-first ordering holds: on a step a tier can cover,
# local beats cloud, and per-step escalation beats all-frontier. These are the
# tunable knobs of the cold-start estimate — see the design doc.
_ENERGY_J = {"local": 5.0, "mid": 0.5, "frontier": 0.5, "search": 0.2}
_LATENCY_S = {"local": 1.5, "mid": 1.5, "frontier": 4.0, "search": 1.0}
_POWER_W = {"local": 20.0, "mid": 5.0, "frontier": 5.0, "search": 5.0}
_SEARCH_COST_USD = 0.001
class Tier(IntEnum):
"""Model capability tier; ordered so a higher tier covers a harder step."""
LOCAL = 0
MID = 1
FRONTIER = 2
class Difficulty(IntEnum):
"""Estimated step difficulty; compared against :class:`Tier` rank."""
EASY = 0
MED = 1
HARD = 2
_TIER_NAME = {Tier.LOCAL: "local", Tier.MID: "mid", Tier.FRONTIER: "frontier"}
def tier_name(tier: Tier) -> str:
return _TIER_NAME[tier]
def tier_telemetry(tier_key: str) -> dict[str, float]:
"""Estimated ``(cost, energy, latency, power, tokens)`` for one step at a tier.
``tier_key`` is one of ``local``/``mid``/``frontier``/``search``.
"""
out_tokens = _TOKENS_OUT[tier_key]
if tier_key == "search":
cost = _SEARCH_COST_USD
else:
cost = _prices.cost(TIER_MODEL[tier_key], _TOKENS_IN, out_tokens)
return {
"cost_usd": cost,
"energy_joules": _ENERGY_J[tier_key],
"latency_seconds": _LATENCY_S[tier_key],
"power_watts": _POWER_W[tier_key],
"tokens": float(_TOKENS_IN + out_tokens),
}
def step_difficulty(
kind: str, content: str, *, retry_signal: bool = False
) -> Difficulty:
"""Heuristic difficulty for a canonical step.
- ``code`` actions are at least MED (real execution / logic).
- very long actions, or any step in a trajectory that hit a retry signal
("your answer is wrong" etc.), bump to HARD.
- everything else (plain messages / short reasoning) is EASY.
"""
base = Difficulty.MED if kind == "code" else Difficulty.EASY
if retry_signal or len(content) > 1500:
return Difficulty.HARD
return base
def covers(tier: Tier, difficulty: Difficulty) -> bool:
"""True iff ``tier`` is capable enough for a step of this ``difficulty``."""
return int(tier) >= int(difficulty)
def min_covering_tier(difficulty: Difficulty) -> Tier:
"""The cheapest tier that still covers ``difficulty`` (local-first optimum)."""
return Tier(int(difficulty))
__all__ = [
"Difficulty",
"Tier",
"TIER_MODEL",
"covers",
"min_covering_tier",
"step_difficulty",
"tier_name",
"tier_telemetry",
]
src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py
+4 -8
View File
@@ -17,19 +17,15 @@ from typing import Any, Dict, List
from openjarvis.agents.hybrid.expert_registry import ExpertTool, build_tool_specs
from openjarvis.agents.hybrid.toolorchestra.rollout import (
RL_ORCHESTRATOR_SYS,
UnifiedRollout,
build_system_prompt,
tool_call_tag,
)
def _system_prompt(tools: List[ExpertTool]) -> str:
specs = build_tool_specs(tools)
return (
RL_ORCHESTRATOR_SYS
+ "\n\nAvailable tools (call one per turn, or answer directly):\n"
+ json.dumps(specs, indent=2)
)
# Faithful ToolOrchestra system prompt (paper's prepare_sft_data.py format).
return build_system_prompt(build_tool_specs(tools))
def trajectory_to_record(
@@ -44,7 +40,7 @@ def trajectory_to_record(
"""Convert a passing :class:`UnifiedRollout` into one SFT JSONL record."""
conversations: List[Dict[str, str]] = [
{"role": "system", "content": _system_prompt(tools)},
{"role": "user", "content": f"Problem: {question}\n\nChoose an appropriate tool."},
{"role": "user", "content": f"Problem: {question}"},
]
for turn in rollout.turns:
src/openjarvis/learning/intelligence/orchestrator/sft_data/verify.py
+254
View File
@@ -0,0 +1,254 @@
"""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"
# --- 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
# Symbolic / exact equality via sympy when available.
try:
import sympy
from sympy.parsing.latex import parse_latex # noqa: F401 (optional)
def _parse(x: str):
try:
return sympy.sympify(x.replace("^", "**"))
except Exception:
try:
return sympy.parsing.latex.parse_latex(x)
except Exception:
return None
pe, ge = _parse(pred), _parse(g)
if pe is not None and ge is not None:
try:
if sympy.simplify(pe - ge) == 0:
return True
except Exception:
pass
except Exception:
pass
# 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("GEMINI_API_KEY")
if not api_key:
return None
try:
from openai import OpenAI
except Exception:
return None
try:
client = OpenAI(api_key=api_key, base_url=GEMINI_BASE_URL)
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.chat.completions.create(
model=GEMINI_MODEL,
messages=[{"role": "user", "content": prompt}],
temperature=0.0,
max_tokens=8,
)
verdict = (resp.choices[0].message.content or "").strip().upper()
if "PASS" in verdict:
return True
if "FAIL" in verdict:
return False
return None
except Exception:
return None
# --- public API --------------------------------------------------------------
def verify_answer(task: Task, prediction: str) -> bool:
"""Domain-dispatched correctness check for ``prediction`` against ``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":
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",
]
src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py
+74 -25
View File
@@ -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,21 +63,27 @@ 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
# ADP cold-start trace synthesis (see sft_data/build.py). Empty configs
# falls back to sft_data.adp_loader.DEFAULT_CONFIGS.
adp_configs: str = ""
distill_max_tasks: int = 2000
# 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"
@@ -272,31 +278,74 @@ class OrchestratorSFTTrainer:
)
def _generate_traces(self) -> None:
"""Generate SFT traces from the ADP corpus (no GPU / no API keys).
"""Generate SFT traces by base Qwen3-8B self-sampling (v1 cold-start).
Builds the THOUGHT/TOOL/INPUT ``conversations`` JSONL by re-tiering
NeuLab ADP trajectories through the local↔cloud paradigms and keeping
the cheapest predicted-correct rendering per task. See
``sft_data/build.py`` and the cold-start design doc.
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.learning.intelligence.orchestrator.sft_data.build import (
build_sft_dataset,
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,
)
trace_path = Path(self.config.trace_cache_path)
configs = (
[c.strip() for c in self.config.adp_configs.split(",") if c.strip()]
or None
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:
stats = build_sft_dataset(
tasks = load_sft_tasks()
if self.config.distill_max_tasks:
tasks = tasks[: self.config.distill_max_tasks]
stats = generate_sft_dataset(
str(trace_path),
max_tasks=self.config.distill_max_tasks,
**({"configs": configs} if configs else {}),
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("ADP trace generation failed (%s); writing empty set", exc)
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()
tests/agents/test_expert_registry.py
+123 -5
View File
@@ -7,10 +7,16 @@ import random
import pytest
from openjarvis.agents.hybrid.expert_registry import (
CATEGORY_BASIC,
CATEGORY_CLOUD_FRONTIER,
CATEGORY_LOCAL_OSS,
ExpertTool,
KIND_MODEL,
KIND_TOOL,
build_tool_specs,
default_catalog,
openjarvis_tool,
orchestrator_catalog,
sample_tool_config,
to_worker_dict,
tools_by_name,
@@ -22,8 +28,8 @@ def test_each_model_is_its_own_tool():
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", "qwen2_5_coder_32b",
"llama_3_3_70b", "claude_opus"):
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
@@ -37,9 +43,10 @@ def test_catalog_names_unique_and_valid():
def test_local_model_included_only_when_served():
assert "local_model" not in {t.name for t in default_catalog()}
# 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)["local_model"]
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
@@ -92,10 +99,121 @@ def test_price_jitter_changes_prices_reproducibly():
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
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 == "vllm"
assert by[n].price_in == 0.0 and by[n].price_out == 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():
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"].base_url == "http://x/v1"
assert by["qwen3_6_27b_fp8"].base_url == "http://y/v1"
# Unmapped local model -> base_url None.
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_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["local_model"])
local = to_worker_dict(by["qwen3_8b"])
assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"
tests/test_orchestrator_learning/sft_data/test_datasets.py
+136
View File
@@ -0,0 +1,136 @@
"""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?"
tests/test_orchestrator_learning/sft_data/test_sft_data.py
-187
View File
@@ -1,187 +0,0 @@
"""Offline tests for the orchestrator SFT cold-start pipeline (no network/GPU).
Covers: ADP transcription -> canonical Episode, difficulty/tier heuristics,
paradigm rendering, reward-ranked selection, serialization round-trip, and the
end-to-end builder via an injected fixture source.
"""
from __future__ import annotations
import json
from pathlib import Path
from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import (
trajectory_rows_to_episode,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.build import (
build_sft_dataset,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
PARADIGMS,
render_all,
)
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
Difficulty,
Tier,
covers,
min_covering_tier,
step_difficulty,
tier_telemetry,
)
from openjarvis.learning.intelligence.orchestrator.types import Episode
# --- fixtures ---------------------------------------------------------------
_EASY_TURNS = [
{"source": "user", "class_": "message_action", "content": "What is 2 + 2?"},
{"source": "agent", "class_": "message_action", "content": "The answer is 4."},
]
_CODE_TURNS = [
{
"source": "user",
"class_": "message_action",
"content": "Sort this list in Python.",
},
{"source": "agent", "class_": "code_action", "content": "sorted([3, 1, 2])"},
{"source": "user", "class_": "message_action", "content": "[1, 2, 3]"},
{"source": "agent", "class_": "message_action", "content": "Sorted: [1, 2, 3]"},
]
def _easy_episode() -> Episode:
ep = trajectory_rows_to_episode("easy-1", _EASY_TURNS)
assert ep is not None
return ep
def _code_episode() -> Episode:
ep = trajectory_rows_to_episode("code-1", _CODE_TURNS)
assert ep is not None
return ep
# --- adp_loader -------------------------------------------------------------
def test_transcribe_keeps_all_steps_and_problem():
ep = _code_episode()
assert ep.initial_prompt == "Sort this list in Python."
assert ep.num_turns() == 2 # two agent turns
assert ep.steps[0].action.tool_name == "code"
assert ep.steps[-1].action.is_final_answer is True
assert ep.correct is True # ADP rows are demonstrated solutions
def test_transcribe_rejects_empty():
assert trajectory_rows_to_episode("x", []) is None
assert trajectory_rows_to_episode("x", [{"source": "user", "content": ""}]) is None
# --- tiers ------------------------------------------------------------------
def test_difficulty_heuristic():
assert step_difficulty("message", "short") == Difficulty.EASY
assert step_difficulty("code", "sorted([])") == Difficulty.MED
assert step_difficulty("message", "x", retry_signal=True) == Difficulty.HARD
def test_min_covering_tier_and_covers():
assert min_covering_tier(Difficulty.EASY) == Tier.LOCAL
assert min_covering_tier(Difficulty.MED) == Tier.MID
assert covers(Tier.FRONTIER, Difficulty.HARD)
assert not covers(Tier.LOCAL, Difficulty.HARD)
def test_tier_telemetry_local_is_free_cloud_costs():
assert tier_telemetry("local")["cost_usd"] == 0.0
assert tier_telemetry("frontier")["cost_usd"] > 0.0
local_e = tier_telemetry("local")["energy_joules"]
mid_e = tier_telemetry("mid")["energy_joules"]
assert local_e > mid_e
# --- paradigms --------------------------------------------------------------
def test_render_all_covers_every_paradigm():
rendered = render_all(_code_episode())
assert {r.paradigm for r in rendered} == set(PARADIGMS)
def test_baseline_local_incorrect_on_code_task():
rendered = {r.paradigm: r for r in render_all(_code_episode())}
# code step is MED -> local can't cover it
assert rendered["baseline_local"].predicted_correct is False
assert rendered["baseline_cloud"].predicted_correct is True
assert rendered["toolorchestra"].predicted_correct is True
def test_toolorchestra_escalates_code_step():
rendered = {r.paradigm: r for r in render_all(_code_episode())}
tiers = rendered["toolorchestra"].step_tiers
assert tiers[0] == "mid" # code step escalated off local
assert rendered["baseline_local"].step_tiers[0] == "local"
# --- select -----------------------------------------------------------------
def test_select_prefers_local_on_easy_task():
best = select_best(render_all(_easy_episode()))
assert best is not None
# all-easy -> baseline_local is cheapest-correct
assert best.paradigm == "baseline_local"
def test_select_drops_unsolved_task():
ep = _easy_episode()
# force every paradigm to be predicted-incorrect
renderings = render_all(ep)
for r in renderings:
r.predicted_correct = False
assert select_best(renderings) is None
# --- serialize --------------------------------------------------------------
def test_serialize_schema_and_final_answer():
best = select_best(render_all(_code_episode()))
assert best is not None
rec = to_record(best, reward=0.5)
convo = rec["conversations"]
assert convo[0]["role"] == "system"
assert convo[1]["role"] == "user"
assert any(
c["role"] == "assistant" and "FINAL_ANSWER:" in c["content"] for c in convo
)
assert rec["paradigm"] == best.paradigm
assert rec["metrics"]["num_steps"] == best.episode.num_turns()
# --- build (end-to-end, injected source) ------------------------------------
def test_build_sft_dataset_end_to_end(tmp_path: Path):
def fake_source(*, max_tasks=None, configs=None):
yield _easy_episode()
yield _code_episode()
out = tmp_path / "traces.jsonl"
stats = build_sft_dataset(str(out), max_tasks=10, source=fake_source)
assert stats["records_written"] == 2
assert stats["tasks_seen"] == 2
lines = out.read_text().strip().splitlines()
assert len(lines) == 2
for line in lines:
rec = json.loads(line)
assert rec["conversations"][0]["role"] == "system"
stats_file = out.with_suffix(out.suffix + ".stats.json")
assert stats_file.exists()
assert sum(stats["paradigm_distribution"].values()) == 2
tests/test_orchestrator_learning/sft_data/test_verify.py
+107
View File
@@ -0,0 +1,107 @@
"""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():
# x^2 + 2x + 1 == (x+1)^2 via sympy (skips gracefully if sympy absent).
t = _math("(x+1)^2")
assert verify_answer(t, "x^2 + 2*x + 1") is True
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
tests/test_orchestrator_learning/test_build_orchestrator_sft.py
+124
View File
@@ -0,0 +1,124 @@
"""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() -> list[Task]:
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="that's the result", 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).
out = tmp_path / "v1.jsonl"
rc = drv.main([
"--out", str(out),
"--samples-per-task", "1",
"--max-tasks", "2",
])
assert rc == 0
lines = out.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"]
tests/test_orchestrator_learning/test_eval_backend.py
+170
View File
@@ -0,0 +1,170 @@
"""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.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.turns = [object()]
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")
tests/test_orchestrator_learning/test_grpo_epoch.py
+408
View File
@@ -0,0 +1,408 @@
"""Offline (no-GPU) tests for OrchestratorGRPOTrainer._train_epoch.
We never touch torch/model/network: ``load_grpo_prompts`` is monkeypatched
to return a handful of fake Tasks, and ``_grpo_step`` is monkeypatched to
return canned metrics. The test asserts ``_train_epoch`` chunks the prompt
pool correctly and aggregates per-chunk metrics.
"""
from __future__ import annotations
from dataclasses import dataclass
import pytest
import openjarvis.learning.intelligence.orchestrator.policy_model as policy_model_mod
import openjarvis.learning.intelligence.orchestrator.sft_data.datasets as datasets
from openjarvis.learning.intelligence.orchestrator.grpo_trainer import (
OrchestratorGRPOConfig,
OrchestratorGRPOTrainer,
RolloutRecord,
)
from openjarvis.learning.intelligence.orchestrator.policy_model import (
OrchestratorPolicyModel,
)
@dataclass
class FakeTask:
task_id: str
question: str
answer: str
domain: str = "math"
@property
def instruction(self) -> str:
return self.question
def _make_tasks(n: int):
return [
FakeTask(task_id=f"t{i}", question=f"q{i}", answer=str(i))
for i in range(n)
]
def _build_trainer(monkeypatch, n_prompts: int, prompts_per_step: int):
"""Build a trainer with no real model and patched data loading.
``from_pretrained`` is stubbed to a model-less policy so construction
touches no GPU / network even when torch happens to be installed.
"""
tasks = _make_tasks(n_prompts)
monkeypatch.setattr(datasets, "load_grpo_prompts", lambda *a, **k: tasks)
monkeypatch.setattr(
policy_model_mod.OrchestratorPolicyModel,
"from_pretrained",
classmethod(lambda cls, *a, **k: OrchestratorPolicyModel(model=None)),
)
cfg = OrchestratorGRPOConfig(
grpo_max_prompts=n_prompts,
prompts_per_step=prompts_per_step,
num_samples_per_prompt=2,
)
trainer = OrchestratorGRPOTrainer(cfg)
return trainer, tasks
def test_train_epoch_chunks_and_aggregates(monkeypatch):
trainer, tasks = _build_trainer(monkeypatch, n_prompts=10, prompts_per_step=4)
calls = []
def fake_step(chunk):
calls.append(len(chunk))
return {
"loss": 1.0,
"reward": 0.5,
"accuracy": 1.0,
"n_prompts": len(chunk),
}
monkeypatch.setattr(trainer, "_grpo_step", fake_step)
metrics = trainer._train_epoch(epoch=0)
# 10 prompts / 4 per step => chunks of [4, 4, 2]
assert calls == [4, 4, 2]
assert metrics["n_batches"] == 3
assert metrics["n_prompts"] == 10
# Weighted means: all chunks report the same per-prompt values.
assert abs(metrics["loss"] - 1.0) < 1e-9
assert abs(metrics["reward"] - 0.5) < 1e-9
assert abs(metrics["accuracy"] - 1.0) < 1e-9
assert metrics["epoch"] == 0
def test_train_epoch_exact_division(monkeypatch):
trainer, _ = _build_trainer(monkeypatch, n_prompts=8, prompts_per_step=4)
chunks = []
monkeypatch.setattr(
trainer,
"_grpo_step",
lambda chunk: (
chunks.append(len(chunk)),
{"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)},
)[1],
)
metrics = trainer._train_epoch(epoch=2)
assert chunks == [4, 4]
assert metrics["n_batches"] == 2
assert metrics["n_prompts"] == 8
assert metrics["epoch"] == 2
def test_train_epoch_weighted_mean(monkeypatch):
trainer, _ = _build_trainer(monkeypatch, n_prompts=6, prompts_per_step=4)
# First chunk (4 prompts) reward=1.0, second chunk (2 prompts) reward=0.0
seq = iter([1.0, 0.0])
def fake_step(chunk):
return {
"loss": 0.0,
"reward": next(seq),
"accuracy": 0.0,
"n_prompts": len(chunk),
}
monkeypatch.setattr(trainer, "_grpo_step", fake_step)
metrics = trainer._train_epoch(epoch=0)
# weighted: (1.0*4 + 0.0*2) / 6 = 0.6667
assert abs(metrics["reward"] - (4.0 / 6.0)) < 1e-9
def test_prompts_cached(monkeypatch):
trainer, _ = _build_trainer(monkeypatch, n_prompts=4, prompts_per_step=4)
n_loads = {"count": 0}
orig = datasets.load_grpo_prompts
def counting(*a, **k):
n_loads["count"] += 1
return orig(*a, **k)
monkeypatch.setattr(datasets, "load_grpo_prompts", counting)
monkeypatch.setattr(
trainer,
"_grpo_step",
lambda chunk: {"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)},
)
trainer._train_epoch(epoch=0)
trainer._train_epoch(epoch=1)
# Loaded once, cached on the trainer thereafter.
assert n_loads["count"] == 1
def test_train_epoch_no_model_still_iterates(monkeypatch):
"""Even with policy.model None (no GPU), epoch iterates via _grpo_step."""
trainer, _ = _build_trainer(monkeypatch, n_prompts=5, prompts_per_step=2)
assert trainer.policy.model is None # CPU/no-torch stub
called = {"n": 0}
def fake_step(chunk):
called["n"] += 1
return {"loss": 0.0, "reward": 0.0, "accuracy": 0.0, "n_prompts": len(chunk)}
monkeypatch.setattr(trainer, "_grpo_step", fake_step)
metrics = trainer._train_epoch(epoch=0)
assert called["n"] == 3 # [2, 2, 1]
assert metrics["n_prompts"] == 5
# ---------------------------------------------------------------------------
# _grpo_step over the unified rollout (offline, fake model + patched paths)
# ---------------------------------------------------------------------------
torch = pytest.importorskip("torch")
class _FakeModel:
"""Stand-in for the HF policy/ref model: trainable enough to backprop.
``parameters()`` returns no params so the NaN-grad scan and
``clip_grad_norm_`` are no-ops; ``train()`` is a no-op. The real
log-prob / generation paths are monkeypatched, so the model is never
actually called.
"""
def parameters(self):
return iter(())
def train(self):
return self
class _FakeOptimizer:
def zero_grad(self):
pass
def step(self):
pass
def _build_step_trainer(monkeypatch, *, group_size=4):
"""Trainer wired for a CPU ``_grpo_step`` run: fake model + optimizer,
device='cpu', and the rollout/logprob paths left for the test to patch.
"""
tasks = _make_tasks(3)
monkeypatch.setattr(datasets, "load_grpo_prompts", lambda *a, **k: tasks)
monkeypatch.setattr(
policy_model_mod.OrchestratorPolicyModel,
"from_pretrained",
classmethod(lambda cls, *a, **k: OrchestratorPolicyModel(model=None)),
)
cfg = OrchestratorGRPOConfig(
grpo_max_prompts=3,
prompts_per_step=3,
num_samples_per_prompt=group_size,
)
trainer = OrchestratorGRPOTrainer(cfg)
# Make it look like a real (CPU) training setup.
trainer.policy.model = _FakeModel()
trainer.ref_policy.model = _FakeModel()
trainer.optimizer = _FakeOptimizer()
trainer.device = torch.device("cpu")
return trainer, tasks
def test_grpo_step_uses_unified_rollout_not_old_env(monkeypatch):
"""_grpo_step samples groups via the policy rollout caller, normalises
advantages within the group, and never touches the old
THOUGHT/TOOL/INPUT environment path."""
# Tripwire: the old environment generation path must NOT be invoked.
import openjarvis.learning.intelligence.orchestrator.policy_model as pm
def _boom(*a, **k): # pragma: no cover - only fires on regression
raise AssertionError(
"GRPO generation hit the old THOUGHT/TOOL/INPUT path "
"(policy_model._parse_output / predict_action)"
)
monkeypatch.setattr(pm.OrchestratorPolicyModel, "_parse_output", _boom)
monkeypatch.setattr(pm.OrchestratorPolicyModel, "predict_action", _boom)
monkeypatch.setattr(pm.OrchestratorPolicyModel, "_build_prompt", _boom)
trainer, tasks = _build_step_trainer(monkeypatch, group_size=4)
# verify_answer: even-indexed rollouts correct, odd incorrect -> spread of
# rewards inside each group so std > 0 and advantages get normalised.
import openjarvis.learning.intelligence.orchestrator.sft_data.verify as verify_mod
seen_answers = []
def fake_verify(task, pred):
seen_answers.append(pred)
return pred.endswith("correct")
monkeypatch.setattr(verify_mod, "verify_answer", fake_verify)
def fake_rollout_group(task):
recs = []
for i in range(trainer.config.num_samples_per_prompt):
ans = f"ans-{i}-correct" if i % 2 == 0 else f"ans-{i}-wrong"
recs.append(
RolloutRecord(
final_answer=ans,
cost_usd=0.01 * (i + 1),
token_ids=[1, 2, 3, 4],
assistant_mask=[0, 1, 0, 1],
)
)
return recs
monkeypatch.setattr(trainer, "_rollout_group", fake_rollout_group)
def fake_logprob(token_ids, assistant_mask, *, ref):
# token_ids/mask flow through unchanged from the records.
assert token_ids == [1, 2, 3, 4]
assert assistant_mask == [0, 1, 0, 1]
# Distinct, grad-bearing values for current vs ref so the ratio is
# finite and backward() builds a real graph.
base = -1.0 if ref else -0.9
return torch.tensor(base, requires_grad=not ref)
monkeypatch.setattr(trainer, "_trajectory_logprob", fake_logprob)
metrics = trainer._grpo_step(tasks)
# Metrics aggregate over the 3 tasks.
assert metrics["n_prompts"] == 3
# Half of each group is correct -> accuracy 0.5.
assert abs(metrics["accuracy"] - 0.5) < 1e-9
assert "loss" in metrics and "reward" in metrics
# The unified-rollout caller produced our canned answers (proves we went
# through the rollout path, not the old env).
assert any(a.endswith("correct") for a in seen_answers)
assert any(a.endswith("wrong") for a in seen_answers)
def test_grpo_step_group_normalised_advantages(monkeypatch):
"""Advantages within a group are mean-0 / std-1 (group-relative)."""
trainer, tasks = _build_step_trainer(monkeypatch, group_size=4)
import openjarvis.learning.intelligence.orchestrator.sft_data.verify as verify_mod
monkeypatch.setattr(
verify_mod, "verify_answer", lambda task, pred: pred == "good"
)
all_groups = []
def fake_rollout_group(task):
# 2 correct, 2 wrong, varied costs -> non-degenerate reward spread.
answers = ["good", "good", "bad", "bad"]
recs = [
RolloutRecord(
final_answer=answers[i],
cost_usd=0.0,
token_ids=[5, 6, 7],
assistant_mask=[0, 1, 1],
)
for i in range(4)
]
all_groups.append(recs)
return recs
monkeypatch.setattr(trainer, "_rollout_group", fake_rollout_group)
def fake_logprob(token_ids, assistant_mask, *, ref):
return torch.tensor(0.0, requires_grad=not ref)
monkeypatch.setattr(trainer, "_trajectory_logprob", fake_logprob)
trainer._grpo_step(tasks)
# 3 tasks, each a group of 4. Read the advantages _grpo_step stamped onto
# each record; every group must be group-normalised (mean ~0, std ~1).
assert len(all_groups) == 3
for recs in all_groups:
group = [r.advantage for r in recs]
assert len(group) == 4
mean = sum(group) / 4
assert abs(mean) < 1e-6
var = sum((x - mean) ** 2 for x in group) / 4
assert abs(var - 1.0) < 1e-6
def test_stitch_trajectory_assistant_mask():
"""Stitching per-turn (prompt_ids, gen_ids) yields a correct mask: 1 only
on the policy-generated tokens, growing prompts diffed across turns."""
trainer = OrchestratorGRPOTrainer.__new__(OrchestratorGRPOTrainer)
# Turn 1: prompt [10,11], gen [20,21]. Turn 2 prompt is the full prior
# transcript [10,11,20,21] + a new observation token [12], gen [22].
turn_buffer = [
([10, 11], [20, 21]),
([10, 11, 20, 21, 12], [22]),
]
token_ids, mask = trainer._stitch_trajectory(turn_buffer)
# Turn1: new prompt [10,11] mask0, gen [20,21] mask1.
# Turn2: new prompt [12] mask0 (the [10,11,20,21] already emitted), gen [22] mask1.
assert token_ids == [10, 11, 20, 21, 12, 22]
assert mask == [0, 0, 1, 1, 0, 1]
def test_trajectory_logprob_masks_to_assistant_positions(monkeypatch):
"""_trajectory_logprob sums log-probs only over assistant-masked tokens,
using logits at position i-1 to predict token i."""
trainer = OrchestratorGRPOTrainer.__new__(OrchestratorGRPOTrainer)
trainer.device = torch.device("cpu")
token_ids = [3, 7, 5, 9]
assistant_mask = [0, 1, 0, 1] # tokens at idx 1 and 3 are assistant
vocab = 10
seq = len(token_ids)
# Deterministic logits so the expected sum is computable.
logits = torch.zeros(1, seq, vocab)
for i in range(seq):
logits[0, i, :] = torch.arange(vocab, dtype=torch.float32) * 0.1 * (i + 1)
class _LogitModel:
def __call__(self, input_ids=None):
class _Out:
pass
o = _Out()
o.logits = logits
return o
def parameters(self):
return iter(())
trainer.policy = OrchestratorPolicyModel(model=_LogitModel())
trainer.ref_policy = OrchestratorPolicyModel(model=_LogitModel())
out = trainer._trajectory_logprob(token_ids, assistant_mask, ref=True)
# Expected: log_softmax(logits[i-1])[token_i] summed over masked i in {1,3}.
expected = 0.0
for i in range(1, seq):
if assistant_mask[i] != 1:
continue
lp = torch.log_softmax(logits[0, i - 1, :], dim=-1)[token_ids[i]]
expected += lp.item()
assert abs(out.item() - expected) < 1e-5
tests/test_orchestrator_learning/test_reward_costaware.py
+86
View File
@@ -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)
tests/test_orchestrator_learning/test_sft_trainer.py
+1 -1
View File
@@ -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