diff --git a/pyproject.toml b/pyproject.toml
index 6e253ed7..03159a49 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,6 +28,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
+ "braintrust>=0.0.150",
"click>=8",
"datasets>=4.5.0",
"ddgs>=9.11.4",
diff --git a/src/openjarvis/agents/hybrid/_base.py b/src/openjarvis/agents/hybrid/_base.py
index 97f315ae..94317420 100644
--- a/src/openjarvis/agents/hybrid/_base.py
+++ b/src/openjarvis/agents/hybrid/_base.py
@@ -445,8 +445,8 @@ class LocalCloudAgent(BaseAgent):
tools: Optional[list] = None,
tool_choice: Optional[dict] = None,
output_config: Optional[dict] = None,
- timeout: float = 600.0,
- max_retries: int = 12,
+ timeout: float = 60.0,
+ max_retries: int = 2,
trace_role: str = "cloud",
) -> Tuple[str, int, int, int]:
"""Single Anthropic call. Returns (text, p_tok, c_tok, n_web_searches).
@@ -536,7 +536,7 @@ class LocalCloudAgent(BaseAgent):
response_format: Optional[dict] = None,
tools: Optional[list] = None,
tool_choice: Optional[Any] = None,
- timeout: float = 600.0,
+ timeout: float = 60.0,
trace_role: str = "cloud",
) -> Tuple[str, int, int]:
"""Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured;
@@ -604,7 +604,7 @@ class LocalCloudAgent(BaseAgent):
system: Optional[str] = None,
max_tokens: int = 4096,
temperature: float = 0.0,
- timeout: float = 600.0,
+ timeout: float = 60.0,
trace_role: str = "cloud",
extra_body: Optional[Dict[str, Any]] = None,
) -> Tuple[str, int, int]:
diff --git a/src/openjarvis/agents/hybrid/_prices.py b/src/openjarvis/agents/hybrid/_prices.py
index 1c8c546b..64e68787 100644
--- a/src/openjarvis/agents/hybrid/_prices.py
+++ b/src/openjarvis/agents/hybrid/_prices.py
@@ -11,6 +11,7 @@ from __future__ import annotations
# USD per million tokens, (input, output). Local models = 0.
PRICES: dict[str, tuple[float, float]] = {
+ "claude-opus-4-8": (5.00, 25.0),
"claude-opus-4-7": (5.00, 25.0),
"claude-sonnet-4-6": (3.00, 15.0),
"claude-haiku-4-5": (1.00, 5.00),
@@ -33,6 +34,17 @@ PRICES: dict[str, tuple[float, float]] = {
"qwen/qwen-2.5-coder-32b-instruct": (0.08, 0.18),
"qwen/qwen3-32b": (0.10, 0.30),
"meta-llama/llama-3.3-70b-instruct": (0.13, 0.39),
+ # OpenRouter slugs for the orchestrator's local-OSS class when routed via
+ # OpenRouter instead of self-hosted vLLM. OpenRouter list price, checked
+ # 2026-07-11. These replace earlier active-param-scaled guesses that were
+ # 2-5x low on output; the cost-aware reward reads these directly, so the
+ # guesses were systematically undercharging the mid/large Qwen experts.
+ "qwen/qwen3.5-9b": (0.10, 0.15),
+ # qwen3.6-27b is not listed on OpenRouter; qwen3.5-27b is the nearest real
+ # quote and is what we charge for it.
+ "qwen/qwen3.6-27b": (0.195, 1.56),
+ "qwen/qwen3.5-122b-a10b": (0.26, 2.08),
+ "qwen/qwen3.5-397b-a17b": (0.385, 2.45),
}
# Models whose API rejects an explicit `temperature` param — callers should
@@ -41,6 +53,10 @@ NO_TEMP_PREFIXES: tuple[str, ...] = (
"claude-opus-4-7",
"claude-sonnet-4-7",
"claude-haiku-4-7",
+ # 4-8 family also rejects `temperature` ("deprecated for this model").
+ "claude-opus-4-8",
+ "claude-sonnet-4-8",
+ "claude-haiku-4-8",
)
diff --git a/src/openjarvis/agents/hybrid/expert_registry.py b/src/openjarvis/agents/hybrid/expert_registry.py
new file mode 100644
index 00000000..1164408f
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/expert_registry.py
@@ -0,0 +1,751 @@
+"""Faithful ToolOrchestra "unified tool calling" registry (arXiv:2511.21689 §3.1).
+
+The paper exposes **every tool AND every model through a single flat tool
+interface** — each is its own named function with a description and a typed
+parameter schema, and for each training instance a *random subset* of tools is
+sampled with *randomized pricing* (§3.3, "General tool configuration"). This is
+unlike the eval-port shortcut in ``toolorchestra.py``, which collapses the whole
+catalog into three meta-tools (``search``/``enhance_reasoning``/``answer``) with
+a ``model`` slot. This module restores the faithful design.
+
+Each :class:`ExpertTool` knows:
+
+* the orchestrator-visible ``name`` / ``description`` / param schema (what goes
+ into the tools JSON the policy conditions on), and
+* the concrete backend (``backend_type`` + ``model`` + ``base_url``) so a caller
+ can turn it into the worker dict that ``toolorchestra._call_worker`` dispatches.
+
+Everything here is pure data + deterministic transforms (no network, no model
+calls), so the spec building, sampling, and pricing logic is offline-testable.
+Dispatch stays in ``toolorchestra.py`` (via :func:`to_worker_dict`) to avoid a
+circular import.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+from openjarvis.agents.hybrid._prices import PRICES
+
+# Kinds of tool in the unified interface.
+KIND_MODEL = "model" # an LLM exposed as a tool (the paper's "models as tools")
+KIND_WEB_SEARCH = "web_search"
+KIND_CODE = "code_interpreter"
+KIND_TOOL = "tool" # a bridged real OpenJarvis tool (custom param schema)
+
+VALID_KINDS = (KIND_MODEL, KIND_WEB_SEARCH, KIND_CODE, KIND_TOOL)
+
+# Flat-catalog category label, surfaced in each tool's spec so the orchestrator
+# can tell tool types apart without us imposing any hierarchy (the menu stays
+# flat — this is just a tag).
+CATEGORY_BASIC = "basic_tool"
+# Two-model-class taxonomy for the orchestrator catalog — the only model tiers.
+CATEGORY_CLOUD_FRONTIER = "cloud_frontier"
+CATEGORY_LOCAL_OSS = "local_open_source"
+
+# Backend dispatch types understood by ``toolorchestra._call_worker`` (plus the
+# ``openjarvis-tool`` bridge, dispatched in ``unified.make_dispatch`` via the
+# OpenJarvis ToolExecutor rather than ``_call_worker``).
+VALID_BACKENDS = (
+ "vllm",
+ "openai",
+ "anthropic",
+ "gemini",
+ "openrouter",
+ "anthropic-web-search",
+ "tavily-search",
+ "modal-python",
+ "openjarvis-tool",
+)
+
+
+@dataclass(frozen=True)
+class ExpertTool:
+ """One entry in the unified tool catalog.
+
+ ``price_in`` / ``price_out`` are USD per 1M tokens (0.0 for local / non-LLM
+ tools). ``latency_s`` is a rough average used only to populate the
+ description's cost/latency line — the orchestrator was trained to read that
+ table, so we surface it verbatim in the spec.
+ """
+
+ name: str
+ kind: str
+ backend_type: str
+ summary: str
+ model: Optional[str] = None
+ base_url: Optional[str] = None
+ price_in: float = 0.0
+ price_out: float = 0.0
+ latency_s: float = 5.0
+ category: str = "" # cloud_frontier | local_open_source | basic_tool
+ # Optional explicit JSON-schema for the tool's arguments. Set for bridged
+ # real OpenJarvis tools (``openjarvis-tool`` backend) whose params don't fit
+ # the fixed kind-based schemas; takes precedence over the kind default.
+ param_schema: Optional[dict] = None
+ # When True, ``description()`` omits the price/latency line — used by the
+ # anonymized catalog so the orchestrator can't route on cost/identity.
+ hide_cost: bool = False
+
+ def __post_init__(self) -> None:
+ if self.kind not in VALID_KINDS:
+ raise ValueError(f"{self.name}: invalid kind {self.kind!r}")
+ if self.backend_type not in VALID_BACKENDS:
+ raise ValueError(f"{self.name}: invalid backend {self.backend_type!r}")
+ if self.kind == KIND_MODEL and not self.model:
+ raise ValueError(f"{self.name}: model-kind tool needs a concrete model")
+
+ # ---- orchestrator-visible spec -------------------------------------
+
+ def _param_schema(self) -> Dict[str, object]:
+ """JSON-schema for the tool's arguments (one typed param per kind).
+
+ An explicit ``param_schema`` (set by :func:`openjarvis_tool` for bridged
+ real tools) overrides the kind-based default.
+ """
+ if self.param_schema is not None:
+ return self.param_schema
+ if self.kind == KIND_WEB_SEARCH:
+ return {
+ "type": "object",
+ "properties": {
+ "query": {
+ "type": "string",
+ "description": "Search query string.",
+ }
+ },
+ "required": ["query"],
+ }
+ if self.kind == KIND_CODE:
+ return {
+ "type": "object",
+ "properties": {
+ "code": {
+ "type": "string",
+ "description": "Python code to execute. Print results.",
+ }
+ },
+ "required": ["code"],
+ }
+ # model tool
+ return {
+ "type": "object",
+ "properties": {
+ "input": {
+ "type": "string",
+ "description": "The sub-question or instruction for this model.",
+ }
+ },
+ "required": ["input"],
+ }
+
+ def description(self) -> str:
+ """Full description incl. the price/latency line (paper bakes this in)."""
+ if self.hide_cost:
+ return self.summary
+ if self.kind == KIND_MODEL:
+ cost_line = (
+ f" Pricing: ${self.price_in:.2f}/1M input, "
+ f"${self.price_out:.2f}/1M output; avg latency ~{self.latency_s:.0f}s."
+ )
+ else:
+ cost_line = f" Avg latency ~{self.latency_s:.0f}s."
+ return self.summary.rstrip(".") + "." + cost_line
+
+ def to_spec(self) -> Dict[str, object]:
+ """OpenAI-style tool spec the orchestrator conditions on.
+
+ Flat list, but each function carries a ``category`` tag so the policy can
+ distinguish generalist vs specialized models vs basic tools.
+ """
+ fn: Dict[str, object] = {
+ "name": self.name,
+ "description": self.description(),
+ "parameters": self._param_schema(),
+ }
+ if self.category:
+ fn["category"] = self.category
+ return {"type": "function", "function": fn}
+
+
+def _price(model: str) -> tuple[float, float]:
+ return PRICES.get(model, (0.0, 0.0))
+
+
+def _tool_name(model: str) -> str:
+ """Tool-safe function name derived from a model id (``qwen3-8b`` -> ``qwen3_8b``).
+
+ Strips any provider prefix and replaces non-alphanumerics with underscores so
+ the catalog exposes one named tool per concrete model.
+ """
+ base = model.split("/")[-1].lower()
+ safe = re.sub(r"[^a-z0-9]+", "_", base).strip("_")
+ return safe or "local_model"
+
+
+_ANON_ALPHABET = "abcdefghijklmnopqrstuvwxyz0123456789"
+
+# Rough size/scale hint per model — a capability signal WITHOUT the brand name,
+# so the orchestrator can route big-vs-small on task difficulty but can't route on
+# a proprietary name/tier prior. Closed models have no public param count, so they
+# get a scale tier; open models get their param count (else parsed from the name).
+_MODEL_SIZE = {
+ # The real orchestrator_catalog: 2 cloud frontier + 4 local OSS Qwen.
+ "gpt-5.5": "frontier-scale",
+ "claude-opus-4-8": "frontier-scale",
+ "Qwen/Qwen3.5-9B": "~9B params",
+ "Qwen/Qwen3.6-27B-FP8": "~27B params",
+ "Qwen/Qwen3.5-122B-A10B-FP8": "~122B total, ~10B active (MoE)",
+ "Qwen/Qwen3.5-397B-A17B-FP8": "~397B total, ~17B active (MoE)",
+ # OpenRouter slugs (the tool.model id when these route through OpenRouter).
+ "qwen/qwen3.5-122b-a10b": "~122B total, ~10B active (MoE)",
+ "qwen/qwen3.5-397b-a17b": "~397B total, ~17B active (MoE)",
+}
+
+
+def _size_hint(model: str) -> str:
+ """Anonymized scale descriptor for an expert (no brand). Falls back to a
+ param count parsed from the model id (``...-9b`` -> ``~9B params``)."""
+ if model in _MODEL_SIZE:
+ return _MODEL_SIZE[model]
+ m = re.search(r"(\d+(?:\.\d+)?)\s*b\b", model.lower())
+ return f"~{m.group(1)}B params" if m else "unspecified scale"
+
+
+def _class_hint(tool) -> str:
+ """Which of the two model classes this expert belongs to (anonymized)."""
+ cat = getattr(tool, "category", "")
+ if cat == CATEGORY_CLOUD_FRONTIER:
+ return "cloud frontier"
+ if cat == CATEGORY_LOCAL_OSS:
+ return "local open-source"
+ return ""
+
+
+def _cost_tier_hint(tool) -> str:
+ """Coarse cost tier (not exact pricing — that's the bias we anonymize away).
+ Shown so the policy learns the strong-but-expensive vs cheap tradeoff instead
+ of always delegating to the biggest model. Driven by class: cloud frontier is
+ expensive, self-hosted open-source is cheap."""
+ cat = getattr(tool, "category", "")
+ if cat == CATEGORY_CLOUD_FRONTIER:
+ return "expensive"
+ if cat == CATEGORY_LOCAL_OSS:
+ return "cheap"
+ if getattr(tool, "backend_type", "") == "vllm" or (tool.price_out or 0) <= 0:
+ return "cheap"
+ po = tool.price_out
+ return "cheap" if po < 5 else "moderate" if po < 20 else "expensive"
+
+
+def anonymize_tools(tools, rng):
+ """Strip model identity for unbiased routing data.
+
+ Each MODEL expert is replaced with an opaque random label
+ (``expert_<4 rand>``), a uniform description, a uniform category and no
+ price/latency line — so the orchestrator cannot route on a model's name,
+ position, cost, or tier (all of which we found dominate the choice). The full
+ list is shuffled to kill position bias. Basic tools (calculator, web_search,
+ …) keep their real names — the policy must still know what they do.
+
+ Returns ``(anon_tools, anon_to_real)`` where ``anon_to_real`` maps each opaque
+ label back to the real tool name (for offline analysis only; never shown to
+ the model). ``.model`` is preserved on each tool so dispatch still reaches the
+ right backend. Pass a fresh ``rng`` per rollout so labels don't stabilise.
+ """
+ from dataclasses import replace
+
+ anon_to_real: Dict[str, str] = {}
+ experts: List[ExpertTool] = []
+ basics: List[ExpertTool] = []
+ for t in tools:
+ if t.kind == KIND_MODEL:
+ tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
+ while tag in anon_to_real:
+ tag = "model_" + "".join(rng.choice(_ANON_ALPHABET) for _ in range(4))
+ anon_to_real[tag] = t.name
+ bits = [
+ b
+ for b in (_class_hint(t), _size_hint(t.model), _cost_tier_hint(t))
+ if b
+ ]
+ experts.append(
+ replace(
+ t,
+ name=tag,
+ summary="Another model — "
+ + ", ".join(bits)
+ + ". Send it a sub-question.",
+ category="model",
+ hide_cost=True,
+ )
+ )
+ else:
+ basics.append(t)
+ # Shuffle WITHIN the experts to kill per-expert position bias, but keep all
+ # experts as a block on top and the basic tools underneath — a clean split
+ # (experts first) rather than experts and utilities interleaved.
+ rng.shuffle(experts)
+ out = experts + basics
+ return out, anon_to_real
+
+
+def openjarvis_tool(
+ registered_name: str,
+ *,
+ summary: str,
+ params: dict,
+ latency_s: float = 5.0,
+) -> ExpertTool:
+ """Build an :class:`ExpertTool` that bridges a real OpenJarvis tool.
+
+ ``registered_name`` is the tool's key in ``ToolRegistry`` (e.g. ``calculator``,
+ ``shell_exec``). ``params`` is the JSON-schema *properties*-style dict for the
+ tool's arguments; it is surfaced verbatim by :meth:`ExpertTool.to_spec`. The
+ resulting tool dispatches through the OpenJarvis ``ToolExecutor`` (backend
+ ``openjarvis-tool``) rather than ``_call_worker``.
+ """
+ return ExpertTool(
+ name=registered_name,
+ kind=KIND_TOOL,
+ backend_type="openjarvis-tool",
+ summary=summary,
+ model=registered_name,
+ latency_s=latency_s,
+ category=CATEGORY_BASIC,
+ param_schema=params,
+ )
+
+
+# Real OpenJarvis tools bridged into the orchestrator catalog as basic tools.
+# Names must match the ``ToolRegistry`` keys (confirmed present: calculator,
+# shell_exec, file_read, file_write, http_request).
+def _openjarvis_basic_tools() -> List[ExpertTool]:
+ def obj(properties: dict, required: List[str]) -> dict:
+ return {"type": "object", "properties": properties, "required": required}
+
+ return [
+ openjarvis_tool(
+ "calculator",
+ summary="Evaluate an arithmetic / math expression and return the result.",
+ params=obj(
+ {
+ "expression": {
+ "type": "string",
+ "description": "Math expression to evaluate.",
+ }
+ },
+ ["expression"],
+ ),
+ latency_s=1.0,
+ ),
+ openjarvis_tool(
+ "shell_exec",
+ summary=(
+ "Run a shell command and return its stdout/stderr. Critical "
+ "for terminal / TerminalBench-style tasks."
+ ),
+ params=obj(
+ {
+ "command": {
+ "type": "string",
+ "description": "Shell command to execute.",
+ }
+ },
+ ["command"],
+ ),
+ latency_s=4.0,
+ ),
+ openjarvis_tool(
+ "file_read",
+ summary="Read the contents of a file at the given path.",
+ params=obj(
+ {
+ "path": {
+ "type": "string",
+ "description": "Path of the file to read.",
+ }
+ },
+ ["path"],
+ ),
+ latency_s=1.0,
+ ),
+ openjarvis_tool(
+ "file_write",
+ summary="Write content to a file at the given path.",
+ params=obj(
+ {
+ "path": {
+ "type": "string",
+ "description": "Path of the file to write.",
+ },
+ "content": {"type": "string", "description": "Content to write."},
+ },
+ ["path", "content"],
+ ),
+ latency_s=1.0,
+ ),
+ openjarvis_tool(
+ "http_request",
+ summary="Make an HTTP request to a URL and return the response body.",
+ params={
+ "type": "object",
+ "properties": {
+ "url": {"type": "string", "description": "Request URL."},
+ "method": {
+ "type": "string",
+ "description": "HTTP method (GET, POST, ...). Default GET.",
+ },
+ },
+ "required": ["url"],
+ },
+ latency_s=4.0,
+ ),
+ openjarvis_tool(
+ "think",
+ summary=(
+ "Record a private reasoning step (scratchpad). No external "
+ "effect; use to plan before acting on hard reasoning tasks."
+ ),
+ params=obj(
+ {
+ "thought": {
+ "type": "string",
+ "description": "Your reasoning or thought process.",
+ }
+ },
+ ["thought"],
+ ),
+ latency_s=0.5,
+ ),
+ openjarvis_tool(
+ "apply_patch",
+ summary=(
+ "Apply a unified-diff patch to a file. Use to edit code for "
+ "terminal / SWE-style tasks."
+ ),
+ params=obj(
+ {
+ "patch": {
+ "type": "string",
+ "description": "The unified diff patch text to apply.",
+ },
+ "path": {
+ "type": "string",
+ "description": "Target file path (auto-detected from the "
+ "patch header if omitted).",
+ },
+ },
+ ["patch"],
+ ),
+ latency_s=2.0,
+ ),
+ openjarvis_tool(
+ "pdf_extract",
+ summary=(
+ "Extract text from a PDF file. Use for GAIA-style tasks with "
+ "PDF attachments."
+ ),
+ params=obj(
+ {
+ "file_path": {
+ "type": "string",
+ "description": "Path to the PDF file.",
+ },
+ "pages": {
+ "type": "string",
+ "description": "Page range, e.g. '1-5' or '1,3,5'. "
+ "Omit for all pages.",
+ },
+ },
+ ["file_path"],
+ ),
+ latency_s=3.0,
+ ),
+ openjarvis_tool(
+ "db_query",
+ summary=(
+ "Run a SQL query against a SQLite/Postgres database and return "
+ "rows. Read-only by default."
+ ),
+ params=obj(
+ {
+ "query": {"type": "string", "description": "SQL query to execute."},
+ "db_path": {
+ "type": "string",
+ "description": "Path to a SQLite DB file. Defaults to "
+ "in-memory.",
+ },
+ "read_only": {
+ "type": "boolean",
+ "description": "Restrict to SELECT/EXPLAIN/PRAGMA. "
+ "Default: true.",
+ },
+ },
+ ["query"],
+ ),
+ latency_s=3.0,
+ ),
+ ]
+
+
+# Local-Cloud Hybrid orchestrator catalog. The menu is two model classes — cloud
+# frontier models and open-source models — plus the basic tools (web search, code
+# interpreter, and the bridged real OpenJarvis tools). The orchestrator model
+# itself is NOT in the catalog.
+#
+# Routing default is OpenRouter for every model tool, so the catalog works with
+# no self-hosted servers. A model is routed to local vLLM instead when its id is
+# in ``local_endpoints`` or ``model_backends`` overrides it. The
+# orchestrator-visible tool *name* is always derived from the canonical id, so
+# routing can change (vLLM <-> OpenRouter <-> native API) without shifting the
+# menu the policy was trained on.
+_LOCAL_OSS_MODELS = (
+ # (canonical_id, openrouter_slug)
+ ("Qwen/Qwen3.5-9B", "qwen/qwen3.5-9b"),
+ ("Qwen/Qwen3.6-27B-FP8", "qwen/qwen3.6-27b"),
+ ("Qwen/Qwen3.5-122B-A10B-FP8", "qwen/qwen3.5-122b-a10b"),
+ ("Qwen/Qwen3.5-397B-A17B-FP8", "qwen/qwen3.5-397b-a17b"),
+)
+
+_CLOUD_FRONTIER_MODELS = (
+ # (canonical, native_backend, openrouter_slug, summary, latency_s)
+ # Neutral, uniform summaries (no capability ranking) so the orchestrator
+ # doesn't just pick whichever model is labelled "strongest" — routing should
+ # be learned from the reward, not hand-labelled here.
+ ("gpt-5.5", "openai", "openai/gpt-5.5", "Expert model (GPT-5.5).", 30.0),
+ (
+ "claude-opus-4-8",
+ "anthropic",
+ "anthropic/claude-opus-4.8",
+ "Expert model (Claude Opus 4.8).",
+ 26.0,
+ ),
+)
+
+
+def _model_tool(
+ canonical: str,
+ *,
+ native_backend: str,
+ or_slug: str,
+ summary: str,
+ lat: float,
+ category: str,
+ local_endpoints: Dict[str, str],
+ model_backends: Dict[str, str],
+ openrouter_slugs: Dict[str, str],
+) -> ExpertTool:
+ """Build one model tool, resolving its backend.
+
+ Backend precedence: explicit ``model_backends[canonical]`` > vLLM if the model
+ has a ``local_endpoints`` entry > the model's NATIVE provider (openai /
+ anthropic / gemini) when it has one > OpenRouter. Frontier models thus hit
+ their first-party API by default (OpenRouter's gpt-5.5 was returning 400s);
+ OSS Qwen experts (native_backend="vllm") fall through to OpenRouter.
+ """
+ backend = (
+ model_backends.get(canonical)
+ or ("vllm" if canonical in local_endpoints else None)
+ or (
+ native_backend
+ if native_backend in ("openai", "anthropic", "gemini")
+ else "openrouter"
+ )
+ )
+ name = _tool_name(canonical)
+ if backend == "vllm":
+ # Self-hosted: free per the cost model.
+ return ExpertTool(
+ name=name,
+ kind=KIND_MODEL,
+ backend_type="vllm",
+ summary=summary,
+ model=canonical,
+ base_url=local_endpoints.get(canonical),
+ price_in=0.0,
+ price_out=0.0,
+ latency_s=lat,
+ category=category,
+ )
+ if backend == "openrouter":
+ slug = openrouter_slugs.get(canonical, or_slug)
+ pi, po = _price(slug)
+ if (pi, po) == (0.0, 0.0): # fall back to the canonical id's price
+ pi, po = _price(canonical)
+ return ExpertTool(
+ name=name,
+ kind=KIND_MODEL,
+ backend_type="openrouter",
+ summary=summary,
+ model=slug,
+ base_url=None,
+ price_in=pi,
+ price_out=po,
+ latency_s=lat,
+ category=category,
+ )
+ # native provider API (openai / anthropic / gemini)
+ pi, po = _price(canonical)
+ return ExpertTool(
+ name=name,
+ kind=KIND_MODEL,
+ backend_type=backend,
+ summary=summary,
+ model=canonical,
+ price_in=pi,
+ price_out=po,
+ latency_s=lat,
+ category=category,
+ )
+
+
+def orchestrator_catalog(
+ *,
+ local_endpoints: Optional[Dict[str, str]] = None,
+ model_backends: Optional[Dict[str, str]] = None,
+ openrouter_slugs: Optional[Dict[str, str]] = None,
+ include_tools: bool = True,
+) -> List[ExpertTool]:
+ """Return the orchestrator's tool catalog: two model classes + basic tools.
+
+ Routing for the model tools defaults to **OpenRouter** (so the catalog works
+ with no self-hosted servers). Overrides, in precedence order:
+
+ * ``model_backends`` maps a canonical model id -> ``"vllm" | "openrouter" |
+ "openai" | "anthropic" | "gemini"`` to force that model's backend.
+ * ``local_endpoints`` maps a canonical id (e.g. ``"Qwen/Qwen3.5-9B"``) to a
+ vLLM ``base_url``; a model present here routes to vLLM (free) unless
+ ``model_backends`` says otherwise.
+ * ``openrouter_slugs`` overrides the per-model OpenRouter slug used when a
+ model routes through OpenRouter.
+
+ ``include_tools`` (default True) appends the basic tools — web search, code
+ interpreter, and the bridged real OpenJarvis tools (calculator, shell_exec,
+ file_read, file_write, http_request).
+ """
+ local_endpoints = local_endpoints or {}
+ model_backends = model_backends or {}
+ openrouter_slugs = openrouter_slugs or {}
+ cat: List[ExpertTool] = []
+
+ # Env-gated expert exclusion: OJ_EXCLUDE_EXPERTS is a comma-separated list of
+ # case-insensitive substrings matched against a model's canonical id. Any
+ # match is skipped from the catalog. Used to temporarily drop unreliable
+ # experts (e.g. OpenRouter giants during a provider outage) without editing
+ # the registry — unset the var to restore them.
+ _excl = {
+ s.strip().lower()
+ for s in os.environ.get("OJ_EXCLUDE_EXPERTS", "").split(",")
+ if s.strip()
+ }
+
+ def _excluded(canonical: str) -> bool:
+ c = canonical.lower()
+ return any(x in c for x in _excl)
+
+ # ---- cloud frontier models ----
+ for canonical, native_backend, or_slug, summary, lat in _CLOUD_FRONTIER_MODELS:
+ if _excluded(canonical):
+ continue
+ cat.append(
+ _model_tool(
+ canonical,
+ native_backend=native_backend,
+ or_slug=or_slug,
+ summary=summary,
+ lat=lat,
+ category=CATEGORY_CLOUD_FRONTIER,
+ local_endpoints=local_endpoints,
+ model_backends=model_backends,
+ openrouter_slugs=openrouter_slugs,
+ )
+ )
+
+ # ---- open-source models (OpenRouter by default; vLLM when an endpoint or
+ # a model_backends override is supplied) ----
+ for canonical, or_slug in _LOCAL_OSS_MODELS:
+ if _excluded(canonical):
+ continue
+ cat.append(
+ _model_tool(
+ canonical,
+ native_backend="vllm",
+ or_slug=or_slug,
+ summary=f"Expert model ({canonical}).",
+ lat=4.0,
+ category=CATEGORY_LOCAL_OSS,
+ local_endpoints=local_endpoints,
+ model_backends=model_backends,
+ openrouter_slugs=openrouter_slugs,
+ )
+ )
+
+ if include_tools:
+ # ---- basic tools ----
+ cat.append(
+ ExpertTool(
+ name="web_search",
+ kind=KIND_WEB_SEARCH,
+ backend_type="tavily-search",
+ summary="Web search (Tavily). Use for facts that need a live lookup.",
+ model="tavily",
+ latency_s=8.0,
+ category=CATEGORY_BASIC,
+ )
+ )
+ cat.append(
+ ExpertTool(
+ name="code_interpreter",
+ kind=KIND_CODE,
+ backend_type="modal-python",
+ summary="Python sandbox. Execute code and return stdout/stderr.",
+ model="modal-python",
+ latency_s=6.0,
+ category=CATEGORY_BASIC,
+ )
+ )
+ cat.extend(_openjarvis_basic_tools())
+
+ return cat
+
+
+def build_tool_specs(tools: List[ExpertTool]) -> List[Dict[str, object]]:
+ """Turn a tool list into the OpenAI-style tools JSON the policy sees."""
+ return [t.to_spec() for t in tools]
+
+
+def tools_by_name(tools: List[ExpertTool]) -> Dict[str, ExpertTool]:
+ return {t.name: t for t in tools}
+
+
+def to_worker_dict(tool: ExpertTool) -> Dict[str, object]:
+ """Convert a tool into the worker dict ``toolorchestra._call_worker`` eats."""
+ d: Dict[str, object] = {
+ "name": tool.name,
+ "type": tool.backend_type,
+ "model": tool.model,
+ }
+ if tool.base_url:
+ d["base_url"] = tool.base_url
+ return d
+
+
+__all__ = [
+ "CATEGORY_BASIC",
+ "CATEGORY_CLOUD_FRONTIER",
+ "CATEGORY_LOCAL_OSS",
+ "ExpertTool",
+ "KIND_CODE",
+ "KIND_MODEL",
+ "KIND_TOOL",
+ "KIND_WEB_SEARCH",
+ "build_tool_specs",
+ "openjarvis_tool",
+ "orchestrator_catalog",
+ "to_worker_dict",
+ "tools_by_name",
+]
diff --git a/src/openjarvis/agents/hybrid/retry.py b/src/openjarvis/agents/hybrid/retry.py
new file mode 100644
index 00000000..ca20c1a1
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/retry.py
@@ -0,0 +1,98 @@
+"""Exponential backoff + jitter for the cloud calls a rollout makes.
+
+The judge already had this (``evals/core/scorer.py``) — it was added after a run
+where 93/100 GAIA judge calls 429'd and zeroed the whole bench. The rollout's own
+two cloud paths did NOT:
+
+* the ORCHESTRATOR call (``toolorchestra/unified.py``), and
+* the EXPERT dispatch (``expert_registry.py``).
+
+Without backoff a single transient 429 makes the expert return an empty/error
+observation, the clean gate then rejects the whole trajectory ("error or empty
+tool observation"), and the rollout is silently wasted. That's tolerable at
+concurrency 12; at 100 it would shred the batch. Same policy as the judge:
+retry only on transient errors, exponential + jittered, then give up and let the
+caller record the failure.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+import time
+from typing import Callable, TypeVar
+
+LOGGER = logging.getLogger(__name__)
+
+MAX_RETRIES = 6
+BASE_DELAY_S = 2.0
+MAX_DELAY_S = 60.0
+
+# Transient / server-side failures. A 400 (bad request) or 401 (bad key) is NOT
+# here on purpose: retrying those just burns time on a deterministic failure.
+RETRYABLE_MARKERS = (
+ "429",
+ "rate_limit",
+ "rate limit",
+ "overloaded",
+ "timeout",
+ "timed out",
+ "503",
+ "502",
+ "500",
+ "connection",
+ "temporarily unavailable",
+ "internalservererror",
+ "apiconnectionerror",
+ # An expert that returns 200-OK with an EMPTY body. The HTTP call "succeeds",
+ # so nothing raises and no retry fires — the rollout just gets an empty
+ # observation and the clean gate then bins the whole trajectory. Seen from the
+ # OpenRouter-hosted Qwen 122B/397B (audit 2026-07-13: 6 of 72 rollouts). We
+ # raise EmptyExpertResponse for it so it retries like any other transient.
+ "empty expert response",
+)
+
+
+class EmptyExpertResponse(RuntimeError):
+ """An expert returned 200-OK with no content — transient, worth retrying."""
+
+ def __init__(self, model: str) -> None:
+ super().__init__(f"empty expert response from {model}")
+
+T = TypeVar("T")
+
+
+def is_retryable(exc: Exception) -> bool:
+ msg = str(exc).lower()
+ return any(marker in msg for marker in RETRYABLE_MARKERS)
+
+
+def with_backoff(fn: Callable[[], T], *, what: str = "cloud call") -> T:
+ """Run ``fn``, retrying transient failures with exponential backoff + jitter.
+
+ Re-raises the last exception on a non-retryable error or once the retry
+ budget is exhausted, so the caller still sees the failure.
+ """
+ last_exc: Exception | None = None
+ for attempt in range(MAX_RETRIES):
+ try:
+ return fn()
+ except Exception as exc: # noqa: BLE001 — re-raised below
+ last_exc = exc
+ if attempt == MAX_RETRIES - 1 or not is_retryable(exc):
+ raise
+ delay = min(BASE_DELAY_S * (2**attempt), MAX_DELAY_S)
+ delay += random.uniform(0.0, delay * 0.25) # jitter: de-sync the herd
+ LOGGER.warning(
+ "%s failed (attempt %d/%d): %s — retrying in %.1fs",
+ what,
+ attempt + 1,
+ MAX_RETRIES,
+ exc,
+ delay,
+ )
+ time.sleep(delay)
+ raise last_exc # type: ignore[misc]
+
+
+__all__ = ["with_backoff", "is_retryable", "MAX_RETRIES"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py
deleted file mode 100644
index fdb6ebb1..00000000
--- a/src/openjarvis/agents/hybrid/toolorchestra.py
+++ /dev/null
@@ -1,1864 +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 configured provider's server-side
- web-search helper when available.
-
- 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 provider-native web search, the configured frontier
-cloud model, and 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,
- GEMINI_SEARCH_COST_PER_CALL,
- OPENAI_WEB_SEARCH_COST_PER_CALL,
- WEB_SEARCH_COST_PER_CALL,
- LocalCloudAgent,
- tavily_search_context,
-)
-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-*` -> provider-native web search when the cloud
- endpoint supports it; otherwise Anthropic
- """
- if slot.startswith("search"):
- ep = (cloud_endpoint or "anthropic").lower()
- if ep == "openai":
- return {
- "name": f"search:{slot}",
- "type": "openai-web-search",
- "model": cloud_model,
- }
- if ep == "gemini":
- return {
- "name": f"search:{slot}",
- "type": "gemini-web-search",
- "model": cloud_model,
- }
- return {
- "name": f"search:{slot}",
- "type": "anthropic-web-search",
- "model": _DEFAULT_WEB_SEARCH_MODEL,
- }
- if slot.endswith("-1") or slot.endswith("-math-1"):
- ep = (cloud_endpoint or "anthropic").lower()
- if ep not in ("anthropic", "openai", "gemini"):
- ep = "anthropic"
- return {
- "name": f"frontier:{slot}",
- "type": ep,
- "model": cloud_model,
- }
- if slot.endswith("-2") or slot.endswith("-math-2"):
- return {
- "name": f"mid:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
- # `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
- # we substitute whatever local model the cell wired up).
- if local_model and local_endpoint:
- return {
- "name": f"local:{slot}",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- }
- # Fallback if no local — gpt-5-mini.
- return {
- "name": f"mid-fallback:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
-
-
-# ============================================================================
-# Paper-match expert mapping (2026-05-19).
-# ============================================================================
-# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
-# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
-# models and routes `search` through Tavily, while `enhance_reasoning` is
-# expected to produce code that the caller pipes through a Modal sandbox
-# (handled at dispatch time, not here).
-#
-# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
-# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
-#
-# reasoner-1 -> GPT-5 (frontier reasoner)
-# reasoner-2 -> GPT-5-mini (mid)
-# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
-# answer-1 -> GPT-5
-# answer-2 -> GPT-5-mini
-# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
-# answer-4 -> local Qwen
-# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
-# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
-# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
-# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
-# search-* -> Tavily search (paper)
-#
-# `enhance_reasoning` is dispatched through the coder specialist regardless of
-# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
-# routes the same way in all three cases, then optionally extracts a python
-# code block and execs it in Modal. (We keep the slot-aware routing inside the
-# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
-# pins the coder regardless. See `_run_rl_paper` dispatch.)
-
-_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
-_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
-
-
-def _paper_expert_for(
- slot: str,
- local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str,
- cloud_endpoint: str = "openai",
-) -> Dict[str, Any]:
- """Paper-match counterpart of ``_expert_for``.
-
- Differs from ``_expert_for``:
- - Search slots go to ``tavily-search`` (not Anthropic web_search).
- - Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
- OpenRouter rather than collapsing onto the local vLLM.
- - Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
- as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
- - ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
- """
- if slot.startswith("search"):
- return {
- "name": f"tavily:{slot}",
- "type": "tavily-search",
- "model": "tavily",
- }
- if slot in ("answer-math-1", "answer-math-2"):
- return {
- "name": f"math-coder:{slot}",
- "type": "openrouter",
- "model": _PAPER_CODER_OPENROUTER,
- }
- if slot == "answer-3":
- return {
- "name": f"generalist-llama:{slot}",
- "type": "openrouter",
- "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
- }
- if slot.endswith("-1"):
- # Tier-1 frontier reasoner / answer — paper uses GPT-5.
- return {
- "name": f"frontier:{slot}",
- "type": "openai",
- "model": "gpt-5",
- }
- if slot.endswith("-2"):
- return {
- "name": f"mid:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
- # `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
- # also serves the local Qwen for the rare local-tier slot).
- if local_model and local_endpoint:
- return {
- "name": f"local:{slot}",
- "type": "vllm",
- "model": local_model,
- "base_url": local_endpoint,
- }
- return {
- "name": f"mid-fallback:{slot}",
- "type": "openai",
- "model": "gpt-5-mini",
- }
-
-
-# ---- Tavily + Modal helpers -------------------------------------------------
-
-
-def _call_tavily_search(
- query: str,
- max_results: int = 5,
-) -> Tuple[str, int, int, float, int]:
- """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0, cost, uses).
-
- 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``).
- """
- res = tavily_search_context(query, max_results=max_results)
- return res["text"], 0, 0, float(res["cost_usd"]), int(res["n_searches"])
-
-
-_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."
- ),
- }
- )
- if ep == "openai":
- search_type = "openai-web-search"
- search_model = cloud_model
- search_desc = "OpenAI hosted web search on the configured frontier model."
- elif ep == "gemini":
- search_type = "gemini-web-search"
- search_model = cloud_model
- search_desc = "Gemini Google Search grounding on the configured frontier model."
- else:
- search_type = "anthropic-web-search"
- search_model = _DEFAULT_WEB_SEARCH_MODEL
- search_desc = "Anthropic server-side web_search."
- pool.append(
- {
- "id": len(pool),
- "name": "web-search",
- "type": search_type,
- "model": search_model,
- "description": (
- f"{search_desc} Use for facts that need a lookup "
- "(recent events, rare names/dates, niche sources). Returns a digest."
- ),
- }
- )
- pool.append(
- {
- "id": len(pool),
- "name": f"frontier-{ep}",
- "type": ep,
- "model": cloud_model,
- "description": (
- "Frontier reasoning model. Use for hard multi-step reasoning, "
- "code review, or a final synthesis pass. Expensive — use sparingly."
- ),
- }
- )
- pool.append(
- {
- "id": len(pool),
- "name": "frontier-openai-mini",
- "type": "openai",
- "model": "gpt-5-mini",
- "description": (
- "Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
- "fraction of frontier cost."
- ),
- }
- )
- return pool
-
-
-# Worker types toolorchestra's `_call_worker` actually dispatches.
-#
-# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
-# `tavily-search` — Tavily API search (the paper's web tool).
-# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
-# Used for the code/math specialists and Llama-3.3-70B /
-# Qwen3-32B generalists.
-# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
-# paper's "Python sandbox" inside `enhance_reasoning`).
-_TOOLORCH_VALID_TYPES = (
- "vllm",
- "openai",
- "anthropic",
- "anthropic-web-search",
- "openai-web-search",
- "gemini",
- "gemini-web-search",
- "tavily-search",
- "openrouter",
- "modal-python",
-)
-_TOOLORCH_SEARCH_TYPES = (
- "anthropic-web-search",
- "openai-web-search",
- "gemini-web-search",
- "tavily-search",
-)
-
-# Default model used when an `anthropic-web-search` entry omits `model`.
-_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
-
-
-def _resolve_worker_pool(
- cfg: Dict[str, Any],
- local_model: Optional[str],
- local_endpoint: Optional[str],
- cloud_model: str,
- cloud_endpoint: str = "anthropic",
-) -> List[Dict[str, Any]]:
- """Return the worker pool for this run.
-
- Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
- default pool is ignored entirely. Falls back to ``_default_pool`` when
- the override is absent.
-
- Each user-supplied entry must be a dict with keys ``id``, ``name``,
- ``type``, and (for non-search types) ``model``. Search worker types are
- ``anthropic-web-search``, ``openai-web-search``, ``gemini-web-search``,
- and ``tavily-search``. ``anthropic-web-search`` entries may omit
- ``model`` — it defaults to ``claude-haiku-4-5``. OpenAI and Gemini
- search workers default to the configured cloud model. Tavily does not
- require a model.
-
- Substitution: ``model = "$local"`` (or ``""``) 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 in _TOOLORCH_SEARCH_TYPES:
- if model in (None, ""):
- if wtype == "anthropic-web-search":
- model = _DEFAULT_WEB_SEARCH_MODEL
- elif wtype in ("openai-web-search", "gemini-web-search"):
- model = cloud_model
- else:
- model = wtype
- entry["model"] = model
- elif not isinstance(model, str):
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
- )
- if (
- wtype in ("openai-web-search", "gemini-web-search")
- and model not in PRICES
- ):
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: model {model!r} "
- f"is not in PRICES (known: {sorted(PRICES)})"
- )
- # Search workers don't satisfy the "needs a solver" requirement.
- else:
- if not isinstance(model, str) or not model:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
- )
- if wtype == "vllm":
- if not entry.get("base_url"):
- if not local_endpoint:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: vllm worker needs "
- "'base_url' (or a configured local_endpoint to fall back to)"
- )
- entry["base_url"] = local_endpoint
- entry.setdefault("api_key", "EMPTY")
- else:
- if model not in PRICES:
- raise ValueError(
- f"Invalid worker_pool entry [{wid}]: model {model!r} "
- f"is not in PRICES (known: {sorted(PRICES)})"
- )
- has_non_search = True
- entry.setdefault(
- "description",
- f"User-supplied {wtype} worker ({model}).",
- )
- resolved.append(entry)
-
- if not has_non_search:
- raise ValueError(
- "Invalid worker_pool entry [-]: worker_pool must contain at least "
- "one non-search worker (vllm / openai / anthropic / gemini)"
- )
- return resolved
-
-
-def _call_worker(
- worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
-) -> Tuple[str, int, int, bool, float, int]:
- """Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
- wtype = worker.get("type", "openai")
- max_tok = int(cfg.get("worker_max_tokens", 4096))
- temp = float(cfg.get("worker_temperature", 0.2))
-
- if wtype == "vllm":
- text, p, c = LocalCloudAgent._call_vllm(
- worker["model"],
- worker["base_url"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- enable_thinking=False,
- )
- return text, p, c, True, 0.0, 0
- if wtype == "openai":
- is_gpt5 = is_gpt5_family(worker["model"])
- eff_temp = 1.0 if is_gpt5 else temp
- # GPT-5 is a reasoning model: hidden reasoning tokens count against
- # `max_completion_tokens`, so a 4096 cap can be fully consumed by
- # reasoning and leave 0 visible content (empty answer). Give the
- # reasoning headroom on top of the answer budget.
- eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
- text, p, c = LocalCloudAgent._call_openai(
- worker["model"],
- user=prompt,
- max_tokens=eff_max_tok,
- temperature=eff_temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "gemini":
- text, p, c = LocalCloudAgent._call_gemini(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "anthropic":
- eff_temp = temp if supports_temperature(worker["model"]) else 0.0
- text, p, c, _ = LocalCloudAgent._call_anthropic(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=eff_temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "anthropic-web-search":
- eff_temp = temp if supports_temperature(worker["model"]) else 0.0
- text, p, c, n_searches = LocalCloudAgent._call_anthropic(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=eff_temp,
- tools=[ANTHROPIC_WEB_SEARCH_TOOL],
- tool_choice={"type": "any"},
- )
- extra = n_searches * WEB_SEARCH_COST_PER_CALL
- return text, p, c, False, extra, n_searches
- if wtype == "openai-web-search":
- eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp
- text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent(
- worker["model"],
- user=prompt,
- max_tokens=max(max_tok, 16384)
- if is_gpt5_family(worker["model"])
- else max_tok,
- temperature=eff_temp,
- )
- extra = n_searches * OPENAI_WEB_SEARCH_COST_PER_CALL
- return text, p, c, False, extra, n_searches
- if wtype == "gemini-web-search":
- text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- )
- extra = n_searches * GEMINI_SEARCH_COST_PER_CALL
- return text, p, c, False, extra, n_searches
- if wtype == "tavily-search":
- max_results = int(cfg.get("tavily_max_results", 5))
- text, p, c, extra, n_searches = _call_tavily_search(
- str(prompt),
- max_results=max_results,
- )
- return text, p, c, False, extra, n_searches
- if wtype == "openrouter":
- text, p, c = LocalCloudAgent._call_openrouter(
- worker["model"],
- user=prompt,
- max_tokens=max_tok,
- temperature=temp,
- )
- return text, p, c, False, 0.0, 0
- if wtype == "modal-python":
- # `prompt` is the python code string to exec.
- timeout_s = int(cfg.get("modal_python_timeout_s", 60))
- out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
- # No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
- # charges per sandbox-second — not tracked here.
- return out, 0, 0, False, 0.0, 0
- raise ValueError(f"unsupported worker type: {wtype!r}")
-
-
-def _swe_call_worker(
- worker: Dict[str, Any],
- prompt: str,
- cfg: Dict[str, Any],
- task: Dict[str, Any],
- workdir: Path,
- turn: int,
-) -> Tuple[str, int, int, bool, float, int, int]:
- """SWE-bench worker dispatch: route solver workers through
- run_swe_agent_loop on a shared workdir. Web-search workers fall back
- to the regular one-shot dispatch (search isn't an agent loop).
-
- Trailing ``bash_turns`` (last element) counts agent-loop turns so the
- caller can surface ``tool_calls`` per row. Fallbacks to one-shot
- workers return 0 bash turns (no agent loop ran)."""
- wtype = worker.get("type", "openai")
- if wtype in _TOOLORCH_SEARCH_TYPES:
- # Search workers stay one-shot.
- text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
- return text, p, c, is_local, extra, n_searches, 0
- if wtype == "vllm":
- backbone = "local"
- endpoint = worker.get("base_url")
- loop_cloud_endpoint = "anthropic" # unused when backbone=local
- elif wtype in ("anthropic", "openai", "gemini"):
- backbone = "cloud"
- endpoint = None
- loop_cloud_endpoint = wtype
- else:
- # Unknown type — one-shot fallback.
- text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
- return text, p, c, is_local, extra, n_searches, 0
- out = run_swe_agent_loop(
- task,
- backbone=backbone,
- backbone_model=worker["model"],
- cloud_endpoint=loop_cloud_endpoint,
- local_endpoint=endpoint,
- initial_prompt=prompt,
- max_turns=int(cfg.get("swe_max_turns", 30)),
- bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
- output_cap=int(cfg.get("swe_output_cap", 10_000)),
- turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
- trace_prefix=f"toolorch_turn{turn}",
- workdir=workdir,
- )
- is_local = backbone == "local"
- return (
- out["final_summary"] or out["answer"],
- out["tokens_in"],
- out["tokens_out"],
- is_local,
- 0.0,
- 0,
- int(out["turns"]),
- )
-
-
-@AgentRegistry.register("toolorchestra")
-class ToolOrchestraAgent(LocalCloudAgent):
- """Multi-turn dispatcher over a mixed worker pool.
-
- Two modes (see module docstring): ``method_cfg.orchestrator_mode``
- is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
- (paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
- """
-
- agent_id = "toolorchestra"
-
- def __init__(self, *args: Any, **kwargs: Any) -> None:
- super().__init__(*args, **kwargs)
- # Validate `method_cfg.worker_pool` early — surfaces config errors
- # at agent construction rather than on the first task. No-op when
- # the override is absent.
- if self._cfg.get("worker_pool") is not None:
- _resolve_worker_pool(
- self._cfg,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
- # Validate `orchestrator_mode` (typo-checked here, not on first task).
- mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
- if mode not in ("prompted", "rl"):
- raise ValueError(
- f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
- f"got {mode!r}"
- )
-
- def _run_paradigm(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
- if mode == "rl":
- return self._run_rl(input, context, **kwargs)
- return self._run_prompted(input, context, **kwargs)
-
- # ------------------------------------------------------------------
- # Legacy prompted-orchestrator path.
- # ------------------------------------------------------------------
- def _run_prompted(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- cfg = self._cfg
- question = input
- # Resolution order (strict replace, no merge):
- # 1. `cfg["workers"]` — legacy direct override, used by tests.
- # 2. `cfg["worker_pool"]` — cell-config override; validated +
- # $local/$cloud substituted.
- # 3. `_default_pool(...)` — heterogeneous default.
- if cfg.get("workers"):
- workers = cfg["workers"]
- else:
- workers = _resolve_worker_pool(
- cfg,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
- if not workers:
- raise RuntimeError("toolorchestra: empty worker pool")
-
- max_turns = int(cfg.get("max_turns", 6))
- orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
-
- task_meta = (context.metadata.get("task") if context is not None else {}) or {}
- swe_mode = (
- bool(cfg.get("swe_use_agent_loop"))
- and bool(task_meta.get("problem_statement"))
- and bool(task_meta.get("repo"))
- and bool(task_meta.get("base_commit"))
- )
- shared_workdir: Optional[Path] = None
- if swe_mode:
- shared_workdir = Path(
- tempfile.mkdtemp(
- prefix=f"toolorch-swe-{task_meta.get('task_id', 'x')}-"
- )
- )
- try:
- _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
- except Exception:
- shutil.rmtree(shared_workdir, ignore_errors=True)
- raise
- self.record_trace_event(
- {
- "kind": "toolorchestra_swe_workdir",
- "workdir": str(shared_workdir),
- "repo": task_meta["repo"],
- "base_commit": task_meta["base_commit"],
- }
- )
-
- # try/finally guards ``shared_workdir`` against exceptions raised
- # anywhere in the turn loop, the worker calls, the fallback, or
- # the diff-extraction step. Without this, at n=500 SWE-bench an
- # exception leaves hundreds of MB of cloned repos in tempdir.
- try:
- history: List[Dict[str, Any]] = []
- tokens_local = 0
- tokens_cloud = 0
- cost = 0.0
- n_web_searches_total = 0
- # tool_calls: bash turns from SWE subloops + web_search uses
- # from GAIA. Orchestrator dispatch turns are NOT counted (they
- # produce text only — calling a worker is one tool call's worth
- # of "delegation" but the actual tool action happens inside).
- tool_calls = 0
- final_answer: Optional[str] = None
- forced_final = False
- parse_failures = 0
-
- for turn in range(1, max_turns + 1):
- sys_prompt = ORCHESTRATOR_SYS
- if turn == max_turns and final_answer is None:
- sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
- forced_final = True
-
- user = _build_user_prompt(question, workers, history)
- text, o_in, o_out = self._call_cloud(
- user=user,
- system=sys_prompt,
- max_tokens=orch_max_tokens,
- temperature=0.0,
- )
- tokens_cloud += o_in + o_out
- cost += self.cost_usd(self._cloud_model, o_in, o_out)
-
- action = _parse_action(text)
- history.append(
- {
- "role": "orchestrator",
- "turn": turn,
- "raw": text,
- "action": action,
- }
- )
- self.record_trace_event(
- {
- "kind": "toolorchestra_action",
- "turn": turn,
- "action": action,
- "raw": text,
- }
- )
-
- if action is None:
- parse_failures += 1
- if parse_failures >= 2 or forced_final:
- final_answer = _extract_final_answer_text(text)
- break
- continue
-
- kind = action.get("action")
- if kind == "final_answer":
- final_answer = str(action.get("answer", "")).strip()
- break
- if kind == "call_worker":
- wid = action.get("worker_id")
- w_input = action.get("input", "")
- if not isinstance(wid, int) or not (0 <= wid < len(workers)):
- parse_failures += 1
- if parse_failures >= 2 or forced_final:
- final_answer = _extract_final_answer_text(text)
- break
- continue
- worker = workers[wid]
- if swe_mode and shared_workdir is not None:
- (
- w_text,
- w_in,
- w_out,
- is_local,
- extra_cost,
- n_searches,
- bash_turns,
- ) = _swe_call_worker(
- worker,
- str(w_input),
- cfg,
- task_meta,
- shared_workdir,
- turn,
- )
- tool_calls += bash_turns
- else:
- w_text, w_in, w_out, is_local, extra_cost, n_searches = (
- _call_worker(worker, str(w_input), cfg)
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- n_web_searches_total += n_searches
- tool_calls += n_searches
- history.append(
- {
- "role": "worker",
- "turn": turn,
- "worker_id": wid,
- "worker_name": worker["name"],
- "worker_model": worker["model"],
- "output": w_text,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "n_web_searches": n_searches,
- }
- )
- continue
- # Unknown action kind — treat as parse failure.
- parse_failures += 1
-
- if final_answer is None:
- # Hard fallback: call the strongest non-search worker directly.
- # "Strongest" = highest output-token price in `_prices.PRICES`,
- # which tracks model capability tier closely enough for this.
- # Search workers are excluded — they answer fact-lookup
- # questions, not synthesis.
- non_search = [
- w for w in workers if w.get("type") not in _TOOLORCH_SEARCH_TYPES
- ] or workers
- worker = max(
- non_search,
- key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
- )
- if swe_mode and shared_workdir is not None:
- (ans, w_in, w_out, is_local, extra_cost, _, bash_turns) = (
- _swe_call_worker(
- worker,
- question,
- cfg,
- task_meta,
- shared_workdir,
- max_turns + 1,
- )
- )
- tool_calls += bash_turns
- else:
- ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
- worker, question, cfg
- )
- if is_local:
- tokens_local += w_in + w_out
- else:
- tokens_cloud += w_in + w_out
- cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
- history.append(
- {
- "role": "worker",
- "turn": max_turns + 1,
- "worker_id": worker["id"],
- "worker_name": worker["name"],
- "worker_model": worker["model"],
- "output": ans,
- "tokens_in": w_in,
- "tokens_out": w_out,
- "fallback": True,
- }
- )
- final_answer = ans
-
- # In SWE mode, the authoritative output is the working-tree diff —
- # frame it (the runner extracts it via the scorer's ```diff fence).
- if swe_mode and shared_workdir is not None:
- patch = _extract_diff(shared_workdir)
- if patch.strip():
- final_answer = (
- f"{final_answer}\n\n```diff\n{patch}```"
- if final_answer
- else f"```diff\n{patch}```"
- )
-
- meta = {
- "tokens_local": tokens_local,
- "tokens_cloud": tokens_cloud,
- "cost_usd": cost,
- "turns": len([h for h in history if h["role"] == "orchestrator"]),
- "web_search_uses": n_web_searches_total,
- "tool_calls": int(tool_calls),
- "traces": {
- "history": history,
- "forced_final": forced_final,
- "parse_failures": parse_failures,
- "workers": workers,
- "n_web_searches": n_web_searches_total,
- "note": (
- "inference-only port; the RL-trained Nemotron-Orchestrator-8B "
- "is NOT in the loop. Results are preliminary."
- ),
- },
- }
- return final_answer, meta
- finally:
- if shared_workdir is not None:
- shutil.rmtree(shared_workdir, ignore_errors=True)
-
- # ------------------------------------------------------------------
- # Paper-faithful Orchestrator-8B path.
- # ------------------------------------------------------------------
- def _run_rl(
- self,
- input: str,
- context: Optional[AgentContext],
- **kwargs: Any,
- ) -> Tuple[str, Dict[str, Any]]:
- cfg = self._cfg
- question = input
-
- # Orchestrator endpoint / model (where the RL'd 8B lives).
- orch_endpoint = str(
- cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
- )
- orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
- max_turns = int(cfg.get("max_turns", 8))
- orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
- orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
-
- # Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
- # replaces `_expert_for` and `enhance_reasoning` is post-processed
- # through a Modal Python sandbox. See module docstring + paper-match
- # doc at `docs/26.5.19/toolorchestra-papermatch.md`.
- paper_mode = str(cfg.get("pool", "")).lower() == "paper"
-
- # SWE-bench detection: same gate as the prompted path. Requires
- # `method_cfg.swe_use_agent_loop = true` AND the task carries the
- # SWE-bench fields. When active, the `enhance_reasoning` and
- # `answer` workers route through `run_swe_agent_loop` on a shared
- # workdir; search workers stay one-shot. At end, the working-tree
- # diff is appended to final_answer so `_score_swebench` can extract
- # it via the ```diff fence.
- task_meta = (context.metadata.get("task") if context is not None else {}) or {}
- swe_mode = (
- bool(cfg.get("swe_use_agent_loop"))
- and bool(task_meta.get("problem_statement"))
- and bool(task_meta.get("repo"))
- and bool(task_meta.get("base_commit"))
- )
- shared_workdir: Optional[Path] = None
- if swe_mode:
- shared_workdir = Path(
- tempfile.mkdtemp(
- prefix=f"toolorch-rl-swe-{task_meta.get('task_id', 'x')}-"
- )
- )
- try:
- _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
- except Exception:
- shutil.rmtree(shared_workdir, ignore_errors=True)
- raise
- self.record_trace_event(
- {
- "kind": "toolorchestra_rl_swe_workdir",
- "workdir": str(shared_workdir),
- "repo": task_meta["repo"],
- "base_commit": task_meta["base_commit"],
- }
- )
-
- # ``context_str`` mirrors the upstream's running context — accumulates
- # search documents and code/exec snippets across turns. We keep this
- # as a single string for prompt simplicity; the upstream uses
- # tokenized cutoffs (we cap at ~24k chars instead).
- context_str = ""
- doc_list: List[str] = []
- history: List[Dict[str, Any]] = []
- tokens_local = 0
- tokens_cloud = 0
- cost = 0.0
- n_web_searches_total = 0
- tool_calls = 0
- final_answer: Optional[str] = None
- parse_failures = 0
-
- # Single outer try/finally guards `shared_workdir` against any
- # exception in the orchestrator loop, the post-loop fallback, or
- # the diff-extraction step. Matches the prompted path's pattern.
- try:
- for turn in range(1, max_turns + 1):
- user = (
- f"Problem: {question}\n\n{context_str}\n\n"
- "Choose an appropriate tool."
- )
-
- # Orchestrator-8B served on local vLLM. We pass the three NVlabs
- # tools verbatim. In paper-mode we use the local helper so we
- # get the SDK-level ``tool_calls`` object back — `_call_vllm`
- # returns just text and loses the call when vLLM's parser
- # caught it. Orchestrator-8B emits its routing decision in the
- # OpenAI-native ``tool_calls`` array with an empty text body,
- # so the legacy `_call_vllm` path saw nothing and silently fell
- # through to the answer-1 fallback (parse_failures: 2 on every
- # non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
- # modes now use `_call_orchestrator_with_tool_calls` so the
- # parser can read structured tool calls; the text-tag path in
- # `_parse_rl_tool_call` is still the fallback when `tool_calls`
- # is empty.
- text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
- orch_model,
- orch_endpoint,
- user=user,
- system=RL_ORCHESTRATOR_SYS,
- max_tokens=orch_max_tokens,
- temperature=orch_temp,
- tools=RL_TOOLS_SPEC,
- )
- self.record_trace_event(
- {
- "kind": "vllm",
- "role": "orchestrator",
- "model": orch_model,
- "endpoint": orch_endpoint,
- "system": RL_ORCHESTRATOR_SYS,
- "user": user,
- "response": text,
- "tool_calls": [
- {
- "id": getattr(tc, "id", None),
- "type": getattr(tc, "type", None),
- "function": {
- "name": getattr(
- getattr(tc, "function", None), "name", None
- ),
- "arguments": getattr(
- getattr(tc, "function", None), "arguments", None
- ),
- },
- }
- for tc in (sdk_tool_calls or [])
- ],
- "tokens_in": o_in,
- "tokens_out": o_out,
- }
- )
- tokens_local += o_in + o_out
-
- action = _parse_rl_tool_call(text, sdk_tool_calls)
- history.append(
- {
- "role": "orchestrator",
- "turn": turn,
- "raw": text,
- "action": action,
- }
- )
- self.record_trace_event(
- {
- "kind": "toolorchestra_rl_action",
- "turn": turn,
- "action": action,
- "raw": text,
- }
- )
-
- if action is None:
- parse_failures += 1
- if parse_failures >= 2:
- break
- continue
-
- name = action["name"]
- args = action.get("arguments", {})
- slot = args.get("model", "")
-
- # Validate against the upstream tool/arg schema.
- valid = (
- name in RL_ALL_TOOLS
- and isinstance(slot, str)
- and (slot in RL_ALL_TOOLS[name]["model"])
- )
- if not valid:
- parse_failures += 1
- if parse_failures >= 2:
- break
- # Replay with a softer nudge in the context.
- context_str += (
- f"\n[Orchestrator emitted invalid tool call "
- f"name={name!r} slot={slot!r} — try again.]\n"
- )
- continue
-
- # Paper-match (`method_cfg.pool == "paper"`) routes through
- # the Tavily/OpenRouter/Modal pool instead of the default
- # Anthropic-web-search-driven mapping. For `search` this
- # also forces the worker prompt to a raw query string
- # (Tavily takes a single search string, not a chat-style
- # framing).
- if paper_mode:
- worker = _paper_expert_for(
- slot,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
- # In paper mode, `enhance_reasoning` is always the coder
- # specialist regardless of the orchestrator's chosen tier.
- # The coder is then expected to emit a python block which
- # we exec in Modal (below).
- if name == "enhance_reasoning":
- worker = {
- "name": f"coder:{slot}",
- "type": "openrouter",
- "model": _PAPER_CODER_OPENROUTER,
- }
- else:
- worker = _expert_for(
- slot,
- self._local_model,
- self._local_endpoint,
- self._cloud_model,
- self._cloud_endpoint,
- )
-
- # Dispatch — the orchestrator only conveys a tool/model
- # choice, NOT a question rewrite; the prompt we send the
- # expert is the same context the orchestrator saw, framed
- # appropriately for the tool.
- if name == "search":
- if paper_mode:
- # Tavily takes a query string. Orchestrator-8B often
- # emits an extra `query` arg (not in the upstream
- # schema but useful) — prefer it; else fall back to
- # the raw question.
- q = args.get("query")
- w_input = q if isinstance(q, str) and q.strip() else question
- else:
- w_input = (
- f"Search the web to gather information that helps answer:\n"
- f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
- )
- elif name == "enhance_reasoning":
- if paper_mode:
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Write a short Python script that computes intermediate "
- "results which help answer the problem. Output ONLY the "
- "code inside one ```python ... ``` fenced block. Print "
- "any results you derive using `print(...)`. The script "
- "must run with the Python stdlib only — no extra pip "
- "installs."
- )
- else:
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Reason carefully. Outline the key intermediate steps and any "
- "computations or facts you can derive. Do NOT give a final "
- "answer — the orchestrator will collect your reasoning and "
- "call the answer tool next."
- )
- else: # name == "answer"
- w_input = (
- f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
- "Provide the final answer to the user. Respect any "
- "answer-format rules in the question (e.g. GAIA's "
- "FINAL ANSWER: 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..d63fd427
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/agent.py
@@ -0,0 +1,848 @@
+"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
+
+Two modes, gated by ``method_cfg.orchestrator_mode``:
+
+* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
+ orchestrator, dispatching to a numbered worker pool via JSON
+ ``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
+ a prompted upper-bound reference point — NOT the paper's setup.
+
+* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
+ served on a local vLLM is the orchestrator. It emits OpenAI-style
+ ``tool_calls`` (or ``{...}`` text blocks when
+ vLLM's tool parser doesn't catch them) for three expert tools —
+ ``enhance_reasoning``, ``answer``, ``search`` — exactly as in the
+ upstream ``evaluation/tools.json``. Each tool's ``model`` arg
+ (``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real
+ backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier
+ Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
+ for `*-3`. Search routes to the Anthropic server-side web_search.
+
+ We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
+ code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
+ Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
+ worker types. Energy-wise, "expert" answers are cloud calls.
+
+Pipeline per task (RL mode):
+
+1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
+ appropriate tool.` with the three tools declared.
+2. It emits one ``tool_call`` per turn — ``search`` updates the
+ context, ``enhance_reasoning`` appends code/exec output (we run the
+ tool as a plain LLM call, no sandbox — the model just gets prose
+ back), ``answer`` produces the final answer and the loop stops.
+3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
+ back to the strongest expert worker.
+
+Prompted-mode pipeline:
+
+1. Orchestrator (cloud) reads question + numbered worker pool.
+2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
+ "input": str}`` or ``{"action": "final_answer", "answer": str}``.
+3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
+ prompt; fallback to strongest worker on parse failure.
+
+Workers come from ``cfg["workers"]`` or a sensible default pool (local
+Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
+gpt-5-mini).
+"""
+
+from __future__ import annotations
+
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from openjarvis.agents._stubs import AgentContext
+from openjarvis.agents.hybrid._base import LocalCloudAgent
+from openjarvis.agents.hybrid._prices import PRICES
+from openjarvis.agents.hybrid.mini_swe_agent import (
+ _clone_repo,
+ _extract_diff,
+)
+from openjarvis.agents.hybrid.toolorchestra.clients import (
+ _call_orchestrator_with_tool_calls,
+)
+from openjarvis.agents.hybrid.toolorchestra.experts import (
+ _PAPER_CODER_OPENROUTER,
+ _expert_for,
+ _paper_expert_for,
+)
+from openjarvis.agents.hybrid.toolorchestra.parsing import (
+ _build_user_prompt,
+ _extract_final_answer_text,
+ _parse_action,
+ _parse_rl_tool_call,
+)
+from openjarvis.agents.hybrid.toolorchestra.prompts import (
+ FORCE_FINAL_PROMPT,
+ ORCHESTRATOR_SYS,
+ RL_ALL_TOOLS,
+ RL_ORCHESTRATOR_SYS,
+ RL_TOOLS_SPEC,
+)
+from openjarvis.agents.hybrid.toolorchestra.sandbox import (
+ _call_modal_python,
+ _extract_first_python_block,
+)
+from openjarvis.agents.hybrid.toolorchestra.workers import (
+ _TOOLORCH_SEARCH_TYPES,
+ _call_worker,
+ _resolve_worker_pool,
+ _swe_call_worker,
+)
+from openjarvis.core.registry import AgentRegistry
+
+
+@AgentRegistry.register("toolorchestra")
+class ToolOrchestraAgent(LocalCloudAgent):
+ """Multi-turn dispatcher over a mixed worker pool.
+
+ Two modes (see module docstring): ``method_cfg.orchestrator_mode``
+ is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
+ (paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
+ """
+
+ agent_id = "toolorchestra"
+
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
+ super().__init__(*args, **kwargs)
+ # Validate `method_cfg.worker_pool` early — surfaces config errors
+ # at agent construction rather than on the first task. No-op when
+ # the override is absent.
+ if self._cfg.get("worker_pool") is not None:
+ _resolve_worker_pool(
+ self._cfg,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+ # Validate `orchestrator_mode` (typo-checked here, not on first task).
+ mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
+ if mode not in ("prompted", "rl"):
+ raise ValueError(
+ f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
+ f"got {mode!r}"
+ )
+
+ def _run_paradigm(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
+ if mode == "rl":
+ return self._run_rl(input, context, **kwargs)
+ return self._run_prompted(input, context, **kwargs)
+
+ # ------------------------------------------------------------------
+ # Legacy prompted-orchestrator path.
+ # ------------------------------------------------------------------
+ def _run_prompted(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ cfg = self._cfg
+ question = input
+ # Resolution order (strict replace, no merge):
+ # 1. `cfg["workers"]` — legacy direct override, used by tests.
+ # 2. `cfg["worker_pool"]` — cell-config override; validated +
+ # $local/$cloud substituted.
+ # 3. `_default_pool(...)` — heterogeneous default.
+ if cfg.get("workers"):
+ workers = cfg["workers"]
+ else:
+ workers = _resolve_worker_pool(
+ cfg,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+ if not workers:
+ raise RuntimeError("toolorchestra: empty worker pool")
+
+ max_turns = int(cfg.get("max_turns", 6))
+ orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
+
+ task_meta = (context.metadata.get("task") if context is not None else {}) or {}
+ swe_mode = (
+ bool(cfg.get("swe_use_agent_loop"))
+ and bool(task_meta.get("problem_statement"))
+ and bool(task_meta.get("repo"))
+ and bool(task_meta.get("base_commit"))
+ )
+ shared_workdir: Optional[Path] = None
+ if swe_mode:
+ shared_workdir = Path(
+ tempfile.mkdtemp(
+ prefix=f"toolorch-swe-{task_meta.get('task_id', 'x')}-"
+ )
+ )
+ try:
+ _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
+ except Exception:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+ raise
+ self.record_trace_event(
+ {
+ "kind": "toolorchestra_swe_workdir",
+ "workdir": str(shared_workdir),
+ "repo": task_meta["repo"],
+ "base_commit": task_meta["base_commit"],
+ }
+ )
+
+ # try/finally guards ``shared_workdir`` against exceptions raised
+ # anywhere in the turn loop, the worker calls, the fallback, or
+ # the diff-extraction step. Without this, at n=500 SWE-bench an
+ # exception leaves hundreds of MB of cloned repos in tempdir.
+ try:
+ history: List[Dict[str, Any]] = []
+ tokens_local = 0
+ tokens_cloud = 0
+ cost = 0.0
+ n_web_searches_total = 0
+ # tool_calls: bash turns from SWE subloops + web_search uses
+ # from GAIA. Orchestrator dispatch turns are NOT counted (they
+ # produce text only — calling a worker is one tool call's worth
+ # of "delegation" but the actual tool action happens inside).
+ tool_calls = 0
+ final_answer: Optional[str] = None
+ forced_final = False
+ parse_failures = 0
+
+ for turn in range(1, max_turns + 1):
+ sys_prompt = ORCHESTRATOR_SYS
+ if turn == max_turns and final_answer is None:
+ sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
+ forced_final = True
+
+ user = _build_user_prompt(question, workers, history)
+ text, o_in, o_out = self._call_cloud(
+ user=user,
+ system=sys_prompt,
+ max_tokens=orch_max_tokens,
+ temperature=0.0,
+ )
+ tokens_cloud += o_in + o_out
+ cost += self.cost_usd(self._cloud_model, o_in, o_out)
+
+ action = _parse_action(text)
+ history.append(
+ {
+ "role": "orchestrator",
+ "turn": turn,
+ "raw": text,
+ "action": action,
+ }
+ )
+ self.record_trace_event(
+ {
+ "kind": "toolorchestra_action",
+ "turn": turn,
+ "action": action,
+ "raw": text,
+ }
+ )
+
+ if action is None:
+ parse_failures += 1
+ if parse_failures >= 2 or forced_final:
+ final_answer = _extract_final_answer_text(text)
+ break
+ continue
+
+ kind = action.get("action")
+ if kind == "final_answer":
+ final_answer = str(action.get("answer", "")).strip()
+ break
+ if kind == "call_worker":
+ wid = action.get("worker_id")
+ w_input = action.get("input", "")
+ if not isinstance(wid, int) or not (0 <= wid < len(workers)):
+ parse_failures += 1
+ if parse_failures >= 2 or forced_final:
+ final_answer = _extract_final_answer_text(text)
+ break
+ continue
+ worker = workers[wid]
+ if swe_mode and shared_workdir is not None:
+ (
+ w_text,
+ w_in,
+ w_out,
+ is_local,
+ extra_cost,
+ n_searches,
+ bash_turns,
+ ) = _swe_call_worker(
+ worker,
+ str(w_input),
+ cfg,
+ task_meta,
+ shared_workdir,
+ turn,
+ )
+ tool_calls += bash_turns
+ else:
+ w_text, w_in, w_out, is_local, extra_cost, n_searches = (
+ _call_worker(worker, str(w_input), cfg)
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ n_web_searches_total += n_searches
+ tool_calls += n_searches
+ history.append(
+ {
+ "role": "worker",
+ "turn": turn,
+ "worker_id": wid,
+ "worker_name": worker["name"],
+ "worker_model": worker["model"],
+ "output": w_text,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "n_web_searches": n_searches,
+ }
+ )
+ continue
+ # Unknown action kind — treat as parse failure.
+ parse_failures += 1
+
+ if final_answer is None:
+ # Hard fallback: call the strongest non-search worker directly.
+ # "Strongest" = highest output-token price in `_prices.PRICES`,
+ # which tracks model capability tier closely enough for this.
+ # Search workers are excluded — they answer fact-lookup
+ # questions, not synthesis.
+ non_search = [
+ w for w in workers if w.get("type") not in _TOOLORCH_SEARCH_TYPES
+ ] or workers
+ worker = max(
+ non_search,
+ key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
+ )
+ if swe_mode and shared_workdir is not None:
+ (ans, w_in, w_out, is_local, extra_cost, _, bash_turns) = (
+ _swe_call_worker(
+ worker,
+ question,
+ cfg,
+ task_meta,
+ shared_workdir,
+ max_turns + 1,
+ )
+ )
+ tool_calls += bash_turns
+ else:
+ ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
+ worker, question, cfg
+ )
+ if is_local:
+ tokens_local += w_in + w_out
+ else:
+ tokens_cloud += w_in + w_out
+ cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
+ history.append(
+ {
+ "role": "worker",
+ "turn": max_turns + 1,
+ "worker_id": worker["id"],
+ "worker_name": worker["name"],
+ "worker_model": worker["model"],
+ "output": ans,
+ "tokens_in": w_in,
+ "tokens_out": w_out,
+ "fallback": True,
+ }
+ )
+ final_answer = ans
+
+ # In SWE mode, the authoritative output is the working-tree diff —
+ # frame it (the runner extracts it via the scorer's ```diff fence).
+ if swe_mode and shared_workdir is not None:
+ patch = _extract_diff(shared_workdir)
+ if patch.strip():
+ final_answer = (
+ f"{final_answer}\n\n```diff\n{patch}```"
+ if final_answer
+ else f"```diff\n{patch}```"
+ )
+
+ meta = {
+ "tokens_local": tokens_local,
+ "tokens_cloud": tokens_cloud,
+ "cost_usd": cost,
+ "turns": len([h for h in history if h["role"] == "orchestrator"]),
+ "web_search_uses": n_web_searches_total,
+ "tool_calls": int(tool_calls),
+ "traces": {
+ "history": history,
+ "forced_final": forced_final,
+ "parse_failures": parse_failures,
+ "workers": workers,
+ "n_web_searches": n_web_searches_total,
+ "note": (
+ "inference-only port; the RL-trained Nemotron-Orchestrator-8B "
+ "is NOT in the loop. Results are preliminary."
+ ),
+ },
+ }
+ return final_answer, meta
+ finally:
+ if shared_workdir is not None:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+
+ # ------------------------------------------------------------------
+ # Paper-faithful Orchestrator-8B path.
+ # ------------------------------------------------------------------
+ def _run_rl(
+ self,
+ input: str,
+ context: Optional[AgentContext],
+ **kwargs: Any,
+ ) -> Tuple[str, Dict[str, Any]]:
+ cfg = self._cfg
+ question = input
+
+ # Orchestrator endpoint / model (where the RL'd 8B lives).
+ orch_endpoint = str(
+ cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
+ )
+ orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
+ max_turns = int(cfg.get("max_turns", 8))
+ orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
+ orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
+
+ # Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
+ # replaces `_expert_for` and `enhance_reasoning` is post-processed
+ # through a Modal Python sandbox. See module docstring + paper-match
+ # doc at `docs/26.5.19/toolorchestra-papermatch.md`.
+ paper_mode = str(cfg.get("pool", "")).lower() == "paper"
+
+ # SWE-bench detection: same gate as the prompted path. Requires
+ # `method_cfg.swe_use_agent_loop = true` AND the task carries the
+ # SWE-bench fields. When active, the `enhance_reasoning` and
+ # `answer` workers route through `run_swe_agent_loop` on a shared
+ # workdir; search workers stay one-shot. At end, the working-tree
+ # diff is appended to final_answer so `_score_swebench` can extract
+ # it via the ```diff fence.
+ task_meta = (context.metadata.get("task") if context is not None else {}) or {}
+ swe_mode = (
+ bool(cfg.get("swe_use_agent_loop"))
+ and bool(task_meta.get("problem_statement"))
+ and bool(task_meta.get("repo"))
+ and bool(task_meta.get("base_commit"))
+ )
+ shared_workdir: Optional[Path] = None
+ if swe_mode:
+ shared_workdir = Path(
+ tempfile.mkdtemp(
+ prefix=f"toolorch-rl-swe-{task_meta.get('task_id', 'x')}-"
+ )
+ )
+ try:
+ _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
+ except Exception:
+ shutil.rmtree(shared_workdir, ignore_errors=True)
+ raise
+ self.record_trace_event(
+ {
+ "kind": "toolorchestra_rl_swe_workdir",
+ "workdir": str(shared_workdir),
+ "repo": task_meta["repo"],
+ "base_commit": task_meta["base_commit"],
+ }
+ )
+
+ # ``context_str`` mirrors the upstream's running context — accumulates
+ # search documents and code/exec snippets across turns. We keep this
+ # as a single string for prompt simplicity; the upstream uses
+ # tokenized cutoffs (we cap at ~24k chars instead).
+ context_str = ""
+ doc_list: List[str] = []
+ history: List[Dict[str, Any]] = []
+ tokens_local = 0
+ tokens_cloud = 0
+ cost = 0.0
+ n_web_searches_total = 0
+ tool_calls = 0
+ final_answer: Optional[str] = None
+ parse_failures = 0
+
+ # Single outer try/finally guards `shared_workdir` against any
+ # exception in the orchestrator loop, the post-loop fallback, or
+ # the diff-extraction step. Matches the prompted path's pattern.
+ try:
+ for turn in range(1, max_turns + 1):
+ user = (
+ f"Problem: {question}\n\n{context_str}\n\n"
+ "Choose an appropriate tool."
+ )
+
+ # Orchestrator-8B served on local vLLM. We pass the three NVlabs
+ # tools verbatim. In paper-mode we use the local helper so we
+ # get the SDK-level ``tool_calls`` object back — `_call_vllm`
+ # returns just text and loses the call when vLLM's parser
+ # caught it. Orchestrator-8B emits its routing decision in the
+ # OpenAI-native ``tool_calls`` array with an empty text body,
+ # so the legacy `_call_vllm` path saw nothing and silently fell
+ # through to the answer-1 fallback (parse_failures: 2 on every
+ # non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
+ # modes now use `_call_orchestrator_with_tool_calls` so the
+ # parser can read structured tool calls; the text-tag path in
+ # `_parse_rl_tool_call` is still the fallback when `tool_calls`
+ # is empty.
+ text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
+ orch_model,
+ orch_endpoint,
+ user=user,
+ system=RL_ORCHESTRATOR_SYS,
+ max_tokens=orch_max_tokens,
+ temperature=orch_temp,
+ tools=RL_TOOLS_SPEC,
+ )
+ self.record_trace_event(
+ {
+ "kind": "vllm",
+ "role": "orchestrator",
+ "model": orch_model,
+ "endpoint": orch_endpoint,
+ "system": RL_ORCHESTRATOR_SYS,
+ "user": user,
+ "response": text,
+ "tool_calls": [
+ {
+ "id": getattr(tc, "id", None),
+ "type": getattr(tc, "type", None),
+ "function": {
+ "name": getattr(
+ getattr(tc, "function", None), "name", None
+ ),
+ "arguments": getattr(
+ getattr(tc, "function", None), "arguments", None
+ ),
+ },
+ }
+ for tc in (sdk_tool_calls or [])
+ ],
+ "tokens_in": o_in,
+ "tokens_out": o_out,
+ }
+ )
+ tokens_local += o_in + o_out
+
+ action = _parse_rl_tool_call(text, sdk_tool_calls)
+ history.append(
+ {
+ "role": "orchestrator",
+ "turn": turn,
+ "raw": text,
+ "action": action,
+ }
+ )
+ self.record_trace_event(
+ {
+ "kind": "toolorchestra_rl_action",
+ "turn": turn,
+ "action": action,
+ "raw": text,
+ }
+ )
+
+ if action is None:
+ parse_failures += 1
+ if parse_failures >= 2:
+ break
+ continue
+
+ name = action["name"]
+ args = action.get("arguments", {})
+ slot = args.get("model", "")
+
+ # Validate against the upstream tool/arg schema.
+ valid = (
+ name in RL_ALL_TOOLS
+ and isinstance(slot, str)
+ and (slot in RL_ALL_TOOLS[name]["model"])
+ )
+ if not valid:
+ parse_failures += 1
+ if parse_failures >= 2:
+ break
+ # Replay with a softer nudge in the context.
+ context_str += (
+ f"\n[Orchestrator emitted invalid tool call "
+ f"name={name!r} slot={slot!r} — try again.]\n"
+ )
+ continue
+
+ # Paper-match (`method_cfg.pool == "paper"`) routes through
+ # the Tavily/OpenRouter/Modal pool instead of the default
+ # Anthropic-web-search-driven mapping. For `search` this
+ # also forces the worker prompt to a raw query string
+ # (Tavily takes a single search string, not a chat-style
+ # framing).
+ if paper_mode:
+ worker = _paper_expert_for(
+ slot,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+ # In paper mode, `enhance_reasoning` is always the coder
+ # specialist regardless of the orchestrator's chosen tier.
+ # The coder is then expected to emit a python block which
+ # we exec in Modal (below).
+ if name == "enhance_reasoning":
+ worker = {
+ "name": f"coder:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_CODER_OPENROUTER,
+ }
+ else:
+ worker = _expert_for(
+ slot,
+ self._local_model,
+ self._local_endpoint,
+ self._cloud_model,
+ self._cloud_endpoint,
+ )
+
+ # Dispatch — the orchestrator only conveys a tool/model
+ # choice, NOT a question rewrite; the prompt we send the
+ # expert is the same context the orchestrator saw, framed
+ # appropriately for the tool.
+ if name == "search":
+ if paper_mode:
+ # Tavily takes a query string. Orchestrator-8B often
+ # emits an extra `query` arg (not in the upstream
+ # schema but useful) — prefer it; else fall back to
+ # the raw question.
+ q = args.get("query")
+ w_input = q if isinstance(q, str) and q.strip() else question
+ else:
+ w_input = (
+ f"Search the web to gather information that helps answer:\n"
+ f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
+ )
+ elif name == "enhance_reasoning":
+ if paper_mode:
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Write a short Python script that computes intermediate "
+ "results which help answer the problem. Output ONLY the "
+ "code inside one ```python ... ``` fenced block. Print "
+ "any results you derive using `print(...)`. The script "
+ "must run with the Python stdlib only — no extra pip "
+ "installs."
+ )
+ else:
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Reason carefully. Outline the key intermediate steps and any "
+ "computations or facts you can derive. Do NOT give a final "
+ "answer — the orchestrator will collect your reasoning and "
+ "call the answer tool next."
+ )
+ else: # name == "answer"
+ w_input = (
+ f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
+ "Provide the final answer to the user. Respect any "
+ "answer-format rules in the question (e.g. GAIA's "
+ "FINAL ANSWER: 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..ec390ac4
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/clients.py
@@ -0,0 +1,51 @@
+"""Orchestrator vLLM tool-call client for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Tuple
+
+
+def _call_orchestrator_with_tool_calls(
+ model: str,
+ endpoint: str,
+ *,
+ user: str,
+ system: str,
+ max_tokens: int,
+ temperature: float,
+ tools: List[Dict[str, Any]],
+ timeout: float = 600.0,
+) -> Tuple[str, int, int, Any]:
+ """Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls).
+
+ Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level
+ ``tool_calls`` object so the RL-mode parser can match against it
+ directly. Otherwise vLLM's tool parser silently swallows the tool call
+ into the SDK field while leaving ``content == ''`` — and the text-tag
+ parser sees nothing, falling through to the answer-1 fallback. (Bug
+ observed 2026-05-19 on the paper-match smoke; same path was buggy on
+ the default pool too, just less reproducibly.)
+ """
+ from openai import OpenAI
+
+ client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
+ messages = [
+ {"role": "system", "content": system},
+ {"role": "user", "content": user},
+ ]
+ resp = client.chat.completions.create(
+ model=model,
+ messages=messages,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ tools=tools,
+ extra_body={"chat_template_kwargs": {"enable_thinking": False}},
+ )
+ choice = resp.choices[0]
+ message = choice.message
+ text = message.content or ""
+ tool_calls = getattr(message, "tool_calls", None)
+ u = resp.usage
+ p = getattr(u, "prompt_tokens", 0) if u else 0
+ c = getattr(u, "completion_tokens", 0) if u else 0
+ return text, p, c, tool_calls
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/experts.py b/src/openjarvis/agents/hybrid/toolorchestra/experts.py
new file mode 100644
index 00000000..da6ff6d8
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/experts.py
@@ -0,0 +1,187 @@
+"""Slot -> worker expert mapping for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+# Default model used when an `anthropic-web-search` entry omits `model`.
+_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
+
+# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
+# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
+# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
+# tiers since we don't have Qwen-Math served.
+#
+# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
+# so the substitution is deferred until we know the cell's resolved local/cloud
+# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
+
+
+def _expert_for(
+ slot: str,
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "anthropic",
+) -> Dict[str, Any]:
+ """Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
+
+ Routing policy:
+ - `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off
+ `cloud_endpoint` ("anthropic"/"openai"/"gemini")
+ - `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's
+ cost tier for mid OpenAI calls)
+ - `*-3` (local tier) -> local vLLM (`local_model`)
+ - `answer-math-*` -> same tiers as the numeric suffix
+ - `search-*` -> provider-native web search when the cloud
+ endpoint supports it; otherwise Anthropic
+ """
+ if slot.startswith("search"):
+ ep = (cloud_endpoint or "anthropic").lower()
+ if ep == "openai":
+ return {
+ "name": f"search:{slot}",
+ "type": "openai-web-search",
+ "model": cloud_model,
+ }
+ if ep == "gemini":
+ return {
+ "name": f"search:{slot}",
+ "type": "gemini-web-search",
+ "model": cloud_model,
+ }
+ return {
+ "name": f"search:{slot}",
+ "type": "anthropic-web-search",
+ "model": _DEFAULT_WEB_SEARCH_MODEL,
+ }
+ if slot.endswith("-1") or slot.endswith("-math-1"):
+ ep = (cloud_endpoint or "anthropic").lower()
+ if ep not in ("anthropic", "openai", "gemini"):
+ ep = "anthropic"
+ return {
+ "name": f"frontier:{slot}",
+ "type": ep,
+ "model": cloud_model,
+ }
+ if slot.endswith("-2") or slot.endswith("-math-2"):
+ return {
+ "name": f"mid:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+ # `*-3` / `*-4` collapse to local vLLM (paper uses Qwen3-32B etc.;
+ # we substitute whatever local model the cell wired up).
+ if local_model and local_endpoint:
+ return {
+ "name": f"local:{slot}",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ }
+ # Fallback if no local — gpt-5-mini.
+ return {
+ "name": f"mid-fallback:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+
+
+# ============================================================================
+# Paper-match expert mapping (2026-05-19).
+# ============================================================================
+# Maps the orchestrator's `model` slot to a paper-match worker spec. Differs
+# from `_expert_for` in that it pulls in OpenRouter-hosted code/math/generalist
+# models and routes `search` through Tavily, while `enhance_reasoning` is
+# expected to produce code that the caller pipes through a Modal sandbox
+# (handled at dispatch time, not here).
+#
+# Slot map (paper-faithful where we can; substitutions noted in toolorchestra
+# paper-match docs `docs/26.5.19/toolorchestra-papermatch.md`):
+#
+# reasoner-1 -> GPT-5 (frontier reasoner)
+# reasoner-2 -> GPT-5-mini (mid)
+# reasoner-3 -> local Qwen (Orchestrator-8B endpoint also serves this)
+# answer-1 -> GPT-5
+# answer-2 -> GPT-5-mini
+# answer-3 -> Llama-3.3-70B (OpenRouter, generalist tier-3 per spec)
+# answer-4 -> local Qwen
+# answer-math-1 -> Qwen-2.5-Coder-32B via OpenRouter
+# (paper uses Qwen-2.5-Math-72B; not on OpenRouter — see doc)
+# answer-math-2 -> Qwen-2.5-Coder-32B via OpenRouter
+# (paper uses Qwen-2.5-Math-7B; not on OpenRouter — see doc)
+# search-* -> Tavily search (paper)
+#
+# `enhance_reasoning` is dispatched through the coder specialist regardless of
+# slot tier — the orchestrator emits one of `reasoner-{1,2,3}` and the caller
+# routes the same way in all three cases, then optionally extracts a python
+# code block and execs it in Modal. (We keep the slot-aware routing inside the
+# `reasoner-*` map above for parity, but the `enhance_reasoning` tool itself
+# pins the coder regardless. See `_run_rl_paper` dispatch.)
+
+_PAPER_CODER_OPENROUTER = "qwen/qwen-2.5-coder-32b-instruct"
+_PAPER_GENERALIST_TIER3_OPENROUTER = "meta-llama/llama-3.3-70b-instruct"
+
+
+def _paper_expert_for(
+ slot: str,
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "openai",
+) -> Dict[str, Any]:
+ """Paper-match counterpart of ``_expert_for``.
+
+ Differs from ``_expert_for``:
+ - Search slots go to ``tavily-search`` (not Anthropic web_search).
+ - Tier-3 generalist answer (``answer-3``) routes to Llama-3.3-70B via
+ OpenRouter rather than collapsing onto the local vLLM.
+ - Math slots route to the OpenRouter code specialist (Qwen-2.5-Coder-32B)
+ as a substitute for the unavailable Qwen-2.5-Math-{72B,7B}.
+ - ``reasoner-1`` / ``answer-1`` route to GPT-5 by default (paper).
+ """
+ if slot.startswith("search"):
+ return {
+ "name": f"tavily:{slot}",
+ "type": "tavily-search",
+ "model": "tavily",
+ }
+ if slot in ("answer-math-1", "answer-math-2"):
+ return {
+ "name": f"math-coder:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_CODER_OPENROUTER,
+ }
+ if slot == "answer-3":
+ return {
+ "name": f"generalist-llama:{slot}",
+ "type": "openrouter",
+ "model": _PAPER_GENERALIST_TIER3_OPENROUTER,
+ }
+ if slot.endswith("-1"):
+ # Tier-1 frontier reasoner / answer — paper uses GPT-5.
+ return {
+ "name": f"frontier:{slot}",
+ "type": "openai",
+ "model": "gpt-5",
+ }
+ if slot.endswith("-2"):
+ return {
+ "name": f"mid:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
+ # `*-3` / `*-4` collapse onto the local vLLM (the orchestrator endpoint
+ # also serves the local Qwen for the rare local-tier slot).
+ if local_model and local_endpoint:
+ return {
+ "name": f"local:{slot}",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ }
+ return {
+ "name": f"mid-fallback:{slot}",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ }
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/parsing.py b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py
new file mode 100644
index 00000000..b36feaa3
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py
@@ -0,0 +1,161 @@
+"""Action / tool-call parsing + prompt assembly for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, List, Optional
+
+# Regex for ``{...}`` 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)
+
+# Fallback for a known failure mode: the SFT'd orchestrator often emits its
+# delegation as ``\boxed{expert_ab12: }`` (math-data habit bleeding
+# into routing) instead of a real ```` JSON block. The name before the
+# colon is a valid (anonymized) tool label; the rest is the sub-question/args.
+#
+# Real outputs use several shapes, all handled here:
+# \boxed{web_search: } -> name, args
+# \boxed{web_search query: } -> the word "query" is dropped
+# \boxed{file_read, path: } -> the ", " hint sets the arg key
+# We require a ``name : `` structure, so a genuine
+# answer like ``\boxed{10}``, ``\boxed{b,e}`` or ``\boxed{The answer is: 42}``
+# (first token isn't followed by a bare ``:`` / ``, key:`` / ``query:``) is left
+# alone. ``group("key")`` is the explicit arg key when the model wrote ``, key:``.
+_BOXED_DELEGATION_RE = re.compile(
+ r"\\boxed\{\s*([A-Za-z_][\w-]*)\s*"
+ r"(?:,\s*(?P\w+)\s*)?" # optional ", key" hint (e.g. file_read, path:)
+ r"(?:query\s*)?" # optional literal "query" word before the colon
+ r":\s*(.+?)\s*\}\s*$",
+ re.DOTALL,
+)
+
+
+def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
+ """Return ``{"name": str, "arguments": dict}`` or None.
+
+ Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
+ back to scraping ``{...}`` tags from the raw
+ content. We take the first tool call only — Orchestrator-8B was trained
+ to emit exactly one per turn.
+ """
+ # SDK-level path.
+ if sdk_tool_calls:
+ first = sdk_tool_calls[0]
+ name = getattr(getattr(first, "function", None), "name", None)
+ args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
+ try:
+ args = json.loads(args_raw)
+ except json.JSONDecodeError:
+ args = {}
+ if isinstance(name, str) and isinstance(args, dict):
+ return {"name": name, "arguments": args}
+ # Text-tag fallback.
+ if not isinstance(content, str):
+ return None
+ m = _TOOL_CALL_TAG_RE.search(content)
+ if m:
+ try:
+ obj = json.loads(m.group(1))
+ except json.JSONDecodeError:
+ obj = None
+ if isinstance(obj, dict):
+ name = obj.get("name")
+ args = obj.get("arguments", {})
+ if isinstance(name, str) and isinstance(args, dict):
+ return {"name": name, "arguments": args}
+ # \boxed{name[, key][ query]: } fallback (see _BOXED_DELEGATION_RE).
+ bm = _BOXED_DELEGATION_RE.search(content.strip())
+ if bm:
+ arg_val = bm.groups()[-1].strip()
+ key = bm.group("key") or "input"
+ return {"name": bm.group(1), "arguments": {key: arg_val}}
+ return None
+
+
+def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
+ return "\n".join(
+ f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
+ )
+
+
+def _build_user_prompt(
+ question: str,
+ workers: List[Dict[str, Any]],
+ history: List[Dict[str, Any]],
+) -> str:
+ pieces = [
+ f"Worker pool:\n{_build_pool_block(workers)}",
+ f"User question:\n{question}",
+ ]
+ if history:
+ pieces.append("Conversation so far (orchestrator turns and worker outputs):")
+ for h in history:
+ if h["role"] == "orchestrator":
+ pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
+ else:
+ pieces.append(
+ f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
+ f"{h['output']}"
+ )
+ pieces.append(
+ "Emit the next JSON action object now — exactly one object, no prose."
+ )
+ return "\n\n".join(pieces)
+
+
+def _strip_fences(s: str) -> str:
+ s = s.strip()
+ if s.startswith("```"):
+ first_nl = s.find("\n")
+ if first_nl != -1:
+ s = s[first_nl + 1 :]
+ if s.endswith("```"):
+ s = s[:-3]
+ s = s.strip()
+ return s
+
+
+def _parse_action(text: str) -> Optional[Dict[str, Any]]:
+ s = _strip_fences(text)
+ # First try direct parse, then balanced-brace extraction.
+ try:
+ obj = json.loads(s)
+ if isinstance(obj, dict) and "action" in obj:
+ return obj
+ except json.JSONDecodeError:
+ pass
+ start = s.find("{")
+ if start == -1:
+ return None
+ depth = 0
+ for i in range(start, len(s)):
+ c = s[i]
+ if c == "{":
+ depth += 1
+ elif c == "}":
+ depth -= 1
+ if depth == 0:
+ try:
+ obj = json.loads(s[start : i + 1])
+ if isinstance(obj, dict) and "action" in obj:
+ return obj
+ except json.JSONDecodeError:
+ return None
+ return None
+
+
+def _extract_final_answer_text(text: str) -> str:
+ """Best-effort: pull the answer string from a malformed action emission.
+
+ Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
+ """
+ m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
+ if m:
+ return m.group(1).encode("utf-8").decode("unicode_escape")
+ m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
+ if m:
+ return m.group(1).strip()
+ return text.strip()
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/prompts.py b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py
new file mode 100644
index 00000000..6c2e9038
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py
@@ -0,0 +1,110 @@
+"""Prompt strings + tool specs for ToolOrchestraAgent (split from toolorchestra.py)."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+ORCHESTRATOR_SYS = """\
+You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
+
+ {"action": "call_worker", "worker_id": , "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..9542fe3e
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py
@@ -0,0 +1,417 @@
+"""Faithful unified-tool rollout loop for ToolOrchestra (arXiv:2511.21689 §2.2).
+
+One reasoning->action->observation loop where the orchestrator picks **a named
+tool** (one per model, from :mod:`expert_registry`) each turn, the environment
+executes it, and the observation is appended to a running context. The rollout
+ends when the orchestrator emits a turn with **no tool call** (its text is the
+final answer) or ``max_turns`` is hit.
+
+The loop is parameterized over two injected callables so it is pure control flow
+(no network) and unit-testable with fakes — the agent supplies real ones:
+
+* ``call_orchestrator(messages, tool_specs) -> (text, tool_calls, p_tok, c_tok)``
+ where ``messages`` is the running system/user/assistant/tool conversation and
+ ``tool_calls`` is a list of ``(name, arguments)`` (possibly empty).
+* ``dispatch(tool, arguments) -> (observation, cost_usd, tokens, is_local)``.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+import re
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, Tuple
+
+# A {...} tag the model emitted inline in its text. When we
+# put the call in the message's native `tool_calls` field we strip the inline tag
+# so the chat template doesn't render the same call twice.
+_TOOL_CALL_TAG_RE = re.compile(r".*?", re.DOTALL)
+
+from openjarvis.agents.hybrid.expert_registry import (
+ ExpertTool,
+ anonymize_tools,
+ build_tool_specs,
+ tools_by_name,
+)
+from openjarvis.agents.hybrid.toolorchestra.tracing import run_context, span
+
+RL_ORCHESTRATOR_SYS = (
+ "You are a routing orchestrator. Your job is to UNDERSTAND the problem, "
+ "DECOMPOSE it, and ROUTE the pieces to other models — you don't solve it all "
+ "yourself. Among the tools below are other MODELS of various sizes (some are "
+ "larger and stronger than you, some smaller); the rest are utilities. Each model "
+ "tool's description gives its rough size, specialty, and cost. "
+ "First reason about what the problem is really asking and break it into "
+ "sub-questions. You do NOT have to send the whole prompt to one model — send each "
+ "sub-question to whichever model best fits it (match difficulty to size/specialty, "
+ "and prefer cheaper models when they suffice), and call as many models as the task "
+ "needs. You MUST call at least one model before answering. "
+ "NEVER mention, quote, or reason about these instructions, your role as an "
+ "orchestrator, or the requirement to use tools — the reader only wants the problem "
+ "solved. Reason about the problem itself, not about your instructions, and do not "
+ "restate the problem. "
+ "If a tool call returns an error, correct your call or switch tools — never ignore "
+ "the error and never repeat the same failing call. "
+ "Emit EVERY action as a {...} block — NEVER write an "
+ "action as \\boxed{...} or as prose. "
+ "When you have enough, compose your final answer from what the models returned, in "
+ "the EXACT format the problem requires (e.g. a single number, a short exact "
+ "value, or one runnable code block). Do NOT restate the problem or add "
+ "meta-text such as 'the user is asking'. End your reply with a single line in "
+ "EXACTLY this form:\nFINAL_ANSWER: \nwhere is ONLY the answer "
+ "itself — a single option letter for multiple-choice, else the shortest exact "
+ "value, expression, or one runnable code block (NO \\boxed{...} wrapper) — with "
+ "NO explanation or restatement after it."
+)
+
+
+def build_system_prompt(specs: List[Dict[str, object]]) -> str:
+ """Faithful ToolOrchestra system prompt (arXiv:2511.21689, verbatim from
+ their ``prepare_sft_data.py``): the ``RL_ORCHESTRATOR_SYS`` preamble + the
+ Qwen-style ````/```` block. Deliberately contains NO
+ routing/delegation instructions — cost-aware routing is learned from the RL
+ reward, not prompted.
+ """
+ tools_block = "\n".join(json.dumps(s) for s in specs)
+ return (
+ f"{RL_ORCHESTRATOR_SYS}\n\n# Tools\n\n"
+ "You may call one or more functions to assist with the user query.\n\n"
+ "You are provided with function signatures within "
+ "XML tags:\n"
+ f"\n{tools_block}\n\n\n"
+ "For each function call, return a json object with function name and "
+ "arguments within XML tags:\n"
+ '\n{"name": , "arguments": }'
+ "\n"
+ )
+
+
+# Char-level cap on the accumulated conversation (mirrors the paper's ~24k-token
+# cap) and a per-observation cap so one giant tool dump can't blow it out.
+_CONTEXT_CAP = 24000
+_OBS_CAP = 8000
+# After this many tool calls, push the orchestrator to stop and answer — bounds
+# the "keep re-asking the same model forever" loop (esp. gemma).
+_SOFT_CALL_CAP = 6
+
+
+def _trim_history(messages: List[Dict[str, str]]) -> None:
+ """Keep the message history under ``_CONTEXT_CAP`` chars by dropping the
+ OLDEST assistant/tool exchange, never the system prompt or the problem."""
+
+ def total() -> int:
+ return sum(len(m.get("content") or "") for m in messages)
+
+ # messages[0]=system, messages[1]=user problem — always keep those two.
+ while total() > _CONTEXT_CAP and len(messages) > 4:
+ del messages[2:4]
+
+
+@dataclass
+class UnifiedTurn:
+ """One orchestrator turn. ``tool_name is None`` marks the final-answer turn."""
+
+ reasoning: str
+ tool_name: Optional[str] = None
+ arguments: Dict[str, object] = field(default_factory=dict)
+ observation: Optional[str] = None
+
+
+@dataclass
+class UnifiedRollout:
+ turns: List[UnifiedTurn]
+ final_answer: str
+ cost_usd: float = 0.0
+ tokens: int = 0
+ num_tool_calls: int = 0
+ parse_failures: int = 0
+ # When experts were anonymized for this rollout, maps the opaque label the
+ # policy saw (e.g. ``expert_a3f9``) back to the real tool name (``gpt_5_5``).
+ anon_map: Optional[Dict[str, str]] = None
+ # The EXACT tool specs the policy saw this rollout (anonymized when
+ # ``anonymize=True``). Serialization MUST build the saved system prompt from
+ # these — not from the real registry — or the prompt's ```` block ends
+ # up with real model names+pricing while the assistant ```` tags
+ # use the anon labels, which both breaks SFT (calls a tool absent from its
+ # list) and re-injects the name/cost bias anonymization removed.
+ tool_specs: Optional[List[Dict[str, object]]] = None
+
+ def tool_calls(self) -> List[Tuple[str, Dict[str, object]]]:
+ return [(t.tool_name, t.arguments) for t in self.turns if t.tool_name]
+
+
+def _tool_prompt(tool: ExpertTool, arguments: Dict[str, object], question: str) -> str:
+ """The text we actually send the dispatched tool, framed by its arg schema."""
+ for key in ("input", "query", "code"):
+ val = arguments.get(key)
+ if isinstance(val, str) and val.strip():
+ return val
+ return question
+
+
+def _run_unified_rollout_inner(
+ question: str,
+ tools: List[ExpertTool],
+ *,
+ call_orchestrator: Callable[
+ ..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
+ ],
+ dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
+ max_turns: int = 50,
+ system: str = RL_ORCHESTRATOR_SYS,
+ anonymize: bool = False,
+) -> UnifiedRollout:
+ """Drive the faithful unified-tool rollout for one task.
+
+ ``anonymize``: replace each model expert with an opaque random label, a
+ uniform description and no cost line, shuffled — so the policy can't route on
+ a model's name/position/cost (which we found dominate the choice). The
+ anon->real mapping is returned on the rollout for offline analysis.
+ """
+ anon_map: Optional[Dict[str, str]] = None
+ if anonymize:
+ tools, anon_map = anonymize_tools(tools, random.Random())
+ specs = build_tool_specs(tools)
+ by_name = tools_by_name(tools)
+
+ # Proper multi-turn conversation (NOT a flattened user blob): the orchestrator
+ # sees its own calls as `assistant` turns and each observation as a `tool`
+ # turn, so it can tell a tool result from user input and doesn't re-derive the
+ # whole problem every turn. This mirrors the serialized SFT format exactly.
+ messages: List[Dict[str, str]] = [
+ {"role": "system", "content": system},
+ {"role": "user", "content": f"Problem: {question}"},
+ ]
+ turns: List[UnifiedTurn] = []
+ cost = 0.0
+ tokens = 0
+ n_tool_calls = 0
+ parse_failures = 0
+ nudges = 0
+ empty_input_nudges = 0
+ final_answer = ""
+
+ for _ in range(max_turns):
+ text, tool_calls, p_tok, c_tok = call_orchestrator(messages, specs)
+ tokens += int(p_tok) + int(c_tok)
+
+ if not tool_calls:
+ # Enforce the "MUST delegate to a model" rule: if the orchestrator
+ # tries to answer before ANY tool call, nudge it to route instead of
+ # accepting a solve-it-yourself answer (which reject-sampling drops).
+ if n_tool_calls == 0 and nudges < 2:
+ nudges += 1
+ messages.append({"role": "assistant", "content": text or ""})
+ messages.append(
+ {
+ "role": "user",
+ "content": (
+ "Make progress by delegating a concrete sub-question to one of the "
+ "models now via a ."
+ ),
+ }
+ )
+ continue
+ # No tool call -> the orchestrator is answering. Terminate.
+ final_answer = (text or "").strip()
+ turns.append(UnifiedTurn(reasoning=text or "", tool_name=None))
+ messages.append({"role": "assistant", "content": text or ""})
+ break
+
+ name, arguments = tool_calls[0]
+ if name not in by_name:
+ parse_failures += 1
+ messages.append({"role": "assistant", "content": text or ""})
+ messages.append(
+ {
+ "role": "user",
+ "content": f"[invalid tool {name!r} — choose one from the provided tool list]",
+ }
+ )
+ if parse_failures >= 2:
+ final_answer = (text or "").strip()
+ break
+ continue
+
+ tool = by_name[name]
+ # Empty-input guard: the small orchestrator often emits a with
+ # a blank/missing input on harder multi-turn tasks. Dispatching it returns
+ # a "no input provided" error observation that poisons the whole (often
+ # otherwise-correct) trajectory — the dominant clean-yield killer on hard
+ # tasks. Instead, nudge the model to resend WITH input and drop the
+ # malformed turn (don't record it), so it never enters the training data.
+ if tool.kind == "model" or tool.backend_type != "openjarvis-tool":
+ _has_input = any(
+ isinstance(arguments.get(k), str) and arguments.get(k).strip()
+ for k in ("input", "query", "code")
+ )
+ if not _has_input and empty_input_nudges < 3:
+ empty_input_nudges += 1
+ messages.append({"role": "assistant", "content": text or ""})
+ messages.append(
+ {
+ "role": "user",
+ "content": (
+ f"Your call to {name} had an empty 'input'. Resend the "
+ " with a non-empty 'input' field containing the "
+ "concrete sub-question to delegate."
+ ),
+ }
+ )
+ continue
+ obs, dcost, dtok, _is_local = dispatch(tool, arguments)
+ cost += float(dcost)
+ tokens += int(dtok)
+ n_tool_calls += 1
+ turns.append(
+ UnifiedTurn(
+ reasoning=text or "",
+ tool_name=name,
+ arguments=dict(arguments),
+ observation=obs,
+ )
+ )
+ # The model's own action as an assistant turn, then the observation as a
+ # distinct `tool` turn it reads as a tool response.
+ #
+ # Use the NATIVE OpenAI tool-call protocol (assistant carries `tool_calls`
+ # with an id; the tool message references it via `tool_call_id`) rather
+ # than embedding the tag as plain text. vLLM tolerates the
+ # text form, but strict OpenAI-compatible providers (Anthropic) reject a
+ # `tool` message with no `tool_call_id` — 400 "tool_call_id: Field
+ # required" — which made cloud teachers unusable as the orchestrator.
+ # The SFT target is unaffected: it's rendered from `roll.turns` by the
+ # serializer, not from this message list.
+ call_id = f"call_{n_tool_calls}"
+ # Strip the tag if the model emitted it inline, so the call isn't rendered
+ # twice (once from content, once from the template's tool_calls block).
+ call_content = _TOOL_CALL_TAG_RE.sub("", text or "").strip()
+ messages.append(
+ {
+ "role": "assistant",
+ "content": call_content,
+ "tool_calls": [
+ {
+ "id": call_id,
+ "type": "function",
+ "function": {
+ "name": name,
+ "arguments": json.dumps(arguments, ensure_ascii=False),
+ },
+ }
+ ],
+ }
+ )
+ obs_text = obs or ""
+ if len(obs_text) > _OBS_CAP:
+ obs_text = obs_text[:_OBS_CAP] + "\n…[truncated]"
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": call_id,
+ "name": name,
+ "content": obs_text,
+ }
+ )
+ if n_tool_calls >= _SOFT_CALL_CAP:
+ messages.append(
+ {
+ "role": "user",
+ "content": (
+ "You now have enough information from the models. Do NOT call any "
+ "more tools — reply with your FINAL_ANSWER line only."
+ ),
+ }
+ )
+ _trim_history(messages)
+ else:
+ # Hit max_turns with no explicit answer: use the last observation/text.
+ final_answer = (
+ (turns[-1].observation or turns[-1].reasoning).strip() if turns else ""
+ )
+
+ return UnifiedRollout(
+ turns=turns,
+ final_answer=final_answer,
+ cost_usd=cost,
+ tokens=tokens,
+ num_tool_calls=n_tool_calls,
+ parse_failures=parse_failures,
+ anon_map=anon_map,
+ tool_specs=specs,
+ )
+
+
+def run_unified_rollout(
+ question: str,
+ tools: List[ExpertTool],
+ *,
+ call_orchestrator: Callable[
+ ..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
+ ],
+ dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
+ max_turns: int = 50,
+ system: str = RL_ORCHESTRATOR_SYS,
+ anonymize: bool = False,
+) -> UnifiedRollout:
+ """One ToolOrchestra rollout = one Braintrust trace. Opens the root span,
+ runs the loop (orchestrator turns + expert/tool dispatches nest inside as
+ child spans), then logs the final answer + cost/tokens. No-op when tracing
+ is disabled (see :mod:`.tracing`)."""
+ _run_meta, _run_tags = run_context()
+ _fields = dict(
+ input={"question": question},
+ metadata={
+ "max_turns": max_turns,
+ "anonymize": anonymize,
+ "n_tools": len(tools),
+ **_run_meta,
+ },
+ )
+ if _run_tags:
+ _fields["tags"] = _run_tags
+ with span("toolorchestra.rollout", span_type="task", **_fields) as _s:
+ roll = _run_unified_rollout_inner(
+ question,
+ tools,
+ call_orchestrator=call_orchestrator,
+ dispatch=dispatch,
+ max_turns=max_turns,
+ system=system,
+ anonymize=anonymize,
+ )
+ _s.log(
+ output=roll.final_answer,
+ metrics={
+ "cost_usd": roll.cost_usd,
+ "tokens": roll.tokens,
+ "num_tool_calls": roll.num_tool_calls,
+ "parse_failures": roll.parse_failures,
+ },
+ metadata={
+ "n_experts_available": len(roll.anon_map or {}),
+ "answered": bool(roll.final_answer),
+ # label -> real model, so every anonymized route span in this
+ # trace can be decoded back to the model that actually ran.
+ "anon_map": roll.anon_map or {},
+ },
+ )
+ return roll
+
+
+def tool_call_tag(name: str, arguments: Dict[str, object]) -> str:
+ """Render a tool call as the ``{...}`` text the model emits."""
+ return (
+ f"{json.dumps({'name': name, 'arguments': arguments})}"
+ )
+
+
+__all__ = [
+ "RL_ORCHESTRATOR_SYS",
+ "UnifiedRollout",
+ "UnifiedTurn",
+ "build_system_prompt",
+ "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..bbe28c6d
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py
@@ -0,0 +1,82 @@
+"""Tavily search + Modal Python sandbox helpers for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+import re
+from typing import Optional, Tuple
+
+# ---- Tavily + Modal helpers -------------------------------------------------
+
+
+def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]:
+ """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0).
+
+ Token counts are reported as zero (no LLM was billed); the OpenJarvis
+ accounting layer separately tallies tool-call counts. Falls back to
+ DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``).
+ """
+ from openjarvis.tools.web_search import WebSearchTool
+
+ tool = WebSearchTool(max_results=max_results)
+ res = tool.execute(query=query, max_results=max_results)
+ text = res.content or ""
+ if not res.success and not text:
+ text = "(no results)"
+ return text, 0, 0
+
+
+_MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox"
+
+
+def _call_modal_python(code: str, timeout_s: int = 60) -> Tuple[str, int]:
+ """Execute a single Python snippet in a fresh Modal Sandbox.
+
+ Returns ``(combined_stdout_stderr, returncode)``. Logs are capped at 8 KiB.
+ Any exception (modal auth, network, sandbox boot failure) is captured into
+ the returned string with a non-zero rc — we never raise back to the
+ orchestrator loop. The sandbox is torn down at the end via ``terminate()``.
+ """
+ try:
+ import modal
+
+ app = modal.App.lookup(_MODAL_APP_NAME, create_if_missing=True)
+ # python:3.12-slim is small + boots fast; the paper uses a generic
+ # Python image too. We rely on stdlib only — no extra pip installs.
+ image = modal.Image.debian_slim(python_version="3.12")
+ sb = modal.Sandbox.create(
+ "python",
+ "-c",
+ code,
+ app=app,
+ image=image,
+ timeout=int(timeout_s),
+ )
+ sb.wait()
+ try:
+ out = sb.stdout.read() or ""
+ except Exception:
+ out = ""
+ try:
+ err = sb.stderr.read() or ""
+ except Exception:
+ err = ""
+ rc = sb.returncode if sb.returncode is not None else -1
+ try:
+ sb.terminate()
+ except Exception:
+ pass
+ combined = out + (("\n" + err) if err else "")
+ if len(combined) > 8192:
+ combined = combined[:8192] + "\n... (output truncated)"
+ return combined, int(rc)
+ except Exception as exc:
+ return f"[modal-python error: {type(exc).__name__}: {exc}]", -1
+
+
+_PY_CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL)
+
+
+def _extract_first_python_block(text: str) -> Optional[str]:
+ """Return the first ```python ... ``` block (or ```...```), or None."""
+ m = _PY_CODE_RE.search(text or "")
+ return m.group(1).strip() if m else None
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/tracing.py b/src/openjarvis/agents/hybrid/toolorchestra/tracing.py
new file mode 100644
index 00000000..86e5935f
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/tracing.py
@@ -0,0 +1,170 @@
+"""Braintrust telemetry for the ToolOrchestra rollout.
+
+Each ``run_unified_rollout`` becomes ONE Braintrust trace (root span); the
+orchestrator turns and every expert/tool dispatch nest inside it, and the
+underlying OpenAI/Anthropic calls (wrapped clients) nest one level deeper — so
+you see the full routing tree with inputs/outputs/tokens/cost per node.
+
+On by default. Degrades to a total no-op (never raises, never changes behavior)
+when: ``OJ_BRAINTRUST=0``, the ``braintrust`` package isn't installed, or
+``BRAINTRUST_API_KEY`` is unset. Concurrency: rollouts run one-per-thread
+(ThreadPoolExecutor); threads don't inherit contextvars, so each rollout opens
+its own independent root trace and same-thread child calls nest correctly.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import os
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_STATE: dict = {"resolved": False, "enabled": False, "bt": None}
+
+
+def _truthy(v: str) -> bool:
+ return v.strip().lower() not in ("0", "false", "no", "off", "")
+
+
+def _resolve() -> bool:
+ """Lazily decide whether tracing is active and init the logger once."""
+ if _STATE["resolved"]:
+ return _STATE["enabled"]
+ _STATE["resolved"] = True
+ if not _truthy(os.getenv("OJ_BRAINTRUST", "1")): # on by default
+ return False
+ if not os.getenv("BRAINTRUST_API_KEY"):
+ logger.info(
+ "braintrust on-by-default but BRAINTRUST_API_KEY unset — tracing disabled"
+ )
+ return False
+ try:
+ import braintrust as _bt
+
+ proj_id = os.getenv("OJ_BRAINTRUST_PROJECT_ID")
+ if proj_id:
+ _bt.init_logger(project_id=proj_id)
+ else:
+ _bt.init_logger(project=os.getenv("OJ_BRAINTRUST_PROJECT", "toolorchestra"))
+ _STATE["bt"] = _bt
+ _STATE["enabled"] = True
+ logger.info(
+ "braintrust tracing ENABLED (%s)",
+ f"project_id={proj_id}"
+ if proj_id
+ else f"project={os.getenv('OJ_BRAINTRUST_PROJECT', 'toolorchestra')}",
+ )
+ except Exception as exc: # missing pkg / bad key / init failure — never crash
+ logger.warning("braintrust init failed (%s) — tracing disabled", exc)
+ return _STATE["enabled"]
+
+
+def enabled() -> bool:
+ return _resolve()
+
+
+def run_context() -> tuple[dict, list]:
+ """Run-level metadata + tags for the ROOT rollout span, sourced from env
+ (set by the generation driver). No-op-safe: returns ({}, []) if nothing is
+ set, and never raises. Env keys:
+
+ - ``OJ_RUN_LABEL`` — human run label (also stamped on the uploaded dataset).
+ - ``OJ_GEN_MODEL`` — SPECIFIC gen model id (e.g. ``Qwen/Qwen3.5-9B``).
+ - ``OJ_RUN_STAGE`` — ``prod``|``smoke`` (inferred from the label if unset).
+ - ``OJ_CFG_*`` — config knobs (temperature/max_turns/anonymize/...).
+ """
+ import datetime
+
+ meta: dict = {}
+ tags: list = []
+ try:
+ label = os.getenv("OJ_RUN_LABEL")
+ gen_model = os.getenv("OJ_GEN_MODEL")
+ stage = os.getenv("OJ_RUN_STAGE")
+ if not stage and label:
+ stage = "smoke" if "smoke" in label.lower() else "prod"
+ if label:
+ meta["run_label"] = label
+ if gen_model:
+ meta["gen_model"] = gen_model
+ if stage:
+ meta["stage"] = stage
+ cfg = {}
+ for env_key, key in (
+ ("OJ_CFG_TEMPERATURE", "temperature"),
+ ("OJ_CFG_MAX_TURNS", "max_turns"),
+ ("OJ_CFG_ANONYMIZE", "anonymize"),
+ ("OJ_CFG_REJECTION_ONLY", "rejection_only"),
+ ):
+ v = os.getenv(env_key)
+ if v not in (None, ""):
+ cfg[key] = v
+ if cfg:
+ meta["config"] = cfg
+ date = datetime.date.today().isoformat()
+ tags = [t for t in (gen_model, stage, date) if t]
+ except Exception: # never let telemetry enrichment break a rollout
+ return {}, []
+ return meta, tags
+
+
+def wrap_client(client: Any) -> Any:
+ """Wrap an OpenAI/Anthropic client so its calls auto-log under the current
+ span. Pass-through (unchanged client) when tracing is off or on any error —
+ so this is always safe to call at client construction."""
+ if not _resolve():
+ return client
+ try:
+ bt = _STATE["bt"]
+ mod = type(client).__module__.lower()
+ if "anthropic" in mod:
+ return bt.wrap_anthropic(client)
+ return bt.wrap_openai(client)
+ except Exception as exc:
+ logger.warning("braintrust wrap_client failed (%s) — using raw client", exc)
+ return client
+
+
+class _NullSpan:
+ """No-op span used when tracing is disabled (keeps call sites branch-free)."""
+
+ def log(self, **_kw: Any) -> None:
+ pass
+
+ def __enter__(self) -> "_NullSpan":
+ return self
+
+ def __exit__(self, *_a: Any) -> bool:
+ return False
+
+
+@contextlib.contextmanager
+def span(name: str, *, span_type: str = "task", **fields: Any):
+ """Open a Braintrust span (or a no-op). Use as::
+
+ with span("toolorchestra.rollout", input=...) as s:
+ ...
+ s.log(output=..., metrics=..., metadata=...)
+
+ ``fields`` (input/metadata/...) are forwarded to ``start_span``.
+ """
+ if not _resolve():
+ yield _NullSpan()
+ return
+ # Guard only span CREATION (so a Braintrust hiccup degrades to a no-op). Do
+ # NOT wrap the yielded body in try/except: an exception from the rollout would
+ # be thrown into this generator, caught here, and masked by a second yield
+ # ("generator didn't stop after throw()"). Let the body's own exceptions
+ # propagate through the `with` untouched — Braintrust records + re-raises.
+ try:
+ cm = _STATE["bt"].start_span(name=name, type=span_type, **fields)
+ except Exception as exc: # span creation failed — run trace-less, never break
+ logger.warning(
+ "braintrust start_span(%s) failed (%s) — continuing untraced", name, exc
+ )
+ yield _NullSpan()
+ return
+ with cm as s:
+ yield s
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/unified.py b/src/openjarvis/agents/hybrid/toolorchestra/unified.py
new file mode 100644
index 00000000..704e18c4
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/unified.py
@@ -0,0 +1,280 @@
+"""Real backends for the unified-tool rollout — the bridge between the pure
+:func:`run_unified_rollout` loop and live model/tool calls.
+
+``make_call_orchestrator`` returns the ``call_orchestrator`` callable (a teacher
+LLM emitting tool calls over the unified spec); ``make_dispatch`` returns the
+``dispatch`` callable (executes a chosen tool via ``_call_worker``). Both are
+import-safe — the OpenAI SDK is imported lazily so this module loads without
+network or keys.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from openjarvis.agents.hybrid._prices import cost as _model_cost
+from openjarvis.agents.hybrid.expert_registry import KIND_MODEL, ExpertTool, to_worker_dict
+from openjarvis.agents.hybrid.retry import EmptyExpertResponse, with_backoff
+from openjarvis.agents.hybrid.toolorchestra.parsing import _parse_rl_tool_call
+from openjarvis.agents.hybrid.toolorchestra.rollout import (
+ build_system_prompt,
+)
+from openjarvis.agents.hybrid.toolorchestra.tracing import span
+from openjarvis.agents.hybrid.toolorchestra.workers import _call_worker
+
+# Guards the lazy, one-time build of the shared ToolExecutor in
+# ``_dispatch_openjarvis_tool``: with concurrent rollouts (parallel rejection
+# sampling) several threads can reach the build at once and would each instantiate
+# the full tool registry. The lock makes it build-once.
+_EXECUTOR_BUILD_LOCK = threading.Lock()
+
+
+def make_call_orchestrator(
+ model: str,
+ *,
+ base_url: Optional[str] = None,
+ api_key: Optional[str] = None,
+ temperature: float = 1.0,
+ max_tokens: int = 4096,
+ timeout: float = 600.0,
+ native_tools: bool = True,
+ frequency_penalty: float = 0.2,
+ presence_penalty: float = 0.0,
+ repetition_penalty: float = 1.15,
+) -> Callable[..., Tuple[str, List[Tuple[str, Dict[str, Any]]], int, int]]:
+ """Teacher-orchestrator caller. ``base_url=None`` → OpenAI cloud; set it to a
+ vLLM endpoint (with ``api_key="EMPTY"``) to drive a local served teacher.
+
+ ``native_tools`` picks the tool-calling convention, and the two modes MUST NOT
+ be mixed on the same model:
+
+ * ``True`` (default — DATA GENERATION with a *base* model): pass the OpenAI
+ ``tools=`` param so the server's native tool template + parser (gemma4 /
+ qwen3_xml / hermes) drive the base model, which only emits reliable tool
+ calls that way. The rollout captures ``(name, arguments)`` and the serializer
+ re-renders them into the canonical JSON ```` form regardless.
+ * ``False`` (SERVING a *fine-tuned* model): bake the ```` block + JSON
+ ```` format into the system prompt (exactly what the serializer
+ trained on) and DROP ``tools=``. If ``tools=`` is set here, the served chat
+ template injects its own XML ```` instructions, which conflict
+ with the JSON format the model learned -> it falls back to ``\\boxed{}`` and
+ never routes. ``_parse_rl_tool_call`` scrapes the JSON tag from raw text.
+ """
+
+ # Create the client ONCE and reuse it for every turn. Creating a fresh OpenAI
+ # client per call (as before) leaks an httpx connection pool each time -> under
+ # heavy parallel generation, sockets pile up in CLOSE-WAIT, the run degrades and
+ # has to be restarted. The orchestrator is the highest-frequency call, so reusing
+ # one client here removes the dominant leak. httpx clients are thread-safe.
+ from openai import OpenAI
+
+ from openjarvis.agents.hybrid.toolorchestra.tracing import wrap_client
+
+ client = wrap_client(
+ OpenAI(base_url=base_url, api_key=api_key or "EMPTY", timeout=timeout)
+ )
+
+ def call_orchestrator(messages: List[Dict[str, Any]], specs: List[Dict[str, Any]]):
+ # ``messages`` is the full running conversation (system/user/assistant/tool)
+ # built by run_unified_rollout — pass it through so the model sees its own
+ # prior calls + tool responses, not a flattened user blob.
+ if native_tools:
+ # Base-model generation: native tool template drives the tool calls.
+ send = messages
+ kwargs = {"tools": specs} if specs else {}
+ else:
+ # Fine-tuned serving: JSON format is baked into the system prompt and
+ # tools= is dropped (see make_call_orchestrator docstring).
+ send = list(messages)
+ if specs and send and send[0].get("role") == "system":
+ send[0] = {"role": "system", "content": build_system_prompt(specs)}
+ kwargs = {}
+ # repetition_penalty is a vLLM extra (not OpenAI-native); only send it to a
+ # local vLLM endpoint (base_url set), never to the cloud frontier APIs.
+ if base_url and repetition_penalty and repetition_penalty != 1.0:
+ kwargs.setdefault("extra_body", {})["repetition_penalty"] = (
+ repetition_penalty
+ )
+ # Retry transient failures (429 / 5xx / connection) with exponential
+ # backoff + jitter. Without this a single rate-limit kills the rollout —
+ # fine at concurrency 12, ruinous at 100.
+ resp = with_backoff(
+ lambda: client.chat.completions.create(
+ model=model,
+ messages=send,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ frequency_penalty=frequency_penalty,
+ presence_penalty=presence_penalty,
+ **kwargs,
+ ),
+ what=f"orchestrator call ({model})",
+ )
+ msg = resp.choices[0].message
+ text = msg.content or ""
+ sdk_tool_calls = getattr(msg, "tool_calls", None)
+ u = resp.usage
+ p = getattr(u, "prompt_tokens", 0) if u else 0
+ c = getattr(u, "completion_tokens", 0) if u else 0
+
+ parsed = _parse_rl_tool_call(text, sdk_tool_calls)
+ tool_calls = [(parsed["name"], parsed["arguments"])] if parsed else []
+ return text, tool_calls, int(p), int(c)
+
+ return call_orchestrator
+
+
+def _dispatch_openjarvis_tool(
+ tool: ExpertTool,
+ arguments: Dict[str, Any],
+ executor_holder: Dict[str, Any],
+) -> Tuple[str, float, int, bool]:
+ """Execute a bridged real OpenJarvis tool via its ``ToolExecutor``.
+
+ Builds the executor lazily (cached in ``executor_holder``) and degrades to a
+ clear error string — never a crash — if the tool registry / executor can't be
+ instantiated. Returns ``(content, cost_usd, total_tokens, is_local=True)``.
+ """
+ try:
+ executor = executor_holder.get("executor")
+ if executor is None:
+ with _EXECUTOR_BUILD_LOCK:
+ # Re-check inside the lock: another thread may have built it while
+ # we waited.
+ executor = executor_holder.get("executor")
+ if executor is None:
+ from openjarvis.core.registry import ToolRegistry
+ from openjarvis.tools._stubs import ToolExecutor
+
+ instances = []
+ for name in ToolRegistry.keys():
+ entry = ToolRegistry.get(name)
+ try:
+ instances.append(
+ entry() if isinstance(entry, type) else entry
+ )
+ except Exception:
+ continue
+ # Auto-approve confirmation-gated tools (shell_exec,
+ # git_commit, ...): rollouts/eval run headless in a sandbox, so
+ # there is no human to confirm. Without this, those tools return
+ # a "requires confirmation" error instead of executing — which
+ # would silently break TerminalBench.
+ executor = ToolExecutor(
+ instances,
+ interactive=True,
+ confirm_callback=lambda _prompt: True,
+ )
+ executor_holder["executor"] = executor
+
+ from openjarvis.core.types import ToolCall
+
+ call = ToolCall(
+ id=f"orch-{tool.name}",
+ name=str(tool.model),
+ arguments=json.dumps(arguments or {}),
+ )
+ result = executor.execute(call)
+ usage = getattr(result, "usage", None) or {}
+ total_tokens = int(usage.get("total_tokens", 0) or 0)
+ return (
+ str(getattr(result, "content", "")),
+ float(getattr(result, "cost_usd", 0.0) or 0.0),
+ total_tokens,
+ True,
+ )
+ except Exception as exc: # never crash the rollout on a tool-bridge failure
+ return (f"[openjarvis-tool error: {tool.model}: {exc}]", 0.0, 0, True)
+
+
+def make_dispatch(
+ cfg: Optional[Dict[str, Any]] = None,
+) -> Callable[[ExpertTool, Dict[str, Any]], Tuple[str, float, int, bool]]:
+ """Tool-execution caller: run the chosen tool and return (obs, cost, tokens, is_local)."""
+ cfg = cfg or {}
+ executor_holder: Dict[str, Any] = {}
+
+ def _dispatch_inner(tool: ExpertTool, arguments: Dict[str, Any]):
+ # Bridged real OpenJarvis tools run through the ToolExecutor, not
+ # the model/worker path.
+ if tool.backend_type == "openjarvis-tool":
+ return _dispatch_openjarvis_tool(tool, arguments or {}, executor_holder)
+
+ worker = to_worker_dict(tool)
+ prompt = ""
+ for key in ("input", "query", "code"):
+ val = arguments.get(key)
+ if isinstance(val, str) and val.strip():
+ prompt = val
+ break
+ if not prompt.strip():
+ # The orchestrator emitted a tool call with no/empty input. Don't hit
+ # the API (it 400s on empty content) — return a usable error so the
+ # rollout keeps going instead of dropping.
+ return (
+ f"[{tool.name}: no input provided — supply a non-empty "
+ "'input' to delegate]",
+ 0.0,
+ 0,
+ True,
+ )
+ # Retry transient expert failures (429 / 5xx / connection) AND the sneaky
+ # one: a 200-OK with an empty body. Nothing raises on that, so no retry
+ # fires, the rollout gets an empty observation, and the clean gate bins the
+ # whole trajectory. The OpenRouter-hosted Qwen 122B/397B do this (audit
+ # 2026-07-13: 6 of 72 rollouts lost).
+ #
+ # MODEL EXPERTS ONLY. A *sandbox* tool returning nothing is not a failure —
+ # `code_interpreter` legitimately prints nothing when the code just defines
+ # a function. Retrying that burned 18 of 72 retry-chains (~2 min of a
+ # blocked worker thread each) and then killed the rollout outright. Only an
+ # expert that answers with silence is broken.
+ is_model_expert = tool.kind == KIND_MODEL
+
+ def _call_expert():
+ out = _call_worker(worker, prompt, cfg)
+ if is_model_expert and not (out[0] or "").strip():
+ raise EmptyExpertResponse(str(tool.model))
+ return out
+
+ text, p, c, is_local, extra_cost, _n = with_backoff(
+ _call_expert,
+ what=f"expert call ({tool.model})",
+ )
+ usd = (0.0 if is_local else _model_cost(str(tool.model), p, c)) + float(
+ extra_cost
+ )
+ return text, usd, int(p) + int(c), bool(is_local)
+
+ def dispatch(tool: ExpertTool, arguments: Dict[str, Any]):
+ # One span per route so the trace tree shows which expert/tool was called,
+ # with the delegated input, the returned observation, and cost/tokens. The
+ # wrapped OpenAI/Anthropic call (if any) nests one level deeper.
+ # Span is titled with the REAL model that actually ran (``tool.model``) so
+ # the trace reads clearly even under anonymization; ``anon_label`` in the
+ # metadata records the opaque label the orchestrator actually saw/chose.
+ real_model = str(tool.model)
+ with span(
+ f"route:{real_model}",
+ span_type="tool",
+ input=arguments,
+ metadata={
+ "real_model": real_model,
+ "anon_label": tool.name,
+ "backend": tool.backend_type,
+ },
+ ) as s:
+ obs, usd, toks, is_local = _dispatch_inner(tool, arguments)
+ s.log(
+ output=obs,
+ metrics={"cost_usd": usd, "tokens": toks},
+ metadata={"is_local": is_local},
+ )
+ return obs, usd, toks, is_local
+
+ return dispatch
+
+
+__all__ = ["make_call_orchestrator", "make_dispatch"]
diff --git a/src/openjarvis/agents/hybrid/toolorchestra/workers.py b/src/openjarvis/agents/hybrid/toolorchestra/workers.py
new file mode 100644
index 00000000..14f2255b
--- /dev/null
+++ b/src/openjarvis/agents/hybrid/toolorchestra/workers.py
@@ -0,0 +1,477 @@
+"""Worker pool resolution + dispatch for ToolOrchestraAgent."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from openjarvis.agents.hybrid._base import (
+ ANTHROPIC_WEB_SEARCH_TOOL,
+ GEMINI_SEARCH_COST_PER_CALL,
+ OPENAI_WEB_SEARCH_COST_PER_CALL,
+ WEB_SEARCH_COST_PER_CALL,
+ LocalCloudAgent,
+)
+from openjarvis.agents.hybrid._prices import (
+ PRICES,
+ is_gpt5_family,
+ supports_temperature,
+)
+from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
+from openjarvis.agents.hybrid.toolorchestra.sandbox import (
+ _call_modal_python,
+ _call_tavily_search,
+)
+
+
+def _default_pool(
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str = "claude-opus-4-7",
+ cloud_endpoint: str = "anthropic",
+) -> List[Dict[str, Any]]:
+ """Default heterogeneous worker pool.
+
+ The frontier worker's ``type`` + ``model`` track the cell's resolved
+ ``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5,
+ gemini-2.5-pro, …) route their frontier slot to the right SDK.
+ """
+ ep = (cloud_endpoint or "anthropic").lower()
+ if ep not in ("anthropic", "openai", "gemini"):
+ ep = "anthropic"
+ pool: List[Dict[str, Any]] = []
+ if local_model and local_endpoint:
+ pool.append(
+ {
+ "id": len(pool),
+ "name": "local-qwen",
+ "type": "vllm",
+ "model": local_model,
+ "base_url": local_endpoint,
+ "description": (
+ "Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
+ "concise extraction, formatting, arithmetic on given data."
+ ),
+ }
+ )
+ if ep == "openai":
+ search_type = "openai-web-search"
+ search_model = cloud_model
+ search_desc = "OpenAI hosted web search on the configured frontier model."
+ elif ep == "gemini":
+ search_type = "gemini-web-search"
+ search_model = cloud_model
+ search_desc = "Gemini Google Search grounding on the configured frontier model."
+ else:
+ search_type = "anthropic-web-search"
+ search_model = _DEFAULT_WEB_SEARCH_MODEL
+ search_desc = "Anthropic server-side web_search."
+ pool.append(
+ {
+ "id": len(pool),
+ "name": "web-search",
+ "type": search_type,
+ "model": search_model,
+ "description": (
+ f"{search_desc} Use for facts that need a lookup "
+ "(recent events, rare names/dates, niche sources). Returns a digest."
+ ),
+ }
+ )
+ pool.append(
+ {
+ "id": len(pool),
+ "name": f"frontier-{ep}",
+ "type": ep,
+ "model": cloud_model,
+ "description": (
+ "Frontier reasoning model. Use for hard multi-step reasoning, "
+ "code review, or a final synthesis pass. Expensive — use sparingly."
+ ),
+ }
+ )
+ pool.append(
+ {
+ "id": len(pool),
+ "name": "frontier-openai-mini",
+ "type": "openai",
+ "model": "gpt-5-mini",
+ "description": (
+ "Mid-tier OpenAI model. Solid general knowledge and reasoning at a "
+ "fraction of frontier cost."
+ ),
+ }
+ )
+ return pool
+
+
+# Worker types toolorchestra's `_call_worker` actually dispatches.
+#
+# Paper-match additions (2026-05-19) — opt in via `method_cfg.pool = "paper"`:
+# `tavily-search` — Tavily API search (the paper's web tool).
+# `openrouter` — OpenAI-compatible client at openrouter.ai/api/v1.
+# Used for the code/math specialists and Llama-3.3-70B /
+# Qwen3-32B generalists.
+# `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the
+# paper's "Python sandbox" inside `enhance_reasoning`).
+_TOOLORCH_VALID_TYPES = (
+ "vllm",
+ "openai",
+ "anthropic",
+ "anthropic-web-search",
+ "openai-web-search",
+ "gemini",
+ "gemini-web-search",
+ "tavily-search",
+ "openrouter",
+ "modal-python",
+)
+_TOOLORCH_SEARCH_TYPES = (
+ "anthropic-web-search",
+ "openai-web-search",
+ "gemini-web-search",
+ "tavily-search",
+)
+
+# Default model used when an `anthropic-web-search` entry omits `model`.
+_DEFAULT_WEB_SEARCH_MODEL = "claude-haiku-4-5"
+
+
+def _resolve_worker_pool(
+ cfg: Dict[str, Any],
+ local_model: Optional[str],
+ local_endpoint: Optional[str],
+ cloud_model: str,
+ cloud_endpoint: str = "anthropic",
+) -> List[Dict[str, Any]]:
+ """Return the worker pool for this run.
+
+ Strict replace, not merge: if ``cfg["worker_pool"]`` is set, the
+ default pool is ignored entirely. Falls back to ``_default_pool`` when
+ the override is absent.
+
+ Each user-supplied entry must be a dict with keys ``id``, ``name``,
+ ``type``, and (for non-search types) ``model``. Search worker types are
+ ``anthropic-web-search``, ``openai-web-search``, ``gemini-web-search``,
+ and ``tavily-search``. ``anthropic-web-search`` entries may omit
+ ``model`` — it defaults to ``claude-haiku-4-5``. OpenAI and Gemini
+ search workers default to the configured cloud model. Tavily does not
+ require a model.
+
+ Substitution: ``model = "$local"`` (or ``""``) 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 in _TOOLORCH_SEARCH_TYPES:
+ if model in (None, ""):
+ if wtype == "anthropic-web-search":
+ model = _DEFAULT_WEB_SEARCH_MODEL
+ elif wtype in ("openai-web-search", "gemini-web-search"):
+ model = cloud_model
+ else:
+ model = wtype
+ entry["model"] = model
+ elif not isinstance(model, str):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
+ )
+ if (
+ wtype in ("openai-web-search", "gemini-web-search")
+ and model not in PRICES
+ ):
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: model {model!r} "
+ f"is not in PRICES (known: {sorted(PRICES)})"
+ )
+ # Search workers don't satisfy the "needs a solver" requirement.
+ else:
+ if not isinstance(model, str) or not model:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
+ )
+ if wtype == "vllm":
+ if not entry.get("base_url"):
+ if not local_endpoint:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: vllm worker needs "
+ "'base_url' (or a configured local_endpoint to fall back to)"
+ )
+ entry["base_url"] = local_endpoint
+ entry.setdefault("api_key", "EMPTY")
+ else:
+ if model not in PRICES:
+ raise ValueError(
+ f"Invalid worker_pool entry [{wid}]: model {model!r} "
+ f"is not in PRICES (known: {sorted(PRICES)})"
+ )
+ has_non_search = True
+ entry.setdefault(
+ "description",
+ f"User-supplied {wtype} worker ({model}).",
+ )
+ resolved.append(entry)
+
+ if not has_non_search:
+ raise ValueError(
+ "Invalid worker_pool entry [-]: worker_pool must contain at least "
+ "one non-search worker (vllm / openai / anthropic / gemini)"
+ )
+ return resolved
+
+
+# Anthropic model id -> OpenRouter slug (for OJ_ANTHROPIC_VIA_OPENROUTER=1).
+_ANTHROPIC_OPENROUTER_SLUGS = {
+ "claude-opus-4-8": "anthropic/claude-opus-4.8",
+ "claude-opus-4-7": "anthropic/claude-opus-4.7",
+ "claude-haiku-4-5-20251001": "anthropic/claude-haiku-4.5",
+}
+
+
+def _call_worker(
+ worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
+) -> Tuple[str, int, int, bool, float, int]:
+ """Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches)."""
+ wtype = worker.get("type", "openai")
+ max_tok = int(
+ cfg.get("worker_max_tokens") or os.environ.get("OJ_WORKER_MAX_TOKENS", "4096")
+ )
+ temp = float(cfg.get("worker_temperature", 0.2))
+
+ if wtype == "vllm":
+ text, p, c = LocalCloudAgent._call_vllm(
+ worker["model"],
+ worker["base_url"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ enable_thinking=False,
+ )
+ return text, p, c, True, 0.0, 0
+ if wtype == "openai":
+ is_gpt5 = is_gpt5_family(worker["model"])
+ eff_temp = 1.0 if is_gpt5 else temp
+ # GPT-5 is a reasoning model: hidden reasoning tokens count against
+ # `max_completion_tokens`, so a 4096 cap can be fully consumed by
+ # reasoning and leave 0 visible content (empty answer). Give the
+ # reasoning headroom on top of the answer budget.
+ eff_max_tok = max(max_tok, 16384) if is_gpt5 else max_tok
+ text, p, c = LocalCloudAgent._call_openai(
+ worker["model"],
+ user=prompt,
+ max_tokens=eff_max_tok,
+ temperature=eff_temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "gemini":
+ text, p, c = LocalCloudAgent._call_gemini(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "anthropic":
+ # Escape hatch for a dead/rotated ANTHROPIC_API_KEY: the same Claude
+ # models are served BY Anthropic through OpenRouter, billed on the
+ # OpenRouter key. Flip OJ_ANTHROPIC_VIA_OPENROUTER=1 and every
+ # anthropic-typed expert keeps working (2026-07-13: key rotated
+ # mid-eval and zeroed a run — the judge 401'd and every answer
+ # defaulted to score 0).
+ if os.environ.get("OJ_ANTHROPIC_VIA_OPENROUTER") == "1":
+ slug = _ANTHROPIC_OPENROUTER_SLUGS.get(
+ worker["model"], f"anthropic/{worker['model']}"
+ )
+ text, p, c = LocalCloudAgent._call_openrouter(
+ slug, user=prompt, max_tokens=max_tok, temperature=temp
+ )
+ return text, p, c, False, 0.0, 0
+ eff_temp = temp if supports_temperature(worker["model"]) else 0.0
+ text, p, c, _ = LocalCloudAgent._call_anthropic(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=eff_temp,
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "anthropic-web-search":
+ eff_temp = temp if supports_temperature(worker["model"]) else 0.0
+ text, p, c, n_searches = LocalCloudAgent._call_anthropic(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=eff_temp,
+ tools=[ANTHROPIC_WEB_SEARCH_TOOL],
+ tool_choice={"type": "any"},
+ )
+ extra = n_searches * WEB_SEARCH_COST_PER_CALL
+ return text, p, c, False, extra, n_searches
+ if wtype == "openai-web-search":
+ eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp
+ text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent(
+ worker["model"],
+ user=prompt,
+ max_tokens=max(max_tok, 16384)
+ if is_gpt5_family(worker["model"])
+ else max_tok,
+ temperature=eff_temp,
+ )
+ extra = n_searches * OPENAI_WEB_SEARCH_COST_PER_CALL
+ return text, p, c, False, extra, n_searches
+ if wtype == "gemini-web-search":
+ text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ )
+ extra = n_searches * GEMINI_SEARCH_COST_PER_CALL
+ return text, p, c, False, extra, n_searches
+ if wtype == "tavily-search":
+ # Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL`
+ # for parity with the Anthropic web-search worker. One call = one
+ # "n_search" for accounting.
+ max_results = int(cfg.get("tavily_max_results", 5))
+ text, p, c = _call_tavily_search(str(prompt), max_results=max_results)
+ return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1
+ if wtype == "openrouter":
+ # Pin to providers measured to actually return completions. OpenRouter
+ # load-balances across upstreams, and several return 200-OK with an EMPTY
+ # body (measured 2026-07-13: Chutes 2/2 empty, SiliconFlow 1/1,
+ # DigitalOcean 1/1, Parasail 1/1 — a blocklist was whack-a-mole, new
+ # broken providers kept appearing). Allowlist verified 12/12 non-empty;
+ # ``allow_fallbacks: False`` stops OpenRouter from silently routing
+ # outside it. Retries only ever "worked" by re-rolling this dice.
+ text, p, c = LocalCloudAgent._call_openrouter(
+ worker["model"],
+ user=prompt,
+ max_tokens=max_tok,
+ temperature=temp,
+ extra_body={
+ "provider": {
+ "order": ["Alibaba", "Novita", "GMICloud", "AtlasCloud"],
+ "allow_fallbacks": False,
+ }
+ },
+ )
+ return text, p, c, False, 0.0, 0
+ if wtype == "modal-python":
+ # `prompt` is the python code string to exec.
+ timeout_s = int(cfg.get("modal_python_timeout_s", 60))
+ out, _rc = _call_modal_python(str(prompt), timeout_s=timeout_s)
+ # No LLM tokens consumed; report 0 in/out. Cost is whatever Modal
+ # charges per sandbox-second — not tracked here.
+ return out, 0, 0, False, 0.0, 0
+ raise ValueError(f"unsupported worker type: {wtype!r}")
+
+
+def _swe_call_worker(
+ worker: Dict[str, Any],
+ prompt: str,
+ cfg: Dict[str, Any],
+ task: Dict[str, Any],
+ workdir: Path,
+ turn: int,
+) -> Tuple[str, int, int, bool, float, int, int]:
+ """SWE-bench worker dispatch: route solver workers through
+ run_swe_agent_loop on a shared workdir. Web-search workers fall back
+ to the regular one-shot dispatch (search isn't an agent loop).
+
+ Trailing ``bash_turns`` (last element) counts agent-loop turns so the
+ caller can surface ``tool_calls`` per row. Fallbacks to one-shot
+ workers return 0 bash turns (no agent loop ran)."""
+ wtype = worker.get("type", "openai")
+ if wtype in _TOOLORCH_SEARCH_TYPES:
+ # Search workers stay one-shot.
+ text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
+ return text, p, c, is_local, extra, n_searches, 0
+ if wtype == "vllm":
+ backbone = "local"
+ endpoint = worker.get("base_url")
+ loop_cloud_endpoint = "anthropic" # unused when backbone=local
+ elif wtype in ("anthropic", "openai", "gemini"):
+ backbone = "cloud"
+ endpoint = None
+ loop_cloud_endpoint = wtype
+ else:
+ # Unknown type — one-shot fallback.
+ text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg)
+ return text, p, c, is_local, extra, n_searches, 0
+ out = run_swe_agent_loop(
+ task,
+ backbone=backbone,
+ backbone_model=worker["model"],
+ cloud_endpoint=loop_cloud_endpoint,
+ local_endpoint=endpoint,
+ initial_prompt=prompt,
+ max_turns=int(cfg.get("swe_max_turns", 30)),
+ bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)),
+ output_cap=int(cfg.get("swe_output_cap", 10_000)),
+ turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)),
+ trace_prefix=f"toolorch_turn{turn}",
+ workdir=workdir,
+ )
+ is_local = backbone == "local"
+ return (
+ out["final_summary"] or out["answer"],
+ out["tokens_in"],
+ out["tokens_out"],
+ is_local,
+ 0.0,
+ 0,
+ int(out["turns"]),
+ )
diff --git a/tests/agents/test_anonymize_tools.py b/tests/agents/test_anonymize_tools.py
new file mode 100644
index 00000000..07a9391d
--- /dev/null
+++ b/tests/agents/test_anonymize_tools.py
@@ -0,0 +1,129 @@
+"""Unit coverage for ``anonymize_tools`` — the identity-stripping step that makes
+routing data unbiased. See ``expert_registry.anonymize_tools``.
+
+The anonymizer takes the orchestrator catalog and replaces every MODEL expert
+with an opaque ``model_xxxx`` label, a uniform brand-free description, no
+price/latency line, and shuffles the expert block — so the policy can't route on
+a model's name, position, cost, or tier. Basic tools keep their real names.
+It returns ``(anon_tools, anon_to_real)``.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+
+from openjarvis.agents.hybrid.expert_registry import (
+ KIND_MODEL,
+ anonymize_tools,
+ build_tool_specs,
+ orchestrator_catalog,
+ tools_by_name,
+)
+
+# Real brand tokens that must never leak into the anonymized, model-facing specs.
+_BRANDS = ("gpt", "claude", "gemini", "qwen")
+
+
+def _model_names(cat):
+ return [t.name for t in cat if t.kind == KIND_MODEL]
+
+
+def _basic_names(cat):
+ return [t.name for t in cat if t.kind != KIND_MODEL]
+
+
+def test_no_brand_names_in_anonymized_specs():
+ cat = orchestrator_catalog()
+ anon, _ = anonymize_tools(cat, random.Random(0))
+ specs = build_tool_specs(anon)
+ # The whole model-facing payload (names + descriptions + categories) must be
+ # brand-free. `.model` is preserved on the tool for dispatch but is NOT part
+ # of the spec the orchestrator conditions on, so it's fine that it still holds
+ # the real id.
+ blob = json.dumps(specs).lower()
+ for brand in _BRANDS:
+ assert brand not in blob, f"brand {brand!r} leaked into anonymized specs"
+ # And specifically the anonymized expert descriptions carry no brand.
+ for t in anon:
+ if t.name.startswith("model_"):
+ assert not any(b in t.description().lower() for b in _BRANDS)
+
+
+def test_hide_cost_removes_price_line_from_descriptions():
+ cat = orchestrator_catalog()
+ # Sanity: the raw model tools DO surface a price line before anonymizing.
+ raw_model = next(t for t in cat if t.kind == KIND_MODEL)
+ assert "Pricing:" in raw_model.description()
+
+ anon, anon_to_real = anonymize_tools(cat, random.Random(1))
+ for t in anon:
+ if t.name in anon_to_real: # an anonymized model expert
+ assert t.hide_cost is True
+ desc = t.description()
+ assert "Pricing:" not in desc
+ assert "$" not in desc
+ assert "/1M" not in desc
+ assert "latency" not in desc.lower()
+
+
+def test_each_model_maps_to_opaque_label_and_round_trips():
+ cat = orchestrator_catalog()
+ orig_models = _model_names(cat)
+ anon, anon_to_real = anonymize_tools(cat, random.Random(2))
+
+ # One opaque label per real model, all in the model_xxxx namespace.
+ assert len(anon_to_real) == len(orig_models)
+ assert all(lbl.startswith("model_") for lbl in anon_to_real)
+
+ # No collisions: labels unique, and each real model recovered exactly once.
+ assert len(set(anon_to_real)) == len(anon_to_real)
+ assert len(set(anon_to_real.values())) == len(anon_to_real)
+ assert set(anon_to_real.values()) == set(orig_models)
+
+ # Round-trip: every anonymized expert's label resolves back to a real name,
+ # and `.model` is preserved so dispatch still reaches the right backend.
+ by_orig = tools_by_name(cat)
+ for t in anon:
+ if t.name.startswith("model_"):
+ real = anon_to_real[t.name]
+ assert real in by_orig
+ assert t.model == by_orig[real].model # backend id untouched
+
+ # Basic tools keep their real names and are untouched by the label map.
+ basics = _basic_names(cat)
+ anon_names = {t.name for t in anon}
+ assert set(basics) <= anon_names
+ assert not (set(basics) & set(anon_to_real))
+
+
+def test_experts_block_on_top_then_basics():
+ cat = orchestrator_catalog()
+ anon, anon_to_real = anonymize_tools(cat, random.Random(3))
+ n_models = len(anon_to_real)
+ # Experts form a contiguous block at the front, basics underneath.
+ assert all(t.name.startswith("model_") for t in anon[:n_models])
+ assert all(not t.name.startswith("model_") for t in anon[n_models:])
+
+
+def test_labels_and_order_are_shuffled_not_identity():
+ cat = orchestrator_catalog()
+ orig_models = _model_names(cat)
+
+ orderings = set()
+ labels_for_first = set()
+ for seed in range(25):
+ anon, anon_to_real = anonymize_tools(cat, random.Random(seed))
+ # Real-model order as it appears in the anonymized expert block.
+ order = tuple(anon_to_real[t.name] for t in anon if t.name.startswith("model_"))
+ orderings.add(order)
+ # Label assigned to the first catalog model varies across rngs.
+ rev = {real: lbl for lbl, real in anon_to_real.items()}
+ labels_for_first.add(rev[orig_models[0]])
+
+ # Not a fixed identity: across seeds the expert order actually varies...
+ assert len(orderings) > 1
+ # ...and at least one ordering differs from the input model order.
+ assert any(order != tuple(orig_models) for order in orderings)
+ # Labels are random per-rng, not a stable function of position.
+ assert len(labels_for_first) > 1
diff --git a/tests/agents/test_expert_registry.py b/tests/agents/test_expert_registry.py
new file mode 100644
index 00000000..303816ab
--- /dev/null
+++ b/tests/agents/test_expert_registry.py
@@ -0,0 +1,221 @@
+"""Tests for the faithful ToolOrchestra unified-tool registry."""
+
+from __future__ import annotations
+
+import pytest
+
+from openjarvis.agents.hybrid.expert_registry import (
+ CATEGORY_BASIC,
+ CATEGORY_CLOUD_FRONTIER,
+ CATEGORY_LOCAL_OSS,
+ KIND_MODEL,
+ KIND_TOOL,
+ ExpertTool,
+ build_tool_specs,
+ openjarvis_tool,
+ orchestrator_catalog,
+ to_worker_dict,
+ tools_by_name,
+)
+
+
+def test_catalog_names_unique_and_valid():
+ cat = orchestrator_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_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 = orchestrator_catalog()
+ specs = build_tool_specs(cat)
+ by = {s["function"]["name"]: s for s in specs}
+ gpt = by["gpt_5_5"]
+ assert gpt["type"] == "function"
+ assert "input" in gpt["function"]["parameters"]["properties"]
+ # Price table is surfaced in the description (the policy is trained on it).
+ assert "/1M input" in gpt["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"]
+
+
+# Two cloud-frontier + four local-OSS model tools, in catalog order.
+_ORCH_MODEL_NAMES = [
+ "gpt_5_5",
+ "claude_opus_4_8",
+ "qwen3_5_9b",
+ "qwen3_6_27b_fp8",
+ "qwen3_5_122b_a10b_fp8",
+ "qwen3_5_397b_a17b_fp8",
+]
+
+# Bridged real OpenJarvis tools (basic) appended after web_search/code_interpreter.
+_ORCH_BASIC_NAMES = [
+ "web_search",
+ "code_interpreter",
+ "calculator",
+ "shell_exec",
+ "file_read",
+ "file_write",
+ "http_request",
+ "think",
+ "apply_patch",
+ "pdf_extract",
+ "db_query",
+]
+
+
+def test_orchestrator_catalog_two_model_classes_plus_basics():
+ cat = orchestrator_catalog()
+ names = [t.name for t in cat]
+ # 6 model tools (2 cloud_frontier + 4 local_oss) come first.
+ assert names[:6] == _ORCH_MODEL_NAMES
+ # then the basic tools.
+ assert set(_ORCH_BASIC_NAMES) <= set(names)
+ assert len(cat) == 6 + len(_ORCH_BASIC_NAMES)
+ by = tools_by_name(cat)
+ assert by["gpt_5_5"].category == CATEGORY_CLOUD_FRONTIER
+ assert by["claude_opus_4_8"].category == CATEGORY_CLOUD_FRONTIER
+ # Default routing for every model tool is OpenRouter (no servers required).
+ for n in (
+ "qwen3_5_9b",
+ "qwen3_6_27b_fp8",
+ "qwen3_5_122b_a10b_fp8",
+ "qwen3_5_397b_a17b_fp8",
+ ):
+ assert by[n].category == CATEGORY_LOCAL_OSS
+ assert by[n].backend_type == "openrouter"
+ assert by[n].base_url is None
+ # OpenRouter routing carries the slug + a (estimated) per-token price.
+ assert "/" in by[n].model and by[n].price_in > 0.0
+
+
+def test_orchestrator_catalog_categories_present():
+ cat = orchestrator_catalog()
+ cats = {t.category for t in cat}
+ assert cats == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS, CATEGORY_BASIC}
+
+
+def test_orchestrator_can_drop_tools():
+ cat = orchestrator_catalog(include_tools=False)
+ assert {t.category for t in cat} == {CATEGORY_CLOUD_FRONTIER, CATEGORY_LOCAL_OSS}
+ assert len(cat) == 6
+
+
+def test_orchestrator_specs_include_category_field():
+ specs = build_tool_specs(orchestrator_catalog())
+ by = {s["function"]["name"]: s for s in specs}
+ assert by["gpt_5_5"]["function"]["category"] == CATEGORY_CLOUD_FRONTIER
+ assert by["qwen3_5_9b"]["function"]["category"] == CATEGORY_LOCAL_OSS
+ assert by["web_search"]["function"]["category"] == CATEGORY_BASIC
+ assert by["shell_exec"]["function"]["category"] == CATEGORY_BASIC
+ # Every tool carries a category tag.
+ assert all("category" in s["function"] for s in specs)
+
+
+def test_orchestrator_local_models_get_base_url_when_provided():
+ # An endpoint switches that model from the OpenRouter default to local vLLM.
+ cat = orchestrator_catalog(
+ local_endpoints={
+ "Qwen/Qwen3.5-9B": "http://x/v1",
+ "Qwen/Qwen3.6-27B-FP8": "http://y/v1",
+ }
+ )
+ by = tools_by_name(cat)
+ assert by["qwen3_5_9b"].backend_type == "vllm"
+ assert by["qwen3_5_9b"].base_url == "http://x/v1"
+ assert by["qwen3_5_9b"].price_in == 0.0 and by["qwen3_5_9b"].price_out == 0.0
+ assert by["qwen3_6_27b_fp8"].base_url == "http://y/v1"
+ # Unmapped local model -> stays on the OpenRouter default (base_url None).
+ assert by["qwen3_5_122b_a10b_fp8"].backend_type == "openrouter"
+ assert by["qwen3_5_122b_a10b_fp8"].base_url is None
+ # Cloud frontier carries real pricing.
+ assert by["claude_opus_4_8"].price_in > 0.0
+ assert by["gpt_5_5"].price_in > 0.0
+
+
+def test_orchestrator_model_backends_override():
+ # Force a cloud model onto OpenRouter and a local model onto vLLM explicitly.
+ cat = orchestrator_catalog(
+ model_backends={
+ "claude-opus-4-8": "openrouter",
+ "Qwen/Qwen3.5-397B-A17B-FP8": "vllm",
+ },
+ local_endpoints={"Qwen/Qwen3.5-397B-A17B-FP8": "http://z/v1"},
+ )
+ by = tools_by_name(cat)
+ assert by["claude_opus_4_8"].backend_type == "openrouter"
+ assert by["claude_opus_4_8"].model == "anthropic/claude-opus-4.8"
+ # No override for gpt-5.5 -> it defaults to its NATIVE first-party API
+ # (openai), not OpenRouter. Frontier models hit their native provider by
+ # default; OpenRouter is only the fallback for OSS/local models or when
+ # explicitly requested via model_backends.
+ assert by["gpt_5_5"].backend_type == "openai" # native default
+ assert by["gpt_5_5"].model == "gpt-5.5"
+ assert by["qwen3_5_397b_a17b_fp8"].backend_type == "vllm"
+ assert by["qwen3_5_397b_a17b_fp8"].base_url == "http://z/v1"
+
+
+def test_orchestrator_openrouter_slug_override():
+ cat = orchestrator_catalog(
+ openrouter_slugs={"Qwen/Qwen3.5-9B": "qwen/qwen3.5-9b-custom"}
+ )
+ by = tools_by_name(cat)
+ assert by["qwen3_5_9b"].model == "qwen/qwen3.5-9b-custom"
+
+
+def test_openjarvis_tool_bridges_real_tool_with_custom_schema():
+ params = {
+ "type": "object",
+ "properties": {"command": {"type": "string"}},
+ "required": ["command"],
+ }
+ t = openjarvis_tool("shell_exec", summary="Run a shell command.", params=params)
+ assert t.kind == KIND_TOOL
+ assert t.backend_type == "openjarvis-tool"
+ assert t.model == "shell_exec"
+ assert t.category == CATEGORY_BASIC
+ spec = t.to_spec()
+ assert spec["function"]["name"] == "shell_exec"
+ assert spec["function"]["parameters"] == params
+ assert spec["function"]["category"] == CATEGORY_BASIC
+
+
+def test_build_tool_specs_includes_category_for_bridged_tools():
+ specs = build_tool_specs(
+ [
+ openjarvis_tool(
+ "calculator",
+ summary="Math.",
+ params={
+ "type": "object",
+ "properties": {"expression": {"type": "string"}},
+ "required": ["expression"],
+ },
+ ),
+ ]
+ )
+ assert specs[0]["function"]["category"] == CATEGORY_BASIC
+ assert "expression" in specs[0]["function"]["parameters"]["properties"]
+
+
+def test_to_worker_dict_maps_backend():
+ cat = orchestrator_catalog(local_endpoints={"Qwen/Qwen3.5-9B": "http://x/v1"})
+ by = tools_by_name(cat)
+ assert to_worker_dict(by["gpt_5_5"]) == {
+ "name": "gpt_5_5",
+ "type": "openai",
+ "model": "gpt-5.5",
+ }
+ local = to_worker_dict(by["qwen3_5_9b"])
+ assert local["type"] == "vllm" and local["base_url"] == "http://x/v1"