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

toolorchestra: split the monolith into a package + anonymized expert registry

Browse Source 
Break the 1.8k-line toolorchestra.py into a focused package (agent,
rollout, workers, experts, unified, parsing, prompts, clients, sandbox,
tracing) and add expert_registry.py, which serves the orchestrator an
anonymized expert catalog (opaque model_xxxx labels, hidden costs) so a
router can't route on brand prior. retry.py adds backoff to the rollout's
orchestrator + expert cloud calls. Prices corrected to OpenRouter list.
This commit is contained in:
Andrew Park
2026-07-15 13:45:07 -07:00
parent 2e68e227b7
commit bc7ac3e14f
19 changed files with 4023 additions and 1868 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+1
View File
@@ -28,6 +28,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"braintrust>=0.0.150",
"click>=8",
"datasets>=4.5.0",
"ddgs>=9.11.4",
src/openjarvis/agents/hybrid/_base.py
+4 -4
View File
@@ -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]:
src/openjarvis/agents/hybrid/_prices.py
+16
View File
@@ -11,6 +11,7 @@ from __future__ import annotations
# USD per million tokens, (input, output). Local models = 0.
PRICES: dict[str, tuple[float, float]] = {
"claude-opus-4-8": (5.00, 25.0),
"claude-opus-4-7": (5.00, 25.0),
"claude-sonnet-4-6": (3.00, 15.0),
"claude-haiku-4-5": (1.00, 5.00),
@@ -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",
)
src/openjarvis/agents/hybrid/expert_registry.py
+751
View File
@@ -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",
]
src/openjarvis/agents/hybrid/retry.py
+98
View File
@@ -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"]
src/openjarvis/agents/hybrid/toolorchestra.py
-1864
View File
File diff suppressed because it is too large Load Diff
src/openjarvis/agents/hybrid/toolorchestra/__init__.py
+20
View File
@@ -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"]
src/openjarvis/agents/hybrid/toolorchestra/agent.py
+848
View File
@@ -0,0 +1,848 @@
"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
Two modes, gated by ``method_cfg.orchestrator_mode``:
* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the
orchestrator, dispatching to a numbered worker pool via JSON
``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as
a prompted upper-bound reference point — NOT the paper's setup.
* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B``
served on a local vLLM is the orchestrator. It emits OpenAI-style
``tool_calls`` (or ``<tool_call>{...}</tool_call>`` text blocks when
vLLM's tool parser doesn't catch them) for three expert tools —
``enhance_reasoning``, ``answer``, ``search`` — exactly as in the
upstream ``evaluation/tools.json``. Each tool's ``model`` arg
(``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real
backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier
Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen
for `*-3`. Search routes to the Anthropic server-side web_search.
We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the
code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B,
Qwen-Math, Qwen-Coder); the expert pool collapses onto our existing
worker types. Energy-wise, "expert" answers are cloud calls.
Pipeline per task (RL mode):
1. Orchestrator-8B reads `Problem: ...\\n\\n{context}\\n\\nChoose an
appropriate tool.` with the three tools declared.
2. It emits one ``tool_call`` per turn — ``search`` updates the
context, ``enhance_reasoning`` appends code/exec output (we run the
tool as a plain LLM call, no sandbox — the model just gets prose
back), ``answer`` produces the final answer and the loop stops.
3. Up to ``max_turns`` (default 8) turns; on parse failure we fall
back to the strongest expert worker.
Prompted-mode pipeline:
1. Orchestrator (cloud) reads question + numbered worker pool.
2. Each turn it emits ``{"action": "call_worker", "worker_id": int,
"input": str}`` or ``{"action": "final_answer", "answer": str}``.
3. Up to ``max_turns`` (default 6) calls before forcing a final-answer
prompt; fallback to strongest worker on parse failure.
Workers come from ``cfg["workers"]`` or a sensible default pool (local
Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7,
gpt-5-mini).
"""
from __future__ import annotations
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from openjarvis.agents._stubs import AgentContext
from openjarvis.agents.hybrid._base import LocalCloudAgent
from openjarvis.agents.hybrid._prices import PRICES
from openjarvis.agents.hybrid.mini_swe_agent import (
_clone_repo,
_extract_diff,
)
from openjarvis.agents.hybrid.toolorchestra.clients import (
_call_orchestrator_with_tool_calls,
)
from openjarvis.agents.hybrid.toolorchestra.experts import (
_PAPER_CODER_OPENROUTER,
_expert_for,
_paper_expert_for,
)
from openjarvis.agents.hybrid.toolorchestra.parsing import (
_build_user_prompt,
_extract_final_answer_text,
_parse_action,
_parse_rl_tool_call,
)
from openjarvis.agents.hybrid.toolorchestra.prompts import (
FORCE_FINAL_PROMPT,
ORCHESTRATOR_SYS,
RL_ALL_TOOLS,
RL_ORCHESTRATOR_SYS,
RL_TOOLS_SPEC,
)
from openjarvis.agents.hybrid.toolorchestra.sandbox import (
_call_modal_python,
_extract_first_python_block,
)
from openjarvis.agents.hybrid.toolorchestra.workers import (
_TOOLORCH_SEARCH_TYPES,
_call_worker,
_resolve_worker_pool,
_swe_call_worker,
)
from openjarvis.core.registry import AgentRegistry
@AgentRegistry.register("toolorchestra")
class ToolOrchestraAgent(LocalCloudAgent):
"""Multi-turn dispatcher over a mixed worker pool.
Two modes (see module docstring): ``method_cfg.orchestrator_mode``
is ``"prompted"`` (default, cloud-as-orchestrator) or ``"rl"``
(paper-faithful, drives ``nvidia/Orchestrator-8B`` on a local vLLM).
"""
agent_id = "toolorchestra"
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
# Validate `method_cfg.worker_pool` early — surfaces config errors
# at agent construction rather than on the first task. No-op when
# the override is absent.
if self._cfg.get("worker_pool") is not None:
_resolve_worker_pool(
self._cfg,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# Validate `orchestrator_mode` (typo-checked here, not on first task).
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
if mode not in ("prompted", "rl"):
raise ValueError(
f"toolorchestra: orchestrator_mode must be 'prompted' or 'rl'; "
f"got {mode!r}"
)
def _run_paradigm(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
mode = str(self._cfg.get("orchestrator_mode", "prompted")).lower()
if mode == "rl":
return self._run_rl(input, context, **kwargs)
return self._run_prompted(input, context, **kwargs)
# ------------------------------------------------------------------
# Legacy prompted-orchestrator path.
# ------------------------------------------------------------------
def _run_prompted(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
cfg = self._cfg
question = input
# Resolution order (strict replace, no merge):
# 1. `cfg["workers"]` — legacy direct override, used by tests.
# 2. `cfg["worker_pool"]` — cell-config override; validated +
# $local/$cloud substituted.
# 3. `_default_pool(...)` — heterogeneous default.
if cfg.get("workers"):
workers = cfg["workers"]
else:
workers = _resolve_worker_pool(
cfg,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
if not workers:
raise RuntimeError("toolorchestra: empty worker pool")
max_turns = int(cfg.get("max_turns", 6))
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024))
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
swe_mode = (
bool(cfg.get("swe_use_agent_loop"))
and bool(task_meta.get("problem_statement"))
and bool(task_meta.get("repo"))
and bool(task_meta.get("base_commit"))
)
shared_workdir: Optional[Path] = None
if swe_mode:
shared_workdir = Path(
tempfile.mkdtemp(
prefix=f"toolorch-swe-{task_meta.get('task_id', 'x')}-"
)
)
try:
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
except Exception:
shutil.rmtree(shared_workdir, ignore_errors=True)
raise
self.record_trace_event(
{
"kind": "toolorchestra_swe_workdir",
"workdir": str(shared_workdir),
"repo": task_meta["repo"],
"base_commit": task_meta["base_commit"],
}
)
# try/finally guards ``shared_workdir`` against exceptions raised
# anywhere in the turn loop, the worker calls, the fallback, or
# the diff-extraction step. Without this, at n=500 SWE-bench an
# exception leaves hundreds of MB of cloned repos in tempdir.
try:
history: List[Dict[str, Any]] = []
tokens_local = 0
tokens_cloud = 0
cost = 0.0
n_web_searches_total = 0
# tool_calls: bash turns from SWE subloops + web_search uses
# from GAIA. Orchestrator dispatch turns are NOT counted (they
# produce text only — calling a worker is one tool call's worth
# of "delegation" but the actual tool action happens inside).
tool_calls = 0
final_answer: Optional[str] = None
forced_final = False
parse_failures = 0
for turn in range(1, max_turns + 1):
sys_prompt = ORCHESTRATOR_SYS
if turn == max_turns and final_answer is None:
sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT
forced_final = True
user = _build_user_prompt(question, workers, history)
text, o_in, o_out = self._call_cloud(
user=user,
system=sys_prompt,
max_tokens=orch_max_tokens,
temperature=0.0,
)
tokens_cloud += o_in + o_out
cost += self.cost_usd(self._cloud_model, o_in, o_out)
action = _parse_action(text)
history.append(
{
"role": "orchestrator",
"turn": turn,
"raw": text,
"action": action,
}
)
self.record_trace_event(
{
"kind": "toolorchestra_action",
"turn": turn,
"action": action,
"raw": text,
}
)
if action is None:
parse_failures += 1
if parse_failures >= 2 or forced_final:
final_answer = _extract_final_answer_text(text)
break
continue
kind = action.get("action")
if kind == "final_answer":
final_answer = str(action.get("answer", "")).strip()
break
if kind == "call_worker":
wid = action.get("worker_id")
w_input = action.get("input", "")
if not isinstance(wid, int) or not (0 <= wid < len(workers)):
parse_failures += 1
if parse_failures >= 2 or forced_final:
final_answer = _extract_final_answer_text(text)
break
continue
worker = workers[wid]
if swe_mode and shared_workdir is not None:
(
w_text,
w_in,
w_out,
is_local,
extra_cost,
n_searches,
bash_turns,
) = _swe_call_worker(
worker,
str(w_input),
cfg,
task_meta,
shared_workdir,
turn,
)
tool_calls += bash_turns
else:
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
_call_worker(worker, str(w_input), cfg)
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
n_web_searches_total += n_searches
tool_calls += n_searches
history.append(
{
"role": "worker",
"turn": turn,
"worker_id": wid,
"worker_name": worker["name"],
"worker_model": worker["model"],
"output": w_text,
"tokens_in": w_in,
"tokens_out": w_out,
"n_web_searches": n_searches,
}
)
continue
# Unknown action kind — treat as parse failure.
parse_failures += 1
if final_answer is None:
# Hard fallback: call the strongest non-search worker directly.
# "Strongest" = highest output-token price in `_prices.PRICES`,
# which tracks model capability tier closely enough for this.
# Search workers are excluded — they answer fact-lookup
# questions, not synthesis.
non_search = [
w for w in workers if w.get("type") not in _TOOLORCH_SEARCH_TYPES
] or workers
worker = max(
non_search,
key=lambda w: PRICES.get(w.get("model", ""), (0.0, 0.0))[1],
)
if swe_mode and shared_workdir is not None:
(ans, w_in, w_out, is_local, extra_cost, _, bash_turns) = (
_swe_call_worker(
worker,
question,
cfg,
task_meta,
shared_workdir,
max_turns + 1,
)
)
tool_calls += bash_turns
else:
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
worker, question, cfg
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
history.append(
{
"role": "worker",
"turn": max_turns + 1,
"worker_id": worker["id"],
"worker_name": worker["name"],
"worker_model": worker["model"],
"output": ans,
"tokens_in": w_in,
"tokens_out": w_out,
"fallback": True,
}
)
final_answer = ans
# In SWE mode, the authoritative output is the working-tree diff —
# frame it (the runner extracts it via the scorer's ```diff fence).
if swe_mode and shared_workdir is not None:
patch = _extract_diff(shared_workdir)
if patch.strip():
final_answer = (
f"{final_answer}\n\n```diff\n{patch}```"
if final_answer
else f"```diff\n{patch}```"
)
meta = {
"tokens_local": tokens_local,
"tokens_cloud": tokens_cloud,
"cost_usd": cost,
"turns": len([h for h in history if h["role"] == "orchestrator"]),
"web_search_uses": n_web_searches_total,
"tool_calls": int(tool_calls),
"traces": {
"history": history,
"forced_final": forced_final,
"parse_failures": parse_failures,
"workers": workers,
"n_web_searches": n_web_searches_total,
"note": (
"inference-only port; the RL-trained Nemotron-Orchestrator-8B "
"is NOT in the loop. Results are preliminary."
),
},
}
return final_answer, meta
finally:
if shared_workdir is not None:
shutil.rmtree(shared_workdir, ignore_errors=True)
# ------------------------------------------------------------------
# Paper-faithful Orchestrator-8B path.
# ------------------------------------------------------------------
def _run_rl(
self,
input: str,
context: Optional[AgentContext],
**kwargs: Any,
) -> Tuple[str, Dict[str, Any]]:
cfg = self._cfg
question = input
# Orchestrator endpoint / model (where the RL'd 8B lives).
orch_endpoint = str(
cfg.get("orchestrator_endpoint", "http://localhost:8003/v1")
)
orch_model = str(cfg.get("orchestrator_model", "orchestrator-8b"))
max_turns = int(cfg.get("max_turns", 8))
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
orch_temp = float(cfg.get("orchestrator_temperature", 1.0))
# Paper-match pool toggle (2026-05-19). When set, `_paper_expert_for`
# replaces `_expert_for` and `enhance_reasoning` is post-processed
# through a Modal Python sandbox. See module docstring + paper-match
# doc at `docs/26.5.19/toolorchestra-papermatch.md`.
paper_mode = str(cfg.get("pool", "")).lower() == "paper"
# SWE-bench detection: same gate as the prompted path. Requires
# `method_cfg.swe_use_agent_loop = true` AND the task carries the
# SWE-bench fields. When active, the `enhance_reasoning` and
# `answer` workers route through `run_swe_agent_loop` on a shared
# workdir; search workers stay one-shot. At end, the working-tree
# diff is appended to final_answer so `_score_swebench` can extract
# it via the ```diff fence.
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
swe_mode = (
bool(cfg.get("swe_use_agent_loop"))
and bool(task_meta.get("problem_statement"))
and bool(task_meta.get("repo"))
and bool(task_meta.get("base_commit"))
)
shared_workdir: Optional[Path] = None
if swe_mode:
shared_workdir = Path(
tempfile.mkdtemp(
prefix=f"toolorch-rl-swe-{task_meta.get('task_id', 'x')}-"
)
)
try:
_clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir)
except Exception:
shutil.rmtree(shared_workdir, ignore_errors=True)
raise
self.record_trace_event(
{
"kind": "toolorchestra_rl_swe_workdir",
"workdir": str(shared_workdir),
"repo": task_meta["repo"],
"base_commit": task_meta["base_commit"],
}
)
# ``context_str`` mirrors the upstream's running context — accumulates
# search documents and code/exec snippets across turns. We keep this
# as a single string for prompt simplicity; the upstream uses
# tokenized cutoffs (we cap at ~24k chars instead).
context_str = ""
doc_list: List[str] = []
history: List[Dict[str, Any]] = []
tokens_local = 0
tokens_cloud = 0
cost = 0.0
n_web_searches_total = 0
tool_calls = 0
final_answer: Optional[str] = None
parse_failures = 0
# Single outer try/finally guards `shared_workdir` against any
# exception in the orchestrator loop, the post-loop fallback, or
# the diff-extraction step. Matches the prompted path's pattern.
try:
for turn in range(1, max_turns + 1):
user = (
f"Problem: {question}\n\n{context_str}\n\n"
"Choose an appropriate tool."
)
# Orchestrator-8B served on local vLLM. We pass the three NVlabs
# tools verbatim. In paper-mode we use the local helper so we
# get the SDK-level ``tool_calls`` object back — `_call_vllm`
# returns just text and loses the call when vLLM's parser
# caught it. Orchestrator-8B emits its routing decision in the
# OpenAI-native ``tool_calls`` array with an empty text body,
# so the legacy `_call_vllm` path saw nothing and silently fell
# through to the answer-1 fallback (parse_failures: 2 on every
# non-opus-gaia cell — see docs/reports/toolorchestra.md). Both
# modes now use `_call_orchestrator_with_tool_calls` so the
# parser can read structured tool calls; the text-tag path in
# `_parse_rl_tool_call` is still the fallback when `tool_calls`
# is empty.
text, o_in, o_out, sdk_tool_calls = _call_orchestrator_with_tool_calls(
orch_model,
orch_endpoint,
user=user,
system=RL_ORCHESTRATOR_SYS,
max_tokens=orch_max_tokens,
temperature=orch_temp,
tools=RL_TOOLS_SPEC,
)
self.record_trace_event(
{
"kind": "vllm",
"role": "orchestrator",
"model": orch_model,
"endpoint": orch_endpoint,
"system": RL_ORCHESTRATOR_SYS,
"user": user,
"response": text,
"tool_calls": [
{
"id": getattr(tc, "id", None),
"type": getattr(tc, "type", None),
"function": {
"name": getattr(
getattr(tc, "function", None), "name", None
),
"arguments": getattr(
getattr(tc, "function", None), "arguments", None
),
},
}
for tc in (sdk_tool_calls or [])
],
"tokens_in": o_in,
"tokens_out": o_out,
}
)
tokens_local += o_in + o_out
action = _parse_rl_tool_call(text, sdk_tool_calls)
history.append(
{
"role": "orchestrator",
"turn": turn,
"raw": text,
"action": action,
}
)
self.record_trace_event(
{
"kind": "toolorchestra_rl_action",
"turn": turn,
"action": action,
"raw": text,
}
)
if action is None:
parse_failures += 1
if parse_failures >= 2:
break
continue
name = action["name"]
args = action.get("arguments", {})
slot = args.get("model", "")
# Validate against the upstream tool/arg schema.
valid = (
name in RL_ALL_TOOLS
and isinstance(slot, str)
and (slot in RL_ALL_TOOLS[name]["model"])
)
if not valid:
parse_failures += 1
if parse_failures >= 2:
break
# Replay with a softer nudge in the context.
context_str += (
f"\n[Orchestrator emitted invalid tool call "
f"name={name!r} slot={slot!r} — try again.]\n"
)
continue
# Paper-match (`method_cfg.pool == "paper"`) routes through
# the Tavily/OpenRouter/Modal pool instead of the default
# Anthropic-web-search-driven mapping. For `search` this
# also forces the worker prompt to a raw query string
# (Tavily takes a single search string, not a chat-style
# framing).
if paper_mode:
worker = _paper_expert_for(
slot,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# In paper mode, `enhance_reasoning` is always the coder
# specialist regardless of the orchestrator's chosen tier.
# The coder is then expected to emit a python block which
# we exec in Modal (below).
if name == "enhance_reasoning":
worker = {
"name": f"coder:{slot}",
"type": "openrouter",
"model": _PAPER_CODER_OPENROUTER,
}
else:
worker = _expert_for(
slot,
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
# Dispatch — the orchestrator only conveys a tool/model
# choice, NOT a question rewrite; the prompt we send the
# expert is the same context the orchestrator saw, framed
# appropriately for the tool.
if name == "search":
if paper_mode:
# Tavily takes a query string. Orchestrator-8B often
# emits an extra `query` arg (not in the upstream
# schema but useful) — prefer it; else fall back to
# the raw question.
q = args.get("query")
w_input = q if isinstance(q, str) and q.strip() else question
else:
w_input = (
f"Search the web to gather information that helps answer:\n"
f"{question}\n\nCurrent context:\n{context_str or '(empty)'}"
)
elif name == "enhance_reasoning":
if paper_mode:
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Write a short Python script that computes intermediate "
"results which help answer the problem. Output ONLY the "
"code inside one ```python ... ``` fenced block. Print "
"any results you derive using `print(...)`. The script "
"must run with the Python stdlib only — no extra pip "
"installs."
)
else:
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Reason carefully. Outline the key intermediate steps and any "
"computations or facts you can derive. Do NOT give a final "
"answer — the orchestrator will collect your reasoning and "
"call the answer tool next."
)
else: # name == "answer"
w_input = (
f"Problem: {question}\n\nContext:\n{context_str or '(empty)'}\n\n"
"Provide the final answer to the user. Respect any "
"answer-format rules in the question (e.g. GAIA's "
"FINAL ANSWER: <value> convention)."
)
# SWE mode: route enhance_reasoning / answer workers through
# the SWE agent loop on the shared workdir so they can read
# files, run tests, and edit the working tree. Search workers
# stay one-shot (no agent loop). The `_swe_call_worker`
# one-shot fallbacks (openai-typed workers, search) return
# bash_turns=0; vllm/anthropic-typed workers run the loop.
bash_turns = 0
if swe_mode and shared_workdir is not None and name != "search":
(
w_text,
w_in,
w_out,
is_local,
extra_cost,
n_searches,
bash_turns,
) = _swe_call_worker(
worker,
w_input,
cfg,
task_meta,
shared_workdir,
turn,
)
else:
w_text, w_in, w_out, is_local, extra_cost, n_searches = (
_call_worker(worker, w_input, cfg)
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
n_web_searches_total += n_searches
# SWE bash turns count as tool calls (each one is a $BASH block
# the agent executed). On non-SWE turns fall back to the
# original "at least one expert call" accounting.
tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches)
# Paper-match: pipe coder output through a Modal sandbox so
# `enhance_reasoning` actually executes the code the coder
# wrote. Append the exec output to the worker's text. No-op
# when no python block is found.
modal_exec_output: Optional[str] = None
modal_exec_rc: Optional[int] = None
if paper_mode and name == "enhance_reasoning" and not swe_mode:
code = _extract_first_python_block(w_text)
if code:
timeout_s = int(cfg.get("modal_python_timeout_s", 60))
modal_exec_output, modal_exec_rc = _call_modal_python(
code,
timeout_s=timeout_s,
)
tool_calls += 1
w_text = (
f"{w_text}\n\n[modal-python stdout/stderr "
f"(rc={modal_exec_rc})]\n{modal_exec_output}"
)
history.append(
{
"role": "worker",
"turn": turn,
"tool": name,
"slot": slot,
"worker_model": worker["model"],
"worker_type": worker["type"],
"output": w_text,
"tokens_in": w_in,
"tokens_out": w_out,
"n_web_searches": n_searches,
"bash_turns": bash_turns,
"modal_exec_rc": modal_exec_rc,
}
)
# Update accumulated context for the next turn.
if name == "search":
# Treat the search worker's response as a document.
doc_list.append(w_text)
ctx_docs = "\n\n".join(
f"Doc {i + 1}: {d}" for i, d in enumerate(doc_list)
)
# Crude char-level cap mirrors the upstream's ~24k token cap.
context_str = ("Documents:\n" + ctx_docs)[-24000:]
elif name == "enhance_reasoning":
snippet = f"\n\nReasoning/exec output:\n{w_text}"
context_str = (context_str + snippet)[-24000:]
else: # answer
final_answer = w_text.strip()
break
if final_answer is None:
# Hard fallback: ask the frontier worker directly. In SWE
# mode route this final call through the agent loop too so
# it can still touch the workdir and emit a diff.
expert_fn = _paper_expert_for if paper_mode else _expert_for
worker = expert_fn(
"answer-1",
self._local_model,
self._local_endpoint,
self._cloud_model,
self._cloud_endpoint,
)
fb_bash_turns = 0
if swe_mode and shared_workdir is not None:
(ans, w_in, w_out, is_local, extra_cost, _, fb_bash_turns) = (
_swe_call_worker(
worker,
question,
cfg,
task_meta,
shared_workdir,
max_turns + 1,
)
)
tool_calls += fb_bash_turns
else:
ans, w_in, w_out, is_local, extra_cost, _ = _call_worker(
worker, question, cfg
)
if is_local:
tokens_local += w_in + w_out
else:
tokens_cloud += w_in + w_out
cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost
history.append(
{
"role": "worker",
"turn": max_turns + 1,
"tool": "answer",
"slot": "answer-1",
"worker_model": worker["model"],
"worker_type": worker["type"],
"output": ans,
"tokens_in": w_in,
"tokens_out": w_out,
"bash_turns": fb_bash_turns,
"fallback": True,
}
)
final_answer = ans
# In SWE mode, the authoritative output is the working-tree diff —
# frame it so `_score_swebench`'s extract_patch picks it up.
if swe_mode and shared_workdir is not None:
patch = _extract_diff(shared_workdir)
if patch.strip():
final_answer = (
f"{final_answer}\n\n```diff\n{patch}```"
if final_answer
else f"```diff\n{patch}```"
)
meta = {
"tokens_local": tokens_local,
"tokens_cloud": tokens_cloud,
"cost_usd": cost,
"turns": len([h for h in history if h["role"] == "orchestrator"]),
"web_search_uses": n_web_searches_total,
"tool_calls": int(tool_calls),
"traces": {
"history": history,
"parse_failures": parse_failures,
"orchestrator_model": orch_model,
"orchestrator_endpoint": orch_endpoint,
"mode": "rl",
"pool": "paper" if paper_mode else "default",
"swe_mode": swe_mode,
"note": (
"RL-trained nvidia/Orchestrator-8B as orchestrator. "
"Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto "
"our hybrid worker types — see toolorchestra.py docstring."
),
},
}
return final_answer, meta
finally:
if shared_workdir is not None:
shutil.rmtree(shared_workdir, ignore_errors=True)
__all__ = ["ToolOrchestraAgent"]
src/openjarvis/agents/hybrid/toolorchestra/clients.py
+51
View File
@@ -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
src/openjarvis/agents/hybrid/toolorchestra/experts.py
+187
View File
@@ -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",
}
src/openjarvis/agents/hybrid/toolorchestra/parsing.py
+161
View File
@@ -0,0 +1,161 @@
"""Action / tool-call parsing + prompt assembly for ToolOrchestraAgent."""
from __future__ import annotations
import json
import re
from typing import Any, Dict, List, Optional
# Regex for ``<tool_call>{...}</tool_call>`` blocks emitted by Orchestrator-8B
# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a
# hermes-style template). Captures the JSON payload.
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
# Fallback for a known failure mode: the SFT'd orchestrator often emits its
# delegation as ``\boxed{expert_ab12: <sub-question>}`` (math-data habit bleeding
# into routing) instead of a real ``<tool_call>`` JSON block. The name before the
# colon is a valid (anonymized) tool label; the rest is the sub-question/args.
#
# Real outputs use several shapes, all handled here:
# \boxed{web_search: <query>} -> name, args
# \boxed{web_search query: <query>} -> the word "query" is dropped
# \boxed{file_read, path: <path>} -> the ", <key>" hint sets the arg key
# We require a ``name <optional , key | query> : <args>`` structure, so a genuine
# answer like ``\boxed{10}``, ``\boxed{b,e}`` or ``\boxed{The answer is: 42}``
# (first token isn't followed by a bare ``:`` / ``, key:`` / ``query:``) is left
# alone. ``group("key")`` is the explicit arg key when the model wrote ``, key:``.
_BOXED_DELEGATION_RE = re.compile(
r"\\boxed\{\s*([A-Za-z_][\w-]*)\s*"
r"(?:,\s*(?P<key>\w+)\s*)?" # optional ", key" hint (e.g. file_read, path:)
r"(?:query\s*)?" # optional literal "query" word before the colon
r":\s*(.+?)\s*\}\s*$",
re.DOTALL,
)
def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]:
"""Return ``{"name": str, "arguments": dict}`` or None.
Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls
back to scraping ``<tool_call>{...}</tool_call>`` tags from the raw
content. We take the first tool call only — Orchestrator-8B was trained
to emit exactly one per turn.
"""
# SDK-level path.
if sdk_tool_calls:
first = sdk_tool_calls[0]
name = getattr(getattr(first, "function", None), "name", None)
args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}"
try:
args = json.loads(args_raw)
except json.JSONDecodeError:
args = {}
if isinstance(name, str) and isinstance(args, dict):
return {"name": name, "arguments": args}
# Text-tag fallback.
if not isinstance(content, str):
return None
m = _TOOL_CALL_TAG_RE.search(content)
if m:
try:
obj = json.loads(m.group(1))
except json.JSONDecodeError:
obj = None
if isinstance(obj, dict):
name = obj.get("name")
args = obj.get("arguments", {})
if isinstance(name, str) and isinstance(args, dict):
return {"name": name, "arguments": args}
# \boxed{name[, key][ query]: <args>} fallback (see _BOXED_DELEGATION_RE).
bm = _BOXED_DELEGATION_RE.search(content.strip())
if bm:
arg_val = bm.groups()[-1].strip()
key = bm.group("key") or "input"
return {"name": bm.group(1), "arguments": {key: arg_val}}
return None
def _build_pool_block(workers: List[Dict[str, Any]]) -> str:
return "\n".join(
f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers
)
def _build_user_prompt(
question: str,
workers: List[Dict[str, Any]],
history: List[Dict[str, Any]],
) -> str:
pieces = [
f"Worker pool:\n{_build_pool_block(workers)}",
f"User question:\n{question}",
]
if history:
pieces.append("Conversation so far (orchestrator turns and worker outputs):")
for h in history:
if h["role"] == "orchestrator":
pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}")
else:
pieces.append(
f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n"
f"{h['output']}"
)
pieces.append(
"Emit the next JSON action object now — exactly one object, no prose."
)
return "\n\n".join(pieces)
def _strip_fences(s: str) -> str:
s = s.strip()
if s.startswith("```"):
first_nl = s.find("\n")
if first_nl != -1:
s = s[first_nl + 1 :]
if s.endswith("```"):
s = s[:-3]
s = s.strip()
return s
def _parse_action(text: str) -> Optional[Dict[str, Any]]:
s = _strip_fences(text)
# First try direct parse, then balanced-brace extraction.
try:
obj = json.loads(s)
if isinstance(obj, dict) and "action" in obj:
return obj
except json.JSONDecodeError:
pass
start = s.find("{")
if start == -1:
return None
depth = 0
for i in range(start, len(s)):
c = s[i]
if c == "{":
depth += 1
elif c == "}":
depth -= 1
if depth == 0:
try:
obj = json.loads(s[start : i + 1])
if isinstance(obj, dict) and "action" in obj:
return obj
except json.JSONDecodeError:
return None
return None
def _extract_final_answer_text(text: str) -> str:
"""Best-effort: pull the answer string from a malformed action emission.
Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line.
"""
m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL)
if m:
return m.group(1).encode("utf-8").decode("unicode_escape")
m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE)
if m:
return m.group(1).strip()
return text.strip()
src/openjarvis/agents/hybrid/toolorchestra/prompts.py
+110
View File
@@ -0,0 +1,110 @@
"""Prompt strings + tool specs for ToolOrchestraAgent (split from toolorchestra.py)."""
from __future__ import annotations
from typing import Any, Dict, List
ORCHESTRATOR_SYS = """\
You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms:
{"action": "call_worker", "worker_id": <int>, "input": "<question or instruction for that worker>"}
{"action": "final_answer", "answer": "<final answer to the user, respecting the question's answer-format rules>"}
Strategy:
- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math).
- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass.
- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase.
- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question.
"""
FORCE_FINAL_PROMPT = (
"Worker-call budget exhausted. Emit `final_answer` now using everything "
"you've learned. Respect the question's answer-format rules."
)
# ============================================================================
# RL-mode constants (Orchestrator-8B, paper-faithful).
# ============================================================================
#
# Verbatim copies of the upstream system prompt / user-prompt template / tools
# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
# edit the description text — Orchestrator-8B was RL-trained against this
# exact wording and pricing/latency table.
RL_ORCHESTRATOR_SYS = "You are good at using tools."
RL_TOOLS_SPEC: List[Dict[str, Any]] = [
{
"type": "function",
"function": {
"name": "enhance_reasoning",
"description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
"parameters": {
"properties": {
"model": {
"description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
{
"type": "function",
"function": {
"name": "answer",
"description": "give the final answer. Not allowed to call if documents is empty.",
"parameters": {
"properties": {
"model": {
"description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
{
"type": "function",
"function": {
"name": "search",
"description": "Search for missing information",
"parameters": {
"properties": {
"model": {
"description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
"type": "string",
}
},
"required": ["model"],
"title": "parameters",
"type": "object",
},
},
},
]
# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
"enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
"answer": {
"model": [
"answer-1",
"answer-2",
"answer-3",
"answer-4",
"answer-math-1",
"answer-math-2",
],
},
"search": {"model": ["search-1", "search-2", "search-3"]},
}
src/openjarvis/agents/hybrid/toolorchestra/rollout.py
+417
View File
@@ -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 <tool_call>{...}</tool_call> 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"<tool_call>.*?</tool_call>", 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 <tool_call>{...}</tool_call> block — NEVER write an "
"action as \\boxed{...} or as prose. "
"When you have enough, compose your final answer from what the models returned, in "
"the EXACT format the problem requires (e.g. a single number, a short exact "
"value, or one runnable code block). Do NOT restate the problem or add "
"meta-text such as 'the user is asking'. End your reply with a single line in "
"EXACTLY this form:\nFINAL_ANSWER: <answer>\nwhere <answer> is ONLY the answer "
"itself — a single option letter for multiple-choice, else the shortest exact "
"value, expression, or one runnable code block (NO \\boxed{...} wrapper) — with "
"NO explanation or restatement after it."
)
def build_system_prompt(specs: List[Dict[str, object]]) -> str:
"""Faithful ToolOrchestra system prompt (arXiv:2511.21689, verbatim from
their ``prepare_sft_data.py``): the ``RL_ORCHESTRATOR_SYS`` preamble + the
Qwen-style ``<tools>``/``<tool_call>`` block. Deliberately contains NO
routing/delegation instructions — cost-aware routing is learned from the RL
reward, not prompted.
"""
tools_block = "\n".join(json.dumps(s) for s in specs)
return (
f"{RL_ORCHESTRATOR_SYS}\n\n# Tools\n\n"
"You may call one or more functions to assist with the user query.\n\n"
"You are provided with function signatures within <tools></tools> "
"XML tags:\n"
f"<tools>\n{tools_block}\n</tools>\n\n"
"For each function call, return a json object with function name and "
"arguments within <tool_call></tool_call> XML tags:\n"
'<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}'
"\n</tool_call>"
)
# Char-level cap on the accumulated conversation (mirrors the paper's ~24k-token
# cap) and a per-observation cap so one giant tool dump can't blow it out.
_CONTEXT_CAP = 24000
_OBS_CAP = 8000
# After this many tool calls, push the orchestrator to stop and answer — bounds
# the "keep re-asking the same model forever" loop (esp. gemma).
_SOFT_CALL_CAP = 6
def _trim_history(messages: List[Dict[str, str]]) -> None:
"""Keep the message history under ``_CONTEXT_CAP`` chars by dropping the
OLDEST assistant/tool exchange, never the system prompt or the problem."""
def total() -> int:
return sum(len(m.get("content") or "") for m in messages)
# messages[0]=system, messages[1]=user problem — always keep those two.
while total() > _CONTEXT_CAP and len(messages) > 4:
del messages[2:4]
@dataclass
class UnifiedTurn:
"""One orchestrator turn. ``tool_name is None`` marks the final-answer turn."""
reasoning: str
tool_name: Optional[str] = None
arguments: Dict[str, object] = field(default_factory=dict)
observation: Optional[str] = None
@dataclass
class UnifiedRollout:
turns: List[UnifiedTurn]
final_answer: str
cost_usd: float = 0.0
tokens: int = 0
num_tool_calls: int = 0
parse_failures: int = 0
# When experts were anonymized for this rollout, maps the opaque label the
# policy saw (e.g. ``expert_a3f9``) back to the real tool name (``gpt_5_5``).
anon_map: Optional[Dict[str, str]] = None
# The EXACT tool specs the policy saw this rollout (anonymized when
# ``anonymize=True``). Serialization MUST build the saved system prompt from
# these — not from the real registry — or the prompt's ``<tools>`` block ends
# up with real model names+pricing while the assistant ``<tool_call>`` tags
# use the anon labels, which both breaks SFT (calls a tool absent from its
# list) and re-injects the name/cost bias anonymization removed.
tool_specs: Optional[List[Dict[str, object]]] = None
def tool_calls(self) -> List[Tuple[str, Dict[str, object]]]:
return [(t.tool_name, t.arguments) for t in self.turns if t.tool_name]
def _tool_prompt(tool: ExpertTool, arguments: Dict[str, object], question: str) -> str:
"""The text we actually send the dispatched tool, framed by its arg schema."""
for key in ("input", "query", "code"):
val = arguments.get(key)
if isinstance(val, str) and val.strip():
return val
return question
def _run_unified_rollout_inner(
question: str,
tools: List[ExpertTool],
*,
call_orchestrator: Callable[
..., Tuple[str, List[Tuple[str, Dict[str, object]]], int, int]
],
dispatch: Callable[[ExpertTool, Dict[str, object]], Tuple[str, float, int, bool]],
max_turns: int = 50,
system: str = RL_ORCHESTRATOR_SYS,
anonymize: bool = False,
) -> UnifiedRollout:
"""Drive the faithful unified-tool rollout for one task.
``anonymize``: replace each model expert with an opaque random label, a
uniform description and no cost line, shuffled — so the policy can't route on
a model's name/position/cost (which we found dominate the choice). The
anon->real mapping is returned on the rollout for offline analysis.
"""
anon_map: Optional[Dict[str, str]] = None
if anonymize:
tools, anon_map = anonymize_tools(tools, random.Random())
specs = build_tool_specs(tools)
by_name = tools_by_name(tools)
# Proper multi-turn conversation (NOT a flattened user blob): the orchestrator
# sees its own calls as `assistant` turns and each observation as a `tool`
# turn, so it can tell a tool result from user input and doesn't re-derive the
# whole problem every turn. This mirrors the serialized SFT format exactly.
messages: List[Dict[str, str]] = [
{"role": "system", "content": system},
{"role": "user", "content": f"Problem: {question}"},
]
turns: List[UnifiedTurn] = []
cost = 0.0
tokens = 0
n_tool_calls = 0
parse_failures = 0
nudges = 0
empty_input_nudges = 0
final_answer = ""
for _ in range(max_turns):
text, tool_calls, p_tok, c_tok = call_orchestrator(messages, specs)
tokens += int(p_tok) + int(c_tok)
if not tool_calls:
# Enforce the "MUST delegate to a model" rule: if the orchestrator
# tries to answer before ANY tool call, nudge it to route instead of
# accepting a solve-it-yourself answer (which reject-sampling drops).
if n_tool_calls == 0 and nudges < 2:
nudges += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": (
"Make progress by delegating a concrete sub-question to one of the "
"models now via a <tool_call>."
),
}
)
continue
# No tool call -> the orchestrator is answering. Terminate.
final_answer = (text or "").strip()
turns.append(UnifiedTurn(reasoning=text or "", tool_name=None))
messages.append({"role": "assistant", "content": text or ""})
break
name, arguments = tool_calls[0]
if name not in by_name:
parse_failures += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": f"[invalid tool {name!r} — choose one from the provided tool list]",
}
)
if parse_failures >= 2:
final_answer = (text or "").strip()
break
continue
tool = by_name[name]
# Empty-input guard: the small orchestrator often emits a <tool_call> with
# a blank/missing input on harder multi-turn tasks. Dispatching it returns
# a "no input provided" error observation that poisons the whole (often
# otherwise-correct) trajectory — the dominant clean-yield killer on hard
# tasks. Instead, nudge the model to resend WITH input and drop the
# malformed turn (don't record it), so it never enters the training data.
if tool.kind == "model" or tool.backend_type != "openjarvis-tool":
_has_input = any(
isinstance(arguments.get(k), str) and arguments.get(k).strip()
for k in ("input", "query", "code")
)
if not _has_input and empty_input_nudges < 3:
empty_input_nudges += 1
messages.append({"role": "assistant", "content": text or ""})
messages.append(
{
"role": "user",
"content": (
f"Your call to {name} had an empty 'input'. Resend the "
"<tool_call> with a non-empty 'input' field containing the "
"concrete sub-question to delegate."
),
}
)
continue
obs, dcost, dtok, _is_local = dispatch(tool, arguments)
cost += float(dcost)
tokens += int(dtok)
n_tool_calls += 1
turns.append(
UnifiedTurn(
reasoning=text or "",
tool_name=name,
arguments=dict(arguments),
observation=obs,
)
)
# The model's own action as an assistant turn, 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 <tool_call> 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 ``<tool_call>{...}</tool_call>`` text the model emits."""
return (
f"<tool_call>{json.dumps({'name': name, 'arguments': arguments})}</tool_call>"
)
__all__ = [
"RL_ORCHESTRATOR_SYS",
"UnifiedRollout",
"UnifiedTurn",
"build_system_prompt",
"run_unified_rollout",
"tool_call_tag",
]
src/openjarvis/agents/hybrid/toolorchestra/sandbox.py
+82
View File
@@ -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
src/openjarvis/agents/hybrid/toolorchestra/tracing.py
+170
View File
@@ -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
src/openjarvis/agents/hybrid/toolorchestra/unified.py
+280
View File
@@ -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 ``<tool_call>`` form regardless.
* ``False`` (SERVING a *fine-tuned* model): bake the ``<tools>`` block + JSON
``<tool_call>`` format into the system prompt (exactly what the serializer
trained on) and DROP ``tools=``. If ``tools=`` is set here, the served chat
template injects its own XML ``<function=>`` instructions, which conflict
with the JSON format the model learned -> it falls back to ``\\boxed{}`` and
never routes. ``_parse_rl_tool_call`` scrapes the JSON tag from raw text.
"""
# Create the client ONCE and reuse it for every turn. Creating a fresh OpenAI
# client per call (as before) leaks an httpx connection pool each time -> under
# heavy parallel generation, sockets pile up in CLOSE-WAIT, the run degrades and
# has to be restarted. The orchestrator is the highest-frequency call, so reusing
# one client here removes the dominant leak. httpx clients are thread-safe.
from openai import OpenAI
from openjarvis.agents.hybrid.toolorchestra.tracing import wrap_client
client = wrap_client(
OpenAI(base_url=base_url, api_key=api_key or "EMPTY", timeout=timeout)
)
def call_orchestrator(messages: List[Dict[str, Any]], specs: List[Dict[str, Any]]):
# ``messages`` is the full running conversation (system/user/assistant/tool)
# built by run_unified_rollout — pass it through so the model sees its own
# prior calls + tool responses, not a flattened user blob.
if native_tools:
# Base-model generation: native tool template drives the tool calls.
send = messages
kwargs = {"tools": specs} if specs else {}
else:
# Fine-tuned serving: JSON format is baked into the system prompt and
# tools= is dropped (see make_call_orchestrator docstring).
send = list(messages)
if specs and send and send[0].get("role") == "system":
send[0] = {"role": "system", "content": build_system_prompt(specs)}
kwargs = {}
# repetition_penalty is a vLLM extra (not OpenAI-native); only send it to a
# local vLLM endpoint (base_url set), never to the cloud frontier APIs.
if base_url and repetition_penalty and repetition_penalty != 1.0:
kwargs.setdefault("extra_body", {})["repetition_penalty"] = (
repetition_penalty
)
# 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"]
src/openjarvis/agents/hybrid/toolorchestra/workers.py
+477
View File
@@ -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 ``"<local>"``) resolves to
``local_model``; ``model = "$cloud"`` / ``"<cloud>"`` to ``cloud_model``.
On any validation failure, raises ``ValueError`` with the message
``"Invalid worker_pool entry [<id>]: <reason>"``. Fails fast at agent
init rather than mid-task.
"""
override = cfg.get("worker_pool")
if override is None:
return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint)
if not isinstance(override, list) or not override:
raise ValueError(
"Invalid worker_pool entry [-]: worker_pool must be a non-empty list"
)
resolved: List[Dict[str, Any]] = []
seen_ids: set = set()
has_non_search = False
for raw in override:
wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?"
if not isinstance(raw, dict):
raise ValueError(
f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict"
)
entry = dict(raw)
wid = entry.get("id")
if not isinstance(wid, int):
raise ValueError(
f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int"
)
if wid in seen_ids:
raise ValueError(f"Invalid worker_pool entry [{wid}]: duplicate id")
seen_ids.add(wid)
if not entry.get("name") or not isinstance(entry["name"], str):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string"
)
wtype = entry.get("type") or entry.get("endpoint")
if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'type' must be one of "
f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})"
)
wtype = wtype.lower()
entry["type"] = wtype
# Substitute $local / $cloud placeholders (before any model check).
model = entry.get("model")
if isinstance(model, str) and model in ("$local", "<local>"):
if not local_model:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model='{model}' "
"requires a local_model to be configured for this cell"
)
model = local_model
entry["model"] = model
elif isinstance(model, str) and model in ("$cloud", "<cloud>"):
model = cloud_model
entry["model"] = model
if wtype in _TOOLORCH_SEARCH_TYPES:
if model in (None, ""):
if wtype == "anthropic-web-search":
model = _DEFAULT_WEB_SEARCH_MODEL
elif wtype in ("openai-web-search", "gemini-web-search"):
model = cloud_model
else:
model = wtype
entry["model"] = model
elif not isinstance(model, str):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set"
)
if (
wtype in ("openai-web-search", "gemini-web-search")
and model not in PRICES
):
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model {model!r} "
f"is not in PRICES (known: {sorted(PRICES)})"
)
# Search workers don't satisfy the "needs a solver" requirement.
else:
if not isinstance(model, str) or not model:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
)
if wtype == "vllm":
if not entry.get("base_url"):
if not local_endpoint:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: vllm worker needs "
"'base_url' (or a configured local_endpoint to fall back to)"
)
entry["base_url"] = local_endpoint
entry.setdefault("api_key", "EMPTY")
else:
if model not in PRICES:
raise ValueError(
f"Invalid worker_pool entry [{wid}]: model {model!r} "
f"is not in PRICES (known: {sorted(PRICES)})"
)
has_non_search = True
entry.setdefault(
"description",
f"User-supplied {wtype} worker ({model}).",
)
resolved.append(entry)
if not has_non_search:
raise ValueError(
"Invalid worker_pool entry [-]: worker_pool must contain at least "
"one non-search worker (vllm / openai / anthropic / gemini)"
)
return resolved
# 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"]),
)
tests/agents/test_anonymize_tools.py
+129
View File
@@ -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
tests/agents/test_expert_registry.py
+221
View File
@@ -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"