From 08916ec4683eec5e56140492106e54c12bf37649 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 17 Jun 2026 11:35:55 -0700 Subject: [PATCH] feat(orchestrator): faithful ToolOrchestra unified-tool + rejection-sampling SFT pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 ( 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. --- scripts/orchestrator/build_unified_sft.py | 97 + .../agents/hybrid/expert_registry.py | 311 +++ src/openjarvis/agents/hybrid/toolorchestra.py | 1662 ----------------- .../agents/hybrid/toolorchestra/__init__.py | 20 + .../agents/hybrid/toolorchestra/agent.py | 771 ++++++++ .../agents/hybrid/toolorchestra/clients.py | 50 + .../agents/hybrid/toolorchestra/experts.py | 170 ++ .../agents/hybrid/toolorchestra/parsing.py | 138 ++ .../agents/hybrid/toolorchestra/prompts.py | 106 ++ .../agents/hybrid/toolorchestra/rollout.py | 142 ++ .../agents/hybrid/toolorchestra/sandbox.py | 79 + .../agents/hybrid/toolorchestra/unified.py | 110 ++ .../agents/hybrid/toolorchestra/workers.py | 421 +++++ .../orchestrator/sft_data/reject_sample.py | 119 ++ .../orchestrator/sft_data/toolscale.py | 128 ++ .../sft_data/unified_serialize.py | 92 + tests/agents/test_expert_registry.py | 101 + .../sft_data/test_rejection_pipeline.py | 152 ++ 18 files changed, 3007 insertions(+), 1662 deletions(-) create mode 100644 scripts/orchestrator/build_unified_sft.py create mode 100644 src/openjarvis/agents/hybrid/expert_registry.py delete mode 100644 src/openjarvis/agents/hybrid/toolorchestra.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/__init__.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/agent.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/clients.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/experts.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/parsing.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/prompts.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/rollout.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/sandbox.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/unified.py create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/workers.py create mode 100644 src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py create mode 100644 src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py create mode 100644 src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py create mode 100644 tests/agents/test_expert_registry.py create mode 100644 tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py diff --git a/scripts/orchestrator/build_unified_sft.py b/scripts/orchestrator/build_unified_sft.py new file mode 100644 index 00000000..e4eea00b --- /dev/null +++ b/scripts/orchestrator/build_unified_sft.py @@ -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()) diff --git a/src/openjarvis/agents/hybrid/expert_registry.py b/src/openjarvis/agents/hybrid/expert_registry.py new file mode 100644 index 00000000..703b2193 --- /dev/null +++ b/src/openjarvis/agents/hybrid/expert_registry.py @@ -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", +] diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py deleted file mode 100644 index d5ebd78d..00000000 --- a/src/openjarvis/agents/hybrid/toolorchestra.py +++ /dev/null @@ -1,1662 +0,0 @@ -"""ToolOrchestraAgent — port of NVlabs ToolOrchestra (arXiv:2511.21689). - -Two modes, gated by ``method_cfg.orchestrator_mode``: - -* ``"prompted"`` (default, legacy): a cloud model (Opus etc.) plays the - orchestrator, dispatching to a numbered worker pool via JSON - ``{"action": "call_worker"|"final_answer", ...}`` actions. Useful as - a prompted upper-bound reference point — NOT the paper's setup. - -* ``"rl"`` (paper-faithful): the RL-trained ``nvidia/Orchestrator-8B`` - served on a local vLLM is the orchestrator. It emits OpenAI-style - ``tool_calls`` (or ``{...}`` text blocks when - vLLM's tool parser doesn't catch them) for three expert tools — - ``enhance_reasoning``, ``answer``, ``search`` — exactly as in the - upstream ``evaluation/tools.json``. Each tool's ``model`` arg - (``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real - backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier - Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen - for `*-3`. Search routes to the 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 re -import shutil -import tempfile -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -from openjarvis.agents._stubs import AgentContext -from openjarvis.agents.hybrid._base import ( - ANTHROPIC_WEB_SEARCH_TOOL, - 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 ( - _clone_repo, - _extract_diff, - run_swe_agent_loop, -) -from openjarvis.core.registry import AgentRegistry - -ORCHESTRATOR_SYS = """\ -You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms: - - {"action": "call_worker", "worker_id": , "input": ""} - - {"action": "final_answer", "answer": ""} - -Strategy: - -- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math). -- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass. -- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase. -- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question. -""" - -FORCE_FINAL_PROMPT = ( - "Worker-call budget exhausted. Emit `final_answer` now using everything " - "you've learned. Respect the question's answer-format rules." -) - - -# ============================================================================ -# RL-mode constants (Orchestrator-8B, paper-faithful). -# ============================================================================ -# -# Verbatim copies of the upstream system prompt / user-prompt template / tools -# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't -# edit the description text — Orchestrator-8B was RL-trained against this -# exact wording and pricing/latency table. - -RL_ORCHESTRATOR_SYS = "You are good at using tools." - -RL_TOOLS_SPEC: List[Dict[str, Any]] = [ - { - "type": "function", - "function": { - "name": "enhance_reasoning", - "description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem", - "parameters": { - "properties": { - "model": { - "description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s", - "type": "string", - } - }, - "required": ["model"], - "title": "parameters", - "type": "object", - }, - }, - }, - { - "type": "function", - "function": { - "name": "answer", - "description": "give the final answer. Not allowed to call if documents is empty.", - "parameters": { - "properties": { - "model": { - "description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s", - "type": "string", - } - }, - "required": ["model"], - "title": "parameters", - "type": "object", - }, - }, - }, - { - "type": "function", - "function": { - "name": "search", - "description": "Search for missing information", - "parameters": { - "properties": { - "model": { - "description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s", - "type": "string", - } - }, - "required": ["model"], - "title": "parameters", - "type": "object", - }, - }, - }, -] - -# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104). -RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = { - "enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]}, - "answer": { - "model": [ - "answer-1", "answer-2", "answer-3", "answer-4", - "answer-math-1", "answer-math-2", - ], - }, - "search": {"model": ["search-1", "search-2", "search-3"]}, -} - -# Map the orchestrator's `model` slot to a concrete OpenJarvis worker spec. -# Tiers ranked by the upstream tools.json table (`*-1` = frontier, -# `*-2` = mid, `*-3` = local). math-1 / math-2 collapse onto the same -# tiers since we don't have Qwen-Math served. -# -# Each entry is a callable `(local_model, local_endpoint, cloud_model) -> worker_dict` -# so the substitution is deferred until we know the cell's resolved local/cloud -# pair. Worker dicts share the schema validated by `_resolve_worker_pool`. - -def _expert_for(slot: str, local_model: Optional[str], - local_endpoint: Optional[str], - cloud_model: str, - cloud_endpoint: str = "anthropic") -> Dict[str, Any]: - """Map an upstream model slot (`answer-1`, `search-3`, …) to a worker spec. - - Routing policy: - - `*-1` (frontier tier) -> cloud (`cloud_model`), wtype keyed off - `cloud_endpoint` ("anthropic"/"openai"/"gemini") - - `*-2` (mid tier) -> cloud `gpt-5-mini` (matches the paper's - cost tier for mid OpenAI calls) - - `*-3` (local tier) -> local vLLM (`local_model`) - - `answer-math-*` -> same tiers as the numeric suffix - - `search-*` -> 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", - } - - -# ---- 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 - - -def _call_orchestrator_with_tool_calls( - model: str, - endpoint: str, - *, - user: str, - system: str, - max_tokens: int, - temperature: float, - tools: List[Dict[str, Any]], - timeout: float = 600.0, -) -> Tuple[str, int, int, Any]: - """Orchestrator-aware vLLM call. Returns (text, p_tok, c_tok, tool_calls). - - Mirrors ``LocalCloudAgent._call_vllm`` but ALSO surfaces the SDK-level - ``tool_calls`` object so the RL-mode parser can match against it - directly. Otherwise vLLM's tool parser silently swallows the tool call - into the SDK field while leaving ``content == ''`` — and the text-tag - parser sees nothing, falling through to the answer-1 fallback. (Bug - observed 2026-05-19 on the paper-match smoke; same path was buggy on - the default pool too, just less reproducibly.) - """ - from openai import OpenAI - - client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout) - messages = [ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ] - resp = client.chat.completions.create( - model=model, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - tools=tools, - extra_body={"chat_template_kwargs": {"enable_thinking": False}}, - ) - choice = resp.choices[0] - message = choice.message - text = message.content or "" - tool_calls = getattr(message, "tool_calls", None) - u = resp.usage - p = getattr(u, "prompt_tokens", 0) if u else 0 - c = getattr(u, "completion_tokens", 0) if u else 0 - return text, p, c, tool_calls - - -def _paper_pool( - local_model: Optional[str], - local_endpoint: Optional[str], -) -> List[Dict[str, Any]]: - """Paper-match worker pool (registered for traces / inspection). - - NOTE: in RL mode the orchestrator dispatches via tool/slot rather than - worker_id, so this list is purely informational — `_paper_expert_for` - is the actual routing function. We still return a list here so the - paradigm's trace metadata has something concrete to log. - """ - pool: List[Dict[str, Any]] = [] - if local_model and local_endpoint: - pool.append({ - "id": len(pool), - "name": "local-qwen", - "type": "vllm", - "model": local_model, - "base_url": local_endpoint, - "description": "Local Qwen vLLM (paper uses Qwen3-32B).", - }) - pool.append({ - "id": len(pool), "name": "tavily-search", - "type": "tavily-search", "model": "tavily", - "description": "Tavily web search.", - }) - pool.append({ - "id": len(pool), "name": "modal-python", - "type": "modal-python", "model": "modal-python", - "description": "Modal Sandbox for one-shot Python exec.", - }) - pool.append({ - "id": len(pool), "name": "code-specialist", - "type": "openrouter", "model": _PAPER_CODER_OPENROUTER, - "description": "Qwen-2.5-Coder-32B via OpenRouter (paper).", - }) - pool.append({ - "id": len(pool), "name": "generalist-llama", - "type": "openrouter", "model": _PAPER_GENERALIST_TIER3_OPENROUTER, - "description": "Llama-3.3-70B-Instruct via OpenRouter (paper tier-3).", - }) - pool.append({ - "id": len(pool), "name": "generalist-gpt5", - "type": "openai", "model": "gpt-5", - "description": "GPT-5 frontier generalist.", - }) - pool.append({ - "id": len(pool), "name": "generalist-gpt5-mini", - "type": "openai", "model": "gpt-5-mini", - "description": "GPT-5-mini mid generalist.", - }) - return pool - - -# Regex for ``{...}`` blocks emitted by Orchestrator-8B -# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a -# hermes-style template). Captures the JSON payload. -_TOOL_CALL_TAG_RE = re.compile( - r"\s*(\{.*?\})\s*", re.DOTALL -) - - -def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]: - """Return ``{"name": str, "arguments": dict}`` or None. - - Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls - back to scraping ``{...}`` tags from the raw - content. We take the first tool call only — Orchestrator-8B was trained - to emit exactly one per turn. - """ - # SDK-level path. - if sdk_tool_calls: - first = sdk_tool_calls[0] - name = getattr(getattr(first, "function", None), "name", None) - args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}" - try: - args = json.loads(args_raw) - except json.JSONDecodeError: - args = {} - if isinstance(name, str) and isinstance(args, dict): - return {"name": name, "arguments": args} - # Text-tag fallback. - if not isinstance(content, str): - return None - m = _TOOL_CALL_TAG_RE.search(content) - if not m: - return None - try: - obj = json.loads(m.group(1)) - except json.JSONDecodeError: - return None - if not isinstance(obj, dict): - return None - name = obj.get("name") - args = obj.get("arguments", {}) - if not isinstance(name, str) or not isinstance(args, dict): - return None - return {"name": name, "arguments": args} - - -def _build_pool_block(workers: List[Dict[str, Any]]) -> str: - return "\n".join( - f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers - ) - - -def _build_user_prompt( - question: str, - workers: List[Dict[str, Any]], - history: List[Dict[str, Any]], -) -> str: - pieces = [ - f"Worker pool:\n{_build_pool_block(workers)}", - f"User question:\n{question}", - ] - if history: - pieces.append("Conversation so far (orchestrator turns and worker outputs):") - for h in history: - if h["role"] == "orchestrator": - pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}") - else: - pieces.append( - f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n" - f"{h['output']}" - ) - pieces.append( - "Emit the next JSON action object now — exactly one object, no prose." - ) - return "\n\n".join(pieces) - - -def _strip_fences(s: str) -> str: - s = s.strip() - if s.startswith("```"): - first_nl = s.find("\n") - if first_nl != -1: - s = s[first_nl + 1:] - if s.endswith("```"): - s = s[:-3] - s = s.strip() - return s - - -def _parse_action(text: str) -> Optional[Dict[str, Any]]: - s = _strip_fences(text) - # First try direct parse, then balanced-brace extraction. - try: - obj = json.loads(s) - if isinstance(obj, dict) and "action" in obj: - return obj - except json.JSONDecodeError: - pass - start = s.find("{") - if start == -1: - return None - depth = 0 - for i in range(start, len(s)): - c = s[i] - if c == "{": - depth += 1 - elif c == "}": - depth -= 1 - if depth == 0: - try: - obj = json.loads(s[start : i + 1]) - if isinstance(obj, dict) and "action" in obj: - return obj - except json.JSONDecodeError: - return None - return None - - -def _extract_final_answer_text(text: str) -> str: - """Best-effort: pull the answer string from a malformed action emission. - - Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line. - """ - m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL) - if m: - return m.group(1).encode("utf-8").decode("unicode_escape") - m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE) - if m: - return m.group(1).strip() - return text.strip() - - -# ---------- Worker pool ---------- - -def _default_pool( - local_model: Optional[str], - local_endpoint: Optional[str], - cloud_model: str = "claude-opus-4-7", - cloud_endpoint: str = "anthropic", -) -> List[Dict[str, Any]]: - """Default heterogeneous worker pool. - - The frontier worker's ``type`` + ``model`` track the cell's resolved - ``(cloud_model, cloud_endpoint)`` pair so non-Anthropic cells (gpt-5, - gemini-2.5-pro, …) route their frontier slot to the right SDK. - """ - ep = (cloud_endpoint or "anthropic").lower() - if ep not in ("anthropic", "openai", "gemini"): - ep = "anthropic" - pool: List[Dict[str, Any]] = [] - if local_model and local_endpoint: - pool.append({ - "id": len(pool), - "name": "local-qwen", - "type": "vllm", - "model": local_model, - "base_url": local_endpoint, - "description": ( - "Open-weights Qwen3.5 served locally. Cheap and fast. Good at " - "concise extraction, formatting, arithmetic on given data." - ), - }) - 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 ``""``) resolves to - ``local_model``; ``model = "$cloud"`` / ``""`` to ``cloud_model``. - - On any validation failure, raises ``ValueError`` with the message - ``"Invalid worker_pool entry []: "``. Fails fast at agent - init rather than mid-task. - """ - override = cfg.get("worker_pool") - if override is None: - return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint) - if not isinstance(override, list) or not override: - raise ValueError( - "Invalid worker_pool entry [-]: worker_pool must be a non-empty list" - ) - - resolved: List[Dict[str, Any]] = [] - seen_ids: set = set() - has_non_search = False - for raw in override: - wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?" - if not isinstance(raw, dict): - raise ValueError( - f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict" - ) - entry = dict(raw) - wid = entry.get("id") - if not isinstance(wid, int): - raise ValueError( - f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int" - ) - if wid in seen_ids: - raise ValueError( - f"Invalid worker_pool entry [{wid}]: duplicate id" - ) - seen_ids.add(wid) - if not entry.get("name") or not isinstance(entry["name"], str): - raise ValueError( - f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string" - ) - wtype = entry.get("type") or entry.get("endpoint") - if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES: - raise ValueError( - f"Invalid worker_pool entry [{wid}]: 'type' must be one of " - f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})" - ) - wtype = wtype.lower() - entry["type"] = wtype - # Substitute $local / $cloud placeholders (before any model check). - model = entry.get("model") - if isinstance(model, str) and model in ("$local", ""): - if not local_model: - raise ValueError( - f"Invalid worker_pool entry [{wid}]: model='{model}' " - "requires a local_model to be configured for this cell" - ) - model = local_model - entry["model"] = model - elif isinstance(model, str) and model in ("$cloud", ""): - model = cloud_model - entry["model"] = model - if wtype == "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"]), - ) - - -@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: convention)." - ) - - # SWE mode: route enhance_reasoning / answer workers through - # the SWE agent loop on the shared workdir so they can read - # files, run tests, and edit the working tree. Search workers - # stay one-shot (no agent loop). The `_swe_call_worker` - # one-shot fallbacks (openai-typed workers, search) return - # bash_turns=0; vllm/anthropic-typed workers run the loop. - bash_turns = 0 - if swe_mode and shared_workdir is not None and name != "search": - (w_text, w_in, w_out, is_local, extra_cost, - n_searches, bash_turns) = _swe_call_worker( - worker, w_input, cfg, task_meta, shared_workdir, turn, - ) - else: - w_text, w_in, w_out, is_local, extra_cost, n_searches = _call_worker( - worker, w_input, cfg - ) - if is_local: - tokens_local += w_in + w_out - else: - tokens_cloud += w_in + w_out - cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost - n_web_searches_total += n_searches - # SWE bash turns count as tool calls (each one is a $BASH block - # the agent executed). On non-SWE turns fall back to the - # original "at least one expert call" accounting. - tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches) - - # Paper-match: pipe coder output through a Modal sandbox so - # `enhance_reasoning` actually executes the code the coder - # wrote. Append the exec output to the worker's text. No-op - # when no python block is found. - modal_exec_output: Optional[str] = None - modal_exec_rc: Optional[int] = None - if (paper_mode and name == "enhance_reasoning" - and not swe_mode): - code = _extract_first_python_block(w_text) - if code: - timeout_s = int(cfg.get("modal_python_timeout_s", 60)) - modal_exec_output, modal_exec_rc = _call_modal_python( - code, timeout_s=timeout_s, - ) - tool_calls += 1 - w_text = ( - f"{w_text}\n\n[modal-python stdout/stderr " - f"(rc={modal_exec_rc})]\n{modal_exec_output}" - ) - - history.append({ - "role": "worker", - "turn": turn, - "tool": name, - "slot": slot, - "worker_model": worker["model"], - "worker_type": worker["type"], - "output": w_text, - "tokens_in": w_in, - "tokens_out": w_out, - "n_web_searches": n_searches, - "bash_turns": bash_turns, - "modal_exec_rc": modal_exec_rc, - }) - - # Update accumulated context for the next turn. - if name == "search": - # Treat the search worker's response as a document. - doc_list.append(w_text) - ctx_docs = "\n\n".join( - f"Doc {i+1}: {d}" for i, d in enumerate(doc_list) - ) - # Crude char-level cap mirrors the upstream's ~24k token cap. - context_str = ("Documents:\n" + ctx_docs)[-24000:] - elif name == "enhance_reasoning": - snippet = f"\n\nReasoning/exec output:\n{w_text}" - context_str = (context_str + snippet)[-24000:] - else: # answer - final_answer = w_text.strip() - break - - if final_answer is None: - # Hard fallback: ask the frontier worker directly. In SWE - # mode route this final call through the agent loop too so - # it can still touch the workdir and emit a diff. - expert_fn = _paper_expert_for if paper_mode else _expert_for - worker = expert_fn( - "answer-1", self._local_model, self._local_endpoint, - self._cloud_model, self._cloud_endpoint, - ) - fb_bash_turns = 0 - if swe_mode and shared_workdir is not None: - (ans, w_in, w_out, is_local, extra_cost, - _, fb_bash_turns) = _swe_call_worker( - worker, question, cfg, task_meta, - shared_workdir, max_turns + 1, - ) - tool_calls += fb_bash_turns - else: - ans, w_in, w_out, is_local, extra_cost, _ = _call_worker( - worker, question, cfg - ) - if is_local: - tokens_local += w_in + w_out - else: - tokens_cloud += w_in + w_out - cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost - history.append({ - "role": "worker", - "turn": max_turns + 1, - "tool": "answer", - "slot": "answer-1", - "worker_model": worker["model"], - "worker_type": worker["type"], - "output": ans, - "tokens_in": w_in, - "tokens_out": w_out, - "bash_turns": fb_bash_turns, - "fallback": True, - }) - final_answer = ans - - # In SWE mode, the authoritative output is the working-tree diff — - # frame it so `_score_swebench`'s extract_patch picks it up. - if swe_mode and shared_workdir is not None: - patch = _extract_diff(shared_workdir) - if patch.strip(): - final_answer = ( - f"{final_answer}\n\n```diff\n{patch}```" - if final_answer else f"```diff\n{patch}```" - ) - - meta = { - "tokens_local": tokens_local, - "tokens_cloud": tokens_cloud, - "cost_usd": cost, - "turns": len([h for h in history if h["role"] == "orchestrator"]), - "web_search_uses": n_web_searches_total, - "tool_calls": int(tool_calls), - "traces": { - "history": history, - "parse_failures": parse_failures, - "orchestrator_model": orch_model, - "orchestrator_endpoint": orch_endpoint, - "mode": "rl", - "pool": "paper" if paper_mode else "default", - "swe_mode": swe_mode, - "note": ( - "RL-trained nvidia/Orchestrator-8B as orchestrator. " - "Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto " - "our hybrid worker types — see toolorchestra.py docstring." - ), - }, - } - return final_answer, meta - finally: - if shared_workdir is not None: - shutil.rmtree(shared_workdir, ignore_errors=True) - - -__all__ = ["ToolOrchestraAgent"] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/__init__.py b/src/openjarvis/agents/hybrid/toolorchestra/__init__.py new file mode 100644 index 00000000..33226abe --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/__init__.py @@ -0,0 +1,20 @@ +"""ToolOrchestraAgent package (split from the former toolorchestra.py module). + +Importing this package registers the agent and re-exports the public surface, +so ``from openjarvis.agents.hybrid.toolorchestra import ToolOrchestraAgent`` +keeps working unchanged. Submodules: + + prompts — system prompts, RL tool specs / arg schema + experts — slot -> backend worker mapping (default + paper-match) + sandbox — Tavily search + Modal Python sandbox helpers + clients — orchestrator vLLM tool-call client + parsing — action / tool-call parsing + user-prompt assembly + workers — worker pool resolution + dispatch (_call_worker etc.) + agent — ToolOrchestraAgent (the registered agent class) +""" + +from __future__ import annotations + +from openjarvis.agents.hybrid.toolorchestra.agent import ToolOrchestraAgent + +__all__ = ["ToolOrchestraAgent"] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/agent.py b/src/openjarvis/agents/hybrid/toolorchestra/agent.py new file mode 100644 index 00000000..b2fc963a --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/agent.py @@ -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 ``{...}`` 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: convention)." + ) + + # SWE mode: route enhance_reasoning / answer workers through + # the SWE agent loop on the shared workdir so they can read + # files, run tests, and edit the working tree. Search workers + # stay one-shot (no agent loop). The `_swe_call_worker` + # one-shot fallbacks (openai-typed workers, search) return + # bash_turns=0; vllm/anthropic-typed workers run the loop. + bash_turns = 0 + if swe_mode and shared_workdir is not None and name != "search": + (w_text, w_in, w_out, is_local, extra_cost, + n_searches, bash_turns) = _swe_call_worker( + worker, w_input, cfg, task_meta, shared_workdir, turn, + ) + else: + w_text, w_in, w_out, is_local, extra_cost, n_searches = _call_worker( + worker, w_input, cfg + ) + if is_local: + tokens_local += w_in + w_out + else: + tokens_cloud += w_in + w_out + cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost + n_web_searches_total += n_searches + # SWE bash turns count as tool calls (each one is a $BASH block + # the agent executed). On non-SWE turns fall back to the + # original "at least one expert call" accounting. + tool_calls += bash_turns if bash_turns > 0 else max(1, n_searches) + + # Paper-match: pipe coder output through a Modal sandbox so + # `enhance_reasoning` actually executes the code the coder + # wrote. Append the exec output to the worker's text. No-op + # when no python block is found. + modal_exec_output: Optional[str] = None + modal_exec_rc: Optional[int] = None + if (paper_mode and name == "enhance_reasoning" + and not swe_mode): + code = _extract_first_python_block(w_text) + if code: + timeout_s = int(cfg.get("modal_python_timeout_s", 60)) + modal_exec_output, modal_exec_rc = _call_modal_python( + code, timeout_s=timeout_s, + ) + tool_calls += 1 + w_text = ( + f"{w_text}\n\n[modal-python stdout/stderr " + f"(rc={modal_exec_rc})]\n{modal_exec_output}" + ) + + history.append({ + "role": "worker", + "turn": turn, + "tool": name, + "slot": slot, + "worker_model": worker["model"], + "worker_type": worker["type"], + "output": w_text, + "tokens_in": w_in, + "tokens_out": w_out, + "n_web_searches": n_searches, + "bash_turns": bash_turns, + "modal_exec_rc": modal_exec_rc, + }) + + # Update accumulated context for the next turn. + if name == "search": + # Treat the search worker's response as a document. + doc_list.append(w_text) + ctx_docs = "\n\n".join( + f"Doc {i+1}: {d}" for i, d in enumerate(doc_list) + ) + # Crude char-level cap mirrors the upstream's ~24k token cap. + context_str = ("Documents:\n" + ctx_docs)[-24000:] + elif name == "enhance_reasoning": + snippet = f"\n\nReasoning/exec output:\n{w_text}" + context_str = (context_str + snippet)[-24000:] + else: # answer + final_answer = w_text.strip() + break + + if final_answer is None: + # Hard fallback: ask the frontier worker directly. In SWE + # mode route this final call through the agent loop too so + # it can still touch the workdir and emit a diff. + expert_fn = _paper_expert_for if paper_mode else _expert_for + worker = expert_fn( + "answer-1", self._local_model, self._local_endpoint, + self._cloud_model, self._cloud_endpoint, + ) + fb_bash_turns = 0 + if swe_mode and shared_workdir is not None: + (ans, w_in, w_out, is_local, extra_cost, + _, fb_bash_turns) = _swe_call_worker( + worker, question, cfg, task_meta, + shared_workdir, max_turns + 1, + ) + tool_calls += fb_bash_turns + else: + ans, w_in, w_out, is_local, extra_cost, _ = _call_worker( + worker, question, cfg + ) + if is_local: + tokens_local += w_in + w_out + else: + tokens_cloud += w_in + w_out + cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost + history.append({ + "role": "worker", + "turn": max_turns + 1, + "tool": "answer", + "slot": "answer-1", + "worker_model": worker["model"], + "worker_type": worker["type"], + "output": ans, + "tokens_in": w_in, + "tokens_out": w_out, + "bash_turns": fb_bash_turns, + "fallback": True, + }) + final_answer = ans + + # In SWE mode, the authoritative output is the working-tree diff — + # frame it so `_score_swebench`'s extract_patch picks it up. + if swe_mode and shared_workdir is not None: + patch = _extract_diff(shared_workdir) + if patch.strip(): + final_answer = ( + f"{final_answer}\n\n```diff\n{patch}```" + if final_answer else f"```diff\n{patch}```" + ) + + meta = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": cost, + "turns": len([h for h in history if h["role"] == "orchestrator"]), + "web_search_uses": n_web_searches_total, + "tool_calls": int(tool_calls), + "traces": { + "history": history, + "parse_failures": parse_failures, + "orchestrator_model": orch_model, + "orchestrator_endpoint": orch_endpoint, + "mode": "rl", + "pool": "paper" if paper_mode else "default", + "swe_mode": swe_mode, + "note": ( + "RL-trained nvidia/Orchestrator-8B as orchestrator. " + "Expert pool collapses Tavily/FAISS/Qwen-Math/Coder onto " + "our hybrid worker types — see toolorchestra.py docstring." + ), + }, + } + return final_answer, meta + finally: + if shared_workdir is not None: + shutil.rmtree(shared_workdir, ignore_errors=True) + + +__all__ = ["ToolOrchestraAgent"] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/clients.py b/src/openjarvis/agents/hybrid/toolorchestra/clients.py new file mode 100644 index 00000000..d941862b --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/clients.py @@ -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 diff --git a/src/openjarvis/agents/hybrid/toolorchestra/experts.py b/src/openjarvis/agents/hybrid/toolorchestra/experts.py new file mode 100644 index 00000000..a4646a63 --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/experts.py @@ -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", + } diff --git a/src/openjarvis/agents/hybrid/toolorchestra/parsing.py b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py new file mode 100644 index 00000000..d24d2e3c --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/parsing.py @@ -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 ``{...}`` blocks emitted by Orchestrator-8B +# when the vLLM tool parser doesn't catch them (e.g. `qwen3_xml` parser on a +# hermes-style template). Captures the JSON payload. +_TOOL_CALL_TAG_RE = re.compile( + r"\s*(\{.*?\})\s*", re.DOTALL +) + + +def _parse_rl_tool_call(content: str, sdk_tool_calls: Any) -> Optional[Dict[str, Any]]: + """Return ``{"name": str, "arguments": dict}`` or None. + + Prefers the SDK-level ``tool_calls`` (when vLLM's parser matched), falls + back to scraping ``{...}`` tags from the raw + content. We take the first tool call only — Orchestrator-8B was trained + to emit exactly one per turn. + """ + # SDK-level path. + if sdk_tool_calls: + first = sdk_tool_calls[0] + name = getattr(getattr(first, "function", None), "name", None) + args_raw = getattr(getattr(first, "function", None), "arguments", None) or "{}" + try: + args = json.loads(args_raw) + except json.JSONDecodeError: + args = {} + if isinstance(name, str) and isinstance(args, dict): + return {"name": name, "arguments": args} + # Text-tag fallback. + if not isinstance(content, str): + return None + m = _TOOL_CALL_TAG_RE.search(content) + if not m: + return None + try: + obj = json.loads(m.group(1)) + except json.JSONDecodeError: + return None + if not isinstance(obj, dict): + return None + name = obj.get("name") + args = obj.get("arguments", {}) + if not isinstance(name, str) or not isinstance(args, dict): + return None + return {"name": name, "arguments": args} + + +def _build_pool_block(workers: List[Dict[str, Any]]) -> str: + return "\n".join( + f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers + ) + + +def _build_user_prompt( + question: str, + workers: List[Dict[str, Any]], + history: List[Dict[str, Any]], +) -> str: + pieces = [ + f"Worker pool:\n{_build_pool_block(workers)}", + f"User question:\n{question}", + ] + if history: + pieces.append("Conversation so far (orchestrator turns and worker outputs):") + for h in history: + if h["role"] == "orchestrator": + pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}") + else: + pieces.append( + f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n" + f"{h['output']}" + ) + pieces.append( + "Emit the next JSON action object now — exactly one object, no prose." + ) + return "\n\n".join(pieces) + + +def _strip_fences(s: str) -> str: + s = s.strip() + if s.startswith("```"): + first_nl = s.find("\n") + if first_nl != -1: + s = s[first_nl + 1:] + if s.endswith("```"): + s = s[:-3] + s = s.strip() + return s + + +def _parse_action(text: str) -> Optional[Dict[str, Any]]: + s = _strip_fences(text) + # First try direct parse, then balanced-brace extraction. + try: + obj = json.loads(s) + if isinstance(obj, dict) and "action" in obj: + return obj + except json.JSONDecodeError: + pass + start = s.find("{") + if start == -1: + return None + depth = 0 + for i in range(start, len(s)): + c = s[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + try: + obj = json.loads(s[start : i + 1]) + if isinstance(obj, dict) and "action" in obj: + return obj + except json.JSONDecodeError: + return None + return None + + +def _extract_final_answer_text(text: str) -> str: + """Best-effort: pull the answer string from a malformed action emission. + + Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line. + """ + m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL) + if m: + return m.group(1).encode("utf-8").decode("unicode_escape") + m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE) + if m: + return m.group(1).strip() + return text.strip() diff --git a/src/openjarvis/agents/hybrid/toolorchestra/prompts.py b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py new file mode 100644 index 00000000..752e22c7 --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/prompts.py @@ -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": , "input": ""} + + {"action": "final_answer", "answer": ""} + +Strategy: + +- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math). +- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass. +- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase. +- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question. +""" + +FORCE_FINAL_PROMPT = ( + "Worker-call budget exhausted. Emit `final_answer` now using everything " + "you've learned. Respect the question's answer-format rules." +) + + +# ============================================================================ +# RL-mode constants (Orchestrator-8B, paper-faithful). +# ============================================================================ +# +# Verbatim copies of the upstream system prompt / user-prompt template / tools +# from `external/ToolOrchestra/evaluation/eval_hle.py` + `tools.json`. Don't +# edit the description text — Orchestrator-8B was RL-trained against this +# exact wording and pricing/latency table. + +RL_ORCHESTRATOR_SYS = "You are good at using tools." + +RL_TOOLS_SPEC: List[Dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "enhance_reasoning", + "description": "tool to enhance answer model reasoning. analyze the problem, write code, execute it and return intermidiate results that will help solve the problem", + "parameters": { + "properties": { + "model": { + "description": "The model used to reason. Choices: ['reasoner-1', 'reasoner-2', 'reasoner-3']. reasoner-1 demonstrates strong understanding and reasoning capabilities, which usually provides reliable insights. reasoner-2 can analyze some problems, but could hallucinate and make mistakes in difficult scenarios. reasoner-3 can reason over the context and reveal the logic. \nModel | price per million input tokens | price per million output tokens | average latency\nreasoner-1 | $1.25 | $10 | 31s\nreasoner-2 | $0.25 | $2 | 25s\nreasoner-3 | $0.8 | $0.8 | 9s", + "type": "string", + } + }, + "required": ["model"], + "title": "parameters", + "type": "object", + }, + }, + }, + { + "type": "function", + "function": { + "name": "answer", + "description": "give the final answer. Not allowed to call if documents is empty.", + "parameters": { + "properties": { + "model": { + "description": "The model used to answer. Choices: ['answer-1', 'answer-2', 'answer-3', 'answer-4', 'answer-math-1', 'answer-math-2']. answer-1 exhibits strong functional calling abilities and performs excellent in most domains (math, physics, social science, etc.). answer-2 presents reasonable solutions in some tasks, but could get stuck in complex reasoning and specific domain knowledge. answer-3 could solve easy to medium tasks, but is not capable of tackling tasks with strong expertise and long-horizon planning. answer-4 demonstrates basic capability: it can understand basic instructions, do simple steps, yet it sometimes misreads details, mixes concepts. answer-math-1 can solve moderate (middle school) math problem, though it becomes incapable in more difficult tasks. answer-math-2 can follow simple instructions and perform easy (primary-level) math problems, but struggle in more complex logic. The table below shows the pricing and latency of each model:\nModel | price per million input tokens | price per million output tokens | average latency\nanswer-1 | $1.25 | $10 | 96s\nanswer-2 | $0.25 | $2 | 27s\nanswer-3 | $0.9 | $0.9 | 15s\nanswer-4 | $0.8 | $0.8 | 11s\nanswer-math-1 | $0.9 | $0.9 | 13s\nanswer-math-2 | $$0.2 | $0.2 | 9s", + "type": "string", + } + }, + "required": ["model"], + "title": "parameters", + "type": "object", + }, + }, + }, + { + "type": "function", + "function": { + "name": "search", + "description": "Search for missing information", + "parameters": { + "properties": { + "model": { + "description": "The model used to search for missing information. Choices: ['search-1', 'search-2', 'search-3']. search-1 usually identifies the missing information and can write concise queries for effective search. search-2 can reason over the context and write queries to find the missing content for answering questions. search-3 can also write queries to find information. The table below shows the pricing and latency:\nModel | price per million input tokens | price per million output tokens | average latency\nsearch-1 | $1.25 | $10 | 22s\nsearch-2 | $0.25 | $2 | 16s\nsearch-3 | $0.8 | $0.8 | 8s", + "type": "string", + } + }, + "required": ["model"], + "title": "parameters", + "type": "object", + }, + }, + }, +] + +# RL_ALL_TOOLS: argument-validation schema (mirrors eval_hle.py:104). +RL_ALL_TOOLS: Dict[str, Dict[str, List[str]]] = { + "enhance_reasoning": {"model": ["reasoner-1", "reasoner-2", "reasoner-3"]}, + "answer": { + "model": [ + "answer-1", "answer-2", "answer-3", "answer-4", + "answer-math-1", "answer-math-2", + ], + }, + "search": {"model": ["search-1", "search-2", "search-3"]}, +} diff --git a/src/openjarvis/agents/hybrid/toolorchestra/rollout.py b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py new file mode 100644 index 00000000..a783f7dc --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/rollout.py @@ -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 ``{...}`` text the model emits.""" + return f"{json.dumps({'name': name, 'arguments': arguments})}" + + +__all__ = [ + "RL_ORCHESTRATOR_SYS", + "UnifiedRollout", + "UnifiedTurn", + "run_unified_rollout", + "tool_call_tag", +] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py b/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py new file mode 100644 index 00000000..30956d14 --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/sandbox.py @@ -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 diff --git a/src/openjarvis/agents/hybrid/toolorchestra/unified.py b/src/openjarvis/agents/hybrid/toolorchestra/unified.py new file mode 100644 index 00000000..4a861c0f --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/unified.py @@ -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"] diff --git a/src/openjarvis/agents/hybrid/toolorchestra/workers.py b/src/openjarvis/agents/hybrid/toolorchestra/workers.py new file mode 100644 index 00000000..a6e227c9 --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra/workers.py @@ -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 ``""``) resolves to + ``local_model``; ``model = "$cloud"`` / ``""`` to ``cloud_model``. + + On any validation failure, raises ``ValueError`` with the message + ``"Invalid worker_pool entry []: "``. Fails fast at agent + init rather than mid-task. + """ + override = cfg.get("worker_pool") + if override is None: + return _default_pool(local_model, local_endpoint, cloud_model, cloud_endpoint) + if not isinstance(override, list) or not override: + raise ValueError( + "Invalid worker_pool entry [-]: worker_pool must be a non-empty list" + ) + + resolved: List[Dict[str, Any]] = [] + seen_ids: set = set() + has_non_search = False + for raw in override: + wid_repr = raw.get("id", "?") if isinstance(raw, dict) else "?" + if not isinstance(raw, dict): + raise ValueError( + f"Invalid worker_pool entry [{wid_repr}]: entry must be a dict" + ) + entry = dict(raw) + wid = entry.get("id") + if not isinstance(wid, int): + raise ValueError( + f"Invalid worker_pool entry [{wid_repr}]: 'id' must be an int" + ) + if wid in seen_ids: + raise ValueError( + f"Invalid worker_pool entry [{wid}]: duplicate id" + ) + seen_ids.add(wid) + if not entry.get("name") or not isinstance(entry["name"], str): + raise ValueError( + f"Invalid worker_pool entry [{wid}]: 'name' must be a non-empty string" + ) + wtype = entry.get("type") or entry.get("endpoint") + if not isinstance(wtype, str) or wtype.lower() not in _TOOLORCH_VALID_TYPES: + raise ValueError( + f"Invalid worker_pool entry [{wid}]: 'type' must be one of " + f"{_TOOLORCH_VALID_TYPES} (got {wtype!r})" + ) + wtype = wtype.lower() + entry["type"] = wtype + # Substitute $local / $cloud placeholders (before any model check). + model = entry.get("model") + if isinstance(model, str) and model in ("$local", ""): + if not local_model: + raise ValueError( + f"Invalid worker_pool entry [{wid}]: model='{model}' " + "requires a local_model to be configured for this cell" + ) + model = local_model + entry["model"] = model + elif isinstance(model, str) and model in ("$cloud", ""): + model = cloud_model + entry["model"] = model + if wtype == "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"]), + ) diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py new file mode 100644 index 00000000..605efecf --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/reject_sample.py @@ -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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py new file mode 100644 index 00000000..6e1aa824 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/toolscale.py @@ -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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py new file mode 100644 index 00000000..6f9761c4 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/unified_serialize.py @@ -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 ``{...}`` 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 ```` 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"] diff --git a/tests/agents/test_expert_registry.py b/tests/agents/test_expert_registry.py new file mode 100644 index 00000000..c635e0ac --- /dev/null +++ b/tests/agents/test_expert_registry.py @@ -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" diff --git a/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py new file mode 100644 index 00000000..e0d35608 --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_rejection_pipeline.py @@ -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 tag (what the parser reads back). + assert any("" 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()