mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-08-14 00:47:52 +00:00
hybrid: six local+cloud paradigm agents (advisors, conductor, minions, archon, skillorchestra, toolorchestra) (#423)
This commit is contained in:
@@ -125,3 +125,8 @@ learning.db
|
||||
**/learning/benchmarks/
|
||||
**/teacher_traces/
|
||||
*.session.json
|
||||
|
||||
# Local dev artifacts (hybrid worker logs + cli debug dumps)
|
||||
minion_logs/
|
||||
*.oj-debug.json
|
||||
oj-debug.*.json
|
||||
|
||||
+4
-13
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "OpenJarvis"
|
||||
version = "1.0.2"
|
||||
version = "0.1.1"
|
||||
description = "OpenJarvis — modular AI assistant backend with composable intelligence primitives"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -81,19 +81,13 @@ server = [
|
||||
"python-multipart>=0.0.9",
|
||||
]
|
||||
openhands = ["openhands-sdk>=1.0; python_version >= '3.12'"]
|
||||
gpu-metrics = ["nvidia-ml-py>=12.560.30"]
|
||||
gpu-metrics = ["pynvml>=12.0"]
|
||||
energy-amd = ["amdsmi>=6.1"]
|
||||
energy-apple = ["zeus-ml[apple]"]
|
||||
energy-all = ["nvidia-ml-py>=12.560.30", "amdsmi>=6.1", "zeus-ml[apple]"]
|
||||
energy-all = ["pynvml>=12.0", "amdsmi>=6.1", "zeus-ml[apple]"]
|
||||
orchestrator-training = ["torch>=2.0", "transformers>=4.40"]
|
||||
learning-dspy = ["dspy>=2.6"]
|
||||
learning-gepa = ["gepa>=0.1"]
|
||||
# ACE (Agentic Context Engineering) is supported via
|
||||
# ``openjarvis.learning.agents.ace_optimizer`` but ACE upstream isn't on
|
||||
# PyPI and isn't structured as an installable Python package as of
|
||||
# v1.0.1, so there's no ``learning-ace`` extra. To use ACE, follow the
|
||||
# manual setup in docs/learning/ace.md (clone the upstream repo, add
|
||||
# its ``src/`` to PYTHONPATH).
|
||||
channel-telegram = ["python-telegram-bot>=21.0"]
|
||||
channel-discord = ["discord.py>=2.3"]
|
||||
channel-slack = ["slack-sdk>=3.27"]
|
||||
@@ -176,6 +170,7 @@ markers = [
|
||||
"nvidia: requires NVIDIA GPU",
|
||||
"slow: long-running test",
|
||||
"live_external: requires HERMES_AGENT_PATH and OPENCLAW_PATH; spawns real foreign-framework subprocesses",
|
||||
"modal: requires Modal token + network; runs real swebench harness on Modal",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
@@ -191,10 +186,6 @@ select = ["E", "F", "I", "W"]
|
||||
# hybrid/ is research code with long prompt strings and paradigm-specific
|
||||
# config dicts — same relaxation as evals research code above.
|
||||
"src/openjarvis/agents/hybrid/*.py" = ["E501"]
|
||||
# research_loop.py carries the multi-paragraph planner system prompt as
|
||||
# inline string literals; line-length wrapping would harm readability of
|
||||
# the prompt itself.
|
||||
"src/openjarvis/agents/research_loop.py" = ["E501"]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
|
||||
@@ -84,11 +84,6 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import openjarvis.agents.proactive_agent # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Hybrid local+cloud paradigm agents (Minions, Conductor, Archon, Advisors,
|
||||
# SkillOrchestra, ToolOrchestra). Each module registers under its own name
|
||||
# via @AgentRegistry.register(). Optional deps may make some unavailable.
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
# Hybrid local+cloud paradigm agents
|
||||
|
||||
Six paradigms ported from
|
||||
[`/matx/u/aspark/hybrid-local-cloud-compute`](../../../../../) — each is
|
||||
registered as a standard OpenJarvis agent so the rest of the platform
|
||||
(SDK, CLI, distillation, evals) can use them like any other agent.
|
||||
Six paradigms ported from the original ``hybrid-local-cloud-compute``
|
||||
harness — each is registered as a standard OpenJarvis agent so the rest
|
||||
of the platform (SDK, CLI, distillation, evals) can use them like any
|
||||
other agent. Results live under ``$OPENJARVIS_HYBRID_EXPERIMENTS_DIR``
|
||||
(defaults to ``~/.openjarvis/experiments/hybrid/``).
|
||||
|
||||
| Agent | Plan shape | Trains what? | Workers |
|
||||
|-------------------|-----------------|----------------------|---------------------------|
|
||||
@@ -45,15 +46,15 @@ structural scorer).
|
||||
## Quickstart
|
||||
|
||||
```bash
|
||||
cd /matx/u/aspark/OpenJarvis
|
||||
cd OpenJarvis
|
||||
source .env # API keys
|
||||
|
||||
# 1. Start vLLM in another shell (see CLAUDE.md for the full recipe)
|
||||
# 1. Start vLLM in another shell (see your local launch recipe)
|
||||
# CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m vllm.entrypoints.openai.api_server \
|
||||
# --model Qwen/Qwen3.5-27B-FP8 --port 8001 ...
|
||||
|
||||
# 2. (Optional) for Minions: install the upstream library
|
||||
.venv/bin/uv pip install -e /matx/u/aspark/hybrid-local-cloud-compute/external/minions
|
||||
.venv/bin/uv pip install -e path/to/minions
|
||||
|
||||
# 3. Run a smoke cell
|
||||
.venv/bin/python -m openjarvis.agents.hybrid.runner \
|
||||
@@ -61,7 +62,7 @@ source .env # API keys
|
||||
```
|
||||
|
||||
Outputs land in
|
||||
`$OPENJARVIS_HYBRID_EXPERIMENTS_DIR/<cell>/{results.jsonl,summary.json,config.json,logs/}`
|
||||
`$OPENJARVIS_HYBRID_EXPERIMENTS_DIR/runs/<cell>/{results.jsonl,summary.json,config.json,logs/}`
|
||||
(defaults to `~/.openjarvis-hybrid/experiments/`). The schema matches the
|
||||
hybrid harness so the existing rescore / dashboard scripts work
|
||||
unmodified.
|
||||
@@ -80,7 +81,7 @@ That appends a `[cells.<name>]` block to
|
||||
## How good is each paradigm?
|
||||
|
||||
Numbers from the upstream hybrid harness
|
||||
(`/matx/u/aspark/hybrid-local-cloud-compute/docs/results.md`) at full N —
|
||||
(`~/.openjarvis/experiments/hybrid/docs/results.md`) at full N —
|
||||
GAIA val n=165, SWE-bench-Verified n=500. Local = Qwen-3.5-27B-FP8, cloud
|
||||
= Opus 4.7. Cloud-only baseline: GAIA 0.570 / $1.09, SWE 0.238 / $0.95.
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ Each module here registers one agent under ``@AgentRegistry.register("<name>")``
|
||||
conductor — zero-shot planner emits a DAG of up to 5 worker calls
|
||||
minions — supervisor (cloud) ↔ worker (local) reactive loop
|
||||
archon — layered (generator → ranker → fuser) inference-time search
|
||||
skillorchestra — skill-aware router picks one agent from a pool
|
||||
skillorchestra — eval orchestrator: skill-routed search→reasoning→answer loop
|
||||
toolorchestra — prompted multi-turn dispatcher over a mixed tool/model pool
|
||||
|
||||
All agents share :class:`LocalCloudAgent` as the base. They are bench-agnostic:
|
||||
@@ -15,9 +15,9 @@ or the bench's native formatter) and hands it in via ``run(input=...)``. Task
|
||||
metadata that the paradigm needs (a problem statement vs. a question, hints,
|
||||
etc.) goes through ``context.metadata``.
|
||||
|
||||
The hybrid harness at ``/matx/u/aspark/hybrid-local-cloud-compute`` is the
|
||||
reference implementation and stays untouched — these ports are the
|
||||
OpenJarvis-native versions of the same paradigms.
|
||||
The original ``hybrid-local-cloud-compute`` harness is the reference
|
||||
implementation and stays untouched — these ports are the OpenJarvis-native
|
||||
versions of the same paradigms.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -35,6 +35,8 @@ for _modname in (
|
||||
"skillorchestra",
|
||||
"toolorchestra",
|
||||
"mini_swe_agent",
|
||||
"baseline_cloud",
|
||||
"baseline_local",
|
||||
):
|
||||
try:
|
||||
__import__(f"openjarvis.agents.hybrid.{_modname}")
|
||||
|
||||
@@ -37,13 +37,16 @@ kwargs (``local_model``, ``local_endpoint``, ``cloud_endpoint``, …) follow.
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from abc import abstractmethod
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any, Deque, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent
|
||||
from openjarvis.agents.hybrid._openai_retry import patch_openai_globally
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
NO_TEMP_PREFIXES,
|
||||
is_gpt5_family,
|
||||
@@ -54,9 +57,29 @@ from openjarvis.agents.hybrid._prices import (
|
||||
)
|
||||
from openjarvis.engine._stubs import InferenceEngine
|
||||
|
||||
# Install OpenAI SDK retry + per-org concurrency cap at import time so
|
||||
# every paradigm (advisors, conductor, minions, mini_swe_agent's cloud
|
||||
# loop, archon, …) inherits the hardening without each call site having
|
||||
# to remember to opt in. See ``_openai_retry.py`` for the env knobs
|
||||
# (default: 4 concurrent, 8 retries, 2s/60s exponential backoff w/ jitter).
|
||||
patch_openai_globally()
|
||||
|
||||
# Anthropic server-side web_search: $10 per 1000 searches.
|
||||
WEB_SEARCH_COST_PER_CALL = 0.01
|
||||
|
||||
# OpenAI Responses-API hosted web_search tool: $10 per 1000 calls
|
||||
# (2025-12 public list price for the `web_search` / `web_search_preview`
|
||||
# tool, billed per tool call). Same shape as Anthropic, so we reuse the
|
||||
# $0.01/call number — kept as a separate constant so it can drift.
|
||||
OPENAI_WEB_SEARCH_COST_PER_CALL = 0.01
|
||||
|
||||
# Gemini Google-Search grounding: billed at $35 per 1000 grounded
|
||||
# *requests* (2025-12 public list price for the Grounding-with-Google-Search
|
||||
# tool, charged once per request that uses the tool regardless of how many
|
||||
# internal queries it issues). We charge per grounded request, not per
|
||||
# `web_search_queries` entry.
|
||||
GEMINI_SEARCH_COST_PER_CALL = 0.035
|
||||
|
||||
ANTHROPIC_WEB_SEARCH_TOOL = {
|
||||
"type": "web_search_20250305",
|
||||
"name": "web_search",
|
||||
@@ -64,6 +87,35 @@ ANTHROPIC_WEB_SEARCH_TOOL = {
|
||||
}
|
||||
|
||||
|
||||
def build_web_search_tool(max_uses: int = 8) -> Dict[str, Any]:
|
||||
"""Build an Anthropic server-side web_search tool block with a custom cap.
|
||||
|
||||
Web_search is server-side: Anthropic runs the searches internally before
|
||||
returning, so ``max_uses`` is the only knob the caller has to bound cost
|
||||
per task. Defaults to 8 (matches ``ANTHROPIC_WEB_SEARCH_TOOL``).
|
||||
"""
|
||||
return {
|
||||
"type": "web_search_20250305",
|
||||
"name": "web_search",
|
||||
"max_uses": int(max_uses),
|
||||
}
|
||||
|
||||
|
||||
def web_search_cfg(method_cfg: Optional[Dict[str, Any]]) -> Tuple[bool, int]:
|
||||
"""Parse ``method_cfg.web_search = { enabled, max_uses }``.
|
||||
|
||||
Defaults: enabled=False, max_uses=8. ``enabled`` defaults to False so
|
||||
existing cells stay one-shot (backwards compat). Cells opting in flip
|
||||
``enabled=true`` in the registry. Returns ``(enabled, max_uses)``.
|
||||
"""
|
||||
if not method_cfg:
|
||||
return False, 8
|
||||
ws = method_cfg.get("web_search")
|
||||
if not isinstance(ws, dict):
|
||||
return False, 8
|
||||
return bool(ws.get("enabled", False)), int(ws.get("max_uses", 8))
|
||||
|
||||
|
||||
# ---------- Thread-local trace buffer ----------
|
||||
#
|
||||
# Every call through ``_call_anthropic`` / ``_call_openai`` / ``_call_vllm``
|
||||
@@ -96,6 +148,126 @@ def _close_trace() -> None:
|
||||
delattr(_TRACE_STATE, "events")
|
||||
|
||||
|
||||
# ---------- Thread-local LLM-call counter ----------
|
||||
#
|
||||
# Parallel to the trace buffer: every cloud SDK call (anthropic / openai /
|
||||
# gemini, including each turn of ``_call_anthropic_agent`` and each turn of
|
||||
# the mini-SWE multi-turn loops) bumps ``cloud``; every local vLLM call
|
||||
# bumps ``local``. ``run()`` opens a fresh pair per task, drops the totals
|
||||
# into ``meta["n_cloud_calls"]`` / ``meta["n_local_calls"]``, then closes.
|
||||
#
|
||||
# Why thread-local and not an instance counter: agents are shared across
|
||||
# the runner's ``ThreadPoolExecutor`` (one agent, N concurrent tasks). An
|
||||
# attribute on ``self`` would race; this matches the trace buffer's
|
||||
# pattern exactly so it composes with the rest of the per-task plumbing.
|
||||
|
||||
_CALL_COUNTS = threading.local()
|
||||
|
||||
|
||||
def _call_counts() -> Optional[Dict[str, int]]:
|
||||
return getattr(_CALL_COUNTS, "counts", None)
|
||||
|
||||
|
||||
def _bump_cloud_calls(n: int = 1) -> None:
|
||||
counts = _call_counts()
|
||||
if counts is not None:
|
||||
counts["cloud"] += int(n)
|
||||
|
||||
|
||||
def _bump_local_calls(n: int = 1) -> None:
|
||||
counts = _call_counts()
|
||||
if counts is not None:
|
||||
counts["local"] += int(n)
|
||||
|
||||
|
||||
def _open_call_counts() -> Dict[str, int]:
|
||||
counts: Dict[str, int] = {"cloud": 0, "local": 0}
|
||||
_CALL_COUNTS.counts = counts
|
||||
return counts
|
||||
|
||||
|
||||
def _close_call_counts() -> None:
|
||||
if hasattr(_CALL_COUNTS, "counts"):
|
||||
delattr(_CALL_COUNTS, "counts")
|
||||
|
||||
|
||||
# ---------- OpenRouter in-process rate limiter ----------
|
||||
#
|
||||
# OpenRouter enforces per-account RPM and concurrency limits. A single agent
|
||||
# process running a wide ThreadPoolExecutor (Conductor's 7-worker pool fanned
|
||||
# across the runner's per-task threads) can trivially blow past those caps,
|
||||
# triggering 429s that the OpenAI SDK retries with backoff — wasted latency
|
||||
# and noisy traces.
|
||||
#
|
||||
# We gate every ``_call_openrouter`` call through two limits, **shared across
|
||||
# threads in this Python process** (singleton; lazy-initialized):
|
||||
#
|
||||
# - **Concurrency** — a ``threading.Semaphore``. Default 20 in-flight calls;
|
||||
# override via ``OJ_OPENROUTER_MAX_CONCURRENT``. Held only around the SDK
|
||||
# ``client.chat.completions.create`` invocation, not the bookkeeping.
|
||||
# - **RPM** — a sliding window of timestamps in a ``deque``. Default 60
|
||||
# requests/minute; override via ``OJ_OPENROUTER_RPM``. If the deque has
|
||||
# already accumulated ``RPM`` timestamps within the last 60 seconds, the
|
||||
# caller sleeps until the oldest entry ages out, then proceeds. Every
|
||||
# completed call appends ``time.time()`` so the call we just made is
|
||||
# counted against the next window.
|
||||
#
|
||||
# Other cloud helpers (``_call_openai`` / ``_call_anthropic`` / ``_call_gemini``)
|
||||
# are NOT rate-limited here — OpenAI and Anthropic have their own concurrency
|
||||
# patch (``_openai_retry``) and Gemini's free-tier RPM is loose enough that
|
||||
# we haven't hit it yet. Add limiters for those one-by-one if needed.
|
||||
|
||||
_OPENROUTER_LIMITER_LOCK = threading.Lock()
|
||||
_OPENROUTER_LIMITER: Optional["_OpenRouterLimiter"] = None
|
||||
|
||||
|
||||
class _OpenRouterLimiter:
|
||||
"""Process-wide concurrency + sliding-window RPM gate for OpenRouter."""
|
||||
|
||||
def __init__(self, max_concurrent: int, rpm: int) -> None:
|
||||
self.max_concurrent = int(max_concurrent)
|
||||
self.rpm = int(rpm)
|
||||
self._sem = threading.Semaphore(self.max_concurrent)
|
||||
self._window: Deque[float] = deque()
|
||||
self._window_lock = threading.Lock()
|
||||
|
||||
def acquire_concurrency(self) -> None:
|
||||
self._sem.acquire()
|
||||
|
||||
def release_concurrency(self) -> None:
|
||||
self._sem.release()
|
||||
|
||||
def wait_for_rpm_slot(self) -> None:
|
||||
"""Block until making one more call would not exceed RPM in 60s."""
|
||||
while True:
|
||||
with self._window_lock:
|
||||
now = time.time()
|
||||
cutoff = now - 60.0
|
||||
while self._window and self._window[0] < cutoff:
|
||||
self._window.popleft()
|
||||
if len(self._window) < self.rpm:
|
||||
return
|
||||
# Sleep until the oldest in-window call ages out, then recheck.
|
||||
sleep_s = 60.0 - (now - self._window[0]) + 0.01
|
||||
if sleep_s > 0:
|
||||
time.sleep(sleep_s)
|
||||
|
||||
def record_call(self) -> None:
|
||||
with self._window_lock:
|
||||
self._window.append(time.time())
|
||||
|
||||
|
||||
def _openrouter_limiter() -> _OpenRouterLimiter:
|
||||
global _OPENROUTER_LIMITER
|
||||
if _OPENROUTER_LIMITER is None:
|
||||
with _OPENROUTER_LIMITER_LOCK:
|
||||
if _OPENROUTER_LIMITER is None:
|
||||
max_concurrent = int(os.environ.get("OJ_OPENROUTER_MAX_CONCURRENT", "20") or 20)
|
||||
rpm = int(os.environ.get("OJ_OPENROUTER_RPM", "60") or 60)
|
||||
_OPENROUTER_LIMITER = _OpenRouterLimiter(max_concurrent, rpm)
|
||||
return _OPENROUTER_LIMITER
|
||||
|
||||
|
||||
def _serialize_block(block: Any) -> Dict[str, Any]:
|
||||
"""Turn an Anthropic content block (text / tool_use / server_tool_use /
|
||||
web_search_tool_result / thinking) into a JSON-safe dict.
|
||||
@@ -223,13 +395,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
tool_choice: Optional[dict] = None,
|
||||
output_config: Optional[dict] = None,
|
||||
timeout: float = 600.0,
|
||||
max_retries: int = 5,
|
||||
max_retries: int = 12,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int, int]:
|
||||
"""Single Anthropic call. Returns (text, p_tok, c_tok, n_web_searches).
|
||||
|
||||
Strips ``temperature`` for Opus 4.7+ (rejected by the API). Captures
|
||||
the call into the active per-task trace if one is open.
|
||||
the call into the active per-task trace if one is open. Bumped
|
||||
default max_retries to 12 (~2 min of backoff) so cells survive
|
||||
sustained Anthropic 529 "Overloaded" windows when many cells share
|
||||
Opus quota.
|
||||
"""
|
||||
import anthropic
|
||||
|
||||
@@ -251,6 +426,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
kwargs["output_config"] = output_config
|
||||
t0 = time.time()
|
||||
msg = client.messages.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
text = "".join(b.text for b in msg.content if hasattr(b, "text"))
|
||||
srv = getattr(msg.usage, "server_tool_use", None)
|
||||
@@ -321,6 +497,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
kwargs["tool_choice"] = tool_choice
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
@@ -352,6 +529,181 @@ class LocalCloudAgent(BaseAgent):
|
||||
})
|
||||
return text, p, c
|
||||
|
||||
@staticmethod
|
||||
def _call_openrouter(
|
||||
model: str,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "cloud",
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, int, int]:
|
||||
"""Single OpenRouter call. Returns (text, p_tok, c_tok).
|
||||
|
||||
OpenRouter is OpenAI-API-compatible; we use the OpenAI SDK with a
|
||||
custom ``base_url`` and ``OPENROUTER_API_KEY``. ``model`` is the
|
||||
OpenRouter slug ``"<provider>/<model>"`` (e.g.
|
||||
``"deepseek/deepseek-r1"``). For convenience the caller may also
|
||||
pass the OpenJarvis-engine-style ``"openrouter/<provider>/<model>"``
|
||||
prefix (see ``src/openjarvis/engine/cloud.py``) — we strip it here.
|
||||
|
||||
Note: unlike ``_call_openai``, we do NOT apply the GPT-5 family
|
||||
``max_completion_tokens`` rewrite or temperature stripping —
|
||||
OpenRouter models have varied parameter support. We pass through
|
||||
whatever the caller supplies; if a specific model errors on
|
||||
``temperature``, that's the cell's problem to handle.
|
||||
|
||||
``extra_body`` is forwarded to the OpenAI SDK as the ``extra_body``
|
||||
kwarg so callers can pass OpenRouter-specific fields (e.g.
|
||||
``{"reasoning": {"effort": "medium"}}`` to enable thinking on Qwen3,
|
||||
or ``{"provider": {...}}`` to pin routing). The SDK serializes
|
||||
``extra_body`` into the JSON request body alongside the standard
|
||||
fields.
|
||||
|
||||
Every call goes through the process-wide
|
||||
``_OpenRouterLimiter`` (concurrency semaphore + sliding-window RPM
|
||||
deque) so a wide ThreadPoolExecutor can't blow past account caps.
|
||||
|
||||
Trace events use ``"kind": "openrouter"`` so the dashboard can
|
||||
distinguish them from native OpenAI calls.
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
if model.startswith("openrouter/"):
|
||||
model = model[len("openrouter/"):]
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
raise RuntimeError(
|
||||
"OPENROUTER_API_KEY is not set; cannot call OpenRouter."
|
||||
)
|
||||
client = OpenAI(
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_key=api_key,
|
||||
timeout=timeout,
|
||||
)
|
||||
messages: list = []
|
||||
if system:
|
||||
messages.append({"role": "system", "content": system})
|
||||
messages.append({"role": "user", "content": user})
|
||||
kwargs: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if extra_body:
|
||||
kwargs["extra_body"] = extra_body
|
||||
limiter = _openrouter_limiter()
|
||||
limiter.wait_for_rpm_slot()
|
||||
limiter.acquire_concurrency()
|
||||
try:
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
finally:
|
||||
limiter.release_concurrency()
|
||||
limiter.record_call()
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
text = message.content or ""
|
||||
tool_calls = _serialize_openai_tool_calls(getattr(message, "tool_calls", None))
|
||||
reasoning = getattr(message, "reasoning_content", None) or getattr(
|
||||
message, "reasoning", None
|
||||
)
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
_record_event({
|
||||
"kind": "openrouter",
|
||||
"role": trace_role,
|
||||
"model": model,
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tool_calls": tool_calls,
|
||||
"reasoning_content": reasoning,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"finish_reason": getattr(choice, "finish_reason", None),
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
return text, p, c
|
||||
|
||||
@staticmethod
|
||||
def _call_gemini(
|
||||
model: str,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int]:
|
||||
"""Single Gemini Developer-API call. Returns (text, p_tok, c_tok).
|
||||
|
||||
Tool-call parity with Anthropic is intentionally NOT implemented —
|
||||
skillorchestra / baseline-cloud only need text generation, and the
|
||||
google-genai tool-config plumbing diverges enough from the
|
||||
Anthropic/OpenAI shape that wiring it would double the surface
|
||||
area of this file. If a future paradigm needs Gemini tool use,
|
||||
extend this helper rather than hand-rolling it in the agent.
|
||||
|
||||
Captures the call into the active per-task trace via the
|
||||
``"gemini"`` kind so the dashboard's trace renderer picks it up.
|
||||
"""
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
client = genai.Client(http_options=types.HttpOptions(timeout=int(timeout * 1000)))
|
||||
cfg = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_tokens,
|
||||
)
|
||||
if system:
|
||||
cfg.system_instruction = system
|
||||
t0 = time.time()
|
||||
resp = client.models.generate_content(
|
||||
model=model,
|
||||
contents=user,
|
||||
config=cfg,
|
||||
)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
# `resp.text` is the convenience accessor that concatenates every
|
||||
# text part in the first candidate. Empty if the model emitted
|
||||
# only non-text parts (which we don't request).
|
||||
text = (resp.text or "") if hasattr(resp, "text") else ""
|
||||
um = getattr(resp, "usage_metadata", None)
|
||||
p = int(getattr(um, "prompt_token_count", 0) or 0) if um else 0
|
||||
c = int(getattr(um, "candidates_token_count", 0) or 0) if um else 0
|
||||
finish_reason = None
|
||||
try:
|
||||
finish_reason = str(resp.candidates[0].finish_reason)
|
||||
except Exception:
|
||||
pass
|
||||
_record_event({
|
||||
"kind": "gemini",
|
||||
"role": trace_role,
|
||||
"model": model,
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"finish_reason": finish_reason,
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
return text, p, c
|
||||
|
||||
@staticmethod
|
||||
def _call_vllm(
|
||||
model: str,
|
||||
@@ -392,6 +744,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
kwargs["tool_choice"] = tool_choice
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
_bump_local_calls()
|
||||
latency = time.time() - t0
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
@@ -426,6 +779,334 @@ class LocalCloudAgent(BaseAgent):
|
||||
})
|
||||
return text, p, c
|
||||
|
||||
@staticmethod
|
||||
def _call_anthropic_agent(
|
||||
model: str,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
tools: Optional[list] = None,
|
||||
max_turns: int = 8,
|
||||
timeout: float = 600.0,
|
||||
max_retries: int = 5,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int, int, int]:
|
||||
"""Multi-turn Anthropic loop with optional tools.
|
||||
|
||||
Returns ``(final_text, prompt_tokens_sum, completion_tokens_sum,
|
||||
n_web_searches_sum, turns)``. The loop appends the assistant
|
||||
response (preserving server-tool blocks so Anthropic's continuation
|
||||
is valid) and re-prompts until the model stops with ``end_turn``
|
||||
or we hit ``max_turns``. Server-side tools (web_search) execute
|
||||
inside a single message and don't require a tool_result echo —
|
||||
the model just continues thinking with the search results
|
||||
already in its context. Client-side tools aren't executed here;
|
||||
if the model emits a client-side ``tool_use`` block we stop
|
||||
(paradigms that want client tools should call ``_call_anthropic``
|
||||
directly and handle their own dispatch).
|
||||
"""
|
||||
import anthropic
|
||||
|
||||
client = anthropic.Anthropic(timeout=timeout, max_retries=max_retries)
|
||||
# Build the conversation. We grow ``messages`` across turns so the
|
||||
# model sees its own prior assistant content (text + server tool
|
||||
# use + web_search_tool_result). For server tools Anthropic is
|
||||
# happy as long as we feed the raw assistant content back; no
|
||||
# synthetic user tool_result block is needed.
|
||||
messages: List[Dict[str, Any]] = [{"role": "user", "content": user}]
|
||||
p_total = 0
|
||||
c_total = 0
|
||||
n_searches_total = 0
|
||||
last_text = ""
|
||||
turns = 0
|
||||
for turn in range(max(1, max_turns)):
|
||||
turns = turn + 1
|
||||
kwargs: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"max_tokens": max_tokens,
|
||||
"messages": messages,
|
||||
}
|
||||
if system:
|
||||
kwargs["system"] = system
|
||||
if supports_temperature(model):
|
||||
kwargs["temperature"] = temperature
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
t0 = time.time()
|
||||
msg = client.messages.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
text = "".join(b.text for b in msg.content if hasattr(b, "text"))
|
||||
srv = getattr(msg.usage, "server_tool_use", None)
|
||||
n_searches = getattr(srv, "web_search_requests", 0) if srv else 0
|
||||
content_blocks = [_serialize_block(b) for b in msg.content]
|
||||
tool_use_blocks = [
|
||||
b for b in content_blocks
|
||||
if b.get("type") in ("tool_use", "server_tool_use")
|
||||
]
|
||||
tool_result_blocks = [
|
||||
b for b in content_blocks
|
||||
if b.get("type") in ("web_search_tool_result", "tool_result")
|
||||
]
|
||||
stop_reason = getattr(msg, "stop_reason", None)
|
||||
_record_event({
|
||||
"kind": "anthropic",
|
||||
"role": trace_role,
|
||||
"model": model,
|
||||
"system": system if turn == 0 else None,
|
||||
"user": user if turn == 0 else None,
|
||||
"turn": turn,
|
||||
"response": text,
|
||||
"content_blocks": content_blocks,
|
||||
"tool_calls": tool_use_blocks,
|
||||
"tool_results": tool_result_blocks,
|
||||
"tokens_in": msg.usage.input_tokens,
|
||||
"tokens_out": msg.usage.output_tokens,
|
||||
"n_web_searches": n_searches,
|
||||
"tools_declared": tools,
|
||||
"stop_reason": stop_reason,
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
p_total += msg.usage.input_tokens
|
||||
c_total += msg.usage.output_tokens
|
||||
n_searches_total += n_searches
|
||||
if text:
|
||||
last_text = text
|
||||
# If the model wants a client-side tool we don't dispatch
|
||||
# here — break and let the caller (or future loop variant)
|
||||
# handle it. Only ``server_tool_use`` blocks (web_search)
|
||||
# are auto-continued by Anthropic itself.
|
||||
client_tool_use = any(
|
||||
b.get("type") == "tool_use" for b in content_blocks
|
||||
)
|
||||
if client_tool_use:
|
||||
break
|
||||
if stop_reason == "end_turn" or stop_reason is None:
|
||||
break
|
||||
# Otherwise: ``stop_reason`` like "max_tokens" or "tool_use"
|
||||
# (server side) — Anthropic returned mid-thought. Append the
|
||||
# assistant turn and ask it to continue.
|
||||
messages.append({"role": "assistant", "content": msg.content})
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": "Continue.",
|
||||
})
|
||||
return last_text, p_total, c_total, n_searches_total, turns
|
||||
|
||||
@staticmethod
|
||||
def _call_openai_agent(
|
||||
model: str,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
tools: Optional[list] = None,
|
||||
max_turns: int = 8,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int, int, int]:
|
||||
"""OpenAI hosted-web-search call via the Responses API.
|
||||
|
||||
Returns ``(final_text, prompt_tokens_sum, completion_tokens_sum,
|
||||
n_web_searches_sum, turns)`` — the same 5-tuple shape as
|
||||
``_call_anthropic_agent`` so callers can dispatch uniformly.
|
||||
|
||||
OpenAI's hosted web search is exposed only through the **Responses
|
||||
API** (``client.responses.create``), not chat completions. The tool
|
||||
is declared as ``{"type": "web_search"}``; older SDK/model combos
|
||||
only accept the legacy ``"web_search_preview"`` name, so we retry
|
||||
once with that on a tool-type rejection. The Responses API runs the
|
||||
search server-side and returns the post-search answer in one call,
|
||||
so ``turns`` is 1 — the ``max_turns`` arg is accepted only for
|
||||
signature parity with ``_call_anthropic_agent``.
|
||||
|
||||
``n_web_searches`` counts ``web_search_call`` items in the response
|
||||
output. Token usage is summed from ``response.usage``
|
||||
(``input_tokens`` / ``output_tokens``).
|
||||
"""
|
||||
from openai import OpenAI
|
||||
|
||||
del max_turns # Responses API resolves search server-side in one call.
|
||||
client = OpenAI(timeout=timeout)
|
||||
|
||||
search_tool_names = ["web_search", "web_search_preview"]
|
||||
kwargs_base: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"input": user,
|
||||
"max_output_tokens": max_tokens,
|
||||
}
|
||||
if system:
|
||||
kwargs_base["instructions"] = system
|
||||
# GPT-5 family ignores `temperature` on the Responses API (reasoning
|
||||
# models reject it); only pass it for non-gpt-5 models.
|
||||
if not is_gpt5_family(model):
|
||||
kwargs_base["temperature"] = temperature
|
||||
|
||||
resp = None
|
||||
last_exc: Optional[BaseException] = None
|
||||
used_tool_name = search_tool_names[0]
|
||||
t0 = time.time()
|
||||
for tool_name in search_tool_names:
|
||||
try:
|
||||
resp = client.responses.create(
|
||||
**kwargs_base,
|
||||
tools=[{"type": tool_name}],
|
||||
)
|
||||
used_tool_name = tool_name
|
||||
break
|
||||
except Exception as exc: # noqa: BLE001
|
||||
# Only fall through to the legacy name on what looks like a
|
||||
# tool-type rejection; re-raise anything else immediately.
|
||||
last_exc = exc
|
||||
msg = str(exc).lower()
|
||||
if "web_search" in msg or "tool" in msg or "unsupported" in msg:
|
||||
continue
|
||||
raise
|
||||
if resp is None:
|
||||
raise last_exc if last_exc is not None else RuntimeError(
|
||||
"openai responses.create failed for all web_search tool names"
|
||||
)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
|
||||
# Extract output text. Prefer the SDK convenience accessor; fall back
|
||||
# to walking the output items for `output_text` content parts.
|
||||
text = ""
|
||||
try:
|
||||
text = resp.output_text or ""
|
||||
except Exception: # noqa: BLE001
|
||||
text = ""
|
||||
output_items = list(getattr(resp, "output", None) or [])
|
||||
if not text:
|
||||
chunks: List[str] = []
|
||||
for item in output_items:
|
||||
if getattr(item, "type", None) != "message":
|
||||
continue
|
||||
for part in getattr(item, "content", None) or []:
|
||||
if getattr(part, "type", None) in ("output_text", "text"):
|
||||
chunks.append(getattr(part, "text", "") or "")
|
||||
text = "".join(chunks)
|
||||
|
||||
n_searches = sum(
|
||||
1 for item in output_items
|
||||
if getattr(item, "type", None) in (
|
||||
"web_search_call", "web_search_tool_call",
|
||||
)
|
||||
)
|
||||
u = getattr(resp, "usage", None)
|
||||
p = int(getattr(u, "input_tokens", 0) or 0) if u else 0
|
||||
c = int(getattr(u, "output_tokens", 0) or 0) if u else 0
|
||||
_record_event({
|
||||
"kind": "openai_agent",
|
||||
"role": trace_role,
|
||||
"model": model,
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"output_items": _jsonable(output_items),
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"n_web_searches": n_searches,
|
||||
"tools_declared": [{"type": used_tool_name}],
|
||||
"stop_reason": getattr(resp, "status", None),
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
return text, p, c, n_searches, 1
|
||||
|
||||
@staticmethod
|
||||
def _call_gemini_agent(
|
||||
model: str,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
tools: Optional[list] = None,
|
||||
max_turns: int = 8,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int, int, int]:
|
||||
"""Gemini call grounded with Google Search.
|
||||
|
||||
Returns ``(final_text, prompt_tokens_sum, completion_tokens_sum,
|
||||
n_web_searches_sum, turns)`` — same shape as ``_call_anthropic_agent``.
|
||||
|
||||
Grounding is wired by adding ``Tool(google_search=GoogleSearch())``
|
||||
to the ``GenerateContentConfig``. Gemini resolves the grounding
|
||||
server-side inside a single ``generate_content`` call, so
|
||||
``turns`` is always 1 and ``max_turns`` / ``tools`` are accepted
|
||||
only for signature parity with ``_call_anthropic_agent``.
|
||||
|
||||
``n_web_searches`` is derived from ``candidates[0].grounding_metadata``
|
||||
— we count the ``web_search_queries`` Gemini reports, falling back to
|
||||
0 when no grounding metadata is present (the model answered without
|
||||
searching).
|
||||
"""
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
del tools, max_turns # grounding is config-level + single-call.
|
||||
client = genai.Client(
|
||||
http_options=types.HttpOptions(timeout=int(timeout * 1000))
|
||||
)
|
||||
cfg = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_tokens,
|
||||
tools=[types.Tool(google_search=types.GoogleSearch())],
|
||||
)
|
||||
if system:
|
||||
cfg.system_instruction = system
|
||||
t0 = time.time()
|
||||
resp = client.models.generate_content(
|
||||
model=model,
|
||||
contents=user,
|
||||
config=cfg,
|
||||
)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
text = (resp.text or "") if hasattr(resp, "text") else ""
|
||||
um = getattr(resp, "usage_metadata", None)
|
||||
p = int(getattr(um, "prompt_token_count", 0) or 0) if um else 0
|
||||
c = int(getattr(um, "candidates_token_count", 0) or 0) if um else 0
|
||||
|
||||
# Derive search count from grounding metadata when present.
|
||||
n_searches = 0
|
||||
web_search_queries: List[str] = []
|
||||
finish_reason = None
|
||||
try:
|
||||
cand0 = resp.candidates[0]
|
||||
finish_reason = str(getattr(cand0, "finish_reason", None))
|
||||
gm = getattr(cand0, "grounding_metadata", None)
|
||||
if gm is not None:
|
||||
queries = getattr(gm, "web_search_queries", None) or []
|
||||
web_search_queries = [str(q) for q in queries]
|
||||
n_searches = len(web_search_queries)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
_record_event({
|
||||
"kind": "gemini_agent",
|
||||
"role": trace_role,
|
||||
"model": model,
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"n_web_searches": n_searches,
|
||||
"web_search_queries": web_search_queries,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"finish_reason": finish_reason,
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
return text, p, c, n_searches, 1
|
||||
|
||||
def _call_cloud(
|
||||
self,
|
||||
*,
|
||||
@@ -460,6 +1141,23 @@ class LocalCloudAgent(BaseAgent):
|
||||
temperature=temperature,
|
||||
**kwargs,
|
||||
)
|
||||
if self._cloud_endpoint == "gemini":
|
||||
# Gemini helper doesn't accept tools / response_format kwargs.
|
||||
# Drop them rather than letting them surface as a TypeError so
|
||||
# paradigms that opportunistically pass these for OpenAI /
|
||||
# Anthropic still work when routed to Gemini.
|
||||
kwargs.pop("tools", None)
|
||||
kwargs.pop("tool_choice", None)
|
||||
kwargs.pop("response_format", None)
|
||||
kwargs.pop("output_config", None)
|
||||
return self._call_gemini(
|
||||
self._cloud_model,
|
||||
user=user,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
**kwargs,
|
||||
)
|
||||
raise ValueError(f"unsupported cloud endpoint: {self._cloud_endpoint!r}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
@@ -478,6 +1176,8 @@ class LocalCloudAgent(BaseAgent):
|
||||
"tokens_cloud": 0,
|
||||
"cost_usd": 0.0,
|
||||
"latency_s": 0.0,
|
||||
"n_cloud_calls": 0,
|
||||
"n_local_calls": 0,
|
||||
"soft_error": reason,
|
||||
"traces": {"soft_error": reason},
|
||||
}
|
||||
@@ -499,6 +1199,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
self._emit_turn_start(input)
|
||||
t0 = time.time()
|
||||
events = _open_trace()
|
||||
counts = _open_call_counts()
|
||||
meta: Dict[str, Any]
|
||||
answer: str = ""
|
||||
soft_reason: Optional[str] = None
|
||||
@@ -514,6 +1215,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
soft_reason = soft
|
||||
meta = self._soft_fail_metadata(soft)
|
||||
finally:
|
||||
# Snapshot call counts BEFORE closing the thread-local state.
|
||||
# Subclasses don't track this themselves — every SDK call site
|
||||
# bumps the thread-local in this file, so the totals are
|
||||
# authoritative as of right now. Overwrites any value
|
||||
# ``_run_paradigm`` happened to set (it shouldn't set them).
|
||||
n_cloud = counts.get("cloud", 0)
|
||||
n_local = counts.get("local", 0)
|
||||
if "meta" in locals():
|
||||
meta["n_cloud_calls"] = int(n_cloud)
|
||||
meta["n_local_calls"] = int(n_local)
|
||||
# Persist the trace before the trace state is closed (and even on
|
||||
# hard failure, so we get a record of what we did before it broke).
|
||||
self._write_trace_log(
|
||||
@@ -521,6 +1232,7 @@ class LocalCloudAgent(BaseAgent):
|
||||
events, soft_reason, exc_obj,
|
||||
)
|
||||
_close_trace()
|
||||
_close_call_counts()
|
||||
meta.setdefault("latency_s", time.time() - t0)
|
||||
if soft_reason is not None:
|
||||
self._emit_turn_end(soft_error=soft_reason)
|
||||
@@ -606,10 +1318,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
|
||||
__all__ = [
|
||||
"ANTHROPIC_WEB_SEARCH_TOOL",
|
||||
"GEMINI_SEARCH_COST_PER_CALL",
|
||||
"LocalCloudAgent",
|
||||
"NO_TEMP_PREFIXES",
|
||||
"OPENAI_WEB_SEARCH_COST_PER_CALL",
|
||||
"WEB_SEARCH_COST_PER_CALL",
|
||||
"_bump_cloud_calls",
|
||||
"_bump_local_calls",
|
||||
"build_web_search_tool",
|
||||
"estimate_cost",
|
||||
"is_gpt5_family",
|
||||
"supports_temperature",
|
||||
"web_search_cfg",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
"""GPU energy collector for hybrid cell runs.
|
||||
|
||||
Samples ``pynvml.nvmlDeviceGetPowerUsage`` at ~2Hz on a background thread,
|
||||
integrates power × dt → joules across the cell's wall-time. Used as a
|
||||
context manager wrapping ``_run_cell_locked`` so the whole cell's GPU
|
||||
energy is attributed to the cell (not per-task — concurrency makes
|
||||
per-task attribution very noisy).
|
||||
|
||||
Failure modes are absorbed: if NVML isn't available (no GPU on this host,
|
||||
container without ``libnvidia-ml.so.1``, permission denied) the collector
|
||||
logs *once* and returns ``0.0``. It must never crash the run.
|
||||
|
||||
Scope notes
|
||||
-----------
|
||||
* Samples **all visible NVIDIA GPUs on the host where the runner runs**.
|
||||
In our setup that's the L40S node that also hosts vLLM
|
||||
(``mkt1`` / ``matx2``), so this captures the local-model serving GPUs.
|
||||
If the runner is invoked from a host without GPUs (e.g. a login node
|
||||
with vLLM on a remote box), the collector logs and yields 0 — set
|
||||
``OPENJARVIS_HYBRID_ENERGY=0`` to silence the warning.
|
||||
* ``CUDA_VISIBLE_DEVICES`` is *honored* for parity with vLLM: only those
|
||||
GPU indices are sampled. Unset = all GPUs on the host.
|
||||
* **Cloud energy is not measured** — cloud calls go over HTTPS to
|
||||
Anthropic/OpenAI/Google, no measurable joules on our side. A future
|
||||
pass could add a per-token J/token estimate (e.g. Patterson et al.
|
||||
2021, Luccioni et al. 2022) but those numbers are vendor-opaque and
|
||||
uncertain — leaving as a TODO until we explicitly decide to estimate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
_NVML_WARNED = False
|
||||
|
||||
|
||||
def _log_once(msg: str) -> None:
|
||||
global _NVML_WARNED
|
||||
if not _NVML_WARNED:
|
||||
print(f"[energy] {msg}", flush=True)
|
||||
_NVML_WARNED = True
|
||||
|
||||
|
||||
def _resolve_gpu_indices(total: int) -> List[int]:
|
||||
"""Honor CUDA_VISIBLE_DEVICES; default to every visible GPU."""
|
||||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if not cvd:
|
||||
return list(range(total))
|
||||
out: List[int] = []
|
||||
for s in cvd.split(","):
|
||||
s = s.strip()
|
||||
if not s:
|
||||
continue
|
||||
try:
|
||||
idx = int(s)
|
||||
except ValueError:
|
||||
continue
|
||||
if 0 <= idx < total:
|
||||
out.append(idx)
|
||||
return out or list(range(total))
|
||||
|
||||
|
||||
class EnergyCollector:
|
||||
"""Background NVML power sampler integrating to joules.
|
||||
|
||||
Use as a context manager::
|
||||
|
||||
with EnergyCollector() as ec:
|
||||
... # run the cell
|
||||
joules = ec.energy_j_total # also available after __exit__
|
||||
|
||||
Safe to instantiate even when NVML is unavailable: ``energy_j_total``
|
||||
will simply be ``0.0`` and a one-time warning will be printed.
|
||||
"""
|
||||
|
||||
def __init__(self, sample_hz: float = 2.0) -> None:
|
||||
self.sample_hz = float(sample_hz)
|
||||
self.energy_j_total: float = 0.0
|
||||
self.samples: int = 0
|
||||
self.gpu_indices: List[int] = []
|
||||
self._stop = threading.Event()
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._t0: float = 0.0
|
||||
self._t1: float = 0.0
|
||||
self._enabled = os.environ.get("OPENJARVIS_HYBRID_ENERGY", "1") != "0"
|
||||
self._pynvml = None
|
||||
self._handles: list = []
|
||||
|
||||
# ---- context manager
|
||||
def __enter__(self) -> "EnergyCollector":
|
||||
self._t0 = time.time()
|
||||
if not self._enabled:
|
||||
return self
|
||||
try:
|
||||
import pynvml # type: ignore[import-not-found]
|
||||
pynvml.nvmlInit()
|
||||
total = pynvml.nvmlDeviceGetCount()
|
||||
self.gpu_indices = _resolve_gpu_indices(total)
|
||||
self._handles = [
|
||||
pynvml.nvmlDeviceGetHandleByIndex(i) for i in self.gpu_indices
|
||||
]
|
||||
# Probe once so we fail fast if power query is unsupported.
|
||||
for h in self._handles:
|
||||
pynvml.nvmlDeviceGetPowerUsage(h)
|
||||
self._pynvml = pynvml
|
||||
except Exception as e: # noqa: BLE001 — NVML failures must never crash the run
|
||||
_log_once(
|
||||
f"NVML unavailable ({type(e).__name__}: {e}); "
|
||||
"energy_j_total will be 0. Set OPENJARVIS_HYBRID_ENERGY=0 to silence."
|
||||
)
|
||||
self._pynvml = None
|
||||
return self
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=self._sample_loop, name="energy-sampler", daemon=True
|
||||
)
|
||||
self._thread.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None:
|
||||
self._t1 = time.time()
|
||||
self._stop.set()
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=5.0)
|
||||
if self._pynvml is not None:
|
||||
try:
|
||||
self._pynvml.nvmlShutdown()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
# ---- sampler thread
|
||||
def _sample_loop(self) -> None:
|
||||
"""Trapezoid-integrate Σ_gpu power_W × dt over the cell run.
|
||||
|
||||
``nvmlDeviceGetPowerUsage`` returns milliwatts. Sample interval
|
||||
is best-effort (~``1/sample_hz`` s); if the host is overloaded
|
||||
and a tick lands late we still use the real ``dt`` so the
|
||||
integral stays honest. Trapezoid rule (mean of consecutive
|
||||
readings) damps jitter vs. left-Riemann.
|
||||
"""
|
||||
pynvml = self._pynvml
|
||||
assert pynvml is not None
|
||||
period = 1.0 / max(self.sample_hz, 0.1)
|
||||
last_t = time.time()
|
||||
last_total_w: Optional[float] = None
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
total_mw = 0
|
||||
for h in self._handles:
|
||||
total_mw += pynvml.nvmlDeviceGetPowerUsage(h)
|
||||
total_w = total_mw / 1000.0
|
||||
except Exception: # noqa: BLE001 — keep going even on transient NVML errors
|
||||
self._stop.wait(period)
|
||||
continue
|
||||
now = time.time()
|
||||
if last_total_w is not None:
|
||||
dt = now - last_t
|
||||
# Trapezoid rule: mean power across the interval × dt.
|
||||
self.energy_j_total += 0.5 * (total_w + last_total_w) * dt
|
||||
last_t = now
|
||||
last_total_w = total_w
|
||||
self.samples += 1
|
||||
self._stop.wait(period)
|
||||
|
||||
# ---- inspection
|
||||
@property
|
||||
def wall_s(self) -> float:
|
||||
return max(self._t1 - self._t0, 0.0) if self._t1 else (time.time() - self._t0)
|
||||
|
||||
|
||||
__all__ = ["EnergyCollector"]
|
||||
@@ -0,0 +1,355 @@
|
||||
"""Process-wide retry + concurrency hardening for cloud OpenAI calls.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
|
||||
When we run the hybrid paradigms (Minions, Advisors, Conductor, …) at
|
||||
n=100 against ``gpt-5`` / ``gpt-5-mini`` over the prepaid OpenAI quota,
|
||||
sustained concurrency walls the org-level rate limit and the OpenAI SDK
|
||||
raises :class:`openai.RateLimitError`. The SDK's own retry path is short
|
||||
(default ``max_retries=2`` on a small backoff) — under sustained pressure
|
||||
every wave of retries hits the same wall and the runner records 19-69
|
||||
errored rows per cell, degenerating the result.
|
||||
|
||||
Mirrors the existing ``_patch_anthropic_globally`` pattern in
|
||||
``minions.py``: monkey-patch the SDK at module level so it applies even
|
||||
to libraries that build their own ``openai.OpenAI()`` clients
|
||||
(HazyResearch Minions's ``OpenAIClient``, Archon's adapters,
|
||||
``mini_swe_agent``, etc.). One call to :func:`patch_openai_globally` from
|
||||
either ``_base.py`` or ``minions._apply_patches_once`` is enough — the
|
||||
patch is idempotent and process-wide.
|
||||
|
||||
What it does
|
||||
------------
|
||||
|
||||
1. Bumps ``openai.OpenAI()`` constructor defaults to
|
||||
``timeout=600.0`` / ``max_retries=8`` (the SDK's own backoff is fine
|
||||
for transient blips; we layer our own loop on top for sustained walls).
|
||||
2. Wraps ``chat.completions.create`` with:
|
||||
|
||||
- A **per-org semaphore** (``OPENJARVIS_OPENAI_MAX_CONCURRENCY``,
|
||||
default 4) that throttles sustained concurrency. Single bursts are
|
||||
fine — prepaid quotas wall on sustained rate, not on a brief spike.
|
||||
The semaphore is **only acquired for cloud calls** (``api.openai.com``);
|
||||
vLLM calls routed through the OpenAI SDK (``base_url`` set to a
|
||||
local endpoint or ``api_key="EMPTY"``) bypass it.
|
||||
- Exponential backoff with jitter on :class:`openai.RateLimitError`,
|
||||
:class:`openai.APITimeoutError`, :class:`openai.APIConnectionError`,
|
||||
and 5xx :class:`openai.APIStatusError`. Starts at 2 s, caps at 60 s,
|
||||
up to 8 attempts (~3.5 min worst case). Honors a ``Retry-After``
|
||||
header when the SDK surfaces one.
|
||||
- On exhaustion, re-raises the last exception (it bubbles up to the
|
||||
runner, which records ``error="RateLimitError: ..."`` in
|
||||
``results.jsonl`` — no silent drop).
|
||||
|
||||
Env knobs
|
||||
---------
|
||||
|
||||
- ``OPENJARVIS_OPENAI_MAX_CONCURRENCY`` (default ``4``) — semaphore
|
||||
capacity. Set to e.g. ``2`` if the wall is still hit; set to ``0`` to
|
||||
disable throttling entirely (passes through to the SDK).
|
||||
- ``OPENJARVIS_OPENAI_MAX_RETRIES`` (default ``8``) — outer retry loop
|
||||
cap (separate from the SDK's own ``max_retries``).
|
||||
- ``OPENJARVIS_OPENAI_RETRY_BASE`` (default ``2.0``) — base seconds for
|
||||
exponential backoff. Schedule is ``min(60, base * 2**attempt) * jitter``.
|
||||
- ``OPENJARVIS_OPENAI_RETRY_CAP`` (default ``60.0``) — max single-step
|
||||
sleep in seconds.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tunables (read once at module load, can be overridden via env)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
v = int(os.environ.get(name, "") or default)
|
||||
return max(0, v)
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return float(os.environ.get(name, "") or default)
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
_MAX_CONCURRENCY = _env_int("OPENJARVIS_OPENAI_MAX_CONCURRENCY", 4)
|
||||
_MAX_RETRIES = _env_int("OPENJARVIS_OPENAI_MAX_RETRIES", 8)
|
||||
_RETRY_BASE = _env_float("OPENJARVIS_OPENAI_RETRY_BASE", 2.0)
|
||||
_RETRY_CAP = _env_float("OPENJARVIS_OPENAI_RETRY_CAP", 60.0)
|
||||
|
||||
|
||||
# Single process-wide semaphore. ``BoundedSemaphore(0)`` would block
|
||||
# forever, so when the env knob is 0 we hand back a no-op context manager.
|
||||
class _NullSem:
|
||||
def __enter__(self) -> "_NullSem":
|
||||
return self
|
||||
|
||||
def __exit__(self, *a: Any) -> None:
|
||||
return None
|
||||
|
||||
|
||||
_SEM: Any
|
||||
if _MAX_CONCURRENCY > 0:
|
||||
_SEM = threading.BoundedSemaphore(_MAX_CONCURRENCY)
|
||||
else:
|
||||
_SEM = _NullSem()
|
||||
|
||||
|
||||
_PATCHED = False
|
||||
_PATCH_LOCK = threading.Lock()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _is_local_endpoint(client: Any) -> bool:
|
||||
"""True if this OpenAI client points at a local vLLM endpoint.
|
||||
|
||||
We detect via either ``api_key == "EMPTY"`` (the convention used by
|
||||
our ``_call_vllm`` and ``mini_swe_agent``) or a ``base_url`` whose
|
||||
hostname resolves to localhost. Either signal is enough; both are
|
||||
cheap to read.
|
||||
"""
|
||||
try:
|
||||
api_key = getattr(client, "api_key", None)
|
||||
if api_key == "EMPTY":
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
base_url = str(getattr(client, "base_url", "") or "")
|
||||
if not base_url:
|
||||
return False
|
||||
host = urlparse(base_url).hostname or ""
|
||||
return host in ("localhost", "127.0.0.1", "0.0.0.0", "::1")
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _extract_retry_after(exc: BaseException) -> Optional[float]:
|
||||
"""Pull a Retry-After header off an APIStatusError if the SDK exposed it.
|
||||
|
||||
OpenAI's SDK keeps the underlying ``httpx.Response`` on
|
||||
``exc.response`` for ``APIStatusError`` subclasses. Header may be a
|
||||
seconds-integer or an HTTP-date; we only handle the integer form
|
||||
(the only thing OpenAI sends in practice).
|
||||
"""
|
||||
resp = getattr(exc, "response", None)
|
||||
if resp is None:
|
||||
return None
|
||||
headers = getattr(resp, "headers", None)
|
||||
if not headers:
|
||||
return None
|
||||
for name in ("retry-after", "Retry-After", "x-ratelimit-reset-requests"):
|
||||
val = headers.get(name) if hasattr(headers, "get") else None
|
||||
if val is None:
|
||||
continue
|
||||
try:
|
||||
secs = float(val)
|
||||
if 0 <= secs <= 600:
|
||||
return secs
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _is_retryable(exc: BaseException) -> bool:
|
||||
"""Whether to retry this OpenAI exception class."""
|
||||
try:
|
||||
import openai
|
||||
except ImportError:
|
||||
return False
|
||||
if isinstance(exc, (
|
||||
openai.RateLimitError,
|
||||
openai.APITimeoutError,
|
||||
openai.APIConnectionError,
|
||||
openai.InternalServerError,
|
||||
)):
|
||||
return True
|
||||
if isinstance(exc, openai.APIStatusError):
|
||||
status = getattr(exc, "status_code", None)
|
||||
# 429 is RateLimitError already; 5xx is retryable; 408 is a
|
||||
# timeout the SDK didn't classify (rare).
|
||||
return status is not None and (status >= 500 or status in (408, 409, 429))
|
||||
return False
|
||||
|
||||
|
||||
def _sleep_for(attempt: int, exc: BaseException) -> float:
|
||||
"""Backoff for attempt index ``attempt`` (0-based)."""
|
||||
hinted = _extract_retry_after(exc)
|
||||
if hinted is not None and hinted > 0:
|
||||
# Respect a server-provided hint, but clamp to our cap so a
|
||||
# pathological header can't stall the run for hours.
|
||||
return min(_RETRY_CAP, hinted) + random.uniform(0, 0.5)
|
||||
base = min(_RETRY_CAP, _RETRY_BASE * (2 ** attempt))
|
||||
# Full jitter — better tail behavior than equal jitter when many
|
||||
# workers wake at the same moment.
|
||||
return random.uniform(0.0, base)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wrapping
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _wrap_create(orig: Callable[..., Any]) -> Callable[..., Any]:
|
||||
"""Wrap a ``chat.completions.create`` (or ``responses.create``) bound
|
||||
method's underlying function with retry + concurrency throttling.
|
||||
|
||||
The wrapper is a regular function that takes ``self`` as the first
|
||||
arg, so it can replace ``Completions.create`` at the class level and
|
||||
still see the bound client through ``self._client``.
|
||||
"""
|
||||
|
||||
def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
client = getattr(self, "_client", None)
|
||||
local = client is not None and _is_local_endpoint(client)
|
||||
# Local vLLM calls bypass the per-org throttle (no rate limit) and
|
||||
# the long retry loop (vLLM is mostly either up or down — a 60s
|
||||
# backoff just delays surfacing the failure). But brief
|
||||
# ConnectionError blips do happen mid-sweep (socket queue, brief
|
||||
# server warmup pause): give the local path a SHORT retry — 3
|
||||
# attempts, 1s/2s/4s — so we don't error an entire row on a
|
||||
# transient refused connection. Anything else (BadRequest etc.)
|
||||
# still raises immediately.
|
||||
if local:
|
||||
local_last_exc: Optional[BaseException] = None
|
||||
for attempt in range(3):
|
||||
try:
|
||||
return orig(self, *args, **kwargs)
|
||||
except BaseException as exc: # noqa: BLE001
|
||||
try:
|
||||
import openai
|
||||
except ImportError:
|
||||
raise
|
||||
if not isinstance(exc, (
|
||||
openai.APIConnectionError,
|
||||
openai.APITimeoutError,
|
||||
openai.InternalServerError,
|
||||
)):
|
||||
raise
|
||||
local_last_exc = exc
|
||||
if attempt >= 2:
|
||||
break
|
||||
time.sleep(2 ** attempt)
|
||||
assert local_last_exc is not None
|
||||
raise local_last_exc
|
||||
|
||||
last_exc: Optional[BaseException] = None
|
||||
for attempt in range(_MAX_RETRIES + 1):
|
||||
try:
|
||||
with _SEM:
|
||||
return orig(self, *args, **kwargs)
|
||||
except BaseException as exc: # noqa: BLE001
|
||||
if not _is_retryable(exc):
|
||||
raise
|
||||
last_exc = exc
|
||||
if attempt >= _MAX_RETRIES:
|
||||
break
|
||||
delay = _sleep_for(attempt, exc)
|
||||
# Stderr, not stdout: heartbeat / progress lines must
|
||||
# stay parseable in the runner log.
|
||||
try:
|
||||
import sys
|
||||
print(
|
||||
f"[openai-retry] attempt {attempt + 1}/{_MAX_RETRIES} "
|
||||
f"{type(exc).__name__}: {str(exc)[:120]} — "
|
||||
f"sleeping {delay:.1f}s",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(delay)
|
||||
# Exhausted. Re-raise so the runner records the row as errored.
|
||||
assert last_exc is not None
|
||||
raise last_exc
|
||||
|
||||
wrapped._hybrid_patched = True # type: ignore[attr-defined]
|
||||
wrapped.__wrapped__ = orig # type: ignore[attr-defined]
|
||||
return wrapped
|
||||
|
||||
|
||||
def patch_openai_globally() -> None:
|
||||
"""Idempotently monkey-patch the OpenAI SDK to add retry + throttling.
|
||||
|
||||
Safe to call from multiple modules — guarded by ``_PATCHED`` under a
|
||||
lock. Patches both the ``Completions.create`` method and the
|
||||
``OpenAI.__init__`` defaults.
|
||||
"""
|
||||
global _PATCHED
|
||||
if _PATCHED:
|
||||
return
|
||||
with _PATCH_LOCK:
|
||||
if _PATCHED:
|
||||
return
|
||||
try:
|
||||
import openai
|
||||
from openai.resources.chat import completions as _comp_mod
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
# Bump constructor defaults so callers that don't pass timeout /
|
||||
# max_retries explicitly still get sensible values. ``setdefault``
|
||||
# so any explicit caller value wins.
|
||||
if not getattr(openai.OpenAI.__init__, "_hybrid_patched", False):
|
||||
_orig_init = openai.OpenAI.__init__
|
||||
|
||||
def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
|
||||
kwargs.setdefault("timeout", 600.0)
|
||||
kwargs.setdefault("max_retries", _MAX_RETRIES)
|
||||
return _orig_init(self, *args, **kwargs)
|
||||
|
||||
_patched_init._hybrid_patched = True # type: ignore[attr-defined]
|
||||
openai.OpenAI.__init__ = _patched_init # type: ignore[assignment]
|
||||
|
||||
# Wrap chat.completions.create. The SDK exposes the bound method
|
||||
# via ``Completions.create``; we replace the class attribute so
|
||||
# every instance (including ones built inside external libs)
|
||||
# sees the wrapped version.
|
||||
if not getattr(_comp_mod.Completions.create, "_hybrid_patched", False):
|
||||
_comp_mod.Completions.create = _wrap_create( # type: ignore[assignment]
|
||||
_comp_mod.Completions.create
|
||||
)
|
||||
|
||||
# Also patch the async variant for completeness (none of our
|
||||
# paradigms use it today, but Archon / future paradigms might).
|
||||
try:
|
||||
from openai.resources.chat import completions as _comp_mod_async
|
||||
|
||||
cls = getattr(_comp_mod_async, "AsyncCompletions", None)
|
||||
if cls is not None and not getattr(
|
||||
cls.create, "_hybrid_patched", False
|
||||
):
|
||||
# Async wrapper is structurally different — only patch
|
||||
# the bumped defaults via __init__; full retry loop on
|
||||
# async would need an async wrapper. Leave that for the
|
||||
# day a paradigm actually uses it.
|
||||
pass
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
_PATCHED = True
|
||||
|
||||
|
||||
def current_settings() -> Tuple[int, int, float, float]:
|
||||
"""For tests / smoke runs: return (concurrency, retries, base, cap)."""
|
||||
return _MAX_CONCURRENCY, _MAX_RETRIES, _RETRY_BASE, _RETRY_CAP
|
||||
|
||||
|
||||
__all__ = ["patch_openai_globally", "current_settings"]
|
||||
@@ -19,7 +19,17 @@ PRICES: dict[str, tuple[float, float]] = {
|
||||
"gpt-5-mini": (0.25, 2.00),
|
||||
"gpt-5-mini-2025-08-07": (0.25, 2.00),
|
||||
"gpt-4o": (0.15, 0.60),
|
||||
# Gemini Developer API prices (USD per 1M tokens), 2025-12 list price.
|
||||
# 2.5 Pro uses tiered pricing (>200K context = $2.50/$15); we charge the
|
||||
# low-context tier since GAIA / SWE-bench prompts stay well under 200K.
|
||||
"gemini-2.5-pro": (1.25, 10.0),
|
||||
"gemini-2.5-flash": (0.30, 2.50),
|
||||
"gemini-2.5-flash-lite": (0.10, 0.40),
|
||||
# OpenRouter slugs (used by toolorchestra paper-match pool).
|
||||
# Prices are OpenRouter list (USD/1M tokens), 2026-05 snapshot.
|
||||
"qwen/qwen-2.5-coder-32b-instruct": (0.08, 0.18),
|
||||
"qwen/qwen3-32b": (0.10, 0.30),
|
||||
"meta-llama/llama-3.3-70b-instruct": (0.13, 0.39),
|
||||
}
|
||||
|
||||
# Models whose API rejects an explicit `temperature` param — callers should
|
||||
@@ -44,3 +54,17 @@ def supports_temperature(model: str) -> bool:
|
||||
def is_gpt5_family(model: str) -> bool:
|
||||
"""GPT-5 series requires ``max_completion_tokens`` and forced temp=1."""
|
||||
return model.startswith("gpt-5")
|
||||
|
||||
|
||||
def is_reasoning_model(model: str) -> bool:
|
||||
"""Models that consume the output-token budget on hidden chain-of-thought
|
||||
before emitting visible answer text. At max_tokens=4096 these silently
|
||||
truncate with empty answers on GAIA (26/100 GPT-5, 18/100 Gemini Pro)."""
|
||||
m = (model or "").lower()
|
||||
return is_gpt5_family(model) or "gemini-2.5-pro" in m
|
||||
|
||||
|
||||
def default_max_output_tokens(model: str) -> int:
|
||||
"""Sane default for ``max_tokens`` per cloud call. Reasoning models get
|
||||
a larger budget so their hidden thinking doesn't crowd out the answer."""
|
||||
return 16384 if is_reasoning_model(model) else 4096
|
||||
|
||||
@@ -28,7 +28,14 @@ import urllib.request
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.hybrid._base import LocalCloudAgent
|
||||
from openjarvis.agents.hybrid._base import (
|
||||
GEMINI_SEARCH_COST_PER_CALL,
|
||||
OPENAI_WEB_SEARCH_COST_PER_CALL,
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
LocalCloudAgent,
|
||||
build_web_search_tool,
|
||||
web_search_cfg,
|
||||
)
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import (
|
||||
run_swe_agent_loop,
|
||||
)
|
||||
@@ -84,6 +91,21 @@ def _resolve_local_model(endpoint: str, registry_model: str) -> str:
|
||||
return served[0] if served else registry_model
|
||||
|
||||
|
||||
# Cloud endpoints that have a server-side web-search agent loop wired in
|
||||
# `_base.py`. Anything else (openrouter, vllm, unknown) cannot ground.
|
||||
_SEARCH_CAPABLE_ENDPOINTS = ("anthropic", "openai", "gemini")
|
||||
|
||||
|
||||
def _search_cost_per_call(endpoint: str) -> float:
|
||||
"""Per-search-call USD cost for the configured cloud endpoint."""
|
||||
if endpoint == "openai":
|
||||
return OPENAI_WEB_SEARCH_COST_PER_CALL
|
||||
if endpoint == "gemini":
|
||||
return GEMINI_SEARCH_COST_PER_CALL
|
||||
# anthropic (and any caller that already validated the endpoint).
|
||||
return WEB_SEARCH_COST_PER_CALL
|
||||
|
||||
|
||||
@AgentRegistry.register("advisors")
|
||||
class AdvisorsAgent(LocalCloudAgent):
|
||||
"""Three-step executor ↔ advisor ↔ executor loop. See module docstring."""
|
||||
@@ -112,13 +134,39 @@ class AdvisorsAgent(LocalCloudAgent):
|
||||
advisor_max_tokens = int(cfg.get("advisor_max_tokens", 1024))
|
||||
advisor_temperature = float(cfg.get("advisor_temperature", 0.2))
|
||||
|
||||
# 1. Initial executor pass
|
||||
initial_resp, e1_in, e1_out = self._call_cloud(
|
||||
user=f"Question:\n{question}",
|
||||
system=EXECUTOR_INITIAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
temperature=0.0,
|
||||
)
|
||||
ws_enabled, ws_max_uses = web_search_cfg(cfg)
|
||||
if ws_enabled and self._cloud_endpoint not in _SEARCH_CAPABLE_ENDPOINTS:
|
||||
raise ValueError(
|
||||
f"web_search.enabled=true but cloud_endpoint={self._cloud_endpoint!r}; "
|
||||
"server-side web_search is wired for anthropic / openai / gemini "
|
||||
"executors only. Route this cell through one of those or disable "
|
||||
"web_search — otherwise search would silently no-op and the "
|
||||
"executor answers blind."
|
||||
)
|
||||
use_ws = ws_enabled
|
||||
gaia_max_turns = int(cfg.get("gaia_max_turns", 8))
|
||||
n_searches_total = 0
|
||||
|
||||
# 1. Initial executor pass — advisor (Qwen) doesn't get tools;
|
||||
# only the cloud executor passes do. With web_search on, dispatch
|
||||
# to the search-capable agent loop for the configured provider.
|
||||
if use_ws:
|
||||
initial_resp, e1_in, e1_out, n_s1, e1_turns = self._executor_search(
|
||||
user=f"Question:\n{question}",
|
||||
system=EXECUTOR_INITIAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
ws_max_uses=ws_max_uses,
|
||||
max_turns=gaia_max_turns,
|
||||
)
|
||||
n_searches_total += n_s1
|
||||
else:
|
||||
initial_resp, e1_in, e1_out = self._call_cloud(
|
||||
user=f"Question:\n{question}",
|
||||
system=EXECUTOR_INITIAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
temperature=0.0,
|
||||
)
|
||||
e1_turns = 1
|
||||
|
||||
# 2. Advisor pass (local)
|
||||
if not self._local_endpoint or not self._local_model:
|
||||
@@ -147,30 +195,105 @@ class AdvisorsAgent(LocalCloudAgent):
|
||||
f"Produce your best final answer now, respecting the question's "
|
||||
f"answer-format rules."
|
||||
)
|
||||
final_answer, e2_in, e2_out = self._call_cloud(
|
||||
user=final_user,
|
||||
system=EXECUTOR_FINAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
temperature=0.0,
|
||||
)
|
||||
if use_ws:
|
||||
final_answer, e2_in, e2_out, n_s2, e2_turns = self._executor_search(
|
||||
user=final_user,
|
||||
system=EXECUTOR_FINAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
ws_max_uses=ws_max_uses,
|
||||
max_turns=gaia_max_turns,
|
||||
)
|
||||
n_searches_total += n_s2
|
||||
else:
|
||||
final_answer, e2_in, e2_out = self._call_cloud(
|
||||
user=final_user,
|
||||
system=EXECUTOR_FINAL_SYS,
|
||||
max_tokens=executor_max_tokens,
|
||||
temperature=0.0,
|
||||
)
|
||||
e2_turns = 1
|
||||
|
||||
tokens_local = adv_in + adv_out
|
||||
tokens_cloud = e1_in + e1_out + e2_in + e2_out
|
||||
cost = self.cost_usd(self._cloud_model, e1_in + e2_in, e1_out + e2_out)
|
||||
cost += n_searches_total * _search_cost_per_call(self._cloud_endpoint)
|
||||
|
||||
meta: Dict[str, Any] = {
|
||||
"tokens_local": tokens_local,
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost,
|
||||
"turns": 3,
|
||||
# executor pass 1 + advisor pass (1) + executor pass 2. With
|
||||
# web_search on, each executor pass is a multi-turn loop, so
|
||||
# this is > 3; one-shot (no search) it's exactly 3.
|
||||
"turns": e1_turns + 1 + e2_turns,
|
||||
"web_search_uses": n_searches_total,
|
||||
# GAIA: only the executor passes invoke a tool (web_search).
|
||||
"tool_calls": int(n_searches_total),
|
||||
"traces": {
|
||||
"initial_response": initial_resp,
|
||||
"advisor_feedback": advisor_text,
|
||||
"web_search_enabled": use_ws,
|
||||
"n_web_searches": n_searches_total,
|
||||
"note": "inference-only advisor (untrained); lower bound on the technique.",
|
||||
},
|
||||
}
|
||||
return final_answer, meta
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Web-search executor dispatch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _executor_search(
|
||||
self,
|
||||
*,
|
||||
user: str,
|
||||
system: str,
|
||||
max_tokens: int,
|
||||
ws_max_uses: int,
|
||||
max_turns: int,
|
||||
) -> Tuple[str, int, int, int, int]:
|
||||
"""Run a search-capable executor pass for the configured cloud.
|
||||
|
||||
Dispatches by ``self._cloud_endpoint`` to the matching ``_base``
|
||||
agent loop. Returns the shared 5-tuple ``(text, p_tok, c_tok,
|
||||
n_searches, turns)``. The endpoint is assumed already validated
|
||||
against ``_SEARCH_CAPABLE_ENDPOINTS`` by the caller.
|
||||
"""
|
||||
if self._cloud_endpoint == "anthropic":
|
||||
return self._call_anthropic_agent(
|
||||
self._cloud_model,
|
||||
user=user,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.0,
|
||||
tools=[build_web_search_tool(ws_max_uses)],
|
||||
max_turns=max_turns,
|
||||
)
|
||||
if self._cloud_endpoint == "openai":
|
||||
return self._call_openai_agent(
|
||||
self._cloud_model,
|
||||
user=user,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.0,
|
||||
max_turns=max_turns,
|
||||
)
|
||||
if self._cloud_endpoint == "gemini":
|
||||
return self._call_gemini_agent(
|
||||
self._cloud_model,
|
||||
user=user,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.0,
|
||||
max_turns=max_turns,
|
||||
)
|
||||
# Genuinely unsupported (openrouter / vllm / unknown). The caller
|
||||
# guard should have caught this; raise defensively.
|
||||
raise ValueError(
|
||||
f"web_search executor pass requested but cloud_endpoint="
|
||||
f"{self._cloud_endpoint!r} has no search wiring."
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# SWE-bench variant: each executor pass is a full mini-SWE-agent run.
|
||||
# ------------------------------------------------------------------
|
||||
@@ -265,6 +388,9 @@ class AdvisorsAgent(LocalCloudAgent):
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost,
|
||||
"turns": initial_out["turns"] + 1 + final_out["turns"],
|
||||
# SWE: sum bash turns from both executor passes; advisor pass
|
||||
# is a local-model critique with no tools.
|
||||
"tool_calls": int(initial_out["turns"] + final_out["turns"]),
|
||||
"traces": {
|
||||
"swe_mode": True,
|
||||
"initial_summary": initial_out["final_summary"],
|
||||
|
||||
@@ -26,9 +26,9 @@ How the hybrid harness wires it (and what we mirror here):
|
||||
- ``ranker_model`` / ``fuser_model`` (default: ``cloud_model`` for both)
|
||||
- ``max_tokens`` (default 2048), ``temperature`` (default 0.7)
|
||||
|
||||
Requires the Archon library (cloned at
|
||||
``hybrid-local-cloud-compute/external/Archon`` — add its ``src`` to
|
||||
``PYTHONPATH`` or pip-install editable). Import is lazy.
|
||||
Requires the Archon library from https://github.com/Stanford-ILIAD/Archon
|
||||
— either pip-install editable or set ``ARCHON_SRC`` to the checkout's
|
||||
``src/`` directory. Import is lazy.
|
||||
|
||||
Ported from ``hybrid-local-cloud-compute/adapters/archon_adapter.py``.
|
||||
"""
|
||||
@@ -42,7 +42,15 @@ import types
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.hybrid._base import LocalCloudAgent, _record_event
|
||||
from openjarvis.agents.hybrid._base import (
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
LocalCloudAgent,
|
||||
_bump_cloud_calls,
|
||||
_bump_local_calls,
|
||||
_record_event,
|
||||
build_web_search_tool,
|
||||
web_search_cfg,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
NO_TEMP_PREFIXES,
|
||||
)
|
||||
@@ -79,13 +87,11 @@ def _stub_archon_imports() -> None:
|
||||
|
||||
|
||||
def _add_archon_to_path() -> None:
|
||||
"""Locate Archon's ``src`` dir. The hybrid clone is the default location;
|
||||
override with ``ARCHON_SRC`` env var if Archon is installed elsewhere."""
|
||||
archon_src = os.environ.get(
|
||||
"ARCHON_SRC",
|
||||
"/matx/u/aspark/hybrid-local-cloud-compute/external/Archon/src",
|
||||
)
|
||||
if archon_src not in sys.path and os.path.isdir(archon_src):
|
||||
"""Locate Archon's ``src`` dir. Set ``ARCHON_SRC`` to point at your
|
||||
Archon checkout (``<repo>/src``); otherwise we assume ``archon`` is on
|
||||
``sys.path`` already (e.g. ``pip install`` from a local clone)."""
|
||||
archon_src = os.environ.get("ARCHON_SRC")
|
||||
if archon_src and os.path.isdir(archon_src) and archon_src not in sys.path:
|
||||
sys.path.insert(0, archon_src)
|
||||
|
||||
|
||||
@@ -125,8 +131,11 @@ def _tally() -> Dict[str, int]:
|
||||
counts = {
|
||||
"cloud_prompt": 0, "cloud_completion": 0,
|
||||
"local_prompt": 0, "local_completion": 0,
|
||||
"n_web_searches": 0,
|
||||
}
|
||||
_TALLY_LOCAL.counts = counts
|
||||
# Backfill for older threads that started before we added the key.
|
||||
counts.setdefault("n_web_searches", 0)
|
||||
return counts
|
||||
|
||||
|
||||
@@ -134,9 +143,24 @@ def _reset_tally() -> None:
|
||||
_TALLY_LOCAL.counts = {
|
||||
"cloud_prompt": 0, "cloud_completion": 0,
|
||||
"local_prompt": 0, "local_completion": 0,
|
||||
"n_web_searches": 0,
|
||||
}
|
||||
|
||||
|
||||
# Per-thread web_search tool: when set, the Anthropic generator declares
|
||||
# it on every call. Set inside ``_run_paradigm`` before invoking Archon
|
||||
# so the ranker/fuser passes pick it up; cleared in ``finally``.
|
||||
_WS_LOCAL = threading.local()
|
||||
|
||||
|
||||
def _set_anthropic_web_search(tool: Optional[Dict[str, Any]]) -> None:
|
||||
_WS_LOCAL.tool = tool
|
||||
|
||||
|
||||
def _get_anthropic_web_search() -> Optional[Dict[str, Any]]:
|
||||
return getattr(_WS_LOCAL, "tool", None)
|
||||
|
||||
|
||||
def _make_local_generator(local_endpoint: str, local_model: str):
|
||||
"""Archon custom-generator signature: (model, messages, max_tokens, temperature) -> str."""
|
||||
from openai import OpenAI
|
||||
@@ -153,6 +177,7 @@ def _make_local_generator(local_endpoint: str, local_model: str):
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
_bump_local_calls()
|
||||
except Exception as e:
|
||||
_record_event({
|
||||
"kind": "archon_local_gen_error",
|
||||
@@ -205,6 +230,7 @@ def _wrap_archon_cloud_generators() -> None:
|
||||
kwargs["max_completion_tokens"] = max_tokens
|
||||
t0 = _time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
u = resp.usage
|
||||
if u:
|
||||
_tally()["cloud_prompt"] += getattr(u, "prompt_tokens", 0) or 0
|
||||
@@ -237,13 +263,20 @@ def _wrap_archon_cloud_generators() -> None:
|
||||
)
|
||||
if not model.startswith(NO_TEMP_PREFIXES):
|
||||
kwargs["temperature"] = temperature
|
||||
ws_tool = _get_anthropic_web_search()
|
||||
if ws_tool is not None:
|
||||
kwargs["tools"] = [ws_tool]
|
||||
t0 = _time.time()
|
||||
resp = client.messages.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
text = "".join(b.text for b in resp.content if hasattr(b, "text"))
|
||||
u = resp.usage
|
||||
if u:
|
||||
_tally()["cloud_prompt"] += getattr(u, "input_tokens", 0) or 0
|
||||
_tally()["cloud_completion"] += getattr(u, "output_tokens", 0) or 0
|
||||
srv = getattr(u, "server_tool_use", None) if u else None
|
||||
n_searches = getattr(srv, "web_search_requests", 0) if srv else 0
|
||||
_tally()["n_web_searches"] += int(n_searches)
|
||||
_record_event({
|
||||
"kind": "archon_cloud_anthropic",
|
||||
"model": model,
|
||||
@@ -252,6 +285,8 @@ def _wrap_archon_cloud_generators() -> None:
|
||||
"response": text.strip(),
|
||||
"tokens_in": getattr(u, "input_tokens", 0) if u else 0,
|
||||
"tokens_out": getattr(u, "output_tokens", 0) if u else 0,
|
||||
"n_web_searches": int(n_searches),
|
||||
"tools_declared": kwargs.get("tools"),
|
||||
"latency_s": _time.time() - t0,
|
||||
"ts": _time.time(),
|
||||
})
|
||||
@@ -355,14 +390,18 @@ def _presets():
|
||||
"samples": 1,
|
||||
}],
|
||||
],
|
||||
"single_local": lambda K, local_model, *_a, **_kw: [
|
||||
# ``single_local`` honors the cfg ``max_tokens`` (passed positionally
|
||||
# like ``ensemble_rank_fuse``). Previously it hard-coded 2048, which
|
||||
# cut Qwen off mid-reasoning before it could emit the GAIA
|
||||
# ``FINAL ANSWER:`` line — the scorer then had nothing to extract.
|
||||
"single_local": lambda K, local_model, ranker_model, fuser_model, max_tokens, temperature: [
|
||||
[{
|
||||
"type": "generator",
|
||||
"model": local_model,
|
||||
"model_type": "vllm_local",
|
||||
"top_k": 1,
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 2048,
|
||||
"max_tokens": max_tokens,
|
||||
"samples": 1,
|
||||
}],
|
||||
],
|
||||
@@ -430,6 +469,14 @@ class ArchonAgent(LocalCloudAgent):
|
||||
archon_cfg = {"name": f"hybrid-archon-{arch}", "layers": layers}
|
||||
|
||||
_reset_tally()
|
||||
# Web_search opt-in: when enabled, declare the native server-side
|
||||
# tool on Anthropic ranker/fuser passes (via thread-local). The
|
||||
# local proposers run on vLLM and don't see it.
|
||||
ws_enabled, ws_max_uses = web_search_cfg(cfg)
|
||||
if ws_enabled:
|
||||
_set_anthropic_web_search(build_web_search_tool(ws_max_uses))
|
||||
else:
|
||||
_set_anthropic_web_search(None)
|
||||
archon = Archon(archon_cfg)
|
||||
|
||||
try:
|
||||
@@ -437,8 +484,16 @@ class ArchonAgent(LocalCloudAgent):
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": input},
|
||||
])
|
||||
except Exception as e:
|
||||
answer = f"[archon error: {e!r}]"
|
||||
except Exception:
|
||||
# Re-raise so the base ``run()`` / runner's ``_run_one_inner``
|
||||
# records this in the row's ``error`` field instead of stashing
|
||||
# the exception string in ``answer`` (where it scores as a wrong
|
||||
# answer and never counts toward ``n_err``). Anthropic 529
|
||||
# overloads on the ranker/fuser pass were silently masked this
|
||||
# way — they are infra failures, not model misses.
|
||||
raise
|
||||
finally:
|
||||
_set_anthropic_web_search(None)
|
||||
|
||||
if isinstance(answer, list):
|
||||
answer = answer[-1] if answer else ""
|
||||
@@ -446,16 +501,21 @@ class ArchonAgent(LocalCloudAgent):
|
||||
|
||||
cp = _tally()["cloud_prompt"]
|
||||
cc = _tally()["cloud_completion"]
|
||||
n_searches = _tally().get("n_web_searches", 0)
|
||||
cost = _cost_cloud(ranker_model, cp, cc)
|
||||
if fuser_model != ranker_model:
|
||||
# Conservative: charge both at the more expensive of the two.
|
||||
cost = max(cost, _cost_cloud(fuser_model, cp, cc))
|
||||
cost += n_searches * WEB_SEARCH_COST_PER_CALL
|
||||
|
||||
meta = {
|
||||
"tokens_local": _tally()["local_prompt"] + _tally()["local_completion"],
|
||||
"tokens_cloud": cp + cc,
|
||||
"cost_usd": cost,
|
||||
"turns": (K + 2) if arch == "ensemble_rank_fuse" else 1,
|
||||
"web_search_uses": n_searches,
|
||||
# GAIA: only ranker/fuser can hit web_search; proposers don't.
|
||||
"tool_calls": int(n_searches),
|
||||
"traces": {
|
||||
"architecture": arch,
|
||||
"n_samples": K,
|
||||
@@ -463,6 +523,8 @@ class ArchonAgent(LocalCloudAgent):
|
||||
"fuser_model": fuser_model,
|
||||
"local_model": self._local_model,
|
||||
"tokens_breakdown": dict(_tally()),
|
||||
"web_search_enabled": ws_enabled,
|
||||
"n_web_searches": n_searches,
|
||||
},
|
||||
}
|
||||
return answer, meta
|
||||
@@ -566,6 +628,9 @@ class ArchonAgent(LocalCloudAgent):
|
||||
"tokens_cloud": total_tokens_cloud,
|
||||
"cost_usd": total_cost,
|
||||
"turns": sum(c["turns"] for c in candidates) + 1,
|
||||
# SWE: total bash turns across the K candidate runs; ranker
|
||||
# is a single text call with no tools.
|
||||
"tool_calls": int(sum(c["turns"] for c in candidates)),
|
||||
"traces": {
|
||||
"swe_mode": True,
|
||||
"K": K,
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
"""BaselineCloudAgent — cloud-only reference for the hybrid ablation.
|
||||
|
||||
Used as the "what does the cloud do alone?" row in the n=100 ablation
|
||||
matrix (see ``.openjarvis/experiments/hybrid/docs/results-table.md``).
|
||||
No local model is involved — ``local_*`` settings are ignored.
|
||||
|
||||
On GAIA the agent makes one cloud call with the formatted prompt (which
|
||||
already carries the ``FINAL ANSWER:`` format reminder from
|
||||
``_prompts.format_gaia``) and returns the text. On SWE-bench-Verified
|
||||
the agent delegates to :func:`run_swe_agent_loop` with ``backbone="cloud"``
|
||||
so the model gets to run bash and read the repo — same wiring as the
|
||||
``mini-swe-agent-swebenchverified-opus-*`` cells. As of 2026-05-15
|
||||
``_loop_cloud`` dispatches to per-endpoint loops for Anthropic, OpenAI,
|
||||
and Gemini so all three cloud backbones get the proper bash-agent loop
|
||||
on SWE (previously OpenAI / Gemini SWE cells silently fell back to a
|
||||
one-shot blind patch — fixed).
|
||||
|
||||
Construction args mirror :class:`LocalCloudAgent`. The ``cloud`` block
|
||||
in the cell registry determines the cloud model + endpoint; ``local``
|
||||
is accepted for schema compatibility but unused.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.hybrid._base import (
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
LocalCloudAgent,
|
||||
build_web_search_tool,
|
||||
web_search_cfg,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import cost as estimate_cost
|
||||
from openjarvis.agents.hybrid._prices import default_max_output_tokens
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
|
||||
@AgentRegistry.register("baseline_cloud")
|
||||
class BaselineCloudAgent(LocalCloudAgent):
|
||||
"""Cloud-only baseline used as a reference in the n=100 ablation.
|
||||
|
||||
Configurable knobs via ``cfg``:
|
||||
|
||||
- ``cloud_max_tokens`` (int, default 4096 / 16384 for reasoning models):
|
||||
max_tokens per GAIA call and per turn of the SWE agent loop. Default
|
||||
jumps to 16384 for GPT-5 family and Gemini 2.5 Pro because those models
|
||||
burn the budget on hidden chain-of-thought before emitting visible
|
||||
answer text; at 4096 they silently truncated 18–26% of GAIA cells with
|
||||
empty answers. Override per-cell via ``method_cfg`` to opt out.
|
||||
- ``swe_max_turns`` (int, default 30): SWE-bench loop turn cap.
|
||||
- ``swe_bash_timeout_s`` (int, default 120): SWE-bench bash timeout.
|
||||
"""
|
||||
|
||||
agent_id = "baseline_cloud"
|
||||
|
||||
def _run_paradigm(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
cfg = self._cfg
|
||||
task: Dict[str, Any] = {}
|
||||
if context is not None:
|
||||
task = context.metadata.get("task") or {}
|
||||
|
||||
is_swe = bool(
|
||||
task.get("problem_statement")
|
||||
and task.get("repo")
|
||||
and task.get("base_commit")
|
||||
)
|
||||
|
||||
if is_swe:
|
||||
out = run_swe_agent_loop(
|
||||
task,
|
||||
backbone="cloud",
|
||||
backbone_model=self._cloud_model,
|
||||
cloud_endpoint=self._cloud_endpoint,
|
||||
initial_prompt=input,
|
||||
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("cloud_max_tokens", default_max_output_tokens(self._cloud_model))),
|
||||
trace_prefix="baseline_cloud",
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": 0,
|
||||
"tokens_cloud": out["tokens_in"] + out["tokens_out"],
|
||||
"cost_usd": out["cost_usd"],
|
||||
"turns": out["turns"],
|
||||
# SWE-bench: one bash invocation per agent turn.
|
||||
"tool_calls": int(out["turns"]),
|
||||
"traces": {
|
||||
"backbone": "cloud",
|
||||
"max_turns_hit": out["max_turns_hit"],
|
||||
"patch_chars": len(out["patch"]),
|
||||
"final_summary": out["final_summary"],
|
||||
},
|
||||
}
|
||||
return out["answer"], meta
|
||||
|
||||
# GAIA branch. If `web_search.enabled` is true AND we're on
|
||||
# Anthropic, run the multi-turn agent loop with the native
|
||||
# server-side web_search tool. Otherwise fall back to the
|
||||
# legacy one-shot call (preserves behavior of every existing
|
||||
# non-opted-in cell).
|
||||
ws_enabled, ws_max_uses = web_search_cfg(cfg)
|
||||
gaia_max_turns = int(cfg.get("gaia_max_turns", 8))
|
||||
if ws_enabled and self._cloud_endpoint == "anthropic":
|
||||
text, p_tok, c_tok, n_searches, turns = self._call_anthropic_agent(
|
||||
self._cloud_model,
|
||||
user=input,
|
||||
max_tokens=int(cfg.get("cloud_max_tokens", default_max_output_tokens(self._cloud_model))),
|
||||
temperature=0.0,
|
||||
tools=[build_web_search_tool(ws_max_uses)],
|
||||
max_turns=gaia_max_turns,
|
||||
)
|
||||
cost = (
|
||||
estimate_cost(self._cloud_model, p_tok, c_tok)
|
||||
+ n_searches * WEB_SEARCH_COST_PER_CALL
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": 0,
|
||||
"tokens_cloud": p_tok + c_tok,
|
||||
"cost_usd": cost,
|
||||
"turns": turns,
|
||||
"web_search_uses": n_searches,
|
||||
# GAIA: the only tool is web_search.
|
||||
"tool_calls": int(n_searches),
|
||||
"traces": {
|
||||
"mode": "anthropic_agent_loop",
|
||||
"is_swe": is_swe,
|
||||
"cloud_endpoint": self._cloud_endpoint,
|
||||
"web_search_enabled": True,
|
||||
"web_search_max_uses": ws_max_uses,
|
||||
"n_web_searches": n_searches,
|
||||
},
|
||||
}
|
||||
return text, meta
|
||||
|
||||
if ws_enabled and self._cloud_endpoint != "anthropic":
|
||||
# OpenAI / Gemini don't have a parity native web_search tool
|
||||
# wired here. Skip cleanly rather than fake one. Cells that
|
||||
# want web_search must run on Anthropic until those backends
|
||||
# are wired.
|
||||
self.record_trace_event({
|
||||
"kind": "web_search_skipped",
|
||||
"reason": "non_anthropic_endpoint",
|
||||
"endpoint": self._cloud_endpoint,
|
||||
})
|
||||
|
||||
# One-shot direct cloud call. GAIA only — SWE goes through the
|
||||
# mini-SWE-agent loop above (now supports anthropic/openai/gemini).
|
||||
text, p_tok, c_tok = self._call_cloud(
|
||||
user=input,
|
||||
max_tokens=int(cfg.get("cloud_max_tokens", default_max_output_tokens(self._cloud_model))),
|
||||
temperature=0.0,
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": 0,
|
||||
"tokens_cloud": p_tok + c_tok,
|
||||
"cost_usd": estimate_cost(self._cloud_model, p_tok, c_tok),
|
||||
"turns": 1,
|
||||
"web_search_uses": 0,
|
||||
# GAIA one-shot: zero tool calls (no bash, no web_search).
|
||||
"tool_calls": 0,
|
||||
"traces": {
|
||||
"mode": "one_shot",
|
||||
"is_swe": is_swe,
|
||||
"cloud_endpoint": self._cloud_endpoint,
|
||||
},
|
||||
}
|
||||
return text, meta
|
||||
|
||||
|
||||
__all__ = ["BaselineCloudAgent"]
|
||||
@@ -0,0 +1,159 @@
|
||||
"""BaselineLocalAgent — local-only reference for the hybrid ablation.
|
||||
|
||||
Mirror of :class:`BaselineCloudAgent` (`baseline_cloud.py`) but the entire
|
||||
trajectory runs on the local vLLM model. No cloud teacher / router / advisor
|
||||
is involved — this is the "what does the local model do by itself?" floor
|
||||
in the n=100 ablation matrix.
|
||||
|
||||
On GAIA the agent makes one local call with the formatted prompt (which
|
||||
already carries the ``FINAL ANSWER:`` reminder from
|
||||
``_prompts.format_gaia``) and returns the text. On SWE-bench-Verified
|
||||
the agent delegates to :func:`run_swe_agent_loop` with ``backbone="local"``
|
||||
so the model gets to run bash and read the repo — same wiring as the
|
||||
``mini-swe-agent`` cells but driven by the local model.
|
||||
|
||||
Construction args mirror :class:`LocalCloudAgent`. The ``local`` block in
|
||||
the cell registry determines the local model + endpoint; ``cloud`` is
|
||||
accepted for schema compatibility but unused (and ``cost_usd`` is always
|
||||
0 — local inference is free).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.agents.hybrid._base import LocalCloudAgent
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
|
||||
@AgentRegistry.register("baseline_local")
|
||||
class BaselineLocalAgent(LocalCloudAgent):
|
||||
"""Local-only baseline. ``cloud_*`` fields are ignored.
|
||||
|
||||
Configurable knobs via ``cfg``:
|
||||
|
||||
- ``local_max_tokens`` (int, default 4096): max_tokens per GAIA call
|
||||
and per turn of the SWE agent loop.
|
||||
- ``local_temperature`` (float, default 0.0): sampling temperature
|
||||
for the local model.
|
||||
- ``swe_use_agent_loop`` (bool, default True for SWE): if False the
|
||||
SWE branch falls back to a one-shot blind patch (not recommended;
|
||||
kept for parity with other agents).
|
||||
- ``swe_max_turns`` (int, default 30): SWE-bench loop turn cap.
|
||||
- ``swe_bash_timeout_s`` (int, default 120): bash timeout per turn.
|
||||
- ``swe_turn_max_tokens`` (int, default 4096): max_tokens per agent
|
||||
turn inside the SWE loop. Falls back to ``local_max_tokens``.
|
||||
"""
|
||||
|
||||
agent_id = "baseline_local"
|
||||
|
||||
def _run_paradigm(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
cfg = self._cfg
|
||||
task: Dict[str, Any] = {}
|
||||
if context is not None:
|
||||
task = context.metadata.get("task") or {}
|
||||
|
||||
is_swe = bool(
|
||||
task.get("problem_statement")
|
||||
and task.get("repo")
|
||||
and task.get("base_commit")
|
||||
)
|
||||
|
||||
local_max_tokens = int(cfg.get("local_max_tokens", 4096))
|
||||
local_temperature = float(cfg.get("local_temperature", 0.0))
|
||||
|
||||
if not self._local_model or not self._local_endpoint:
|
||||
raise ValueError(
|
||||
"baseline_local requires `local.model` and `local.endpoint` "
|
||||
"in the cell registry — got "
|
||||
f"model={self._local_model!r}, endpoint={self._local_endpoint!r}"
|
||||
)
|
||||
|
||||
if is_swe:
|
||||
use_loop = bool(cfg.get("swe_use_agent_loop", True))
|
||||
if use_loop:
|
||||
out = run_swe_agent_loop(
|
||||
task,
|
||||
backbone="local",
|
||||
backbone_model=self._local_model,
|
||||
local_endpoint=self._local_endpoint,
|
||||
initial_prompt=input,
|
||||
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", local_max_tokens)
|
||||
),
|
||||
trace_prefix="baseline_local",
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": out["tokens_in"] + out["tokens_out"],
|
||||
"tokens_cloud": 0,
|
||||
"cost_usd": 0.0,
|
||||
"turns": out["turns"],
|
||||
# SWE-bench: one bash invocation per agent turn.
|
||||
"tool_calls": int(out["turns"]),
|
||||
"traces": {
|
||||
"backbone": "local",
|
||||
"max_turns_hit": out["max_turns_hit"],
|
||||
"patch_chars": len(out["patch"]),
|
||||
"final_summary": out["final_summary"],
|
||||
},
|
||||
}
|
||||
return out["answer"], meta
|
||||
|
||||
# One-shot blind patch fallback on SWE (no bash).
|
||||
text, p_tok, c_tok = self._call_vllm(
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
user=input,
|
||||
max_tokens=local_max_tokens,
|
||||
temperature=local_temperature,
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": p_tok + c_tok,
|
||||
"tokens_cloud": 0,
|
||||
"cost_usd": 0.0,
|
||||
"turns": 1,
|
||||
"tool_calls": 0,
|
||||
"traces": {
|
||||
"mode": "one_shot_swe",
|
||||
"backbone": "local",
|
||||
},
|
||||
}
|
||||
return text, meta
|
||||
|
||||
# GAIA branch — one-shot local call. Matches baseline_cloud's
|
||||
# GAIA one-shot path (no web_search; baseline_cloud's web_search
|
||||
# is Anthropic server-side only and has no local equivalent).
|
||||
text, p_tok, c_tok = self._call_vllm(
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
user=input,
|
||||
max_tokens=local_max_tokens,
|
||||
temperature=local_temperature,
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": p_tok + c_tok,
|
||||
"tokens_cloud": 0,
|
||||
"cost_usd": 0.0,
|
||||
"turns": 1,
|
||||
"web_search_uses": 0,
|
||||
"tool_calls": 0,
|
||||
"traces": {
|
||||
"mode": "one_shot",
|
||||
"is_swe": is_swe,
|
||||
"backbone": "local",
|
||||
},
|
||||
}
|
||||
return text, meta
|
||||
|
||||
|
||||
__all__ = ["BaselineLocalAgent"]
|
||||
@@ -31,6 +31,7 @@ from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
@@ -39,8 +40,16 @@ 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._base import (
|
||||
GEMINI_SEARCH_COST_PER_CALL,
|
||||
OPENAI_WEB_SEARCH_COST_PER_CALL,
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
LocalCloudAgent,
|
||||
build_web_search_tool,
|
||||
web_search_cfg,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
PRICES,
|
||||
is_gpt5_family,
|
||||
supports_temperature,
|
||||
)
|
||||
@@ -175,57 +184,327 @@ def _vllm_alive(base_url: str) -> bool:
|
||||
|
||||
|
||||
def _default_pool(local_model: Optional[str], local_endpoint: Optional[str]) -> List[Dict[str, Any]]:
|
||||
"""Default worker pool — faithful to the Sakana Conductor paper (arXiv 2512.04388).
|
||||
|
||||
The paper composes a heterogeneous 7-worker pool spanning three frontier
|
||||
cloud models (Gemini-2.5-Pro, Claude Sonnet-4, GPT-5) and four open-weights
|
||||
workers routed via OpenRouter (DeepSeek-R1-Distill-Qwen-32B, Gemma3-27B-it,
|
||||
Qwen3-32B with reasoning off, Qwen3-32B with reasoning on). No local vLLM
|
||||
worker is in the paper's default — cells that want one should supply it
|
||||
explicitly via ``cfg["worker_pool"]``.
|
||||
|
||||
Each provider is gated by an ``OJ_CONDUCTOR_DISABLE_*`` env var so a cell
|
||||
that lacks one set of credentials can still run the rest of the pool. If
|
||||
every provider is disabled the result is the empty list — the caller's
|
||||
"empty worker pool" check will surface it.
|
||||
|
||||
``local_model`` / ``local_endpoint`` are accepted for signature stability
|
||||
(callers still pass them) but no longer consulted here; the local vLLM is
|
||||
only included when the user opts in via ``cfg["worker_pool"]``.
|
||||
"""
|
||||
del local_model, local_endpoint # paper default carries no local worker
|
||||
pool: List[Dict[str, Any]] = []
|
||||
if local_model and local_endpoint and _vllm_alive(local_endpoint):
|
||||
if not os.environ.get("OJ_CONDUCTOR_DISABLE_GEMINI"):
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "local-qwen",
|
||||
"endpoint": "vllm",
|
||||
"model": local_model,
|
||||
"base_url": local_endpoint,
|
||||
"api_key": "EMPTY",
|
||||
"name": "gemini-pro",
|
||||
"endpoint": "gemini",
|
||||
"model": "gemini-2.5-pro",
|
||||
"description": (
|
||||
"Open-weights Qwen3.5 served locally. Cheap and fast. Good at "
|
||||
"concise extraction, formatting, arithmetic on given data; "
|
||||
"weaker at open-domain factual recall and complex reasoning."
|
||||
"Google Gemini 2.5 Pro. Frontier multimodal reasoner with a "
|
||||
"very large context window. Strong at long-document synthesis, "
|
||||
"multi-hop factual reasoning, and tasks that benefit from "
|
||||
"wide retrieval. Slower and pricier than mid-tier workers."
|
||||
),
|
||||
})
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "frontier-anthropic",
|
||||
"endpoint": "anthropic",
|
||||
"model": "claude-opus-4-7",
|
||||
"description": (
|
||||
"Frontier reasoning model. Strongest at multi-step reasoning, "
|
||||
"careful instruction following, code, and writing. Expensive; "
|
||||
"use sparingly for hard or decisive steps."
|
||||
),
|
||||
})
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "frontier-openai-mini",
|
||||
"endpoint": "openai",
|
||||
"model": "gpt-5-mini",
|
||||
"description": (
|
||||
"Mid-tier OpenAI model. Solid general knowledge and reasoning at "
|
||||
"a fraction of frontier cost. Good default for retrieval-style or "
|
||||
"broad-knowledge questions."
|
||||
),
|
||||
})
|
||||
if not os.environ.get("OJ_CONDUCTOR_DISABLE_ANTHROPIC"):
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "claude-sonnet-4",
|
||||
"endpoint": "anthropic",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"description": (
|
||||
"Anthropic Claude Sonnet 4. Strong general-purpose reasoner "
|
||||
"with careful instruction following and reliable formatting. "
|
||||
"Good default for code, structured writing, and decisive "
|
||||
"steps where accuracy matters more than raw throughput."
|
||||
),
|
||||
})
|
||||
if not os.environ.get("OJ_CONDUCTOR_DISABLE_OPENAI"):
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "gpt-5",
|
||||
"endpoint": "openai",
|
||||
"model": "gpt-5",
|
||||
"description": (
|
||||
"OpenAI GPT-5. Frontier-tier broad-knowledge model. Best for "
|
||||
"open-domain factual recall, creative generation, and "
|
||||
"ambiguous questions where coverage matters. Expensive; use "
|
||||
"for steps where breadth of world knowledge is the bottleneck."
|
||||
),
|
||||
})
|
||||
if not os.environ.get("OJ_CONDUCTOR_DISABLE_OPENROUTER"):
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "deepseek-r1-distill-qwen-32b",
|
||||
"endpoint": "openrouter",
|
||||
"model": "deepseek/deepseek-r1-distill-qwen-32b",
|
||||
"description": (
|
||||
"DeepSeek R1 distilled into Qwen-32B (open weights via "
|
||||
"OpenRouter). Specialized for chain-of-thought math, logic, "
|
||||
"and competitive-programming-style problems. Verbose; "
|
||||
"produces extensive reasoning traces before the final answer."
|
||||
),
|
||||
})
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "gemma3-27b-it",
|
||||
"endpoint": "openrouter",
|
||||
"model": "google/gemma-3-27b-it",
|
||||
"description": (
|
||||
"Google Gemma 3 27B Instruct (open weights via OpenRouter). "
|
||||
"Mid-size instruction-tuned model. Cheap and fast; solid at "
|
||||
"concise summarization, extraction, and short-form Q&A on "
|
||||
"given context. Weaker than the frontier workers on multi-step "
|
||||
"reasoning."
|
||||
),
|
||||
})
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "qwen3-32b",
|
||||
"endpoint": "openrouter",
|
||||
"model": "qwen/qwen3-32b",
|
||||
"description": (
|
||||
"Qwen3-32B in non-thinking mode (open weights via OpenRouter). "
|
||||
"Fast general-purpose dialogue and instruction following. "
|
||||
"Use when the step is straightforward generation, "
|
||||
"summarization, or formatting — does NOT spend tokens on "
|
||||
"internal reasoning."
|
||||
),
|
||||
})
|
||||
pool.append({
|
||||
"id": len(pool),
|
||||
"name": "qwen3-32b-thinking",
|
||||
"endpoint": "openrouter",
|
||||
"model": "qwen/qwen3-32b",
|
||||
"extra_body": {"reasoning": {"effort": "medium"}},
|
||||
"description": (
|
||||
"Qwen3-32B with reasoning enabled (open weights via "
|
||||
"OpenRouter). Same backbone as 'qwen3-32b' but spends tokens "
|
||||
"on an internal chain of thought before answering. Stronger "
|
||||
"on math, code, and multi-step logic; slower and consumes "
|
||||
"more completion tokens. Prefer this for hard reasoning "
|
||||
"steps; prefer the non-thinking variant for plain dialogue."
|
||||
),
|
||||
})
|
||||
# Reassign ids contiguously in case env-gates skipped some entries.
|
||||
for new_id, entry in enumerate(pool):
|
||||
entry["id"] = new_id
|
||||
return pool
|
||||
|
||||
|
||||
# Endpoints conductor's `_call_worker` actually knows how to dispatch to.
|
||||
# Web-search is NOT supported here — toolorchestra has the web-search
|
||||
# dispatcher. OpenRouter is OpenAI-compatible; Gemini is text-only (no
|
||||
# tool-call parity with Anthropic — see `_base._call_gemini`). Both are
|
||||
# opt-in via cfg["worker_pool"] (not added to _default_pool).
|
||||
_CONDUCTOR_VALID_ENDPOINTS = ("vllm", "openai", "anthropic", "openrouter", "gemini")
|
||||
|
||||
|
||||
def _resolve_worker_pool(
|
||||
cfg: Dict[str, Any],
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str,
|
||||
) -> 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``,
|
||||
``endpoint``, and ``model``. ``endpoint`` must be one of
|
||||
``vllm`` / ``openai`` / ``anthropic`` / ``openrouter`` / ``gemini`` —
|
||||
conductor does not wire web-search workers. OpenRouter workers route
|
||||
through the OpenAI-compatible OpenRouter proxy; Gemini workers call
|
||||
the google-genai SDK (text-only, no tool use). Both are opt-in via
|
||||
``cfg["worker_pool"]`` (not in the default pool).
|
||||
|
||||
Substitution: ``model = "$local"`` (or ``"<local>"``) resolves to
|
||||
``local_model``; ``model = "$cloud"`` / ``"<cloud>"`` to ``cloud_model``.
|
||||
|
||||
On any validation failure, raises ``ValueError`` with the message
|
||||
``"Invalid worker_pool entry [<id>]: <reason>"``. Fails fast at agent
|
||||
init rather than mid-task.
|
||||
"""
|
||||
override = cfg.get("worker_pool")
|
||||
if override is None:
|
||||
return _default_pool(local_model, local_endpoint)
|
||||
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"
|
||||
)
|
||||
endpoint = entry.get("endpoint") or entry.get("type")
|
||||
if not isinstance(endpoint, str) or endpoint.lower() not in _CONDUCTOR_VALID_ENDPOINTS:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'endpoint' must be one of "
|
||||
f"{_CONDUCTOR_VALID_ENDPOINTS} (got {endpoint!r})"
|
||||
)
|
||||
endpoint = endpoint.lower()
|
||||
entry["endpoint"] = endpoint
|
||||
# Substitute $local / $cloud placeholders.
|
||||
model = entry.get("model")
|
||||
if isinstance(model, str) and model in ("$local", "<local>"):
|
||||
if not local_model:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: model='{model}' "
|
||||
"requires a local_model to be configured for this cell"
|
||||
)
|
||||
model = local_model
|
||||
entry["model"] = model
|
||||
elif isinstance(model, str) and model in ("$cloud", "<cloud>"):
|
||||
model = cloud_model
|
||||
entry["model"] = model
|
||||
if not isinstance(model, str) or not model:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: 'model' must be a non-empty string"
|
||||
)
|
||||
if endpoint == "vllm":
|
||||
if not entry.get("base_url"):
|
||||
# Default to the local endpoint if not specified — matches
|
||||
# how _default_pool wires it.
|
||||
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")
|
||||
# Local also counts as a non-search worker for the
|
||||
# "must have at least one solver" check.
|
||||
has_non_search = True
|
||||
else:
|
||||
# Cloud workers: model must be priced (any unknown model would
|
||||
# silently cost $0, which masks billing mistakes downstream).
|
||||
# OpenRouter is exempt — its model space is huge and varies
|
||||
# per-provider; cost reporting for openrouter workers is 0
|
||||
# (same as vllm). Cells that need accurate billing for an
|
||||
# openrouter worker should add it to PRICES themselves.
|
||||
if endpoint != "openrouter" and model not in PRICES:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [{wid}]: model {model!r} is "
|
||||
f"not in PRICES (known: {sorted(PRICES)})"
|
||||
)
|
||||
has_non_search = True
|
||||
entry.setdefault(
|
||||
"description",
|
||||
f"User-supplied {endpoint} 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 / openrouter / gemini)"
|
||||
)
|
||||
|
||||
# The planner picks a worker by its `id` (it sees "Model <id> (<name>)"),
|
||||
# but execution dispatches via `workers[model_id]` — list *position*. If a
|
||||
# config supplies non-contiguous or out-of-order ids those two diverge and
|
||||
# a plan that names worker N silently runs a different worker. Enforce
|
||||
# contiguous 0..N-1 ids that equal list position so id == index always.
|
||||
expected = list(range(len(resolved)))
|
||||
actual = [w["id"] for w in resolved]
|
||||
if actual != expected:
|
||||
raise ValueError(
|
||||
f"Invalid worker_pool entry [-]: worker ids must be contiguous "
|
||||
f"0..{len(resolved) - 1} in list order so the planner's model_id "
|
||||
f"matches the dispatch index (got ids {actual}, expected {expected})"
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def _format_worker_pool(workers: List[Dict[str, Any]]) -> str:
|
||||
return "\n".join(
|
||||
f"Model {w['id']} ({w['name']}): {w['description']}" for w in workers
|
||||
)
|
||||
|
||||
|
||||
def _build_conductor_prompt(question: str, workers: List[Dict[str, Any]]) -> str:
|
||||
return (
|
||||
def _search_capable_indices(workers: List[Dict[str, Any]]) -> List[int]:
|
||||
"""Indices of workers whose endpoint can run server-side web search."""
|
||||
return [
|
||||
w["id"] for w in workers
|
||||
if (w.get("endpoint") or "openai").lower()
|
||||
in _SEARCH_CAPABLE_WORKER_ENDPOINTS
|
||||
]
|
||||
|
||||
|
||||
def _build_conductor_prompt(
|
||||
question: str,
|
||||
workers: List[Dict[str, Any]],
|
||||
*,
|
||||
web_search_enabled: bool = False,
|
||||
) -> str:
|
||||
"""Build the planner prompt.
|
||||
|
||||
When ``web_search_enabled`` is set (GAIA cells with web_search on),
|
||||
append an explicit routing constraint: only the listed model indices
|
||||
can perform web research, so any step that needs to look something up
|
||||
on the web MUST be routed to one of them. Steps routed elsewhere
|
||||
answer blind from parametric memory.
|
||||
"""
|
||||
base = (
|
||||
f"Available models:\n{_format_worker_pool(workers)}\n\n"
|
||||
f"User question:\n{question}\n"
|
||||
)
|
||||
if not web_search_enabled:
|
||||
return base
|
||||
capable = _search_capable_indices(workers)
|
||||
if capable:
|
||||
cap_str = ", ".join(str(i) for i in capable)
|
||||
constraint = (
|
||||
"\n\nWEB SEARCH CONSTRAINT:\n"
|
||||
f"Only these model indices can perform live web search: [{cap_str}]. "
|
||||
"Any step that needs to look up facts, current events, or other "
|
||||
"information not reliably known from memory MUST be routed to one "
|
||||
"of those indices. Steps routed to any other model can only use "
|
||||
"their parametric memory and will answer such questions blind.\n"
|
||||
)
|
||||
else:
|
||||
# No search-capable worker at all — the run-level guard raises
|
||||
# before we get here, but keep the prompt honest just in case.
|
||||
constraint = (
|
||||
"\n\nWEB SEARCH CONSTRAINT:\n"
|
||||
"No model in this pool can perform live web search; rely on the "
|
||||
"models' own knowledge.\n"
|
||||
)
|
||||
return base + constraint
|
||||
|
||||
|
||||
def _build_step_prompt(
|
||||
@@ -250,13 +529,42 @@ def _build_step_prompt(
|
||||
|
||||
# ---------- Worker invocation ----------
|
||||
|
||||
# Worker endpoints that can run server-side web search via a `_base`
|
||||
# agent loop. openrouter / vllm cannot ground.
|
||||
_SEARCH_CAPABLE_WORKER_ENDPOINTS = ("anthropic", "openai", "gemini")
|
||||
|
||||
|
||||
def _worker_search_cost_per_call(endpoint: str) -> float:
|
||||
"""Per-search-call USD cost for a worker's cloud endpoint."""
|
||||
if endpoint == "openai":
|
||||
return OPENAI_WEB_SEARCH_COST_PER_CALL
|
||||
if endpoint == "gemini":
|
||||
return GEMINI_SEARCH_COST_PER_CALL
|
||||
return WEB_SEARCH_COST_PER_CALL
|
||||
|
||||
|
||||
def _call_worker(
|
||||
worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any]
|
||||
) -> Tuple[str, int, int, bool]:
|
||||
"""Returns (text, p_tok, c_tok, is_local)."""
|
||||
worker: Dict[str, Any],
|
||||
prompt: str,
|
||||
cfg: Dict[str, Any],
|
||||
*,
|
||||
web_search_tool: Optional[Dict[str, Any]] = None,
|
||||
web_search_max_uses: int = 8,
|
||||
) -> Tuple[str, int, int, bool, int]:
|
||||
"""Returns (text, p_tok, c_tok, is_local, n_web_searches).
|
||||
|
||||
``web_search_tool``: a truthy marker that web_search is enabled for
|
||||
this run. When set AND the worker endpoint is search-capable
|
||||
(anthropic / openai / gemini), the worker call is routed through the
|
||||
matching ``_base`` agent loop so the worker can ground its answer.
|
||||
``n_web_searches`` is the actual count the provider ran.
|
||||
``web_search_max_uses`` caps the Anthropic search tool.
|
||||
Search-incapable workers (openrouter, vllm) ignore it and return 0.
|
||||
"""
|
||||
ep = (worker.get("endpoint") or "openai").lower()
|
||||
max_tok = int(cfg.get("worker_max_tokens", 4096))
|
||||
temp = float(cfg.get("worker_temperature", 0.2))
|
||||
use_ws = web_search_tool is not None
|
||||
|
||||
if ep == "vllm":
|
||||
text, p, c = LocalCloudAgent._call_vllm(
|
||||
@@ -267,24 +575,71 @@ def _call_worker(
|
||||
temperature=temp,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return text, p, c, True
|
||||
return text, p, c, True, 0
|
||||
if ep == "openai":
|
||||
if use_ws:
|
||||
text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=(1.0 if is_gpt5_family(worker["model"]) else temp),
|
||||
)
|
||||
return text, p, c, False, n_searches
|
||||
text, p, c = LocalCloudAgent._call_openai(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=(1.0 if is_gpt5_family(worker["model"]) else temp),
|
||||
)
|
||||
return text, p, c, False
|
||||
return text, p, c, False, 0
|
||||
if ep == "openrouter":
|
||||
# OpenRouter is OpenAI-compatible; the helper handles the
|
||||
# base_url + OPENROUTER_API_KEY plumbing. No server-side web
|
||||
# search wired here. is_local=False so tokens count as cloud.
|
||||
# ``worker["extra_body"]`` (e.g. {"reasoning": {"effort": "medium"}})
|
||||
# is forwarded to the SDK so paper-faithful workers like
|
||||
# "qwen3-32b-thinking" can toggle reasoning per-call.
|
||||
extra_body = worker.get("extra_body")
|
||||
text, p, c = LocalCloudAgent._call_openrouter(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
extra_body=extra_body if isinstance(extra_body, dict) else None,
|
||||
)
|
||||
return text, p, c, False, 0
|
||||
if ep == "anthropic":
|
||||
eff_temp = temp if supports_temperature(worker["model"]) else 0.0
|
||||
text, p, c, _ = LocalCloudAgent._call_anthropic(
|
||||
worker["model"],
|
||||
anthropic_kwargs: Dict[str, Any] = dict(
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=eff_temp,
|
||||
)
|
||||
return text, p, c, False
|
||||
if web_search_tool is not None:
|
||||
anthropic_kwargs["tools"] = [web_search_tool]
|
||||
text, p, c, n_searches = LocalCloudAgent._call_anthropic(
|
||||
worker["model"], **anthropic_kwargs
|
||||
)
|
||||
return text, p, c, False, n_searches
|
||||
if ep == "gemini":
|
||||
# Gemini Developer API via google-genai. With web_search on, route
|
||||
# through the Google-Search-grounded agent loop; otherwise plain
|
||||
# text generation. is_local=False so tokens count as cloud.
|
||||
if use_ws:
|
||||
text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
)
|
||||
return text, p, c, False, n_searches
|
||||
text, p, c = LocalCloudAgent._call_gemini(
|
||||
worker["model"],
|
||||
user=prompt,
|
||||
max_tokens=max_tok,
|
||||
temperature=temp,
|
||||
)
|
||||
return text, p, c, False, 0
|
||||
raise ValueError(f"unsupported worker endpoint: {ep!r}")
|
||||
|
||||
|
||||
@@ -295,30 +650,35 @@ def _swe_worker_step(
|
||||
cfg: Dict[str, Any],
|
||||
workdir: Path,
|
||||
step_idx: int,
|
||||
) -> Tuple[str, int, int, bool]:
|
||||
) -> Tuple[str, int, int, bool, int, int]:
|
||||
"""Run one Conductor worker step as a mini-SWE-agent subloop on a shared
|
||||
workdir. Returns (final_summary_or_diff, tokens_in, tokens_out, is_local)
|
||||
in the same shape as ``_call_worker``."""
|
||||
workdir. Returns (final_summary_or_diff, tokens_in, tokens_out, is_local,
|
||||
n_web_searches, bash_turns). SWE workers don't use web_search (the bash
|
||||
tool is the only tool they need); ``bash_turns`` counts the agent-loop
|
||||
turns so the caller can surface ``tool_calls``."""
|
||||
ep = (worker.get("endpoint") or "openai").lower()
|
||||
if ep == "vllm":
|
||||
backbone, model, endpoint, is_local = (
|
||||
"local", worker["model"], worker.get("base_url"), True,
|
||||
)
|
||||
cloud_endpoint = "anthropic" # unused on the local path
|
||||
elif ep == "anthropic":
|
||||
backbone, model, endpoint, is_local = (
|
||||
"cloud", worker["model"], None, False,
|
||||
)
|
||||
cloud_endpoint = "anthropic"
|
||||
else:
|
||||
# OpenAI workers (gpt-5-mini etc.) aren't supported as agent-loop
|
||||
# backbones today (the loop's tool-call format is Anthropic- or
|
||||
# OpenAI-via-vllm-shaped only). Fall back to one-shot for those —
|
||||
# SWE-bench-wise they were already weak; this preserves behavior.
|
||||
return _call_worker(worker, prompt, cfg)
|
||||
text, p, c, is_local, n_searches = _call_worker(worker, prompt, cfg)
|
||||
return text, p, c, is_local, n_searches, 0
|
||||
out = run_swe_agent_loop(
|
||||
task,
|
||||
backbone=backbone,
|
||||
backbone_model=model,
|
||||
cloud_endpoint="anthropic" if backbone == "cloud" else "anthropic",
|
||||
cloud_endpoint=cloud_endpoint,
|
||||
local_endpoint=endpoint,
|
||||
initial_prompt=prompt,
|
||||
max_turns=int(cfg.get("swe_max_turns", 30)),
|
||||
@@ -328,7 +688,10 @@ def _swe_worker_step(
|
||||
trace_prefix=f"conductor_step{step_idx}",
|
||||
workdir=workdir,
|
||||
)
|
||||
return out["final_summary"] or out["answer"], out["tokens_in"], out["tokens_out"], is_local
|
||||
return (
|
||||
out["final_summary"] or out["answer"],
|
||||
out["tokens_in"], out["tokens_out"], is_local, 0, int(out["turns"]),
|
||||
)
|
||||
|
||||
|
||||
@AgentRegistry.register("conductor")
|
||||
@@ -337,6 +700,20 @@ class ConductorAgent(LocalCloudAgent):
|
||||
|
||||
agent_id = "conductor"
|
||||
|
||||
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 (default pool is built later, lazily,
|
||||
# because `_vllm_alive` needs a live network probe).
|
||||
if self._cfg.get("worker_pool") is not None:
|
||||
_resolve_worker_pool(
|
||||
self._cfg,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
)
|
||||
|
||||
def _run_paradigm(
|
||||
self,
|
||||
input: str,
|
||||
@@ -345,14 +722,45 @@ class ConductorAgent(LocalCloudAgent):
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
question = input
|
||||
cfg = self._cfg
|
||||
workers = cfg.get("workers") or _default_pool(
|
||||
self._local_model, self._local_endpoint
|
||||
)
|
||||
# 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 (Opus +
|
||||
# gpt-5-mini + optional local Qwen).
|
||||
if cfg.get("workers"):
|
||||
workers = cfg["workers"]
|
||||
else:
|
||||
workers = _resolve_worker_pool(
|
||||
cfg,
|
||||
self._local_model,
|
||||
self._local_endpoint,
|
||||
self._cloud_model,
|
||||
)
|
||||
if not workers:
|
||||
raise RuntimeError("conductor: empty worker pool")
|
||||
|
||||
# 1. Plan
|
||||
user = _build_conductor_prompt(question, workers)
|
||||
# Determine swe_mode up front — needed so the planner prompt can
|
||||
# carry the GAIA web-search routing constraint (Task 3). SWE tasks
|
||||
# use the bash tool, not web_search, so the constraint is GAIA-only.
|
||||
task_meta_early = (
|
||||
context.metadata.get("task") if context is not None else {}
|
||||
) or {}
|
||||
swe_mode_early = (
|
||||
bool(cfg.get("swe_use_agent_loop"))
|
||||
and bool(task_meta_early.get("problem_statement"))
|
||||
and bool(task_meta_early.get("repo"))
|
||||
and bool(task_meta_early.get("base_commit"))
|
||||
)
|
||||
ws_enabled, ws_max_uses = web_search_cfg(cfg)
|
||||
planner_ws = ws_enabled and not swe_mode_early
|
||||
|
||||
# 1. Plan — when web_search is on (GAIA), the prompt names which
|
||||
# worker indices can actually search, so the planner routes
|
||||
# research steps to a search-capable worker.
|
||||
user = _build_conductor_prompt(
|
||||
question, workers, web_search_enabled=planner_ws,
|
||||
)
|
||||
plan_text, p_in, p_out = self._call_cloud(
|
||||
user=user,
|
||||
system=CONDUCTOR_SYS,
|
||||
@@ -400,19 +808,54 @@ class ConductorAgent(LocalCloudAgent):
|
||||
# every worker step runs through run_swe_agent_loop on a SHARED
|
||||
# workdir so step N+1 builds on step N's edits. The final patch is
|
||||
# whatever `git diff` produces after the last step.
|
||||
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"))
|
||||
)
|
||||
# ``task_meta`` / ``swe_mode`` were computed up front (see Task-3
|
||||
# planner constraint above) — reuse them.
|
||||
task_meta = task_meta_early
|
||||
swe_mode = swe_mode_early
|
||||
steps: List[Dict[str, Any]] = []
|
||||
tokens_local = 0
|
||||
tokens_cloud = 0
|
||||
cost = 0.0
|
||||
final_answer = ""
|
||||
shared_workdir: Optional[Path] = None
|
||||
n_web_searches_total = 0
|
||||
# tool_calls aggregator: bash turns on SWE (per worker subloop) +
|
||||
# web_search uses on GAIA. Conductor planner is text-only.
|
||||
tool_calls = 0
|
||||
|
||||
# Web_search opt-in: when enabled, search-capable workers
|
||||
# (anthropic / openai / gemini) route through their `_base` agent
|
||||
# loop in `_call_worker` and ground their answers. GAIA-only — SWE
|
||||
# workers use bash not web_search. openrouter / vllm workers can't
|
||||
# ground and ignore the flag. If the worker pool has NO
|
||||
# search-capable worker, web_search can never run on any step —
|
||||
# every step would answer the GAIA question blind from parametric
|
||||
# memory. Fail loud instead of degrading silently.
|
||||
# ``ws_enabled`` / ``ws_max_uses`` computed up front for the planner
|
||||
# constraint — reuse them here.
|
||||
if ws_enabled and not swe_mode:
|
||||
search_workers = [
|
||||
w for w in workers
|
||||
if (w.get("endpoint") or "openai").lower()
|
||||
in _SEARCH_CAPABLE_WORKER_ENDPOINTS
|
||||
]
|
||||
if not search_workers:
|
||||
endpoints = sorted({
|
||||
(w.get("endpoint") or "openai").lower() for w in workers
|
||||
})
|
||||
raise ValueError(
|
||||
f"web_search.enabled=true but the worker pool has no "
|
||||
f"search-capable worker (endpoints present: {endpoints}); "
|
||||
"server-side web_search is wired only for anthropic / "
|
||||
"openai / gemini workers in conductor's _call_worker. "
|
||||
"Add one of those to the pool or disable web_search — "
|
||||
"otherwise every step answers blind from parametric memory."
|
||||
)
|
||||
# ``ws_tool`` doubles as the enable marker passed to `_call_worker`
|
||||
# (truthy => route search-capable workers through their agent loop).
|
||||
ws_tool = (
|
||||
build_web_search_tool(ws_max_uses) if ws_enabled else None
|
||||
)
|
||||
|
||||
try:
|
||||
if swe_mode:
|
||||
@@ -444,18 +887,52 @@ class ConductorAgent(LocalCloudAgent):
|
||||
"swe_mode": swe_mode,
|
||||
})
|
||||
|
||||
worker_ep = (worker.get("endpoint") or "openai").lower()
|
||||
# Post-hoc routing check: if web_search is on but the
|
||||
# planner routed this step to a search-incapable worker,
|
||||
# record a warning into the trace (don't crash — the step
|
||||
# may legitimately not need search; see Task-3 planner
|
||||
# constraint that tries to prevent this upfront).
|
||||
if (
|
||||
ws_enabled and not swe_mode
|
||||
and worker_ep not in _SEARCH_CAPABLE_WORKER_ENDPOINTS
|
||||
):
|
||||
self.record_trace_event({
|
||||
"kind": "conductor_search_routing_warning",
|
||||
"step_idx": i,
|
||||
"worker_id": mid,
|
||||
"worker_name": worker["name"],
|
||||
"worker_endpoint": worker_ep,
|
||||
"warning": (
|
||||
f"web_search enabled but step {i} routed to "
|
||||
f"search-incapable worker {worker['name']!r} "
|
||||
f"(endpoint {worker_ep!r}); this step cannot "
|
||||
"ground and may answer blind."
|
||||
),
|
||||
})
|
||||
|
||||
if swe_mode:
|
||||
text, w_in, w_out, is_local = _swe_worker_step(
|
||||
worker, task_meta, prompt, cfg, shared_workdir, i,
|
||||
text, w_in, w_out, is_local, n_searches, bash_turns = (
|
||||
_swe_worker_step(
|
||||
worker, task_meta, prompt, cfg, shared_workdir, i,
|
||||
)
|
||||
)
|
||||
tool_calls += bash_turns
|
||||
else:
|
||||
text, w_in, w_out, is_local = _call_worker(worker, prompt, cfg)
|
||||
text, w_in, w_out, is_local, n_searches = _call_worker(
|
||||
worker, prompt, cfg,
|
||||
web_search_tool=ws_tool,
|
||||
web_search_max_uses=ws_max_uses,
|
||||
)
|
||||
|
||||
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)
|
||||
cost += n_searches * _worker_search_cost_per_call(worker_ep)
|
||||
n_web_searches_total += n_searches
|
||||
tool_calls += n_searches
|
||||
steps.append({
|
||||
"step_idx": i,
|
||||
"model_id": mid,
|
||||
@@ -497,10 +974,14 @@ class ConductorAgent(LocalCloudAgent):
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost,
|
||||
"turns": len(steps) + 1, # planner + N execution steps
|
||||
"web_search_uses": n_web_searches_total,
|
||||
"tool_calls": int(tool_calls),
|
||||
"traces": {
|
||||
"steps": traces,
|
||||
"plan": plan,
|
||||
"fallback_used": fallback_used,
|
||||
"web_search_enabled": ws_enabled,
|
||||
"n_web_searches": n_web_searches_total,
|
||||
"parse_attempts": parse_attempts,
|
||||
"workers": [
|
||||
{k: v for k, v in w.items() if k != "api_key"}
|
||||
|
||||
@@ -35,7 +35,10 @@ Differences vs. the upstream
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
@@ -43,15 +46,39 @@ 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, _record_event
|
||||
from openjarvis.agents.hybrid._base import (
|
||||
LocalCloudAgent,
|
||||
_bump_cloud_calls,
|
||||
_bump_local_calls,
|
||||
_record_event,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
cost as estimate_cost,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import (
|
||||
is_gpt5_family,
|
||||
supports_temperature,
|
||||
)
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
|
||||
# Gemini's FunctionDeclaration.parameters expects a Schema-shaped dict (or
|
||||
# Schema object) with capitalized type strings ("OBJECT", "STRING"). The
|
||||
# OpenAI/Anthropic JSON-Schema lower-case form is silently dropped — the
|
||||
# model then can't call the tool. Build a fresh dict instead of reusing
|
||||
# BASH_TOOL_OPENAI's parameters.
|
||||
BASH_TOOL_GEMINI_PARAMETERS: Dict[str, Any] = {
|
||||
"type": "OBJECT",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "STRING",
|
||||
"description": "The bash command to run.",
|
||||
},
|
||||
},
|
||||
"required": ["command"],
|
||||
}
|
||||
|
||||
|
||||
SYSTEM_PROMPT = """\
|
||||
You are an expert software engineer fixing a bug in a Python repository. \
|
||||
You have one tool, `bash`, that runs a shell command and returns stdout, \
|
||||
@@ -111,6 +138,36 @@ BASH_TOOL_OPENAI = {
|
||||
|
||||
# ---------- Workdir / bash plumbing ----------
|
||||
|
||||
# Models trained on SWE-bench Docker images (Qwen especially) reflexively
|
||||
# prefix commands with ``cd /testbed`` — the standard container repo path.
|
||||
# Our harness has no ``/testbed``; the repo is cloned into a per-task
|
||||
# tempdir and bash already runs with ``cwd`` set to it. An un-rewritten
|
||||
# ``cd /testbed`` errors with "No such file or directory" and, chained with
|
||||
# ``&&``, aborts the whole command — so the agent burns every turn and
|
||||
# never lands an edit. We rewrite ``/testbed`` references to the real
|
||||
# workdir so those commands run as intended.
|
||||
_TESTBED_CD_RE = re.compile(r"^\s*cd\s+/testbed(?:/\S*)?\s*(?:&&|;)\s*")
|
||||
|
||||
|
||||
def _rewrite_testbed_paths(command: str, workdir: Path) -> str:
|
||||
"""Neutralize hard-coded ``/testbed`` paths in a model-issued command.
|
||||
|
||||
- A leading ``cd /testbed && ...`` (or ``;``) is stripped — bash already
|
||||
runs in the repo root, so the rest of the command is correct as-is.
|
||||
- Any remaining ``/testbed`` occurrences (e.g. ``cat /testbed/foo.py``)
|
||||
are rewritten to the real workdir.
|
||||
"""
|
||||
wd = str(workdir)
|
||||
new = _TESTBED_CD_RE.sub("", command)
|
||||
# Bare ``cd /testbed`` with nothing after it → no-op into the workdir.
|
||||
if re.fullmatch(r"\s*cd\s+/testbed/?\s*", new):
|
||||
new = f"cd {wd}"
|
||||
# Replace any other /testbed path references (word-boundary so we don't
|
||||
# clobber e.g. /testbedrock).
|
||||
new = re.sub(r"/testbed(?=/|\b)", wd, new)
|
||||
return new
|
||||
|
||||
|
||||
def _clone_repo(repo: str, base_commit: str, dest: Path) -> None:
|
||||
"""Shallow-fetch the SWE-bench repo at the right commit into ``dest``."""
|
||||
url = f"https://github.com/{repo}.git"
|
||||
@@ -124,27 +181,89 @@ def _clone_repo(repo: str, base_commit: str, dest: Path) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _decode_bash_output(raw: bytes, exit_code: int) -> str:
|
||||
"""Safely decode bash stdout/stderr bytes into a str the LLM can read.
|
||||
|
||||
The model sometimes runs commands that produce binary output (``cat``
|
||||
on a ``.pyc`` / ``.png`` / packed extension, a ``find`` that pipes a
|
||||
binary blob, ``xxd`` on a libc, ...). Decoding those with strict UTF-8
|
||||
raises ``UnicodeDecodeError`` deep inside ``Popen.communicate()`` and
|
||||
crashes the whole agent loop (1 errored row in the n=100 sweep per
|
||||
such command). We:
|
||||
|
||||
1. Decode with ``errors="replace"`` so partial / mixed output is
|
||||
always recoverable as a str.
|
||||
2. If the result looks predominantly binary — contains a NUL byte OR
|
||||
has more than ~5% U+FFFD replacement chars after decode — replace
|
||||
it with a one-line stub so the model can keep working without
|
||||
drowning in Mojibake. Threshold checked against the *bytes* length
|
||||
(no NUL byte ⇒ probably text-ish; LLMs handle the occasional
|
||||
replacement char fine).
|
||||
"""
|
||||
if not raw:
|
||||
return ""
|
||||
decoded = raw.decode("utf-8", errors="replace")
|
||||
if b"\x00" in raw or decoded.count("�") * 20 > len(decoded):
|
||||
return f"[binary output: {len(raw)} bytes, exit={exit_code}]"
|
||||
return decoded
|
||||
|
||||
|
||||
def _run_bash(
|
||||
command: str, workdir: Path, *, timeout: int = 120, output_cap: int = 10_000
|
||||
) -> Dict[str, Any]:
|
||||
"""Run one shell command in ``workdir``. Returns dict with stdout, stderr,
|
||||
exit_code, and a ``truncated`` flag if output was clamped."""
|
||||
exit_code, and a ``truncated`` flag if output was clamped.
|
||||
|
||||
The command is launched in its own process group (``start_new_session``)
|
||||
so a model-issued command that backgrounds a long-lived child (a dev
|
||||
server, ``sleep``, a hung test runner) can be killed *as a tree* on
|
||||
timeout. Plain ``subprocess.run(..., capture_output=True, timeout=...)``
|
||||
only kills the direct child and then re-blocks on ``communicate()``
|
||||
draining the pipe — which a surviving grandchild holds open forever,
|
||||
silently wedging the whole agent loop.
|
||||
"""
|
||||
t0 = time.time()
|
||||
command = _rewrite_testbed_paths(command, workdir)
|
||||
# Capture as bytes (no ``text=True``) so a tool invocation that emits
|
||||
# binary output (compiled artifact, image, PDF, gzipped tarball) can't
|
||||
# crash the loop on a strict UTF-8 decode mid-``communicate()``. We
|
||||
# decode below with ``errors="replace"`` and, if the result looks
|
||||
# binary (null byte or >5% replacement chars), substitute a stub so
|
||||
# the model doesn't waste tokens / context on Mojibake.
|
||||
proc = subprocess.Popen(
|
||||
["bash", "-lc", command],
|
||||
cwd=str(workdir),
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
start_new_session=True,
|
||||
)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
["bash", "-lc", command],
|
||||
cwd=str(workdir),
|
||||
capture_output=True, text=True, timeout=timeout,
|
||||
)
|
||||
stdout = proc.stdout
|
||||
stderr = proc.stderr
|
||||
stdout_b, stderr_b = proc.communicate(timeout=timeout)
|
||||
exit_code = proc.returncode
|
||||
timed_out = False
|
||||
except subprocess.TimeoutExpired as e:
|
||||
stdout = (e.stdout or "") if isinstance(e.stdout, str) else ""
|
||||
stderr = (e.stderr or "") if isinstance(e.stderr, str) else ""
|
||||
except subprocess.TimeoutExpired:
|
||||
# Kill the whole process group so backgrounded grandchildren can't
|
||||
# keep the stdout/stderr pipe open and deadlock the drain below.
|
||||
for sig in (signal.SIGTERM, signal.SIGKILL):
|
||||
try:
|
||||
os.killpg(proc.pid, sig)
|
||||
except (ProcessLookupError, PermissionError):
|
||||
break
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
break
|
||||
except subprocess.TimeoutExpired:
|
||||
continue
|
||||
try:
|
||||
stdout_b, stderr_b = proc.communicate(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
stdout_b, stderr_b = b"", b"";
|
||||
stdout_b = stdout_b or b""
|
||||
stderr_b = stderr_b or b""
|
||||
exit_code = -1
|
||||
timed_out = True
|
||||
stdout = _decode_bash_output(stdout_b, exit_code)
|
||||
stderr = _decode_bash_output(stderr_b, exit_code)
|
||||
truncated = False
|
||||
if len(stdout) > output_cap:
|
||||
stdout = stdout[:output_cap] + f"\n…[+{len(stdout) - output_cap} chars truncated]"
|
||||
@@ -212,10 +331,13 @@ def run_swe_agent_loop(
|
||||
initial_prompt: Optional[str] = None,
|
||||
max_turns: int = 30,
|
||||
bash_timeout: int = 120,
|
||||
bash_timeout_s: Optional[int] = None,
|
||||
output_cap: int = 10_000,
|
||||
turn_max_tokens: int = 4096,
|
||||
trace_prefix: str = "mini_swe",
|
||||
workdir: Optional[Path] = None,
|
||||
compact_at_tokens: int = 24_000,
|
||||
compact_keep_last: int = 4,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run a mini-SWE-agent loop for one SWE-bench task. Returns:
|
||||
|
||||
@@ -254,6 +376,8 @@ def run_swe_agent_loop(
|
||||
want to chain multiple subloops over the same working tree can
|
||||
manage their own workdir.
|
||||
"""
|
||||
if bash_timeout_s is not None:
|
||||
bash_timeout = int(bash_timeout_s)
|
||||
repo = task.get("repo") or ""
|
||||
base_commit = task.get("base_commit") or ""
|
||||
if not repo or not base_commit:
|
||||
@@ -310,6 +434,8 @@ def run_swe_agent_loop(
|
||||
output_cap=output_cap,
|
||||
turn_max_tokens=turn_max_tokens,
|
||||
trace_prefix=trace_prefix,
|
||||
compact_at_tokens=compact_at_tokens,
|
||||
compact_keep_last=compact_keep_last,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported backbone: {backbone!r}")
|
||||
@@ -340,7 +466,7 @@ def run_swe_agent_loop(
|
||||
shutil.rmtree(workdir, ignore_errors=True)
|
||||
|
||||
|
||||
# ---------- Cloud loop (Anthropic multi-turn with tools) ----------
|
||||
# ---------- Cloud loop (dispatcher → per-endpoint multi-turn tool loops) ----------
|
||||
|
||||
def _loop_cloud(
|
||||
problem: str,
|
||||
@@ -354,11 +480,47 @@ def _loop_cloud(
|
||||
turn_max_tokens: int,
|
||||
trace_prefix: str,
|
||||
) -> Dict[str, Any]:
|
||||
if cloud_endpoint != "anthropic":
|
||||
raise ValueError(
|
||||
f"mini-SWE-agent cloud backbone currently supports anthropic only; "
|
||||
f"got {cloud_endpoint!r}"
|
||||
"""Route to the right per-endpoint cloud loop. Anthropic path is the
|
||||
original byte-identical implementation (16 cells in the n=100 sweep
|
||||
depend on its exact behavior). OpenAI / Gemini paths added 2026-05-15
|
||||
to unblock the 8 SWE cells that were stuck on Anthropic-only support."""
|
||||
if cloud_endpoint == "anthropic":
|
||||
return _loop_cloud_anthropic(
|
||||
problem, workdir,
|
||||
model=model, max_turns=max_turns,
|
||||
bash_timeout=bash_timeout, output_cap=output_cap,
|
||||
turn_max_tokens=turn_max_tokens, trace_prefix=trace_prefix,
|
||||
)
|
||||
if cloud_endpoint == "openai":
|
||||
return _loop_cloud_openai(
|
||||
problem, workdir,
|
||||
model=model, max_turns=max_turns,
|
||||
bash_timeout=bash_timeout, output_cap=output_cap,
|
||||
turn_max_tokens=turn_max_tokens, trace_prefix=trace_prefix,
|
||||
)
|
||||
if cloud_endpoint == "gemini":
|
||||
return _loop_cloud_gemini(
|
||||
problem, workdir,
|
||||
model=model, max_turns=max_turns,
|
||||
bash_timeout=bash_timeout, output_cap=output_cap,
|
||||
turn_max_tokens=turn_max_tokens, trace_prefix=trace_prefix,
|
||||
)
|
||||
raise ValueError(
|
||||
f"mini-SWE-agent cloud backbone unsupported endpoint: {cloud_endpoint!r}"
|
||||
)
|
||||
|
||||
|
||||
def _loop_cloud_anthropic(
|
||||
problem: str,
|
||||
workdir: Path,
|
||||
*,
|
||||
model: str,
|
||||
max_turns: int,
|
||||
bash_timeout: int,
|
||||
output_cap: int,
|
||||
turn_max_tokens: int,
|
||||
trace_prefix: str,
|
||||
) -> Dict[str, Any]:
|
||||
import anthropic
|
||||
client = anthropic.Anthropic(timeout=600.0, max_retries=5)
|
||||
messages: List[Dict[str, Any]] = [{"role": "user", "content": problem}]
|
||||
@@ -380,6 +542,7 @@ def _loop_cloud(
|
||||
kwargs["temperature"] = 0.0
|
||||
t0 = time.time()
|
||||
msg = client.messages.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
tokens_in += msg.usage.input_tokens
|
||||
tokens_out += msg.usage.output_tokens
|
||||
@@ -457,8 +620,596 @@ def _loop_cloud(
|
||||
}
|
||||
|
||||
|
||||
# ---------- Cloud loop (OpenAI multi-turn with function tools) ----------
|
||||
|
||||
def _loop_cloud_openai(
|
||||
problem: str,
|
||||
workdir: Path,
|
||||
*,
|
||||
model: str,
|
||||
max_turns: int,
|
||||
bash_timeout: int,
|
||||
output_cap: int,
|
||||
turn_max_tokens: int,
|
||||
trace_prefix: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""OpenAI Chat Completions multi-turn loop. Mirrors the Anthropic
|
||||
branch: each turn the model either calls ``bash`` (one or more parallel
|
||||
tool_calls) or produces a no-tool-call final message that terminates
|
||||
the loop. The final message's text is returned as ``final_summary``;
|
||||
the patch comes from ``git diff`` on the workdir.
|
||||
|
||||
Quirks:
|
||||
- GPT-5 family rejects ``temperature`` and uses ``max_completion_tokens``
|
||||
instead of ``max_tokens``. We branch on ``is_gpt5_family`` for both.
|
||||
- ``tool_calls`` arguments arrive as JSON-string blobs; we tolerate
|
||||
malformed JSON by treating it as an empty arg dict (matches the
|
||||
``_loop_local`` behavior).
|
||||
"""
|
||||
from openai import OpenAI
|
||||
client = OpenAI(timeout=600.0)
|
||||
|
||||
messages: List[Dict[str, Any]] = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": problem},
|
||||
]
|
||||
tokens_in = 0
|
||||
tokens_out = 0
|
||||
final_text = ""
|
||||
turns = 0
|
||||
for turn in range(1, max_turns + 1):
|
||||
turns = turn
|
||||
kwargs: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"tools": [BASH_TOOL_OPENAI],
|
||||
"tool_choice": "auto",
|
||||
}
|
||||
if is_gpt5_family(model):
|
||||
kwargs["max_completion_tokens"] = turn_max_tokens
|
||||
else:
|
||||
kwargs["max_tokens"] = turn_max_tokens
|
||||
kwargs["temperature"] = 0.0
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
u = resp.usage
|
||||
tokens_in += getattr(u, "prompt_tokens", 0) if u else 0
|
||||
tokens_out += getattr(u, "completion_tokens", 0) if u else 0
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
tool_calls = list(getattr(message, "tool_calls", None) or [])
|
||||
text = message.content or ""
|
||||
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_turn",
|
||||
"turn": turn,
|
||||
"endpoint": "openai",
|
||||
"finish_reason": choice.finish_reason,
|
||||
"tokens_in": getattr(u, "prompt_tokens", 0) if u else 0,
|
||||
"tokens_out": getattr(u, "completion_tokens", 0) if u else 0,
|
||||
"latency_s": latency,
|
||||
"text": text,
|
||||
"tool_calls": [
|
||||
{"id": tc.id, "name": tc.function.name, "arguments": tc.function.arguments}
|
||||
for tc in tool_calls
|
||||
],
|
||||
"ts": time.time(),
|
||||
})
|
||||
|
||||
# Append the assistant turn (including any tool_calls) so the
|
||||
# follow-up tool messages have the right call ids to reference.
|
||||
# OpenAI Chat Completions rejects ``content: null`` with a 400
|
||||
# ("expected a string, got null") on the next turn when this
|
||||
# message gets replayed. Use ``""`` — explicitly allowed by the
|
||||
# schema when ``tool_calls`` is present, and equivalent to
|
||||
# "assistant had no visible text, only tool calls".
|
||||
assistant_msg: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": text or "",
|
||||
}
|
||||
if tool_calls:
|
||||
assistant_msg["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id, "type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in tool_calls
|
||||
]
|
||||
messages.append(assistant_msg)
|
||||
|
||||
if not tool_calls:
|
||||
# Silent-truncation recovery: ``finish_reason='length'`` means
|
||||
# the response was cut mid-generation. If a tool call was
|
||||
# forming when the cap hit, ``tool_calls`` is empty AND ``text``
|
||||
# is empty/short — treating that as "model done" exits the loop
|
||||
# with a useless empty final_summary (observed on gpt-5-mini SWE
|
||||
# cells, 2026-05-15). Inject a one-shot recovery nudge and let
|
||||
# the loop continue; only terminate naturally on ``stop``.
|
||||
if (
|
||||
choice.finish_reason == "length"
|
||||
and not text.strip()
|
||||
and turn < max_turns
|
||||
):
|
||||
messages.append({
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Your previous response was truncated by the token limit "
|
||||
"before producing a tool call or final summary. Retry: "
|
||||
"either issue ONE bash tool call (short command, no large "
|
||||
"output) or send a brief one-line final summary with no "
|
||||
"tool calls to end the loop."
|
||||
),
|
||||
})
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_recover",
|
||||
"turn": turn, "reason": "length_truncation_no_tool_call",
|
||||
"ts": time.time(),
|
||||
})
|
||||
continue
|
||||
# No tool call → the model is done. Same termination rule as
|
||||
# the Anthropic branch.
|
||||
final_text = text.strip()
|
||||
break
|
||||
|
||||
for tc in tool_calls:
|
||||
try:
|
||||
args = json.loads(tc.function.arguments or "{}")
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
if tc.function.name != "bash":
|
||||
obs = f"unknown tool: {tc.function.name!r}"
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_unknown_tool",
|
||||
"turn": turn, "name": tc.function.name, "input": args,
|
||||
"ts": time.time(),
|
||||
})
|
||||
else:
|
||||
command = str(args.get("command", ""))
|
||||
result = _run_bash(
|
||||
command, workdir,
|
||||
timeout=bash_timeout, output_cap=output_cap,
|
||||
)
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_bash",
|
||||
"turn": turn, "command": command,
|
||||
**result, "ts": time.time(),
|
||||
})
|
||||
obs = _format_observation(result)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"content": obs,
|
||||
})
|
||||
|
||||
return {
|
||||
"tokens_in": tokens_in,
|
||||
"tokens_out": tokens_out,
|
||||
"turns": turns,
|
||||
"final_summary": final_text,
|
||||
"max_turns_hit": turns == max_turns and not final_text,
|
||||
}
|
||||
|
||||
|
||||
# ---------- Cloud loop (Gemini multi-turn with function tools) ----------
|
||||
|
||||
def _loop_cloud_gemini(
|
||||
problem: str,
|
||||
workdir: Path,
|
||||
*,
|
||||
model: str,
|
||||
max_turns: int,
|
||||
bash_timeout: int,
|
||||
output_cap: int,
|
||||
turn_max_tokens: int,
|
||||
trace_prefix: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""google-genai multi-turn loop (Gemini Developer API, v1.64).
|
||||
|
||||
Tool-use plumbing diverges from OpenAI/Anthropic enough to warrant a
|
||||
parallel branch:
|
||||
|
||||
- Parameters must be a Schema-shaped dict with capitalized type names
|
||||
("OBJECT", "STRING"). The lower-case JSON-Schema form is silently
|
||||
dropped — the model never produces a call.
|
||||
- We explicitly disable ``automatic_function_calling`` so the SDK
|
||||
stops at the FunctionCall part and we drive the loop ourselves.
|
||||
- Termination heuristic: stop when the model's response has zero
|
||||
function_call parts. Same intent as Anthropic ("no tool_use blocks")
|
||||
and OpenAI ("empty tool_calls"). Gemini occasionally emits a turn
|
||||
with both text and function_call parts; we still treat that as a
|
||||
tool turn (matches Anthropic's behavior with mixed content).
|
||||
- System prompt goes on the config, not the contents list (Gemini
|
||||
convention).
|
||||
- The function-response part body is a free-form dict; we wrap the
|
||||
bash observation in ``{"output": str}``.
|
||||
"""
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
client = genai.Client(http_options=types.HttpOptions(timeout=600_000))
|
||||
bash_tool = types.Tool(function_declarations=[
|
||||
types.FunctionDeclaration(
|
||||
name="bash",
|
||||
description=BASH_TOOL_ANTHROPIC["description"],
|
||||
parameters=BASH_TOOL_GEMINI_PARAMETERS,
|
||||
),
|
||||
])
|
||||
|
||||
contents: List[types.Content] = [
|
||||
types.Content(role="user", parts=[types.Part(text=problem)]),
|
||||
]
|
||||
tokens_in = 0
|
||||
tokens_out = 0
|
||||
final_text = ""
|
||||
turns = 0
|
||||
for turn in range(1, max_turns + 1):
|
||||
turns = turn
|
||||
cfg = types.GenerateContentConfig(
|
||||
temperature=0.0,
|
||||
max_output_tokens=turn_max_tokens,
|
||||
system_instruction=SYSTEM_PROMPT,
|
||||
tools=[bash_tool],
|
||||
automatic_function_calling=types.AutomaticFunctionCallingConfig(
|
||||
disable=True,
|
||||
),
|
||||
)
|
||||
t0 = time.time()
|
||||
resp = client.models.generate_content(
|
||||
model=model, contents=contents, config=cfg,
|
||||
)
|
||||
_bump_cloud_calls()
|
||||
latency = time.time() - t0
|
||||
um = getattr(resp, "usage_metadata", None)
|
||||
p = int(getattr(um, "prompt_token_count", 0) or 0) if um else 0
|
||||
c = int(getattr(um, "candidates_token_count", 0) or 0) if um else 0
|
||||
tokens_in += p
|
||||
tokens_out += c
|
||||
|
||||
# Pull parts out of the first candidate. Defensive against the
|
||||
# zero-candidate / safety-filtered case (returns empty parts and
|
||||
# we'll terminate on the next branch).
|
||||
cand_parts: List[Any] = []
|
||||
try:
|
||||
cand_content = resp.candidates[0].content
|
||||
cand_parts = list(getattr(cand_content, "parts", None) or [])
|
||||
except Exception:
|
||||
cand_content = None
|
||||
cand_parts = []
|
||||
|
||||
text_parts: List[str] = []
|
||||
function_calls: List[Tuple[str, Dict[str, Any]]] = []
|
||||
for part in cand_parts:
|
||||
fc = getattr(part, "function_call", None)
|
||||
if fc is not None:
|
||||
fc_args = dict(getattr(fc, "args", None) or {})
|
||||
function_calls.append((fc.name, fc_args))
|
||||
elif getattr(part, "text", None):
|
||||
text_parts.append(part.text)
|
||||
|
||||
finish_reason = None
|
||||
try:
|
||||
finish_reason = str(resp.candidates[0].finish_reason)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_turn",
|
||||
"turn": turn,
|
||||
"endpoint": "gemini",
|
||||
"finish_reason": finish_reason,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"latency_s": latency,
|
||||
"text": "\n".join(text_parts),
|
||||
"tool_calls": [
|
||||
{"name": name, "arguments": args}
|
||||
for name, args in function_calls
|
||||
],
|
||||
"ts": time.time(),
|
||||
})
|
||||
|
||||
# Append the model's content as-is so the next turn sees its own
|
||||
# prior function_call parts (Gemini requires this for the
|
||||
# function_response to bind).
|
||||
if cand_content is not None:
|
||||
contents.append(cand_content)
|
||||
else:
|
||||
# Safety-filtered / empty — synthesize an empty model turn so
|
||||
# the conversation stays well-formed and exit the loop.
|
||||
contents.append(types.Content(role="model", parts=[]))
|
||||
|
||||
if not function_calls:
|
||||
# Silent-failure recovery for Gemini's quirky finish reasons:
|
||||
# MALFORMED_FUNCTION_CALL — model wanted to call bash but
|
||||
# produced unparseable args (24/100 of broken n=100 SWE
|
||||
# cells, 2026-05-15). Empty text + empty function_calls →
|
||||
# loop would exit with no final summary.
|
||||
# MAX_TOKENS — same shape as OpenAI's ``length``; truncated
|
||||
# mid-generation, no tool call landed.
|
||||
# Inject a recovery nudge and let the loop continue; only
|
||||
# treat genuine ``STOP`` with text as a final answer.
|
||||
fr_str = str(finish_reason or "")
|
||||
empty_text = not any(t.strip() for t in text_parts)
|
||||
recoverable = empty_text and turn < max_turns and (
|
||||
"MALFORMED_FUNCTION_CALL" in fr_str
|
||||
or "MAX_TOKENS" in fr_str
|
||||
)
|
||||
if recoverable:
|
||||
contents.append(types.Content(
|
||||
role="user",
|
||||
parts=[types.Part(text=(
|
||||
"Your previous response had no parsable function call "
|
||||
"and no final text (finish_reason="
|
||||
f"{fr_str}). Retry: either issue ONE well-formed "
|
||||
"`bash` function call (short command, valid JSON-ish "
|
||||
"args) or send a brief final text message with no "
|
||||
"function call to end the loop."
|
||||
))],
|
||||
))
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_recover",
|
||||
"turn": turn,
|
||||
"reason": f"empty_response_{fr_str}",
|
||||
"ts": time.time(),
|
||||
})
|
||||
continue
|
||||
final_text = "\n".join(text_parts).strip()
|
||||
break
|
||||
|
||||
response_parts: List[Any] = []
|
||||
for name, args in function_calls:
|
||||
if name != "bash":
|
||||
obs = f"unknown tool: {name!r}"
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_unknown_tool",
|
||||
"turn": turn, "name": name, "input": args,
|
||||
"ts": time.time(),
|
||||
})
|
||||
else:
|
||||
command = str(args.get("command", ""))
|
||||
result = _run_bash(
|
||||
command, workdir,
|
||||
timeout=bash_timeout, output_cap=output_cap,
|
||||
)
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_bash",
|
||||
"turn": turn, "command": command,
|
||||
**result, "ts": time.time(),
|
||||
})
|
||||
obs = _format_observation(result)
|
||||
response_parts.append(types.Part.from_function_response(
|
||||
name=name, response={"output": obs},
|
||||
))
|
||||
contents.append(types.Content(role="user", parts=response_parts))
|
||||
|
||||
return {
|
||||
"tokens_in": tokens_in,
|
||||
"tokens_out": tokens_out,
|
||||
"turns": turns,
|
||||
"final_summary": final_text,
|
||||
"max_turns_hit": turns == max_turns and not final_text,
|
||||
}
|
||||
|
||||
|
||||
# ---------- Local loop (vLLM, OpenAI-compatible multi-turn with tools) ----------
|
||||
|
||||
_COMPACT_PROMPT = (
|
||||
"Summarize the SWE-bench agent trajectory so far in under 2000 characters. "
|
||||
"Preserve: filenames touched, hypotheses tested, what worked, what failed, "
|
||||
"and the current plan. Be terse — no preamble, no quoted output, just facts."
|
||||
)
|
||||
|
||||
_TIKTOKEN_ENC = None
|
||||
_TIKTOKEN_WARNED = False
|
||||
|
||||
|
||||
def _get_tiktoken_enc() -> Any:
|
||||
global _TIKTOKEN_ENC, _TIKTOKEN_WARNED
|
||||
if _TIKTOKEN_ENC is not None:
|
||||
return _TIKTOKEN_ENC
|
||||
try:
|
||||
import tiktoken
|
||||
_TIKTOKEN_ENC = tiktoken.get_encoding("cl100k_base")
|
||||
except Exception as exc:
|
||||
if not _TIKTOKEN_WARNED:
|
||||
print(f"[mini_swe_agent] tiktoken unavailable ({exc!r}); falling back to len(s)//4", flush=True)
|
||||
_TIKTOKEN_WARNED = True
|
||||
_TIKTOKEN_ENC = False
|
||||
return _TIKTOKEN_ENC
|
||||
|
||||
|
||||
def _estimate_prompt_tokens(messages: List[Dict[str, Any]]) -> int:
|
||||
enc = _get_tiktoken_enc()
|
||||
total = 0
|
||||
for m in messages:
|
||||
total += 4 # per-message overhead
|
||||
c = m.get("content")
|
||||
if isinstance(c, str):
|
||||
s = c
|
||||
elif isinstance(c, list):
|
||||
parts = []
|
||||
for block in c:
|
||||
if isinstance(block, dict):
|
||||
parts.append(str(block.get("content") or block.get("text") or ""))
|
||||
s = "\n".join(parts)
|
||||
else:
|
||||
s = ""
|
||||
for tc in (m.get("tool_calls") or []):
|
||||
try:
|
||||
s += "\n" + (tc["function"]["arguments"] or "")
|
||||
s += "\n" + (tc["function"].get("name") or "")
|
||||
except (KeyError, TypeError):
|
||||
pass
|
||||
tcid = m.get("tool_call_id")
|
||||
if tcid:
|
||||
s += "\n" + str(tcid)
|
||||
if enc:
|
||||
total += len(enc.encode(s, disallowed_special=()))
|
||||
else:
|
||||
total += len(s) // 4
|
||||
return total
|
||||
|
||||
|
||||
_EXIT_PATTERNS = (
|
||||
re.compile(r"exit_code\s*[=:]\s*(-?\d+)"),
|
||||
re.compile(r"returncode\s*[=:]\s*(-?\d+)"),
|
||||
re.compile(r"\bexit\s+(-?\d+)\b"),
|
||||
)
|
||||
|
||||
|
||||
def _parse_exit_code(content: Any) -> str:
|
||||
if not isinstance(content, str):
|
||||
return "?"
|
||||
for pat in _EXIT_PATTERNS:
|
||||
m = pat.search(content)
|
||||
if m:
|
||||
return m.group(1)
|
||||
return "?"
|
||||
|
||||
|
||||
def _identify_turns(messages: List[Dict[str, Any]]) -> List[Tuple[int, int]]:
|
||||
"""Return list of (start_idx, end_idx_exclusive) for each assistant+tools turn.
|
||||
|
||||
A turn = one assistant message (with or without tool_calls) plus any
|
||||
immediately-following tool messages. System + initial user are skipped.
|
||||
"""
|
||||
turns: List[Tuple[int, int]] = []
|
||||
i = 0
|
||||
n = len(messages)
|
||||
while i < n:
|
||||
role = messages[i].get("role")
|
||||
if role == "assistant":
|
||||
j = i + 1
|
||||
while j < n and messages[j].get("role") == "tool":
|
||||
j += 1
|
||||
turns.append((i, j))
|
||||
i = j
|
||||
else:
|
||||
i += 1
|
||||
return turns
|
||||
|
||||
|
||||
def _compact_local_messages(
|
||||
messages: List[Dict[str, Any]],
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
keep_last: int,
|
||||
trace_prefix: str,
|
||||
compact_at_tokens: int = 24_000,
|
||||
) -> List[Dict[str, Any]]:
|
||||
if len(messages) < 2:
|
||||
return messages
|
||||
system_msg = messages[0]
|
||||
initial_user = messages[1]
|
||||
|
||||
turns = _identify_turns(messages)
|
||||
if len(turns) <= keep_last:
|
||||
return messages
|
||||
|
||||
keep_turns = turns[-keep_last:]
|
||||
old_turns = turns[:-keep_last]
|
||||
keep_start = keep_turns[0][0]
|
||||
|
||||
# Stage 1: elide tool observations in old turns.
|
||||
before_tokens = _estimate_prompt_tokens(messages)
|
||||
new_messages: List[Dict[str, Any]] = list(messages)
|
||||
n_tool_elided = 0
|
||||
for (s, e) in old_turns:
|
||||
for k in range(s, e):
|
||||
m = new_messages[k]
|
||||
if m.get("role") != "tool":
|
||||
continue
|
||||
orig = m.get("content")
|
||||
if not isinstance(orig, str):
|
||||
continue
|
||||
n_chars = len(orig)
|
||||
if n_chars <= 200:
|
||||
continue
|
||||
exit_code = _parse_exit_code(orig)
|
||||
stub = f"[tool output elided: {n_chars} chars, exit={exit_code}]"
|
||||
new_messages[k] = {
|
||||
"role": "tool",
|
||||
"tool_call_id": m.get("tool_call_id"),
|
||||
"content": stub,
|
||||
}
|
||||
n_tool_elided += 1
|
||||
|
||||
after_stage1_tokens = _estimate_prompt_tokens(new_messages)
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_compact",
|
||||
"stage": "1",
|
||||
"msgs_before": len(messages),
|
||||
"msgs_after": len(new_messages),
|
||||
"before_tokens": before_tokens,
|
||||
"after_tokens": after_stage1_tokens,
|
||||
"n_tool_elided": n_tool_elided,
|
||||
"n_turns_folded": 0,
|
||||
"ts": time.time(),
|
||||
})
|
||||
|
||||
if after_stage1_tokens <= compact_at_tokens:
|
||||
return new_messages
|
||||
|
||||
# Stage 2: fold old turns into a single synthetic system summary.
|
||||
middle = new_messages[2:keep_start]
|
||||
tail = new_messages[keep_start:]
|
||||
if not middle:
|
||||
return new_messages
|
||||
|
||||
summary_input = [
|
||||
{"role": "system", "content": _COMPACT_PROMPT},
|
||||
{"role": "user", "content": json.dumps(
|
||||
[{"role": m.get("role"),
|
||||
"content": m.get("content") if isinstance(m.get("content"), str) else str(m.get("content"))[:4000]}
|
||||
for m in middle],
|
||||
default=str,
|
||||
)[:60_000]},
|
||||
]
|
||||
summary = ""
|
||||
try:
|
||||
if client is not None:
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=summary_input,
|
||||
temperature=0.0,
|
||||
max_tokens=1024,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
_bump_local_calls()
|
||||
summary = (resp.choices[0].message.content or "").strip()[:2000]
|
||||
except Exception as exc:
|
||||
summary = f"[compaction summary failed: {exc!r}; older turns dropped]"
|
||||
if not summary:
|
||||
summary = "[no summary produced; older turns dropped]"
|
||||
|
||||
n_turns_folded = len(old_turns)
|
||||
synthetic = {
|
||||
"role": "system",
|
||||
"content": f"[turns 1–{n_turns_folded} elided: {summary}]",
|
||||
}
|
||||
folded = [system_msg, initial_user, synthetic, *tail]
|
||||
after_stage2_tokens = _estimate_prompt_tokens(folded)
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_compact",
|
||||
"stage": "2",
|
||||
"msgs_before": len(new_messages),
|
||||
"msgs_after": len(folded),
|
||||
"before_tokens": after_stage1_tokens,
|
||||
"after_tokens": after_stage2_tokens,
|
||||
"n_tool_elided": n_tool_elided,
|
||||
"n_turns_folded": n_turns_folded,
|
||||
"summary_chars": len(summary),
|
||||
"ts": time.time(),
|
||||
})
|
||||
return folded
|
||||
|
||||
|
||||
def _loop_local(
|
||||
problem: str,
|
||||
workdir: Path,
|
||||
@@ -470,7 +1221,16 @@ def _loop_local(
|
||||
output_cap: int,
|
||||
turn_max_tokens: int,
|
||||
trace_prefix: str,
|
||||
compact_at_tokens: int = 22_000,
|
||||
compact_keep_last: int = 3,
|
||||
) -> Dict[str, Any]:
|
||||
# Qwen-27B has a 32k context. With ``max_tokens=turn_max_tokens`` reserved
|
||||
# for output (default 4096) plus ~1k for the bash tool schema + system
|
||||
# prompt + format overhead, the practical input ceiling is ~27k. We
|
||||
# compact at 22k so there's slack for one more tool result before the
|
||||
# next turn's pre-call check fires again. Earlier we used 24k + keep=4
|
||||
# but still saw 28k-input 400s on the n=100 SWE sweep (the keep window
|
||||
# alone routinely exceeded the budget once bash outputs piled up).
|
||||
from openai import OpenAI
|
||||
client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=600.0)
|
||||
|
||||
@@ -484,16 +1244,60 @@ def _loop_local(
|
||||
turns = 0
|
||||
for turn in range(1, max_turns + 1):
|
||||
turns = turn
|
||||
if compact_at_tokens > 0 and _estimate_prompt_tokens(messages) > compact_at_tokens:
|
||||
messages = _compact_local_messages(
|
||||
messages, client=client, model=model,
|
||||
keep_last=compact_keep_last, trace_prefix=trace_prefix,
|
||||
compact_at_tokens=compact_at_tokens,
|
||||
)
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=turn_max_tokens,
|
||||
tools=[BASH_TOOL_OPENAI],
|
||||
tool_choice="auto",
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
try:
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=turn_max_tokens,
|
||||
tools=[BASH_TOOL_OPENAI],
|
||||
tool_choice="auto",
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
except Exception as exc:
|
||||
# Emergency compaction on a context-length 400 from vLLM /
|
||||
# OpenAI ("maximum context length is N tokens"). Our pre-call
|
||||
# estimator can undercount when tool args / tool_call_ids /
|
||||
# template overhead spike, so the budget check missed and the
|
||||
# server walled the call. Compact aggressively (keep_last=1)
|
||||
# and retry once. Re-raise on anything else or on a second
|
||||
# failure — the runner records the row as errored.
|
||||
msg = str(exc)
|
||||
is_ctx = (
|
||||
"maximum context length" in msg
|
||||
or "context length" in msg.lower() and "exceed" in msg.lower()
|
||||
)
|
||||
if not is_ctx:
|
||||
raise
|
||||
_record_event({
|
||||
"kind": f"{trace_prefix}_emergency_compact",
|
||||
"turn": turn,
|
||||
"error": msg[:300],
|
||||
"tokens_before": _estimate_prompt_tokens(messages),
|
||||
"ts": time.time(),
|
||||
})
|
||||
messages = _compact_local_messages(
|
||||
messages, client=client, model=model,
|
||||
keep_last=1, trace_prefix=trace_prefix,
|
||||
compact_at_tokens=max(8_000, compact_at_tokens // 2),
|
||||
)
|
||||
resp = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=turn_max_tokens,
|
||||
tools=[BASH_TOOL_OPENAI],
|
||||
tool_choice="auto",
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
_bump_local_calls()
|
||||
latency = time.time() - t0
|
||||
u = resp.usage
|
||||
tokens_in += getattr(u, "prompt_tokens", 0) if u else 0
|
||||
@@ -518,10 +1322,17 @@ def _loop_local(
|
||||
"ts": time.time(),
|
||||
})
|
||||
|
||||
messages.append({
|
||||
# Match the OpenAI cloud branch: content="" (not None) when only
|
||||
# tool_calls are present; omit ``tool_calls`` entirely when there
|
||||
# are none (vs. setting it to None) so the message validates
|
||||
# against the strict OpenAI schema if it ever gets replayed by
|
||||
# the compactor's summarizer call.
|
||||
assistant_local_msg: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": text or None,
|
||||
"tool_calls": [
|
||||
"content": text or "",
|
||||
}
|
||||
if tool_calls:
|
||||
assistant_local_msg["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id, "type": "function",
|
||||
"function": {
|
||||
@@ -530,8 +1341,8 @@ def _loop_local(
|
||||
},
|
||||
}
|
||||
for tc in tool_calls
|
||||
] if tool_calls else None,
|
||||
})
|
||||
]
|
||||
messages.append(assistant_local_msg)
|
||||
|
||||
if not tool_calls:
|
||||
final_text = text.strip()
|
||||
|
||||
@@ -17,9 +17,9 @@ acc / $0.67 (vs $1.09).
|
||||
|
||||
Requires the ``minions`` library from
|
||||
https://github.com/HazyResearch/minions installed in the same env (e.g.
|
||||
``uv pip install -e /matx/u/aspark/hybrid-local-cloud-compute/external/minions``).
|
||||
Import is lazy — the agent class registers without ``minions`` available,
|
||||
and the import error only fires on ``run()``.
|
||||
``uv pip install -e path/to/minions``). Import is lazy — the agent class
|
||||
registers without ``minions`` available, and the import error only fires
|
||||
on ``run()``.
|
||||
|
||||
Compatibility patches applied at first ``run()`` (idempotent):
|
||||
|
||||
@@ -45,9 +45,13 @@ 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,
|
||||
build_web_search_tool,
|
||||
web_search_cfg,
|
||||
)
|
||||
from openjarvis.agents.hybrid._openai_retry import (
|
||||
patch_openai_globally as _patch_openai_globally,
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import NO_TEMP_PREFIXES
|
||||
from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop
|
||||
@@ -200,6 +204,13 @@ def _patch_anthropic_globally() -> None:
|
||||
|
||||
def make_patched(orig): # type: ignore[no-untyped-def]
|
||||
def patched(self, **kwargs): # type: ignore[no-untyped-def]
|
||||
# External Minions's AnthropicClient.chat passes
|
||||
# `cache_control={"type":"ephemeral"}` as a top-level kwarg
|
||||
# (clients/anthropic.py:207). Newer Anthropic SDKs reject that
|
||||
# — cache_control belongs on individual content blocks, not on
|
||||
# Messages.create itself. Strip it; we're not relying on the
|
||||
# ephemeral hint for correctness in these short minions turns.
|
||||
kwargs.pop("cache_control", None)
|
||||
model = kwargs.get("model", "")
|
||||
if model.startswith(NO_TEMP_PREFIXES):
|
||||
kwargs.pop("temperature", None)
|
||||
@@ -215,6 +226,96 @@ def _patch_anthropic_globally() -> None:
|
||||
cls.create = make_patched(orig) # type: ignore[assignment]
|
||||
|
||||
|
||||
def _patch_gemini_client_usage() -> None:
|
||||
"""Patch vendored ``GeminiClient.schat`` for Gemini 2.5 quirks.
|
||||
|
||||
Two issues in the upstream client:
|
||||
|
||||
1. ``response.usage_metadata.candidates_token_count`` is sometimes ``None``
|
||||
(empty / thinking-only responses on 2.5 Pro), and the upstream code
|
||||
does ``total_token_count - candidates_token_count`` raw → ``TypeError``.
|
||||
|
||||
2. ``response.text`` raises if the model only emitted a non-text part
|
||||
(e.g. safety block, thinking-only). We swallow it as empty.
|
||||
|
||||
Both fixes are idempotent and bypass the original ``schat`` body only
|
||||
on the value-extraction lines — the API call itself is unchanged.
|
||||
"""
|
||||
from minions.clients.gemini import GeminiClient # type: ignore[import-not-found]
|
||||
from minions.usage import Usage # type: ignore[import-not-found]
|
||||
|
||||
if getattr(GeminiClient.schat, "_hybrid_patched", False):
|
||||
return
|
||||
|
||||
_orig_schat = GeminiClient.schat
|
||||
|
||||
def _safe_int(x): # type: ignore[no-untyped-def]
|
||||
try:
|
||||
return int(x) if x is not None else 0
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
def patched_schat(self, messages, **kwargs): # type: ignore[no-untyped-def]
|
||||
# Mirror the upstream "native" branch by hand, but defensively.
|
||||
# Skip the OpenAI-compat branch — Minions paradigm never sets that.
|
||||
if self.use_openai_api:
|
||||
return _orig_schat(self, messages, **kwargs)
|
||||
if isinstance(messages, dict):
|
||||
messages = [messages]
|
||||
contents, system_instruction = self._format_content(messages)
|
||||
if not system_instruction:
|
||||
system_instruction = self.system_instruction
|
||||
tools = self._prepare_tools(messages=messages)
|
||||
config_kwargs = {
|
||||
"temperature": self.temperature,
|
||||
"max_output_tokens": self.max_tokens,
|
||||
}
|
||||
if self.thinking_budget is not None or self.thinking_level is not None:
|
||||
tc = {}
|
||||
if self.thinking_budget is not None:
|
||||
tc["thinking_budget"] = self.thinking_budget
|
||||
if self.thinking_level is not None:
|
||||
tc["thinking_level"] = self.thinking_level
|
||||
config_kwargs["thinking_config"] = self.types.ThinkingConfig(**tc)
|
||||
if tools:
|
||||
config_kwargs["tools"] = tools
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
config = self.types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
response = self.client.models.generate_content(
|
||||
model=self.model_name,
|
||||
contents=contents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Defensive text accessor — upstream `response.text` can raise when
|
||||
# the model only emitted a non-text part.
|
||||
try:
|
||||
text = response.text or ""
|
||||
except Exception:
|
||||
try:
|
||||
parts = response.candidates[0].content.parts or []
|
||||
text = "".join(getattr(p, "text", "") or "" for p in parts)
|
||||
except Exception:
|
||||
text = ""
|
||||
|
||||
um = getattr(response, "usage_metadata", None)
|
||||
total = _safe_int(getattr(um, "total_token_count", 0)) if um else 0
|
||||
comp = _safe_int(getattr(um, "candidates_token_count", 0)) if um else 0
|
||||
prompt = _safe_int(getattr(um, "prompt_token_count", 0)) if um else 0
|
||||
# Prefer the explicit prompt count if present; fall back to (total - comp).
|
||||
if not prompt and total:
|
||||
prompt = max(total - comp, 0)
|
||||
usage = Usage(prompt_tokens=prompt, completion_tokens=comp)
|
||||
|
||||
if self.local:
|
||||
return [text], usage, ["stop"]
|
||||
return [text], usage
|
||||
|
||||
patched_schat._hybrid_patched = True # type: ignore[attr-defined]
|
||||
GeminiClient.schat = patched_schat # type: ignore[assignment]
|
||||
|
||||
|
||||
def _patch_minions_extract_json() -> None:
|
||||
"""Minions's ``_extract_json`` uses a non-greedy regex that grabs the
|
||||
first short bracket pair and prefers ```json``` fences. With structured
|
||||
@@ -245,13 +346,23 @@ def _apply_patches_once() -> None:
|
||||
return
|
||||
_stub_missing_imports()
|
||||
_patch_anthropic_globally()
|
||||
# Mirror the Anthropic patch for OpenAI so the Minions library's own
|
||||
# ``OpenAIClient`` instances pick up retry + per-org concurrency caps.
|
||||
# Idempotent — also applied at ``_base`` import time.
|
||||
_patch_openai_globally()
|
||||
_patch_gemini_client_usage()
|
||||
_patch_minions_extract_json()
|
||||
_PATCHES_APPLIED = True
|
||||
|
||||
|
||||
# ---------- Pre-fetch helper (GAIA only) ----------
|
||||
|
||||
def _prefetch_context(question: str, cloud_endpoint: str, cloud_model: str) -> Dict[str, Any]:
|
||||
def _prefetch_context(
|
||||
question: str,
|
||||
cloud_endpoint: str,
|
||||
cloud_model: str,
|
||||
max_uses: int = 8,
|
||||
) -> Dict[str, Any]:
|
||||
"""Use Anthropic web_search to fetch real source material the worker can read.
|
||||
|
||||
Minions's premise is "worker reads a doc, asks cloud for help" — but GAIA
|
||||
@@ -278,7 +389,7 @@ def _prefetch_context(question: str, cloud_endpoint: str, cloud_model: str) -> D
|
||||
cloud_model,
|
||||
user=prompt,
|
||||
max_tokens=8192,
|
||||
tools=[ANTHROPIC_WEB_SEARCH_TOOL],
|
||||
tools=[build_web_search_tool(max_uses)],
|
||||
tool_choice={"type": "any"},
|
||||
)
|
||||
from openjarvis.agents.hybrid._prices import cost as _cost_usd
|
||||
@@ -362,6 +473,9 @@ class MinionsAgent(LocalCloudAgent):
|
||||
from minions.clients.anthropic import (
|
||||
AnthropicClient, # type: ignore[import-not-found]
|
||||
)
|
||||
from minions.clients.gemini import (
|
||||
GeminiClient, # type: ignore[import-not-found]
|
||||
)
|
||||
from minions.clients.openai import (
|
||||
OpenAIClient, # type: ignore[import-not-found]
|
||||
)
|
||||
@@ -397,6 +511,17 @@ class MinionsAgent(LocalCloudAgent):
|
||||
temperature=0.0,
|
||||
max_tokens=4096,
|
||||
)
|
||||
elif self._cloud_endpoint == "gemini":
|
||||
# The vendored Minion library already special-cases GeminiClient
|
||||
# in minion.py: it passes response_mime_type=application/json plus
|
||||
# a Pydantic response_schema so the supervisor reply parses with
|
||||
# the same {decision, message, answer} shape Opus/GPT use. We just
|
||||
# have to hand it a GeminiClient instance — no extra plumbing.
|
||||
cloud_client = GeminiClient(
|
||||
model_name=self._cloud_model,
|
||||
temperature=0.0,
|
||||
max_tokens=4096,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported cloud endpoint: {self._cloud_endpoint!r}")
|
||||
|
||||
@@ -412,12 +537,29 @@ class MinionsAgent(LocalCloudAgent):
|
||||
# GAIA-shape only: prefetch a web_search digest so the worker has
|
||||
# something real to read. SWE-bench (problem_statement only) already
|
||||
# ships its own doc.
|
||||
#
|
||||
# Honors the new ``method_cfg.web_search`` schema:
|
||||
# - omitted → prefetch ON (legacy default for minions GAIA)
|
||||
# - enabled = false → prefetch OFF
|
||||
# - enabled = true → prefetch ON (honors max_uses)
|
||||
prefetch: Dict[str, Any] = {
|
||||
"text": "", "tokens": 0, "cost_usd": 0.0, "n_searches": 0,
|
||||
}
|
||||
if task_meta.get("question"):
|
||||
ws_block = cfg.get("web_search") if isinstance(cfg.get("web_search"), dict) else None
|
||||
ws_enabled, ws_max_uses = web_search_cfg(cfg)
|
||||
# If the cell explicitly set web_search.enabled = false, honor that.
|
||||
# If it set web_search.enabled = true, honor max_uses. If it didn't
|
||||
# set a web_search block at all, keep legacy prefetch ON.
|
||||
prefetch_on = (
|
||||
ws_block is None # legacy default
|
||||
or ws_enabled
|
||||
)
|
||||
if task_meta.get("question") and prefetch_on:
|
||||
prefetch = _prefetch_context(
|
||||
task_meta["question"], self._cloud_endpoint, self._cloud_model
|
||||
task_meta["question"],
|
||||
self._cloud_endpoint,
|
||||
self._cloud_model,
|
||||
max_uses=ws_max_uses,
|
||||
)
|
||||
|
||||
if prefetch.get("text"):
|
||||
@@ -462,6 +604,10 @@ class MinionsAgent(LocalCloudAgent):
|
||||
"tokens_cloud": (rp + rc) + prefetch["tokens"],
|
||||
"cost_usd": self.cost_usd(self._cloud_model, rp, rc) + prefetch["cost_usd"],
|
||||
"turns": cfg.get("max_rounds", 3),
|
||||
"web_search_uses": prefetch["n_searches"],
|
||||
# GAIA: only countable tool surface is the prefetch web_search.
|
||||
# The Minions protocol itself is supervisor↔worker text, no tools.
|
||||
"tool_calls": int(prefetch["n_searches"]),
|
||||
"traces": {
|
||||
"mode": mode,
|
||||
"supervisor_messages": out.get("supervisor_messages"),
|
||||
@@ -540,6 +686,8 @@ class MinionsAgent(LocalCloudAgent):
|
||||
"tokens_cloud": p_in + p_out,
|
||||
"cost_usd": supervisor_cost,
|
||||
"turns": 1 + out["turns"],
|
||||
# SWE: only the worker invokes tools (bash); supervisor is text-only.
|
||||
"tool_calls": int(out["turns"]),
|
||||
"traces": {
|
||||
"swe_mode": True,
|
||||
"supervisor_plan": plan_text,
|
||||
|
||||
@@ -11,7 +11,7 @@ bench = "gaia"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-9B", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { web_search = { enabled = true, max_uses = 8 } }
|
||||
|
||||
[cells.advisors-swebenchverified-qwen9b-opus-3]
|
||||
method = "advisors"
|
||||
@@ -28,7 +28,7 @@ bench = "gaia"
|
||||
n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { web_search = { enabled = true, max_uses = 8 } }
|
||||
concurrency = 4
|
||||
|
||||
[cells.advisors-swebenchverified-qwen9b-opus-30]
|
||||
|
||||
@@ -25,7 +25,7 @@ bench = "gaia"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { architecture = "single_local", max_tokens = 1024 }
|
||||
method_cfg = { architecture = "single_local", max_tokens = 8192 }
|
||||
|
||||
[cells.archon-swebenchverified-qwen27b-opus-ensemble-K3-3]
|
||||
method = "archon"
|
||||
@@ -54,7 +54,7 @@ bench = "gaia"
|
||||
n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { architecture = "single_local", max_tokens = 2048 }
|
||||
method_cfg = { architecture = "single_local", max_tokens = 8192 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.archon-gaia-gemma31b-opus-ensemble-K5-30]
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# n=100 ablation — local-only baseline (no cloud teacher / router).
|
||||
#
|
||||
# Paradigm is `baseline_local` (see baseline_local.py). On GAIA the agent
|
||||
# makes one local vLLM call with the FINAL-ANSWER-formatted prompt. On
|
||||
# SWE-bench-Verified the agent runs the mini-SWE bash agent loop with
|
||||
# `backbone="local"` so the local model drives turns directly.
|
||||
#
|
||||
# Cell `cloud` is set to a dummy model purely so the hybrid runner doesn't
|
||||
# choke on missing fields; `baseline_local` ignores `cloud_*` entirely
|
||||
# and cost_usd is always 0.0 (local inference is free).
|
||||
|
||||
# ============================================================
|
||||
# :8004 — Gemma-4-31B (TP=4 on GPUs 0-3)
|
||||
# ============================================================
|
||||
|
||||
[cells.baseline-local-gemma31b-gaia-n100]
|
||||
method = "baseline_local"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.baseline-local-gemma31b-swe-n100]
|
||||
method = "baseline_local"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
|
||||
# ============================================================
|
||||
# :8002 — Qwen-3.6-27B-FP8 (PP=2 on GPUs 5-6)
|
||||
# ============================================================
|
||||
|
||||
[cells.baseline-local-qwen36-gaia-n100]
|
||||
method = "baseline_local"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.6-27B-FP8", endpoint = "http://localhost:8002/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.baseline-local-qwen36-swe-n100]
|
||||
method = "baseline_local"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.6-27B-FP8", endpoint = "http://localhost:8002/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
|
||||
# ============================================================
|
||||
# :8001 — Qwen-3.5-27B-FP8 (GPU 7) — the Qwen-3.5 solo floor
|
||||
# ============================================================
|
||||
|
||||
[cells.baseline-local-qwen27b-gaia-n100]
|
||||
method = "baseline_local"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.baseline-local-qwen27b-swe-n100]
|
||||
method = "baseline_local"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { local_max_tokens = 4096, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 4
|
||||
@@ -1,9 +1,15 @@
|
||||
# Conductor cells. Inference-only repro of arXiv 2512.04388 (Sakana, 2025).
|
||||
#
|
||||
# Stage-1 substitutes the (untrained) Qwen2.5-7B conductor with a zero-shot
|
||||
# frontier planner (Opus). Workers default to local Qwen3.5-27B + Opus +
|
||||
# gpt-5-mini; the adapter auto-drops the local worker if vLLM is down.
|
||||
# Override `method_cfg.workers` to customize the pool.
|
||||
# frontier planner (Opus). Workers come from `method_cfg.worker_pool` — an
|
||||
# explicit, deterministic pool. The n100 cells below pin a 3-worker pool
|
||||
# (0=local, 1=Opus, 2=gpt-5-mini) via `$local` / `$cloud` substitution so
|
||||
# the named local + cloud models are actually exercised.
|
||||
#
|
||||
# DO NOT rely on the implicit default pool: `_default_pool` is the paper's
|
||||
# 7-worker cloud-only pool (no local, no Opus). Cells that want hybrid
|
||||
# local+cloud MUST set `worker_pool` explicitly (race-free, no per-task
|
||||
# vLLM probe).
|
||||
#
|
||||
# Naming: conductor-<bench>-<conductor-short>-<N>
|
||||
|
||||
@@ -52,3 +58,48 @@ local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1"
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { worker_max_tokens = 8192, worker_temperature = 0.2, conductor_max_tokens = 4096 }
|
||||
concurrency = 8
|
||||
|
||||
# ============================================================
|
||||
# n=100 ablation cells (added 2026-05-18)
|
||||
# ============================================================
|
||||
|
||||
# Deterministic 3-worker pool: 0=local (Qwen3.5-27B via vLLM), 1=Opus 4.7
|
||||
# (cloud), 2=gpt-5-mini. `$local` / `$cloud` resolve against the cell's
|
||||
# `local` / `cloud` blocks. GAIA cells enable web_search so the search-
|
||||
# capable workers (anthropic / openai) can ground their answers.
|
||||
[cells.conductor-qwen27b-opus47-gaia-n100]
|
||||
method = "conductor"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
concurrency = 4
|
||||
[cells.conductor-qwen27b-opus47-gaia-n100.method_cfg]
|
||||
worker_max_tokens = 4096
|
||||
worker_temperature = 0.2
|
||||
conductor_max_tokens = 2048
|
||||
web_search = { enabled = true, max_uses = 8 }
|
||||
worker_pool = [
|
||||
{ id = 0, name = "local-qwen", endpoint = "vllm", model = "$local", description = "Open-weights Qwen3.5-27B served locally via vLLM. Cheap and fast. Good at concise extraction, formatting, and arithmetic on given data; weaker at open-domain factual recall and complex multi-step reasoning. Cannot perform live web search." },
|
||||
{ id = 1, name = "frontier-anthropic", endpoint = "anthropic", model = "$cloud", description = "Anthropic Claude Opus 4.7. Frontier reasoning model — strongest at multi-step reasoning, careful instruction following, code, and writing. Can perform live web search. Expensive; use for hard or decisive steps." },
|
||||
{ id = 2, name = "frontier-openai-mini", endpoint = "openai", model = "gpt-5-mini", description = "OpenAI gpt-5-mini. Mid-tier model with solid general knowledge and reasoning at a fraction of frontier cost. Can perform live web search. Good default for retrieval-style or broad-knowledge questions." },
|
||||
]
|
||||
|
||||
[cells.conductor-qwen27b-opus47-swe-n100]
|
||||
method = "conductor"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
concurrency = 4
|
||||
[cells.conductor-qwen27b-opus47-swe-n100.method_cfg]
|
||||
worker_max_tokens = 8192
|
||||
worker_temperature = 0.2
|
||||
conductor_max_tokens = 4096
|
||||
worker_pool = [
|
||||
{ id = 0, name = "local-qwen", endpoint = "vllm", model = "$local", description = "Open-weights Qwen3.5-27B served locally via vLLM. Cheap and fast. Good at concise extraction, formatting, and arithmetic; weaker at complex multi-step reasoning and large code patches." },
|
||||
{ id = 1, name = "frontier-anthropic", endpoint = "anthropic", model = "$cloud", description = "Anthropic Claude Opus 4.7. Frontier reasoning model — strongest at multi-step reasoning, careful instruction following, and code. Writes tight, well-formed unified diffs. Expensive; use for hard or decisive steps." },
|
||||
{ id = 2, name = "frontier-openai-mini", endpoint = "openai", model = "gpt-5-mini", description = "OpenAI gpt-5-mini. Mid-tier model with solid general knowledge and reasoning at a fraction of frontier cost. Good default for broad-knowledge or straightforward code steps." },
|
||||
]
|
||||
|
||||
@@ -1,15 +1,27 @@
|
||||
# SkillOrchestra cells (arXiv:2602.19672).
|
||||
#
|
||||
# Deployment-time skill-aware routing only. The paper's offline pipeline
|
||||
# (explore → learn → select over an SGLang model pool + FAISS wiki index)
|
||||
# isn't reproduced here — we don't have the SGLang serving stack, the
|
||||
# retriever index, or a training split to learn a routing policy on. What
|
||||
# we DO reproduce is the inference-time orchestrator: an Opus router
|
||||
# analyzes the question into skill weights, picks one of two agents
|
||||
# (local Qwen-27B vs cloud Opus) under an explicit cost trade-off, then
|
||||
# the chosen agent answers. Mirrors the eval_orchestrator paradigm in
|
||||
# external/SkillOrchestra/skillorchestra/prompts/eval_orchestrator.py.
|
||||
# See adapters/skillorchestra_adapter.py for full notes.
|
||||
# Faithful port of the eval-orchestrator runtime — the multi-round
|
||||
# search -> reasoning -> answer ReAct loop, verbatim eval_orchestrator
|
||||
# prompts, the StageSkillHandbook + 5 routing strategies, real Python
|
||||
# subprocess execution, and a model-alias pool. See the package
|
||||
# README at agents/hybrid/skillorchestra/README.md.
|
||||
#
|
||||
# method_cfg knobs (all optional):
|
||||
# routing_strategy none | router_decides | analyze_model_decide
|
||||
# | weighted_avg | weakest_skill | strongest_skill
|
||||
# handbook_path StageSkillHandbook JSON; relative -> resolved inside
|
||||
# the skillorchestra package. weighted/weakest/strongest
|
||||
# need one; router_decides / none do not.
|
||||
# max_rounds orchestrator loop cap (default 6).
|
||||
# retriever_url FAISS wiki retriever; absent -> Anthropic web_search.
|
||||
# model_pool per-alias { model, endpoint } overrides.
|
||||
# orchestrator_model / orchestrator_endpoint pin the orchestrator LLM.
|
||||
#
|
||||
# GAIA cells run the full skill-routing loop against the hand-authored
|
||||
# seed handbook (handbook_seed.json — NOT learned; swap in a learned one
|
||||
# when the explore->learn->select pipeline has run). SWE-bench is out of
|
||||
# scope for the original QA orchestrator: those cells run the cloud
|
||||
# backbone through the shared mini SWE agent loop.
|
||||
|
||||
[cells.skillorchestra-gaia-qwen27b-opus-3]
|
||||
method = "skillorchestra"
|
||||
@@ -17,7 +29,7 @@ bench = "gaia"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { routing_strategy = "weighted_avg", handbook_path = "handbook_seed.json", max_rounds = 6 }
|
||||
|
||||
[cells.skillorchestra-swebenchverified-qwen27b-opus-3]
|
||||
method = "skillorchestra"
|
||||
@@ -25,7 +37,7 @@ bench = "swebench-verified"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { swe_use_agent_loop = true, swe_max_turns = 50, swe_bash_timeout_s = 120 }
|
||||
concurrency = 3
|
||||
|
||||
[cells.skillorchestra-gaia-qwen27b-opus-30]
|
||||
@@ -34,7 +46,7 @@ bench = "gaia"
|
||||
n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { routing_strategy = "weighted_avg", handbook_path = "handbook_seed.json", max_rounds = 6 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.skillorchestra-swebenchverified-qwen27b-opus-30]
|
||||
@@ -43,5 +55,5 @@ bench = "swebench-verified"
|
||||
n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = {}
|
||||
method_cfg = { swe_use_agent_loop = true, swe_max_turns = 50, swe_bash_timeout_s = 120 }
|
||||
concurrency = 3
|
||||
|
||||
@@ -35,7 +35,7 @@ method_cfg = { swe_use_agent_loop = true, conductor_max_tokens = 2048, swe_max_t
|
||||
method = "advisors"
|
||||
bench = "swebench-verified"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-9B", endpoint = "http://localhost:8001/v1" }
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { swe_use_agent_loop = true, swe_max_turns = 25, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096, advisor_max_tokens = 2048 }
|
||||
|
||||
@@ -45,7 +45,7 @@ method_cfg = { swe_use_agent_loop = true, swe_max_turns = 25, swe_bash_timeout_s
|
||||
method = "skillorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8002/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { swe_use_agent_loop = true, router_max_tokens = 1024, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
|
||||
@@ -55,7 +55,7 @@ method_cfg = { swe_use_agent_loop = true, router_max_tokens = 1024, swe_max_turn
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8002/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { swe_use_agent_loop = true, max_turns = 4, orchestrator_max_tokens = 1024, swe_max_turns = 25, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
|
||||
@@ -65,6 +65,6 @@ method_cfg = { swe_use_agent_loop = true, max_turns = 4, orchestrator_max_tokens
|
||||
method = "archon"
|
||||
bench = "swebench-verified"
|
||||
n = 3
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8002/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { swe_use_agent_loop = true, n_samples = 3, swe_max_turns = 25, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096, ranker_max_tokens = 1024 }
|
||||
|
||||
@@ -1,9 +1,29 @@
|
||||
# ToolOrchestra cells. Prompted port — uses a cloud model (Opus) as the
|
||||
# orchestrator, NOT the RL-trained Nemotron-Orchestrator-8B from the paper
|
||||
# (arXiv:2511.21689). Treat results as preliminary until a real
|
||||
# Orchestrator-8B deployment is wired up.
|
||||
# ToolOrchestra cells (arXiv:2511.21689).
|
||||
#
|
||||
# Naming: toolorchestra-<bench>-<local-short>-<cloud-short>-<N>
|
||||
# Two modes via `method_cfg.orchestrator_mode`:
|
||||
#
|
||||
# "prompted" (default) -- cloud model (Opus etc.) plays the orchestrator
|
||||
# over a heterogeneous worker pool. Useful as a
|
||||
# prompted upper-bound; NOT what the paper does.
|
||||
#
|
||||
# "rl" -- paper-faithful path. The RL-trained
|
||||
# `nvidia/Orchestrator-8B` served on a local vLLM
|
||||
# (default :8003) is the orchestrator and emits
|
||||
# OpenAI-style tool_calls for the three NVlabs
|
||||
# tools (enhance_reasoning / answer / search).
|
||||
# Expert pool maps tool-`model` slots to our
|
||||
# workers: *-1 -> cloud, *-2 -> gpt-5-mini,
|
||||
# *-3 -> local vLLM, search -> Anthropic web_search.
|
||||
# See `toolorchestra.py` module docstring for
|
||||
# what's collapsed vs. the upstream (no Tavily /
|
||||
# FAISS-wiki / Qwen-Coder code-interpreter).
|
||||
#
|
||||
# Naming: toolorchestra-<orch-short>-<cloud-short>-<bench>-<N>
|
||||
# (legacy prompted cells keep their original naming scheme.)
|
||||
|
||||
# ============================================================
|
||||
# Legacy prompted-mode cells (cloud-as-orchestrator).
|
||||
# ============================================================
|
||||
|
||||
[cells.toolorchestra-gaia-qwen27b-opus-3]
|
||||
method = "toolorchestra"
|
||||
@@ -28,3 +48,182 @@ n = 30
|
||||
local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { max_turns = 8, orchestrator_max_tokens = 1024, worker_max_tokens = 4096 }
|
||||
|
||||
# ============================================================
|
||||
# RL-mode cells -- Orchestrator-8B drives the loop (n=100).
|
||||
# ============================================================
|
||||
#
|
||||
# `local` carries the orchestrator's vLLM (nvidia/Orchestrator-8B on :8003).
|
||||
# Tier-3 expert calls (`answer-3`, `reasoner-3`, …) route to this same model.
|
||||
# `cloud` is the frontier worker invoked for tier-1 slots.
|
||||
# `gpt-5-mini` is hard-coded as the tier-2 mid worker (matches the paper's
|
||||
# `gpt-5-mini` slots in `eval_hle.py` MODEL_MAPPING).
|
||||
|
||||
[cells.toolorchestra-orch8b-opus47-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 4
|
||||
|
||||
[cells.toolorchestra-orch8b-opus47-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 3
|
||||
|
||||
# ============================================================
|
||||
# Paper-match pool smoke (2026-05-19). n=5 GAIA. Opts into the
|
||||
# closer-to-paper worker pool via `method_cfg.pool = "paper"`:
|
||||
# - search -> Tavily API (TAVILY_API_KEY in OpenJarvis/.env)
|
||||
# - enhance_.. -> Qwen2.5-Coder-32B (OpenRouter) + Modal Python sandbox
|
||||
# - answer-1 -> GPT-5 answer-math-* -> Qwen2.5-Coder-32B
|
||||
# - answer-2 -> GPT-5-mini answer-3 -> Llama-3.3-70B (OR)
|
||||
# - answer-4 -> local Qwen (the Orchestrator-8B endpoint at :8003)
|
||||
# Skipped vs. paper: FAISS RAG, Qwen2.5-Math-{72B,7B} (not on OpenRouter).
|
||||
[cells.toolorchestra-papermatch-orch8b-opus47-gaia-n5]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 5
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2, pool = "paper", modal_python_timeout_s = 60, tavily_max_results = 5 }
|
||||
concurrency = 1
|
||||
|
||||
# Smoke cell for the RL-mode SWE wiring (2026-05-19). n=2, no subset so it
|
||||
# just grabs the first two SWE-bench Verified tasks. Used to verify
|
||||
# end-to-end: workdir clone, _swe_call_worker dispatch through Orchestrator-8B,
|
||||
# diff extraction, Modal-harness scoring. Promote to n=100 once green.
|
||||
[cells.toolorchestra-orch8b-opus47-swe-smoke2]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 2
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-opus-4-7", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 6, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 15, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 1
|
||||
|
||||
# ============================================================
|
||||
# ToolOrchestra × cloud-worker ablation (n=100, 2026-05-19)
|
||||
# Cloud-axis sweep: Haiku 4.5, GPT-5, GPT-5 mini, Gemini 2.5 Pro,
|
||||
# Gemini 2.5 Flash. Mirrors the skillorchestra cloud-axis cells.
|
||||
# All share Orchestrator-8B on :8003 as the local orchestrator.
|
||||
# concurrency=3 for Anthropic/Google, 2 for OpenAI (prepaid-quota wall).
|
||||
# ============================================================
|
||||
|
||||
# ---------- Anthropic Haiku 4.5 ----------
|
||||
|
||||
[cells.toolorchestra-orch8b-haiku45-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-haiku-4-5", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 3
|
||||
|
||||
[cells.toolorchestra-orch8b-haiku45-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "claude-haiku-4-5", endpoint = "anthropic" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 3
|
||||
|
||||
# ---------- OpenAI GPT-5 ----------
|
||||
|
||||
[cells.toolorchestra-orch8b-gpt5-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gpt-5", endpoint = "openai" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 2
|
||||
|
||||
[cells.toolorchestra-orch8b-gpt5-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gpt-5", endpoint = "openai" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 2
|
||||
|
||||
# ---------- OpenAI GPT-5 mini ----------
|
||||
|
||||
[cells.toolorchestra-orch8b-gpt5mini-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gpt-5-mini", endpoint = "openai" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 2
|
||||
|
||||
[cells.toolorchestra-orch8b-gpt5mini-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gpt-5-mini", endpoint = "openai" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 2
|
||||
|
||||
# ---------- Google Gemini 2.5 Pro ----------
|
||||
|
||||
[cells.toolorchestra-orch8b-gemini25pro-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gemini-2.5-pro", endpoint = "gemini" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 3
|
||||
|
||||
[cells.toolorchestra-orch8b-gemini25pro-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gemini-2.5-pro", endpoint = "gemini" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 3
|
||||
|
||||
# ---------- Google Gemini 2.5 Flash ----------
|
||||
|
||||
[cells.toolorchestra-orch8b-gemini25flash-gaia-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "gaia"
|
||||
n = 100
|
||||
subset = "gaia_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gemini-2.5-flash", endpoint = "gemini" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 4096, worker_temperature = 0.2 }
|
||||
concurrency = 3
|
||||
|
||||
[cells.toolorchestra-orch8b-gemini25flash-swe-n100]
|
||||
method = "toolorchestra"
|
||||
bench = "swebench-verified"
|
||||
n = 100
|
||||
subset = "swebench_verified_n100_seed42.json"
|
||||
local = { model = "nvidia/Orchestrator-8B", endpoint = "http://localhost:8003/v1" }
|
||||
cloud = { model = "gemini-2.5-flash", endpoint = "gemini" }
|
||||
method_cfg = { orchestrator_mode = "rl", orchestrator_endpoint = "http://localhost:8003/v1", orchestrator_model = "nvidia/Orchestrator-8B", max_turns = 8, orchestrator_max_tokens = 4096, worker_max_tokens = 8192, worker_temperature = 0.2, swe_use_agent_loop = true, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 }
|
||||
concurrency = 3
|
||||
|
||||
@@ -8,7 +8,7 @@ Reads a cell definition from ``registry/<method>.toml`` (bundled with this
|
||||
package or pointed at by ``OPENJARVIS_HYBRID_REGISTRY_DIR``), constructs
|
||||
the registered agent, loads bench tasks via OpenJarvis's existing dataset
|
||||
providers, runs every task, scores it, and writes
|
||||
``<EXPERIMENTS_DIR>/<cell>/results.jsonl`` + ``summary.json``.
|
||||
``<EXPERIMENTS_DIR>/runs/<cell>/results.jsonl`` + ``summary.json``.
|
||||
|
||||
The output schema matches ``hybrid-local-cloud-compute/runner.py`` so the
|
||||
existing rescore / dashboard scripts can read OpenJarvis cells without
|
||||
@@ -26,6 +26,7 @@ import threading
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from concurrent.futures import TimeoutError as FuturesTimeoutError
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
@@ -36,6 +37,7 @@ except ModuleNotFoundError:
|
||||
import tomli as tomllib # type: ignore[import-not-found,no-redef]
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext, AgentResult
|
||||
from openjarvis.agents.hybrid._energy import EnergyCollector
|
||||
from openjarvis.agents.hybrid._prompts import format_prompt as _format_prompt
|
||||
|
||||
PACKAGE_DIR = Path(__file__).parent
|
||||
@@ -43,13 +45,61 @@ DEFAULT_REGISTRY_DIR = PACKAGE_DIR / "registry"
|
||||
DEFAULT_EXPERIMENTS_DIR = Path(
|
||||
os.environ.get(
|
||||
"OPENJARVIS_HYBRID_EXPERIMENTS_DIR",
|
||||
str(Path.home() / ".openjarvis-hybrid" / "experiments"),
|
||||
Path.home() / ".openjarvis" / "experiments" / "hybrid",
|
||||
)
|
||||
)
|
||||
DEFAULT_SUBSETS_DIR = DEFAULT_EXPERIMENTS_DIR / "subsets"
|
||||
DEFAULT_RUNS_DIR = DEFAULT_EXPERIMENTS_DIR / "runs"
|
||||
|
||||
# Hard per-task wall-clock cap. Even when every individual network /
|
||||
# subprocess call has its own timeout, a pathological chain (SDK retries
|
||||
# stacking on top of a hung connection, a Modal harness subprocess whose
|
||||
# grandchildren keep its stdout pipe open, etc.) can leave one task
|
||||
# blocked indefinitely — and with the runner's ThreadPoolExecutor that
|
||||
# wedges the whole cell (`as_completed` never advances). SWE-bench tasks
|
||||
# legitimately run ~15-20 min, so the default cap is 30 min: long enough
|
||||
# never to abort a healthy task, short enough that a frozen one is
|
||||
# abandoned and recorded as an error row (which the resume logic re-runs)
|
||||
# instead of silently killing the process. Override with
|
||||
# ``OPENJARVIS_HYBRID_TASK_TIMEOUT_S`` (0 / negative disables).
|
||||
DEFAULT_TASK_TIMEOUT_S = float(
|
||||
os.environ.get("OPENJARVIS_HYBRID_TASK_TIMEOUT_S", "1800") or 1800
|
||||
)
|
||||
|
||||
|
||||
# ---------- Registry ----------
|
||||
|
||||
_SWE_BENCHES = {"swebench-verified", "swebench_verified", "swebench"}
|
||||
|
||||
|
||||
def _validate_cells(cells: Dict[str, Dict[str, Any]]) -> None:
|
||||
"""Catch registry mistakes that would silently degrade behaviour.
|
||||
|
||||
Currently: skillorchestra on a SWE bench MUST have
|
||||
``method_cfg.swe_use_agent_loop = true``. Without the flag,
|
||||
skillorchestra.py falls back to a one-shot cloud call even for
|
||||
SWE-bench tasks, which is rarely what the experimenter wants and is
|
||||
invisible at runtime (Bug 5, 2026-05-15).
|
||||
"""
|
||||
bad: List[str] = []
|
||||
for name, cell in cells.items():
|
||||
if cell.get("method") != "skillorchestra":
|
||||
continue
|
||||
if cell.get("bench") not in _SWE_BENCHES:
|
||||
continue
|
||||
mcfg = cell.get("method_cfg") or {}
|
||||
if not bool(mcfg.get("swe_use_agent_loop")):
|
||||
bad.append(name)
|
||||
if bad:
|
||||
raise ValueError(
|
||||
"skillorchestra SWE cells missing required "
|
||||
"`method_cfg.swe_use_agent_loop = true`: "
|
||||
+ ", ".join(sorted(bad))
|
||||
+ ". Without this flag the cell silently falls back to a "
|
||||
"one-shot cloud call for SWE-bench tasks."
|
||||
)
|
||||
|
||||
|
||||
def load_registry(registry_dir: Optional[Path] = None) -> Dict[str, Dict[str, Any]]:
|
||||
"""Merge every ``<registry_dir>/*.toml``. Cell names must be unique."""
|
||||
base = registry_dir or DEFAULT_REGISTRY_DIR
|
||||
@@ -67,6 +117,7 @@ def load_registry(registry_dir: Optional[Path] = None) -> Dict[str, Dict[str, An
|
||||
f"duplicate cell {name!r} (already defined before {p.name})"
|
||||
)
|
||||
cells[name] = cell
|
||||
_validate_cells(cells)
|
||||
return cells
|
||||
|
||||
|
||||
@@ -81,12 +132,17 @@ def _load_gaia_tasks(n: Optional[int]) -> List[Dict[str, Any]]:
|
||||
out: List[Dict[str, Any]] = []
|
||||
for rec in ds.iter_records():
|
||||
# rec.problem is the formatted question prompt; rec.metadata carries
|
||||
# the GAIA-specific fields including any reference answer.
|
||||
# the GAIA-specific fields including any reference answer. Prefer the
|
||||
# upstream GAIA `task_id` field (bare uuid) over rec.record_id (which
|
||||
# OpenJarvis prefixes with `gaia-`) so subsets keyed by the upstream
|
||||
# id round-trip.
|
||||
md = rec.metadata or {}
|
||||
task_id = md.get("task_id") or rec.record_id
|
||||
out.append({
|
||||
"task_id": rec.record_id,
|
||||
"question": rec.metadata.get("question", rec.problem),
|
||||
"task_id": task_id,
|
||||
"question": md.get("question", rec.problem),
|
||||
"reference": rec.reference,
|
||||
"metadata": dict(rec.metadata),
|
||||
"metadata": dict(md),
|
||||
})
|
||||
return out
|
||||
|
||||
@@ -95,7 +151,7 @@ def _load_swebench_tasks(n: Optional[int]) -> List[Dict[str, Any]]:
|
||||
"""SWE-bench-Verified test. Each task carries patch-evaluation fields."""
|
||||
from openjarvis.evals.datasets.swebench import SWEBenchDataset
|
||||
|
||||
ds = SWEBenchDataset()
|
||||
ds = SWEBenchDataset(variant="verified")
|
||||
ds.load(max_samples=n)
|
||||
out: List[Dict[str, Any]] = []
|
||||
for rec in ds.iter_records():
|
||||
@@ -124,36 +180,151 @@ def load_tasks(bench: str, n: Optional[int]) -> List[Dict[str, Any]]:
|
||||
raise ValueError(f"unknown bench: {bench!r}")
|
||||
|
||||
|
||||
def _load_subset_file(subset_path: str) -> Dict[str, Any]:
|
||||
"""Resolve a cell's ``subset`` field to a parsed JSON dict.
|
||||
|
||||
Resolution order:
|
||||
1. Absolute path → use as-is.
|
||||
2. Bare filename / relative path → look up under
|
||||
``<experiments>/subsets/`` (matches where ``make_subset.py``
|
||||
writes its output).
|
||||
|
||||
Accepts both list-of-ids and dict-with-task_ids shapes; the legacy
|
||||
harness wrote the dict shape and we preserve that. Returns a dict
|
||||
with at least a ``task_ids`` list so callers don't have to branch.
|
||||
"""
|
||||
p = Path(subset_path)
|
||||
if not p.is_absolute():
|
||||
p = DEFAULT_SUBSETS_DIR / p
|
||||
if not p.exists():
|
||||
raise FileNotFoundError(f"subset file not found: {p}")
|
||||
data = json.loads(p.read_text())
|
||||
if isinstance(data, list):
|
||||
return {"task_ids": list(data)}
|
||||
if isinstance(data, dict):
|
||||
if "task_ids" not in data:
|
||||
raise ValueError(
|
||||
f"subset {p.name} has no 'task_ids' field; got keys {list(data.keys())}"
|
||||
)
|
||||
return data
|
||||
raise ValueError(f"subset {p.name} must be a list or dict; got {type(data).__name__}")
|
||||
|
||||
|
||||
def _apply_subset(
|
||||
tasks: List[Dict[str, Any]],
|
||||
subset: Dict[str, Any],
|
||||
cell: Dict[str, Any],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Filter ``tasks`` to the subset's task IDs, preserving subset order.
|
||||
|
||||
Hard-errors if the cell's ``n`` doesn't equal ``len(task_ids)`` so a
|
||||
typo in the registry can't silently shrink the eval. Also errors if
|
||||
any subset ID is missing from the dataset (caller's bench wiring is
|
||||
broken).
|
||||
"""
|
||||
ids: List[str] = list(subset["task_ids"])
|
||||
cell_n = int(cell["n"])
|
||||
if cell_n != len(ids):
|
||||
raise ValueError(
|
||||
f"subset n={len(ids)} ≠ cell n={cell_n} — refusing to silently "
|
||||
"change scope. Fix the registry's `n` to match the subset file."
|
||||
)
|
||||
if "bench" in subset and subset["bench"] != cell["bench"]:
|
||||
raise ValueError(
|
||||
f"subset bench={subset['bench']!r} ≠ cell bench={cell['bench']!r}"
|
||||
)
|
||||
order = {tid: i for i, tid in enumerate(ids)}
|
||||
allow = set(ids)
|
||||
kept = [t for t in tasks if t["task_id"] in allow]
|
||||
kept.sort(key=lambda t: order[t["task_id"]])
|
||||
missing = allow - {t["task_id"] for t in kept}
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"subset references {len(missing)} task_ids not in dataset "
|
||||
f"(e.g. {next(iter(missing))!r})"
|
||||
)
|
||||
return kept
|
||||
|
||||
|
||||
# ---------- Scoring ----------
|
||||
|
||||
def _score_gaia(task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
"""Exact-match-with-format-normalization GAIA scorer.
|
||||
_GAIA_SCORER = None
|
||||
_GAIA_SCORER_LOCK = threading.Lock()
|
||||
|
||||
Lightweight version: extracts the final-answer line and string-compares
|
||||
against the reference. Use the OpenJarvis gaia_exact scorer for the
|
||||
judge-tiebreaker path.
|
||||
|
||||
def _get_gaia_scorer():
|
||||
"""Lazily build the shared GAIA scorer (normalized exact-match + LLM judge).
|
||||
|
||||
Judge model defaults to ``gpt-5-mini-2025-08-07`` (override via
|
||||
``OPENJARVIS_GAIA_JUDGE_MODEL``); the judge backend is the ``cloud``
|
||||
engine, so ``OPENJARVIS_CONFIG`` needs a ``[engine.cloud]`` section.
|
||||
"""
|
||||
import re
|
||||
global _GAIA_SCORER
|
||||
if _GAIA_SCORER is None:
|
||||
with _GAIA_SCORER_LOCK:
|
||||
if _GAIA_SCORER is None:
|
||||
from openjarvis.evals.backends.jarvis_direct import (
|
||||
JarvisDirectBackend,
|
||||
)
|
||||
from openjarvis.evals.scorers.gaia_exact import GAIAScorer
|
||||
|
||||
judge_model = os.environ.get(
|
||||
"OPENJARVIS_GAIA_JUDGE_MODEL", "gpt-5-mini-2025-08-07"
|
||||
)
|
||||
try:
|
||||
backend = JarvisDirectBackend(engine_key="cloud")
|
||||
except Exception: # noqa: BLE001
|
||||
backend = None
|
||||
_GAIA_SCORER = GAIAScorer(backend, judge_model)
|
||||
return _GAIA_SCORER
|
||||
|
||||
|
||||
def _score_gaia(task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
"""GAIA scorer — normalized exact-match with an LLM-judge fallback.
|
||||
|
||||
Uses the shared OpenJarvis :class:`GAIAScorer`. The previous version
|
||||
only credited answers that emitted a literal ``FINAL ANSWER:`` line and
|
||||
string-matched it; a verbose answer that stated the right answer in
|
||||
prose silently scored 0. Opus emits the marker ~92% of the time but
|
||||
GPT-5-mini / Haiku almost never do, so their GAIA cells were badly
|
||||
undercounted. The judge recovers the answer from prose instead.
|
||||
"""
|
||||
from openjarvis.evals.core.types import EvalRecord
|
||||
|
||||
ref = (task.get("reference") or "").strip()
|
||||
if not ref:
|
||||
return {"success": False, "score": 0.0, "details": {"reason": "no_reference"}}
|
||||
m = re.search(
|
||||
r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$",
|
||||
answer,
|
||||
re.IGNORECASE | re.MULTILINE,
|
||||
|
||||
record = EvalRecord(
|
||||
record_id=str(task.get("task_id") or ""),
|
||||
problem=str(task.get("question") or ""),
|
||||
reference=ref,
|
||||
category="agentic",
|
||||
metadata=dict(task.get("metadata") or {}),
|
||||
)
|
||||
pred = (m.group(1).strip() if m else answer.strip()).rstrip(".").strip()
|
||||
success = pred.lower() == ref.lower()
|
||||
is_correct, details = _get_gaia_scorer().score(record, answer or "")
|
||||
details = dict(details or {})
|
||||
details.setdefault("reference", ref)
|
||||
return {
|
||||
"success": success,
|
||||
"score": 1.0 if success else 0.0,
|
||||
"details": {"prediction": pred, "reference": ref},
|
||||
"success": bool(is_correct),
|
||||
"score": 1.0 if is_correct else 0.0,
|
||||
"details": details,
|
||||
}
|
||||
|
||||
|
||||
def _score_swebench(task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
"""Modal-backed SWE-bench Verified harness scorer."""
|
||||
def _score_swebench(
|
||||
task: Dict[str, Any],
|
||||
answer: str,
|
||||
cell_name: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Modal-backed SWE-bench Verified harness scorer.
|
||||
|
||||
``cell_name`` is passed through to :class:`SWEBenchHarnessScorer` so the
|
||||
underlying ``run_id`` is unique per (cell, instance). Without it,
|
||||
concurrent hybrid cells scoring the same task collide on the shared
|
||||
swebench harness cache and the second cell silently scores 0 with
|
||||
``reason: no_report`` (or reads the first cell's verdict).
|
||||
"""
|
||||
from openjarvis.evals.core.types import EvalRecord
|
||||
from openjarvis.evals.scorers.swebench_harness import (
|
||||
SWEBenchHarnessScorer,
|
||||
@@ -171,7 +342,10 @@ def _score_swebench(task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
category="agentic",
|
||||
metadata={"instance_id": task["task_id"]},
|
||||
)
|
||||
scorer = SWEBenchHarnessScorer(timeout_s=int(os.environ.get("SWEBENCH_TIMEOUT_S", "1800")))
|
||||
scorer = SWEBenchHarnessScorer(
|
||||
timeout_s=int(os.environ.get("SWEBENCH_TIMEOUT_S", "1800")),
|
||||
cell_name=cell_name,
|
||||
)
|
||||
is_correct, details = scorer.score(record, answer)
|
||||
return {
|
||||
"success": bool(is_correct),
|
||||
@@ -180,11 +354,16 @@ def _score_swebench(task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def score(bench: str, task: Dict[str, Any], answer: str) -> Dict[str, Any]:
|
||||
def score(
|
||||
bench: str,
|
||||
task: Dict[str, Any],
|
||||
answer: str,
|
||||
cell_name: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
if bench == "gaia":
|
||||
return _score_gaia(task, answer)
|
||||
if bench in ("swebench-verified", "swebench_verified", "swebench"):
|
||||
return _score_swebench(task, answer)
|
||||
return _score_swebench(task, answer, cell_name=cell_name)
|
||||
raise ValueError(f"unknown bench: {bench!r}")
|
||||
|
||||
|
||||
@@ -256,7 +435,26 @@ def _build_agent(cell: Dict[str, Any]):
|
||||
)
|
||||
|
||||
|
||||
def _run_one(agent, bench: str, task: Dict[str, Any], log_dir: str) -> Dict[str, Any]:
|
||||
def _error_row(task: Dict[str, Any], t0: float, error: str) -> Dict[str, Any]:
|
||||
"""Build a hybrid-shape error row (kept null-shaped like the catch path
|
||||
in :func:`_run_one` so the resume logic re-runs it)."""
|
||||
return {
|
||||
"task_id": task["task_id"],
|
||||
"answer": "",
|
||||
"tokens_local": 0, "tokens_cloud": 0,
|
||||
"cost_usd": 0.0, "latency_s": time.time() - t0,
|
||||
"web_search_uses": 0,
|
||||
"tool_calls": 0,
|
||||
"n_cloud_calls": 0,
|
||||
"n_local_calls": 0,
|
||||
"traces": {},
|
||||
"error": error,
|
||||
}
|
||||
|
||||
|
||||
def _run_one_inner(
|
||||
agent, bench: str, task: Dict[str, Any], log_dir: str
|
||||
) -> Dict[str, Any]:
|
||||
"""Run the agent on one task. Returns a hybrid-shape row."""
|
||||
prompt = _format_prompt(task)
|
||||
ctx = AgentContext(metadata={
|
||||
@@ -275,20 +473,80 @@ def _run_one(agent, bench: str, task: Dict[str, Any], log_dir: str) -> Dict[str,
|
||||
"tokens_cloud": int(meta.get("tokens_cloud", 0)),
|
||||
"cost_usd": float(meta.get("cost_usd", 0.0)),
|
||||
"latency_s": float(meta.get("latency_s", time.time() - t0)),
|
||||
"web_search_uses": int(meta.get("web_search_uses", 0)),
|
||||
"tool_calls": int(meta.get("tool_calls", 0)),
|
||||
"n_cloud_calls": int(meta.get("n_cloud_calls", 0)),
|
||||
"n_local_calls": int(meta.get("n_local_calls", 0)),
|
||||
"traces": meta.get("traces", {}),
|
||||
}
|
||||
if "soft_error" in meta:
|
||||
out["soft_error"] = meta["soft_error"]
|
||||
return {**out, "error": None}
|
||||
except Exception as e:
|
||||
return {
|
||||
"task_id": task["task_id"],
|
||||
"answer": "",
|
||||
"tokens_local": 0, "tokens_cloud": 0,
|
||||
"cost_usd": 0.0, "latency_s": time.time() - t0,
|
||||
"traces": {},
|
||||
"error": f"{type(e).__name__}: {e}\n{traceback.format_exc()}",
|
||||
}
|
||||
return _error_row(
|
||||
task, t0, f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
)
|
||||
|
||||
|
||||
def _run_one(
|
||||
agent,
|
||||
bench: str,
|
||||
task: Dict[str, Any],
|
||||
log_dir: str,
|
||||
*,
|
||||
task_timeout_s: float = DEFAULT_TASK_TIMEOUT_S,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run one task under a hard wall-clock cap.
|
||||
|
||||
``_run_one_inner`` runs on a dedicated **daemon** thread; if it doesn't
|
||||
finish within ``task_timeout_s`` we give up on it and record a
|
||||
``TaskTimeout`` error row. The worker thread is then *abandoned* — a
|
||||
truly wedged task (hung socket read with no enforced timeout, a Modal
|
||||
harness subprocess deadlocked draining pipes) cannot be killed
|
||||
cooperatively in CPython, so leaking the thread is the only safe
|
||||
option. It's a daemon thread, so it never blocks process exit, and the
|
||||
leak is bounded (one per timed-out task) — far cheaper than letting the
|
||||
whole cell freeze on the runner's ``as_completed`` join. The error row
|
||||
makes the resume logic re-run the task on the next invocation.
|
||||
|
||||
``task_timeout_s <= 0`` disables the cap (runs inline, legacy behavior).
|
||||
"""
|
||||
t0 = time.time()
|
||||
if task_timeout_s <= 0:
|
||||
return _run_one_inner(agent, bench, task, log_dir)
|
||||
|
||||
box: Dict[str, Any] = {}
|
||||
|
||||
def _target() -> None:
|
||||
try:
|
||||
box["row"] = _run_one_inner(agent, bench, task, log_dir)
|
||||
except BaseException as e: # noqa: BLE001 — never let the worker die silently
|
||||
box["row"] = _error_row(
|
||||
task, t0, f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
)
|
||||
|
||||
worker = threading.Thread(
|
||||
target=_target,
|
||||
name=f"hybrid-task-{task['task_id']}",
|
||||
daemon=True,
|
||||
)
|
||||
worker.start()
|
||||
worker.join(timeout=task_timeout_s)
|
||||
if worker.is_alive():
|
||||
print(
|
||||
f"[timeout] task={task['task_id']} exceeded "
|
||||
f"{task_timeout_s/60:.1f}m — abandoning worker, recording error row",
|
||||
flush=True,
|
||||
)
|
||||
return _error_row(
|
||||
task, t0,
|
||||
f"TaskTimeout: task exceeded the {task_timeout_s:.0f}s hybrid "
|
||||
"per-task wall-clock cap (likely a hung network or Modal-harness "
|
||||
"call); worker thread abandoned, task left for resume.",
|
||||
)
|
||||
return box.get("row") or _error_row(
|
||||
task, t0, "TaskError: worker thread exited without producing a row."
|
||||
)
|
||||
|
||||
|
||||
def _heartbeat(done: int, total: int, row: Dict[str, Any], t_start: float) -> None:
|
||||
@@ -312,6 +570,8 @@ def _write_summary(
|
||||
cell: Dict[str, Any],
|
||||
tasks: List[Dict[str, Any]],
|
||||
t_start: float,
|
||||
n_processed: int = -1,
|
||||
energy_j_session: float = 0.0,
|
||||
) -> None:
|
||||
results_path = out_dir / "results.jsonl"
|
||||
rows = [
|
||||
@@ -326,8 +586,38 @@ def _write_summary(
|
||||
total_cost = sum(r.get("cost_usd", 0.0) for r in rows)
|
||||
total_local = sum(r.get("tokens_local", 0) for r in rows)
|
||||
total_cloud = sum(r.get("tokens_cloud", 0) for r in rows)
|
||||
total_web_searches = sum(int(r.get("web_search_uses", 0) or 0) for r in rows)
|
||||
total_tool_calls = sum(int(r.get("tool_calls", 0) or 0) for r in rows)
|
||||
total_cloud_calls = sum(int(r.get("n_cloud_calls", 0) or 0) for r in rows)
|
||||
total_local_calls = sum(int(r.get("n_local_calls", 0) or 0) for r in rows)
|
||||
elapsed = time.time() - t_start
|
||||
|
||||
# Preserve prior wall_time_s on no-op resume so we don't clobber the
|
||||
# original run's runtime. A resume that did zero work (everything was
|
||||
# already cached in results.jsonl) records ~seconds of elapsed time;
|
||||
# writing that as wall_time_s makes the cell look 20× faster than it
|
||||
# really was. If we did process anything this session, accumulate so
|
||||
# partial-resumes still report total wall time honestly.
|
||||
summary_path = out_dir / "summary.json"
|
||||
prior_wall = 0.0
|
||||
prior_energy = 0.0
|
||||
if summary_path.exists():
|
||||
try:
|
||||
prior = json.loads(summary_path.read_text())
|
||||
prior_wall = float(prior.get("wall_time_s", 0.0) or 0.0)
|
||||
prior_energy = float(prior.get("energy_j_total", 0.0) or 0.0)
|
||||
except Exception:
|
||||
prior_wall = 0.0
|
||||
prior_energy = 0.0
|
||||
if n_processed == 0 and prior_wall > 0:
|
||||
wall = prior_wall
|
||||
# No work done this session → keep prior energy total (don't add a
|
||||
# spurious idle-load reading from a resume that processed nothing).
|
||||
energy_j = prior_energy
|
||||
else:
|
||||
wall = prior_wall + elapsed
|
||||
energy_j = prior_energy + float(energy_j_session or 0.0)
|
||||
|
||||
summary = {
|
||||
"cell": cell_name,
|
||||
"method": cell["method"],
|
||||
@@ -338,14 +628,27 @@ def _write_summary(
|
||||
"accuracy": acc,
|
||||
"tokens_local_total": total_local,
|
||||
"tokens_cloud_total": total_cloud,
|
||||
"web_search_uses_total": total_web_searches,
|
||||
"tool_calls_total": total_tool_calls,
|
||||
"n_cloud_calls_total": total_cloud_calls,
|
||||
"n_local_calls_total": total_local_calls,
|
||||
"cost_usd_total": total_cost,
|
||||
"wall_time_s": elapsed,
|
||||
"wall_time_s": wall,
|
||||
# GPU energy integrated over the cell's wall-time across the GPUs
|
||||
# visible to the runner host. Cloud energy is **not** included —
|
||||
# see ``_energy.py``. Joules; sum of session + any prior resumes.
|
||||
# TODO: decide whether to add a cloud J/token estimate; for now 0
|
||||
# cloud contribution. (Patterson 2021 / Luccioni 2022 are options.)
|
||||
"energy_j_total": energy_j,
|
||||
"task_count": len(tasks),
|
||||
}
|
||||
(out_dir / "summary.json").write_text(json.dumps(summary, indent=2))
|
||||
summary_path.write_text(json.dumps(summary, indent=2))
|
||||
print(
|
||||
f"[summary] {cell_name}: n={n_done}/{cell['n']} err={n_err} "
|
||||
f"acc={acc:.3f} cost=${total_cost:.2f} time={elapsed/60:.1f}m",
|
||||
f"acc={acc:.3f} cost=${total_cost:.2f} time={wall/60:.1f}m "
|
||||
f"energy={energy_j/1000:.1f}kJ "
|
||||
f"(session +{elapsed/60:.1f}m +{energy_j_session/1000:.1f}kJ, "
|
||||
f"processed={n_processed})",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
@@ -358,7 +661,7 @@ def run_cell(
|
||||
resume: bool = True,
|
||||
root: Optional[Path] = None,
|
||||
) -> None:
|
||||
out_root = root or DEFAULT_EXPERIMENTS_DIR
|
||||
out_root = root or DEFAULT_RUNS_DIR
|
||||
out_dir = _cell_dir(cell_name, out_root)
|
||||
with _cell_lock(out_dir, cell_name):
|
||||
_run_cell_locked(
|
||||
@@ -398,49 +701,140 @@ def _run_cell_locked(
|
||||
flush=True,
|
||||
)
|
||||
|
||||
tasks = load_tasks(cell["bench"], n=cell["n"])
|
||||
print(f"[load] {cell['bench']} → {len(tasks)} tasks", flush=True)
|
||||
subset_path = cell.get("subset")
|
||||
if subset_path:
|
||||
subset = _load_subset_file(subset_path)
|
||||
# Load the full bench (n=None) so we can pick the exact subset IDs.
|
||||
# The dataset providers are cached, so this isn't a re-fetch.
|
||||
all_tasks = load_tasks(cell["bench"], n=None)
|
||||
tasks = _apply_subset(all_tasks, subset, cell)
|
||||
print(
|
||||
f"[load] {cell['bench']} subset={Path(subset_path).name} "
|
||||
f"→ {len(tasks)} tasks",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
tasks = load_tasks(cell["bench"], n=cell["n"])
|
||||
print(f"[load] {cell['bench']} → {len(tasks)} tasks", flush=True)
|
||||
|
||||
pending = [t for t in tasks if t["task_id"] not in done_ids]
|
||||
concurrency = max(1, int(cell.get("concurrency", 1)))
|
||||
if concurrency > 1:
|
||||
print(f"[concurrency] {concurrency} workers", flush=True)
|
||||
|
||||
# Hard per-task wall-clock cap. A cell may override it via the registry
|
||||
# (``method_cfg.task_timeout_s``); otherwise the process-wide default
|
||||
# (env ``OPENJARVIS_HYBRID_TASK_TIMEOUT_S``, 1800s) applies. 0 disables.
|
||||
mcfg = cell.get("method_cfg") or {}
|
||||
task_timeout_s = float(mcfg.get("task_timeout_s", DEFAULT_TASK_TIMEOUT_S))
|
||||
if task_timeout_s > 0:
|
||||
print(f"[task-timeout] {task_timeout_s/60:.1f}m per task", flush=True)
|
||||
|
||||
agent = _build_agent(cell)
|
||||
|
||||
t_start = time.time()
|
||||
write_lock = threading.Lock()
|
||||
completed = [0]
|
||||
written_ok_ids: set = set()
|
||||
log_dir = str(out_dir / "logs")
|
||||
|
||||
def _process(task: Dict[str, Any]) -> None:
|
||||
row = _run_one(agent, cell["bench"], task, log_dir)
|
||||
row = _run_one(
|
||||
agent, cell["bench"], task, log_dir,
|
||||
task_timeout_s=task_timeout_s,
|
||||
)
|
||||
scored: Optional[Dict[str, Any]] = None
|
||||
if do_score and row.get("error") is None:
|
||||
try:
|
||||
scored = score(cell["bench"], task, row["answer"])
|
||||
scored = score(
|
||||
cell["bench"], task, row["answer"], cell_name=cell_name,
|
||||
)
|
||||
except Exception as e:
|
||||
scored = {
|
||||
"success": False, "score": 0.0,
|
||||
"details": {"score_error": str(e)},
|
||||
}
|
||||
full_row = {**row, "score": scored}
|
||||
with write_lock, results_path.open("a") as f:
|
||||
f.write(json.dumps(full_row) + "\n")
|
||||
f.flush()
|
||||
with write_lock:
|
||||
# Idempotency guard: a Modal retry can re-run the same task within
|
||||
# one process. Skip appending once a non-error row exists for this
|
||||
# task_id so results.jsonl never carries duplicate rows.
|
||||
if full_row["task_id"] in written_ok_ids:
|
||||
return
|
||||
with results_path.open("a") as f:
|
||||
f.write(json.dumps(full_row) + "\n")
|
||||
f.flush()
|
||||
if full_row.get("error") is None:
|
||||
written_ok_ids.add(full_row["task_id"])
|
||||
completed[0] += 1
|
||||
_heartbeat(completed[0], len(tasks), full_row, t_start)
|
||||
|
||||
if concurrency == 1:
|
||||
for task in pending:
|
||||
_process(task)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as ex:
|
||||
futures = [ex.submit(_process, t) for t in pending]
|
||||
for fut in as_completed(futures):
|
||||
fut.result()
|
||||
# Hard watchdog. The per-task ``worker.join(timeout=task_timeout_s)`` in
|
||||
# ``_run_one`` is supposed to bound any single task, but in practice we've
|
||||
# seen the main loop wedge in ``futex_wait_queue`` despite the join
|
||||
# returning — likely a daemon thread holding a non-Python lock that the
|
||||
# GC / atexit handler trips on. Defense in depth: if no row hits
|
||||
# ``results.jsonl`` in ``watchdog_stale_s = 2 * task_timeout_s + 600``
|
||||
# seconds, ``os._exit(2)`` the whole process. The wrapper script's
|
||||
# resume logic will pick up unscored tasks on the next invocation, and
|
||||
# we don't burn another 30+ min on a wedge. Disabled when
|
||||
# ``task_timeout_s <= 0`` (matches the legacy in-process timeout knob).
|
||||
watchdog_stale_s = (2 * task_timeout_s + 600) if task_timeout_s > 0 else 0
|
||||
watchdog_stop = threading.Event()
|
||||
|
||||
_write_summary(out_dir, cell_name, cell, tasks, t_start)
|
||||
def _watchdog() -> None:
|
||||
baseline_mtime = (
|
||||
results_path.stat().st_mtime if results_path.exists() else time.time()
|
||||
)
|
||||
last_seen = baseline_mtime
|
||||
while not watchdog_stop.wait(60.0):
|
||||
try:
|
||||
cur = (
|
||||
results_path.stat().st_mtime
|
||||
if results_path.exists()
|
||||
else last_seen
|
||||
)
|
||||
except Exception:
|
||||
cur = last_seen
|
||||
if cur > last_seen:
|
||||
last_seen = cur
|
||||
continue
|
||||
if time.time() - last_seen > watchdog_stale_s:
|
||||
print(
|
||||
f"[watchdog] results.jsonl stale for "
|
||||
f"{int(time.time() - last_seen)}s "
|
||||
f"(threshold {int(watchdog_stale_s)}s) — hard-exit, "
|
||||
f"resume on next invocation.",
|
||||
flush=True,
|
||||
)
|
||||
os._exit(2)
|
||||
|
||||
if watchdog_stale_s > 0:
|
||||
threading.Thread(
|
||||
target=_watchdog, name="hybrid-runner-watchdog", daemon=True
|
||||
).start()
|
||||
|
||||
# GPU energy sampler covers the same wall-clock window as ``wall_time_s``
|
||||
# so the two numbers can be divided into an effective Watts figure.
|
||||
# Sampler is best-effort: NVML failures degrade to ``energy_j_total=0``
|
||||
# without crashing the run (see ``_energy.py``).
|
||||
with EnergyCollector() as energy:
|
||||
if concurrency == 1:
|
||||
for task in pending:
|
||||
_process(task)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as ex:
|
||||
futures = [ex.submit(_process, t) for t in pending]
|
||||
for fut in as_completed(futures):
|
||||
fut.result()
|
||||
|
||||
watchdog_stop.set()
|
||||
|
||||
_write_summary(
|
||||
out_dir, cell_name, cell, tasks, t_start,
|
||||
n_processed=len(pending),
|
||||
energy_j_session=energy.energy_j_total,
|
||||
)
|
||||
|
||||
|
||||
# ---------- CLI ----------
|
||||
@@ -492,6 +886,7 @@ if __name__ == "__main__":
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_EXPERIMENTS_DIR",
|
||||
"DEFAULT_RUNS_DIR",
|
||||
"DEFAULT_REGISTRY_DIR",
|
||||
"load_registry",
|
||||
"load_tasks",
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
# skillorchestra — faithful port of SkillOrchestra (arXiv:2602.19672)
|
||||
|
||||
This package replaces the old single-file `skillorchestra.py` (a 2-agent
|
||||
JSON router that shared none of the original's prompts or structure). It
|
||||
restructures the agent to **be** the original's eval-orchestrator runtime.
|
||||
|
||||
## What's faithful
|
||||
|
||||
- **Prompts** — `prompts/` is a verbatim copy of the upstream
|
||||
`skillorchestra/prompts/` package (`eval_orchestrator.py`,
|
||||
`model_routing.py`, `learning.py`).
|
||||
- **Handbook + routing** — `stage_router.py` and `types.py` are verbatim
|
||||
copies of upstream `adapters/stage_router.py` and `core/types.py`:
|
||||
`StageSkillHandbook`, `parse_skill_analysis`, and all 5 routing
|
||||
strategies (`router_decides`, `analyze_model_decide`, `weighted_avg`,
|
||||
`weakest_skill`, `strongest_skill`).
|
||||
- **Loop** — `orchestrator.py` ports `eval_frames.py:run_single`: the
|
||||
multi-round `search -> reasoning -> answer` ReAct loop, the verbatim
|
||||
worker prompts, `<skill_analysis>` parsing, alias-based model routing,
|
||||
last-round forced answer.
|
||||
- **Code tool** — `tools.run_code` runs model-generated Python in a real
|
||||
`subprocess` with a timeout, exactly as upstream.
|
||||
|
||||
## What can't match without infrastructure
|
||||
|
||||
The original runtime needs three things this cluster doesn't have. Each
|
||||
degrades gracefully and is configurable:
|
||||
|
||||
| Original | Here |
|
||||
|---|---|
|
||||
| FAISS wiki retriever for `search` | `method_cfg.retriever_url` POSTs the same `/retrieve` payload; absent → Anthropic `web_search` |
|
||||
| 6+ SGLang-served pool models | alias tiers collapse onto the cell's local/cloud pair; override per alias via `method_cfg.model_pool` |
|
||||
| Learned `handbook.json` from explore→learn→select | `method_cfg.handbook_path` (a hand-authored `handbook_seed.json` ships here); absent → `routing_strategy="none"`, the original's baseline mode |
|
||||
|
||||
The offline explore→learn→select pipeline is **not** ported — it needs
|
||||
the served model pool + FRAMES/NQ datasets to run. The handbook *schema*
|
||||
it produces is fully supported by `StageSkillHandbook.load`, so a learned
|
||||
handbook can be dropped in later with no code change.
|
||||
|
||||
SWE-bench is out of scope for the original (a QA orchestrator). SWE cells
|
||||
run the cloud backbone through the shared `mini_swe_agent` loop.
|
||||
|
||||
## method_cfg
|
||||
|
||||
`routing_strategy`, `handbook_path`, `max_rounds`, `retriever_url`,
|
||||
`model_pool`, `orchestrator_model` / `orchestrator_endpoint`,
|
||||
`code_timeout_s`, `answer_max_tokens`, `context_char_cap`. `router_model`
|
||||
/ `router_endpoint` are accepted as back-compat aliases.
|
||||
@@ -0,0 +1,22 @@
|
||||
"""SkillOrchestra — faithful port of the eval-orchestrator (arXiv:2602.19672).
|
||||
|
||||
Importing this package registers the ``skillorchestra`` agent. Layout
|
||||
mirrors the upstream repo (``external/SkillOrchestra``):
|
||||
|
||||
* :mod:`.prompts` — verbatim ``eval_orchestrator`` / ``model_routing``
|
||||
/ ``learning`` prompt templates.
|
||||
* :mod:`.stage_router` — verbatim ``StageSkillHandbook``,
|
||||
``parse_skill_analysis``, the 5 routing strategies.
|
||||
* :mod:`.types` — verbatim learning-time data types (BetaCompetence,
|
||||
Skill, AgentProfile, ...).
|
||||
* :mod:`.pool` — model-alias -> local/cloud resolution.
|
||||
* :mod:`.tools` — search / enhance_reasoning / answer executors.
|
||||
* :mod:`.orchestrator` — the multi-round search->code->answer loop.
|
||||
* :mod:`.agent` — :class:`SkillOrchestraAgent`, the harness entry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .agent import SkillOrchestraAgent
|
||||
|
||||
__all__ = ["SkillOrchestraAgent"]
|
||||
@@ -0,0 +1,153 @@
|
||||
"""SkillOrchestraAgent — the OpenJarvis harness entry point.
|
||||
|
||||
A faithful port of the SkillOrchestra eval orchestrator (arXiv:2602.19672,
|
||||
``orchestration/eval_frames.py``). The agent runs the multi-round
|
||||
search -> reasoning -> answer loop in :mod:`.orchestrator`, using the
|
||||
verbatim ``eval_orchestrator`` prompts, the ``StageSkillHandbook`` +
|
||||
``RoutingStrategy`` machinery, real Python subprocess execution, and a
|
||||
model-alias pool collapsed onto the cell's local/cloud pair.
|
||||
|
||||
Three things the original needs that this environment does not have, and
|
||||
how each is handled (see ``README.md`` in this package for the full
|
||||
note):
|
||||
|
||||
* **Learned handbook** — produced offline by the explore->learn->select
|
||||
pipeline. With no handbook the orchestrator runs ``routing_strategy =
|
||||
"none"`` (the original's baseline mode). Point ``method_cfg.handbook_path``
|
||||
at a ``StageSkillHandbook`` JSON to enable skill routing.
|
||||
* **FAISS wiki retriever** — the ``search`` tool POSTs to it when
|
||||
``method_cfg.retriever_url`` is set; otherwise it falls back to
|
||||
Anthropic ``web_search``.
|
||||
* **6+ model pool** — the alias tiers collapse onto the cell's local +
|
||||
cloud models; override per alias with ``method_cfg.model_pool``.
|
||||
|
||||
SWE-bench cells are out of scope for the original (it is a QA
|
||||
orchestrator). They run the cloud backbone through the shared mini SWE
|
||||
agent loop instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from openjarvis.agents._stubs import AgentContext
|
||||
from openjarvis.core.registry import AgentRegistry
|
||||
|
||||
from .._base import LocalCloudAgent
|
||||
from ..mini_swe_agent import run_swe_agent_loop
|
||||
from .orchestrator import run_orchestrator
|
||||
from .stage_router import StageSkillHandbook
|
||||
|
||||
_VALID_STRATEGIES = {
|
||||
"none", "router_decides", "analyze_model_decide",
|
||||
"weighted_avg", "weakest_skill", "strongest_skill",
|
||||
}
|
||||
|
||||
|
||||
@AgentRegistry.register("skillorchestra")
|
||||
class SkillOrchestraAgent(LocalCloudAgent):
|
||||
"""Inference-time skill-aware orchestrator. See module docstring."""
|
||||
|
||||
agent_id = "skillorchestra"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _is_soft_failure(self, exc: BaseException) -> Optional[str]:
|
||||
# Malformed orchestrator / router JSON -> soft-fail row, matching
|
||||
# the rest of the hybrid family.
|
||||
if isinstance(exc, (ValueError, json.JSONDecodeError)):
|
||||
return f"{type(exc).__name__}: {str(exc)[:120]}"
|
||||
return None
|
||||
|
||||
def _load_handbook(self) -> Optional[StageSkillHandbook]:
|
||||
"""Load the StageSkillHandbook from ``method_cfg.handbook_path``.
|
||||
|
||||
A relative path resolves against this package directory so the
|
||||
shipped ``handbook_seed.json`` works out of the box. Any load
|
||||
failure degrades to ``None`` (orchestrator runs baseline mode).
|
||||
"""
|
||||
path = self._cfg.get("handbook_path")
|
||||
if not path:
|
||||
return None
|
||||
p = Path(path)
|
||||
if not p.is_absolute():
|
||||
p = Path(__file__).parent / p
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
return StageSkillHandbook.load(str(p))
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_paradigm(
|
||||
self,
|
||||
input: str,
|
||||
context: Optional[AgentContext],
|
||||
**kwargs: Any,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
cfg = self._cfg
|
||||
task_meta = (context.metadata.get("task") if context is not None else {}) or {}
|
||||
|
||||
# SWE-bench: the original SkillOrchestra has no code-repo mode.
|
||||
# Run the cloud backbone through the shared mini SWE agent loop.
|
||||
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"))
|
||||
)
|
||||
if swe_mode:
|
||||
out = run_swe_agent_loop(
|
||||
task_meta,
|
||||
backbone="cloud",
|
||||
backbone_model=self._cloud_model,
|
||||
cloud_endpoint=self._cloud_endpoint,
|
||||
initial_prompt=input,
|
||||
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="skillorch_swe",
|
||||
)
|
||||
meta = {
|
||||
"tokens_local": 0,
|
||||
"tokens_cloud": out["tokens_in"] + out["tokens_out"],
|
||||
"cost_usd": out["cost_usd"],
|
||||
"turns": int(out["turns"]),
|
||||
"tool_calls": int(out["turns"]),
|
||||
"web_search_uses": 0,
|
||||
"traces": {
|
||||
"mode": "swe_agent_loop",
|
||||
"backbone_model": self._cloud_model,
|
||||
"note": "original SkillOrchestra is QA-only; SWE uses the cloud backbone",
|
||||
},
|
||||
}
|
||||
return out["answer"], meta
|
||||
|
||||
# QA path — the faithful eval orchestrator.
|
||||
handbook = self._load_handbook()
|
||||
strategy = str(cfg.get("routing_strategy", "none"))
|
||||
if strategy not in _VALID_STRATEGIES:
|
||||
raise ValueError(
|
||||
f"routing_strategy {strategy!r} unknown; valid: "
|
||||
f"{sorted(_VALID_STRATEGIES)}"
|
||||
)
|
||||
if handbook is None and strategy in ("router_decides", "analyze_model_decide"):
|
||||
# These two strategies just honor the orchestrator's own model
|
||||
# pick — they need no learned skill data, so an empty handbook
|
||||
# is enough to run the skill-orchestrator prompt + loop.
|
||||
handbook = StageSkillHandbook()
|
||||
if handbook is None:
|
||||
# No handbook -> baseline routing, exactly like the original.
|
||||
strategy = "none"
|
||||
|
||||
return run_orchestrator(
|
||||
self, input, cfg=cfg, handbook=handbook, strategy=strategy,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["SkillOrchestraAgent"]
|
||||
@@ -0,0 +1,500 @@
|
||||
{
|
||||
"version": "0.0.1-seed",
|
||||
"created_at": "2026-05-19",
|
||||
"updated_at": "2026-05-19",
|
||||
"_note": "HAND-AUTHORED SEED \u2014 not produced by the explore->learn->select pipeline. Priors are calibrated to Qwen-27B vs Opus on GAIA-style QA. Replace with a learned StageSkillHandbook when available.",
|
||||
"skills": {
|
||||
"search": {
|
||||
"search.entity_lookup": {
|
||||
"skill_id": "search.entity_lookup",
|
||||
"name": "Entity Lookup",
|
||||
"description": "Find a specific named entity, date, or attribute (person, place, work, organization).",
|
||||
"stage": "search",
|
||||
"examples": [
|
||||
"Who directed the 1997 film Titanic?"
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
},
|
||||
"search.recent_events": {
|
||||
"skill_id": "search.recent_events",
|
||||
"name": "Recent Events",
|
||||
"description": "Find recent or time-sensitive facts unlikely to be in a small model's parametric memory.",
|
||||
"stage": "search",
|
||||
"examples": [
|
||||
"Which team won the 2025 Cricket World Cup?"
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
},
|
||||
"search.multi_hop": {
|
||||
"skill_id": "search.multi_hop",
|
||||
"name": "Multi Hop",
|
||||
"description": "Locate intermediate facts that chain together to answer a compositional question.",
|
||||
"stage": "search",
|
||||
"examples": [
|
||||
"What is the capital of the country that won the most 2024 Olympic gold medals?"
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
}
|
||||
},
|
||||
"code": {
|
||||
"code.arithmetic": {
|
||||
"skill_id": "code.arithmetic",
|
||||
"name": "Arithmetic",
|
||||
"description": "Exact numeric computation, unit conversion, and counting over given values.",
|
||||
"stage": "code",
|
||||
"examples": [
|
||||
"How many days are between 1999-03-01 and 2001-07-15?"
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
},
|
||||
"code.data_transform": {
|
||||
"skill_id": "code.data_transform",
|
||||
"name": "Data Transform",
|
||||
"description": "Parse, filter, sort, or aggregate structured data to derive an intermediate result.",
|
||||
"stage": "code",
|
||||
"examples": [
|
||||
"Given this table, which row has the third-highest revenue?"
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
}
|
||||
},
|
||||
"answer": {
|
||||
"answer.synthesis": {
|
||||
"skill_id": "answer.synthesis",
|
||||
"name": "Synthesis",
|
||||
"description": "Combine retrieved documents and code results into a single correct answer.",
|
||||
"stage": "answer",
|
||||
"examples": [
|
||||
"Combine the search results to state the final figure."
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
},
|
||||
"answer.format_compliance": {
|
||||
"skill_id": "answer.format_compliance",
|
||||
"name": "Format Compliance",
|
||||
"description": "Emit the answer in the exact required format (no units, exact casing, list shape).",
|
||||
"stage": "answer",
|
||||
"examples": [
|
||||
"Answer with a comma-separated list, no units."
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
},
|
||||
"answer.long_context": {
|
||||
"skill_id": "answer.long_context",
|
||||
"name": "Long Context",
|
||||
"description": "Read a long accumulated context and extract the precise answer span.",
|
||||
"stage": "answer",
|
||||
"examples": [
|
||||
"Extract the requested value from the document above."
|
||||
],
|
||||
"discovered_from_problems": []
|
||||
}
|
||||
}
|
||||
},
|
||||
"model_profiles": {
|
||||
"search-1": {
|
||||
"model_alias": "search-1",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "search",
|
||||
"skill_scores": {
|
||||
"search.entity_lookup": 0.82,
|
||||
"search.recent_events": 0.78,
|
||||
"search.multi_hop": 0.8
|
||||
},
|
||||
"skill_attempts": {
|
||||
"search.entity_lookup": 24,
|
||||
"search.recent_events": 24,
|
||||
"search.multi_hop": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"search.entity_lookup": 20,
|
||||
"search.recent_events": 19,
|
||||
"search.multi_hop": 19
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.014,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"search-2": {
|
||||
"model_alias": "search-2",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "search",
|
||||
"skill_scores": {
|
||||
"search.entity_lookup": 0.82,
|
||||
"search.recent_events": 0.78,
|
||||
"search.multi_hop": 0.8
|
||||
},
|
||||
"skill_attempts": {
|
||||
"search.entity_lookup": 24,
|
||||
"search.recent_events": 24,
|
||||
"search.multi_hop": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"search.entity_lookup": 20,
|
||||
"search.recent_events": 19,
|
||||
"search.multi_hop": 19
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.014,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"search-3": {
|
||||
"model_alias": "search-3",
|
||||
"actual_model": "Qwen/Qwen3.5-27B-FP8",
|
||||
"stage": "search",
|
||||
"skill_scores": {
|
||||
"search.entity_lookup": 0.55,
|
||||
"search.recent_events": 0.3,
|
||||
"search.multi_hop": 0.42
|
||||
},
|
||||
"skill_attempts": {
|
||||
"search.entity_lookup": 24,
|
||||
"search.recent_events": 24,
|
||||
"search.multi_hop": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"search.entity_lookup": 13,
|
||||
"search.recent_events": 7,
|
||||
"search.multi_hop": 10
|
||||
},
|
||||
"overall_success_rate": 0.34,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 8,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 280.0,
|
||||
"avg_cost_usd": 0.0,
|
||||
"strengths": [
|
||||
"strict formatting",
|
||||
"arithmetic on given values"
|
||||
],
|
||||
"weaknesses": [
|
||||
"recent / rare facts",
|
||||
"long-context extraction"
|
||||
]
|
||||
},
|
||||
"reasoner-1": {
|
||||
"model_alias": "reasoner-1",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "code",
|
||||
"skill_scores": {
|
||||
"code.arithmetic": 0.85,
|
||||
"code.data_transform": 0.84
|
||||
},
|
||||
"skill_attempts": {
|
||||
"code.arithmetic": 24,
|
||||
"code.data_transform": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"code.arithmetic": 20,
|
||||
"code.data_transform": 20
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.014,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"reasoner-2": {
|
||||
"model_alias": "reasoner-2",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "code",
|
||||
"skill_scores": {
|
||||
"code.arithmetic": 0.85,
|
||||
"code.data_transform": 0.84
|
||||
},
|
||||
"skill_attempts": {
|
||||
"code.arithmetic": 24,
|
||||
"code.data_transform": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"code.arithmetic": 20,
|
||||
"code.data_transform": 20
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.014,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"reasoner-3": {
|
||||
"model_alias": "reasoner-3",
|
||||
"actual_model": "Qwen/Qwen3.5-27B-FP8",
|
||||
"stage": "code",
|
||||
"skill_scores": {
|
||||
"code.arithmetic": 0.78,
|
||||
"code.data_transform": 0.66
|
||||
},
|
||||
"skill_attempts": {
|
||||
"code.arithmetic": 24,
|
||||
"code.data_transform": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"code.arithmetic": 19,
|
||||
"code.data_transform": 16
|
||||
},
|
||||
"overall_success_rate": 0.34,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 8,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 280.0,
|
||||
"avg_cost_usd": 0.0,
|
||||
"strengths": [
|
||||
"strict formatting",
|
||||
"arithmetic on given values"
|
||||
],
|
||||
"weaknesses": [
|
||||
"recent / rare facts",
|
||||
"long-context extraction"
|
||||
]
|
||||
},
|
||||
"answer-1": {
|
||||
"model_alias": "answer-1",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.84,
|
||||
"answer.format_compliance": 0.8,
|
||||
"answer.long_context": 0.86
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 20,
|
||||
"answer.format_compliance": 19,
|
||||
"answer.long_context": 21
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.018,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"answer-2": {
|
||||
"model_alias": "answer-2",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.84,
|
||||
"answer.format_compliance": 0.8,
|
||||
"answer.long_context": 0.86
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 20,
|
||||
"answer.format_compliance": 19,
|
||||
"answer.long_context": 21
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.018,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"answer-3": {
|
||||
"model_alias": "answer-3",
|
||||
"actual_model": "Qwen/Qwen3.5-27B-FP8",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.52,
|
||||
"answer.format_compliance": 0.74,
|
||||
"answer.long_context": 0.58
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 12,
|
||||
"answer.format_compliance": 18,
|
||||
"answer.long_context": 14
|
||||
},
|
||||
"overall_success_rate": 0.34,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 8,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 280.0,
|
||||
"avg_cost_usd": 0.0,
|
||||
"strengths": [
|
||||
"strict formatting",
|
||||
"arithmetic on given values"
|
||||
],
|
||||
"weaknesses": [
|
||||
"recent / rare facts",
|
||||
"long-context extraction"
|
||||
]
|
||||
},
|
||||
"answer-4": {
|
||||
"model_alias": "answer-4",
|
||||
"actual_model": "Qwen/Qwen3.5-27B-FP8",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.52,
|
||||
"answer.format_compliance": 0.74,
|
||||
"answer.long_context": 0.58
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 12,
|
||||
"answer.format_compliance": 18,
|
||||
"answer.long_context": 14
|
||||
},
|
||||
"overall_success_rate": 0.34,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 8,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 280.0,
|
||||
"avg_cost_usd": 0.0,
|
||||
"strengths": [
|
||||
"strict formatting",
|
||||
"arithmetic on given values"
|
||||
],
|
||||
"weaknesses": [
|
||||
"recent / rare facts",
|
||||
"long-context extraction"
|
||||
]
|
||||
},
|
||||
"answer-math-1": {
|
||||
"model_alias": "answer-math-1",
|
||||
"actual_model": "claude-opus-4-7",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.84,
|
||||
"answer.format_compliance": 0.8,
|
||||
"answer.long_context": 0.86
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 20,
|
||||
"answer.format_compliance": 19,
|
||||
"answer.long_context": 21
|
||||
},
|
||||
"overall_success_rate": 0.6,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 14,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 320.0,
|
||||
"avg_cost_usd": 0.018,
|
||||
"strengths": [
|
||||
"multi-hop reasoning",
|
||||
"rare-fact recall"
|
||||
],
|
||||
"weaknesses": [
|
||||
"higher cost"
|
||||
]
|
||||
},
|
||||
"answer-math-2": {
|
||||
"model_alias": "answer-math-2",
|
||||
"actual_model": "Qwen/Qwen3.5-27B-FP8",
|
||||
"stage": "answer",
|
||||
"skill_scores": {
|
||||
"answer.synthesis": 0.52,
|
||||
"answer.format_compliance": 0.74,
|
||||
"answer.long_context": 0.58
|
||||
},
|
||||
"skill_attempts": {
|
||||
"answer.synthesis": 24,
|
||||
"answer.format_compliance": 24,
|
||||
"answer.long_context": 24
|
||||
},
|
||||
"skill_successes": {
|
||||
"answer.synthesis": 12,
|
||||
"answer.format_compliance": 18,
|
||||
"answer.long_context": 14
|
||||
},
|
||||
"overall_success_rate": 0.34,
|
||||
"total_attempts": 24,
|
||||
"total_successes": 8,
|
||||
"avg_prompt_tokens": 1200.0,
|
||||
"avg_completion_tokens": 280.0,
|
||||
"avg_cost_usd": 0.0,
|
||||
"strengths": [
|
||||
"strict formatting",
|
||||
"arithmetic on given values"
|
||||
],
|
||||
"weaknesses": [
|
||||
"recent / rare facts",
|
||||
"long-context extraction"
|
||||
]
|
||||
}
|
||||
},
|
||||
"usage_patterns": {
|
||||
"stages": {},
|
||||
"guidelines": {},
|
||||
"models": {},
|
||||
"raw": {}
|
||||
},
|
||||
"learning_history": [],
|
||||
"routing_insights": [
|
||||
"Route recent/rare-fact searches to a cloud search model; the local model lacks that knowledge.",
|
||||
"Arithmetic and strict-format answers can go to the cheaper local model without losing accuracy.",
|
||||
"Switch to the answer stage once retrieved documents cover every sub-fact of the question."
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
"""The SkillOrchestra eval-orchestrator loop.
|
||||
|
||||
Faithful port of ``orchestration/eval_frames.py:run_single`` — the
|
||||
multi-round search -> reasoning -> answer ReAct loop. Each round:
|
||||
|
||||
1. Build a context string from accumulated docs / code results / attempts.
|
||||
2. Ask the orchestrator model (with the 3 tools) for the next stage. The
|
||||
prompt is the verbatim ``build_skill_orchestrator_prompt`` when a
|
||||
handbook is loaded, else the baseline ``"Problem: ... Choose an
|
||||
appropriate tool."`` string.
|
||||
3. Parse the tool call + any ``<skill_analysis>`` block; route the worker
|
||||
model alias through the configured ``RoutingStrategy``.
|
||||
4. Execute the tool. ``answer`` ends the loop; the last round force-calls
|
||||
``answer``.
|
||||
|
||||
The orchestrator step does raw SDK calls (Anthropic / OpenAI) because it
|
||||
needs the parsed ``tool_use`` blocks back — the same thing
|
||||
``extract_response_content_and_tool_calls`` does in the original.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .._prices import is_gpt5_family, supports_temperature
|
||||
from .pool import ModelSpec, build_pool
|
||||
from .prompts import build_skill_orchestrator_prompt
|
||||
from .stage_router import (
|
||||
StageSkillHandbook,
|
||||
get_routing_strategy,
|
||||
parse_skill_analysis,
|
||||
)
|
||||
from .tools import anthropic_tools, openai_tools, run_answer, run_code, run_search
|
||||
|
||||
# tool name -> routing stage (stage_router uses "reasoning" for code).
|
||||
_TOOL_STAGE = {
|
||||
"search": "search",
|
||||
"enhance_reasoning": "reasoning",
|
||||
"code": "reasoning",
|
||||
"answer": "answer",
|
||||
}
|
||||
_STAGE_DEFAULT_ALIAS = {
|
||||
"search": "search-1",
|
||||
"reasoning": "reasoner-1",
|
||||
"answer": "answer-1",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orchestrator decision step (raw SDK — needs tool_use blocks back)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _orchestrate_step(
|
||||
agent: Any,
|
||||
*,
|
||||
user: str,
|
||||
model: str,
|
||||
endpoint: str,
|
||||
max_tokens: int,
|
||||
) -> Tuple[str, List[Dict[str, Any]], int, int, float]:
|
||||
"""One orchestrator turn. Returns (text, tool_calls, p_tok, c_tok, cost).
|
||||
|
||||
``tool_calls`` is a list of ``{"name", "input"}`` dicts.
|
||||
"""
|
||||
endpoint = endpoint.lower()
|
||||
if endpoint == "anthropic":
|
||||
import anthropic
|
||||
|
||||
client = anthropic.Anthropic(timeout=600.0, max_retries=12)
|
||||
kwargs: Dict[str, Any] = dict(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
messages=[{"role": "user", "content": user}],
|
||||
tools=anthropic_tools(),
|
||||
)
|
||||
if supports_temperature(model):
|
||||
kwargs["temperature"] = 1.0
|
||||
msg = client.messages.create(**kwargs)
|
||||
text = "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
|
||||
tool_calls = [
|
||||
{"name": b.name, "input": dict(b.input or {})}
|
||||
for b in msg.content
|
||||
if getattr(b, "type", "") == "tool_use"
|
||||
]
|
||||
p = getattr(msg.usage, "input_tokens", 0)
|
||||
c = getattr(msg.usage, "output_tokens", 0)
|
||||
|
||||
elif endpoint == "openai":
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(timeout=600.0)
|
||||
kwargs = dict(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": user}],
|
||||
tools=openai_tools(),
|
||||
tool_choice="auto",
|
||||
)
|
||||
if is_gpt5_family(model):
|
||||
kwargs["max_completion_tokens"] = max_tokens
|
||||
kwargs["temperature"] = 1.0
|
||||
else:
|
||||
kwargs["max_tokens"] = max_tokens
|
||||
kwargs["temperature"] = 1.0
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
choice = resp.choices[0].message
|
||||
text = choice.content or ""
|
||||
tool_calls = []
|
||||
for tc in getattr(choice, "tool_calls", None) or []:
|
||||
try:
|
||||
args = json.loads(tc.function.arguments or "{}")
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
tool_calls.append({"name": tc.function.name, "input": args})
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
else:
|
||||
raise ValueError(
|
||||
f"orchestrator endpoint {endpoint!r} unsupported — route the "
|
||||
"orchestrator through anthropic/openai (set method_cfg."
|
||||
"orchestrator_endpoint)."
|
||||
)
|
||||
|
||||
cost = agent.cost_usd(model, p, c)
|
||||
agent.record_trace_event({
|
||||
"kind": "skillorchestra_orchestrate",
|
||||
"model": model,
|
||||
"endpoint": endpoint,
|
||||
"prompt": user,
|
||||
"response": text,
|
||||
"tool_calls": tool_calls,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
})
|
||||
return text, tool_calls, p, c, cost
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Context assembly — eval_frames.py:1305-1351
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _build_context(
|
||||
doc_list: List[Tuple[str, str]],
|
||||
code_list: List[Tuple[str, str]],
|
||||
attempt_list: List[Tuple[str, str]],
|
||||
*,
|
||||
char_cap: int,
|
||||
) -> str:
|
||||
parts: List[str] = []
|
||||
if doc_list:
|
||||
blk = ["## Retrieved Information"]
|
||||
for i, (query, txt) in enumerate(doc_list):
|
||||
blk.append(f"### Search {i + 1} — query: {query}\n{txt}")
|
||||
parts.append("\n\n".join(blk))
|
||||
if code_list:
|
||||
blk = ["## Code Execution Results"]
|
||||
for i, (code, out) in enumerate(code_list):
|
||||
blk.append(
|
||||
f"### Code {i + 1}\n```python\n{code}\n```\n"
|
||||
f"Output:\n{out if out else '(no output)'}"
|
||||
)
|
||||
parts.append("\n\n".join(blk))
|
||||
if attempt_list:
|
||||
blk = ["## Previous Answer Attempts"]
|
||||
for who, ans in attempt_list:
|
||||
blk.append(f"- {who}: {ans}")
|
||||
parts.append("\n".join(blk))
|
||||
ctx = "\n\n".join(parts)
|
||||
if len(ctx) > char_cap:
|
||||
# Keep the tail — most recent docs/code/attempts matter most.
|
||||
ctx = "...[earlier context truncated]...\n" + ctx[-char_cap:]
|
||||
return ctx
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main loop — eval_frames.py:run_single
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_orchestrator(
|
||||
agent: Any,
|
||||
problem: str,
|
||||
*,
|
||||
cfg: Dict[str, Any],
|
||||
handbook: Optional[StageSkillHandbook],
|
||||
strategy: str,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""Run the eval orchestrator on one problem. Returns (answer, metadata)."""
|
||||
max_rounds = int(cfg.get("max_rounds", 6))
|
||||
char_cap = int(cfg.get("context_char_cap", 24000))
|
||||
retriever_url = cfg.get("retriever_url")
|
||||
code_timeout = int(cfg.get("code_timeout_s", 60))
|
||||
answer_max_tokens = int(cfg.get("answer_max_tokens", 40000))
|
||||
ws_max_uses = int(cfg.get("web_search_max_uses", 5))
|
||||
|
||||
# The orchestrator model: a fixed model per run (the original's
|
||||
# MODEL_NAME). Defaults to the cell's cloud model when that endpoint
|
||||
# supports tool calls, else Opus. ``router_model`` / ``router_endpoint``
|
||||
# are accepted as back-compat aliases (pre-restructure cfg key names).
|
||||
orch_endpoint = (cfg.get("orchestrator_endpoint")
|
||||
or cfg.get("router_endpoint")
|
||||
or agent._cloud_endpoint).lower()
|
||||
orch_model = (cfg.get("orchestrator_model")
|
||||
or cfg.get("router_model")
|
||||
or agent._cloud_model)
|
||||
if orch_endpoint not in ("anthropic", "openai"):
|
||||
orch_endpoint, orch_model = "anthropic", "claude-opus-4-7"
|
||||
orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096))
|
||||
|
||||
pool = build_pool(
|
||||
local_model=agent._local_model,
|
||||
local_endpoint=agent._local_endpoint,
|
||||
cloud_model=agent._cloud_model,
|
||||
cloud_endpoint=agent._cloud_endpoint,
|
||||
overrides=cfg.get("model_pool"),
|
||||
)
|
||||
|
||||
doc_list: List[Tuple[str, str]] = []
|
||||
code_list: List[Tuple[str, str]] = []
|
||||
attempt_list: List[Tuple[str, str]] = []
|
||||
route_log: List[Dict[str, Any]] = []
|
||||
|
||||
tokens_local = 0
|
||||
tokens_cloud = 0
|
||||
cost_usd = 0.0
|
||||
tool_calls_n = 0
|
||||
web_uses = 0
|
||||
final_pred = ""
|
||||
used_rounds = 0
|
||||
|
||||
def _route(stage: str, tool_alias: Optional[str], orch_text: str) -> str:
|
||||
"""Resolve the worker alias for ``stage`` via the routing strategy."""
|
||||
if handbook is not None and strategy != "none":
|
||||
sa = parse_skill_analysis(orch_text)
|
||||
rr = get_routing_strategy(strategy, handbook).select_model(
|
||||
stage, sa, tool_call_model=tool_alias,
|
||||
)
|
||||
return rr.model_alias
|
||||
return tool_alias or _STAGE_DEFAULT_ALIAS[stage]
|
||||
|
||||
for step in range(max_rounds):
|
||||
used_rounds = step + 1
|
||||
is_last = step == max_rounds - 1
|
||||
context_str = _build_context(
|
||||
doc_list, code_list, attempt_list, char_cap=char_cap,
|
||||
)
|
||||
|
||||
if handbook is not None and strategy != "none":
|
||||
user = build_skill_orchestrator_prompt(
|
||||
problem=problem,
|
||||
context_str=context_str,
|
||||
strategy=strategy,
|
||||
handbook=handbook,
|
||||
)
|
||||
else:
|
||||
user = (
|
||||
f"Problem: {problem}\n\n{context_str}\n\n"
|
||||
"Choose an appropriate tool."
|
||||
)
|
||||
|
||||
text, tcalls, p, c, ocost = _orchestrate_step(
|
||||
agent, user=user, model=orch_model,
|
||||
endpoint=orch_endpoint, max_tokens=orch_max_tokens,
|
||||
)
|
||||
tokens_cloud += p + c
|
||||
cost_usd += ocost
|
||||
|
||||
# The orchestrator may answer directly in <answer> tags.
|
||||
if not tcalls and "<answer>" in text and "</answer>" in text:
|
||||
final_pred = text.split("<answer>")[-1].split("</answer>")[0].strip()
|
||||
break
|
||||
|
||||
# Last round: force the answer tool (eval_frames.py:1373-1380).
|
||||
if is_last:
|
||||
ans_alias = None
|
||||
for tc in tcalls:
|
||||
if tc["name"] == "answer":
|
||||
ans_alias = (tc.get("input") or {}).get("model")
|
||||
tcalls = [{"name": "answer", "input": {"model": ans_alias or "answer-1"}}]
|
||||
elif not tcalls:
|
||||
# No tool, no answer — record the text and continue.
|
||||
if text.strip():
|
||||
attempt_list.append(("orchestrator", text.strip()[:2000]))
|
||||
continue
|
||||
|
||||
finish = False
|
||||
for tc in tcalls:
|
||||
tool = tc["name"]
|
||||
tool_alias = (tc.get("input") or {}).get("model")
|
||||
stage = _TOOL_STAGE.get(tool, "answer")
|
||||
chosen_alias = _route(stage, tool_alias, text)
|
||||
spec: ModelSpec = pool.get(chosen_alias) or pool[
|
||||
_STAGE_DEFAULT_ALIAS[stage]
|
||||
]
|
||||
route_log.append({
|
||||
"step": step,
|
||||
"tool": tool,
|
||||
"orchestrator_alias": tool_alias,
|
||||
"routed_alias": chosen_alias,
|
||||
"routed_model": spec.model,
|
||||
"is_local": spec.is_local,
|
||||
})
|
||||
tool_calls_n += 1
|
||||
|
||||
if tool == "search":
|
||||
res = run_search(
|
||||
agent, spec, context_str=context_str, problem=problem,
|
||||
retriever_url=retriever_url, web_search_max_uses=ws_max_uses,
|
||||
)
|
||||
docs = res["search_results_data"]
|
||||
joined = "\n---\n".join(d for d in docs if d)[:char_cap]
|
||||
doc_list.append((res["query"], joined or "(no results)"))
|
||||
web_uses += res.get("web_search_uses", 0)
|
||||
elif tool in ("enhance_reasoning", "code"):
|
||||
res = run_code(
|
||||
agent, spec, context_str=context_str, problem=problem,
|
||||
bash_timeout_s=code_timeout,
|
||||
)
|
||||
code_list.append((res["generated_code"], res["exec_result"]))
|
||||
else: # answer
|
||||
res = run_answer(
|
||||
agent, spec, context_str=context_str, problem=problem,
|
||||
max_tokens=answer_max_tokens,
|
||||
)
|
||||
final_pred = res["pred"]
|
||||
attempt_list.append((res["alias"], final_pred))
|
||||
finish = True
|
||||
|
||||
if res["is_local"]:
|
||||
tokens_local += res["tokens_in"] + res["tokens_out"]
|
||||
else:
|
||||
tokens_cloud += res["tokens_in"] + res["tokens_out"]
|
||||
cost_usd += res["cost_usd"]
|
||||
|
||||
if finish:
|
||||
break
|
||||
|
||||
agent.record_trace_event({
|
||||
"kind": "skillorchestra_route_log",
|
||||
"strategy": strategy,
|
||||
"rounds_used": used_rounds,
|
||||
"routes": route_log,
|
||||
})
|
||||
|
||||
meta = {
|
||||
"tokens_local": tokens_local,
|
||||
"tokens_cloud": tokens_cloud,
|
||||
"cost_usd": cost_usd,
|
||||
"turns": used_rounds,
|
||||
"tool_calls": tool_calls_n,
|
||||
"web_search_uses": web_uses,
|
||||
"traces": {
|
||||
"strategy": strategy,
|
||||
"handbook_loaded": handbook is not None,
|
||||
"rounds_used": used_rounds,
|
||||
"routes": route_log,
|
||||
},
|
||||
}
|
||||
return final_pred, meta
|
||||
@@ -0,0 +1,151 @@
|
||||
"""Model-alias pool for the SkillOrchestra eval orchestrator.
|
||||
|
||||
The original SkillOrchestra (``config/models.py`` + ``config/pool_config.json``)
|
||||
maps stage aliases — ``search-1/2/3``, ``reasoner-1/2/3``,
|
||||
``answer-1/2/3/4``, ``answer-math-1/2`` — onto a pool of 6+ models served
|
||||
via SGLang. OpenJarvis runs a 2-model world (one local vLLM student + one
|
||||
cloud model), so the default pool *collapses* the alias tiers onto
|
||||
local/cloud by cost rank: the dearer ``-1`` / ``-2`` aliases (and
|
||||
``answer-math-1``) route to the cloud model, the cheaper ``-3`` / ``-4``
|
||||
aliases (and ``answer-math-2``) route to the local model. This mirrors
|
||||
``stage_router.WeightedAverageStrategy.COST_TIERS``.
|
||||
|
||||
A cell overrides any alias through ``method_cfg.model_pool``::
|
||||
|
||||
method_cfg.model_pool = {
|
||||
"search-1" = { model = "claude-opus-4-7", endpoint = "anthropic" },
|
||||
"search-3" = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" },
|
||||
...
|
||||
}
|
||||
|
||||
``endpoint`` is ``anthropic`` / ``openai`` / ``gemini`` for a cloud model,
|
||||
or an OpenAI-compatible base URL (``http://...``) for a local vLLM model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
# Stage -> ordered alias list. Matches stage_router._get_models_for_stage
|
||||
# and orchestration/tools.json exactly.
|
||||
STAGE_ALIASES: Dict[str, List[str]] = {
|
||||
"search": ["search-1", "search-2", "search-3"],
|
||||
"reasoning": ["reasoner-1", "reasoner-2", "reasoner-3"],
|
||||
"answer": ["answer-1", "answer-2", "answer-3", "answer-4",
|
||||
"answer-math-1", "answer-math-2"],
|
||||
}
|
||||
|
||||
# Every alias the orchestrator can emit, flat.
|
||||
ALL_ALIASES: List[str] = [a for aliases in STAGE_ALIASES.values() for a in aliases]
|
||||
|
||||
# Default tier: which aliases collapse onto the cloud model vs the local
|
||||
# model. Dearer ``-1``/``-2`` (+ answer-math-1) -> cloud; cheaper -> local.
|
||||
_CLOUD_ALIASES = {
|
||||
"search-1", "search-2",
|
||||
"reasoner-1", "reasoner-2",
|
||||
"answer-1", "answer-2", "answer-math-1",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
"""A resolved alias: which concrete model on which endpoint."""
|
||||
|
||||
alias: str
|
||||
model: str
|
||||
endpoint: str # "anthropic" | "openai" | "gemini" | "http://..."
|
||||
kind: str # "cloud" | "local"
|
||||
|
||||
@property
|
||||
def is_local(self) -> bool:
|
||||
return self.kind == "local"
|
||||
|
||||
|
||||
def _endpoint_kind(endpoint: str) -> str:
|
||||
return "local" if endpoint.startswith("http") else "cloud"
|
||||
|
||||
|
||||
def build_pool(
|
||||
*,
|
||||
local_model: Optional[str],
|
||||
local_endpoint: Optional[str],
|
||||
cloud_model: str,
|
||||
cloud_endpoint: str,
|
||||
overrides: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
) -> Dict[str, ModelSpec]:
|
||||
"""Resolve every alias to a :class:`ModelSpec`.
|
||||
|
||||
Default mapping collapses the alias tiers onto the cell's local/cloud
|
||||
pair; ``overrides`` (from ``method_cfg.model_pool``) wins per alias.
|
||||
"""
|
||||
pool: Dict[str, ModelSpec] = {}
|
||||
have_local = bool(local_model and local_endpoint)
|
||||
for alias in ALL_ALIASES:
|
||||
to_cloud = alias in _CLOUD_ALIASES or not have_local
|
||||
if to_cloud:
|
||||
pool[alias] = ModelSpec(alias, cloud_model, cloud_endpoint, "cloud")
|
||||
else:
|
||||
pool[alias] = ModelSpec(
|
||||
alias, local_model, local_endpoint, "local" # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
for alias, spec in (overrides or {}).items():
|
||||
if alias not in pool:
|
||||
continue
|
||||
model = spec.get("model")
|
||||
endpoint = spec.get("endpoint")
|
||||
if not model or not endpoint:
|
||||
continue
|
||||
pool[alias] = ModelSpec(alias, model, endpoint, _endpoint_kind(endpoint))
|
||||
return pool
|
||||
|
||||
|
||||
def call_alias(
|
||||
agent: Any,
|
||||
spec: ModelSpec,
|
||||
*,
|
||||
user: str,
|
||||
system: Optional[str] = None,
|
||||
max_tokens: int = 8000,
|
||||
temperature: float = 1.0,
|
||||
) -> Tuple[str, int, int, float]:
|
||||
"""Single text-generation call through a resolved alias.
|
||||
|
||||
Returns ``(text, tokens_in, tokens_out, cost_usd)``. Dispatches to the
|
||||
right :class:`LocalCloudAgent` SDK helper by endpoint. ``temperature``
|
||||
defaults to 1.0 — the value the original ``call_tool`` uses for every
|
||||
worker call (``eval_frames.py:657``).
|
||||
"""
|
||||
if spec.is_local:
|
||||
text, p, c = agent._call_vllm(
|
||||
spec.model,
|
||||
spec.endpoint,
|
||||
user=user,
|
||||
system=system,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
enable_thinking=False,
|
||||
trace_role="local",
|
||||
)
|
||||
return text, p, c, 0.0
|
||||
|
||||
ep = spec.endpoint.lower()
|
||||
if ep == "anthropic":
|
||||
text, p, c, _ = agent._call_anthropic(
|
||||
spec.model, user=user, system=system,
|
||||
max_tokens=max_tokens, temperature=temperature, trace_role="cloud",
|
||||
)
|
||||
elif ep == "openai":
|
||||
text, p, c = agent._call_openai(
|
||||
spec.model, user=user, system=system,
|
||||
max_tokens=max_tokens, temperature=temperature, trace_role="cloud",
|
||||
)
|
||||
elif ep == "gemini":
|
||||
text, p, c = agent._call_gemini(
|
||||
spec.model, user=user, system=system,
|
||||
max_tokens=max_tokens, temperature=temperature, trace_role="cloud",
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported pool endpoint: {spec.endpoint!r}")
|
||||
return text, p, c, agent.cost_usd(spec.model, p, c)
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Prompt templates for SkillOrchestra.
|
||||
|
||||
Centralized prompts for:
|
||||
- eval_orchestrator: FRAMES orchestrator (search/code/answer)
|
||||
- learning: handbook discovery, refinement, profiler
|
||||
- model_routing: QA benchmarks (skill-based and baseline routing)
|
||||
"""
|
||||
|
||||
from .eval_orchestrator import (
|
||||
build_skill_orchestrator_prompt,
|
||||
format_baseline_tool_info,
|
||||
SKILL_ORCHESTRATOR_PROMPT,
|
||||
SKILL_ANALYSIS_ORCHESTRATOR_PROMPT,
|
||||
)
|
||||
from .learning import (
|
||||
SKILL_DISCOVERY_PROMPT,
|
||||
AGENT_ORCHESTRATION_DISCOVERY_PROMPT,
|
||||
SKILL_IDENTIFICATION_PROMPT,
|
||||
MODE_INSIGHT_PROMPT,
|
||||
PROFILE_SUMMARY_PROMPT,
|
||||
SKILL_SPLIT_PROMPT,
|
||||
SKILL_MERGE_PROMPT,
|
||||
AGENT_ORCHESTRATION_SPLIT_PROMPT,
|
||||
AGENT_ORCHESTRATION_MERGE_PROMPT,
|
||||
FAILURE_DRIVEN_REFINEMENT_PROMPT,
|
||||
)
|
||||
from .model_routing import SKILL_ANALYSIS_PROMPT, BASELINE_PROMPT
|
||||
|
||||
__all__ = [
|
||||
# Eval orchestrator
|
||||
"build_skill_orchestrator_prompt",
|
||||
"format_baseline_tool_info",
|
||||
"SKILL_ORCHESTRATOR_PROMPT",
|
||||
"SKILL_ANALYSIS_ORCHESTRATOR_PROMPT",
|
||||
# Learning
|
||||
"SKILL_DISCOVERY_PROMPT",
|
||||
"AGENT_ORCHESTRATION_DISCOVERY_PROMPT",
|
||||
"SKILL_IDENTIFICATION_PROMPT",
|
||||
"MODE_INSIGHT_PROMPT",
|
||||
"PROFILE_SUMMARY_PROMPT",
|
||||
"SKILL_SPLIT_PROMPT",
|
||||
"SKILL_MERGE_PROMPT",
|
||||
"AGENT_ORCHESTRATION_SPLIT_PROMPT",
|
||||
"AGENT_ORCHESTRATION_MERGE_PROMPT",
|
||||
"FAILURE_DRIVEN_REFINEMENT_PROMPT",
|
||||
# Model routing
|
||||
"SKILL_ANALYSIS_PROMPT",
|
||||
"BASELINE_PROMPT",
|
||||
]
|
||||
@@ -0,0 +1,233 @@
|
||||
"""
|
||||
Eval orchestrator prompts for skill-based agent orchestration.
|
||||
"""
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
# =============================================================================
|
||||
# Baseline Tool Info
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def format_baseline_tool_info() -> str:
|
||||
"""Format baseline tool information for the orchestrator prompt."""
|
||||
return """- Tool: search, Models: search-1 ($10/M output), search-2 ($2/M output), search-3 ($0.8/M output)
|
||||
Description: Search for missing information
|
||||
- Tool: code|enhance_reasoning, Models: reasoner-1 ($10/M output), reasoner-2 ($2/M output), reasoner-3 ($0.8/M output)
|
||||
Description: Write and execute Python code to solve the problem
|
||||
- Tool: answer, Models: answer-1 ($10/M output), answer-2 ($2/M output), answer-3 ($0.9/M output), answer-4 ($0.8/M output), answer-math-1 ($0.9/M output), answer-math-2 ($0.2/M output)
|
||||
Description: Extract the final answer if you think you have enough information to answer the problem"""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Skill-Enhanced Orchestrator Prompt (router_decides strategy)
|
||||
# =============================================================================
|
||||
|
||||
SKILL_ORCHESTRATOR_PROMPT = """You are a skill-based orchestrator for multi-step question answering. Choose the best tool and model for each step.
|
||||
|
||||
## Available Tools and Models (Baseline)
|
||||
|
||||
{baseline_tool_info}
|
||||
|
||||
## Learned Skill Definitions
|
||||
|
||||
### Search Skills
|
||||
{search_skills}
|
||||
|
||||
### Code|Enhance Reasoning Skills
|
||||
{reasoning_skills}
|
||||
|
||||
### Answer Skills
|
||||
{answer_skills}
|
||||
|
||||
## Model Performance (learned from validation)
|
||||
|
||||
### Search Models
|
||||
{search_model_performance}
|
||||
|
||||
### Code|Enhance Reasoning Models
|
||||
{reasoning_model_performance}
|
||||
|
||||
### Answer Models
|
||||
{answer_model_performance}
|
||||
|
||||
## Your Task
|
||||
1. Analyze the problem and current context
|
||||
2. Identify which skills are needed for the next step
|
||||
3. Choose the appropriate tool (search/enhance_reasoning/answer)
|
||||
4. Select the best model for that tool based on skill match and cost
|
||||
|
||||
Consider cost-efficiency: if multiple models can handle it, prefer cheaper ones.
|
||||
|
||||
You must first reason inside <think>...</think> about:
|
||||
- What information is missing or what computation is needed
|
||||
- Which skills from the catalog are required
|
||||
- Which model is best suited based on performance data
|
||||
|
||||
**IMPORTANT**: When calling a tool, you MUST specify the model parameter using the model alias (e.g., "answer-1", "search-1", "reasoner-1"). Use the exact model names from the Available Models section above.
|
||||
|
||||
Problem: {problem}
|
||||
|
||||
{context_str}
|
||||
|
||||
Choose an appropriate tool."""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Skill Analysis Prompt (for weighted_avg, analyze_model_decide strategies)
|
||||
# =============================================================================
|
||||
|
||||
SKILL_ANALYSIS_ORCHESTRATOR_PROMPT = """You are a skill-based orchestrator for multi-step question answering. You select the best tool (search|code|answer) and model by analyzing required skills.
|
||||
|
||||
## Problem to Solve
|
||||
{problem}
|
||||
|
||||
## Current Context
|
||||
{context_str}
|
||||
|
||||
---
|
||||
|
||||
## Quick Reference: What You Need to Do
|
||||
|
||||
**CRITICAL REQUIREMENT**: Before making ANY tool call, you MUST:
|
||||
1. Inside <think>...</think>, analyze required skills and output in <skill_analysis> tags
|
||||
2. Then choose the appropriate tool with the selected model based on the skill analysis.
|
||||
|
||||
**The context above may be long - scroll back to see the problem and context, then follow the instructions below.**
|
||||
|
||||
---
|
||||
|
||||
## Available Tools and Models (Baseline)
|
||||
|
||||
{baseline_tool_info}
|
||||
|
||||
## Learned Skill Definitions
|
||||
|
||||
### Search Skills
|
||||
{search_skills}
|
||||
|
||||
### Reasoning Skills
|
||||
{reasoning_skills}
|
||||
|
||||
### Answer Skills
|
||||
{answer_skills}
|
||||
Use general performance of answer models to select the best model for the answer stage if you think you have enough information to answer the problem.
|
||||
|
||||
## Model Performance (learned from validation)
|
||||
|
||||
### Search Models
|
||||
{search_model_performance}
|
||||
|
||||
### Reasoning Models
|
||||
{reasoning_model_performance}
|
||||
|
||||
### Answer Models
|
||||
{answer_model_performance}
|
||||
---
|
||||
|
||||
## Detailed Instructions
|
||||
|
||||
**STEP 1 - REQUIRED**: Based on the Problem and Context shown at the top, think about what should be the next stage (search|code|answer).
|
||||
Search stage is to find missing information that you think is needed to answer the problem.
|
||||
Code stage is to write and execute Python code to solve the problem.
|
||||
Answer stage is to synthesize all gathered information into a final answer.
|
||||
|
||||
**STEP 2 - REQUIRED FORMAT**:
|
||||
After deciding the next stage, analyze the skills needed for the next stage and provide the detailed skill analysis needed for the next stage.
|
||||
Reason inside <think>...</think> about why these skills are needed and their relative importance. We will use this skill analysis to select the best model for the next stage.
|
||||
Then output your analysis in the following format inside <skill_analysis> tags:
|
||||
<skill_analysis>
|
||||
{{ "required_skills": [ {{"skill_id": "skill.id", "percentage": 50}}, {{"skill_id": "skill.id", "percentage": 30}}, ... ], "reasoning": "Brief explanation of why these skills are needed" }}
|
||||
</skill_analysis>
|
||||
|
||||
**STEP 3**: Choose the appropriate tool with the selected model based on the skill analysis.
|
||||
|
||||
---
|
||||
|
||||
## Final Reminders
|
||||
|
||||
**CRITICAL**: The <skill_analysis> block is MANDATORY and must appear BEFORE your tool call. Without it, the routing system cannot function properly.
|
||||
|
||||
**IMPORTANT**: When calling a tool, you MUST specify the model parameter using the model alias (e.g., "answer-1", "search-1", "reasoner-1"). Use the exact model names from the Available Models section above.
|
||||
|
||||
Now, based on the Problem and Context shown at the top, analyze what should be the next stage (search|code|answer), provide the detailed skill analysis needed for the next stage in the <skill_analysis> tags and then choose an appropriate tool.
|
||||
"""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Prompt Builder
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def build_skill_orchestrator_prompt(
|
||||
problem: str,
|
||||
context_str: str,
|
||||
strategy: str = "router_decides",
|
||||
handbook: Any = None,
|
||||
search_skills: Optional[str] = None,
|
||||
reasoning_skills: Optional[str] = None,
|
||||
answer_skills: Optional[str] = None,
|
||||
search_model_performance: Optional[str] = None,
|
||||
reasoning_model_performance: Optional[str] = None,
|
||||
answer_model_performance: Optional[str] = None,
|
||||
baseline_tool_info: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build enhanced orchestrator prompt with skill catalog and model performance.
|
||||
|
||||
Args:
|
||||
problem: The question/problem to solve
|
||||
context_str: Current context (documents, code results, etc.)
|
||||
strategy: Routing strategy - "router_decides" or "analyze_model_decide" etc.
|
||||
handbook: SkillHandbook object with format_skills(stage/mode) and format_model_performance(stage/mode)
|
||||
search_skills: Override - skill definitions for search stage
|
||||
reasoning_skills: Override - skill definitions for reasoning stage
|
||||
answer_skills: Override - skill definitions for answer stage
|
||||
search_model_performance: Override - model performance for search
|
||||
reasoning_model_performance: Override - model performance for reasoning
|
||||
answer_model_performance: Override - model performance for answer
|
||||
baseline_tool_info: Override - baseline tool descriptions
|
||||
|
||||
Returns:
|
||||
Formatted prompt string
|
||||
"""
|
||||
if baseline_tool_info is None:
|
||||
baseline_tool_info = format_baseline_tool_info()
|
||||
|
||||
if handbook:
|
||||
if search_skills is None:
|
||||
search_skills = handbook.format_skills("search")
|
||||
if reasoning_skills is None:
|
||||
reasoning_skills = handbook.format_skills("code")
|
||||
if answer_skills is None:
|
||||
answer_skills = handbook.format_skills("answer")
|
||||
if search_model_performance is None:
|
||||
search_model_performance = handbook.format_model_performance("search")
|
||||
if reasoning_model_performance is None:
|
||||
reasoning_model_performance = handbook.format_model_performance("code")
|
||||
if answer_model_performance is None:
|
||||
answer_model_performance = handbook.format_model_performance("answer")
|
||||
|
||||
search_skills = search_skills or "No skills defined"
|
||||
reasoning_skills = reasoning_skills or "No skills defined"
|
||||
answer_skills = answer_skills or "No skills defined"
|
||||
search_model_performance = search_model_performance or "No performance data"
|
||||
reasoning_model_performance = reasoning_model_performance or "No performance data"
|
||||
answer_model_performance = answer_model_performance or "No performance data"
|
||||
|
||||
if strategy == "router_decides":
|
||||
template = SKILL_ORCHESTRATOR_PROMPT
|
||||
else:
|
||||
template = SKILL_ANALYSIS_ORCHESTRATOR_PROMPT
|
||||
|
||||
return template.format(
|
||||
baseline_tool_info=baseline_tool_info,
|
||||
search_skills=search_skills,
|
||||
reasoning_skills=reasoning_skills,
|
||||
answer_skills=answer_skills,
|
||||
search_model_performance=search_model_performance,
|
||||
reasoning_model_performance=reasoning_model_performance,
|
||||
answer_model_performance=answer_model_performance,
|
||||
problem=problem,
|
||||
context_str=context_str if context_str else "(No context yet)",
|
||||
)
|
||||
@@ -0,0 +1,448 @@
|
||||
"""
|
||||
LLM prompt templates for Skill Handbook learning.
|
||||
|
||||
All prompts used during the learning pipeline:
|
||||
- Skill discovery from trajectory contrast
|
||||
- Skill identification for a query
|
||||
- Mode-level insight distillation
|
||||
- Profile summarization (strengths/weaknesses)
|
||||
- Skill split/merge analysis
|
||||
- Failure-driven refinement
|
||||
"""
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1a: Skill Discovery - Model Routing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SKILL_DISCOVERY_PROMPT = """You are a skill taxonomist analyzing QA problems and model performance data to discover what skills are needed for effective model routing.
|
||||
|
||||
## Task
|
||||
Analyze the sample problems below along with per-model success/failure data. Your goal: Propose a HIERARCHICAL skill taxonomy with:
|
||||
1. HIGH-LEVEL CATEGORIES (3-6): Broad skill areas that differentiate problems
|
||||
2. FINE-GRAINED SKILLS (2-4 per category): Specific capabilities within each category
|
||||
|
||||
## Requirements
|
||||
- Skills should capture what makes problems DIFFERENT from each other
|
||||
- Skills should explain what makes MODELS perform DIFFERENTLY on those problems
|
||||
- Skills should be SPECIFIC and MEASURABLE (not vague like "intelligence" or "reasoning")
|
||||
- Include INDICATORS (keywords/patterns that suggest a skill is needed)
|
||||
- Include EXAMPLES from the sample problems
|
||||
- Skill IDs should follow the pattern: category_name.specific_skill_name
|
||||
|
||||
## Sample Problems with Model Performance
|
||||
{sample_problems}
|
||||
|
||||
## Contrastive Evidence (where models disagree)
|
||||
{contrastive_evidence}
|
||||
|
||||
## Existing Skills (avoid duplicates)
|
||||
{existing_skills}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object with:
|
||||
{{
|
||||
"categories": [
|
||||
{{
|
||||
"name": "category_name",
|
||||
"description": "What this category covers",
|
||||
"skills": [
|
||||
{{
|
||||
"skill_id": "category_name.skill_name",
|
||||
"name": "Human-readable Skill Name",
|
||||
"description": "What this specific skill involves and why models differ on it",
|
||||
"indicators": ["keyword1", "pattern2", "phrase3"],
|
||||
"examples": ["Example query requiring this skill"],
|
||||
"mode": "answer"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
Aim for a taxonomy that covers the FULL DIVERSITY of the sample problems, not just one narrow topic. If problems span temporal facts, entity lookups, numeric data, relational facts, etc., the taxonomy should reflect all of those."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1a: Skill Discovery - Agent Orchestration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
AGENT_ORCHESTRATION_DISCOVERY_PROMPT = """You are a skill taxonomist analyzing QA problems to discover the underlying skills required to solve them.
|
||||
|
||||
## Stage Information
|
||||
We have 3 stages in our pipeline:
|
||||
- **search**: Web search to retrieve factual information (tool: search)
|
||||
- **code**: Code generation and execution for calculations (tool: enhance_reasoning)
|
||||
- **answer**: Generate final answer from context (tool: answer)
|
||||
|
||||
## Sample Problems
|
||||
{sample_problems}
|
||||
|
||||
## Task
|
||||
Analyze these problems and DISCOVER what skills are needed.
|
||||
Propose a HIERARCHICAL skill taxonomy with:
|
||||
1. HIGH-LEVEL CATEGORIES (3-5): Broad skill areas that differentiate problems
|
||||
2. FINE-GRAINED SKILLS (2-4 per category): Specific capabilities within each category
|
||||
|
||||
## Requirements
|
||||
- Skills should capture what makes problems DIFFERENT and what makes MODELS perform differently
|
||||
- Skills should be SPECIFIC and MEASURABLE (not vague like "intelligence")
|
||||
- Include INDICATORS (keywords/patterns that suggest a skill is needed)
|
||||
- Use hierarchical IDs: stage.category.specific_skill
|
||||
|
||||
## Output Format (JSON)
|
||||
```json
|
||||
{{
|
||||
"categories": [
|
||||
{{
|
||||
"stage": "search|code|answer",
|
||||
"name": "category_name",
|
||||
"description": "What this category covers",
|
||||
"skills": [
|
||||
{{
|
||||
"id": "stage.category.skill_name",
|
||||
"name": "Human Readable Name",
|
||||
"description": "What this specific skill involves",
|
||||
"indicators": ["keyword1", "pattern2", "phrase3"],
|
||||
"examples": ["Example query requiring this skill"]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
Respond with JSON only."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1b: Skill Identification (which skills are active for a query)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SKILL_IDENTIFICATION_PROMPT = """You are an expert at identifying which skills from a catalog are required to handle a given query.
|
||||
|
||||
## Query
|
||||
{query}
|
||||
|
||||
## Ground Truth
|
||||
{ground_truth}
|
||||
|
||||
## Operational Mode
|
||||
{mode}
|
||||
|
||||
## Model Results (all models' outputs and whether they succeeded)
|
||||
Use this contrastive evidence: where models differ in success/failure, the outputs reveal what skills matter for this query.
|
||||
|
||||
{model_results}
|
||||
|
||||
## Available Skills for this Mode
|
||||
{mode_skills}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"active_skills": [
|
||||
{{
|
||||
"skill_id": "the skill id",
|
||||
"weight": 0.0 to 1.0,
|
||||
"reasoning": "brief explanation"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
Weights should sum to approximately 1.0. Only include skills that are genuinely relevant to this specific query and mode. Use the model outputs to infer which skills differentiate successful vs failed attempts."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1b: Mode-level Insight Distillation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MODE_INSIGHT_PROMPT = """You are an expert at analyzing execution patterns to derive reusable routing insights.
|
||||
|
||||
## Task
|
||||
Analyze the execution patterns below and derive mode-level routing insights. These insights should help an orchestrator decide WHEN to use each mode and HOW to transition between modes.
|
||||
|
||||
## Execution Patterns
|
||||
{execution_patterns}
|
||||
|
||||
## Modes
|
||||
{modes}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"insights": [
|
||||
{{
|
||||
"mode": "search|code|answer",
|
||||
"content": "The routing insight as a clear, actionable rule",
|
||||
"insight_type": "transition|usage|constraint",
|
||||
"confidence": 0.0 to 1.0
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
Focus on patterns that generalize across queries, not query-specific observations. Examples:
|
||||
- "If multiple arithmetic operations are needed, switch to code mode instead of search"
|
||||
- "Prefer search-3 for multi-hop queries requiring entity tracking"
|
||||
- "Switch to answer mode once all required facts have been gathered"
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1b: Agent Profile Summarization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PROFILE_SUMMARY_PROMPT = """You are an expert at summarizing agent capabilities from performance data.
|
||||
|
||||
## Agent
|
||||
{agent_id} (model: {model_name}, mode: {mode})
|
||||
|
||||
## Performance Data
|
||||
{performance_data}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"strengths": ["strength1", "strength2"],
|
||||
"weaknesses": ["weakness1", "weakness2"],
|
||||
"routing_signals": ["when to use this agent", "when to avoid"]
|
||||
}}
|
||||
|
||||
Be specific and evidence-based. Reference skill categories where applicable."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 2: Skill Split Analysis
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SKILL_SPLIT_PROMPT = """You are an expert at analyzing whether a skill should be split into finer-grained sub-skills.
|
||||
|
||||
## Skill Under Review
|
||||
ID: {skill_id}
|
||||
Name: {skill_name}
|
||||
Description: {skill_description}
|
||||
Mode: {mode}
|
||||
|
||||
## Evidence for Splitting
|
||||
The agents have highly VARIABLE performance on this skill, suggesting it may conflate distinct capabilities:
|
||||
|
||||
{performance_evidence}
|
||||
|
||||
## Sample Queries Where Agents Disagree
|
||||
{sample_queries}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"should_split": true/false,
|
||||
"rationale": "explanation",
|
||||
"proposed_splits": [
|
||||
{{
|
||||
"skill_id": "mode.category.new_name",
|
||||
"name": "New Skill Name",
|
||||
"description": "What this sub-skill captures",
|
||||
"indicators": ["indicator1", "indicator2"],
|
||||
"distinguishing_feature": "What separates this from sibling skills"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
Only recommend splitting if there is clear evidence that the skill conflates genuinely different capabilities. The splits should be actionable for routing decisions."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 2: Skill Merge Analysis
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SKILL_MERGE_PROMPT = """You are an expert at analyzing whether two skills should be merged.
|
||||
|
||||
## Skills Under Review
|
||||
|
||||
### Skill 1
|
||||
ID: {skill_1_id}
|
||||
Name: {skill_1_name}
|
||||
Description: {skill_1_description}
|
||||
|
||||
### Skill 2
|
||||
ID: {skill_2_id}
|
||||
Name: {skill_2_name}
|
||||
Description: {skill_2_description}
|
||||
|
||||
## Evidence for Merging
|
||||
All agents have statistically INDISTINGUISHABLE performance between these two skills, suggesting they are redundant for routing purposes:
|
||||
|
||||
{performance_evidence}
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"should_merge": true/false,
|
||||
"rationale": "Why merge or not",
|
||||
"merged_skill": {{
|
||||
"skill_id": "mode.category.merged_name",
|
||||
"name": "Merged Skill Name",
|
||||
"description": "Combined description",
|
||||
"indicators": ["combined indicators"]
|
||||
}},
|
||||
"alternative_explanation": "If not merging, explain why they should remain separate"
|
||||
}}
|
||||
|
||||
Only recommend merging if the skills truly capture the same capability from a routing perspective, even if they differ semantically."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 2: Agent Orchestration Split/Merge
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
AGENT_ORCHESTRATION_SPLIT_PROMPT = """You are analyzing whether a skill should be split into more fine-grained skills.
|
||||
|
||||
## Skill to Analyze
|
||||
{skill_definition}
|
||||
|
||||
## Performance Data
|
||||
{performance_data}
|
||||
|
||||
## Sample Queries
|
||||
### High Performance Queries (models succeeded)
|
||||
{high_perf_queries}
|
||||
|
||||
### Low Performance Queries (models failed)
|
||||
{low_perf_queries}
|
||||
|
||||
### Divergent Performance Queries (some models succeeded, others failed)
|
||||
{divergent_queries}
|
||||
|
||||
## Sample Trajectories (if available)
|
||||
### Successful Trajectories
|
||||
{success_trajectories}
|
||||
|
||||
### Failed Trajectories
|
||||
{failure_trajectories}
|
||||
|
||||
## Task
|
||||
Analyze whether this skill should be split:
|
||||
1. Does the skill have high variance across models? (suggests splitting)
|
||||
2. Do different query types show different performance patterns? (suggests splitting)
|
||||
3. Are there clear sub-skills that could be distinguished? (suggests splitting)
|
||||
|
||||
## Output Format (JSON)
|
||||
```json
|
||||
{{
|
||||
"should_split": true/false,
|
||||
"rationale": "Why split or not",
|
||||
"proposed_splits": [
|
||||
{{
|
||||
"skill_id": "stage.category.subskill1",
|
||||
"name": "Sub-skill Name",
|
||||
"description": "What this sub-skill covers",
|
||||
"indicators": ["indicator1", "indicator2"],
|
||||
"distinguishing_feature": "What distinguishes this from sibling skills"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
```
|
||||
|
||||
Respond with JSON only."""
|
||||
|
||||
|
||||
AGENT_ORCHESTRATION_MERGE_PROMPT = """You are analyzing whether two skills should be merged.
|
||||
|
||||
## Skills to Analyze
|
||||
{skills_definitions}
|
||||
|
||||
## Performance Correlation
|
||||
{performance_correlation}
|
||||
|
||||
## Sample Queries
|
||||
### Skill 1 Queries
|
||||
{skill1_queries}
|
||||
|
||||
### Skill 2 Queries
|
||||
{skill2_queries}
|
||||
|
||||
## Task
|
||||
Analyze whether these skills should be merged:
|
||||
1. Do they have nearly identical performance patterns across models? (suggests merge)
|
||||
2. Are they conceptually similar or overlapping? (suggests merge)
|
||||
3. Would merging simplify routing without losing important distinctions? (suggests merge)
|
||||
|
||||
## Output Format (JSON)
|
||||
```json
|
||||
{{
|
||||
"should_merge": true/false,
|
||||
"rationale": "Why merge or not",
|
||||
"merged_skill": {{
|
||||
"skill_id": "stage.category.merged_skill",
|
||||
"name": "Merged Skill Name",
|
||||
"description": "Combined description",
|
||||
"indicators": ["indicator1", "indicator2"]
|
||||
}},
|
||||
"alternative_explanation": "If not merging, explain why they should remain separate"
|
||||
}}
|
||||
```
|
||||
|
||||
Respond with JSON only."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Failure-driven refinement (when skill routing < oracle on training set)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
FAILURE_DRIVEN_REFINEMENT_PROMPT = """You are an expert at analyzing why skill-based model routing fails and how to improve the skill taxonomy.
|
||||
|
||||
## Context
|
||||
We have a skill-based routing system that selects which LLM to call based on identified skills. On the training set, we achieved:
|
||||
- **Oracle accuracy**: {oracle_accuracy:.1%} (best possible — if we always picked the correct model per query)
|
||||
- **Skill-based accuracy**: {skill_accuracy:.1%} (our current routing)
|
||||
|
||||
We failed to achieve oracle-level performance. This suggests either:
|
||||
1. **Missing skills**: Queries require skills not in our catalog
|
||||
2. **Skills too coarse**: Existing skills conflate distinct capabilities and lead to wrong model selection
|
||||
3. **Skill identification gaps**: The router fails to identify the right skills for some query types
|
||||
|
||||
## Current Skill Catalog
|
||||
{skill_catalog}
|
||||
|
||||
## Failed Queries (oracle would have been correct, but we routed wrong)
|
||||
For each failed query we show: the question, what model(s) would have been correct (oracle), what model we routed to, and whether any model got it right.
|
||||
|
||||
{failed_queries}
|
||||
|
||||
## Task
|
||||
Reflect on why the routing failed for these queries. Consider:
|
||||
1. What skills are **missing** from the catalog that would have helped route correctly?
|
||||
2. Which existing skills might need to be **split** into finer-grained sub-skills?
|
||||
3. What **indicators** or patterns in the failed queries suggest new or refined skills?
|
||||
|
||||
## Output Format
|
||||
Return a JSON object:
|
||||
{{
|
||||
"rationale": "Your overall reflection on why routing failed and what the main gaps are",
|
||||
"proposed_new_skills": [
|
||||
{{
|
||||
"skill_id": "category.specific_skill_name",
|
||||
"name": "Human-readable name",
|
||||
"description": "What this skill captures and why it matters for routing",
|
||||
"indicators": ["keyword1", "pattern2", "phrase3"],
|
||||
"example_queries": ["Example from failed queries that would match this skill"]
|
||||
}}
|
||||
],
|
||||
"proposed_splits": [
|
||||
{{
|
||||
"parent_skill_id": "existing.skill.id",
|
||||
"rationale": "Why this skill should be split",
|
||||
"proposed_sub_skills": [
|
||||
{{
|
||||
"skill_id": "existing.new_sub_name",
|
||||
"name": "Sub-skill name",
|
||||
"description": "What this sub-skill captures",
|
||||
"indicators": ["indicator1"],
|
||||
"distinguishing_feature": "What separates this from sibling sub-skills"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
- Only propose new skills or splits that are clearly supported by the failed queries
|
||||
- Skill IDs should follow category.specific_name pattern
|
||||
- Be specific: tie each proposal to concrete failed queries"""
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Prompt templates for model routing (skill-based and baseline)."""
|
||||
|
||||
SKILL_ANALYSIS_PROMPT = """You are a skill-based model router. You are selecting the best model to answer a question by analyzing a question to identify required skills and their importance related to this question.
|
||||
|
||||
## Learned Skill Definitions (from validation)
|
||||
{skill_catalog}
|
||||
|
||||
## Model Performance (learned from validation)
|
||||
{model_performance}
|
||||
|
||||
## Cost Tiers (cheapest to most expensive)
|
||||
- Cheap: Qwen2.5-7B-Instruct, LLaMA-3.1-8B-Instruct, Mistral-7B-Instruct
|
||||
- Medium: Gemma-2-27B-Instruct
|
||||
- Expensive: LLaMA-3.1-70B-Instruct, Mixtral-8x22B-Instruct
|
||||
|
||||
## Task
|
||||
1. First, analyze the question below and identify which skills are needed, along with the percentage/weight of each skill (how important each skill is for answering this question).
|
||||
|
||||
**IMPORTANT: Output your skill analysis FIRST, before any <think> tags.** Use the exact skill_id values from the catalog above (e.g. "disambiguation_and_scope.ambiguous_media_title_resolution").
|
||||
|
||||
<skill_analysis>
|
||||
{{
|
||||
"required_skills": [
|
||||
{{"skill_id": "category.skill_id", "percentage": 50}},
|
||||
{{"skill_id": "category.skill_id", "percentage": 30}},
|
||||
...
|
||||
],
|
||||
"reasoning": "Brief explanation of why these skills are needed"
|
||||
}}
|
||||
</skill_analysis>
|
||||
|
||||
The percentages should sum to approximately 100 (they don't need to be exact, but should reflect relative importance).
|
||||
|
||||
2. After providing the skill analysis, reflect on which model is best suited based on the skills required and model performance data above.
|
||||
|
||||
3. Route to that model using <search> tags and provide final answer in <answer>...</answer>
|
||||
|
||||
Every time you receive new information, you must first conduct reasoning inside <think> ... </think>. \
|
||||
After reasoning, if you find you lack some knowledge, you can call a specialized LLM by writing a query inside <search> LLM-Name:Your-Query </search>. \
|
||||
|
||||
!!! STRICT FORMAT RULES for <search>: !!!
|
||||
+ You MUST replace LLM-Name with the EXACT name of a model selected from [Qwen2.5-7B-Instruct, LLaMA-3.1-8B-Instruct, LLaMA-3.1-70B-Instruct, Mistral-7B-Instruct, Mixtral-8x22B-Instruct, Gemma-2-27B-Instruct]. \
|
||||
+ You MUST replace Your-Query with the EXACT same question as the original question below (DO NOT CHANGE IT). \
|
||||
+ NEVER copy or paste model descriptions into <search>.
|
||||
+ NEVER output the placeholder format <search> LLM-Name:Your-Query </search>. Always replace both parts correctly. \
|
||||
|
||||
Before each LLM call, you MUST explicitly reason inside <think> ... </think> about: \
|
||||
+ Why external information is needed. \
|
||||
+ Which skills from the catalog are required for this question. \
|
||||
+ Which model is best suited based on the model performance data above. \
|
||||
|
||||
When you call an LLM, the response will be returned between <information> and </information>. \
|
||||
You are encouraged to explore and utilize different LLMs to better understand their respective strengths and weaknesses. \
|
||||
|
||||
If you find that no further external knowledge is needed, you can directly provide your final answer to the original question inside <answer> ... </answer>, without additional explanation or illustration. \
|
||||
For example: <answer> Beijing </answer>. \
|
||||
+ Important: You must not output the placeholder text "<answer> and </answer>" alone. \
|
||||
+ You must insert your actual answer between <answer> and </answer>, following the correct format. \
|
||||
+ You must not output the model name or query between <answer> and </answer>. \
|
||||
|
||||
If you think none of the models listed have the necessary skills to answer this question directly, you can route to the model with the highest overall pass rate of models in the pool to get more information.
|
||||
|
||||
Question: {question}
|
||||
"""
|
||||
|
||||
|
||||
BASELINE_PROMPT = """Answer the given question. \
|
||||
Every time you receive new information, you must first conduct reasoning inside <think> ... </think>. \
|
||||
After reasoning, if you find you lack some knowledge, you can call a specialized LLM by writing a query inside <search> LLM-Name:Your-Query </search>. \
|
||||
|
||||
!!! STRICT FORMAT RULES for <search>: !!!
|
||||
+ You MUST replace LLM-Name with the EXACT name of a model selected from [Qwen2.5-7B-Instruct, LLaMA-3.1-8B-Instruct, LLaMA-3.1-70B-Instruct, Mistral-7B-Instruct, Mixtral-8x22B-Instruct, Gemma-2-27B-Instruct]. \
|
||||
+ You MUST replace Your-Query with a CONCRETE QUESTION that helps answer the original question below. \
|
||||
+ NEVER copy or paste model descriptions into <search>.
|
||||
+ NEVER output the placeholder format <search> LLM-Name:Your-Query </search>. Always replace both parts correctly. \
|
||||
|
||||
Before each LLM call, you MUST explicitly reason inside <think> ... </think> about: \
|
||||
+ Why external information is needed. \
|
||||
+ Which model is best suited for answering it, based on the LLMs' abilities (described below). \
|
||||
|
||||
When you call an LLM, the response will be returned between <information> and </information>. \
|
||||
You must not limit yourself to repeatedly calling a single LLM (unless its provided information is consistently the most effective and informative). \
|
||||
You are encouraged to explore and utilize different LLMs to better understand their respective strengths and weaknesses. \
|
||||
It is also acceptable—and recommended—to call different LLMs multiple times for the same input question to gather more comprehensive information. \
|
||||
|
||||
|
||||
#### The Descriptions of Each LLM \
|
||||
|
||||
Qwen2.5-7B-Instruct:\
|
||||
Qwen2.5-7B-Instruct is a powerful Chinese-English instruction-tuned large language model designed for tasks in language, \
|
||||
coding, mathematics, and reasoning. As part of the Qwen2.5 series, it features enhanced knowledge, stronger coding and \
|
||||
math abilities, improved instruction following, better handling of long and structured texts, and supports up to 128K \
|
||||
context tokens. It also offers multilingual capabilities across over 29 languages.\
|
||||
|
||||
|
||||
LLaMA-3.1-8B-Instruct:\
|
||||
LLaMA-3.1-8B-Instruct is an 8-billion-parameter instruction-tuned language model optimized for multilingual dialogue. \
|
||||
It provides strong language understanding, reasoning, and text generation performance, outperforming many open-source \
|
||||
and closed-source models on standard industry benchmarks.\
|
||||
|
||||
|
||||
LLaMA-3.1-70B-Instruct:\
|
||||
LLaMA-3.1-70B-Instruct is a 70-billion-parameter state-of-the-art language model designed for advanced multilingual \
|
||||
dialogue tasks. It excels in language comprehension, complex reasoning, and high-quality text generation, setting a new \
|
||||
standard against both open and closed models in benchmark evaluations.\
|
||||
|
||||
|
||||
Mistral-7B-Instruct:\
|
||||
Mistral-7B-Instruct is a fine-tuned version of the Mistral-7B-v0.3 language model designed to follow instructions, \
|
||||
complete user requests, and generate creative text. It was trained on diverse public conversation datasets to enhance \
|
||||
its ability to handle interactive tasks effectively.\
|
||||
|
||||
|
||||
Mixtral-8x22B-Instruct:\
|
||||
Mixtral-8x22B-Instruct is a cutting-edge sparse Mixture-of-Experts (SMoE) large language model from MistralAI. It \
|
||||
efficiently uses 39B active parameters out of 141B total, delivering high performance at lower costs. The model excels \
|
||||
at following instructions, completing tasks, and generating creative text, with strong skills in multiple languages \
|
||||
(English, French, Italian, German, Spanish), mathematics, and coding. It also supports native function calling and \
|
||||
handles long contexts up to 64K tokens for better information recall.\
|
||||
|
||||
|
||||
Gemma-2-27B-Instruct:\
|
||||
Gemma-2-27B-Instruct is a cutting-edge, instruction-tuned text generation model developed by Google. Built using the \
|
||||
same technology as Gemini, it excels at text understanding, transformation, and code generation. As a lightweight, \
|
||||
decoder-only model with open weights, it is ideal for tasks like question answering, summarization, and reasoning. \
|
||||
Its compact size enables deployment on laptops, desktops, or private cloud setups, making powerful AI more accessible.\
|
||||
|
||||
|
||||
If you find that no further external knowledge is needed, you can directly provide your final answer inside <answer> ... </answer>, without additional explanation or illustration. \
|
||||
For example: <answer> Beijing </answer>. \
|
||||
+ Important: You must not output the placeholder text "<answer> and </answer>" alone. \
|
||||
+ You must insert your actual answer between <answer> and </answer>, following the correct format. \
|
||||
Question: {question}
|
||||
"""
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
Adapter for orchestration eval script.
|
||||
|
||||
Provides: StageSkillHandbook (load from JSON), parse_skill_analysis,
|
||||
get_routing_strategy. Compatible with JSON produced by to_stage_router.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Handbook dataclasses (compatible with stage_router JSON)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class Skill:
|
||||
"""A skill that models can have."""
|
||||
|
||||
skill_id: str
|
||||
name: str
|
||||
description: str
|
||||
stage: str
|
||||
examples: List[str] = field(default_factory=list)
|
||||
discovered_from_problems: List[Dict[str, str]] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "Skill":
|
||||
return cls(
|
||||
skill_id=data.get("skill_id", ""),
|
||||
name=data.get("name", ""),
|
||||
description=data.get("description", ""),
|
||||
stage=data.get("stage", ""),
|
||||
examples=data.get("examples", []),
|
||||
discovered_from_problems=data.get("discovered_from_problems", []),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelProfile:
|
||||
"""Performance profile for a model alias."""
|
||||
|
||||
model_alias: str
|
||||
actual_model: str
|
||||
stage: str
|
||||
skill_scores: Dict[str, float] = field(default_factory=dict)
|
||||
skill_attempts: Dict[str, int] = field(default_factory=dict)
|
||||
skill_successes: Dict[str, int] = field(default_factory=dict)
|
||||
overall_success_rate: float = 0.5
|
||||
total_attempts: int = 0
|
||||
total_successes: int = 0
|
||||
avg_prompt_tokens: float = 0.0
|
||||
avg_completion_tokens: float = 0.0
|
||||
avg_cost_usd: float = 0.0
|
||||
strengths: List[str] = field(default_factory=list)
|
||||
weaknesses: List[str] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "ModelProfile":
|
||||
return cls(
|
||||
model_alias=data.get("model_alias", ""),
|
||||
actual_model=data.get("actual_model", ""),
|
||||
stage=data.get("stage", ""),
|
||||
skill_scores=data.get("skill_scores", {}),
|
||||
skill_attempts=data.get("skill_attempts", {}),
|
||||
skill_successes=data.get("skill_successes", {}),
|
||||
overall_success_rate=data.get("overall_success_rate", 0.5),
|
||||
total_attempts=data.get("total_attempts", 0),
|
||||
total_successes=data.get("total_successes", 0),
|
||||
avg_prompt_tokens=data.get("avg_prompt_tokens", 0.0),
|
||||
avg_completion_tokens=data.get("avg_completion_tokens", 0.0),
|
||||
avg_cost_usd=data.get("avg_cost_usd", 0.0),
|
||||
strengths=data.get("strengths", []),
|
||||
weaknesses=data.get("weaknesses", []),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# StageSkillHandbook
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class StageSkillHandbook:
|
||||
"""Handbook for eval routing. Load from JSON produced by to_stage_router."""
|
||||
|
||||
def __init__(self, load_defaults: bool = True):
|
||||
self.skills: Dict[str, Dict[str, Skill]] = {
|
||||
"search": {},
|
||||
"code": {},
|
||||
"answer": {},
|
||||
}
|
||||
self.model_profiles: Dict[str, ModelProfile] = {}
|
||||
self.usage_patterns: Dict[str, Any] = {"stages": {}, "guidelines": {}, "models": {}, "raw": {}}
|
||||
self.routing_insights: List[str] = []
|
||||
self.learning_history: List[Dict[str, Any]] = []
|
||||
self.version = "1.0.0"
|
||||
self.created_at = ""
|
||||
self.updated_at = ""
|
||||
|
||||
def get_model_skill_scores(self) -> Dict[str, Dict[str, float]]:
|
||||
return {alias: profile.skill_scores for alias, profile in self.model_profiles.items()}
|
||||
|
||||
def get_models_for_stage(self, stage: str) -> List[ModelProfile]:
|
||||
return [p for p in self.model_profiles.values() if p.stage == stage]
|
||||
|
||||
def format_skills(self, stage: str) -> str:
|
||||
catalog_skills = self.skills.get(stage, {})
|
||||
models = self.get_models_for_stage(stage)
|
||||
skills_with_performance = set()
|
||||
for model in models:
|
||||
skills_with_performance.update(model.skill_scores.keys())
|
||||
|
||||
lines = []
|
||||
shown_skills = set()
|
||||
for skill_id, skill in catalog_skills.items():
|
||||
if skill_id in skills_with_performance:
|
||||
shown_skills.add(skill_id)
|
||||
lines.append(f"- {skill_id}: {skill.description}")
|
||||
if skill.examples:
|
||||
lines.append(f" Examples: {', '.join(skill.examples[:2])}")
|
||||
|
||||
orphaned = skills_with_performance - set(catalog_skills.keys())
|
||||
if orphaned:
|
||||
if shown_skills:
|
||||
lines.append("")
|
||||
lines.append("# Additional skills with performance data available:")
|
||||
for skill_id in sorted(orphaned):
|
||||
lines.append(f"- {skill_id}")
|
||||
|
||||
return "\n".join(lines) if lines else "No skills defined"
|
||||
|
||||
def format_model_performance(self, stage: str) -> str:
|
||||
profiles = self.get_models_for_stage(stage)
|
||||
valid_prefixes = {"search": ["search-"], "code": ["reasoner-", "code-"], "answer": ["answer-"]}
|
||||
prefixes = valid_prefixes.get(stage, [])
|
||||
|
||||
lines = []
|
||||
for p in profiles:
|
||||
if not any(p.model_alias.startswith(prefix) for prefix in prefixes):
|
||||
continue
|
||||
has_data = (p.skill_scores and len(p.skill_scores) > 0) or p.strengths or p.weaknesses
|
||||
if p.total_attempts > 0 or has_data:
|
||||
lines.append(f"\n### {p.model_alias} ({p.actual_model})")
|
||||
if p.total_attempts > 0:
|
||||
rate = p.total_successes / p.total_attempts
|
||||
lines.append(f"Overall: {rate:.0%} success ({p.total_successes}/{p.total_attempts})")
|
||||
else:
|
||||
lines.append("Overall: 0% overall")
|
||||
if p.skill_scores:
|
||||
lines.append("Skill scores:")
|
||||
stage_skill_scores = {
|
||||
sid: s
|
||||
for sid, s in p.skill_scores.items()
|
||||
if (sid.split(".")[0] if "." in sid else sid) in ("code", stage)
|
||||
}
|
||||
for skill_id, score in sorted(stage_skill_scores.items(), key=lambda x: x[1], reverse=True):
|
||||
lines.append(f" - {skill_id}: {score:.0%}")
|
||||
if p.strengths:
|
||||
lines.append(f"Strengths: {', '.join(p.strengths[:3])}")
|
||||
if p.weaknesses:
|
||||
lines.append(f"Weaknesses: {', '.join(p.weaknesses[:3])}")
|
||||
|
||||
return "\n".join(lines) if lines else "No model performance data learned yet."
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str) -> "StageSkillHandbook":
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
handbook = cls(load_defaults=False)
|
||||
handbook.version = data.get("version", "1.0.0")
|
||||
handbook.created_at = data.get("created_at", "")
|
||||
handbook.updated_at = data.get("updated_at", "")
|
||||
for stage, skills in data.get("skills", {}).items():
|
||||
if stage not in handbook.skills:
|
||||
handbook.skills[stage] = {}
|
||||
for sid, sdata in skills.items():
|
||||
handbook.skills[stage][sid] = Skill.from_dict(sdata)
|
||||
for alias, pdata in data.get("model_profiles", {}).items():
|
||||
handbook.model_profiles[alias] = ModelProfile.from_dict(pdata)
|
||||
raw_usage = data.get("usage_patterns", {})
|
||||
handbook.usage_patterns = {
|
||||
"stages": raw_usage.get("stages", {}),
|
||||
"guidelines": raw_usage.get("guidelines", {}),
|
||||
"models": raw_usage.get("models", {}),
|
||||
"raw": raw_usage.get("raw", {}),
|
||||
}
|
||||
handbook.learning_history = data.get("learning_history", [])
|
||||
handbook.routing_insights = data.get("routing_insights", [])
|
||||
return handbook
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Skill analysis parsing
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class SkillWeight:
|
||||
skill_id: str
|
||||
percentage: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class SkillAnalysis:
|
||||
stage: str
|
||||
required_skills: List[SkillWeight] = field(default_factory=list)
|
||||
reasoning: str = ""
|
||||
raw_json: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def parse_skill_analysis(output: str) -> Optional[SkillAnalysis]:
|
||||
pattern = r"<skill_analysis>\s*(.*?)\s*</skill_analysis>"
|
||||
match = re.search(pattern, output, re.DOTALL)
|
||||
if not match:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(match.group(1).strip())
|
||||
required_skills = [
|
||||
SkillWeight(skill_id=s.get("skill_id", ""), percentage=float(s.get("percentage", 0)))
|
||||
for s in data.get("required_skills", [])
|
||||
]
|
||||
return SkillAnalysis(
|
||||
stage=data.get("stage", ""),
|
||||
required_skills=required_skills,
|
||||
reasoning=data.get("reasoning", ""),
|
||||
raw_json=data,
|
||||
)
|
||||
except (json.JSONDecodeError, KeyError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Routing strategies
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelRoutingResult:
|
||||
model_alias: str
|
||||
decision_logic: str
|
||||
confidence: float = 0.0
|
||||
all_scores: Dict[str, float] = field(default_factory=dict)
|
||||
|
||||
|
||||
class RoutingStrategy:
|
||||
def __init__(self, handbook: Optional[StageSkillHandbook] = None):
|
||||
self.handbook = handbook
|
||||
self._model_skill_scores: Dict[str, Dict[str, float]] = {}
|
||||
if handbook:
|
||||
self._model_skill_scores = handbook.get_model_skill_scores()
|
||||
|
||||
def _find_skill_id(self, stage: str, skill_id_or_name: str) -> Optional[str]:
|
||||
if not self.handbook:
|
||||
return None
|
||||
handbook_stage = "code" if stage == "reasoning" else stage
|
||||
stage_skills = self.handbook.skills.get(handbook_stage, {})
|
||||
if skill_id_or_name in stage_skills:
|
||||
return skill_id_or_name
|
||||
lower = skill_id_or_name.lower()
|
||||
for sid, skill in stage_skills.items():
|
||||
if skill.name.lower() == lower or lower in skill.name.lower():
|
||||
return sid
|
||||
return None
|
||||
|
||||
def _get_models_for_stage(self, stage: str) -> List[str]:
|
||||
if stage == "search":
|
||||
return ["search-1", "search-2", "search-3"]
|
||||
if stage == "reasoning":
|
||||
return ["reasoner-1", "reasoner-2", "reasoner-3"]
|
||||
if stage == "answer":
|
||||
return ["answer-1", "answer-2", "answer-3", "answer-4", "answer-math-1", "answer-math-2"]
|
||||
return []
|
||||
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class RouterDecidesStrategy(RoutingStrategy):
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
if tool_call_model:
|
||||
return ModelRoutingResult(tool_call_model, "router_decides_from_tool_call", 1.0)
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "router_decides_fallback", 0.5)
|
||||
|
||||
|
||||
class AnalyzeModelDecideStrategy(RoutingStrategy):
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
if tool_call_model:
|
||||
return ModelRoutingResult(tool_call_model, "analyze_model_decide_with_skill_analysis", 1.0)
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "analyze_model_decide_fallback", 0.5)
|
||||
|
||||
|
||||
class WeightedAverageStrategy(RoutingStrategy):
|
||||
COST_TIERS = {
|
||||
"search-3": 1, "search-2": 2, "search-1": 3,
|
||||
"reasoner-3": 1, "reasoner-2": 2, "reasoner-1": 3,
|
||||
"answer-math-2": 1, "answer-4": 1, "answer-3": 2,
|
||||
"answer-math-1": 2, "answer-2": 3, "answer-1": 4,
|
||||
}
|
||||
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
if not skill_analysis or not skill_analysis.required_skills:
|
||||
if tool_call_model:
|
||||
return ModelRoutingResult(tool_call_model, "weighted_avg_no_skills_use_tool_call", 0.7)
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "weighted_avg_no_skills_fallback", 0.5)
|
||||
|
||||
models = self._get_models_for_stage(stage)
|
||||
model_scores = {}
|
||||
for model in models:
|
||||
scores = self._model_skill_scores.get(model, {})
|
||||
weighted_sum = total_weight = 0.0
|
||||
for sw in skill_analysis.required_skills:
|
||||
weight = sw.percentage / 100.0
|
||||
sid = self._find_skill_id(stage, sw.skill_id) or sw.skill_id
|
||||
score = scores.get(sid, 0.0)
|
||||
weighted_sum += weight * score
|
||||
total_weight += weight
|
||||
model_scores[model] = weighted_sum / total_weight if total_weight > 0 else 0.5
|
||||
|
||||
if not model_scores:
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "weighted_avg_no_model_scores", 0.5)
|
||||
max_score = max(model_scores.values())
|
||||
best = [m for m, s in model_scores.items() if abs(s - max_score) < 0.001]
|
||||
best.sort(key=lambda m: self.COST_TIERS.get(m, 999))
|
||||
return ModelRoutingResult(best[0], "weighted_avg_from_skill_analysis", max_score, model_scores)
|
||||
|
||||
|
||||
class WeakestSkillStrategy(RoutingStrategy):
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
if not skill_analysis or not skill_analysis.required_skills:
|
||||
if tool_call_model:
|
||||
return ModelRoutingResult(tool_call_model, "weakest_skill_no_skills_use_tool_call", 0.7)
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "weakest_skill_no_skills_fallback", 0.5)
|
||||
weakest = min(skill_analysis.required_skills, key=lambda s: s.percentage)
|
||||
sid = self._find_skill_id(stage, weakest.skill_id) or weakest.skill_id
|
||||
models = self._get_models_for_stage(stage)
|
||||
model_scores = {m: self._model_skill_scores.get(m, {}).get(sid, 0.5) for m in models}
|
||||
if not model_scores:
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "weakest_skill_no_model_scores", 0.5)
|
||||
best = max(model_scores, key=model_scores.get)
|
||||
return ModelRoutingResult(best, f"weakest_skill_{weakest.skill_id}", model_scores[best], model_scores)
|
||||
|
||||
|
||||
class StrongestSkillStrategy(RoutingStrategy):
|
||||
def select_model(
|
||||
self,
|
||||
stage: str,
|
||||
skill_analysis: Optional[SkillAnalysis] = None,
|
||||
tool_call_model: Optional[str] = None,
|
||||
) -> ModelRoutingResult:
|
||||
if not skill_analysis or not skill_analysis.required_skills:
|
||||
if tool_call_model:
|
||||
return ModelRoutingResult(tool_call_model, "strongest_skill_no_skills_use_tool_call", 0.7)
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "strongest_skill_no_skills_fallback", 0.5)
|
||||
strongest = max(skill_analysis.required_skills, key=lambda s: s.percentage)
|
||||
sid = self._find_skill_id(stage, strongest.skill_id) or strongest.skill_id
|
||||
models = self._get_models_for_stage(stage)
|
||||
model_scores = {m: self._model_skill_scores.get(m, {}).get(sid, 0.5) for m in models}
|
||||
if not model_scores:
|
||||
defaults = {"search": "search-1", "reasoning": "reasoner-1", "answer": "answer-1"}
|
||||
return ModelRoutingResult(defaults.get(stage, "answer-1"), "strongest_skill_no_model_scores", 0.5)
|
||||
best = max(model_scores, key=model_scores.get)
|
||||
return ModelRoutingResult(best, f"strongest_skill_{strongest.skill_id}", model_scores[best], model_scores)
|
||||
|
||||
|
||||
ROUTING_STRATEGIES = {
|
||||
"router_decides": RouterDecidesStrategy,
|
||||
"analyze_model_decide": AnalyzeModelDecideStrategy,
|
||||
"weighted_avg": WeightedAverageStrategy,
|
||||
"weakest_skill": WeakestSkillStrategy,
|
||||
"strongest_skill": StrongestSkillStrategy,
|
||||
}
|
||||
|
||||
|
||||
def get_routing_strategy(
|
||||
strategy_name: str, handbook: Optional[StageSkillHandbook] = None
|
||||
) -> RoutingStrategy:
|
||||
if strategy_name not in ROUTING_STRATEGIES:
|
||||
raise ValueError(
|
||||
f"Unknown routing strategy: {strategy_name}. "
|
||||
f"Available: {list(ROUTING_STRATEGIES.keys())}"
|
||||
)
|
||||
return ROUTING_STRATEGIES[strategy_name](handbook)
|
||||
@@ -0,0 +1,376 @@
|
||||
"""The three SkillOrchestra tools: search, enhance_reasoning (code), answer.
|
||||
|
||||
Faithful port of ``orchestration/eval_frames.py:call_tool`` — same worker
|
||||
prompts, same extraction, same Python subprocess execution. Two deltas,
|
||||
both forced by the OpenJarvis environment and documented inline:
|
||||
|
||||
* ``search`` — the original POSTs to a FAISS wiki retriever service. We
|
||||
honor ``method_cfg.retriever_url`` and POST the exact same payload when
|
||||
it's set; with no retriever configured we fall back to Anthropic's
|
||||
server-side ``web_search`` tool so the stage still grounds.
|
||||
* in-tool correctness check — the original ``answer`` tool LLM-judges the
|
||||
prediction against the gold answer inside ``call_tool``. OpenJarvis
|
||||
scores with its own harness judge downstream, so we only return the
|
||||
prediction; no gold answer is threaded in.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .._base import (
|
||||
GEMINI_SEARCH_COST_PER_CALL,
|
||||
OPENAI_WEB_SEARCH_COST_PER_CALL,
|
||||
WEB_SEARCH_COST_PER_CALL,
|
||||
build_web_search_tool,
|
||||
)
|
||||
from .pool import ModelSpec, call_alias
|
||||
|
||||
# Cloud endpoints with a server-side web-search agent loop wired in
|
||||
# `_base.py`. Anything else (openrouter, vllm, unknown) can't ground.
|
||||
_SEARCH_CAPABLE_ENDPOINTS = ("anthropic", "openai", "gemini")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool schemas — orchestration/tools.json, in Anthropic + OpenAI shapes.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SEARCH_DESC = "Search for missing information."
|
||||
_CODE_DESC = (
|
||||
"Write and execute Python code to compute intermediate results for "
|
||||
"the problem."
|
||||
)
|
||||
_ANSWER_DESC = (
|
||||
"Extract the final answer when you have gathered enough information "
|
||||
"to answer the problem."
|
||||
)
|
||||
|
||||
_ENUMS = {
|
||||
"search": ["search-1", "search-2", "search-3"],
|
||||
"enhance_reasoning": ["reasoner-1", "reasoner-2", "reasoner-3"],
|
||||
"answer": ["answer-1", "answer-2", "answer-3", "answer-4",
|
||||
"answer-math-1", "answer-math-2"],
|
||||
}
|
||||
|
||||
|
||||
def _model_prop(tool: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"type": "string",
|
||||
"description": (
|
||||
f"Model alias for the {tool} tool. Choose one of: "
|
||||
+ ", ".join(_ENUMS[tool])
|
||||
),
|
||||
"enum": _ENUMS[tool],
|
||||
}
|
||||
|
||||
|
||||
def anthropic_tools() -> List[Dict[str, Any]]:
|
||||
"""The 3 orchestrator tools in Anthropic ``input_schema`` shape."""
|
||||
out = []
|
||||
for name, desc in (
|
||||
("search", _SEARCH_DESC),
|
||||
("enhance_reasoning", _CODE_DESC),
|
||||
("answer", _ANSWER_DESC),
|
||||
):
|
||||
out.append({
|
||||
"name": name,
|
||||
"description": desc,
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"model": _model_prop(name)},
|
||||
"required": ["model"],
|
||||
},
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
def openai_tools() -> List[Dict[str, Any]]:
|
||||
"""The 3 orchestrator tools in OpenAI ``function`` shape."""
|
||||
out = []
|
||||
for name, desc in (
|
||||
("search", _SEARCH_DESC),
|
||||
("enhance_reasoning", _CODE_DESC),
|
||||
("answer", _ANSWER_DESC),
|
||||
):
|
||||
out.append({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"description": desc,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"model": _model_prop(name)},
|
||||
"required": ["model"],
|
||||
},
|
||||
},
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# enhance_reasoning / code — eval_frames.py:659-812
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_code(
|
||||
agent: Any,
|
||||
spec: ModelSpec,
|
||||
*,
|
||||
context_str: str,
|
||||
problem: str,
|
||||
bash_timeout_s: int = 60,
|
||||
) -> Dict[str, Any]:
|
||||
"""Generate self-contained Python with ``spec``, execute it, return stdout.
|
||||
|
||||
Mirrors the original worker prompt and ``subprocess.run(['python', ...],
|
||||
timeout=60)`` verbatim. Execution failures yield empty ``exec_result``
|
||||
rather than raising — the orchestrator learns the model can't code.
|
||||
"""
|
||||
prompt = (
|
||||
context_str.strip() + "\n\n"
|
||||
+ f"Question: {problem}\nInstead of directly answering the question, "
|
||||
"please write additional python code that will give intermidiate "
|
||||
"results after execution. Wrap the code within ```python and ```. "
|
||||
"The code should be self-contained with all the import and "
|
||||
"initialization."
|
||||
)
|
||||
text, p, c, cost = call_alias(
|
||||
agent, spec, user=prompt, max_tokens=8000, temperature=1.0,
|
||||
)
|
||||
generated_code = ""
|
||||
if "```python" in text:
|
||||
generated_code = text.split("```python")[-1].split("```")[0]
|
||||
|
||||
exec_result = ""
|
||||
if generated_code.strip():
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
code_path = Path(td) / "exec_code.py"
|
||||
code_path.write_text(generated_code)
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[sys.executable, str(code_path)],
|
||||
timeout=bash_timeout_s,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
exec_result = proc.stdout
|
||||
except Exception:
|
||||
exec_result = ""
|
||||
return {
|
||||
"tool": "enhance_reasoning",
|
||||
"model": spec.model,
|
||||
"alias": spec.alias,
|
||||
"generated_code": generated_code,
|
||||
"exec_result": exec_result,
|
||||
"response": text,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"cost_usd": cost,
|
||||
"is_local": spec.is_local,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# answer — eval_frames.py:814-997
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_answer(
|
||||
agent: Any,
|
||||
spec: ModelSpec,
|
||||
*,
|
||||
context_str: str,
|
||||
problem: str,
|
||||
max_tokens: int = 40000,
|
||||
) -> Dict[str, Any]:
|
||||
"""Generate the final answer with ``spec`` and extract the prediction.
|
||||
|
||||
The original branches the prompt by model family: Qwen-3 / Qwen-math
|
||||
get a ``\\boxed{}`` system prompt; GPT-5 / Claude (and we extend this
|
||||
to every other model) get the ``<think>/<answer>`` instruction. The
|
||||
in-tool LLM correctness check is dropped — OpenJarvis scores
|
||||
downstream.
|
||||
"""
|
||||
base = context_str.strip() + "\n\n" + problem
|
||||
model_l = spec.model.lower()
|
||||
system: Optional[str] = None
|
||||
boxed = False
|
||||
|
||||
if "qwen3" in model_l and "235" not in model_l:
|
||||
system = "Please reason step by step, and put your final answer within \\boxed{}."
|
||||
user = base
|
||||
boxed = True
|
||||
elif "qwen2.5-math" in model_l or "qwen-2.5-math" in model_l:
|
||||
system = "Please reason step by step, and put your final answer within \\boxed{}."
|
||||
user = base
|
||||
boxed = True
|
||||
else:
|
||||
user = base + (
|
||||
"\n\nTake a deep breath and think hard with high reasoning, wrap "
|
||||
"the thoughts within <think> and </think>, and wrap only the "
|
||||
"exact answer without any explanation within <answer> and "
|
||||
"</answer>.Output using the following format:\n<think>\n...\n"
|
||||
"</think>\n<answer>\n...\n</answer>"
|
||||
)
|
||||
|
||||
text, p, c, cost = call_alias(
|
||||
agent, spec, user=user, system=system,
|
||||
max_tokens=max_tokens, temperature=1.0,
|
||||
)
|
||||
|
||||
pred = ""
|
||||
if boxed and "\\boxed{" in text:
|
||||
pred = "}".join(text.split("\\boxed{")[-1].split("}")[:-1]).strip()
|
||||
elif "<answer>" in text:
|
||||
pred = text.split("<answer>")[-1].split("</answer>")[0].strip()
|
||||
else:
|
||||
pred = text.strip()
|
||||
# Original: a >500-word "answer" is treated as a non-answer.
|
||||
if len(pred.split()) > 500:
|
||||
pred = ""
|
||||
|
||||
return {
|
||||
"tool": "answer",
|
||||
"model": spec.model,
|
||||
"alias": spec.alias,
|
||||
"pred": pred,
|
||||
"response": text,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"cost_usd": cost,
|
||||
"is_local": spec.is_local,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# search — eval_frames.py:999-1096
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def run_search(
|
||||
agent: Any,
|
||||
spec: ModelSpec,
|
||||
*,
|
||||
context_str: str,
|
||||
problem: str,
|
||||
retriever_url: Optional[str] = None,
|
||||
topk: int = 150,
|
||||
web_search_max_uses: int = 5,
|
||||
) -> Dict[str, Any]:
|
||||
"""Write a search query with ``spec``, then retrieve documents.
|
||||
|
||||
Query generation is the original verbatim worker prompt. Retrieval:
|
||||
if ``retriever_url`` is set we POST the original ``/retrieve`` payload;
|
||||
otherwise we fall back to Anthropic ``web_search`` (the documented
|
||||
OpenJarvis substitution for the missing FAISS wiki index).
|
||||
"""
|
||||
prompt = (
|
||||
context_str.strip() + "\n\n"
|
||||
+ f"Question: {problem}\nInstead of directly answering the question, "
|
||||
"please think hard and write a concise query to search Wikipedia. "
|
||||
"Wrap the query within <query> and </query>."
|
||||
)
|
||||
text, p, c, cost = call_alias(
|
||||
agent, spec, user=prompt, max_tokens=8000, temperature=1.0,
|
||||
)
|
||||
if "<query>" in text:
|
||||
query = text.split("<query>")[-1].split("</query>")[0].strip()
|
||||
else:
|
||||
query = ""
|
||||
if len(query) < 10:
|
||||
query = problem
|
||||
|
||||
contents: List[str] = []
|
||||
search_uses = 0
|
||||
|
||||
if retriever_url:
|
||||
# Faithful path — the original FAISS retriever service.
|
||||
import requests
|
||||
|
||||
payload = {
|
||||
"queries": [query[:390]],
|
||||
"topk": topk,
|
||||
"return_scores": True,
|
||||
}
|
||||
try:
|
||||
results = requests.post(
|
||||
f"{retriever_url.rstrip('/')}/retrieve", json=payload, timeout=120,
|
||||
).json()
|
||||
for r in results[0]:
|
||||
doc = r.get("document", {})
|
||||
if "content" in doc:
|
||||
contents.append(doc["content"])
|
||||
elif "contents" in doc:
|
||||
contents.append(doc["contents"])
|
||||
except Exception as exc: # noqa: BLE001
|
||||
contents.append(f"[retriever error: {exc}]")
|
||||
else:
|
||||
# Substitution path — server-side web search via the cloud's
|
||||
# `_base` agent loop. The search-capable helpers all talk to
|
||||
# ``agent._cloud_model`` with their provider SDK. If this cell
|
||||
# routes the cloud through an endpoint with no search wiring
|
||||
# (openrouter / vllm), the search stage would produce nothing and
|
||||
# the orchestrator answers the GAIA question blind — fail loud.
|
||||
endpoint = agent._cloud_endpoint
|
||||
if endpoint not in _SEARCH_CAPABLE_ENDPOINTS:
|
||||
raise ValueError(
|
||||
f"skillorchestra search fell back to web_search but "
|
||||
f"cloud_endpoint={endpoint!r}; server-side web_search is "
|
||||
"wired for anthropic / openai / gemini executors only. "
|
||||
"Set method_cfg.retriever_url to a FAISS retriever, route "
|
||||
"this cell's cloud through one of those endpoints, or "
|
||||
"override the search-* aliases in method_cfg.model_pool — "
|
||||
"otherwise the search stage produces nothing and the "
|
||||
"orchestrator answers blind."
|
||||
)
|
||||
search_user = f"Search the web and report findings for: {query}"
|
||||
try:
|
||||
if endpoint == "anthropic":
|
||||
ws_text, wp, wc, n_searches, _ = agent._call_anthropic_agent(
|
||||
agent._cloud_model,
|
||||
user=search_user,
|
||||
max_tokens=4096,
|
||||
temperature=1.0,
|
||||
tools=[build_web_search_tool(web_search_max_uses)],
|
||||
max_turns=4,
|
||||
)
|
||||
ws_cost_per_call = WEB_SEARCH_COST_PER_CALL
|
||||
elif endpoint == "openai":
|
||||
ws_text, wp, wc, n_searches, _ = agent._call_openai_agent(
|
||||
agent._cloud_model,
|
||||
user=search_user,
|
||||
max_tokens=4096,
|
||||
temperature=1.0,
|
||||
max_turns=4,
|
||||
)
|
||||
ws_cost_per_call = OPENAI_WEB_SEARCH_COST_PER_CALL
|
||||
else: # gemini
|
||||
ws_text, wp, wc, n_searches, _ = agent._call_gemini_agent(
|
||||
agent._cloud_model,
|
||||
user=search_user,
|
||||
max_tokens=4096,
|
||||
temperature=1.0,
|
||||
max_turns=4,
|
||||
)
|
||||
ws_cost_per_call = GEMINI_SEARCH_COST_PER_CALL
|
||||
contents.append(ws_text)
|
||||
p += wp
|
||||
c += wc
|
||||
search_uses = n_searches
|
||||
cost += agent.cost_usd(agent._cloud_model, wp, wc)
|
||||
cost += n_searches * ws_cost_per_call
|
||||
except Exception as exc: # noqa: BLE001
|
||||
contents.append(f"[web_search error: {exc}]")
|
||||
|
||||
return {
|
||||
"tool": "search",
|
||||
"model": spec.model,
|
||||
"alias": spec.alias,
|
||||
"query": query,
|
||||
"search_results_data": contents,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"cost_usd": cost,
|
||||
"web_search_uses": search_uses,
|
||||
"is_local": spec.is_local,
|
||||
}
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
Core data types for SkillOrchestra.
|
||||
|
||||
- Skill
|
||||
- AgentProfile
|
||||
- BetaCompetence
|
||||
- ModeMetadata
|
||||
- RoutingInsight
|
||||
- CostStats
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BetaCompetence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class BetaCompetence:
|
||||
"""Bayesian competence estimate for an agent on a specific skill.
|
||||
|
||||
skill_scores / get_competence use empirical_rate (successes/attempts)
|
||||
"""
|
||||
|
||||
alpha: float = 1.0
|
||||
beta: float = 1.0
|
||||
|
||||
@property
|
||||
def mean(self) -> float:
|
||||
return self.alpha / (self.alpha + self.beta)
|
||||
|
||||
@property
|
||||
def empirical_rate(self) -> float:
|
||||
"""Empirical success rate: successes/attempts."""
|
||||
n = self.total_observations
|
||||
if n <= 0:
|
||||
return 0.0
|
||||
successes = max(0, int(self.alpha - 1))
|
||||
return successes / n
|
||||
|
||||
@property
|
||||
def variance(self) -> float:
|
||||
"""Posterior variance."""
|
||||
total = self.alpha + self.beta
|
||||
return (self.alpha * self.beta) / (total * total * (total + 1))
|
||||
|
||||
@property
|
||||
def std(self) -> float:
|
||||
return math.sqrt(self.variance)
|
||||
|
||||
@property
|
||||
def total_observations(self) -> int:
|
||||
"""Total observations (excluding prior)"""
|
||||
return max(0, int(self.alpha + self.beta - 2))
|
||||
|
||||
def update(self, success: bool) -> None:
|
||||
if success:
|
||||
self.alpha += 1.0
|
||||
else:
|
||||
self.beta += 1.0
|
||||
|
||||
def update_batch(self, successes: int, failures: int) -> None:
|
||||
self.alpha += successes
|
||||
self.beta += failures
|
||||
|
||||
def to_dict(self) -> Dict[str, float]:
|
||||
return {"alpha": self.alpha, "beta": self.beta}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, float]) -> BetaCompetence:
|
||||
return cls(alpha=d["alpha"], beta=d["beta"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CostStats
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class CostStats:
|
||||
"""Execution cost statistics for an agent under a specific mode.
|
||||
|
||||
Tracks both total cost (prompt + completion) and completion-only cost
|
||||
separately, since completion cost is the variable component that
|
||||
differs most between models (prompt cost is roughly constant for the
|
||||
same query).
|
||||
"""
|
||||
|
||||
avg_prompt_tokens: float = 0.0
|
||||
avg_completion_tokens: float = 0.0
|
||||
avg_latency_s: float = 0.0
|
||||
avg_cost_usd: float = 0.0
|
||||
avg_completion_cost_usd: float = 0.0
|
||||
avg_prompt_cost_usd: float = 0.0
|
||||
total_executions: int = 0
|
||||
|
||||
def update(
|
||||
self,
|
||||
prompt_tokens: float,
|
||||
completion_tokens: float,
|
||||
latency_s: float,
|
||||
cost_usd: float,
|
||||
completion_cost_usd: float = 0.0,
|
||||
prompt_cost_usd: float = 0.0,
|
||||
) -> None:
|
||||
"""Incremental running-average update."""
|
||||
n = self.total_executions
|
||||
self.avg_prompt_tokens = (self.avg_prompt_tokens * n + prompt_tokens) / (n + 1)
|
||||
self.avg_completion_tokens = (self.avg_completion_tokens * n + completion_tokens) / (n + 1)
|
||||
self.avg_latency_s = (self.avg_latency_s * n + latency_s) / (n + 1)
|
||||
self.avg_cost_usd = (self.avg_cost_usd * n + cost_usd) / (n + 1)
|
||||
self.avg_completion_cost_usd = (self.avg_completion_cost_usd * n + completion_cost_usd) / (n + 1)
|
||||
self.avg_prompt_cost_usd = (self.avg_prompt_cost_usd * n + prompt_cost_usd) / (n + 1)
|
||||
self.total_executions = n + 1
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"avg_prompt_tokens": self.avg_prompt_tokens,
|
||||
"avg_completion_tokens": self.avg_completion_tokens,
|
||||
"avg_latency_s": self.avg_latency_s,
|
||||
"avg_cost_usd": self.avg_cost_usd,
|
||||
"avg_completion_cost_usd": self.avg_completion_cost_usd,
|
||||
"avg_prompt_cost_usd": self.avg_prompt_cost_usd,
|
||||
"total_executions": self.total_executions,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> CostStats:
|
||||
return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RoutingInsight
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class RoutingInsight:
|
||||
"""A single routing insight learned from execution traces"""
|
||||
|
||||
insight_id: str = field(default_factory=lambda: uuid.uuid4().hex[:8])
|
||||
content: str = ""
|
||||
insight_type: str = "" # "transition", "usage", "constraint", "agent_preference"
|
||||
evidence_query_ids: List[str] = field(default_factory=list)
|
||||
confidence: float = 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"insight_id": self.insight_id,
|
||||
"content": self.content,
|
||||
"insight_type": self.insight_type,
|
||||
"evidence_query_ids": self.evidence_query_ids,
|
||||
"confidence": self.confidence,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> RoutingInsight:
|
||||
return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ModeMetadata
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class ModeMetadata:
|
||||
"""Mode-level routing metadata."""
|
||||
|
||||
mode: str = ""
|
||||
description: str = ""
|
||||
insights: List[RoutingInsight] = field(default_factory=list)
|
||||
|
||||
def add_insight(self, insight: RoutingInsight) -> None:
|
||||
self.insights.append(insight)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"mode": self.mode,
|
||||
"description": self.description,
|
||||
"insights": [i.to_dict() for i in self.insights],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> ModeMetadata:
|
||||
insights = [RoutingInsight.from_dict(i) for i in d.get("insights", [])]
|
||||
return cls(
|
||||
mode=d.get("mode", ""),
|
||||
description=d.get("description", ""),
|
||||
insights=insights,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Skill
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class SkillProvenance:
|
||||
"""Tracks how and why a skill was discovered."""
|
||||
|
||||
discovered_from_queries: List[str] = field(default_factory=list)
|
||||
positive_trajectories: List[str] = field(default_factory=list)
|
||||
negative_trajectories: List[str] = field(default_factory=list)
|
||||
discovery_round: int = 0
|
||||
refinement_history: List[Dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"discovered_from_queries": self.discovered_from_queries,
|
||||
"positive_trajectories": self.positive_trajectories,
|
||||
"negative_trajectories": self.negative_trajectories,
|
||||
"discovery_round": self.discovery_round,
|
||||
"refinement_history": self.refinement_history,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> SkillProvenance:
|
||||
return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})
|
||||
|
||||
|
||||
@dataclass
|
||||
class Skill:
|
||||
"""A reusable capability abstraction."""
|
||||
|
||||
skill_id: str = ""
|
||||
name: str = ""
|
||||
description: str = ""
|
||||
indicators: List[str] = field(default_factory=list)
|
||||
examples: List[str] = field(default_factory=list)
|
||||
mode: str = ""
|
||||
parent_skill_id: Optional[str] = None # for hierarchical skills
|
||||
provenance: SkillProvenance = field(default_factory=SkillProvenance)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"skill_id": self.skill_id,
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"indicators": self.indicators,
|
||||
"examples": self.examples,
|
||||
"mode": self.mode,
|
||||
"parent_skill_id": self.parent_skill_id,
|
||||
"provenance": self.provenance.to_dict(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> Skill:
|
||||
provenance = SkillProvenance.from_dict(d.get("provenance", {}))
|
||||
return cls(
|
||||
skill_id=d.get("skill_id", ""),
|
||||
name=d.get("name", ""),
|
||||
description=d.get("description", ""),
|
||||
indicators=d.get("indicators", []),
|
||||
examples=d.get("examples", []),
|
||||
mode=d.get("mode", ""),
|
||||
parent_skill_id=d.get("parent_skill_id"),
|
||||
provenance=provenance,
|
||||
)
|
||||
|
||||
def get_children(self, all_skills: Dict[str, Skill]) -> List[Skill]:
|
||||
"""Get child skills in the hierarchy."""
|
||||
return [s for s in all_skills.values() if s.parent_skill_id == self.skill_id]
|
||||
|
||||
def is_leaf(self, all_skills: Dict[str, Skill]) -> bool:
|
||||
"""True if this skill has no children."""
|
||||
return len(self.get_children(all_skills)) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AgentProfile
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class AgentProfile:
|
||||
"""Agent profile for skill-aware orchestration."""
|
||||
|
||||
agent_id: str = ""
|
||||
mode: str = ""
|
||||
model_name: str = ""
|
||||
tools: List[str] = field(default_factory=list)
|
||||
|
||||
skill_competence: Dict[str, BetaCompetence] = field(default_factory=dict)
|
||||
|
||||
total_attempts: int = 0
|
||||
total_successes: int = 0
|
||||
|
||||
cost_stats: CostStats = field(default_factory=CostStats)
|
||||
|
||||
routing_signals: List[str] = field(default_factory=list)
|
||||
strengths: List[str] = field(default_factory=list)
|
||||
weaknesses: List[str] = field(default_factory=list)
|
||||
|
||||
def get_competence(self, skill_id: str) -> float:
|
||||
"""Get empirical success rate for a skill. Returns 0 if unseen."""
|
||||
if skill_id in self.skill_competence:
|
||||
return self.skill_competence[skill_id].empirical_rate
|
||||
return 0.0
|
||||
|
||||
def get_competence_dist(self, skill_id: str) -> BetaCompetence:
|
||||
"""Get full Beta distribution for a skill, creating with prior if unseen."""
|
||||
if skill_id not in self.skill_competence:
|
||||
self.skill_competence[skill_id] = BetaCompetence()
|
||||
return self.skill_competence[skill_id]
|
||||
|
||||
def update_competence(self, skill_id: str, success: bool) -> None:
|
||||
"""Update competence estimate for a skill."""
|
||||
self.get_competence_dist(skill_id).update(success)
|
||||
|
||||
def weighted_competence(
|
||||
self, skill_weights: Dict[str, float]
|
||||
) -> float:
|
||||
"""Compute weighted competence: sum w_{t,sigma} * alpha/(alpha+beta)."""
|
||||
if not skill_weights:
|
||||
return 0.5
|
||||
total = 0.0
|
||||
for skill_id, weight in skill_weights.items():
|
||||
total += weight * self.get_competence(skill_id)
|
||||
return total
|
||||
|
||||
def category_competence(self, category_prefix: str) -> float:
|
||||
"""Aggregate competence on all skills under a category (skill_id prefix).
|
||||
|
||||
E.g. category_competence('entertainment_knowledge') = avg of
|
||||
get_competence(s) for all s where s.startswith('entertainment_knowledge.').
|
||||
"""
|
||||
prefix = category_prefix.rstrip(".") + "."
|
||||
scores = [
|
||||
self.get_competence(sid)
|
||||
for sid in self.skill_competence
|
||||
if sid.startswith(prefix)
|
||||
]
|
||||
return sum(scores) / len(scores) if scores else 0.0
|
||||
|
||||
def category_competence_for_skills(
|
||||
self, active_skill_ids: List[str]
|
||||
) -> float:
|
||||
"""Category-level competence for hierarchical tie-breaking.
|
||||
|
||||
Extracts parent categories from active_skill_ids (e.g. 'entertainment_knowledge'
|
||||
from 'entertainment_knowledge.episodic_competition_outcome'), computes
|
||||
category_competence for each, returns average.
|
||||
"""
|
||||
categories: set = set()
|
||||
for sid in active_skill_ids:
|
||||
cat = sid.rsplit(".", 1)[0] if "." in sid else sid
|
||||
categories.add(cat)
|
||||
if not categories:
|
||||
return 0.0
|
||||
return sum(self.category_competence(cat) for cat in categories) / len(categories)
|
||||
|
||||
@property
|
||||
def overall_success_rate(self) -> float:
|
||||
"""Overall success rate (trajectory-level when available, else skill-level)."""
|
||||
if self.total_attempts > 0:
|
||||
return self.total_successes / self.total_attempts
|
||||
total_attempts = 0
|
||||
total_successes = 0
|
||||
for bc in self.skill_competence.values():
|
||||
n = bc.total_observations
|
||||
s = max(0, int(bc.alpha - 1))
|
||||
total_attempts += n
|
||||
total_successes += s
|
||||
return total_successes / total_attempts if total_attempts > 0 else 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
skill_scores = {}
|
||||
skill_attempts = {}
|
||||
skill_successes = {}
|
||||
for sid, bc in self.skill_competence.items():
|
||||
obs = bc.total_observations
|
||||
successes = max(0, int(bc.alpha - 1))
|
||||
skill_attempts[sid] = obs
|
||||
skill_successes[sid] = successes
|
||||
skill_scores[sid] = round(successes / obs, 4) if obs > 0 else 0.0
|
||||
|
||||
skill_total_attempts = sum(skill_attempts.values())
|
||||
skill_total_successes = sum(skill_successes.values())
|
||||
|
||||
return {
|
||||
"agent_id": self.agent_id,
|
||||
"mode": self.mode,
|
||||
"model_name": self.model_name,
|
||||
"tools": self.tools,
|
||||
"skill_competence": {
|
||||
sid: bc.to_dict() for sid, bc in self.skill_competence.items()
|
||||
},
|
||||
"skill_scores": skill_scores,
|
||||
"skill_attempts": skill_attempts,
|
||||
"skill_successes": skill_successes,
|
||||
"total_attempts": self.total_attempts if self.total_attempts > 0 else skill_total_attempts,
|
||||
"total_successes": self.total_successes if self.total_attempts > 0 else skill_total_successes,
|
||||
"cost_stats": self.cost_stats.to_dict(),
|
||||
"routing_signals": self.routing_signals,
|
||||
"strengths": self.strengths,
|
||||
"weaknesses": self.weaknesses,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: Dict[str, Any]) -> AgentProfile:
|
||||
skill_competence = {
|
||||
sid: BetaCompetence.from_dict(bc)
|
||||
for sid, bc in d.get("skill_competence", {}).items()
|
||||
}
|
||||
cost_stats = CostStats.from_dict(d.get("cost_stats", {}))
|
||||
return cls(
|
||||
agent_id=d.get("agent_id", ""),
|
||||
mode=d.get("mode", ""),
|
||||
model_name=d.get("model_name", ""),
|
||||
tools=d.get("tools", []),
|
||||
skill_competence=skill_competence,
|
||||
total_attempts=d.get("total_attempts", 0),
|
||||
total_successes=d.get("total_successes", 0),
|
||||
cost_stats=cost_stats,
|
||||
routing_signals=d.get("routing_signals", []),
|
||||
strengths=d.get("strengths", []),
|
||||
weaknesses=d.get("weaknesses", []),
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -63,6 +63,11 @@ class SWEBenchDataset(DatasetProvider):
|
||||
_default_split = "test"
|
||||
|
||||
def __init__(self, variant: str = "verified_mini") -> None:
|
||||
# NOTE: default is the 50-task mini variant. The full 500-task set is
|
||||
# "verified" (princeton-nlp/SWE-bench_Verified). If your subset JSON
|
||||
# references task_ids from the full set (e.g. subsets/swebench_*_n100*),
|
||||
# pass variant="verified" explicitly — otherwise the 450 missing tasks
|
||||
# are silently dropped. See agents/hybrid/runner.py:_load_swebench_tasks.
|
||||
if variant not in _HF_PATHS:
|
||||
raise ValueError(
|
||||
f"Unknown SWE-bench variant {variant!r}; "
|
||||
|
||||
@@ -19,7 +19,12 @@ with two upstream-`swebench` patches applied at import time:
|
||||
``swebench/harness/modal_eval/run_evaluation_modal.py:66`` writes to
|
||||
``/sys/fs/cgroup/cpu/cpu.shares`` (cgroup v1). Modal v2 sandboxes use
|
||||
cgroup v2 — the path doesn't exist and every sandbox dies on the write.
|
||||
Wrap the write in try/except.
|
||||
Wrap the write in try/except. In swebench 4.x the call site is
|
||||
``ModalSandboxRuntime.__init__`` → ``self.write_file(...)`` →
|
||||
``self.sandbox.open(path, "w")``; in older swebench it was a free
|
||||
``set_cpu_quota`` function. We patch both: ``write_file`` swallows
|
||||
FileNotFoundError for cgroup paths, and ``set_cpu_quota`` (if present)
|
||||
is wrapped too.
|
||||
|
||||
2. **Rescore `*_ids` fix**: older harness rescore code read
|
||||
``resolved_instances`` / ``unresolved_instances`` / ``error_instances``
|
||||
@@ -37,6 +42,7 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
@@ -49,6 +55,71 @@ from openjarvis.evals.core.types import EvalRecord
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _run_subprocess_hard_timeout(
|
||||
cmd: list,
|
||||
*,
|
||||
timeout_s: int,
|
||||
cwd: str,
|
||||
) -> "subprocess.CompletedProcess":
|
||||
"""Run ``cmd`` with a timeout that is actually enforced.
|
||||
|
||||
``subprocess.run(..., capture_output=True, timeout=...)`` has a
|
||||
well-known deadlock: on timeout it kills only the *direct* child, then
|
||||
calls ``communicate()`` again to drain the pipes — but if that child
|
||||
spawned grandchildren that inherited the stdout/stderr fds (the Modal
|
||||
``swebench`` harness does exactly this), those grandchildren keep the
|
||||
pipe open and the drain blocks **forever**. The nominal timeout never
|
||||
fires; the runner freezes.
|
||||
|
||||
This helper avoids that by:
|
||||
|
||||
1. Launching the child in its own process group (``start_new_session``)
|
||||
so we can signal the whole tree, not just the direct child.
|
||||
2. On timeout, ``SIGTERM`` then ``SIGKILL`` the entire group so no
|
||||
grandchild survives to hold a pipe open.
|
||||
3. Draining output with a *bounded* ``communicate()`` after the kill so
|
||||
even a stubborn drain can't hang us.
|
||||
|
||||
Raises :class:`subprocess.TimeoutExpired` (same contract as
|
||||
``subprocess.run``) so callers can keep their existing except clause.
|
||||
"""
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
cwd=cwd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
start_new_session=True, # own process group → killable as a tree
|
||||
)
|
||||
try:
|
||||
stdout, stderr = proc.communicate(timeout=timeout_s)
|
||||
return subprocess.CompletedProcess(
|
||||
cmd, proc.returncode, stdout, stderr
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
# Kill the whole group, not just the direct child — Modal harness
|
||||
# subprocesses fork workers that would otherwise keep pipes open.
|
||||
for sig in (signal.SIGTERM, signal.SIGKILL):
|
||||
try:
|
||||
os.killpg(proc.pid, sig)
|
||||
except (ProcessLookupError, PermissionError):
|
||||
break
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
break
|
||||
except subprocess.TimeoutExpired:
|
||||
continue
|
||||
# Drain whatever is left, but never block on it again — the group
|
||||
# is dead, so this returns promptly; the short cap is just paranoia.
|
||||
try:
|
||||
stdout, stderr = proc.communicate(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
stdout, stderr = "", ""
|
||||
raise subprocess.TimeoutExpired(
|
||||
cmd, timeout_s, output=stdout, stderr=stderr
|
||||
)
|
||||
|
||||
|
||||
# ---------- Patch tracking ----------
|
||||
|
||||
_PATCHES_APPLIED = False
|
||||
@@ -100,11 +171,126 @@ def _patch_modal_cgroup_v2() -> None:
|
||||
_m._hybrid_cgroup_patched = True # type: ignore[attr-defined]
|
||||
|
||||
|
||||
_CGROUP_SOURCE_SENTINEL = "_OPENJARVIS_CGROUP_V2_PATCH_APPLIED"
|
||||
|
||||
|
||||
def _patch_modal_sandbox_source() -> None:
|
||||
"""Patch ``run_evaluation_modal.py`` on disk so subprocesses inherit it.
|
||||
|
||||
``_run_harness`` shells out to ``python -m swebench.harness.run_evaluation``,
|
||||
which means our in-process monkey-patches don't help. We do a one-time
|
||||
idempotent textual rewrite of the swebench module file in the venv:
|
||||
|
||||
- Replace the bare ``self.write_file("/sys/fs/cgroup/cpu/cpu.shares", "2048")``
|
||||
with a try/except FileNotFoundError. Marked with a sentinel so we
|
||||
don't reapply on every call.
|
||||
|
||||
Only fires when the original unwrapped line is present and the sentinel
|
||||
isn't — safe to run repeatedly. No-op if upstream ever fixes this.
|
||||
"""
|
||||
try:
|
||||
from swebench.harness.modal_eval import run_evaluation_modal as _m # type: ignore[import-not-found]
|
||||
except Exception:
|
||||
return
|
||||
src_path = getattr(_m, "__file__", None)
|
||||
if not src_path:
|
||||
return
|
||||
try:
|
||||
src = Path(src_path).read_text()
|
||||
except Exception:
|
||||
return
|
||||
if _CGROUP_SOURCE_SENTINEL in src:
|
||||
return
|
||||
needle = ' self.write_file("/sys/fs/cgroup/cpu/cpu.shares", "2048")'
|
||||
if needle not in src:
|
||||
# Upstream changed the line — bail rather than apply blindly.
|
||||
return
|
||||
replacement = (
|
||||
' # ' + _CGROUP_SOURCE_SENTINEL + '\n'
|
||||
' try:\n'
|
||||
' self.write_file("/sys/fs/cgroup/cpu/cpu.shares", "2048")\n'
|
||||
' except FileNotFoundError:\n'
|
||||
' pass # cgroup v2 Modal sandbox — path missing is fine\n'
|
||||
)
|
||||
new_src = src.replace(needle + "\n", replacement, 1)
|
||||
try:
|
||||
Path(src_path).write_text(new_src)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def _patch_modal_sandbox_write_file() -> None:
|
||||
"""Make ``ModalSandboxRuntime.write_file`` survive cgroup-v2 sandboxes.
|
||||
|
||||
swebench 4.x removed ``set_cpu_quota`` and inlined the cgroup write in
|
||||
``ModalSandboxRuntime.__init__`` as
|
||||
``self.write_file("/sys/fs/cgroup/cpu/cpu.shares", "2048")``. Modal v1's
|
||||
``sandbox.open(path, "w")`` raises ``FileNotFoundError`` because the
|
||||
sandbox image is cgroup-v2 and the parent dir doesn't exist, which kills
|
||||
the whole constructor before any patch can be applied. We wrap
|
||||
``write_file`` to swallow that specific failure for cgroup paths, while
|
||||
still letting real write failures (patch/eval script) surface.
|
||||
"""
|
||||
try:
|
||||
from swebench.harness.modal_eval import run_evaluation_modal as _m # type: ignore[import-not-found]
|
||||
except Exception:
|
||||
return
|
||||
runtime = getattr(_m, "ModalSandboxRuntime", None)
|
||||
if runtime is None:
|
||||
return
|
||||
if getattr(runtime, "_hybrid_write_file_patched", False):
|
||||
return
|
||||
orig_write = runtime.write_file
|
||||
|
||||
def patched_write_file(self, file_path: str, content: str): # type: ignore[no-untyped-def]
|
||||
try:
|
||||
return orig_write(self, file_path, content)
|
||||
except FileNotFoundError:
|
||||
# cgroup-v1 paths don't exist in Modal v2 sandboxes — skip
|
||||
# silently for those, re-raise for everything else.
|
||||
if isinstance(file_path, str) and file_path.startswith("/sys/fs/cgroup/"):
|
||||
return None
|
||||
raise
|
||||
|
||||
runtime.write_file = patched_write_file # type: ignore[assignment]
|
||||
runtime._hybrid_write_file_patched = True # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _sentinel_present_on_disk() -> bool:
|
||||
"""Return True iff the cgroup-v2 sentinel is in the installed swebench file.
|
||||
|
||||
Used to detect that a ``uv sync`` / pip reinstall has reverted the textual
|
||||
patch out from under us while the process is still running. The in-process
|
||||
monkey-patches survive that (they live on the imported module object), but
|
||||
the subprocess fork in :func:`_run_harness` reads the file fresh and would
|
||||
silently regress to the broken version.
|
||||
"""
|
||||
try:
|
||||
from swebench.harness.modal_eval import run_evaluation_modal as _m # type: ignore[import-not-found]
|
||||
except Exception:
|
||||
return False
|
||||
src_path = getattr(_m, "__file__", None)
|
||||
if not src_path:
|
||||
return False
|
||||
try:
|
||||
return _CGROUP_SOURCE_SENTINEL in Path(src_path).read_text()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _apply_patches_once() -> None:
|
||||
"""Apply all swebench patches; idempotent and resilient to disk reverts.
|
||||
|
||||
The in-process flag short-circuits the common case, but if the on-disk
|
||||
sentinel is missing we force a re-apply (covers ``uv sync`` / pip
|
||||
reinstall clobbering the textual rewrite while the process is alive).
|
||||
"""
|
||||
global _PATCHES_APPLIED
|
||||
if _PATCHES_APPLIED:
|
||||
if _PATCHES_APPLIED and _sentinel_present_on_disk():
|
||||
return
|
||||
_patch_modal_cgroup_v2()
|
||||
_patch_modal_sandbox_write_file()
|
||||
_patch_modal_sandbox_source()
|
||||
_PATCHES_APPLIED = True
|
||||
|
||||
|
||||
@@ -151,8 +337,8 @@ def _find_report(cache: Path, instance_id: str, run_id: str) -> Optional[Dict[st
|
||||
"""Find the harness's report JSON for one instance.
|
||||
|
||||
swebench writes ``<model_name_or_path>.<run_id>.json`` inside the
|
||||
subprocess CWD. We use ``model_name_or_path="openjarvis-harness"``,
|
||||
``run_id=f"oj-{instance_id}"`` in :func:`_run_harness`.
|
||||
subprocess CWD. We use ``model_name_or_path="openjarvis-harness"``;
|
||||
``run_id`` is built by :func:`_build_run_id`.
|
||||
"""
|
||||
fname = f"openjarvis-harness.{run_id}.json"
|
||||
p = cache / fname
|
||||
@@ -164,7 +350,53 @@ def _find_report(cache: Path, instance_id: str, run_id: str) -> Optional[Dict[st
|
||||
return None
|
||||
|
||||
|
||||
def _run_harness(instance_id: str, patch: str, timeout_s: int) -> Dict[str, Any]:
|
||||
_RUN_ID_SAFE_RE = re.compile(r"[^A-Za-z0-9._-]+")
|
||||
|
||||
|
||||
def _sanitize_run_id_part(s: str) -> str:
|
||||
"""Reduce a free-form string to filesystem-safe ``[A-Za-z0-9._-]+``.
|
||||
|
||||
Both the harness summary filename (``<model>.<run_id>.json``) and the
|
||||
per-instance log subtree (``logs/run_evaluation/<run_id>/...``) are
|
||||
keyed on ``run_id``, so any character that breaks paths or globs will
|
||||
silently corrupt the score. Strip leading/trailing dashes too — those
|
||||
look fine but make filenames awkward to manage by hand.
|
||||
"""
|
||||
return _RUN_ID_SAFE_RE.sub("-", s).strip("-")
|
||||
|
||||
|
||||
def _build_run_id(instance_id: str, cell_name: Optional[str]) -> str:
|
||||
"""Construct a swebench ``run_id`` unique per (cell, instance).
|
||||
|
||||
The harness keys both its "already run, skipping" cache and its report
|
||||
file path on ``run_id`` alone, so two concurrent cells scoring the
|
||||
same ``instance_id`` with the same ``run_id`` collide: the second
|
||||
cell's harness invocation finds the first's report on disk, skips
|
||||
actual execution, and our caller silently reads the wrong verdict (or
|
||||
``no_report`` if the two cells race on the summary file write). See
|
||||
:func:`_run_harness` for the full failure mode.
|
||||
|
||||
With ``cell_name`` we emit ``oj-<cell>-<instance>``, which keeps the
|
||||
intra-cell resume cache working (same cell + same instance → same
|
||||
run_id → harness cache hit) while making inter-cell collisions
|
||||
impossible. Without ``cell_name`` we fall back to the legacy
|
||||
``oj-<instance>`` form for backwards compat with single-cell callers.
|
||||
"""
|
||||
safe_instance = _sanitize_run_id_part(instance_id)
|
||||
if not cell_name:
|
||||
return f"oj-{safe_instance}"
|
||||
safe_cell = _sanitize_run_id_part(cell_name)
|
||||
if not safe_cell:
|
||||
return f"oj-{safe_instance}"
|
||||
return f"oj-{safe_cell}-{safe_instance}"
|
||||
|
||||
|
||||
def _run_harness(
|
||||
instance_id: str,
|
||||
patch: str,
|
||||
timeout_s: int,
|
||||
cell_name: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Hand one prediction to ``python -m swebench.harness.run_evaluation``.
|
||||
|
||||
Returns ``{"success": bool, "score": float, "details": dict}``.
|
||||
@@ -172,7 +404,7 @@ def _run_harness(instance_id: str, patch: str, timeout_s: int) -> Dict[str, Any]
|
||||
_apply_patches_once()
|
||||
backend = os.environ.get("SWEBENCH_BACKEND", "modal").lower()
|
||||
cache = _harness_cache_dir()
|
||||
run_id = f"oj-{instance_id}"
|
||||
run_id = _build_run_id(instance_id, cell_name)
|
||||
|
||||
# Defend against stale reports: ``run_id`` is deterministic per
|
||||
# instance, the cache dir is shared across runs, and ``_find_report``
|
||||
@@ -206,10 +438,25 @@ def _run_harness(instance_id: str, patch: str, timeout_s: int) -> Dict[str, Any]
|
||||
if backend == "modal":
|
||||
cmd += ["--modal", "true"]
|
||||
|
||||
proc = subprocess.run(
|
||||
cmd, capture_output=True, text=True,
|
||||
timeout=timeout_s, cwd=str(cache),
|
||||
)
|
||||
try:
|
||||
proc = _run_subprocess_hard_timeout(
|
||||
cmd, timeout_s=timeout_s, cwd=str(cache),
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
# The harness subprocess (and its Modal grandchildren) exceeded
|
||||
# the cap and were force-killed as a process group. Record an
|
||||
# error verdict rather than letting the exception bubble — the
|
||||
# caller's row stays well-formed and the cell keeps moving.
|
||||
return {
|
||||
"success": False,
|
||||
"score": 0.0,
|
||||
"details": {
|
||||
"reason": "harness_timeout",
|
||||
"timeout_s": timeout_s,
|
||||
"stdout": (exc.stdout or "")[-2000:],
|
||||
"stderr": (exc.stderr or "")[-2000:],
|
||||
},
|
||||
}
|
||||
|
||||
report = _find_report(cache, instance_id, run_id)
|
||||
if report is None:
|
||||
@@ -261,10 +508,17 @@ class SWEBenchHarnessScorer(Scorer):
|
||||
self,
|
||||
*,
|
||||
timeout_s: int = 1800,
|
||||
cell_name: Optional[str] = None,
|
||||
judge_backend: object = None, # noqa: ARG002 — CLI factory compat
|
||||
judge_model: str = "", # noqa: ARG002 — CLI factory compat
|
||||
) -> None:
|
||||
self._timeout_s = int(timeout_s)
|
||||
# ``cell_name`` namespaces the ``run_id`` so concurrent cells scoring
|
||||
# the same SWE instance don't collide on the harness's shared cache.
|
||||
# See :func:`_build_run_id` for the failure mode this prevents. Pass
|
||||
# the hybrid cell name (e.g. ``"skillorchestra-qwen36-opus47-swe-n100"``)
|
||||
# or leave as ``None`` for single-cell callers.
|
||||
self._cell_name = cell_name
|
||||
|
||||
def score(
|
||||
self,
|
||||
@@ -286,10 +540,17 @@ class SWEBenchHarnessScorer(Scorer):
|
||||
if not instance_id:
|
||||
return False, {"reason": "missing_instance_id"}
|
||||
|
||||
result = _run_harness(instance_id, patch, self._timeout_s)
|
||||
result = _run_harness(
|
||||
instance_id, patch, self._timeout_s, cell_name=self._cell_name,
|
||||
)
|
||||
details = dict(result.get("details", {}))
|
||||
details["patch"] = patch
|
||||
return bool(result["success"]), details
|
||||
|
||||
|
||||
__all__ = ["SWEBenchHarnessScorer", "extract_patch"]
|
||||
__all__ = [
|
||||
"SWEBenchHarnessScorer",
|
||||
"extract_patch",
|
||||
"_build_run_id",
|
||||
"_sanitize_run_id_part",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user