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