diff --git a/src/openjarvis/agents/hybrid/registry/toolorchestra.toml b/src/openjarvis/agents/hybrid/registry/toolorchestra.toml index c6893725..ca97953a 100644 --- a/src/openjarvis/agents/hybrid/registry/toolorchestra.toml +++ b/src/openjarvis/agents/hybrid/registry/toolorchestra.toml @@ -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---- +# 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---- +# (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 diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py index 4eec4eb0..8d5c8494 100644 --- a/src/openjarvis/agents/hybrid/toolorchestra.py +++ b/src/openjarvis/agents/hybrid/toolorchestra.py @@ -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 ``{...}`` 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 ``{...}`` blocks emitted by Orchestrator-8B +# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a +# hermes-style template). Captures the JSON payload. +_TOOL_CALL_TAG_RE = re.compile( + r"\s*(\{.*?\})\s*", re.DOTALL +) + + +def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]: + """Return ``{"name": str, "arguments": dict}`` or None. + + Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls + back to scraping ``{...}`` tags from the raw + content. We take the first tool call only — Orchestrator-8B was trained + to emit exactly one per turn. + """ + # SDK-level path. + if sdk_tool_calls: + first = sdk_tool_calls[0] + name = getattr(getattr(first, "function", None), "name", None) + args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}" + try: + args = json.loads(args_raw) + except json.JSONDecodeError: + args = {} + if isinstance(name, str) and isinstance(args, dict): + return {"name": name, "arguments": args} + # Text-tag fallback. + if not isinstance(content, str): + return None + m = _TOOL_CALL_TAG_RE.search(content) + if not m: + return None + try: + obj = json.loads(m.group(1)) + except json.JSONDecodeError: + return None + if not isinstance(obj, dict): + return None + name = obj.get("name") + args = obj.get("arguments", {}) + if not isinstance(name, str) or not isinstance(args, dict): + return None + return {"name": name, "arguments": args} + + def _build_pool_block(workers: List[Dict[str, Any]]) -> str: return "\n".join( f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers @@ -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: 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"]