diff --git a/scripts/orchestrator/build_unified_sft.py b/scripts/orchestrator/build_unified_sft.py
new file mode 100644
index 00000000..e4eea00b
--- /dev/null
+++ b/scripts/orchestrator/build_unified_sft.py
@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+"""Build ToolOrchestra SFT data by rejection sampling over ToolScale.
+
+For each ToolScale task, a teacher orchestrator is rolled out N times over the
+faithful unified tool catalog (one tool per model); passing trajectories are
+serialized into the conversations JSONL the SFT trainer consumes.
+
+Needs API access for the teacher + expert models (set the usual env keys).
+
+Example:
+ uv run python scripts/orchestrator/build_unified_sft.py \
+ --out data/orchestrator_unified_sft.jsonl \
+ --teacher-model gpt-5 --max-tasks 200 --samples-per-task 4 \
+ --local-model qwen3:8b --local-endpoint http://localhost:8001/v1
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from typing import Optional
+
+from openjarvis.agents.hybrid.expert_registry import default_catalog
+from openjarvis.agents.hybrid.toolorchestra.unified import (
+ make_call_orchestrator,
+ make_dispatch,
+)
+from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout
+from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
+ generate_sft_dataset,
+ gold_coverage_verify,
+)
+from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
+ load_toolscale,
+)
+
+
+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("--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)
+ p.add_argument("--max-turns", type=int, default=50)
+ p.add_argument("--temperature", type=float, default=1.0)
+ p.add_argument("--local-model", default=None)
+ p.add_argument("--local-endpoint", default=None)
+ args = p.parse_args(argv)
+
+ logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+ tools = default_catalog(
+ local_model=args.local_model, local_endpoint=args.local_endpoint,
+ )
+ logging.info("Tool catalog (%d): %s", len(tools), [t.name for t in tools])
+
+ call_orch = make_call_orchestrator(
+ args.teacher_model,
+ base_url=args.teacher_base_url,
+ api_key=os.environ.get("OPENAI_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_toolscale(max_tasks=args.max_tasks)
+ stats = generate_sft_dataset(
+ args.out,
+ tasks=tasks,
+ tools=tools,
+ rollout_fn=rollout_fn,
+ verify_fn=gold_coverage_verify,
+ 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())
diff --git a/src/openjarvis/agents/hybrid/expert_registry.py b/src/openjarvis/agents/hybrid/expert_registry.py
new file mode 100644
index 00000000..703b2193
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/expert_registry.py
@@ -0,0 +1,311 @@
+"""Faithful ToolOrchestra "unified tool calling" registry (arXiv:2511.21689 §3.1).
+
+The paper exposes **every tool AND every model through a single flat tool
+interface** — each is its own named function with a description and a typed
+parameter schema, and for each training instance a *random subset* of tools is
+sampled with *randomized pricing* (§3.3, "General tool configuration"). This is
+unlike the eval-port shortcut in ``toolorchestra.py``, which collapses the whole
+catalog into three meta-tools (``search``/``enhance_reasoning``/``answer``) with
+a ``model`` slot. This module restores the faithful design.
+
+Each :class:`ExpertTool` knows:
+
+* the orchestrator-visible ``name`` / ``description`` / param schema (what goes
+ into the tools JSON the policy conditions on), and
+* the concrete backend (``backend_type`` + ``model`` + ``base_url``) so a caller
+ can turn it into the worker dict that ``toolorchestra._call_worker`` dispatches.
+
+Everything here is pure data + deterministic transforms (no network, no model
+calls), so the spec building, sampling, and pricing logic is offline-testable.
+Dispatch stays in ``toolorchestra.py`` (via :func:`to_worker_dict`) to avoid a
+circular import.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+from openjarvis.agents.hybrid._prices import PRICES
+
+# Kinds of tool in the unified interface.
+KIND_MODEL = "model" # an LLM exposed as a tool (the paper's "models as tools")
+KIND_WEB_SEARCH = "web_search"
+KIND_LOCAL_SEARCH = "local_search"
+KIND_CODE = "code_interpreter"
+
+VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_LOCAL_SEARCH, KIND_CODE)
+
+# Backend dispatch types understood by ``toolorchestra._call_worker``.
+VALID_BACKENDS = (
+ "vllm", "openai", "anthropic", "gemini", "openrouter",
+ "anthropic-web-search", "tavily-search", "modal-python",
+)
+
+
+@dataclass(frozen=True)
+class ExpertTool:
+ """One entry in the unified tool catalog.
+
+ ``price_in`` / ``price_out`` are USD per 1M tokens (0.0 for local / non-LLM
+ tools). ``latency_s`` is a rough average used only to populate the
+ description's cost/latency line — the orchestrator was trained to read that
+ table, so we surface it verbatim in the spec.
+ """
+
+ name: str
+ kind: str
+ backend_type: str
+ summary: str
+ model: Optional[str] = None
+ base_url: Optional[str] = None
+ price_in: float = 0.0
+ price_out: float = 0.0
+ latency_s: float = 5.0
+
+ def __post_init__(self) -> None:
+ if self.kind not in VALID_KINDS:
+ raise ValueError(f"{self.name}: invalid kind {self.kind!r}")
+ if self.backend_type not in VALID_BACKENDS:
+ raise ValueError(f"{self.name}: invalid backend {self.backend_type!r}")
+ if self.kind == KIND_MODEL and not self.model:
+ raise ValueError(f"{self.name}: model-kind tool needs a concrete model")
+
+ # ---- orchestrator-visible spec -------------------------------------
+
+ def _param_schema(self) -> Dict[str, object]:
+ """JSON-schema for the tool's arguments (one typed param per kind)."""
+ if self.kind == KIND_WEB_SEARCH or self.kind == KIND_LOCAL_SEARCH:
+ return {
+ "type": "object",
+ "properties": {
+ "query": {
+ "type": "string",
+ "description": "Search query string.",
+ }
+ },
+ "required": ["query"],
+ }
+ if self.kind == KIND_CODE:
+ return {
+ "type": "object",
+ "properties": {
+ "code": {
+ "type": "string",
+ "description": "Python code to execute. Print results.",
+ }
+ },
+ "required": ["code"],
+ }
+ # model tool
+ return {
+ "type": "object",
+ "properties": {
+ "input": {
+ "type": "string",
+ "description": "The sub-question or instruction for this model.",
+ }
+ },
+ "required": ["input"],
+ }
+
+ def description(self) -> str:
+ """Full description incl. the price/latency line (paper bakes this in)."""
+ if self.kind == KIND_MODEL:
+ cost_line = (
+ f" Pricing: ${self.price_in:.2f}/1M input, "
+ f"${self.price_out:.2f}/1M output; avg latency ~{self.latency_s:.0f}s."
+ )
+ else:
+ cost_line = f" Avg latency ~{self.latency_s:.0f}s."
+ return self.summary.rstrip(".") + "." + cost_line
+
+ def to_spec(self) -> Dict[str, object]:
+ """OpenAI-style tool spec the orchestrator conditions on."""
+ return {
+ "type": "function",
+ "function": {
+ "name": self.name,
+ "description": self.description(),
+ "parameters": self._param_schema(),
+ },
+ }
+
+
+def _price(model: str) -> tuple[float, float]:
+ return PRICES.get(model, (0.0, 0.0))
+
+
+# Default catalog: the paper's tool categories, mapped onto the models/tools
+# OpenJarvis can actually call. One named tool per model (faithful §3.1).
+def default_catalog(
+ *,
+ local_model: Optional[str] = None,
+ local_endpoint: Optional[str] = None,
+) -> List[ExpertTool]:
+ """Return the full unified tool catalog.
+
+ ``local_model`` / ``local_endpoint`` wire the on-device vLLM tool when a
+ local backbone is served; omitted → the local model tool is left out.
+ """
+ cat: List[ExpertTool] = []
+
+ # ---- generalist / frontier models ----
+ for name, model, summary, lat in [
+ ("gpt_5", "gpt-5",
+ "Frontier generalist (GPT-5). Strongest reasoning across domains.", 30.0),
+ ("gpt_5_mini", "gpt-5-mini",
+ "Mid-tier generalist (GPT-5-mini). Solid reasoning, much cheaper.", 15.0),
+ ("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",
+ "Frontier generalist (Gemini 2.5 Pro). Strong multimodal reasoning.", 20.0),
+ ("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",
+ "Open generalist (Llama-3.3-70B). Decent general knowledge, low cost.", 10.0),
+ ("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"))
+ pi, po = _price(model)
+ cat.append(ExpertTool(
+ name=name, kind=KIND_MODEL, backend_type=ep, summary=summary,
+ model=model, price_in=pi, price_out=po, latency_s=lat,
+ ))
+
+ # ---- specialized: code ----
+ pi, po = _price("qwen/qwen-2.5-coder-32b-instruct")
+ cat.append(ExpertTool(
+ name="qwen2_5_coder_32b", 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,
+ ))
+
+ # ---- local backbone as a tool (on-device vLLM), if served ----
+ 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."),
+ model=local_model, base_url=local_endpoint,
+ price_in=0.0, price_out=0.0, latency_s=2.0,
+ ))
+
+ # ---- 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,
+ ))
+ 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,
+ ))
+
+ return cat
+
+
+def build_tool_specs(tools: List[ExpertTool]) -> List[Dict[str, object]]:
+ """Turn a tool list into the OpenAI-style tools JSON the policy sees."""
+ return [t.to_spec() for t in tools]
+
+
+def tools_by_name(tools: List[ExpertTool]) -> Dict[str, ExpertTool]:
+ return {t.name: t for t in tools}
+
+
+def sample_tool_config(
+ catalog: List[ExpertTool],
+ *,
+ rng: random.Random,
+ min_tools: int = 4,
+ max_tools: Optional[int] = None,
+ price_jitter: float = 0.0,
+) -> List[ExpertTool]:
+ """Sample a random tool subset with optional price randomization (§3.3).
+
+ Guarantees at least one ``model`` tool and at least one non-model (basic)
+ tool so every instance can both reason and act. ``price_jitter`` (e.g. 0.5)
+ multiplies each model's prices by a per-tool factor drawn uniformly from
+ ``[1-jitter, 1+jitter]``, modeling heterogeneous pricing across users.
+ Deterministic given ``rng``.
+ """
+ if not catalog:
+ raise ValueError("empty catalog")
+ models = [t for t in catalog if t.kind == KIND_MODEL]
+ basics = [t for t in catalog if t.kind != KIND_MODEL]
+ if not models:
+ raise ValueError("catalog has no model tools")
+
+ hi = max_tools if max_tools is not None else len(catalog)
+ hi = min(hi, len(catalog))
+ lo = min(max(min_tools, 2), hi)
+ k = rng.randint(lo, hi)
+
+ # Always include >=1 model; include >=1 basic if any exist.
+ chosen: List[ExpertTool] = [rng.choice(models)]
+ if basics:
+ chosen.append(rng.choice(basics))
+ pool = [t for t in catalog if t not in chosen]
+ rng.shuffle(pool)
+ for t in pool:
+ if len(chosen) >= k:
+ break
+ chosen.append(t)
+
+ # Re-order to catalog order for stable specs.
+ order = {t.name: i for i, t in enumerate(catalog)}
+ chosen.sort(key=lambda t: order[t.name])
+
+ if price_jitter > 0.0:
+ jittered: List[ExpertTool] = []
+ for t in chosen:
+ if t.kind == KIND_MODEL and (t.price_in or t.price_out):
+ f = rng.uniform(1.0 - price_jitter, 1.0 + price_jitter)
+ jittered.append(ExpertTool(
+ name=t.name, kind=t.kind, backend_type=t.backend_type,
+ summary=t.summary, model=t.model, base_url=t.base_url,
+ price_in=round(t.price_in * f, 4),
+ price_out=round(t.price_out * f, 4),
+ latency_s=t.latency_s,
+ ))
+ else:
+ jittered.append(t)
+ return jittered
+ return chosen
+
+
+def to_worker_dict(tool: ExpertTool) -> Dict[str, object]:
+ """Convert a tool into the worker dict ``toolorchestra._call_worker`` eats."""
+ d: Dict[str, object] = {
+ "name": tool.name,
+ "type": tool.backend_type,
+ "model": tool.model,
+ }
+ if tool.base_url:
+ d["base_url"] = tool.base_url
+ return d
+
+
+__all__ = [
+ "ExpertTool",
+ "KIND_CODE",
+ "KIND_LOCAL_SEARCH",
+ "KIND_MODEL",
+ "KIND_WEB_SEARCH",
+ "build_tool_specs",
+ "default_catalog",
+ "sample_tool_config",
+ "to_worker_dict",
+ "tools_by_name",
+]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py
deleted file mode 100644
index d5ebd78d..00000000
--- a/src/openjarvis/agents/hybrid/toolorchestra.py
+++ /dev/null
@@ -1,1662 +0,0 @@
-"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
-
-Two modes, gated by ``method_cfg.orchestrator_mode``:
-
-* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
- orchestrator, dispatching to a numbered worker pool via JSON
- ``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
- a prompted upper-bound reference point — NOT the paper's setup.
-
-* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
- served on a local vLLM is the orchestrator. It emits OpenAI-style
- ``tool_calls`` (or ``{...}`` text blocks when
- vLLM's tool parser doesn't catch them) for three expert tools —
- ``enhance_reasoning``, ``answer``, ``search`` — exactly as in the
- upstream ``evaluation/tools.json``. Each tool's ``model`` arg
- (``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real
- backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier
- Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
- for `*-3`. Search routes to the Anthropic server-side web_search.
-
- We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
- code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
- Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
- worker types. Energy-wise, "expert" answers are cloud calls.
-
-Pipeline per task (RL mode):
-
-1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
- appropriate tool.` with the three tools declared.
-2. It emits one ``tool_call`` per turn — ``search`` updates the
- context, ``enhance_reasoning`` appends code/exec output (we run the
- tool as a plain LLM call, no sandbox — the model just gets prose
- back), ``answer`` produces the final answer and the loop stops.
-3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
- back to the strongest expert worker.
-
-Prompted-mode pipeline:
-
-1. Orchestrator (cloud) reads question + numbered worker pool.
-2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
- "input": str}`` or ``{"action": "final_answer", "answer": str}``.
-3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
- prompt; fallback to strongest worker on parse failure.
-
-Workers come from ``cfg["workers"]`` or a sensible default pool (local
-Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
-gpt-5-mini).
-"""
-
-from __future__ import annotations
-
-import json
-import re
-import shutil
-import tempfile
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-from openjarvis.agents._stubs import AgentContext
-from openjarvis.agents.hybrid._base import (
- ANTHROPIC_WEB_SEARCH_TOOL,
- WEB_SEARCH_COST_PER_CALL,
- LocalCloudAgent,
-)
-from openjarvis.agents.hybrid._prices import (
- PRICES,
- is_gpt5_family,
- supports_temperature,
-)
-from openjarvis.agents.hybrid.mini_swe_agent import (
- _clone_repo,
- _extract_diff,
- run_swe_agent_loop,
-)
-from openjarvis.core.registry import AgentRegistry
-
-ORCHESTRATOR_SYS = """\
-You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
-
- {"action": "call_worker", "worker_id": , "input": ""}
-
- {"action": "final_answer", "answer": ""}
-
-Strategy:
-
-- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math).
-- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass.
-- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase.
-- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question.
-"""
-
-FORCE_FINAL_PROMPT = (
- "Worker-call budget exhausted. Emit `final_answer` now using everything "
- "you've learned. Respect the question's answer-format rules."
-)
-
-
-# ============================================================================
-# RL-mode constants (Orchestrator-8B, paper-faithful).
-# ============================================================================
-#
-# Verbatim copies of the upstream system prompt / user-prompt template / tools
-# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
-# edit the description text — Orchestrator-8B was RL-trained against this
-# exact wording and pricing/latency table.
-
-RL_ORCHESTRATOR_SYS = "You are good at using tools."
-
-RL_TOOLS_SPEC: List[Dict[str, Any]] = [
- {
- "type": "function",
- "function": {
- "name": "enhance_reasoning",
- "description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
- "parameters": {
- "properties": {
- "model": {
- "description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
- "type": "string",
- }
- },
- "required": ["model"],
- "title": "parameters",
- "type": "object",
- },
- },
- },
- {
- "type": "function",
- "function": {
- "name": "answer",
- "description": "give the final answer. Not allowed to call if documents is empty.",
- "parameters": {
- "properties": {
- "model": {
- "description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
- "type": "string",
- }
- },
- "required": ["model"],
- "title": "parameters",
- "type": "object",
- },
- },
- },
- {
- "type": "function",
- "function": {
- "name": "search",
- "description": "Search for missing information",
- "parameters": {
- "properties": {
- "model": {
- "description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
- "type": "string",
- }
- },
- "required": ["model"],
- "title": "parameters",
- "type": "object",
- },
- },
- },
-]
-
-# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
-RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
- "enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
- "answer": {
- "model": [
- "answer-1", "answer-2", "answer-3", "answer-4",
- "answer-math-1", "answer-math-2",
- ],
- },
- "search": {"model": ["search-1", "search-2", "search-3"]},
-}
-
-# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
-# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
-# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
-# tiers since we don't have Qwen-Math served.
-#
-# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
-# so the substitution is deferred until we know the cell's resolved local/cloud
-# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
-
-def _expert_for(slot: str, local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str,
- cloud_endpoint: str = "anthropic") -> Dict[str, Any]:
- """Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
-
- Routing policy:
- - `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off
- `cloud_endpoint` ("anthropic"/"openai"/"gemini")
- - `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's
- cost tier for mid OpenAI calls)
- - `*-3` (local tier) -> local vLLM (`local_model`)
- - `answer-math-*` -> same tiers as the numeric suffix
- - `search-*` -> always the Anthropic web_search tool (the
- upstream uses Tavily; we have web_search)
- """
- if slot.startswith("search"):
- return {
- "name": f"search:{slot}",
- "type": "anthropic-web-search",
- "model": _DEFAULT_WEB_SEARCH_MODEL,
- }
- if slot.endswith("-1") or slot.endswith("-math-1"):
- ep = (cloud_endpoint or "anthropic").lower()
- if ep not in ("anthropic", "openai", "gemini"):
- ep = "anthropic"
- return {
- "name": f"frontier:{slot}",
- "type": ep,
- "model": cloud_model,
- }
- if slot.endswith("-2") or slot.endswith("-math-2"):
- return {
- "name": f"mid:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
- # `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
- # we substitute whatever local model the cell wired up).
- if local_model and local_endpoint:
- return {
- "name": f"local:{slot}",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- }
- # Fallback if no local — gpt-5-mini.
- return {
- "name": f"mid-fallback:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
-
-
-# ============================================================================
-# Paper-match expert mapping (2026-05-19).
-# ============================================================================
-# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
-# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
-# models and routes `search` through Tavily, while `enhance_reasoning` is
-# expected to produce code that the caller pipes through a Modal sandbox
-# (handled at dispatch time, not here).
-#
-# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
-# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
-#
-# reasoner-1 -> GPT-5 (frontier reasoner)
-# reasoner-2 -> GPT-5-mini (mid)
-# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
-# answer-1 -> GPT-5
-# answer-2 -> GPT-5-mini
-# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
-# answer-4 -> local Qwen
-# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
-# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
-# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
-# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
-# search-* -> Tavily search (paper)
-#
-# `enhance_reasoning` is dispatched through the coder specialist regardless of
-# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
-# routes the same way in all three cases, then optionally extracts a python
-# code block and execs it in Modal. (We keep the slot-aware routing inside the
-# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
-# pins the coder regardless. See `_run_rl_paper` dispatch.)
-
-_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
-_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
-
-
-def _paper_expert_for(
- slot: str,
- local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str,
- cloud_endpoint: str = "openai",
-) -> Dict[str, Any]:
- """Paper-match counterpart of ``_expert_for``.
-
- Differs from ``_expert_for``:
- - Search slots go to ``tavily-search`` (not Anthropic web_search).
- - Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
- OpenRouter rather than collapsing onto the local vLLM.
- - Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
- as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
- - ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
- """
- if slot.startswith("search"):
- return {
- "name": f"tavily:{slot}",
- "type": "tavily-search",
- "model": "tavily",
- }
- if slot in ("answer-math-1", "answer-math-2"):
- return {
- "name": f"math-coder:{slot}",
- "type": "openrouter",
- "model": _PAPER_CODER_OPENROUTER,
- }
- if slot == "answer-3":
- return {
- "name": f"generalist-llama:{slot}",
- "type": "openrouter",
- "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
- }
- if slot.endswith("-1"):
- # Tier-1 frontier reasoner / answer — paper uses GPT-5.
- return {
- "name": f"frontier:{slot}",
- "type": "openai",
- "model": "gpt-5",
- }
- if slot.endswith("-2"):
- return {
- "name": f"mid:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
- # `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
- # also serves the local Qwen for the rare local-tier slot).
- if local_model and local_endpoint:
- return {
- "name": f"local:{slot}",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- }
- return {
- "name": f"mid-fallback:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
-
-
-# ---- Tavily + Modal helpers -------------------------------------------------
-
-def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]:
- """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0).
-
- Token counts are reported as zero (no LLM was billed); the OpenJarvis
- accounting layer separately tallies tool-call counts. Falls back to
- DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``).
- """
- from openjarvis.tools.web_search import WebSearchTool
-
- tool = WebSearchTool(max_results=max_results)
- res = tool.execute(query=query, max_results=max_results)
- text = res.content or ""
- if not res.success and not text:
- text = "(no results)"
- return text, 0, 0
-
-
-_MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox"
-
-
-def _call_modal_python(code: str, timeout_s: int = 60) -> Tuple[str, int]:
- """Execute a single Python snippet in a fresh Modal Sandbox.
-
- Returns ``(combined_stdout_stderr, returncode)``. Logs are capped at 8 KiB.
- Any exception (modal auth, network, sandbox boot failure) is captured into
- the returned string with a non-zero rc — we never raise back to the
- orchestrator loop. The sandbox is torn down at the end via ``terminate()``.
- """
- try:
- import modal
-
- app = modal.App.lookup(_MODAL_APP_NAME, create_if_missing=True)
- # python:3.12-slim is small + boots fast; the paper uses a generic
- # Python image too. We rely on stdlib only — no extra pip installs.
- image = modal.Image.debian_slim(python_version="3.12")
- sb = modal.Sandbox.create(
- "python", "-c", code,
- app=app,
- image=image,
- timeout=int(timeout_s),
- )
- sb.wait()
- try:
- out = sb.stdout.read() or ""
- except Exception:
- out = ""
- try:
- err = sb.stderr.read() or ""
- except Exception:
- err = ""
- rc = sb.returncode if sb.returncode is not None else -1
- try:
- sb.terminate()
- except Exception:
- pass
- combined = out + (("\n" + err) if err else "")
- if len(combined) > 8192:
- combined = combined[:8192] + "\n... (output truncated)"
- return combined, int(rc)
- except Exception as exc:
- return f"[modal-python error: {type(exc).__name__}: {exc}]", -1
-
-
-_PY_CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL)
-
-
-def _extract_first_python_block(text: str) -> Optional[str]:
- """Return the first ```python ... ``` block (or ```...```), or None."""
- m = _PY_CODE_RE.search(text or "")
- return m.group(1).strip() if m else None
-
-
-def _call_orchestrator_with_tool_calls(
- model: str,
- endpoint: str,
- *,
- user: str,
- system: str,
- max_tokens: int,
- temperature: float,
- tools: List[Dict[str, Any]],
- timeout: float = 600.0,
-) -> Tuple[str, int, int, Any]:
- """Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls).
-
- Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level
- ``tool_calls`` object so the RL-mode parser can match against it
- directly. Otherwise vLLM's tool parser silently swallows the tool call
- into the SDK field while leaving ``content == ''`` — and the text-tag
- parser sees nothing, falling through to the answer-1 fallback. (Bug
- observed 2026-05-19 on the paper-match smoke; same path was buggy on
- the default pool too, just less reproducibly.)
- """
- from openai import OpenAI
-
- client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
- messages = [
- {"role": "system", "content": system},
- {"role": "user", "content": user},
- ]
- resp = client.chat.completions.create(
- model=model,
- messages=messages,
- temperature=temperature,
- max_tokens=max_tokens,
- tools=tools,
- extra_body={"chat_template_kwargs": {"enable_thinking": False}},
- )
- choice = resp.choices[0]
- message = choice.message
- text = message.content or ""
- tool_calls = getattr(message, "tool_calls", None)
- u = resp.usage
- p = getattr(u, "prompt_tokens", 0) if u else 0
- c = getattr(u, "completion_tokens", 0) if u else 0
- return text, p, c, tool_calls
-
-
-def _paper_pool(
- local_model: Optional[str],
- local_endpoint: Optional[str],
-) -> List[Dict[str, Any]]:
- """Paper-match worker pool (registered for traces / inspection).
-
- NOTE: in RL mode the orchestrator dispatches via tool/slot rather than
- worker_id, so this list is purely informational — `_paper_expert_for`
- is the actual routing function. We still return a list here so the
- paradigm's trace metadata has something concrete to log.
- """
- pool: List[Dict[str, Any]] = []
- if local_model and local_endpoint:
- pool.append({
- "id": len(pool),
- "name": "local-qwen",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- "description": "Local Qwen vLLM (paper uses Qwen3-32B).",
- })
- pool.append({
- "id": len(pool), "name": "tavily-search",
- "type": "tavily-search", "model": "tavily",
- "description": "Tavily web search.",
- })
- pool.append({
- "id": len(pool), "name": "modal-python",
- "type": "modal-python", "model": "modal-python",
- "description": "Modal Sandbox for one-shot Python exec.",
- })
- pool.append({
- "id": len(pool), "name": "code-specialist",
- "type": "openrouter", "model": _PAPER_CODER_OPENROUTER,
- "description": "Qwen-2.5-Coder-32B via OpenRouter (paper).",
- })
- pool.append({
- "id": len(pool), "name": "generalist-llama",
- "type": "openrouter", "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
- "description": "Llama-3.3-70B-Instruct via OpenRouter (paper tier-3).",
- })
- pool.append({
- "id": len(pool), "name": "generalist-gpt5",
- "type": "openai", "model": "gpt-5",
- "description": "GPT-5 frontier generalist.",
- })
- pool.append({
- "id": len(pool), "name": "generalist-gpt5-mini",
- "type": "openai", "model": "gpt-5-mini",
- "description": "GPT-5-mini mid generalist.",
- })
- return pool
-
-
-# Regex for ``{...}`` blocks emitted by Orchestrator-8B
-# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a
-# hermes-style template). Captures the JSON payload.
-_TOOL_CALL_TAG_RE = re.compile(
- r"\s*(\{.*?\})\s*", re.DOTALL
-)
-
-
-def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
- """Return ``{"name": str, "arguments": dict}`` or None.
-
- Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
- back to scraping ``{...}`` tags from the raw
- content. We take the first tool call only — Orchestrator-8B was trained
- to emit exactly one per turn.
- """
- # SDK-level path.
- if sdk_tool_calls:
- first = sdk_tool_calls[0]
- name = getattr(getattr(first, "function", None), "name", None)
- args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
- try:
- args = json.loads(args_raw)
- except json.JSONDecodeError:
- args = {}
- if isinstance(name, str) and isinstance(args, dict):
- return {"name": name, "arguments": args}
- # Text-tag fallback.
- if not isinstance(content, str):
- return None
- m = _TOOL_CALL_TAG_RE.search(content)
- if not m:
- return None
- try:
- obj = json.loads(m.group(1))
- except json.JSONDecodeError:
- return None
- if not isinstance(obj, dict):
- return None
- name = obj.get("name")
- args = obj.get("arguments", {})
- if not isinstance(name, str) or not isinstance(args, dict):
- return None
- return {"name": name, "arguments": args}
-
-
-def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
- return "\n".join(
- f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
- )
-
-
-def _build_user_prompt(
- question: str,
- workers: List[Dict[str, Any]],
- history: List[Dict[str, Any]],
-) -> str:
- pieces = [
- f"Worker pool:\n{_build_pool_block(workers)}",
- f"User question:\n{question}",
- ]
- if history:
- pieces.append("Conversation so far (orchestrator turns and worker outputs):")
- for h in history:
- if h["role"] == "orchestrator":
- pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
- else:
- pieces.append(
- f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
- f"{h['output']}"
- )
- pieces.append(
- "Emit the next JSON action object now — exactly one object, no prose."
- )
- return "\n\n".join(pieces)
-
-
-def _strip_fences(s: str) -> str:
- s = s.strip()
- if s.startswith("```"):
- first_nl = s.find("\n")
- if first_nl != -1:
- s = s[first_nl + 1:]
- if s.endswith("```"):
- s = s[:-3]
- s = s.strip()
- return s
-
-
-def _parse_action(text: str) -> Optional[Dict[str, Any]]:
- s = _strip_fences(text)
- # First try direct parse, then balanced-brace extraction.
- try:
- obj = json.loads(s)
- if isinstance(obj, dict) and "action" in obj:
- return obj
- except json.JSONDecodeError:
- pass
- start = s.find("{")
- if start == -1:
- return None
- depth = 0
- for i in range(start, len(s)):
- c = s[i]
- if c == "{":
- depth += 1
- elif c == "}":
- depth -= 1
- if depth == 0:
- try:
- obj = json.loads(s[start : i + 1])
- if isinstance(obj, dict) and "action" in obj:
- return obj
- except json.JSONDecodeError:
- return None
- return None
-
-
-def _extract_final_answer_text(text: str) -> str:
- """Best-effort: pull the answer string from a malformed action emission.
-
- Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
- """
- m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
- if m:
- return m.group(1).encode("utf-8").decode("unicode_escape")
- m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
- if m:
- return m.group(1).strip()
- return text.strip()
-
-
-# ---------- Worker pool ----------
-
-def _default_pool(
- local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str = "claude-opus-4-7",
- cloud_endpoint: str = "anthropic",
-) -> List[Dict[str, Any]]:
- """Default heterogeneous worker pool.
-
- The frontier worker's ``type`` + ``model`` track the cell's resolved
- ``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5,
- gemini-2.5-pro, …) route their frontier slot to the right SDK.
- """
- ep = (cloud_endpoint or "anthropic").lower()
- if ep not in ("anthropic", "openai", "gemini"):
- ep = "anthropic"
- pool: List[Dict[str, Any]] = []
- if local_model and local_endpoint:
- pool.append({
- "id": len(pool),
- "name": "local-qwen",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- "description": (
- "Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
- "concise extraction, formatting, arithmetic on given data."
- ),
- })
- pool.append({
- "id": len(pool),
- "name": "web-search",
- "type": "anthropic-web-search",
- "model": "claude-haiku-4-5",
- "description": (
- "Anthropic server-side web_search. Use for facts that need a lookup "
- "(recent events, rare names/dates, niche sources). Returns a digest."
- ),
- })
- pool.append({
- "id": len(pool),
- "name": f"frontier-{ep}",
- "type": ep,
- "model": cloud_model,
- "description": (
- "Frontier reasoning model. Use for hard multi-step reasoning, "
- "code review, or a final synthesis pass. Expensive — use sparingly."
- ),
- })
- pool.append({
- "id": len(pool),
- "name": "frontier-openai-mini",
- "type": "openai",
- "model": "gpt-5-mini",
- "description": (
- "Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
- "fraction of frontier cost."
- ),
- })
- return pool
-
-
-# Worker types toolorchestra's `_call_worker` actually dispatches.
-#
-# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
-# `tavily-search` — Tavily API search (the paper's web tool).
-# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
-# Used for the code/math specialists and Llama-3.3-70B /
-# Qwen3-32B generalists.
-# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
-# paper's "Python sandbox" inside `enhance_reasoning`).
-_TOOLORCH_VALID_TYPES = (
- "vllm", "openai", "anthropic", "anthropic-web-search", "gemini",
- "tavily-search", "openrouter", "modal-python",
-)
-
-# Default model used when an `anthropic-web-search` entry omits `model`.
-_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
-
-
-def _resolve_worker_pool(
- cfg: Dict[str, Any],
- local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str,
- cloud_endpoint: str = "anthropic",
-) -> List[Dict[str, Any]]:
- """Return the worker pool for this run.
-
- Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
- default pool is ignored entirely. Falls back to ``_default_pool`` when
- the override is absent.
-
- Each user-supplied entry must be a dict with keys ``id``, ``name``,
- ``type``, and (for non-search types) ``model``. ``type`` must be one
- of ``vllm`` / ``openai`` / ``anthropic`` / ``anthropic-web-search``.
- ``anthropic-web-search`` entries may omit ``model`` — it defaults to
- ``claude-haiku-4-5``.
-
- Substitution: ``model = "$local"`` (or ``""``) resolves to
- ``local_model``; ``model = "$cloud"`` / ``""`` to ``cloud_model``.
-
- On any validation failure, raises ``ValueError`` with the message
- ``"Invalid worker_pool entry []: "``. Fails fast at agent
- init rather than mid-task.
- """
- override = cfg.get("worker_pool")
- if override is None:
- return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint)
- if not isinstance(override, list) or not override:
- raise ValueError(
- "Invalid worker_pool entry [-]: worker_pool must be a non-empty list"
- )
-
- resolved: List[Dict[str, Any]] = []
- seen_ids: set = set()
- has_non_search = False
- for raw in override:
- wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?"
- if not isinstance(raw, dict):
- raise ValueError(
- f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict"
- )
- entry = dict(raw)
- wid = entry.get("id")
- if not isinstance(wid, int):
- raise ValueError(
- f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int"
- )
- if wid in seen_ids:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: duplicate id"
- )
- seen_ids.add(wid)
- if not entry.get("name") or not isinstance(entry["name"], str):
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string"
- )
- wtype = entry.get("type") or entry.get("endpoint")
- if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'type' must be one of "
- f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})"
- )
- wtype = wtype.lower()
- entry["type"] = wtype
- # Substitute $local / $cloud placeholders (before any model check).
- model = entry.get("model")
- if isinstance(model, str) and model in ("$local", ""):
- if not local_model:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: model='{model}' "
- "requires a local_model to be configured for this cell"
- )
- model = local_model
- entry["model"] = model
- elif isinstance(model, str) and model in ("$cloud", ""):
- model = cloud_model
- entry["model"] = model
- if wtype == "anthropic-web-search":
- if model in (None, ""):
- model = _DEFAULT_WEB_SEARCH_MODEL
- entry["model"] = model
- elif not isinstance(model, str):
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
- )
- # Search workers don't satisfy the "needs a solver" requirement.
- else:
- if not isinstance(model, str) or not model:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
- )
- if wtype == "vllm":
- if not entry.get("base_url"):
- if not local_endpoint:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: vllm worker needs "
- "'base_url' (or a configured local_endpoint to fall back to)"
- )
- entry["base_url"] = local_endpoint
- entry.setdefault("api_key", "EMPTY")
- else:
- if model not in PRICES:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: model {model!r} "
- f"is not in PRICES (known: {sorted(PRICES)})"
- )
- has_non_search = True
- entry.setdefault(
- "description",
- f"User-supplied {wtype} worker ({model}).",
- )
- resolved.append(entry)
-
- if not has_non_search:
- raise ValueError(
- "Invalid worker_pool entry [-]: worker_pool must contain at least "
- "one non-search worker (vllm / openai / anthropic)"
- )
- return resolved
-
-
-def _call_worker(
- worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
-) -> Tuple[str, int, int, bool, float, int]:
- """Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
- wtype = worker.get("type", "openai")
- max_tok = int(cfg.get("worker_max_tokens", 4096))
- temp = float(cfg.get("worker_temperature", 0.2))
-
- if wtype == "vllm":
- text, p, c = LocalCloudAgent._call_vllm(
- worker["model"],
- worker["base_url"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- enable_thinking=False,
- )
- return text, p, c, True, 0.0, 0
- if wtype == "openai":
- is_gpt5 = is_gpt5_family(worker["model"])
- eff_temp = 1.0 if is_gpt5 else temp
- # GPT-5 is a reasoning model: hidden reasoning tokens count against
- # `max_completion_tokens`, so a 4096 cap can be fully consumed by
- # reasoning and leave 0 visible content (empty answer). Give the
- # reasoning headroom on top of the answer budget.
- eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
- text, p, c = LocalCloudAgent._call_openai(
- worker["model"],
- user=prompt,
- max_tokens=eff_max_tok,
- temperature=eff_temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "gemini":
- text, p, c = LocalCloudAgent._call_gemini(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "anthropic":
- eff_temp = temp if supports_temperature(worker["model"]) else 0.0
- text, p, c, _ = LocalCloudAgent._call_anthropic(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=eff_temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "anthropic-web-search":
- eff_temp = temp if supports_temperature(worker["model"]) else 0.0
- text, p, c, n_searches = LocalCloudAgent._call_anthropic(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=eff_temp,
- tools=[ANTHROPIC_WEB_SEARCH_TOOL],
- tool_choice={"type": "any"},
- )
- extra = n_searches * WEB_SEARCH_COST_PER_CALL
- return text, p, c, False, extra, n_searches
- if wtype == "tavily-search":
- # Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL`
- # for parity with the Anthropic web-search worker. One call = one
- # "n_search" for accounting.
- max_results = int(cfg.get("tavily_max_results", 5))
- text, p, c = _call_tavily_search(str(prompt), max_results=max_results)
- return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1
- if wtype == "openrouter":
- text, p, c = LocalCloudAgent._call_openrouter(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "modal-python":
- # `prompt` is the python code string to exec.
- timeout_s = int(cfg.get("modal_python_timeout_s", 60))
- out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
- # No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
- # charges per sandbox-second — not tracked here.
- return out, 0, 0, False, 0.0, 0
- raise ValueError(f"unsupported worker type: {wtype!r}")
-
-
-def _swe_call_worker(
- worker: Dict[str, Any],
- prompt: str,
- cfg: Dict[str, Any],
- task: Dict[str, Any],
- workdir: Path,
- turn: int,
-) -> Tuple[str, int, int, bool, float, int, int]:
- """SWE-bench worker dispatch: route solver workers through
- run_swe_agent_loop on a shared workdir. Web-search workers fall back
- to the regular one-shot dispatch (search isn't an agent loop).
-
- Trailing ``bash_turns`` (last element) counts agent-loop turns so the
- caller can surface ``tool_calls`` per row. Fallbacks to one-shot
- workers return 0 bash turns (no agent loop ran)."""
- wtype = worker.get("type", "openai")
- if wtype == "anthropic-web-search":
- # Search workers stay one-shot.
- text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
- return text, p, c, is_local, extra, n_searches, 0
- if wtype == "vllm":
- backbone = "local"
- endpoint = worker.get("base_url")
- loop_cloud_endpoint = "anthropic" # unused when backbone=local
- elif wtype in ("anthropic", "openai", "gemini"):
- backbone = "cloud"
- endpoint = None
- loop_cloud_endpoint = wtype
- else:
- # Unknown type — one-shot fallback.
- text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
- return text, p, c, is_local, extra, n_searches, 0
- out = run_swe_agent_loop(
- task,
- backbone=backbone,
- backbone_model=worker["model"],
- cloud_endpoint=loop_cloud_endpoint,
- local_endpoint=endpoint,
- initial_prompt=prompt,
- max_turns=int(cfg.get("swe_max_turns", 30)),
- bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
- output_cap=int(cfg.get("swe_output_cap", 10_000)),
- turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
- trace_prefix=f"toolorch_turn{turn}",
- workdir=workdir,
- )
- is_local = backbone == "local"
- return (
- out["final_summary"] or out["answer"],
- out["tokens_in"], out["tokens_out"],
- is_local, 0.0, 0, int(out["turns"]),
- )
-
-
-@AgentRegistry.register("toolorchestra")
-class ToolOrchestraAgent(LocalCloudAgent):
- """Multi-turn dispatcher over a mixed worker pool.
-
- Two modes (see module docstring): ``method_cfg.orchestrator_mode``
- is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
- (paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
- """
-
- agent_id = "toolorchestra"
-
- def __init__(self, *args: Any, **kwargs: Any) -> None:
- super().__init__(*args, **kwargs)
- # Validate `method_cfg.worker_pool` early — surfaces config errors
- # at agent construction rather than on the first task. No-op when
- # the override is absent.
- if self._cfg.get("worker_pool") is not None:
- _resolve_worker_pool(
- self._cfg,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
- # Validate `orchestrator_mode` (typo-checked here, not on first task).
- mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
- if mode not in ("prompted", "rl"):
- raise ValueError(
- f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
- f"got {mode!r}"
- )
-
- def _run_paradigm(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
- if mode == "rl":
- return self._run_rl(input, context, **kwargs)
- return self._run_prompted(input, context, **kwargs)
-
- # ------------------------------------------------------------------
- # Legacy prompted-orchestrator path.
- # ------------------------------------------------------------------
- def _run_prompted(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- cfg = self._cfg
- question = input
- # Resolution order (strict replace, no merge):
- # 1. `cfg["workers"]` — legacy direct override, used by tests.
- # 2. `cfg["worker_pool"]` — cell-config override; validated +
- # $local/$cloud substituted.
- # 3. `_default_pool(...)` — heterogeneous default.
- if cfg.get("workers"):
- workers = cfg["workers"]
- else:
- workers = _resolve_worker_pool(
- cfg,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
- if not workers:
- raise RuntimeError("toolorchestra: empty worker pool")
-
- max_turns = int(cfg.get("max_turns", 6))
- orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
-
- task_meta = (context.metadata.get("task") if context is not None else {}) or {}
- swe_mode = (
- bool(cfg.get("swe_use_agent_loop"))
- and bool(task_meta.get("problem_statement"))
- and bool(task_meta.get("repo"))
- and bool(task_meta.get("base_commit"))
- )
- shared_workdir: Optional[Path] = None
- if swe_mode:
- shared_workdir = Path(tempfile.mkdtemp(
- prefix=f"toolorch-swe-{task_meta.get('task_id','x')}-"
- ))
- try:
- _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
- except Exception:
- shutil.rmtree(shared_workdir, ignore_errors=True)
- raise
- self.record_trace_event({
- "kind": "toolorchestra_swe_workdir",
- "workdir": str(shared_workdir),
- "repo": task_meta["repo"],
- "base_commit": task_meta["base_commit"],
- })
-
- # try/finally guards ``shared_workdir`` against exceptions raised
- # anywhere in the turn loop, the worker calls, the fallback, or
- # the diff-extraction step. Without this, at n=500 SWE-bench an
- # exception leaves hundreds of MB of cloned repos in tempdir.
- try:
- history: List[Dict[str, Any]] = []
- tokens_local = 0
- tokens_cloud = 0
- cost = 0.0
- n_web_searches_total = 0
- # tool_calls: bash turns from SWE subloops + web_search uses
- # from GAIA. Orchestrator dispatch turns are NOT counted (they
- # produce text only — calling a worker is one tool call's worth
- # of "delegation" but the actual tool action happens inside).
- tool_calls = 0
- final_answer: Optional[str] = None
- forced_final = False
- parse_failures = 0
-
- for turn in range(1, max_turns + 1):
- sys_prompt = ORCHESTRATOR_SYS
- if turn == max_turns and final_answer is None:
- sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
- forced_final = True
-
- user = _build_user_prompt(question, workers, history)
- text, o_in, o_out = self._call_cloud(
- user=user,
- system=sys_prompt,
- max_tokens=orch_max_tokens,
- temperature=0.0,
- )
- tokens_cloud += o_in + o_out
- cost += self.cost_usd(self._cloud_model, o_in, o_out)
-
- action = _parse_action(text)
- history.append({
- "role": "orchestrator", "turn": turn, "raw": text, "action": action,
- })
- self.record_trace_event({
- "kind": "toolorchestra_action",
- "turn": turn,
- "action": action,
- "raw": text,
- })
-
- if action is None:
- parse_failures += 1
- if parse_failures >= 2 or forced_final:
- final_answer = _extract_final_answer_text(text)
- break
- continue
-
- kind = action.get("action")
- if kind == "final_answer":
- final_answer = str(action.get("answer", "")).strip()
- break
- if kind == "call_worker":
- wid = action.get("worker_id")
- w_input = action.get("input", "")
- if not isinstance(wid, int) or not (0 <= wid < len(workers)):
- parse_failures += 1
- if parse_failures >= 2 or forced_final:
- final_answer = _extract_final_answer_text(text)
- break
- continue
- worker = workers[wid]
- if swe_mode and shared_workdir is not None:
- (w_text, w_in, w_out, is_local, extra_cost,
- n_searches, bash_turns) = (
- _swe_call_worker(
- worker, str(w_input), cfg, task_meta,
- shared_workdir, turn,
- )
- )
- tool_calls += bash_turns
- else:
- w_text, w_in, w_out, is_local, extra_cost, n_searches = (
- _call_worker(worker, str(w_input), cfg)
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- n_web_searches_total += n_searches
- tool_calls += n_searches
- history.append({
- "role": "worker",
- "turn": turn,
- "worker_id": wid,
- "worker_name": worker["name"],
- "worker_model": worker["model"],
- "output": w_text,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "n_web_searches": n_searches,
- })
- continue
- # Unknown action kind — treat as parse failure.
- parse_failures += 1
-
- if final_answer is None:
- # Hard fallback: call the strongest non-search worker directly.
- # "Strongest" = highest output-token price in `_prices.PRICES`,
- # which tracks model capability tier closely enough for this.
- # Search workers are excluded — they answer fact-lookup
- # questions, not synthesis.
- non_search = [
- w for w in workers if w.get("type") != "anthropic-web-search"
- ] or workers
- worker = max(
- non_search,
- key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
- )
- if swe_mode and shared_workdir is not None:
- (ans, w_in, w_out, is_local, extra_cost, _,
- bash_turns) = _swe_call_worker(
- worker, question, cfg, task_meta,
- shared_workdir, max_turns + 1,
- )
- tool_calls += bash_turns
- else:
- ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
- worker, question, cfg
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- history.append({
- "role": "worker",
- "turn": max_turns + 1,
- "worker_id": worker["id"],
- "worker_name": worker["name"],
- "worker_model": worker["model"],
- "output": ans,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "fallback": True,
- })
- final_answer = ans
-
- # In SWE mode, the authoritative output is the working-tree diff —
- # frame it (the runner extracts it via the scorer's ```diff fence).
- if swe_mode and shared_workdir is not None:
- patch = _extract_diff(shared_workdir)
- if patch.strip():
- final_answer = (
- f"{final_answer}\n\n```diff\n{patch}```"
- if final_answer else f"```diff\n{patch}```"
- )
-
- meta = {
- "tokens_local": tokens_local,
- "tokens_cloud": tokens_cloud,
- "cost_usd": cost,
- "turns": len([h for h in history if h["role"] == "orchestrator"]),
- "web_search_uses": n_web_searches_total,
- "tool_calls": int(tool_calls),
- "traces": {
- "history": history,
- "forced_final": forced_final,
- "parse_failures": parse_failures,
- "workers": workers,
- "n_web_searches": n_web_searches_total,
- "note": (
- "inference-only port; the RL-trained Nemotron-Orchestrator-8B "
- "is NOT in the loop. Results are preliminary."
- ),
- },
- }
- return final_answer, meta
- finally:
- if shared_workdir is not None:
- shutil.rmtree(shared_workdir, ignore_errors=True)
-
- # ------------------------------------------------------------------
- # Paper-faithful Orchestrator-8B path.
- # ------------------------------------------------------------------
- def _run_rl(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- cfg = self._cfg
- question = input
-
- # Orchestrator endpoint / model (where the RL'd 8B lives).
- orch_endpoint = str(
- cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
- )
- orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
- max_turns = int(cfg.get("max_turns", 8))
- orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
- orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
-
- # Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
- # replaces `_expert_for` and `enhance_reasoning` is post-processed
- # through a Modal Python sandbox. See module docstring + paper-match
- # doc at `docs/26.5.19/toolorchestra-papermatch.md`.
- paper_mode = str(cfg.get("pool", "")).lower() == "paper"
-
- # SWE-bench detection: same gate as the prompted path. Requires
- # `method_cfg.swe_use_agent_loop = true` AND the task carries the
- # SWE-bench fields. When active, the `enhance_reasoning` and
- # `answer` workers route through `run_swe_agent_loop` on a shared
- # workdir; search workers stay one-shot. At end, the working-tree
- # diff is appended to final_answer so `_score_swebench` can extract
- # it via the ```diff fence.
- task_meta = (context.metadata.get("task") if context is not None else {}) or {}
- swe_mode = (
- bool(cfg.get("swe_use_agent_loop"))
- and bool(task_meta.get("problem_statement"))
- and bool(task_meta.get("repo"))
- and bool(task_meta.get("base_commit"))
- )
- shared_workdir: Optional[Path] = None
- if swe_mode:
- shared_workdir = Path(tempfile.mkdtemp(
- prefix=f"toolorch-rl-swe-{task_meta.get('task_id','x')}-"
- ))
- try:
- _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
- except Exception:
- shutil.rmtree(shared_workdir, ignore_errors=True)
- raise
- self.record_trace_event({
- "kind": "toolorchestra_rl_swe_workdir",
- "workdir": str(shared_workdir),
- "repo": task_meta["repo"],
- "base_commit": task_meta["base_commit"],
- })
-
- # ``context_str`` mirrors the upstream's running context — accumulates
- # search documents and code/exec snippets across turns. We keep this
- # as a single string for prompt simplicity; the upstream uses
- # tokenized cutoffs (we cap at ~24k chars instead).
- context_str = ""
- doc_list: List[str] = []
- history: List[Dict[str, Any]] = []
- tokens_local = 0
- tokens_cloud = 0
- cost = 0.0
- n_web_searches_total = 0
- tool_calls = 0
- final_answer: Optional[str] = None
- parse_failures = 0
-
- # Single outer try/finally guards `shared_workdir` against any
- # exception in the orchestrator loop, the post-loop fallback, or
- # the diff-extraction step. Matches the prompted path's pattern.
- try:
- for turn in range(1, max_turns + 1):
- user = (
- f"Problem: {question}\n\n{context_str}\n\n"
- "Choose an appropriate tool."
- )
-
- # Orchestrator-8B served on local vLLM. We pass the three NVlabs
- # tools verbatim. In paper-mode we use the local helper so we
- # get the SDK-level ``tool_calls`` object back — `_call_vllm`
- # returns just text and loses the call when vLLM's parser
- # caught it. Orchestrator-8B emits its routing decision in the
- # OpenAI-native ``tool_calls`` array with an empty text body,
- # so the legacy `_call_vllm` path saw nothing and silently fell
- # through to the answer-1 fallback (parse_failures: 2 on every
- # non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
- # modes now use `_call_orchestrator_with_tool_calls` so the
- # parser can read structured tool calls; the text-tag path in
- # `_parse_rl_tool_call` is still the fallback when `tool_calls`
- # is empty.
- text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
- orch_model,
- orch_endpoint,
- user=user,
- system=RL_ORCHESTRATOR_SYS,
- max_tokens=orch_max_tokens,
- temperature=orch_temp,
- tools=RL_TOOLS_SPEC,
- )
- self.record_trace_event({
- "kind": "vllm",
- "role": "orchestrator",
- "model": orch_model,
- "endpoint": orch_endpoint,
- "system": RL_ORCHESTRATOR_SYS,
- "user": user,
- "response": text,
- "tool_calls": [
- {
- "id": getattr(tc, "id", None),
- "type": getattr(tc, "type", None),
- "function": {
- "name": getattr(getattr(tc, "function", None), "name", None),
- "arguments": getattr(getattr(tc, "function", None), "arguments", None),
- },
- }
- for tc in (sdk_tool_calls or [])
- ],
- "tokens_in": o_in,
- "tokens_out": o_out,
- })
- tokens_local += o_in + o_out
-
- action = _parse_rl_tool_call(text, sdk_tool_calls)
- history.append({
- "role": "orchestrator", "turn": turn, "raw": text, "action": action,
- })
- self.record_trace_event({
- "kind": "toolorchestra_rl_action",
- "turn": turn,
- "action": action,
- "raw": text,
- })
-
- if action is None:
- parse_failures += 1
- if parse_failures >= 2:
- break
- continue
-
- name = action["name"]
- args = action.get("arguments", {})
- slot = args.get("model", "")
-
- # Validate against the upstream tool/arg schema.
- valid = name in RL_ALL_TOOLS and isinstance(slot, str) and (
- slot in RL_ALL_TOOLS[name]["model"]
- )
- if not valid:
- parse_failures += 1
- if parse_failures >= 2:
- break
- # Replay with a softer nudge in the context.
- context_str += (
- f"\n[Orchestrator emitted invalid tool call "
- f"name={name!r} slot={slot!r} — try again.]\n"
- )
- continue
-
- # Paper-match (`method_cfg.pool == "paper"`) routes through
- # the Tavily/OpenRouter/Modal pool instead of the default
- # Anthropic-web-search-driven mapping. For `search` this
- # also forces the worker prompt to a raw query string
- # (Tavily takes a single search string, not a chat-style
- # framing).
- if paper_mode:
- worker = _paper_expert_for(
- slot, self._local_model, self._local_endpoint,
- self._cloud_model, self._cloud_endpoint,
- )
- # In paper mode, `enhance_reasoning` is always the coder
- # specialist regardless of the orchestrator's chosen tier.
- # The coder is then expected to emit a python block which
- # we exec in Modal (below).
- if name == "enhance_reasoning":
- worker = {
- "name": f"coder:{slot}",
- "type": "openrouter",
- "model": _PAPER_CODER_OPENROUTER,
- }
- else:
- worker = _expert_for(
- slot, self._local_model, self._local_endpoint, self._cloud_model,
- self._cloud_endpoint,
- )
-
- # Dispatch — the orchestrator only conveys a tool/model
- # choice, NOT a question rewrite; the prompt we send the
- # expert is the same context the orchestrator saw, framed
- # appropriately for the tool.
- if name == "search":
- if paper_mode:
- # Tavily takes a query string. Orchestrator-8B often
- # emits an extra `query` arg (not in the upstream
- # schema but useful) — prefer it; else fall back to
- # the raw question.
- q = args.get("query")
- w_input = q if isinstance(q, str) and q.strip() else question
- else:
- w_input = (
- f"Search the web to gather information that helps answer:\n"
- f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
- )
- elif name == "enhance_reasoning":
- if paper_mode:
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Write a short Python script that computes intermediate "
- "results which help answer the problem. Output ONLY the "
- "code inside one ```python ... ``` fenced block. Print "
- "any results you derive using `print(...)`. The script "
- "must run with the Python stdlib only — no extra pip "
- "installs."
- )
- else:
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Reason carefully. Outline the key intermediate steps and any "
- "computations or facts you can derive. Do NOT give a final "
- "answer — the orchestrator will collect your reasoning and "
- "call the answer tool next."
- )
- else: # name == "answer"
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Provide the final answer to the user. Respect any "
- "answer-format rules in the question (e.g. GAIA's "
- "FINAL ANSWER: convention)."
- )
-
- # SWE mode: route enhance_reasoning / answer workers through
- # the SWE agent loop on the shared workdir so they can read
- # files, run tests, and edit the working tree. Search workers
- # stay one-shot (no agent loop). The `_swe_call_worker`
- # one-shot fallbacks (openai-typed workers, search) return
- # bash_turns=0; vllm/anthropic-typed workers run the loop.
- bash_turns = 0
- if swe_mode and shared_workdir is not None and name != "search":
- (w_text, w_in, w_out, is_local, extra_cost,
- n_searches, bash_turns) = _swe_call_worker(
- worker, w_input, cfg, task_meta, shared_workdir, turn,
- )
- else:
- w_text, w_in, w_out, is_local, extra_cost, n_searches = _call_worker(
- worker, w_input, cfg
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- n_web_searches_total += n_searches
- # SWE bash turns count as tool calls (each one is a $BASH block
- # the agent executed). On non-SWE turns fall back to the
- # original "at least one expert call" accounting.
- tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches)
-
- # Paper-match: pipe coder output through a Modal sandbox so
- # `enhance_reasoning` actually executes the code the coder
- # wrote. Append the exec output to the worker's text. No-op
- # when no python block is found.
- modal_exec_output: Optional[str] = None
- modal_exec_rc: Optional[int] = None
- if (paper_mode and name == "enhance_reasoning"
- and not swe_mode):
- code = _extract_first_python_block(w_text)
- if code:
- timeout_s = int(cfg.get("modal_python_timeout_s", 60))
- modal_exec_output, modal_exec_rc = _call_modal_python(
- code, timeout_s=timeout_s,
- )
- tool_calls += 1
- w_text = (
- f"{w_text}\n\n[modal-python stdout/stderr "
- f"(rc={modal_exec_rc})]\n{modal_exec_output}"
- )
-
- history.append({
- "role": "worker",
- "turn": turn,
- "tool": name,
- "slot": slot,
- "worker_model": worker["model"],
- "worker_type": worker["type"],
- "output": w_text,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "n_web_searches": n_searches,
- "bash_turns": bash_turns,
- "modal_exec_rc": modal_exec_rc,
- })
-
- # Update accumulated context for the next turn.
- if name == "search":
- # Treat the search worker's response as a document.
- doc_list.append(w_text)
- ctx_docs = "\n\n".join(
- f"Doc {i+1}: {d}" for i, d in enumerate(doc_list)
- )
- # Crude char-level cap mirrors the upstream's ~24k token cap.
- context_str = ("Documents:\n" + ctx_docs)[-24000:]
- elif name == "enhance_reasoning":
- snippet = f"\n\nReasoning/exec output:\n{w_text}"
- context_str = (context_str + snippet)[-24000:]
- else: # answer
- final_answer = w_text.strip()
- break
-
- if final_answer is None:
- # Hard fallback: ask the frontier worker directly. In SWE
- # mode route this final call through the agent loop too so
- # it can still touch the workdir and emit a diff.
- expert_fn = _paper_expert_for if paper_mode else _expert_for
- worker = expert_fn(
- "answer-1", self._local_model, self._local_endpoint,
- self._cloud_model, self._cloud_endpoint,
- )
- fb_bash_turns = 0
- if swe_mode and shared_workdir is not None:
- (ans, w_in, w_out, is_local, extra_cost,
- _, fb_bash_turns) = _swe_call_worker(
- worker, question, cfg, task_meta,
- shared_workdir, max_turns + 1,
- )
- tool_calls += fb_bash_turns
- else:
- ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
- worker, question, cfg
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- history.append({
- "role": "worker",
- "turn": max_turns + 1,
- "tool": "answer",
- "slot": "answer-1",
- "worker_model": worker["model"],
- "worker_type": worker["type"],
- "output": ans,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "bash_turns": fb_bash_turns,
- "fallback": True,
- })
- final_answer = ans
-
- # In SWE mode, the authoritative output is the working-tree diff —
- # frame it so `_score_swebench`'s extract_patch picks it up.
- if swe_mode and shared_workdir is not None:
- patch = _extract_diff(shared_workdir)
- if patch.strip():
- final_answer = (
- f"{final_answer}\n\n```diff\n{patch}```"
- if final_answer else f"```diff\n{patch}```"
- )
-
- meta = {
- "tokens_local": tokens_local,
- "tokens_cloud": tokens_cloud,
- "cost_usd": cost,
- "turns": len([h for h in history if h["role"] == "orchestrator"]),
- "web_search_uses": n_web_searches_total,
- "tool_calls": int(tool_calls),
- "traces": {
- "history": history,
- "parse_failures": parse_failures,
- "orchestrator_model": orch_model,
- "orchestrator_endpoint": orch_endpoint,
- "mode": "rl",
- "pool": "paper" if paper_mode else "default",
- "swe_mode": swe_mode,
- "note": (
- "RL-trained nvidia/Orchestrator-8B as orchestrator. "
- "Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto "
- "our hybrid worker types — see toolorchestra.py docstring."
- ),
- },
- }
- return final_answer, meta
- finally:
- if shared_workdir is not None:
- shutil.rmtree(shared_workdir, ignore_errors=True)
-
-
-__all__ = ["ToolOrchestraAgent"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/__init__.py b/src/openjarvis/agents/hybrid/toolorchestra/__init__.py
new file mode 100644
index 00000000..33226abe
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/__init__.py
@@ -0,0 +1,20 @@
+"""ToolOrchestraAgent package (split from the former toolorchestra.py module).
+
+Importing this package registers the agent and re-exports the public surface,
+so ``from openjarvis.agents.hybrid.toolorchestra import ToolOrchestraAgent``
+keeps working unchanged. Submodules:
+
+ prompts — system prompts, RL tool specs / arg schema
+ experts — slot -> backend worker mapping (default + paper-match)
+ sandbox — Tavily search + Modal Python sandbox helpers
+ clients — orchestrator vLLM tool-call client
+ parsing — action / tool-call parsing + user-prompt assembly
+ workers — worker pool resolution + dispatch (_call_worker etc.)
+ agent — ToolOrchestraAgent (the registered agent class)
+"""
+
+from __future__ import annotations
+
+from openjarvis.agents.hybrid.toolorchestra.agent import ToolOrchestraAgent
+
+__all__ = ["ToolOrchestraAgent"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/agent.py b/src/openjarvis/agents/hybrid/toolorchestra/agent.py
new file mode 100644
index 00000000..b2fc963a
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/agent.py
@@ -0,0 +1,771 @@
+"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
+
+Two modes, gated by ``method_cfg.orchestrator_mode``:
+
+* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
+ orchestrator, dispatching to a numbered worker pool via JSON
+ ``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
+ a prompted upper-bound reference point — NOT the paper's setup.
+
+* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
+ served on a local vLLM is the orchestrator. It emits OpenAI-style
+ ``tool_calls`` (or ``{...}`` text blocks when
+ vLLM's tool parser doesn't catch them) for three expert tools —
+ ``enhance_reasoning``, ``answer``, ``search`` — exactly as in the
+ upstream ``evaluation/tools.json``. Each tool's ``model`` arg
+ (``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real
+ backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier
+ Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
+ for `*-3`. Search routes to the Anthropic server-side web_search.
+
+ We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
+ code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
+ Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
+ worker types. Energy-wise, "expert" answers are cloud calls.
+
+Pipeline per task (RL mode):
+
+1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
+ appropriate tool.` with the three tools declared.
+2. It emits one ``tool_call`` per turn — ``search`` updates the
+ context, ``enhance_reasoning`` appends code/exec output (we run the
+ tool as a plain LLM call, no sandbox — the model just gets prose
+ back), ``answer`` produces the final answer and the loop stops.
+3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
+ back to the strongest expert worker.
+
+Prompted-mode pipeline:
+
+1. Orchestrator (cloud) reads question + numbered worker pool.
+2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
+ "input": str}`` or ``{"action": "final_answer", "answer": str}``.
+3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
+ prompt; fallback to strongest worker on parse failure.
+
+Workers come from ``cfg["workers"]`` or a sensible default pool (local
+Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
+gpt-5-mini).
+"""
+
+
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from openjarvis.agents._stubs import AgentContext
+from openjarvis.agents.hybrid._base import LocalCloudAgent
+from openjarvis.agents.hybrid._prices import PRICES
+from openjarvis.agents.hybrid.mini_swe_agent import (
+ _clone_repo,
+ _extract_diff,
+ run_swe_agent_loop,
+)
+from openjarvis.core.registry import AgentRegistry
+
+from openjarvis.agents.hybrid.toolorchestra.prompts import (
+ FORCE_FINAL_PROMPT,
+ ORCHESTRATOR_SYS,
+ RL_ALL_TOOLS,
+ RL_ORCHESTRATOR_SYS,
+ RL_TOOLS_SPEC,
+)
+from openjarvis.agents.hybrid.toolorchestra.experts import (
+ _PAPER_CODER_OPENROUTER,
+ _expert_for,
+ _paper_expert_for,
+)
+from openjarvis.agents.hybrid.toolorchestra.sandbox import (
+ _call_modal_python,
+ _extract_first_python_block,
+)
+from openjarvis.agents.hybrid.toolorchestra.clients import (
+ _call_orchestrator_with_tool_calls,
+)
+from openjarvis.agents.hybrid.toolorchestra.parsing import (
+ _build_user_prompt,
+ _extract_final_answer_text,
+ _parse_action,
+ _parse_rl_tool_call,
+)
+from openjarvis.agents.hybrid.toolorchestra.workers import (
+ _call_worker,
+ _default_pool,
+ _resolve_worker_pool,
+ _swe_call_worker,
+)
+
+@AgentRegistry.register("toolorchestra")
+class ToolOrchestraAgent(LocalCloudAgent):
+ """Multi-turn dispatcher over a mixed worker pool.
+
+ Two modes (see module docstring): ``method_cfg.orchestrator_mode``
+ is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
+ (paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
+ """
+
+ agent_id = "toolorchestra"
+
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ super().__init__(*args, **kwargs)
+ # Validate `method_cfg.worker_pool` early — surfaces config errors
+ # at agent construction rather than on the first task. No-op when
+ # the override is absent.
+ if self._cfg.get("worker_pool") is not None:
+ _resolve_worker_pool(
+ self._cfg,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+ # Validate `orchestrator_mode` (typo-checked here, not on first task).
+ mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
+ if mode not in ("prompted", "rl"):
+ raise ValueError(
+ f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
+ f"got {mode!r}"
+ )
+
+ def _run_paradigm(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
+ if mode == "rl":
+ return self._run_rl(input, context, **kwargs)
+ return self._run_prompted(input, context, **kwargs)
+
+ # ------------------------------------------------------------------
+ # Legacy prompted-orchestrator path.
+ # ------------------------------------------------------------------
+ def _run_prompted(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ cfg = self._cfg
+ question = input
+ # Resolution order (strict replace, no merge):
+ # 1. `cfg["workers"]` — legacy direct override, used by tests.
+ # 2. `cfg["worker_pool"]` — cell-config override; validated +
+ # $local/$cloud substituted.
+ # 3. `_default_pool(...)` — heterogeneous default.
+ if cfg.get("workers"):
+ workers = cfg["workers"]
+ else:
+ workers = _resolve_worker_pool(
+ cfg,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+ if not workers:
+ raise RuntimeError("toolorchestra: empty worker pool")
+
+ max_turns = int(cfg.get("max_turns", 6))
+ orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
+
+ task_meta = (context.metadata.get("task") if context is not None else {}) or {}
+ swe_mode = (
+ bool(cfg.get("swe_use_agent_loop"))
+ and bool(task_meta.get("problem_statement"))
+ and bool(task_meta.get("repo"))
+ and bool(task_meta.get("base_commit"))
+ )
+ shared_workdir: Optional[Path] = None
+ if swe_mode:
+ shared_workdir = Path(tempfile.mkdtemp(
+ prefix=f"toolorch-swe-{task_meta.get('task_id','x')}-"
+ ))
+ try:
+ _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
+ except Exception:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+ raise
+ self.record_trace_event({
+ "kind": "toolorchestra_swe_workdir",
+ "workdir": str(shared_workdir),
+ "repo": task_meta["repo"],
+ "base_commit": task_meta["base_commit"],
+ })
+
+ # try/finally guards ``shared_workdir`` against exceptions raised
+ # anywhere in the turn loop, the worker calls, the fallback, or
+ # the diff-extraction step. Without this, at n=500 SWE-bench an
+ # exception leaves hundreds of MB of cloned repos in tempdir.
+ try:
+ history: List[Dict[str, Any]] = []
+ tokens_local = 0
+ tokens_cloud = 0
+ cost = 0.0
+ n_web_searches_total = 0
+ # tool_calls: bash turns from SWE subloops + web_search uses
+ # from GAIA. Orchestrator dispatch turns are NOT counted (they
+ # produce text only — calling a worker is one tool call's worth
+ # of "delegation" but the actual tool action happens inside).
+ tool_calls = 0
+ final_answer: Optional[str] = None
+ forced_final = False
+ parse_failures = 0
+
+ for turn in range(1, max_turns + 1):
+ sys_prompt = ORCHESTRATOR_SYS
+ if turn == max_turns and final_answer is None:
+ sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
+ forced_final = True
+
+ user = _build_user_prompt(question, workers, history)
+ text, o_in, o_out = self._call_cloud(
+ user=user,
+ system=sys_prompt,
+ max_tokens=orch_max_tokens,
+ temperature=0.0,
+ )
+ tokens_cloud += o_in + o_out
+ cost += self.cost_usd(self._cloud_model, o_in, o_out)
+
+ action = _parse_action(text)
+ history.append({
+ "role": "orchestrator", "turn": turn, "raw": text, "action": action,
+ })
+ self.record_trace_event({
+ "kind": "toolorchestra_action",
+ "turn": turn,
+ "action": action,
+ "raw": text,
+ })
+
+ if action is None:
+ parse_failures += 1
+ if parse_failures >= 2 or forced_final:
+ final_answer = _extract_final_answer_text(text)
+ break
+ continue
+
+ kind = action.get("action")
+ if kind == "final_answer":
+ final_answer = str(action.get("answer", "")).strip()
+ break
+ if kind == "call_worker":
+ wid = action.get("worker_id")
+ w_input = action.get("input", "")
+ if not isinstance(wid, int) or not (0 <= wid < len(workers)):
+ parse_failures += 1
+ if parse_failures >= 2 or forced_final:
+ final_answer = _extract_final_answer_text(text)
+ break
+ continue
+ worker = workers[wid]
+ if swe_mode and shared_workdir is not None:
+ (w_text, w_in, w_out, is_local, extra_cost,
+ n_searches, bash_turns) = (
+ _swe_call_worker(
+ worker, str(w_input), cfg, task_meta,
+ shared_workdir, turn,
+ )
+ )
+ tool_calls += bash_turns
+ else:
+ w_text, w_in, w_out, is_local, extra_cost, n_searches = (
+ _call_worker(worker, str(w_input), cfg)
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ n_web_searches_total += n_searches
+ tool_calls += n_searches
+ history.append({
+ "role": "worker",
+ "turn": turn,
+ "worker_id": wid,
+ "worker_name": worker["name"],
+ "worker_model": worker["model"],
+ "output": w_text,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "n_web_searches": n_searches,
+ })
+ continue
+ # Unknown action kind — treat as parse failure.
+ parse_failures += 1
+
+ if final_answer is None:
+ # Hard fallback: call the strongest non-search worker directly.
+ # "Strongest" = highest output-token price in `_prices.PRICES`,
+ # which tracks model capability tier closely enough for this.
+ # Search workers are excluded — they answer fact-lookup
+ # questions, not synthesis.
+ non_search = [
+ w for w in workers if w.get("type") != "anthropic-web-search"
+ ] or workers
+ worker = max(
+ non_search,
+ key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
+ )
+ if swe_mode and shared_workdir is not None:
+ (ans, w_in, w_out, is_local, extra_cost, _,
+ bash_turns) = _swe_call_worker(
+ worker, question, cfg, task_meta,
+ shared_workdir, max_turns + 1,
+ )
+ tool_calls += bash_turns
+ else:
+ ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
+ worker, question, cfg
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ history.append({
+ "role": "worker",
+ "turn": max_turns + 1,
+ "worker_id": worker["id"],
+ "worker_name": worker["name"],
+ "worker_model": worker["model"],
+ "output": ans,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "fallback": True,
+ })
+ final_answer = ans
+
+ # In SWE mode, the authoritative output is the working-tree diff —
+ # frame it (the runner extracts it via the scorer's ```diff fence).
+ if swe_mode and shared_workdir is not None:
+ patch = _extract_diff(shared_workdir)
+ if patch.strip():
+ final_answer = (
+ f"{final_answer}\n\n```diff\n{patch}```"
+ if final_answer else f"```diff\n{patch}```"
+ )
+
+ meta = {
+ "tokens_local": tokens_local,
+ "tokens_cloud": tokens_cloud,
+ "cost_usd": cost,
+ "turns": len([h for h in history if h["role"] == "orchestrator"]),
+ "web_search_uses": n_web_searches_total,
+ "tool_calls": int(tool_calls),
+ "traces": {
+ "history": history,
+ "forced_final": forced_final,
+ "parse_failures": parse_failures,
+ "workers": workers,
+ "n_web_searches": n_web_searches_total,
+ "note": (
+ "inference-only port; the RL-trained Nemotron-Orchestrator-8B "
+ "is NOT in the loop. Results are preliminary."
+ ),
+ },
+ }
+ return final_answer, meta
+ finally:
+ if shared_workdir is not None:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+
+ # ------------------------------------------------------------------
+ # Paper-faithful Orchestrator-8B path.
+ # ------------------------------------------------------------------
+ def _run_rl(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ cfg = self._cfg
+ question = input
+
+ # Orchestrator endpoint / model (where the RL'd 8B lives).
+ orch_endpoint = str(
+ cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
+ )
+ orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
+ max_turns = int(cfg.get("max_turns", 8))
+ orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
+ orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
+
+ # Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
+ # replaces `_expert_for` and `enhance_reasoning` is post-processed
+ # through a Modal Python sandbox. See module docstring + paper-match
+ # doc at `docs/26.5.19/toolorchestra-papermatch.md`.
+ paper_mode = str(cfg.get("pool", "")).lower() == "paper"
+
+ # SWE-bench detection: same gate as the prompted path. Requires
+ # `method_cfg.swe_use_agent_loop = true` AND the task carries the
+ # SWE-bench fields. When active, the `enhance_reasoning` and
+ # `answer` workers route through `run_swe_agent_loop` on a shared
+ # workdir; search workers stay one-shot. At end, the working-tree
+ # diff is appended to final_answer so `_score_swebench` can extract
+ # it via the ```diff fence.
+ task_meta = (context.metadata.get("task") if context is not None else {}) or {}
+ swe_mode = (
+ bool(cfg.get("swe_use_agent_loop"))
+ and bool(task_meta.get("problem_statement"))
+ and bool(task_meta.get("repo"))
+ and bool(task_meta.get("base_commit"))
+ )
+ shared_workdir: Optional[Path] = None
+ if swe_mode:
+ shared_workdir = Path(tempfile.mkdtemp(
+ prefix=f"toolorch-rl-swe-{task_meta.get('task_id','x')}-"
+ ))
+ try:
+ _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
+ except Exception:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+ raise
+ self.record_trace_event({
+ "kind": "toolorchestra_rl_swe_workdir",
+ "workdir": str(shared_workdir),
+ "repo": task_meta["repo"],
+ "base_commit": task_meta["base_commit"],
+ })
+
+ # ``context_str`` mirrors the upstream's running context — accumulates
+ # search documents and code/exec snippets across turns. We keep this
+ # as a single string for prompt simplicity; the upstream uses
+ # tokenized cutoffs (we cap at ~24k chars instead).
+ context_str = ""
+ doc_list: List[str] = []
+ history: List[Dict[str, Any]] = []
+ tokens_local = 0
+ tokens_cloud = 0
+ cost = 0.0
+ n_web_searches_total = 0
+ tool_calls = 0
+ final_answer: Optional[str] = None
+ parse_failures = 0
+
+ # Single outer try/finally guards `shared_workdir` against any
+ # exception in the orchestrator loop, the post-loop fallback, or
+ # the diff-extraction step. Matches the prompted path's pattern.
+ try:
+ for turn in range(1, max_turns + 1):
+ user = (
+ f"Problem: {question}\n\n{context_str}\n\n"
+ "Choose an appropriate tool."
+ )
+
+ # Orchestrator-8B served on local vLLM. We pass the three NVlabs
+ # tools verbatim. In paper-mode we use the local helper so we
+ # get the SDK-level ``tool_calls`` object back — `_call_vllm`
+ # returns just text and loses the call when vLLM's parser
+ # caught it. Orchestrator-8B emits its routing decision in the
+ # OpenAI-native ``tool_calls`` array with an empty text body,
+ # so the legacy `_call_vllm` path saw nothing and silently fell
+ # through to the answer-1 fallback (parse_failures: 2 on every
+ # non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
+ # modes now use `_call_orchestrator_with_tool_calls` so the
+ # parser can read structured tool calls; the text-tag path in
+ # `_parse_rl_tool_call` is still the fallback when `tool_calls`
+ # is empty.
+ text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
+ orch_model,
+ orch_endpoint,
+ user=user,
+ system=RL_ORCHESTRATOR_SYS,
+ max_tokens=orch_max_tokens,
+ temperature=orch_temp,
+ tools=RL_TOOLS_SPEC,
+ )
+ self.record_trace_event({
+ "kind": "vllm",
+ "role": "orchestrator",
+ "model": orch_model,
+ "endpoint": orch_endpoint,
+ "system": RL_ORCHESTRATOR_SYS,
+ "user": user,
+ "response": text,
+ "tool_calls": [
+ {
+ "id": getattr(tc, "id", None),
+ "type": getattr(tc, "type", None),
+ "function": {
+ "name": getattr(getattr(tc, "function", None), "name", None),
+ "arguments": getattr(getattr(tc, "function", None), "arguments", None),
+ },
+ }
+ for tc in (sdk_tool_calls or [])
+ ],
+ "tokens_in": o_in,
+ "tokens_out": o_out,
+ })
+ tokens_local += o_in + o_out
+
+ action = _parse_rl_tool_call(text, sdk_tool_calls)
+ history.append({
+ "role": "orchestrator", "turn": turn, "raw": text, "action": action,
+ })
+ self.record_trace_event({
+ "kind": "toolorchestra_rl_action",
+ "turn": turn,
+ "action": action,
+ "raw": text,
+ })
+
+ if action is None:
+ parse_failures += 1
+ if parse_failures >= 2:
+ break
+ continue
+
+ name = action["name"]
+ args = action.get("arguments", {})
+ slot = args.get("model", "")
+
+ # Validate against the upstream tool/arg schema.
+ valid = name in RL_ALL_TOOLS and isinstance(slot, str) and (
+ slot in RL_ALL_TOOLS[name]["model"]
+ )
+ if not valid:
+ parse_failures += 1
+ if parse_failures >= 2:
+ break
+ # Replay with a softer nudge in the context.
+ context_str += (
+ f"\n[Orchestrator emitted invalid tool call "
+ f"name={name!r} slot={slot!r} — try again.]\n"
+ )
+ continue
+
+ # Paper-match (`method_cfg.pool == "paper"`) routes through
+ # the Tavily/OpenRouter/Modal pool instead of the default
+ # Anthropic-web-search-driven mapping. For `search` this
+ # also forces the worker prompt to a raw query string
+ # (Tavily takes a single search string, not a chat-style
+ # framing).
+ if paper_mode:
+ worker = _paper_expert_for(
+ slot, self._local_model, self._local_endpoint,
+ self._cloud_model, self._cloud_endpoint,
+ )
+ # In paper mode, `enhance_reasoning` is always the coder
+ # specialist regardless of the orchestrator's chosen tier.
+ # The coder is then expected to emit a python block which
+ # we exec in Modal (below).
+ if name == "enhance_reasoning":
+ worker = {
+ "name": f"coder:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_CODER_OPENROUTER,
+ }
+ else:
+ worker = _expert_for(
+ slot, self._local_model, self._local_endpoint, self._cloud_model,
+ self._cloud_endpoint,
+ )
+
+ # Dispatch — the orchestrator only conveys a tool/model
+ # choice, NOT a question rewrite; the prompt we send the
+ # expert is the same context the orchestrator saw, framed
+ # appropriately for the tool.
+ if name == "search":
+ if paper_mode:
+ # Tavily takes a query string. Orchestrator-8B often
+ # emits an extra `query` arg (not in the upstream
+ # schema but useful) — prefer it; else fall back to
+ # the raw question.
+ q = args.get("query")
+ w_input = q if isinstance(q, str) and q.strip() else question
+ else:
+ w_input = (
+ f"Search the web to gather information that helps answer:\n"
+ f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
+ )
+ elif name == "enhance_reasoning":
+ if paper_mode:
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Write a short Python script that computes intermediate "
+ "results which help answer the problem. Output ONLY the "
+ "code inside one ```python ... ``` fenced block. Print "
+ "any results you derive using `print(...)`. The script "
+ "must run with the Python stdlib only — no extra pip "
+ "installs."
+ )
+ else:
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Reason carefully. Outline the key intermediate steps and any "
+ "computations or facts you can derive. Do NOT give a final "
+ "answer — the orchestrator will collect your reasoning and "
+ "call the answer tool next."
+ )
+ else: # name == "answer"
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Provide the final answer to the user. Respect any "
+ "answer-format rules in the question (e.g. GAIA's "
+ "FINAL ANSWER: convention)."
+ )
+
+ # SWE mode: route enhance_reasoning / answer workers through
+ # the SWE agent loop on the shared workdir so they can read
+ # files, run tests, and edit the working tree. Search workers
+ # stay one-shot (no agent loop). The `_swe_call_worker`
+ # one-shot fallbacks (openai-typed workers, search) return
+ # bash_turns=0; vllm/anthropic-typed workers run the loop.
+ bash_turns = 0
+ if swe_mode and shared_workdir is not None and name != "search":
+ (w_text, w_in, w_out, is_local, extra_cost,
+ n_searches, bash_turns) = _swe_call_worker(
+ worker, w_input, cfg, task_meta, shared_workdir, turn,
+ )
+ else:
+ w_text, w_in, w_out, is_local, extra_cost, n_searches = _call_worker(
+ worker, w_input, cfg
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ n_web_searches_total += n_searches
+ # SWE bash turns count as tool calls (each one is a $BASH block
+ # the agent executed). On non-SWE turns fall back to the
+ # original "at least one expert call" accounting.
+ tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches)
+
+ # Paper-match: pipe coder output through a Modal sandbox so
+ # `enhance_reasoning` actually executes the code the coder
+ # wrote. Append the exec output to the worker's text. No-op
+ # when no python block is found.
+ modal_exec_output: Optional[str] = None
+ modal_exec_rc: Optional[int] = None
+ if (paper_mode and name == "enhance_reasoning"
+ and not swe_mode):
+ code = _extract_first_python_block(w_text)
+ if code:
+ timeout_s = int(cfg.get("modal_python_timeout_s", 60))
+ modal_exec_output, modal_exec_rc = _call_modal_python(
+ code, timeout_s=timeout_s,
+ )
+ tool_calls += 1
+ w_text = (
+ f"{w_text}\n\n[modal-python stdout/stderr "
+ f"(rc={modal_exec_rc})]\n{modal_exec_output}"
+ )
+
+ history.append({
+ "role": "worker",
+ "turn": turn,
+ "tool": name,
+ "slot": slot,
+ "worker_model": worker["model"],
+ "worker_type": worker["type"],
+ "output": w_text,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "n_web_searches": n_searches,
+ "bash_turns": bash_turns,
+ "modal_exec_rc": modal_exec_rc,
+ })
+
+ # Update accumulated context for the next turn.
+ if name == "search":
+ # Treat the search worker's response as a document.
+ doc_list.append(w_text)
+ ctx_docs = "\n\n".join(
+ f"Doc {i+1}: {d}" for i, d in enumerate(doc_list)
+ )
+ # Crude char-level cap mirrors the upstream's ~24k token cap.
+ context_str = ("Documents:\n" + ctx_docs)[-24000:]
+ elif name == "enhance_reasoning":
+ snippet = f"\n\nReasoning/exec output:\n{w_text}"
+ context_str = (context_str + snippet)[-24000:]
+ else: # answer
+ final_answer = w_text.strip()
+ break
+
+ if final_answer is None:
+ # Hard fallback: ask the frontier worker directly. In SWE
+ # mode route this final call through the agent loop too so
+ # it can still touch the workdir and emit a diff.
+ expert_fn = _paper_expert_for if paper_mode else _expert_for
+ worker = expert_fn(
+ "answer-1", self._local_model, self._local_endpoint,
+ self._cloud_model, self._cloud_endpoint,
+ )
+ fb_bash_turns = 0
+ if swe_mode and shared_workdir is not None:
+ (ans, w_in, w_out, is_local, extra_cost,
+ _, fb_bash_turns) = _swe_call_worker(
+ worker, question, cfg, task_meta,
+ shared_workdir, max_turns + 1,
+ )
+ tool_calls += fb_bash_turns
+ else:
+ ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
+ worker, question, cfg
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ history.append({
+ "role": "worker",
+ "turn": max_turns + 1,
+ "tool": "answer",
+ "slot": "answer-1",
+ "worker_model": worker["model"],
+ "worker_type": worker["type"],
+ "output": ans,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "bash_turns": fb_bash_turns,
+ "fallback": True,
+ })
+ final_answer = ans
+
+ # In SWE mode, the authoritative output is the working-tree diff —
+ # frame it so `_score_swebench`'s extract_patch picks it up.
+ if swe_mode and shared_workdir is not None:
+ patch = _extract_diff(shared_workdir)
+ if patch.strip():
+ final_answer = (
+ f"{final_answer}\n\n```diff\n{patch}```"
+ if final_answer else f"```diff\n{patch}```"
+ )
+
+ meta = {
+ "tokens_local": tokens_local,
+ "tokens_cloud": tokens_cloud,
+ "cost_usd": cost,
+ "turns": len([h for h in history if h["role"] == "orchestrator"]),
+ "web_search_uses": n_web_searches_total,
+ "tool_calls": int(tool_calls),
+ "traces": {
+ "history": history,
+ "parse_failures": parse_failures,
+ "orchestrator_model": orch_model,
+ "orchestrator_endpoint": orch_endpoint,
+ "mode": "rl",
+ "pool": "paper" if paper_mode else "default",
+ "swe_mode": swe_mode,
+ "note": (
+ "RL-trained nvidia/Orchestrator-8B as orchestrator. "
+ "Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto "
+ "our hybrid worker types — see toolorchestra.py docstring."
+ ),
+ },
+ }
+ return final_answer, meta
+ finally:
+ if shared_workdir is not None:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+
+
+__all__ = ["ToolOrchestraAgent"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/clients.py b/src/openjarvis/agents/hybrid/toolorchestra/clients.py
new file mode 100644
index 00000000..d941862b
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/clients.py
@@ -0,0 +1,50 @@
+"""Orchestrator vLLM tool-call client for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Tuple
+
+def _call_orchestrator_with_tool_calls(
+ model: str,
+ endpoint: str,
+ *,
+ user: str,
+ system: str,
+ max_tokens: int,
+ temperature: float,
+ tools: List[Dict[str, Any]],
+ timeout: float = 600.0,
+) -> Tuple[str, int, int, Any]:
+ """Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls).
+
+ Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level
+ ``tool_calls`` object so the RL-mode parser can match against it
+ directly. Otherwise vLLM's tool parser silently swallows the tool call
+ into the SDK field while leaving ``content == ''`` — and the text-tag
+ parser sees nothing, falling through to the answer-1 fallback. (Bug
+ observed 2026-05-19 on the paper-match smoke; same path was buggy on
+ the default pool too, just less reproducibly.)
+ """
+ from openai import OpenAI
+
+ client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
+ messages = [
+ {"role": "system", "content": system},
+ {"role": "user", "content": user},
+ ]
+ resp = client.chat.completions.create(
+ model=model,
+ messages=messages,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ tools=tools,
+ extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+ )
+ choice = resp.choices[0]
+ message = choice.message
+ text = message.content or ""
+ tool_calls = getattr(message, "tool_calls", None)
+ u = resp.usage
+ p = getattr(u, "prompt_tokens", 0) if u else 0
+ c = getattr(u, "completion_tokens", 0) if u else 0
+ return text, p, c, tool_calls
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/experts.py b/src/openjarvis/agents/hybrid/toolorchestra/experts.py
new file mode 100644
index 00000000..a4646a63
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/experts.py
@@ -0,0 +1,170 @@
+"""Slot -> worker expert mapping for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+# Default model used when an `anthropic-web-search` entry omits `model`.
+_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
+
+# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
+# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
+# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
+# tiers since we don't have Qwen-Math served.
+#
+# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
+# so the substitution is deferred until we know the cell's resolved local/cloud
+# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
+
+def _expert_for(slot: str, local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "anthropic") -> Dict[str, Any]:
+ """Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
+
+ Routing policy:
+ - `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off
+ `cloud_endpoint` ("anthropic"/"openai"/"gemini")
+ - `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's
+ cost tier for mid OpenAI calls)
+ - `*-3` (local tier) -> local vLLM (`local_model`)
+ - `answer-math-*` -> same tiers as the numeric suffix
+ - `search-*` -> always the Anthropic web_search tool (the
+ upstream uses Tavily; we have web_search)
+ """
+ if slot.startswith("search"):
+ return {
+ "name": f"search:{slot}",
+ "type": "anthropic-web-search",
+ "model": _DEFAULT_WEB_SEARCH_MODEL,
+ }
+ if slot.endswith("-1") or slot.endswith("-math-1"):
+ ep = (cloud_endpoint or "anthropic").lower()
+ if ep not in ("anthropic", "openai", "gemini"):
+ ep = "anthropic"
+ return {
+ "name": f"frontier:{slot}",
+ "type": ep,
+ "model": cloud_model,
+ }
+ if slot.endswith("-2") or slot.endswith("-math-2"):
+ return {
+ "name": f"mid:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+ # `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
+ # we substitute whatever local model the cell wired up).
+ if local_model and local_endpoint:
+ return {
+ "name": f"local:{slot}",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ }
+ # Fallback if no local — gpt-5-mini.
+ return {
+ "name": f"mid-fallback:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+
+
+# ============================================================================
+# Paper-match expert mapping (2026-05-19).
+# ============================================================================
+# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
+# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
+# models and routes `search` through Tavily, while `enhance_reasoning` is
+# expected to produce code that the caller pipes through a Modal sandbox
+# (handled at dispatch time, not here).
+#
+# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
+# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
+#
+# reasoner-1 -> GPT-5 (frontier reasoner)
+# reasoner-2 -> GPT-5-mini (mid)
+# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
+# answer-1 -> GPT-5
+# answer-2 -> GPT-5-mini
+# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
+# answer-4 -> local Qwen
+# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
+# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
+# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
+# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
+# search-* -> Tavily search (paper)
+#
+# `enhance_reasoning` is dispatched through the coder specialist regardless of
+# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
+# routes the same way in all three cases, then optionally extracts a python
+# code block and execs it in Modal. (We keep the slot-aware routing inside the
+# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
+# pins the coder regardless. See `_run_rl_paper` dispatch.)
+
+_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
+_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
+
+
+def _paper_expert_for(
+ slot: str,
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "openai",
+) -> Dict[str, Any]:
+ """Paper-match counterpart of ``_expert_for``.
+
+ Differs from ``_expert_for``:
+ - Search slots go to ``tavily-search`` (not Anthropic web_search).
+ - Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
+ OpenRouter rather than collapsing onto the local vLLM.
+ - Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
+ as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
+ - ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
+ """
+ if slot.startswith("search"):
+ return {
+ "name": f"tavily:{slot}",
+ "type": "tavily-search",
+ "model": "tavily",
+ }
+ if slot in ("answer-math-1", "answer-math-2"):
+ return {
+ "name": f"math-coder:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_CODER_OPENROUTER,
+ }
+ if slot == "answer-3":
+ return {
+ "name": f"generalist-llama:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
+ }
+ if slot.endswith("-1"):
+ # Tier-1 frontier reasoner / answer — paper uses GPT-5.
+ return {
+ "name": f"frontier:{slot}",
+ "type": "openai",
+ "model": "gpt-5",
+ }
+ if slot.endswith("-2"):
+ return {
+ "name": f"mid:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+ # `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
+ # also serves the local Qwen for the rare local-tier slot).
+ if local_model and local_endpoint:
+ return {
+ "name": f"local:{slot}",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ }
+ return {
+ "name": f"mid-fallback:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/parsing.py b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py
new file mode 100644
index 00000000..d24d2e3c
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py
@@ -0,0 +1,138 @@
+"""Action / tool-call parsing + prompt assembly for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, List, Optional
+
+# Regex for ``{...}`` blocks emitted by Orchestrator-8B
+# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a
+# hermes-style template). Captures the JSON payload.
+_TOOL_CALL_TAG_RE = re.compile(
+ r"\s*(\{.*?\})\s*", re.DOTALL
+)
+
+
+def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
+ """Return ``{"name": str, "arguments": dict}`` or None.
+
+ Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
+ back to scraping ``{...}`` tags from the raw
+ content. We take the first tool call only — Orchestrator-8B was trained
+ to emit exactly one per turn.
+ """
+ # SDK-level path.
+ if sdk_tool_calls:
+ first = sdk_tool_calls[0]
+ name = getattr(getattr(first, "function", None), "name", None)
+ args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
+ try:
+ args = json.loads(args_raw)
+ except json.JSONDecodeError:
+ args = {}
+ if isinstance(name, str) and isinstance(args, dict):
+ return {"name": name, "arguments": args}
+ # Text-tag fallback.
+ if not isinstance(content, str):
+ return None
+ m = _TOOL_CALL_TAG_RE.search(content)
+ if not m:
+ return None
+ try:
+ obj = json.loads(m.group(1))
+ except json.JSONDecodeError:
+ return None
+ if not isinstance(obj, dict):
+ return None
+ name = obj.get("name")
+ args = obj.get("arguments", {})
+ if not isinstance(name, str) or not isinstance(args, dict):
+ return None
+ return {"name": name, "arguments": args}
+
+
+def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
+ return "\n".join(
+ f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
+ )
+
+
+def _build_user_prompt(
+ question: str,
+ workers: List[Dict[str, Any]],
+ history: List[Dict[str, Any]],
+) -> str:
+ pieces = [
+ f"Worker pool:\n{_build_pool_block(workers)}",
+ f"User question:\n{question}",
+ ]
+ if history:
+ pieces.append("Conversation so far (orchestrator turns and worker outputs):")
+ for h in history:
+ if h["role"] == "orchestrator":
+ pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
+ else:
+ pieces.append(
+ f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
+ f"{h['output']}"
+ )
+ pieces.append(
+ "Emit the next JSON action object now — exactly one object, no prose."
+ )
+ return "\n\n".join(pieces)
+
+
+def _strip_fences(s: str) -> str:
+ s = s.strip()
+ if s.startswith("```"):
+ first_nl = s.find("\n")
+ if first_nl != -1:
+ s = s[first_nl + 1:]
+ if s.endswith("```"):
+ s = s[:-3]
+ s = s.strip()
+ return s
+
+
+def _parse_action(text: str) -> Optional[Dict[str, Any]]:
+ s = _strip_fences(text)
+ # First try direct parse, then balanced-brace extraction.
+ try:
+ obj = json.loads(s)
+ if isinstance(obj, dict) and "action" in obj:
+ return obj
+ except json.JSONDecodeError:
+ pass
+ start = s.find("{")
+ if start == -1:
+ return None
+ depth = 0
+ for i in range(start, len(s)):
+ c = s[i]
+ if c == "{":
+ depth += 1
+ elif c == "}":
+ depth -= 1
+ if depth == 0:
+ try:
+ obj = json.loads(s[start : i + 1])
+ if isinstance(obj, dict) and "action" in obj:
+ return obj
+ except json.JSONDecodeError:
+ return None
+ return None
+
+
+def _extract_final_answer_text(text: str) -> str:
+ """Best-effort: pull the answer string from a malformed action emission.
+
+ Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
+ """
+ m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
+ if m:
+ return m.group(1).encode("utf-8").decode("unicode_escape")
+ m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
+ if m:
+ return m.group(1).strip()
+ return text.strip()
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/prompts.py b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py
new file mode 100644
index 00000000..752e22c7
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py
@@ -0,0 +1,106 @@
+"""Prompt strings + tool specs for ToolOrchestraAgent (split from toolorchestra.py)."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+ORCHESTRATOR_SYS = """\
+You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
+
+ {"action": "call_worker", "worker_id": , "input": ""}
+
+ {"action": "final_answer", "answer": ""}
+
+Strategy:
+
+- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math).
+- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass.
+- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase.
+- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question.
+"""
+
+FORCE_FINAL_PROMPT = (
+ "Worker-call budget exhausted. Emit `final_answer` now using everything "
+ "you've learned. Respect the question's answer-format rules."
+)
+
+
+# ============================================================================
+# RL-mode constants (Orchestrator-8B, paper-faithful).
+# ============================================================================
+#
+# Verbatim copies of the upstream system prompt / user-prompt template / tools
+# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
+# edit the description text — Orchestrator-8B was RL-trained against this
+# exact wording and pricing/latency table.
+
+RL_ORCHESTRATOR_SYS = "You are good at using tools."
+
+RL_TOOLS_SPEC: List[Dict[str, Any]] = [
+ {
+ "type": "function",
+ "function": {
+ "name": "enhance_reasoning",
+ "description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
+ "parameters": {
+ "properties": {
+ "model": {
+ "description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
+ "type": "string",
+ }
+ },
+ "required": ["model"],
+ "title": "parameters",
+ "type": "object",
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "answer",
+ "description": "give the final answer. Not allowed to call if documents is empty.",
+ "parameters": {
+ "properties": {
+ "model": {
+ "description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
+ "type": "string",
+ }
+ },
+ "required": ["model"],
+ "title": "parameters",
+ "type": "object",
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "search",
+ "description": "Search for missing information",
+ "parameters": {
+ "properties": {
+ "model": {
+ "description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
+ "type": "string",
+ }
+ },
+ "required": ["model"],
+ "title": "parameters",
+ "type": "object",
+ },
+ },
+ },
+]
+
+# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
+RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
+ "enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
+ "answer": {
+ "model": [
+ "answer-1", "answer-2", "answer-3", "answer-4",
+ "answer-math-1", "answer-math-2",
+ ],
+ },
+ "search": {"model": ["search-1", "search-2", "search-3"]},
+}
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/rollout.py b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py
new file mode 100644
index 00000000..a783f7dc
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py
@@ -0,0 +1,142 @@
+"""Faithful unified-tool rollout loop for ToolOrchestra (arXiv:2511.21689 §2.2).
+
+One reasoning->action->observation loop where the orchestrator picks **a named
+tool** (one per model, from :mod:`expert_registry`) each turn, the environment
+executes it, and the observation is appended to a running context. The rollout
+ends when the orchestrator emits a turn with **no tool call** (its text is the
+final answer) or ``max_turns`` is hit.
+
+The loop is parameterized over two injected callables so it is pure control flow
+(no network) and unit-testable with fakes — the agent supplies real ones:
+
+* ``call_orchestrator(system, user, tool_specs) -> (text, tool_calls, p_tok, c_tok)``
+ where ``tool_calls`` is a list of ``(name, arguments)`` (possibly empty).
+* ``dispatch(tool, arguments) -> (observation, cost_usd, tokens, is_local)``.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, Tuple
+
+from openjarvis.agents.hybrid.expert_registry import (
+ ExpertTool,
+ build_tool_specs,
+ tools_by_name,
+)
+
+RL_ORCHESTRATOR_SYS = "You are good at using tools."
+
+# Char-level cap on the accumulated context (mirrors the paper's ~24k-token cap).
+_CONTEXT_CAP = 24000
+
+
+@dataclass
+class UnifiedTurn:
+ """One orchestrator turn. ``tool_name is None`` marks the final-answer turn."""
+
+ reasoning: str
+ tool_name: Optional[str] = None
+ arguments: Dict[str, object] = field(default_factory=dict)
+ observation: Optional[str] = None
+
+
+@dataclass
+class UnifiedRollout:
+ turns: List[UnifiedTurn]
+ final_answer: str
+ cost_usd: float = 0.0
+ tokens: int = 0
+ num_tool_calls: int = 0
+ parse_failures: int = 0
+
+ def tool_calls(self) -> List[Tuple[str, Dict[str, object]]]:
+ return [(t.tool_name, t.arguments) for t in self.turns if t.tool_name]
+
+
+def _tool_prompt(tool: ExpertTool, arguments: Dict[str, object], question: str) -> str:
+ """The text we actually send the dispatched tool, framed by its arg schema."""
+ for key in ("input", "query", "code"):
+ val = arguments.get(key)
+ if isinstance(val, str) and val.strip():
+ return val
+ return question
+
+
+def run_unified_rollout(
+ question: str,
+ tools: List[ExpertTool],
+ *,
+ call_orchestrator: Callable[..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]],
+ dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
+ max_turns: int = 50,
+ system: str = RL_ORCHESTRATOR_SYS,
+) -> UnifiedRollout:
+ """Drive the faithful unified-tool rollout for one task."""
+ specs = build_tool_specs(tools)
+ by_name = tools_by_name(tools)
+
+ context = ""
+ turns: List[UnifiedTurn] = []
+ cost = 0.0
+ tokens = 0
+ n_tool_calls = 0
+ parse_failures = 0
+ final_answer = ""
+
+ for _ in range(max_turns):
+ user = (
+ f"Problem: {question}\n\n{context or '(no context yet)'}\n\n"
+ "Choose an appropriate tool, or answer directly if you have enough."
+ )
+ text, tool_calls, p_tok, c_tok = call_orchestrator(system, user, specs)
+ tokens += int(p_tok) + int(c_tok)
+
+ if not tool_calls:
+ # No tool call -> the orchestrator is answering. Terminate.
+ final_answer = (text or "").strip()
+ turns.append(UnifiedTurn(reasoning=text or "", tool_name=None))
+ break
+
+ name, arguments = tool_calls[0]
+ if name not in by_name:
+ parse_failures += 1
+ context = (context + f"\n[invalid tool {name!r} — choose from the list]")[-_CONTEXT_CAP:]
+ if parse_failures >= 2:
+ final_answer = (text or "").strip()
+ break
+ continue
+
+ tool = by_name[name]
+ obs, dcost, dtok, _is_local = dispatch(tool, arguments)
+ cost += float(dcost)
+ tokens += int(dtok)
+ n_tool_calls += 1
+ turns.append(UnifiedTurn(
+ reasoning=text or "", tool_name=name, arguments=dict(arguments),
+ observation=obs,
+ ))
+ context = (context + f"\n[{name}] {obs}")[-_CONTEXT_CAP:]
+ else:
+ # Hit max_turns with no explicit answer: use the last observation/text.
+ final_answer = (turns[-1].observation or turns[-1].reasoning).strip() if turns else ""
+
+ return UnifiedRollout(
+ turns=turns, final_answer=final_answer, cost_usd=cost, tokens=tokens,
+ num_tool_calls=n_tool_calls, parse_failures=parse_failures,
+ )
+
+
+def tool_call_tag(name: str, arguments: Dict[str, object]) -> str:
+ """Render a tool call as the ``{...}`` text the model emits."""
+ return f"{json.dumps({'name': name, 'arguments': arguments})}"
+
+
+__all__ = [
+ "RL_ORCHESTRATOR_SYS",
+ "UnifiedRollout",
+ "UnifiedTurn",
+ "run_unified_rollout",
+ "tool_call_tag",
+]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py b/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py
new file mode 100644
index 00000000..30956d14
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py
@@ -0,0 +1,79 @@
+"""Tavily search + Modal Python sandbox helpers for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+import re
+from typing import Optional, Tuple
+
+# ---- Tavily + Modal helpers -------------------------------------------------
+
+def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]:
+ """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0).
+
+ Token counts are reported as zero (no LLM was billed); the OpenJarvis
+ accounting layer separately tallies tool-call counts. Falls back to
+ DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``).
+ """
+ from openjarvis.tools.web_search import WebSearchTool
+
+ tool = WebSearchTool(max_results=max_results)
+ res = tool.execute(query=query, max_results=max_results)
+ text = res.content or ""
+ if not res.success and not text:
+ text = "(no results)"
+ return text, 0, 0
+
+
+_MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox"
+
+
+def _call_modal_python(code: str, timeout_s: int = 60) -> Tuple[str, int]:
+ """Execute a single Python snippet in a fresh Modal Sandbox.
+
+ Returns ``(combined_stdout_stderr, returncode)``. Logs are capped at 8 KiB.
+ Any exception (modal auth, network, sandbox boot failure) is captured into
+ the returned string with a non-zero rc — we never raise back to the
+ orchestrator loop. The sandbox is torn down at the end via ``terminate()``.
+ """
+ try:
+ import modal
+
+ app = modal.App.lookup(_MODAL_APP_NAME, create_if_missing=True)
+ # python:3.12-slim is small + boots fast; the paper uses a generic
+ # Python image too. We rely on stdlib only — no extra pip installs.
+ image = modal.Image.debian_slim(python_version="3.12")
+ sb = modal.Sandbox.create(
+ "python", "-c", code,
+ app=app,
+ image=image,
+ timeout=int(timeout_s),
+ )
+ sb.wait()
+ try:
+ out = sb.stdout.read() or ""
+ except Exception:
+ out = ""
+ try:
+ err = sb.stderr.read() or ""
+ except Exception:
+ err = ""
+ rc = sb.returncode if sb.returncode is not None else -1
+ try:
+ sb.terminate()
+ except Exception:
+ pass
+ combined = out + (("\n" + err) if err else "")
+ if len(combined) > 8192:
+ combined = combined[:8192] + "\n... (output truncated)"
+ return combined, int(rc)
+ except Exception as exc:
+ return f"[modal-python error: {type(exc).__name__}: {exc}]", -1
+
+
+_PY_CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL)
+
+
+def _extract_first_python_block(text: str) -> Optional[str]:
+ """Return the first ```python ... ``` block (or ```...```), or None."""
+ m = _PY_CODE_RE.search(text or "")
+ return m.group(1).strip() if m else None
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/unified.py b/src/openjarvis/agents/hybrid/toolorchestra/unified.py
new file mode 100644
index 00000000..4a861c0f
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/unified.py
@@ -0,0 +1,110 @@
+"""Real backends for the unified-tool rollout — the bridge between the pure
+:func:`run_unified_rollout` loop and live model/tool calls.
+
+``make_call_orchestrator`` returns the ``call_orchestrator`` callable (a teacher
+LLM emitting tool calls over the unified spec); ``make_dispatch`` returns the
+``dispatch`` callable (executes a chosen tool via ``_call_worker``). Both are
+import-safe — the OpenAI SDK is imported lazily so this module loads without
+network or keys.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from openjarvis.agents.hybrid._prices import cost as _model_cost
+from openjarvis.agents.hybrid.expert_registry import ExpertTool, to_worker_dict
+from openjarvis.agents.hybrid.toolorchestra.parsing import _parse_rl_tool_call
+from openjarvis.agents.hybrid.toolorchestra.rollout import (
+ UnifiedRollout,
+ run_unified_rollout,
+)
+from openjarvis.agents.hybrid.toolorchestra.workers import _call_worker
+
+
+def make_call_orchestrator(
+ model: str,
+ *,
+ base_url: Optional[str] = None,
+ api_key: Optional[str] = None,
+ temperature: float = 1.0,
+ max_tokens: int = 4096,
+ timeout: float = 600.0,
+) -> Callable[..., Tuple[str, List[Tuple[str, Dict[str, Any]]], int, int]]:
+ """Teacher-orchestrator caller. ``base_url=None`` → OpenAI cloud; set it to a
+ vLLM endpoint (with ``api_key="EMPTY"``) to drive a local served teacher.
+ """
+
+ def call_orchestrator(system: str, user: str, specs: List[Dict[str, Any]]):
+ from openai import OpenAI
+
+ client = OpenAI(base_url=base_url, api_key=api_key or "EMPTY", timeout=timeout)
+ resp = client.chat.completions.create(
+ model=model,
+ messages=[
+ {"role": "system", "content": system},
+ {"role": "user", "content": user},
+ ],
+ temperature=temperature,
+ max_tokens=max_tokens,
+ tools=specs,
+ )
+ msg = resp.choices[0].message
+ text = msg.content or ""
+ sdk_tool_calls = getattr(msg, "tool_calls", None)
+ u = resp.usage
+ p = getattr(u, "prompt_tokens", 0) if u else 0
+ c = getattr(u, "completion_tokens", 0) if u else 0
+
+ parsed = _parse_rl_tool_call(text, sdk_tool_calls)
+ tool_calls = [(parsed["name"], parsed["arguments"])] if parsed else []
+ return text, tool_calls, int(p), int(c)
+
+ return call_orchestrator
+
+
+def 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 {}
+
+ def dispatch(tool: ExpertTool, arguments: Dict[str, Any]):
+ worker = to_worker_dict(tool)
+ prompt = ""
+ for key in ("input", "query", "code"):
+ val = arguments.get(key)
+ if isinstance(val, str) and val.strip():
+ prompt = val
+ break
+ text, p, c, is_local, extra_cost, _n = _call_worker(worker, prompt, cfg)
+ usd = (0.0 if is_local else _model_cost(str(tool.model), p, c)) + float(extra_cost)
+ return text, usd, int(p) + int(c), bool(is_local)
+
+ return dispatch
+
+
+def teacher_rollout(
+ question: str,
+ tools: List[ExpertTool],
+ *,
+ teacher_model: str,
+ base_url: Optional[str] = None,
+ api_key: Optional[str] = None,
+ temperature: float = 1.0,
+ max_turns: int = 50,
+ cfg: Optional[Dict[str, Any]] = None,
+) -> UnifiedRollout:
+ """Convenience: one full teacher rollout with real backends."""
+ return run_unified_rollout(
+ question,
+ tools,
+ call_orchestrator=make_call_orchestrator(
+ teacher_model, base_url=base_url, api_key=api_key, temperature=temperature,
+ ),
+ dispatch=make_dispatch(cfg),
+ max_turns=max_turns,
+ )
+
+
+__all__ = ["make_call_orchestrator", "make_dispatch", "teacher_rollout"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/workers.py b/src/openjarvis/agents/hybrid/toolorchestra/workers.py
new file mode 100644
index 00000000..a6e227c9
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/workers.py
@@ -0,0 +1,421 @@
+"""Worker pool resolution + dispatch for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from openjarvis.agents.hybrid._base import (
+ ANTHROPIC_WEB_SEARCH_TOOL,
+ WEB_SEARCH_COST_PER_CALL,
+ LocalCloudAgent,
+)
+from openjarvis.agents.hybrid._prices import (
+ PRICES,
+ is_gpt5_family,
+ supports_temperature,
+)
+from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
+from openjarvis.agents.hybrid.toolorchestra.experts import (
+ _PAPER_CODER_OPENROUTER,
+ _PAPER_GENERALIST_TIER3_OPENROUTER,
+)
+from openjarvis.agents.hybrid.toolorchestra.sandbox import (
+ _call_modal_python,
+ _call_tavily_search,
+)
+
+def _paper_pool(
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+) -> List[Dict[str, Any]]:
+ """Paper-match worker pool (registered for traces / inspection).
+
+ NOTE: in RL mode the orchestrator dispatches via tool/slot rather than
+ worker_id, so this list is purely informational — `_paper_expert_for`
+ is the actual routing function. We still return a list here so the
+ paradigm's trace metadata has something concrete to log.
+ """
+ pool: List[Dict[str, Any]] = []
+ if local_model and local_endpoint:
+ pool.append({
+ "id": len(pool),
+ "name": "local-qwen",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ "description": "Local Qwen vLLM (paper uses Qwen3-32B).",
+ })
+ pool.append({
+ "id": len(pool), "name": "tavily-search",
+ "type": "tavily-search", "model": "tavily",
+ "description": "Tavily web search.",
+ })
+ pool.append({
+ "id": len(pool), "name": "modal-python",
+ "type": "modal-python", "model": "modal-python",
+ "description": "Modal Sandbox for one-shot Python exec.",
+ })
+ pool.append({
+ "id": len(pool), "name": "code-specialist",
+ "type": "openrouter", "model": _PAPER_CODER_OPENROUTER,
+ "description": "Qwen-2.5-Coder-32B via OpenRouter (paper).",
+ })
+ pool.append({
+ "id": len(pool), "name": "generalist-llama",
+ "type": "openrouter", "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
+ "description": "Llama-3.3-70B-Instruct via OpenRouter (paper tier-3).",
+ })
+ pool.append({
+ "id": len(pool), "name": "generalist-gpt5",
+ "type": "openai", "model": "gpt-5",
+ "description": "GPT-5 frontier generalist.",
+ })
+ pool.append({
+ "id": len(pool), "name": "generalist-gpt5-mini",
+ "type": "openai", "model": "gpt-5-mini",
+ "description": "GPT-5-mini mid generalist.",
+ })
+ return pool
+
+def _default_pool(
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str = "claude-opus-4-7",
+ cloud_endpoint: str = "anthropic",
+) -> List[Dict[str, Any]]:
+ """Default heterogeneous worker pool.
+
+ The frontier worker's ``type`` + ``model`` track the cell's resolved
+ ``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5,
+ gemini-2.5-pro, …) route their frontier slot to the right SDK.
+ """
+ ep = (cloud_endpoint or "anthropic").lower()
+ if ep not in ("anthropic", "openai", "gemini"):
+ ep = "anthropic"
+ pool: List[Dict[str, Any]] = []
+ if local_model and local_endpoint:
+ pool.append({
+ "id": len(pool),
+ "name": "local-qwen",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ "description": (
+ "Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
+ "concise extraction, formatting, arithmetic on given data."
+ ),
+ })
+ pool.append({
+ "id": len(pool),
+ "name": "web-search",
+ "type": "anthropic-web-search",
+ "model": "claude-haiku-4-5",
+ "description": (
+ "Anthropic server-side web_search. Use for facts that need a lookup "
+ "(recent events, rare names/dates, niche sources). Returns a digest."
+ ),
+ })
+ pool.append({
+ "id": len(pool),
+ "name": f"frontier-{ep}",
+ "type": ep,
+ "model": cloud_model,
+ "description": (
+ "Frontier reasoning model. Use for hard multi-step reasoning, "
+ "code review, or a final synthesis pass. Expensive — use sparingly."
+ ),
+ })
+ pool.append({
+ "id": len(pool),
+ "name": "frontier-openai-mini",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ "description": (
+ "Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
+ "fraction of frontier cost."
+ ),
+ })
+ return pool
+
+
+# Worker types toolorchestra's `_call_worker` actually dispatches.
+#
+# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
+# `tavily-search` — Tavily API search (the paper's web tool).
+# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
+# Used for the code/math specialists and Llama-3.3-70B /
+# Qwen3-32B generalists.
+# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
+# paper's "Python sandbox" inside `enhance_reasoning`).
+_TOOLORCH_VALID_TYPES = (
+ "vllm", "openai", "anthropic", "anthropic-web-search", "gemini",
+ "tavily-search", "openrouter", "modal-python",
+)
+
+# Default model used when an `anthropic-web-search` entry omits `model`.
+_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
+
+
+def _resolve_worker_pool(
+ cfg: Dict[str, Any],
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "anthropic",
+) -> List[Dict[str, Any]]:
+ """Return the worker pool for this run.
+
+ Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
+ default pool is ignored entirely. Falls back to ``_default_pool`` when
+ the override is absent.
+
+ Each user-supplied entry must be a dict with keys ``id``, ``name``,
+ ``type``, and (for non-search types) ``model``. ``type`` must be one
+ of ``vllm`` / ``openai`` / ``anthropic`` / ``anthropic-web-search``.
+ ``anthropic-web-search`` entries may omit ``model`` — it defaults to
+ ``claude-haiku-4-5``.
+
+ Substitution: ``model = "$local"`` (or ``""``) resolves to
+ ``local_model``; ``model = "$cloud"`` / ``""`` to ``cloud_model``.
+
+ On any validation failure, raises ``ValueError`` with the message
+ ``"Invalid worker_pool entry []: "``. Fails fast at agent
+ init rather than mid-task.
+ """
+ override = cfg.get("worker_pool")
+ if override is None:
+ return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint)
+ if not isinstance(override, list) or not override:
+ raise ValueError(
+ "Invalid worker_pool entry [-]: worker_pool must be a non-empty list"
+ )
+
+ resolved: List[Dict[str, Any]] = []
+ seen_ids: set = set()
+ has_non_search = False
+ for raw in override:
+ wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?"
+ if not isinstance(raw, dict):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict"
+ )
+ entry = dict(raw)
+ wid = entry.get("id")
+ if not isinstance(wid, int):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int"
+ )
+ if wid in seen_ids:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: duplicate id"
+ )
+ seen_ids.add(wid)
+ if not entry.get("name") or not isinstance(entry["name"], str):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string"
+ )
+ wtype = entry.get("type") or entry.get("endpoint")
+ if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'type' must be one of "
+ f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})"
+ )
+ wtype = wtype.lower()
+ entry["type"] = wtype
+ # Substitute $local / $cloud placeholders (before any model check).
+ model = entry.get("model")
+ if isinstance(model, str) and model in ("$local", ""):
+ if not local_model:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: model='{model}' "
+ "requires a local_model to be configured for this cell"
+ )
+ model = local_model
+ entry["model"] = model
+ elif isinstance(model, str) and model in ("$cloud", ""):
+ model = cloud_model
+ entry["model"] = model
+ if wtype == "anthropic-web-search":
+ if model in (None, ""):
+ model = _DEFAULT_WEB_SEARCH_MODEL
+ entry["model"] = model
+ elif not isinstance(model, str):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
+ )
+ # Search workers don't satisfy the "needs a solver" requirement.
+ else:
+ if not isinstance(model, str) or not model:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
+ )
+ if wtype == "vllm":
+ if not entry.get("base_url"):
+ if not local_endpoint:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: vllm worker needs "
+ "'base_url' (or a configured local_endpoint to fall back to)"
+ )
+ entry["base_url"] = local_endpoint
+ entry.setdefault("api_key", "EMPTY")
+ else:
+ if model not in PRICES:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: model {model!r} "
+ f"is not in PRICES (known: {sorted(PRICES)})"
+ )
+ has_non_search = True
+ entry.setdefault(
+ "description",
+ f"User-supplied {wtype} worker ({model}).",
+ )
+ resolved.append(entry)
+
+ if not has_non_search:
+ raise ValueError(
+ "Invalid worker_pool entry [-]: worker_pool must contain at least "
+ "one non-search worker (vllm / openai / anthropic)"
+ )
+ return resolved
+
+
+def _call_worker(
+ worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
+) -> Tuple[str, int, int, bool, float, int]:
+ """Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
+ wtype = worker.get("type", "openai")
+ max_tok = int(cfg.get("worker_max_tokens", 4096))
+ temp = float(cfg.get("worker_temperature", 0.2))
+
+ if wtype == "vllm":
+ text, p, c = LocalCloudAgent._call_vllm(
+ worker["model"],
+ worker["base_url"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ enable_thinking=False,
+ )
+ return text, p, c, True, 0.0, 0
+ if wtype == "openai":
+ is_gpt5 = is_gpt5_family(worker["model"])
+ eff_temp = 1.0 if is_gpt5 else temp
+ # GPT-5 is a reasoning model: hidden reasoning tokens count against
+ # `max_completion_tokens`, so a 4096 cap can be fully consumed by
+ # reasoning and leave 0 visible content (empty answer). Give the
+ # reasoning headroom on top of the answer budget.
+ eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
+ text, p, c = LocalCloudAgent._call_openai(
+ worker["model"],
+ user=prompt,
+ max_tokens=eff_max_tok,
+ temperature=eff_temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "gemini":
+ text, p, c = LocalCloudAgent._call_gemini(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "anthropic":
+ eff_temp = temp if supports_temperature(worker["model"]) else 0.0
+ text, p, c, _ = LocalCloudAgent._call_anthropic(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=eff_temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "anthropic-web-search":
+ eff_temp = temp if supports_temperature(worker["model"]) else 0.0
+ text, p, c, n_searches = LocalCloudAgent._call_anthropic(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=eff_temp,
+ tools=[ANTHROPIC_WEB_SEARCH_TOOL],
+ tool_choice={"type": "any"},
+ )
+ extra = n_searches * WEB_SEARCH_COST_PER_CALL
+ return text, p, c, False, extra, n_searches
+ if wtype == "tavily-search":
+ # Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL`
+ # for parity with the Anthropic web-search worker. One call = one
+ # "n_search" for accounting.
+ max_results = int(cfg.get("tavily_max_results", 5))
+ text, p, c = _call_tavily_search(str(prompt), max_results=max_results)
+ return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1
+ if wtype == "openrouter":
+ text, p, c = LocalCloudAgent._call_openrouter(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "modal-python":
+ # `prompt` is the python code string to exec.
+ timeout_s = int(cfg.get("modal_python_timeout_s", 60))
+ out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
+ # No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
+ # charges per sandbox-second — not tracked here.
+ return out, 0, 0, False, 0.0, 0
+ raise ValueError(f"unsupported worker type: {wtype!r}")
+
+
+def _swe_call_worker(
+ worker: Dict[str, Any],
+ prompt: str,
+ cfg: Dict[str, Any],
+ task: Dict[str, Any],
+ workdir: Path,
+ turn: int,
+) -> Tuple[str, int, int, bool, float, int, int]:
+ """SWE-bench worker dispatch: route solver workers through
+ run_swe_agent_loop on a shared workdir. Web-search workers fall back
+ to the regular one-shot dispatch (search isn't an agent loop).
+
+ Trailing ``bash_turns`` (last element) counts agent-loop turns so the
+ caller can surface ``tool_calls`` per row. Fallbacks to one-shot
+ workers return 0 bash turns (no agent loop ran)."""
+ wtype = worker.get("type", "openai")
+ if wtype == "anthropic-web-search":
+ # Search workers stay one-shot.
+ text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
+ return text, p, c, is_local, extra, n_searches, 0
+ if wtype == "vllm":
+ backbone = "local"
+ endpoint = worker.get("base_url")
+ loop_cloud_endpoint = "anthropic" # unused when backbone=local
+ elif wtype in ("anthropic", "openai", "gemini"):
+ backbone = "cloud"
+ endpoint = None
+ loop_cloud_endpoint = wtype
+ else:
+ # Unknown type — one-shot fallback.
+ text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
+ return text, p, c, is_local, extra, n_searches, 0
+ out = run_swe_agent_loop(
+ task,
+ backbone=backbone,
+ backbone_model=worker["model"],
+ cloud_endpoint=loop_cloud_endpoint,
+ local_endpoint=endpoint,
+ initial_prompt=prompt,
+ max_turns=int(cfg.get("swe_max_turns", 30)),
+ bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
+ output_cap=int(cfg.get("swe_output_cap", 10_000)),
+ turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
+ trace_prefix=f"toolorch_turn{turn}",
+ workdir=workdir,
+ )
+ is_local = backbone == "local"
+ return (
+ out["final_summary"] or out["answer"],
+ out["tokens_in"], out["tokens_out"],
+ is_local, 0.0, 0, int(out["turns"]),
+ )
diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py
new file mode 100644
index 00000000..605efecf
--- /dev/null
+++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py
@@ -0,0 +1,119 @@
+"""Rejection-sampling SFT-data generator (the ToolOrchestra cold-start).
+
+For each ToolScale task: roll out a teacher orchestrator N times, verify each
+trajectory, keep the passing ones (optionally just the cheapest), and serialize
+them into the unified-tool ``conversations`` JSONL the SFT trainer consumes.
+
+The expensive/network parts are injected so the orchestration is pure and
+offline-testable:
+
+* ``rollout_fn(task) -> UnifiedRollout`` — one teacher rollout (temperature>0).
+* ``verify_fn(task, rollout) -> bool`` — did the trajectory solve the task?
+
+:func:`gold_coverage_verify` is a dependency-free default verifier (checks the
+trajectory's tool calls cover the task's golden action names); a real run should
+compose it with an LLM judge on the final answer.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections import Counter
+from pathlib import Path
+from typing import Callable, Iterable, List, Optional
+
+from openjarvis.agents.hybrid.expert_registry import ExpertTool
+from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout
+from openjarvis.learning.intelligence.orchestrator.sft_data.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]
+
+
+def gold_coverage_verify(task: ToolScaleTask, rollout: UnifiedRollout) -> bool:
+ """Dependency-free proxy verifier: trajectory must (a) produce a non-empty
+ answer and (b) call tools covering every golden action name.
+
+ This is the offline stand-in for ToolScale's execution-correctness checker
+ (which needs the DB simulator). Compose with an LLM judge for real runs.
+ """
+ if not rollout.final_answer.strip():
+ return False
+ gold = set(task.gold_action_names())
+ if not gold:
+ return True
+ called = {name for name, _ in rollout.tool_calls()}
+ return gold.issubset(called)
+
+
+def generate_sft_dataset(
+ out_path: str,
+ *,
+ tasks: Iterable[ToolScaleTask],
+ tools: List[ExpertTool],
+ rollout_fn: RolloutFn,
+ verify_fn: VerifyFn = gold_coverage_verify,
+ samples_per_task: int = 4,
+ max_keep_per_task: int = 1,
+ reward_fn: Optional[Callable[[UnifiedRollout], float]] = None,
+) -> dict:
+ """Run rejection sampling over ``tasks`` and write the SFT JSONL.
+
+ ``max_keep_per_task`` caps records kept per task; when >1 the cheapest
+ passing trajectories are kept first. Returns stats + writes a ``.stats.json``.
+ """
+ out = Path(out_path)
+ out.parent.mkdir(parents=True, exist_ok=True)
+
+ seen = 0
+ written = 0
+ dropped = 0
+ domain_counts: Counter[str] = Counter()
+
+ with out.open("w") as fh:
+ for task in tasks:
+ seen += 1
+ passing: List[UnifiedRollout] = []
+ for _ in range(samples_per_task):
+ roll = rollout_fn(task)
+ if roll is None:
+ continue
+ if verify_fn(task, roll):
+ passing.append(roll)
+ if not passing:
+ dropped += 1
+ continue
+ # Keep cheapest-first.
+ passing.sort(key=lambda r: r.cost_usd)
+ for roll in passing[:max_keep_per_task]:
+ reward = reward_fn(roll) if reward_fn else 0.0
+ record = trajectory_to_record(
+ task.task_id, task.instruction, tools, roll,
+ reward=reward, domain=task.domain,
+ )
+ fh.write(json.dumps(record) + "\n")
+ written += 1
+ domain_counts[task.domain] += 1
+
+ stats = {
+ "out_path": str(out),
+ "tasks_seen": seen,
+ "records_written": written,
+ "tasks_dropped": dropped,
+ "samples_per_task": samples_per_task,
+ "domain_distribution": dict(domain_counts),
+ }
+ out.with_suffix(out.suffix + ".stats.json").write_text(json.dumps(stats, indent=2))
+ logger.info("Wrote %d SFT records to %s", written, out)
+ return stats
+
+
+__all__ = ["generate_sft_dataset", "gold_coverage_verify"]
diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py
new file mode 100644
index 00000000..6e1aa824
--- /dev/null
+++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py
@@ -0,0 +1,128 @@
+"""Loader for NVIDIA ToolScale (``nvidia/ToolScale``) — the ToolOrchestra
+RL/SFT task source (arXiv:2511.21689 §3.3).
+
+Each row is a synthetic user-agent-tool task: an instruction ``I``, golden
+function calls ``A`` (the ground-truth tool sequence), and short info ``o`` that
+must be communicated. We normalize the raw HF row into :class:`ToolScaleTask`.
+
+``load_toolscale`` streams via the HuggingFace ``datasets`` library; tests pass
+``source=`` an iterable of raw row dicts so normalization is exercised offline.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterable, Iterator, List, Optional
+
+DATASET_ID = "nvidia/ToolScale"
+
+
+@dataclass
+class GoldAction:
+ name: str
+ arguments: Dict[str, Any] = field(default_factory=dict)
+ action_id: Optional[str] = None
+
+
+@dataclass
+class ToolScaleTask:
+ task_id: str
+ domain: str
+ instruction: str
+ gold_actions: List[GoldAction] = field(default_factory=list)
+ required_info: List[str] = field(default_factory=list)
+ nl_assertions: List[str] = field(default_factory=list)
+
+ def gold_action_names(self) -> List[str]:
+ return [a.name for a in self.gold_actions]
+
+
+def _as_list(v: Any) -> List[Any]:
+ if v is None:
+ return []
+ if isinstance(v, list):
+ return v
+ return [v]
+
+
+def _str_list(v: Any) -> List[str]:
+ out: List[str] = []
+ for item in _as_list(v):
+ if isinstance(item, str):
+ out.append(item)
+ elif isinstance(item, dict):
+ # communicate_info entries are sometimes {"info": "..."} dicts.
+ for key in ("info", "content", "text", "value"):
+ if isinstance(item.get(key), str):
+ out.append(item[key])
+ break
+ return out
+
+
+def normalize_row(row: Dict[str, Any], *, index: int = 0) -> ToolScaleTask:
+ """Turn one raw ToolScale row into a :class:`ToolScaleTask` (pure)."""
+ scenario = row.get("user_scenario") or {}
+ instructions = scenario.get("instructions") or {}
+ instruction = (
+ instructions.get("task_instructions")
+ or instructions.get("reason_for_call")
+ or row.get("task")
+ or row.get("instruction")
+ or ""
+ )
+ domain = scenario.get("domain") or row.get("domain") or "unknown"
+
+ crit = row.get("evaluation_criteria") or {}
+ gold: List[GoldAction] = []
+ for a in _as_list(crit.get("actions")):
+ if isinstance(a, dict) and a.get("name"):
+ gold.append(GoldAction(
+ name=str(a["name"]),
+ arguments=a.get("arguments") or a.get("args") or {},
+ action_id=a.get("action_id"),
+ ))
+
+ required = _str_list(crit.get("communicate_info"))
+ nl = _str_list(crit.get("nl_assertions"))
+
+ task_id = str(row.get("id") or row.get("task_id") or f"toolscale-{index}")
+ return ToolScaleTask(
+ task_id=task_id, domain=str(domain), instruction=str(instruction),
+ gold_actions=gold, required_info=required, nl_assertions=nl,
+ )
+
+
+def load_toolscale(
+ *,
+ max_tasks: Optional[int] = None,
+ split: str = "train",
+ source: Optional[Iterable[Dict[str, Any]]] = None,
+) -> Iterator[ToolScaleTask]:
+ """Yield normalized ToolScale tasks.
+
+ ``source`` overrides the HF stream with an iterable of raw row dicts (tests).
+ When ``source`` is None, streams ``nvidia/ToolScale`` via ``datasets``.
+ """
+ if source is None:
+ from datasets import load_dataset # lazy: optional dep / network
+
+ source = load_dataset(DATASET_ID, split=split, streaming=True)
+
+ n = 0
+ for i, row in enumerate(source):
+ if max_tasks is not None and n >= max_tasks:
+ break
+ task = normalize_row(dict(row), index=i)
+ if not task.instruction.strip():
+ continue
+ yield task
+ n += 1
+
+
+__all__ = [
+ "DATASET_ID",
+ "GoldAction",
+ "ToolScaleTask",
+ "load_toolscale",
+ "normalize_row",
+]
diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py
new file mode 100644
index 00000000..6f9761c4
--- /dev/null
+++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py
@@ -0,0 +1,92 @@
+"""Serialize a verified unified-tool rollout into an SFT ``conversations`` record.
+
+Output matches what ``OrchestratorSFTDataset`` consumes, and trains the model to
+emit the ``{...}`` text form that
+``toolorchestra.parsing._parse_rl_tool_call`` already reads back. One record =
+one passing trajectory.
+
+Roles: ``system`` (the unified tool catalog), ``user`` (the running ``Problem``
+prompt), ``assistant`` (reasoning + a ```` tag, or the final answer),
+``tool`` (the executed observation).
+"""
+
+from __future__ import annotations
+
+import json
+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,
+ 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)
+ )
+
+
+def trajectory_to_record(
+ task_id: str,
+ question: str,
+ tools: List[ExpertTool],
+ rollout: UnifiedRollout,
+ *,
+ reward: float = 0.0,
+ domain: str = "unknown",
+) -> Dict[str, Any]:
+ """Convert a passing :class:`UnifiedRollout` into one SFT JSONL record."""
+ conversations: List[Dict[str, str]] = [
+ {"role": "system", "content": _system_prompt(tools)},
+ {"role": "user", "content": f"Problem: {question}\n\nChoose an appropriate tool."},
+ ]
+
+ for turn in rollout.turns:
+ if turn.tool_name is None:
+ # Final-answer turn.
+ conversations.append({
+ "role": "assistant",
+ "content": (turn.reasoning or "").rstrip()
+ + f"\nFINAL_ANSWER: {rollout.final_answer}",
+ })
+ continue
+ tag = tool_call_tag(turn.tool_name, turn.arguments)
+ reasoning = (turn.reasoning or "").rstrip()
+ conversations.append({
+ "role": "assistant",
+ "content": (reasoning + "\n" + tag).strip(),
+ })
+ conversations.append({
+ "role": "tool",
+ "name": turn.tool_name,
+ "content": turn.observation or "",
+ })
+
+ # If the rollout terminated on max_turns (no None turn), append the answer.
+ if not rollout.turns or rollout.turns[-1].tool_name is not None:
+ conversations.append({
+ "role": "assistant",
+ "content": f"FINAL_ANSWER: {rollout.final_answer}",
+ })
+
+ return {
+ "conversations": conversations,
+ "task_id": task_id,
+ "domain": domain,
+ "reward": reward,
+ "metrics": {
+ "cost_usd": rollout.cost_usd,
+ "tokens": rollout.tokens,
+ "num_tool_calls": rollout.num_tool_calls,
+ "num_turns": len(rollout.turns),
+ },
+ }
+
+
+__all__ = ["trajectory_to_record"]
diff --git a/tests/agents/test_expert_registry.py b/tests/agents/test_expert_registry.py
new file mode 100644
index 00000000..c635e0ac
--- /dev/null
+++ b/tests/agents/test_expert_registry.py
@@ -0,0 +1,101 @@
+"""Tests for the faithful ToolOrchestra unified-tool registry."""
+
+from __future__ import annotations
+
+import random
+
+import pytest
+
+from openjarvis.agents.hybrid.expert_registry import (
+ ExpertTool,
+ KIND_MODEL,
+ build_tool_specs,
+ default_catalog,
+ sample_tool_config,
+ to_worker_dict,
+ tools_by_name,
+)
+
+
+def test_each_model_is_its_own_tool():
+ """Faithful §3.1: one named tool per model, not a meta-tool + slot."""
+ cat = default_catalog()
+ names = {t.name for t in cat}
+ # Distinct model tools, each with its own name.
+ for n in ("gpt_5", "gpt_5_mini", "qwen3_32b", "qwen2_5_coder_32b",
+ "llama_3_3_70b", "claude_opus"):
+ assert n in names, f"missing model tool {n}"
+ # No meta-tool / slot vocabulary leaks in.
+ assert "answer" not in names and "enhance_reasoning" not in names
+
+
+def test_catalog_names_unique_and_valid():
+ cat = default_catalog()
+ names = [t.name for t in cat]
+ assert len(names) == len(set(names))
+ assert all(isinstance(t, ExpertTool) for t in cat)
+
+
+def test_local_model_included_only_when_served():
+ assert "local_model" 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"]
+ assert local.backend_type == "vllm"
+ assert local.base_url == "http://x/v1"
+ assert local.price_in == 0.0 and local.price_out == 0.0
+
+
+def test_invalid_tool_rejected():
+ with pytest.raises(ValueError):
+ ExpertTool(name="x", kind="bogus", backend_type="openai", summary="", model="m")
+ with pytest.raises(ValueError):
+ ExpertTool(name="x", kind=KIND_MODEL, backend_type="openai", summary="", model=None)
+
+
+def test_specs_shape_and_pricing_in_description():
+ cat = default_catalog()
+ specs = build_tool_specs(cat)
+ by = {s["function"]["name"]: s for s in specs}
+ gpt5 = by["gpt_5"]
+ assert gpt5["type"] == "function"
+ assert "input" in gpt5["function"]["parameters"]["properties"]
+ # Price table is surfaced in the description (the policy is trained on it).
+ assert "$1.25/1M input" in gpt5["function"]["description"]
+ # Search tool takes a query, code takes code.
+ assert "query" in by["web_search"]["function"]["parameters"]["properties"]
+ assert "code" in by["code_interpreter"]["function"]["parameters"]["properties"]
+
+
+def test_sample_is_deterministic_and_well_formed():
+ cat = default_catalog(local_model="qwen3:8b", local_endpoint="http://x/v1")
+ a = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
+ b = sample_tool_config(cat, rng=random.Random(0), min_tools=4)
+ assert [t.name for t in a] == [t.name for t in b] # deterministic
+ assert len(a) >= 4
+ assert any(t.kind == KIND_MODEL for t in a) # can reason
+ assert any(t.kind != KIND_MODEL for t in a) # can act
+ assert {t.name for t in a} <= {t.name for t in cat} # subset
+
+
+def test_price_jitter_changes_prices_reproducibly():
+ cat = default_catalog()
+ base = {t.name: t for t in sample_tool_config(cat, rng=random.Random(3), min_tools=8)}
+ jit = {t.name: t for t in sample_tool_config(
+ cat, rng=random.Random(3), min_tools=8, price_jitter=0.5)}
+ # Same subset (same seed/sequence up to jitter draws), but model prices move.
+ moved = [n for n in base
+ if base[n].kind == KIND_MODEL and base[n].price_in
+ and n in jit and jit[n].price_in != base[n].price_in]
+ assert moved, "expected jitter to change at least one model price"
+ for n in moved:
+ f = jit[n].price_in / base[n].price_in
+ assert 0.5 <= f <= 1.5
+
+
+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"])
+ assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"
diff --git a/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py
new file mode 100644
index 00000000..e0d35608
--- /dev/null
+++ b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py
@@ -0,0 +1,152 @@
+"""Offline tests for the rejection-sampling SFT pipeline (unified tools)."""
+
+from __future__ import annotations
+
+import json
+
+from openjarvis.agents.hybrid.expert_registry import default_catalog, tools_by_name
+from openjarvis.agents.hybrid.toolorchestra.rollout import (
+ UnifiedRollout,
+ UnifiedTurn,
+ run_unified_rollout,
+)
+from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
+ generate_sft_dataset,
+ gold_coverage_verify,
+)
+from openjarvis.learning.intelligence.orchestrator.sft_data.toolscale import (
+ normalize_row,
+)
+from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import (
+ trajectory_to_record,
+)
+
+# A representative raw ToolScale row.
+_RAW_ROW = {
+ "id": "movie-001",
+ "user_scenario": {
+ "domain": "entertainment",
+ "instructions": {"task_instructions": "Cancel ticket A03 and refund the user."},
+ },
+ "evaluation_criteria": {
+ "actions": [
+ {"name": "cancel", "arguments": {"booking": "A03"}, "action_id": "x1"},
+ {"name": "refund", "arguments": {"user": "8612"}, "action_id": "x2"},
+ ],
+ "communicate_info": ["refund amount is $20.90"],
+ "nl_assertions": ["the ticket is cancelled"],
+ },
+}
+
+
+def test_normalize_row():
+ t = normalize_row(_RAW_ROW)
+ assert t.task_id == "movie-001"
+ assert t.domain == "entertainment"
+ assert "Cancel ticket A03" in t.instruction
+ assert t.gold_action_names() == ["cancel", "refund"]
+ assert t.required_info == ["refund amount is $20.90"]
+
+
+def test_run_unified_rollout_terminates_on_no_tool_call():
+ tools = default_catalog()
+ by = tools_by_name(tools)
+ name = "qwen3_32b"
+ assert name in by
+
+ scripted = [
+ (f"reason 1\n", [(name, {"input": "do step 1"})], 5, 5),
+ ("here is the answer", [], 3, 3), # no tool call -> terminate
+ ]
+ calls = iter(scripted)
+
+ def call_orch(system, user, specs):
+ return next(calls)
+
+ def dispatch(tool, args):
+ return (f"OBS for {tool.name}", 0.01, 10, False)
+
+ roll = run_unified_rollout(
+ "What is X?", tools, call_orchestrator=call_orch, dispatch=dispatch, max_turns=5,
+ )
+ assert roll.final_answer == "here is the answer"
+ assert roll.num_tool_calls == 1
+ assert roll.tool_calls() == [(name, {"input": "do step 1"})]
+ assert abs(roll.cost_usd - 0.01) < 1e-9
+
+
+def test_serialize_record_shape_and_tool_call_tags():
+ tools = default_catalog()
+ roll = UnifiedRollout(
+ turns=[
+ UnifiedTurn(reasoning="think", tool_name="qwen3_32b",
+ arguments={"input": "q"}, observation="obs"),
+ UnifiedTurn(reasoning="done", tool_name=None),
+ ],
+ final_answer="42", cost_usd=0.02, tokens=30, num_tool_calls=1,
+ )
+ rec = trajectory_to_record("t1", "Q?", tools, roll, reward=0.5, domain="math")
+ roles = [m["role"] for m in rec["conversations"]]
+ assert roles[0] == "system" and roles[1] == "user"
+ assert "tool" in roles and roles[-1] == "assistant"
+ # Tool call is emitted as a tag (what the parser reads back).
+ assert any("" in m["content"] and "qwen3_32b" in m["content"]
+ for m in rec["conversations"] if m["role"] == "assistant")
+ assert "FINAL_ANSWER: 42" in rec["conversations"][-1]["content"]
+ assert rec["reward"] == 0.5 and rec["domain"] == "math"
+
+
+def test_gold_coverage_verify():
+ t = normalize_row(_RAW_ROW)
+ good = UnifiedRollout(
+ turns=[
+ UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
+ UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
+ ],
+ final_answer="done",
+ )
+ missing = UnifiedRollout(
+ turns=[UnifiedTurn("", "cancel", {}, "ok")], final_answer="done")
+ empty_ans = UnifiedRollout(
+ turns=[UnifiedTurn("", "cancel", {}, "ok"),
+ UnifiedTurn("", "refund", {}, "ok")], final_answer="")
+ assert gold_coverage_verify(t, good) is True
+ assert gold_coverage_verify(t, missing) is False
+ assert gold_coverage_verify(t, empty_ans) is False
+
+
+def test_generate_sft_dataset_end_to_end(tmp_path):
+ tools = default_catalog()
+ tasks = [normalize_row(_RAW_ROW), normalize_row({
+ **_RAW_ROW, "id": "unsolvable",
+ "evaluation_criteria": {"actions": [{"name": "never_called"}]},
+ })]
+
+ def rollout_fn(task):
+ # Solve the first task; always miss the gold action of the second.
+ if task.task_id == "movie-001":
+ return UnifiedRollout(
+ turns=[
+ UnifiedTurn("", "cancel", {"booking": "A03"}, "ok"),
+ UnifiedTurn("", "refund", {"user": "8612"}, "ok"),
+ UnifiedTurn("done", None),
+ ],
+ final_answer="refunded $20.90", cost_usd=0.03,
+ )
+ return UnifiedRollout(turns=[UnifiedTurn("x", "cancel", {}, "ok")],
+ final_answer="nope", cost_usd=0.05)
+
+ out = tmp_path / "sft.jsonl"
+ stats = generate_sft_dataset(
+ str(out), tasks=tasks, tools=tools, rollout_fn=rollout_fn,
+ samples_per_task=2,
+ )
+ assert stats["tasks_seen"] == 2
+ assert stats["records_written"] == 1 # only the solvable task
+ assert stats["tasks_dropped"] == 1
+ lines = out.read_text().strip().splitlines()
+ assert len(lines) == 1
+ rec = json.loads(lines[0])
+ assert rec["task_id"] == "movie-001"
+ assert rec["domain"] == "entertainment"
+ assert (tmp_path / "sft.jsonl.stats.json").exists()