mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
hybrid: wire toolorchestra RL mode into SWE-bench loop
The toolorchestra cell that produced 0.560/$31.58 on GAIA (orch8b-opus47)
only ran the RL orchestrator-8B path on GAIA-shaped questions; SWE-bench
fell back to a one-shot cloud call. This unblocks the toolorchestra x SWE
column in the hybrid ablation: now the same n=100 orchestrator-8b +
opus-4-7 cell can run on swebench-verified, gated by
method_cfg.swe_use_agent_loop = true.
_run_rl now mirrors the prompted path's SWE handling — clones a shared
workdir per task, routes enhance_reasoning / answer worker dispatches
through run_swe_agent_loop on that workdir, falls back to the frontier
worker through the SWE loop too if the orchestrator never picks
`answer`, and appends the working-tree diff to the final answer so the
SWE-bench Modal scorer can extract the patch. bash turns count as
tool_calls per row (matches the 7f432453 contract). All cleanup runs
through a try/finally guard so /tmp doesn't leak cloned repos.
Smoke cell `toolorchestra-orch8b-opus47-swe-smoke2` (n=2) runs 2/2
without errors, score 1.0/0.0 (one diff fix lands, one doesn't —
expected for n=2). tool_calls populated, swe_mode=True in traces,
workdirs cleaned up. The n=100 cell now carries the
swe_use_agent_loop=true + matching swe_* knobs.
This commit is contained in:
Andrew Park
@@ -1,9 +1,29 @@
|
||||
# ToolOrchestra cells. Prompted port — uses a cloud model (Opus) as the
|
||||
# orchestrator, NOT the RL-trained Nemotron-Orchestrator-8B from the paper
|
||||
# (arXiv:2511.21689). Treat results as preliminary until a real
|
||||
# Orchestrator-8B deployment is wired up.
|
||||
# ToolOrchestra cells (arXiv:2511.21689).
|
||||
#
|
||||
# Naming: toolorchestra-<bench>-<local-short>-<cloud-short>-<N>
|
||||
# Two modes via `method_cfg.orchestrator_mode`:
|
||||
#
|
||||
# "prompted" (default) -- cloud model (Opus etc.) plays the orchestrator
|
||||
# over a heterogeneous worker pool. Useful as a
|
||||
# prompted upper-bound; NOT what the paper does.
|
||||
#
|
||||
# "rl" -- paper-faithful path. The RL-trained
|
||||
# `nvidia/Orchestrator-8B` served on a local vLLM
|
||||
# (default :8003) is the orchestrator and emits
|
||||
# OpenAI-style tool_calls for the three NVlabs
|
||||
# tools (enhance_reasoning / answer / search).
|
||||
# Expert pool maps tool-`model` slots to our
|
||||
# workers: *-1 -> cloud, *-2 -> gpt-5-mini,
|
||||
# *-3 -> local vLLM, search -> Anthropic web_search.
|
||||
# See `toolorchestra.py` module docstring for
|
||||
# what's collapsed vs. the upstream (no Tavily /
|
||||
# FAISS-wiki / Qwen-Coder code-interpreter).
|
||||
#
|
||||
# Naming: toolorchestra-<orch-short>-<cloud-short>-<bench>-<N>
|
||||
# (legacy prompted cells keep their original naming scheme.)
|
||||
|
||||
# ============================================================
|
||||
# Legacy prompted-mode cells (cloud-as-orchestrator).
|
||||
# ============================================================
|
||||
|
||||
[cells.toolorchestra-gaia-qwen27b-opus-3]
|
||||
method = "toolorchestra"
|
||||
@@ -28,3 +48,46 @@ n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { max_turns = 8, orchestrator_max_tokens = 1024, worker_max_tokens = 4096 }
|
||||
|
||||
# ============================================================
|
||||
# RL-mode cells -- Orchestrator-8B drives the loop (n=100).
|
||||
# ============================================================
|
||||
#
|
||||
# `local` carries the orchestrator's vLLM (nvidia/Orchestrator-8B on :8003).
|
||||
# Tier-3 expert calls (`answer-3`, `reasoner-3`, …) route to this same model.
|
||||
# `cloud` is the frontier worker invoked for tier-1 slots.
|
||||
# `gpt-5-mini` is hard-coded as the tier-2 mid worker (matches the paper's
|
||||
# `gpt-5-mini` slots in `eval_hle.py` MODEL_MAPPING).
|
||||
|
||||
[cells.toolorchestra-orch8b-opus47-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "orchestrator-8b", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "orchestrator-8b", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.toolorchestra-orch8b-opus47-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "orchestrator-8b", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "orchestrator-8b", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 3
|
||||
|
||||
# Smoke cell for the RL-mode SWE wiring (2026-05-19). n=2, no subset so it
|
||||
# just grabs the first two SWE-bench Verified tasks. Used to verify
|
||||
# end-to-end: workdir clone, _swe_call_worker dispatch through Orchestrator-8B,
|
||||
# diff extraction, Modal-harness scoring. Promote to n=100 once green.
|
||||
[cells.toolorchestra-orch8b-opus47-swe-smoke2]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 2
|
||||
local = { model = "orchestrator-8b", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "orchestrator-8b", max_turns = 6, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 15, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 1
|
||||
|
||||
@@ -1,26 +1,40 @@
|
||||
"""ToolOrchestraAgent — prompted port of NVlabs ToolOrchestra (arXiv:2511.21689).
|
||||
"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689).
|
||||
|
||||
The paper RL-trains an 8B Orchestrator (``nvidia/Nemotron-Orchestrator-8B``)
|
||||
to coordinate basic tools + specialist LLMs + generalist LLMs in
|
||||
multi-turn agentic loops, ranked #1 on GAIA at release.
|
||||
Two modes, gated by ``method_cfg.orchestrator_mode``:
|
||||
|
||||
The hybrid harness adapter for ToolOrchestra is a documented stub —
|
||||
running the real thing needs a separate vLLM srun for the Orchestrator-8B
|
||||
checkpoint, a FAISS wiki retriever, a Tavily API key, and a refactor of
|
||||
the upstream eval scripts. None of that fits in our cluster allocation.
|
||||
* ``"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.
|
||||
|
||||
This port keeps the same scope discipline: **inference-time only,
|
||||
prompted, no RL**. A cloud model plays the role of the orchestrator,
|
||||
dispatching to a pool of `(tool | specialist_llm | generalist_llm)`
|
||||
workers in a reactive loop. The loop is the paradigm; the orchestrator
|
||||
weights are not.
|
||||
* ``"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.
|
||||
|
||||
Why ship this at all if it's not the "real" thing? Because the prompted
|
||||
upper-bound is useful as a reference point alongside the other paradigms,
|
||||
and because the OpenJarvis registry needs all six entries for the
|
||||
distillation pipeline to slot ToolOrchestra in alongside the rest.
|
||||
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:
|
||||
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,
|
||||
@@ -31,10 +45,6 @@ Pipeline per task:
|
||||
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).
|
||||
|
||||
Not yet validated end-to-end in the hybrid harness — the hybrid adapter
|
||||
``raise NotImplementedError``s. Treat results from this paradigm as
|
||||
preliminary until we have a real ToolOrchestra-8B deployment.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -85,6 +95,190 @@ FORCE_FINAL_PROMPT = (
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# RL-mode constants (Orchestrator-8B, paper-faithful).
|
||||
# ============================================================================
|
||||
#
|
||||
# Verbatim copies of the upstream system prompt / user-prompt template / tools
|
||||
# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't
|
||||
# edit the description text — Orchestrator-8B was RL-trained against this
|
||||
# exact wording and pricing/latency table.
|
||||
|
||||
RL_ORCHESTRATOR_SYS = "You are good at using tools."
|
||||
|
||||
RL_TOOLS_SPEC: List[Dict[str, Any]] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "enhance_reasoning",
|
||||
"description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "answer",
|
||||
"description": "give the final answer. Not allowed to call if documents is empty.",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"description": "Search for missing information",
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"model": {
|
||||
"description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["model"],
|
||||
"title": "parameters",
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104).
|
||||
RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = {
|
||||
"enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]},
|
||||
"answer": {
|
||||
"model": [
|
||||
"answer-1", "answer-2", "answer-3", "answer-4",
|
||||
"answer-math-1", "answer-math-2",
|
||||
],
|
||||
},
|
||||
"search": {"model": ["search-1", "search-2", "search-3"]},
|
||||
}
|
||||
|
||||
# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec.
|
||||
# Tiers ranked by the upstream tools.json table (`*-1` = frontier,
|
||||
# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same
|
||||
# tiers since we don't have Qwen-Math served.
|
||||
#
|
||||
# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict`
|
||||
# so the substitution is deferred until we know the cell's resolved local/cloud
|
||||
# pair. Worker dicts share the schema validated by `_resolve_worker_pool`.
|
||||
|
||||
def _expert_for(slot: str, local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str) -> Dict[str, Any]:
|
||||
"""Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec.
|
||||
|
||||
Routing policy:
|
||||
- `*-1` (frontier tier) -> cloud (`cloud_model`)
|
||||
- `*-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"):
|
||||
return {
|
||||
"name": f"frontier:{slot}",
|
||||
"type": "anthropic",
|
||||
"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",
|
||||
}
|
||||
|
||||
|
||||
# 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
|
||||
@@ -454,10 +648,11 @@ def _swe_call_worker(
|
||||
|
||||
@AgentRegistry.register("toolorchestra")
|
||||
class ToolOrchestraAgent(LocalCloudAgent):
|
||||
"""Prompted multi-turn dispatcher over a mixed worker pool.
|
||||
"""Multi-turn dispatcher over a mixed worker pool.
|
||||
|
||||
Inference-only port — does NOT use the RL-trained Nemotron-Orchestrator-8B.
|
||||
See module docstring for what's missing relative to the published paper.
|
||||
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"
|
||||
@@ -474,12 +669,33 @@ class ToolOrchestraAgent(LocalCloudAgent):
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
)
|
||||
# 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
|
||||
@@ -705,5 +921,301 @@ class ToolOrchestraAgent(LocalCloudAgent):
|
||||
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))
|
||||
|
||||
# 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. ``trace_role="orchestrator"`` keeps these
|
||||
# rows distinguishable from worker calls in the trace log.
|
||||
text, o_in, o_out = LocalCloudAgent._call_vllm(
|
||||
orch_model,
|
||||
orch_endpoint,
|
||||
user=user,
|
||||
system=RL_ORCHESTRATOR_SYS,
|
||||
max_tokens=orch_max_tokens,
|
||||
temperature=orch_temp,
|
||||
enable_thinking=False,
|
||||
tools=RL_TOOLS_SPEC,
|
||||
trace_role="orchestrator",
|
||||
)
|
||||
tokens_local += o_in + o_out
|
||||
|
||||
# The trace event recorded by ``_call_vllm`` already contains
|
||||
# any SDK-level tool_calls. Re-parse from text for the tag-style
|
||||
# fallback (qwen3_xml parser misses hermes-style tags).
|
||||
# We don't have direct access to the SDK tool_calls object here —
|
||||
# parse from text only. (`_call_vllm` returns just (text, p, c).)
|
||||
action = _parse_rl_tool_call(text, None)
|
||||
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
|
||||
|
||||
worker = _expert_for(
|
||||
slot, self._local_model, self._local_endpoint, self._cloud_model
|
||||
)
|
||||
|
||||
# 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":
|
||||
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":
|
||||
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)
|
||||
|
||||
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,
|
||||
})
|
||||
|
||||
# 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.
|
||||
worker = _expert_for(
|
||||
"answer-1", self._local_model, self._local_endpoint,
|
||||
self._cloud_model,
|
||||
)
|
||||
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",
|
||||
"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"]
|
||||
|
||||
Reference in New Issue
Block a user