diff --git a/pyproject.toml b/pyproject.toml index 2f568fc8..7c0c67b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,6 +179,9 @@ select = ["E", "F", "I", "W"] [tool.ruff.lint.per-file-ignores] "src/openjarvis/evals/datasets/*.py" = ["E501"] "src/openjarvis/evals/scorers/*.py" = ["E501"] +# 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"] [dependency-groups] dev = [ diff --git a/src/openjarvis/agents/__init__.py b/src/openjarvis/agents/__init__.py index c9c7b1bf..9fbbc23e 100644 --- a/src/openjarvis/agents/__init__.py +++ b/src/openjarvis/agents/__init__.py @@ -79,6 +79,14 @@ try: 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. +try: + import openjarvis.agents.hybrid # noqa: F401 +except ImportError: + pass + # Registry alias: "react" -> NativeReActAgent (for backward compat) try: from openjarvis.core.registry import AgentRegistry diff --git a/src/openjarvis/agents/hybrid/README.md b/src/openjarvis/agents/hybrid/README.md new file mode 100644 index 00000000..f07be5ef --- /dev/null +++ b/src/openjarvis/agents/hybrid/README.md @@ -0,0 +1,100 @@ +# 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. + +| Agent | Plan shape | Trains what? | Workers | +|-------------------|-----------------|----------------------|---------------------------| +| `minions` | reactive loop | nothing | 1 local + 1 cloud | +| `conductor` | static DAG | (paper: 7B planner) | up to 5 frontier+open | +| `archon` | static recipe | nothing (search) | K local + cloud rank/fuse | +| `advisors` | reactive loop | (paper: local model) | 1 local + 1 cloud | +| `skillorchestra` | per-query pick | (paper: profiler) | 1 local + 1 cloud | +| `toolorchestra` | reactive loop | (paper: 8B RL) | local + tools + LLM pool | + +Items in parentheses are what the *paper* trains. These OpenJarvis ports +are **inference-only** — none modify weights. The trained variants (advisor +RL, Orchestrator-8B, SkillOrchestra learn-phase) stay TODOs; the prompted +lower-bounds get you 80-90% of the headline accuracy at zero training cost. + +## What's where + +``` +src/openjarvis/agents/hybrid/ +├── _base.py LocalCloudAgent ABC + SDK helpers +├── _prices.py cloud-model pricing + temp-strip quirks +├── _prompts.py GAIA / SWE-bench answer-format instructions +├── advisors.py executor ↔ advisor ↔ executor (3-step) +├── conductor.py static DAG planner +├── minions.py HazyResearch Minions wrapper +├── archon.py Archon (generator → ranker → fuser) +├── skillorchestra.py skill-aware router +├── toolorchestra.py prompted multi-turn tool dispatcher +├── runner.py CLI: python -m ...hybrid.runner --cell NAME +├── registry/ .toml — one cell per (bench, local, cloud, N) +└── scripts/ + └── new_experiment.sh scaffold a new cell, run instructions +``` + +The Modal-backed SWE-bench-Verified scorer is in +`src/openjarvis/evals/scorers/swebench_harness.py` (next to the existing +structural scorer). + +## Quickstart + +```bash +cd /matx/u/aspark/OpenJarvis +source .env # API keys + +# 1. Start vLLM in another shell (see CLAUDE.md for the full 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 + +# 3. Run a smoke cell +.venv/bin/python -m openjarvis.agents.hybrid.runner \ + --cell minions-gaia-qwen27b-opus-3 +``` + +Outputs land in +`$OPENJARVIS_HYBRID_EXPERIMENTS_DIR//{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. + +## Adding a cell + +```bash +src/openjarvis/agents/hybrid/scripts/new_experiment.sh \ + --method conductor --bench gaia \ + --local qwen3.5-27b --cloud claude-opus-4-7 --n 30 +``` + +That appends a `[cells.]` block to +`registry/conductor.toml` and prints the runner invocation. + +## How good is each paradigm? + +Numbers from the upstream hybrid harness +(`/matx/u/aspark/hybrid-local-cloud-compute/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. + +| paradigm | shape | GAIA acc / $ | SWE acc / $ | verdict | +|--------------------|----------------------|-----------------|-----------------|--------------------------------------------------| +| **minions** | reactive 1L+1C loop | 0.576 / $0.67 | 0.276 / $0.09 | **keep** — matches cloud, ~10× cheaper on SWE | +| **skillorchestra** | per-query router | 0.570 / $0.02 | 0.298 / $0.05 | **keep** — cloud-tier acc at 1/50× cost | +| conductor | static DAG planner | 0.503 / $0.03 | 0.296 / $0.07 | mixed — wins SWE, −7pp on GAIA | +| advisors | exec ↔ advisor loop | 0.497 / $0.02 | 0.302 / $0.07 | mixed — wins SWE, −7pp on GAIA | +| archon | gen → rank → fuse | 0.376 / $0.14 | 0.288 / $0.17 | dominated on GAIA (−19pp); only mid on SWE | +| toolorchestra | RL'd 8B + tool pool | — | — | port lands but untested at full N (heavy infra) | + +Cell configs in `registry/` are copies of the hybrid harness's +`experiments/registry/` — same models, same N, same `method_cfg` — so +these OpenJarvis cells should reproduce the harness numbers within +noise. Until that's validated, the harness stays the authoritative +reference. diff --git a/src/openjarvis/agents/hybrid/__init__.py b/src/openjarvis/agents/hybrid/__init__.py new file mode 100644 index 00000000..87e67091 --- /dev/null +++ b/src/openjarvis/agents/hybrid/__init__.py @@ -0,0 +1,42 @@ +"""Hybrid local+cloud paradigms — ported from `hybrid-local-cloud-compute`. + +Each module here registers one agent under ``@AgentRegistry.register("")``: + + advisors — executor (cloud) ↔ advisor (local) ↔ executor (cloud) + 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 + toolorchestra — prompted multi-turn dispatcher over a mixed tool/model pool + +All agents share :class:`LocalCloudAgent` as the base. They are bench-agnostic: +the caller formats the prompt (using ``hybrid_prompts.format_prompt(task, bench)`` +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. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + +# Import each paradigm to trigger its @AgentRegistry.register() decorator. +for _modname in ( + "advisors", + "conductor", + "minions", + "archon", + "skillorchestra", + "toolorchestra", + "mini_swe_agent", +): + try: + __import__(f"openjarvis.agents.hybrid.{_modname}") + except Exception as exc: # pragma: no cover — optional deps may be missing + logger.debug("hybrid agent %s skipped: %s", _modname, exc) diff --git a/src/openjarvis/agents/hybrid/_base.py b/src/openjarvis/agents/hybrid/_base.py new file mode 100644 index 00000000..15ffd114 --- /dev/null +++ b/src/openjarvis/agents/hybrid/_base.py @@ -0,0 +1,615 @@ +"""LocalCloudAgent — shared base for hybrid local+cloud paradigm agents. + +The hybrid paradigms (Minions, Conductor, Archon, Advisors, SkillOrchestra, +ToolOrchestra) all coordinate at least two models: a small **local** model +served by vLLM over an OpenAI-compatible endpoint, and a **cloud** model +reached via the Anthropic or OpenAI SDK. + +Why not just use OpenJarvis's :class:`InferenceEngine` for both? Two reasons: + +1. The reference hybrid adapters (``hybrid-local-cloud-compute/adapters/``) make + raw SDK calls because some of them (Minions, Archon) construct external + library objects that themselves create their own SDK clients. We mirror that + here so the n=500 numbers stay reproducible during the port. +2. Cloud-side quirks (Opus 4.7 temperature stripping, GPT-5 family + ``max_completion_tokens``) are paradigm-shaped — Minions needs structured + outputs on the supervisor turn, SkillOrchestra needs them on the router, + baseline_cloud does not. Keeping the SDK calls in the agent layer lets each + paradigm decide the schema rather than fighting a shared engine API. + +The base class therefore provides only: + +- Standard ``run()`` contract returning an :class:`AgentResult` whose + ``metadata`` carries the hybrid-result fields (``tokens_local``, + ``tokens_cloud``, ``cost_usd``, ``latency_s``, ``traces``). +- ``_call_anthropic`` / ``_call_openai`` / ``_call_vllm`` helpers that handle + Opus 4.7 temperature stripping, GPT-5 token-arg naming, vLLM + ``enable_thinking`` kwargs, and basic token bookkeeping. +- ``_soft_fail_metadata`` for deterministic failure rows (e.g. Qwen JSON + malformation) so the runner doesn't crash the whole cell. + +Agents register themselves with ``@AgentRegistry.register("name")`` and become +discoverable via the existing SDK / CLI flow. The runner constructs them with +the cloud ``(engine, model)`` as the canonical pair, and paradigm-specific +kwargs (``local_model``, ``local_endpoint``, ``cloud_endpoint``, …) follow. +""" + +from __future__ import annotations + +import json +import threading +import time +from abc import abstractmethod +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from openjarvis.agents._stubs import AgentContext, AgentResult, BaseAgent +from openjarvis.agents.hybrid._prices import ( + NO_TEMP_PREFIXES, + is_gpt5_family, + supports_temperature, +) +from openjarvis.agents.hybrid._prices import ( + cost as estimate_cost, +) +from openjarvis.engine._stubs import InferenceEngine + +# Anthropic server-side web_search: $10 per 1000 searches. +WEB_SEARCH_COST_PER_CALL = 0.01 + +ANTHROPIC_WEB_SEARCH_TOOL = { + "type": "web_search_20250305", + "name": "web_search", + "max_uses": 8, +} + + +# ---------- Thread-local trace buffer ---------- +# +# Every call through ``_call_anthropic`` / ``_call_openai`` / ``_call_vllm`` +# appends an event to the active trace if one is open. ``run()`` opens a fresh +# trace per task and writes the digested log to +# ``/.json`` when it's done. Thread-local so concurrent +# tasks in the runner's ThreadPoolExecutor don't stomp each other's trace. + +_TRACE_STATE = threading.local() + + +def _trace_events() -> Optional[List[Dict[str, Any]]]: + return getattr(_TRACE_STATE, "events", None) + + +def _record_event(event: Dict[str, Any]) -> None: + events = _trace_events() + if events is not None: + events.append(event) + + +def _open_trace() -> List[Dict[str, Any]]: + events: List[Dict[str, Any]] = [] + _TRACE_STATE.events = events + return events + + +def _close_trace() -> None: + if hasattr(_TRACE_STATE, "events"): + delattr(_TRACE_STATE, "events") + + +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. + + Each block type carries different fields; we extract everything we can + so the per-task log file is a complete record of what the model emitted + (including every tool call request and tool result body). + """ + out: Dict[str, Any] = {"type": getattr(block, "type", type(block).__name__)} + for attr in ( + "id", "name", "input", "text", "thinking", "signature", + "tool_use_id", "content", + ): + if hasattr(block, attr): + val = getattr(block, attr) + # Nested content (e.g. web_search_tool_result.content is a list of + # citation/result blocks). Recurse for completeness. + if attr == "content" and isinstance(val, list): + out[attr] = [_serialize_block(b) for b in val] + else: + out[attr] = _jsonable(val) + return out + + +def _serialize_openai_tool_calls(tool_calls: Any) -> List[Dict[str, Any]]: + """vLLM / OpenAI returns ChatCompletionMessageToolCall objects. Pull out + id, type, name, and the (JSON-string) arguments so they're round-trippable.""" + out: List[Dict[str, Any]] = [] + if not tool_calls: + return out + for tc in tool_calls: + fn = getattr(tc, "function", None) + out.append({ + "id": getattr(tc, "id", None), + "type": getattr(tc, "type", "function"), + "function": { + "name": getattr(fn, "name", None) if fn else None, + "arguments": getattr(fn, "arguments", None) if fn else None, + }, + }) + return out + + +def _jsonable(v: Any) -> Any: + """Best-effort JSON-friendly conversion. Pydantic models → .model_dump(), + dataclasses untouched (json.dumps handles them via default=str).""" + if hasattr(v, "model_dump"): + try: + return v.model_dump() + except Exception: + pass + if isinstance(v, (list, tuple)): + return [_jsonable(x) for x in v] + if isinstance(v, dict): + return {k: _jsonable(x) for k, x in v.items()} + return v + + +class LocalCloudAgent(BaseAgent): + """Base for paradigm agents that coordinate a local + cloud model pair. + + Subclasses implement :meth:`_run_paradigm` rather than ``run`` so the + base can wrap timing, metadata shaping, and soft-fail handling + uniformly. + + The :meth:`run` contract takes the formatted task prompt as ``input`` + and reads paradigm-shaped data from ``context.metadata``: + + - ``context.metadata["task"]``: optional dict (the bench's raw task + row, used by paradigms that look at hints / problem_statement / etc.). + - ``context.metadata["task_id"]``: optional string identifier. + + Construction args: + - ``engine``, ``model``: the cloud engine + model id (satisfies + :class:`BaseAgent`'s contract; only used incidentally — we make raw + SDK calls). + - ``local_model``, ``local_endpoint``: vLLM-served local model and its + OpenAI-compatible endpoint, e.g. ``"http://localhost:8001/v1"``. + - ``cloud_endpoint``: ``"anthropic"`` or ``"openai"`` — picks the + cloud SDK. + - ``cfg``: paradigm-specific knobs (max_tokens, schemas, mode, …). + """ + + accepts_tools: bool = False + + def __init__( + self, + engine: InferenceEngine, + model: str, + *, + local_model: Optional[str] = None, + local_endpoint: Optional[str] = None, + cloud_endpoint: str = "anthropic", + cfg: Optional[Dict[str, Any]] = None, + bus: Optional[Any] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + ) -> None: + super().__init__( + engine, + model, + bus=bus, + temperature=temperature, + max_tokens=max_tokens, + ) + self._cloud_model = model + self._cloud_endpoint = (cloud_endpoint or "anthropic").lower() + self._local_model = local_model + self._local_endpoint = local_endpoint + self._cfg: Dict[str, Any] = dict(cfg or {}) + + # ------------------------------------------------------------------ + # SDK call helpers — raw clients, paradigm-shaped quirks applied + # ------------------------------------------------------------------ + + @staticmethod + def _call_anthropic( + model: str, + *, + user: str, + system: Optional[str] = None, + max_tokens: int = 4096, + temperature: float = 0.0, + tools: Optional[list] = None, + tool_choice: Optional[dict] = None, + output_config: Optional[dict] = None, + timeout: float = 600.0, + max_retries: int = 5, + 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. + """ + import anthropic + + client = anthropic.Anthropic(timeout=timeout, max_retries=max_retries) + kwargs: Dict[str, Any] = { + "model": model, + "max_tokens": max_tokens, + "messages": [{"role": "user", "content": user}], + } + if system: + kwargs["system"] = system + if supports_temperature(model): + kwargs["temperature"] = temperature + if tools: + kwargs["tools"] = tools + if tool_choice: + kwargs["tool_choice"] = tool_choice + if output_config: + kwargs["output_config"] = output_config + t0 = time.time() + msg = client.messages.create(**kwargs) + 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", + )] + _record_event({ + "kind": "anthropic", + "role": trace_role, + "model": model, + "system": system, + "user": user, + "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, + "tool_choice": tool_choice, + "output_config": output_config, + "stop_reason": getattr(msg, "stop_reason", None), + "latency_s": latency, + "ts": time.time(), + }) + return text, msg.usage.input_tokens, msg.usage.output_tokens, n_searches + + @staticmethod + def _call_openai( + model: str, + *, + user: str, + system: Optional[str] = None, + max_tokens: int = 4096, + temperature: float = 0.0, + response_format: Optional[dict] = None, + tools: Optional[list] = None, + tool_choice: Optional[Any] = None, + timeout: float = 600.0, + trace_role: str = "cloud", + ) -> Tuple[str, int, int]: + """Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured; + also records any tool_calls the model emits.""" + from openai import OpenAI + + client = OpenAI(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} + if is_gpt5_family(model): + kwargs["max_completion_tokens"] = max_tokens + else: + kwargs["max_tokens"] = max_tokens + kwargs["temperature"] = temperature + if response_format is not None: + kwargs["response_format"] = response_format + if tools: + kwargs["tools"] = tools + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice + t0 = time.time() + resp = client.chat.completions.create(**kwargs) + 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": "openai", + "role": trace_role, + "model": model, + "system": system, + "user": user, + "response": text, + "tool_calls": tool_calls, + "reasoning_content": reasoning, + "tokens_in": p, + "tokens_out": c, + "response_format": response_format, + "tools_declared": tools, + "tool_choice": tool_choice, + "finish_reason": getattr(choice, "finish_reason", None), + "latency_s": latency, + "ts": time.time(), + }) + return text, p, c + + @staticmethod + def _call_vllm( + model: str, + endpoint: str, + *, + user: str, + system: Optional[str] = None, + max_tokens: int = 4096, + temperature: float = 0.0, + enable_thinking: bool = False, + tools: Optional[list] = None, + tool_choice: Optional[Any] = None, + timeout: float = 600.0, + trace_role: str = "local", + ) -> Tuple[str, int, int]: + """Local vLLM (OpenAI-compatible) call. Returns (text, p_tok, c_tok). + Captures the full response into the trace — including any tool_calls + the local model emits (vLLM exposes them in + ``resp.choices[0].message.tool_calls`` when ``--enable-auto-tool-choice`` + is on).""" + from openai import OpenAI + + client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout) + messages: list = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": user}) + kwargs: Dict[str, Any] = dict( + model=model, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + extra_body={"chat_template_kwargs": {"enable_thinking": enable_thinking}}, + ) + if tools: + kwargs["tools"] = tools + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice + t0 = time.time() + resp = client.chat.completions.create(**kwargs) + 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": "vllm", + "role": trace_role, + "model": model, + "endpoint": endpoint, + "system": system, + "user": user, + "response": text, + "tool_calls": tool_calls, + "reasoning_content": reasoning, + "tokens_in": p, + "tokens_out": c, + "temperature": temperature, + "max_tokens": max_tokens, + "enable_thinking": enable_thinking, + "tools_declared": tools, + "tool_choice": tool_choice, + "finish_reason": getattr(choice, "finish_reason", None), + "latency_s": latency, + "ts": time.time(), + }) + return text, p, c + + def _call_cloud( + self, + *, + user: str, + system: Optional[str] = None, + max_tokens: int = 4096, + temperature: float = 0.0, + **kwargs: Any, + ) -> Tuple[str, int, int]: + """Dispatch a single cloud call by ``self._cloud_endpoint``. + + Returns (text, p_tok, c_tok). For Anthropic, the web_search count + is discarded — paradigms that care should call ``_call_anthropic`` + directly. + """ + if self._cloud_endpoint == "anthropic": + text, p, c, _ = self._call_anthropic( + self._cloud_model, + user=user, + system=system, + max_tokens=max_tokens, + temperature=temperature, + **kwargs, + ) + return text, p, c + if self._cloud_endpoint == "openai": + return self._call_openai( + self._cloud_model, + user=user, + system=system, + max_tokens=max_tokens, + temperature=temperature, + **kwargs, + ) + raise ValueError(f"unsupported cloud endpoint: {self._cloud_endpoint!r}") + + # ------------------------------------------------------------------ + # Result shaping + # ------------------------------------------------------------------ + + @staticmethod + def _soft_fail_metadata(reason: str) -> Dict[str, Any]: + """Metadata for a soft-fail row (Qwen JSON broke, Anthropic 400, etc.). + + The agent still returns an :class:`AgentResult` with empty content; + the runner records it as score=0 without crashing the cell. + """ + return { + "tokens_local": 0, + "tokens_cloud": 0, + "cost_usd": 0.0, + "latency_s": 0.0, + "soft_error": reason, + "traces": {"soft_error": reason}, + } + + @staticmethod + def cost_usd(model: str, prompt_tokens: int, completion_tokens: int) -> float: + return estimate_cost(model, prompt_tokens, completion_tokens) + + # ------------------------------------------------------------------ + # Run contract + # ------------------------------------------------------------------ + + def run( + self, + input: str, + context: Optional[AgentContext] = None, + **kwargs: Any, + ) -> AgentResult: + self._emit_turn_start(input) + t0 = time.time() + events = _open_trace() + meta: Dict[str, Any] + answer: str = "" + soft_reason: Optional[str] = None + exc_obj: Optional[BaseException] = None + try: + try: + answer, meta = self._run_paradigm(input, context, **kwargs) + except Exception as exc: + soft = self._is_soft_failure(exc) + if soft is None: + exc_obj = exc + raise + soft_reason = soft + meta = self._soft_fail_metadata(soft) + finally: + # 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( + context, input, answer, meta if "meta" in locals() else {}, + events, soft_reason, exc_obj, + ) + _close_trace() + meta.setdefault("latency_s", time.time() - t0) + if soft_reason is not None: + self._emit_turn_end(soft_error=soft_reason) + return AgentResult(content="", metadata=meta, turns=0) + self._emit_turn_end(**{k: v for k, v in meta.items() if k != "traces"}) + return AgentResult( + content=answer, + metadata=meta, + turns=int(meta.get("turns", 0) or 0), + ) + + @staticmethod + def record_trace_event(event: Dict[str, Any]) -> None: + """Public hook for paradigm code that bypasses the SDK helpers + (Minions's protocol loop, Archon's layer pipeline, …) to drop a + custom event into the current task's trace. + """ + _record_event({**event, "ts": event.get("ts", time.time())}) + + def _write_trace_log( + self, + context: Optional[AgentContext], + input: str, + answer: str, + meta: Dict[str, Any], + events: List[Dict[str, Any]], + soft_reason: Optional[str], + exc: Optional[BaseException], + ) -> None: + log_dir = None + task_id = "unknown" + if context is not None: + log_dir = context.metadata.get("log_dir") + task_id = context.metadata.get("task_id") or task_id + if not log_dir: + return + try: + out_dir = Path(log_dir) + out_dir.mkdir(parents=True, exist_ok=True) + blob = { + "task_id": task_id, + "agent": self.agent_id, + "cloud_model": self._cloud_model, + "cloud_endpoint": self._cloud_endpoint, + "local_model": self._local_model, + "local_endpoint": self._local_endpoint, + "cfg": self._cfg, + "input": input, + "answer": answer, + "metadata": meta, + "events": events, + "soft_error": soft_reason, + "error": ( + f"{type(exc).__name__}: {exc}" if exc is not None else None + ), + } + (out_dir / f"{task_id}.json").write_text( + json.dumps(blob, indent=2, default=str) + ) + except Exception: + # Logging must never break a run. + pass + + @abstractmethod + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + """Run the paradigm. Return ``(final_answer, metadata)``. + + Metadata should include the hybrid-shape fields: + ``tokens_local``, ``tokens_cloud``, ``cost_usd``, optional + ``latency_s`` (the base fills it if absent), and a ``traces`` dict. + """ + + # Subclasses override to declare deterministic failure modes they + # want the base to swallow into a soft-fail row. + def _is_soft_failure(self, exc: BaseException) -> Optional[str]: + return None + + +__all__ = [ + "ANTHROPIC_WEB_SEARCH_TOOL", + "LocalCloudAgent", + "NO_TEMP_PREFIXES", + "WEB_SEARCH_COST_PER_CALL", + "estimate_cost", + "is_gpt5_family", + "supports_temperature", +] diff --git a/src/openjarvis/agents/hybrid/_prices.py b/src/openjarvis/agents/hybrid/_prices.py new file mode 100644 index 00000000..b0eeba0b --- /dev/null +++ b/src/openjarvis/agents/hybrid/_prices.py @@ -0,0 +1,46 @@ +"""Cloud-model pricing + per-family quirks for hybrid paradigm agents. + +Ported verbatim from ``hybrid-local-cloud-compute/prices.py``. Kept as a +sibling to the agents rather than merged into ``engine/cloud.py``'s PRICING +on purpose: the hybrid harness is the authoritative cost reference for the +n=500 numbers in ``hybrid-local-cloud-compute/docs/results.md`` and we want +the OpenJarvis ports to charge identically. +""" + +from __future__ import annotations + +# USD per million tokens, (input, output). Local models = 0. +PRICES: dict[str, tuple[float, float]] = { + "claude-opus-4-7": (5.00, 25.0), + "claude-sonnet-4-6": (3.00, 15.0), + "claude-haiku-4-5": (1.00, 5.00), + "claude-haiku-4-5-20251001": (1.00, 5.00), + "gpt-5": (1.25, 10.0), + "gpt-5-mini": (0.25, 2.00), + "gpt-5-mini-2025-08-07": (0.25, 2.00), + "gpt-4o": (0.15, 0.60), + "gemini-2.5-pro": (1.25, 10.0), +} + +# Models whose API rejects an explicit `temperature` param — callers should +# omit it for any model whose name starts with one of these prefixes. +NO_TEMP_PREFIXES: tuple[str, ...] = ( + "claude-opus-4-7", + "claude-sonnet-4-7", + "claude-haiku-4-7", +) + + +def cost(model: str, prompt_tokens: int, completion_tokens: int) -> float: + """USD cost for one call. Unknown models price at 0 (e.g. local vLLM).""" + pi, po = PRICES.get(model, (0.0, 0.0)) + return (prompt_tokens / 1_000_000) * pi + (completion_tokens / 1_000_000) * po + + +def supports_temperature(model: str) -> bool: + return not model.startswith(NO_TEMP_PREFIXES) + + +def is_gpt5_family(model: str) -> bool: + """GPT-5 series requires ``max_completion_tokens`` and forced temp=1.""" + return model.startswith("gpt-5") diff --git a/src/openjarvis/agents/hybrid/_prompts.py b/src/openjarvis/agents/hybrid/_prompts.py new file mode 100644 index 00000000..cf6b4135 --- /dev/null +++ b/src/openjarvis/agents/hybrid/_prompts.py @@ -0,0 +1,68 @@ +"""Bench-specific answer-format instructions for hybrid paradigm agents. + +Ported from ``hybrid-local-cloud-compute/benches/{gaia,swebench_verified}/``. + +Adapters in the hybrid harness call ``benches.format_prompt(task)`` to build +the full user prompt — the paradigms here take the formatted prompt directly, +so callers (the experiment runner or eval driver) use these helpers to build +the prompt once and hand the string to the agent's ``run(input=...)``. +""" + +from __future__ import annotations + +from typing import Any, Dict + +GAIA_INSTRUCTION = ( + "GAIA expects a single short, exact answer. Reason as needed, then commit to " + "one best answer. Your reply MUST end with exactly one line of the form:\n" + "FINAL ANSWER: \n" + "Formatting rules for :\n" + " - If a number, write digits only (no commas, no units, no $ or %). " + "e.g. `34689`, not `34,689 papers`.\n" + " - If a string, use as few words as possible, no articles, no abbreviations.\n" + " - If a list, comma-separate it and apply the rules above to each element.\n" + "Nothing should follow the FINAL ANSWER line." +) + +SWEBENCH_INSTRUCTION = ( + "You are fixing a real Python bug from SWE-bench Verified. Your reply MUST " + "end with a unified diff patch that fixes the bug. The patch MUST:\n" + " - Start with a `diff --git a/ b/` line (one per file changed).\n" + " - Include the standard `--- a/` and `+++ b/` lines below it.\n" + " - Use proper hunk headers `@@ -start,len +start,len @@`.\n" + " - Apply cleanly against the repository at the given base commit.\n" + "Wrap the patch in a ```diff ... ``` code fence. Do not include any prose " + "after the fence — only the closing ```." +) + + +def format_gaia(task: Dict[str, Any]) -> str: + return f"{GAIA_INSTRUCTION}\n\nQuestion:\n{task.get('question', '')}" + + +def format_swebench(task: Dict[str, Any]) -> str: + header = SWEBENCH_INSTRUCTION + "\n\n" + if task.get("repo"): + header += f"Repository: {task['repo']}\n" + if task.get("base_commit"): + header += f"Base commit: {task['base_commit']}\n" + header += "\nGitHub issue:\n" + return header + task.get("problem_statement", "") + + +def format_prompt(task: Dict[str, Any]) -> str: + """Dispatch on task shape — same contract as ``benches.format_prompt``.""" + if task.get("question"): + return format_gaia(task) + if task.get("problem_statement"): + return format_swebench(task) + return "" + + +__all__ = [ + "GAIA_INSTRUCTION", + "SWEBENCH_INSTRUCTION", + "format_gaia", + "format_prompt", + "format_swebench", +] diff --git a/src/openjarvis/agents/hybrid/advisors.py b/src/openjarvis/agents/hybrid/advisors.py new file mode 100644 index 00000000..35797759 --- /dev/null +++ b/src/openjarvis/agents/hybrid/advisors.py @@ -0,0 +1,280 @@ +"""AdvisorsAgent — inference-only port of advisor-models (Asawa et al., 2026). + +Paper: arXiv:2510.02453. A small open-source advisor model writes feedback +that *steers* a black-box cloud executor. The paper trains the advisor with +RL; we don't have a released checkpoint, so this agent is the **inference- +only lower bound**: an untrained Qwen advisor zero-shot prompted with the +paper's structure. + +Pipeline (mirrors ``advisor_models/math/env.py``): + +1. **Executor (cloud)** answers the question. +2. **Advisor (local)** reads question + initial response and writes + critique / hint text. +3. **Executor (cloud)** re-answers given question + its own initial + response + advisor feedback. This final answer is what we score. + +Results from the hybrid harness (n=30 GAIA): +``advisors-gaia-qwen9b-opus-30`` = 0.533, $0.02/task — within 3pp of +baseline-cloud at 30× cheaper. The RL-trained variant would land higher. + +Ported from ``hybrid-local-cloud-compute/adapters/advisors_adapter.py``. +""" + +from __future__ import annotations + +import json +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.mini_swe_agent import ( + run_swe_agent_loop, +) +from openjarvis.core.registry import AgentRegistry + +# Prompts paraphrased from advisor-models/{math,template}/config.py. + +EXECUTOR_INITIAL_SYS = ( + "You are a careful problem-solver. Read the question, reason step by step " + "as needed, then commit to one best answer following any answer-format " + "instructions in the question itself." +) + +EXECUTOR_FINAL_SYS = ( + "You are a careful problem-solver. You previously gave an initial answer " + "to a question and an advisor has reviewed it. Incorporate the advisor's " + "feedback where it improves correctness; ignore it where it is wrong. " + "Produce your best final answer, following any answer-format instructions " + "in the question itself." +) + +ADVISOR_TEMPLATE = """You are an expert advisor reviewing another model's draft answer to a user question. Your job is NOT to answer the question yourself; it is to give the answering model concrete, actionable feedback so it can improve its next attempt. + +The user question was: +{question} + +The answering model's initial response was: +{initial_response} + +Provide feedback focused on: +1. Specific errors in reasoning, calculation, factual recall, or formatting. +2. Concrete corrections or alternative approaches it should consider. +3. What evidence or sub-step it should verify before committing. + +Be concise (a short paragraph or a few bullet points). Do NOT restate the question. Do NOT provide a complete answer — only the feedback the model needs to improve its own next answer.""" + + +def _resolve_local_model(endpoint: str, registry_model: str) -> str: + """If the registry-listed model isn't being served by vLLM, fall back to + whatever the server reports first. Avoids 404s when cell config names + a model id (e.g. ``Qwen3.5-9B``) that's different from what's loaded. + """ + try: + with urllib.request.urlopen( + endpoint.rstrip("/") + "/models", timeout=5 + ) as r: + data = json.loads(r.read()) + served = [m["id"] for m in data.get("data", [])] + except Exception: + return registry_model + if registry_model in served: + return registry_model + return served[0] if served else registry_model + + +@AgentRegistry.register("advisors") +class AdvisorsAgent(LocalCloudAgent): + """Three-step executor ↔ advisor ↔ executor loop. See module docstring.""" + + agent_id = "advisors" + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + question = input + cfg = self._cfg + 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")) + ) + if swe_mode: + return self._run_swe(question, task_meta, cfg) + + executor_max_tokens = int(cfg.get("executor_max_tokens", 4096)) + 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, + ) + + # 2. Advisor pass (local) + if not self._local_endpoint or not self._local_model: + raise ValueError( + "AdvisorsAgent needs local_model + local_endpoint; got " + f"model={self._local_model!r} endpoint={self._local_endpoint!r}" + ) + local_model = _resolve_local_model(self._local_endpoint, self._local_model) + advisor_prompt = ADVISOR_TEMPLATE.format( + question=question, initial_response=initial_resp, + ) + advisor_text, adv_in, adv_out = self._call_vllm( + local_model, + self._local_endpoint, + user=advisor_prompt, + max_tokens=advisor_max_tokens, + temperature=advisor_temperature, + enable_thinking=False, + ) + + # 3. Final executor pass with advisor's hints folded in + final_user = ( + f"Question:\n{question}\n\n" + f"Your initial response was:\n{initial_resp}\n\n" + f"Advisor feedback:\n{advisor_text}\n\n" + 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, + ) + + 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) + + meta: Dict[str, Any] = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": cost, + "turns": 3, + "traces": { + "initial_response": initial_resp, + "advisor_feedback": advisor_text, + "note": "inference-only advisor (untrained); lower bound on the technique.", + }, + } + return final_answer, meta + + # ------------------------------------------------------------------ + # SWE-bench variant: each executor pass is a full mini-SWE-agent run. + # ------------------------------------------------------------------ + + def _run_swe( + self, + question: str, + task: Dict[str, Any], + cfg: Dict[str, Any], + ) -> Tuple[str, Dict[str, Any]]: + if not self._local_endpoint or not self._local_model: + raise ValueError( + "AdvisorsAgent (swe mode) still needs local_model + local_endpoint " + "for the advisor critique step." + ) + + # Each executor pass is its own mini-SWE-agent run with a FRESH + # workdir — the advisor doesn't get a workdir, just the patch text. + 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)) + + # 1. Initial executor pass + initial_out = run_swe_agent_loop( + task, + backbone="cloud", + backbone_model=self._cloud_model, + cloud_endpoint=self._cloud_endpoint, + initial_prompt=question, + max_turns=max_turns, + bash_timeout=bash_timeout, + output_cap=output_cap, + turn_max_tokens=turn_max_tokens, + trace_prefix="advisors_executor1", + ) + + # 2. Advisor pass — local model critiques the produced patch + local_model = _resolve_local_model(self._local_endpoint, self._local_model) + advisor_prompt = ADVISOR_TEMPLATE.format( + question=question, + initial_response=( + f"Summary: {initial_out['final_summary']}\n\n" + f"Patch produced:\n```diff\n{initial_out['patch']}```" + ), + ) + advisor_text, adv_in, adv_out = self._call_vllm( + local_model, + self._local_endpoint, + user=advisor_prompt, + max_tokens=int(cfg.get("advisor_max_tokens", 2048)), + temperature=float(cfg.get("advisor_temperature", 0.2)), + enable_thinking=False, + ) + + # 3. Final executor pass — FRESH workdir, advisor feedback folded + # into the initial prompt. (Don't reuse initial workdir — that + # would smuggle the bad-fix into the new attempt; we want the + # final pass to incorporate ONLY the parts of the advisor's + # feedback that hold up.) + final_prompt = ( + f"{question}\n\n" + f"-----\n" + f"You previously attempted a fix. Your initial summary was:\n" + f"{initial_out['final_summary']}\n\n" + f"An advisor reviewed your attempt and wrote this feedback:\n" + f"{advisor_text}\n\n" + f"Incorporate the advisor's feedback where it improves correctness; " + f"ignore where it is wrong. Produce your best fix now." + ) + final_out = run_swe_agent_loop( + task, + backbone="cloud", + backbone_model=self._cloud_model, + cloud_endpoint=self._cloud_endpoint, + initial_prompt=final_prompt, + max_turns=max_turns, + bash_timeout=bash_timeout, + output_cap=output_cap, + turn_max_tokens=turn_max_tokens, + trace_prefix="advisors_executor2", + ) + + tokens_local = adv_in + adv_out + tokens_cloud = ( + initial_out["tokens_in"] + initial_out["tokens_out"] + + final_out["tokens_in"] + final_out["tokens_out"] + ) + cost = initial_out["cost_usd"] + final_out["cost_usd"] + meta: Dict[str, Any] = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": cost, + "turns": initial_out["turns"] + 1 + final_out["turns"], + "traces": { + "swe_mode": True, + "initial_summary": initial_out["final_summary"], + "initial_patch_chars": len(initial_out["patch"]), + "advisor_feedback": advisor_text, + "final_summary": final_out["final_summary"], + "final_patch_chars": len(final_out["patch"]), + }, + } + return final_out["answer"], meta + + +__all__ = ["AdvisorsAgent"] diff --git a/src/openjarvis/agents/hybrid/archon.py b/src/openjarvis/agents/hybrid/archon.py new file mode 100644 index 00000000..9c0482cf --- /dev/null +++ b/src/openjarvis/agents/hybrid/archon.py @@ -0,0 +1,584 @@ +"""ArchonAgent — port of ScalingIntelligence/Archon. + +Inference-time architecture search: layered (generator → ranker → fuser) +sampling where a generator proposes K candidates, a ranker scores them, +and a fuser synthesizes a final answer. Paper: arXiv:2409.15254. + +How the hybrid harness wires it (and what we mirror here): + +- **Local proposers** (generator layer): K samples from vLLM via an + OpenAI-compatible client at ``local_endpoint``. Injected as a custom + ``vllm_local`` model_type into Archon's ``GENERATE_MAP`` — that's the + only way Ranker/Fuser can pick up custom backends (they re-instantiate + Generator without ``custom_generators``). +- **Cloud ranker + fuser**: Archon's built-in ``OpenAI_API`` / + ``Anthropic_API``. Patched at import time to strip ``temperature`` for + Opus 4.7+ and to tally token usage (Archon ignores ``usage`` by default). + +``cfg`` knobs: + +- ``n_samples`` (int, default 5) — K proposers +- ``architecture`` (str, default ``"ensemble_rank_fuse"``) + + - ``"ensemble_rank_fuse"`` → [K local generators, 1 cloud ranker, 1 cloud fuser] + - ``"single_local"`` → [1 local generator] (debug) + +- ``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. + +Ported from ``hybrid-local-cloud-compute/adapters/archon_adapter.py``. +""" + +from __future__ import annotations + +import os +import sys +import threading +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._prices import ( + NO_TEMP_PREFIXES, +) +from openjarvis.agents.hybrid._prices import ( + cost as _cost_cloud, +) +from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop +from openjarvis.core.registry import AgentRegistry + +ARCHON_SWE_RANKER_SYS = ( + "You are ranking K candidate patches for a SWE-bench bug. For each " + "candidate, you'll see the candidate's summary text and the unified " + "diff. Pick the index (0-based) of the candidate most likely to " + "actually fix the issue. Prefer minimal, targeted patches; reject " + "candidates with no patch or with patches that touch unrelated files. " + "Respond with a single integer on its own line — the index — " + "followed by a one-line justification." +) + + +# ---------- Stubs for Archon's eager-imported heavy deps we don't need ---------- + +def _stub_archon_imports() -> None: + """``utils.py`` imports groq/google/litellm/dotenv at module load. Stub + the ones we don't use so the import chain doesn't fail when those + libraries aren't installed in the OpenJarvis venv.""" + for name in ("groq", "google", "google.generativeai", "litellm"): + if name in sys.modules: + continue + sys.modules[name] = types.ModuleType(name) + sys.modules["groq"].Groq = type("Groq", (), {}) # type: ignore[attr-defined] + sys.modules["litellm"].completion = lambda *a, **k: None # type: ignore[attr-defined] + sys.modules["google.generativeai"] = types.ModuleType("google.generativeai") + + +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): + sys.path.insert(0, archon_src) + + +# ---------- Anthropic patch for Opus 4.7 ---------- + +def _patch_anthropic_for_opus() -> None: + from anthropic.resources.messages import messages as _msgs_mod + + cls = _msgs_mod.Messages + if getattr(cls.create, "_hybrid_archon_patched", False): + return + orig = cls.create + + def patched(self, **kwargs): # type: ignore[no-untyped-def] + if str(kwargs.get("model", "")).startswith(NO_TEMP_PREFIXES): + kwargs.pop("temperature", None) + return orig(self, **kwargs) + + patched._hybrid_archon_patched = True # type: ignore[attr-defined] + cls.create = patched # type: ignore[assignment] + + +# ---------- Per-run token tally + custom generators ---------- + +# Token tallies live in thread-local storage because the runner reuses one +# ArchonAgent across a ``ThreadPoolExecutor`` (``concurrency > 1``). A plain +# module-global dict would let concurrent ``_run_paradigm`` calls reset and +# read each other's counters, producing wrong ``cost_usd`` and +# ``tokens_local`` / ``tokens_cloud`` in per-task ``meta``. +_TALLY_LOCAL = threading.local() + + +def _tally() -> Dict[str, int]: + """Return the calling thread's token tally, creating it on first touch.""" + counts = getattr(_TALLY_LOCAL, "counts", None) + if counts is None: + counts = { + "cloud_prompt": 0, "cloud_completion": 0, + "local_prompt": 0, "local_completion": 0, + } + _TALLY_LOCAL.counts = counts + return counts + + +def _reset_tally() -> None: + _TALLY_LOCAL.counts = { + "cloud_prompt": 0, "cloud_completion": 0, + "local_prompt": 0, "local_completion": 0, + } + + +def _make_local_generator(local_endpoint: str, local_model: str): + """Archon custom-generator signature: (model, messages, max_tokens, temperature) -> str.""" + from openai import OpenAI + + client = OpenAI(base_url=local_endpoint, api_key="EMPTY") + + def local_gen(model, messages, max_tokens=2048, temperature=0.7, **_kw): # type: ignore[no-untyped-def] + import time as _time + t0 = _time.time() + try: + resp = client.chat.completions.create( + model=local_model, + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + ) + except Exception as e: + _record_event({ + "kind": "archon_local_gen_error", + "model": local_model, + "messages": messages, + "error": f"{type(e).__name__}: {e}", + "ts": _time.time(), + }) + return f"[local-vllm error: {e!r}]" + u = resp.usage + if u: + _tally()["local_prompt"] += getattr(u, "prompt_tokens", 0) or 0 + _tally()["local_completion"] += getattr(u, "completion_tokens", 0) or 0 + text = (resp.choices[0].message.content or "").strip() + _record_event({ + "kind": "archon_local_gen", + "model": local_model, + "messages": messages, + "response": text, + "tokens_in": getattr(u, "prompt_tokens", 0) if u else 0, + "tokens_out": getattr(u, "completion_tokens", 0) if u else 0, + "temperature": temperature, + "max_tokens": max_tokens, + "latency_s": _time.time() - t0, + "ts": _time.time(), + }) + return text + + return local_gen + + +def _wrap_archon_cloud_generators() -> None: + """Replace Archon's GENERATE_MAP entries for OpenAI_API / Anthropic_API + with token-tallying wrappers. GENERATE_MAP holds function references; + we update them in place so Ranker/Fuser see the wrapped versions.""" + import anthropic as _anth + from openai import OpenAI as _OAI + + def gen_openai(model, messages, max_tokens=2048, temperature=0.7, **_kw): # type: ignore[no-untyped-def] + import time as _time + client = _OAI() + kwargs: Dict[str, Any] = dict( + model=model, messages=messages, + max_tokens=max_tokens, temperature=temperature, + ) + # GPT-5/o1/o3 reject non-default temperature and use max_completion_tokens. + if model.startswith(("gpt-5", "o1", "o3")): + kwargs.pop("temperature", None) + kwargs.pop("max_tokens", None) + kwargs["max_completion_tokens"] = max_tokens + t0 = _time.time() + resp = client.chat.completions.create(**kwargs) + u = resp.usage + if u: + _tally()["cloud_prompt"] += getattr(u, "prompt_tokens", 0) or 0 + _tally()["cloud_completion"] += getattr(u, "completion_tokens", 0) or 0 + text = (resp.choices[0].message.content or "").strip() + _record_event({ + "kind": "archon_cloud_openai", + "model": model, + "messages": messages, + "response": text, + "tokens_in": getattr(u, "prompt_tokens", 0) if u else 0, + "tokens_out": getattr(u, "completion_tokens", 0) if u else 0, + "latency_s": _time.time() - t0, + "ts": _time.time(), + }) + return text + + def gen_anthropic(model, messages, max_tokens=2048, temperature=0.7, **_kw): # type: ignore[no-untyped-def] + import time as _time + client = _anth.Anthropic(timeout=600.0) + system = "" + msgs = [] + for m in messages: + if m["role"] == "system" and not system: + system = m["content"] + else: + msgs.append(m) + kwargs: Dict[str, Any] = dict( + model=model, system=system, messages=msgs, max_tokens=max_tokens, + ) + if not model.startswith(NO_TEMP_PREFIXES): + kwargs["temperature"] = temperature + t0 = _time.time() + resp = client.messages.create(**kwargs) + 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 + _record_event({ + "kind": "archon_cloud_anthropic", + "model": model, + "system": system, + "messages": msgs, + "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, + "latency_s": _time.time() - t0, + "ts": _time.time(), + }) + return text.strip() + + from archon.completions.components.Generator import ( + GENERATE_MAP as _GMAP, # type: ignore[import-not-found] + ) + _GMAP["OpenAI_API"] = gen_openai + _GMAP["Anthropic_API"] = gen_anthropic + + +_FUSER_FORMAT_REMINDER = ( + "\n\nIMPORTANT — output format: your synthesized response MUST honor the " + "output-format requirements stated in the original user query above. " + "If the query requires ending with `FINAL ANSWER: ` (GAIA), end " + "with exactly that one line and nothing after it. If the query requires " + "a unified diff inside a ```diff … ``` fence (SWE-bench), end with that " + "fence and nothing after the closing ```. Do not produce free-form " + "analysis, markdown headers, or commentary that breaks the required " + "format — the candidate responses may have done so, but your fused " + "answer must not." +) + + +def _patch_archon_prompts() -> None: + """Make Archon's fuser bench-aware. + + ``Fuser.py`` drops the system message we hand to ``archon.generate()`` and + builds its own user message via ``make_fuser_prompt``. That template tells + the fuser to "synthesize into a refined, well-structured response" with + no reminder of the format the user originally asked for, so on SWE-bench + Opus drifts to a markdown bug report and scores 0. We append an explicit + format-honor reminder to the fuser prompt.""" + from archon.completions.components import ( + prompts as _p, # type: ignore[import-not-found] + ) + + if getattr(_p.make_fuser_prompt, "_hybrid_format_patched", False): + return + orig = _p.make_fuser_prompt + + def patched(conv, references, critiques=None, length_control=False): # type: ignore[no-untyped-def] + base = orig(conv, references, critiques=critiques, length_control=length_control) + return base + _FUSER_FORMAT_REMINDER + + patched._hybrid_format_patched = True # type: ignore[attr-defined] + _p.make_fuser_prompt = patched + # Fuser imports the symbol by name at class-body time — rebind there too. + from archon.completions.components import ( + Fuser as _F, # type: ignore[import-not-found] + ) + _F.make_fuser_prompt = patched + + +_PATCHES_APPLIED = False + + +def _apply_patches_once() -> None: + global _PATCHES_APPLIED + if _PATCHES_APPLIED: + return + _stub_archon_imports() + _add_archon_to_path() + _patch_anthropic_for_opus() + # Trigger Archon imports so GENERATE_MAP exists. + import archon.completions.components.Generator # type: ignore[import-not-found] # noqa: F401 + _wrap_archon_cloud_generators() + _patch_archon_prompts() + _PATCHES_APPLIED = True + + +# ---------- Architecture presets ---------- + +def _presets(): + return { + "ensemble_rank_fuse": lambda K, local_model, ranker_model, fuser_model, max_tokens, temperature: [ + [{ + "type": "generator", + "model": local_model, + "model_type": "vllm_local", + "top_k": 1, + "temperature": temperature, + "max_tokens": max_tokens, + "samples": K, + }], + [{ + "type": "ranker", + "model": ranker_model, + "model_type": "Anthropic_API" if ranker_model.startswith("claude") else "OpenAI_API", + "top_k": min(K, 5), + "temperature": 0.0, + "max_tokens": max_tokens, + }], + [{ + "type": "fuser", + "model": fuser_model, + "model_type": "Anthropic_API" if fuser_model.startswith("claude") else "OpenAI_API", + "temperature": 0.0, + "max_tokens": max_tokens, + "samples": 1, + }], + ], + "single_local": lambda K, local_model, *_a, **_kw: [ + [{ + "type": "generator", + "model": local_model, + "model_type": "vllm_local", + "top_k": 1, + "temperature": 0.0, + "max_tokens": 2048, + "samples": 1, + }], + ], + } + + +@AgentRegistry.register("archon") +class ArchonAgent(LocalCloudAgent): + """Layered (generator → ranker → fuser) inference-time search.""" + + agent_id = "archon" + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + if "OPENAI_API_KEY" not in os.environ and "ANTHROPIC_API_KEY" not in os.environ: + raise RuntimeError("Archon needs OPENAI_API_KEY and/or ANTHROPIC_API_KEY") + + cfg = self._cfg + 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")) + ) + if swe_mode: + return self._run_swe(input, task_meta, cfg) + + _apply_patches_once() + from archon.completions import Archon # type: ignore[import-not-found] + from archon.completions.components.Generator import ( + GENERATE_MAP as _GMAP, # type: ignore[import-not-found] + ) + + arch = cfg.get("architecture", "ensemble_rank_fuse") + presets = _presets() + if arch not in presets: + raise ValueError( + f"unknown archon architecture {arch!r}, choose from {sorted(presets)}" + ) + if not self._local_endpoint or not self._local_model: + raise ValueError( + "ArchonAgent needs local_model + local_endpoint; got " + f"model={self._local_model!r} endpoint={self._local_endpoint!r}" + ) + + K = int(cfg.get("n_samples", 5)) + max_tokens = int(cfg.get("max_tokens", 2048)) + temperature = float(cfg.get("temperature", 0.7)) + ranker_model = cfg.get("ranker_model", self._cloud_model) + fuser_model = cfg.get("fuser_model", self._cloud_model) + + # Register our local-vLLM generator per-run (endpoint can vary per cell). + _GMAP["vllm_local"] = _make_local_generator( + self._local_endpoint, self._local_model + ) + + layers = presets[arch]( + K, self._local_model, ranker_model, fuser_model, max_tokens, temperature, + ) + archon_cfg = {"name": f"hybrid-archon-{arch}", "layers": layers} + + _reset_tally() + archon = Archon(archon_cfg) + + try: + answer = archon.generate([ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": input}, + ]) + except Exception as e: + answer = f"[archon error: {e!r}]" + + if isinstance(answer, list): + answer = answer[-1] if answer else "" + answer = str(answer) + + cp = _tally()["cloud_prompt"] + cc = _tally()["cloud_completion"] + 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)) + + 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, + "traces": { + "architecture": arch, + "n_samples": K, + "ranker_model": ranker_model, + "fuser_model": fuser_model, + "local_model": self._local_model, + "tokens_breakdown": dict(_tally()), + }, + } + return answer, meta + + # ------------------------------------------------------------------ + # SWE-bench variant: K diverse mini-SWE-agent runs, cloud ranker picks + # ------------------------------------------------------------------ + + def _run_swe( + self, + input: str, + task: Dict[str, Any], + cfg: Dict[str, Any], + ) -> Tuple[str, Dict[str, Any]]: + if not self._local_endpoint or not self._local_model: + raise ValueError( + "ArchonAgent (swe mode) needs local_model + local_endpoint" + ) + + K = int(cfg.get("n_samples", 3)) + 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)) + + # K independent runs, each on a FRESH workdir → K diverse patches. + candidates: List[Dict[str, Any]] = [] + total_tokens_local = 0 + total_tokens_cloud = 0 + total_cost = 0.0 + for k in range(K): + out = run_swe_agent_loop( + task, + backbone="local", + backbone_model=self._local_model, + local_endpoint=self._local_endpoint, + initial_prompt=input, + max_turns=max_turns, + bash_timeout=bash_timeout, + output_cap=output_cap, + turn_max_tokens=turn_max_tokens, + trace_prefix=f"archon_gen{k}", + ) + candidates.append({ + "idx": k, + "summary": out["final_summary"], + "patch": out["patch"], + "framed": out["answer"], + "tokens_in": out["tokens_in"], + "tokens_out": out["tokens_out"], + "turns": out["turns"], + }) + total_tokens_local += out["tokens_in"] + out["tokens_out"] + self.record_trace_event({ + "kind": "archon_swe_candidate", + "idx": k, + "patch_chars": len(out["patch"]), + "summary": out["final_summary"], + }) + + # Ranker: cloud picks the best candidate. + ranker_user = ( + f"Issue:\n{task.get('problem_statement','')}\n\n" + f"K = {K} candidate patches:\n\n" + + "\n\n".join( + f"=== Candidate {c['idx']} ===\nSummary: {c['summary']}\n" + f"Patch (chars={len(c['patch'])}):\n```diff\n{c['patch']}```" + for c in candidates + ) + ) + ranker_text, r_in, r_out = self._call_cloud( + user=ranker_user, + system=ARCHON_SWE_RANKER_SYS, + max_tokens=int(cfg.get("ranker_max_tokens", 1024)), + temperature=0.0, + ) + total_tokens_cloud += r_in + r_out + total_cost += self.cost_usd(self._cloud_model, r_in, r_out) + + # Parse — first integer on its own line. + chosen_idx = 0 + for line in ranker_text.splitlines(): + stripped = line.strip() + if stripped.isdigit(): + chosen_idx = int(stripped) + break + if not (0 <= chosen_idx < K): + chosen_idx = 0 + chosen = candidates[chosen_idx] + + self.record_trace_event({ + "kind": "archon_swe_rank", + "chosen_idx": chosen_idx, + "ranker_raw": ranker_text, + "tokens_in": r_in, + "tokens_out": r_out, + }) + + meta = { + "tokens_local": total_tokens_local, + "tokens_cloud": total_tokens_cloud, + "cost_usd": total_cost, + "turns": sum(c["turns"] for c in candidates) + 1, + "traces": { + "swe_mode": True, + "K": K, + "candidates": [ + {"idx": c["idx"], "summary": c["summary"], + "patch_chars": len(c["patch"]), "turns": c["turns"]} + for c in candidates + ], + "chosen_idx": chosen_idx, + "ranker_text": ranker_text, + }, + } + return chosen["framed"], meta + + +__all__ = ["ArchonAgent"] diff --git a/src/openjarvis/agents/hybrid/conductor.py b/src/openjarvis/agents/hybrid/conductor.py new file mode 100644 index 00000000..f21e47e9 --- /dev/null +++ b/src/openjarvis/agents/hybrid/conductor.py @@ -0,0 +1,514 @@ +"""ConductorAgent — static-DAG planner (Sakana AI, arXiv 2512.04388). + +Stage-1 inference-only repro. The paper's trained Qwen2.5-7B conductor is +not released; we substitute a strong zero-shot cloud planner (default +Opus) and run the same plan-then-execute machinery. + +Pipeline per task: + +1. **Plan** — the conductor reads the question + numbered worker pool + and emits three lists ``(model_id, subtasks, access_list)`` in JSON, + up to 5 steps. +2. **Execute** — for each step ``i``: build the worker prompt from + ``subtasks[i]`` + the concatenated prior ``(subtask, output)`` + messages selected by ``access_list[i]``; call worker + ``model_id[i]``; the final answer is the output of the last step. + +On plan parse failure: retry once with a stricter "JSON only" prompt; +on second failure, fall back to a single call to the strongest available +worker (last in the pool by convention). + +Workers come from ``cfg["workers"]`` or a sensible default pool +(local Qwen if vLLM is up, plus Opus 4.7 and gpt-5-mini). + +Hybrid harness result: ``conductor-swebenchverified-opusplan-30`` = 0.367 +acc / $0.22 per task — +10pp vs baseline-cloud at ~15× cheaper. + +Ported from ``hybrid-local-cloud-compute/adapters/conductor_adapter.py``. +""" + +from __future__ import annotations + +import ast +import json +import re +import shutil +import tempfile +import urllib.request +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from openjarvis.agents._stubs import AgentContext +from openjarvis.agents.hybrid._base import LocalCloudAgent +from openjarvis.agents.hybrid._prices import ( + is_gpt5_family, + supports_temperature, +) +from openjarvis.agents.hybrid.mini_swe_agent import ( + _clone_repo, + _extract_diff, + run_swe_agent_loop, +) +from openjarvis.core.registry import AgentRegistry + +CONDUCTOR_SYS = """\ +Your role as an assistant involves obtaining answers to questions by an iterative \ +process of querying powerful language models, each with a different skillset. \ +You will be given the user question and a list of available numbered language \ +models with their metadata. + +Plan up to 5 workflow steps. Output THREE lists of equal length: + + model_id: integers (0..N-1) selecting which numbered model handles each step. + subtasks: natural-language instructions (one string per step) for that model. + access_list: for each step, a list of prior step indices whose (subtask, output) + should be included in that step's prompt; use the string "all" to + include every prior step, or [] for none. The first step must use []. + +Pick the smallest number of steps that will reliably produce a correct final answer. \ +The user only sees the output of the LAST step, so make sure the last step both \ +solves the task and produces the final user-facing answer in the requested format. + +Respond ONLY with a single JSON object (no prose, no markdown fence) with exactly \ +these three keys: + + {"model_id": [...], "subtasks": [...], "access_list": [...]} +""" + +CONDUCTOR_STRICTER = ( + "Your previous response was not valid JSON or was missing required fields. " + "Reply with ONLY a single JSON object — no prose, no code fences, no commentary " + "— containing exactly the three keys model_id (list[int]), subtasks (list[str]), " + "and access_list (list[list[int] or \"all\"]) of equal length, at most 5 entries, " + "and access_list[0] must be [] (an empty list)." +) + + +# ---------- Plan parsing ---------- + +def _strip_fences(s: str) -> str: + s = s.strip() + if s.startswith("```"): + first_nl = s.find("\n") + if first_nl != -1: + s = s[first_nl + 1:] + if s.endswith("```"): + s = s[:-3] + s = s.strip() + return s + + +def _try_json(s: str): + try: + return json.loads(s) + except Exception: + return None + + +def _try_literal(s: str): + """Fallback for the paper's literal Python-list output style.""" + out = {} + for key in ("model_id", "subtasks", "access_list"): + m = re.search( + rf"{key}\s*=\s*(\[[^\]]*\](?:\s*\+\s*\[[^\]]*\])*)", s, re.DOTALL + ) + if not m: + return None + try: + out[key] = ast.literal_eval(m.group(1)) + except Exception: + return None + return out + + +def _validate_plan(plan: Any, n_workers: int) -> Optional[str]: + if not isinstance(plan, dict): + return "plan is not a dict" + for k in ("model_id", "subtasks", "access_list"): + if k not in plan: + return f"missing key {k!r}" + mi, st, al = plan["model_id"], plan["subtasks"], plan["access_list"] + if not (isinstance(mi, list) and isinstance(st, list) and isinstance(al, list)): + return "fields must be lists" + if not (len(mi) == len(st) == len(al)): + return f"unequal lengths: {len(mi)}/{len(st)}/{len(al)}" + if not (1 <= len(mi) <= 5): + return f"need 1..5 steps, got {len(mi)}" + for i, m in enumerate(mi): + if not isinstance(m, int) or not (0 <= m < n_workers): + return f"model_id[{i}]={m!r} out of range 0..{n_workers - 1}" + for i, s in enumerate(st): + if not isinstance(s, str) or not s.strip(): + return f"subtasks[{i}] must be non-empty string" + for i, a in enumerate(al): + if a == "all": + continue + if not isinstance(a, list): + return f"access_list[{i}] must be list or \"all\"" + for j in a: + if not isinstance(j, int) or not (0 <= j < i): + return f"access_list[{i}] has bad ref {j!r}" + return None + + +def _parse_plan(text: str, n_workers: int): + s = _strip_fences(text) + for candidate in (_try_json(s), _try_literal(s)): + if candidate is None: + continue + err = _validate_plan(candidate, n_workers) + if err is None: + return candidate, None + return None, "could not parse a valid plan" + + +# ---------- Worker pool ---------- + +def _vllm_alive(base_url: str) -> bool: + try: + with urllib.request.urlopen( + base_url.rstrip("/") + "/models", timeout=3 + ) as r: + return r.status == 200 + except Exception: + return False + + +def _default_pool(local_model: Optional[str], local_endpoint: Optional[str]) -> List[Dict[str, Any]]: + pool: List[Dict[str, Any]] = [] + if local_model and local_endpoint and _vllm_alive(local_endpoint): + pool.append({ + "id": len(pool), + "name": "local-qwen", + "endpoint": "vllm", + "model": local_model, + "base_url": local_endpoint, + "api_key": "EMPTY", + "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." + ), + }) + 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." + ), + }) + return pool + + +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 ( + f"Available models:\n{_format_worker_pool(workers)}\n\n" + f"User question:\n{question}\n" + ) + + +def _build_step_prompt( + question: str, + subtask: str, + prior_steps: List[Dict[str, Any]], + access: "list[int] | str", +) -> str: + indices = list(range(len(prior_steps))) if access == "all" else list(access) + pieces = [f"User question:\n{question}\n"] + if indices: + pieces.append("Previous routing messages:") + for j in indices: + ps = prior_steps[j] + pieces.append( + f"[Step {j} subtask]\n{ps['subtask']}\n" + f"[Step {j} response]\n{ps['output']}" + ) + pieces.append(f"Your subtask:\n{subtask}") + return "\n\n".join(pieces) + + +# ---------- Worker invocation ---------- + +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).""" + ep = (worker.get("endpoint") or "openai").lower() + max_tok = int(cfg.get("worker_max_tokens", 4096)) + temp = float(cfg.get("worker_temperature", 0.2)) + + if ep == "vllm": + text, p, c = LocalCloudAgent._call_vllm( + worker["model"], + worker["base_url"], + user=prompt, + max_tokens=max_tok, + temperature=temp, + enable_thinking=False, + ) + return text, p, c, True + if ep == "openai": + 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 + if ep == "anthropic": + eff_temp = temp if supports_temperature(worker["model"]) else 0.0 + text, p, c, _ = LocalCloudAgent._call_anthropic( + worker["model"], + user=prompt, + max_tokens=max_tok, + temperature=eff_temp, + ) + return text, p, c, False + raise ValueError(f"unsupported worker endpoint: {ep!r}") + + +def _swe_worker_step( + worker: Dict[str, Any], + task: Dict[str, Any], + prompt: str, + cfg: Dict[str, Any], + workdir: Path, + step_idx: int, +) -> Tuple[str, int, int, bool]: + """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``.""" + ep = (worker.get("endpoint") or "openai").lower() + if ep == "vllm": + backbone, model, endpoint, is_local = ( + "local", worker["model"], worker.get("base_url"), True, + ) + elif ep == "anthropic": + backbone, model, endpoint, is_local = ( + "cloud", worker["model"], None, False, + ) + 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) + out = run_swe_agent_loop( + task, + backbone=backbone, + backbone_model=model, + cloud_endpoint="anthropic" if backbone == "cloud" else "anthropic", + local_endpoint=endpoint, + initial_prompt=prompt, + max_turns=int(cfg.get("swe_max_turns", 30)), + bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)), + output_cap=int(cfg.get("swe_output_cap", 10_000)), + turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)), + trace_prefix=f"conductor_step{step_idx}", + workdir=workdir, + ) + return out["final_summary"] or out["answer"], out["tokens_in"], out["tokens_out"], is_local + + +@AgentRegistry.register("conductor") +class ConductorAgent(LocalCloudAgent): + """Plan-then-execute static DAG over a worker pool. See module docstring.""" + + agent_id = "conductor" + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + question = input + cfg = self._cfg + workers = cfg.get("workers") or _default_pool( + self._local_model, self._local_endpoint + ) + if not workers: + raise RuntimeError("conductor: empty worker pool") + + # 1. Plan + user = _build_conductor_prompt(question, workers) + plan_text, p_in, p_out = self._call_cloud( + user=user, + system=CONDUCTOR_SYS, + max_tokens=int(cfg.get("conductor_max_tokens", 2048)), + temperature=0.0, + ) + plan, err = _parse_plan(plan_text, len(workers)) + parse_attempts = [{"text": plan_text, "error": err}] + conductor_p_in, conductor_p_out = p_in, p_out + + if plan is None: + plan_text2, p_in2, p_out2 = self._call_cloud( + user=user, + system=CONDUCTOR_SYS + "\n\n" + CONDUCTOR_STRICTER, + max_tokens=int(cfg.get("conductor_max_tokens", 2048)), + temperature=0.0, + ) + conductor_p_in += p_in2 + conductor_p_out += p_out2 + plan, err2 = _parse_plan(plan_text2, len(workers)) + parse_attempts.append({"text": plan_text2, "error": err2}) + + fallback_used = False + if plan is None: + fallback_used = True + plan = { + "model_id": [len(workers) - 1], + "subtasks": [question], + "access_list": [[]], + } + + self.record_trace_event({ + "kind": "conductor_plan", + "plan": plan, + "fallback_used": fallback_used, + "parse_attempts": parse_attempts, + "workers": [ + {k: v for k, v in w.items() if k != "api_key"} + for w in workers + ], + }) + + # 2. Execute + # If we're on a SWE-bench task AND cfg["swe_use_agent_loop"] is on, + # 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")) + ) + steps: List[Dict[str, Any]] = [] + tokens_local = 0 + tokens_cloud = 0 + cost = 0.0 + final_answer = "" + shared_workdir: Optional[Path] = None + + try: + if swe_mode: + shared_workdir = Path(tempfile.mkdtemp( + prefix=f"conductor-swe-{task_meta.get('task_id','x')}-" + )) + _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir) + self.record_trace_event({ + "kind": "conductor_swe_workdir", + "workdir": str(shared_workdir), + "repo": task_meta["repo"], + "base_commit": task_meta["base_commit"], + }) + + for i, (mid, subtask, access) in enumerate( + zip(plan["model_id"], plan["subtasks"], plan["access_list"]) + ): + worker = workers[mid] + prompt = _build_step_prompt(question, subtask, steps, access) + self.record_trace_event({ + "kind": "conductor_step_dispatch", + "step_idx": i, + "worker_id": mid, + "worker_name": worker["name"], + "worker_model": worker["model"], + "subtask": subtask, + "access": access, + "prompt": prompt, + "swe_mode": swe_mode, + }) + + if swe_mode: + text, w_in, w_out, is_local = _swe_worker_step( + worker, task_meta, prompt, cfg, shared_workdir, i, + ) + else: + text, w_in, w_out, is_local = _call_worker(worker, prompt, cfg) + + if is_local: + tokens_local += w_in + w_out + else: + tokens_cloud += w_in + w_out + cost += self.cost_usd(worker["model"], w_in, w_out) + steps.append({ + "step_idx": i, + "model_id": mid, + "worker_name": worker["name"], + "worker_model": worker["model"], + "subtask": subtask, + "access": access, + "output": text, + "tokens_in": w_in, + "tokens_out": w_out, + }) + final_answer = text + + # For SWE mode, the authoritative patch is whatever lives in + # the working tree at the end — replace whatever the last + # worker emitted with the full diff (so scoring picks it up). + if swe_mode and shared_workdir is not None: + patch = _extract_diff(shared_workdir) + if patch.strip(): + final_answer = ( + f"{final_answer}\n\n```diff\n{patch}```" + if final_answer else f"```diff\n{patch}```" + ) + finally: + if shared_workdir is not None: + shutil.rmtree(shared_workdir, ignore_errors=True) + + # Conductor (planner) cost goes into cloud bucket + cost += self.cost_usd(self._cloud_model, conductor_p_in, conductor_p_out) + tokens_cloud += conductor_p_in + conductor_p_out + + traces = [ + (s["step_idx"], s["model_id"], s["subtask"], s["output"]) + for s in steps + ] + + meta = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": cost, + "turns": len(steps) + 1, # planner + N execution steps + "traces": { + "steps": traces, + "plan": plan, + "fallback_used": fallback_used, + "parse_attempts": parse_attempts, + "workers": [ + {k: v for k, v in w.items() if k != "api_key"} + for w in workers + ], + }, + } + return final_answer, meta + + +__all__ = ["ConductorAgent"] diff --git a/src/openjarvis/agents/hybrid/mini_swe_agent.py b/src/openjarvis/agents/hybrid/mini_swe_agent.py new file mode 100644 index 00000000..7fad36c7 --- /dev/null +++ b/src/openjarvis/agents/hybrid/mini_swe_agent.py @@ -0,0 +1,635 @@ +"""MiniSWEAgent — vendored, ~330-line port of mini-SWE-agent v2. + +Single-LLM agent loop with a ``bash`` tool, run inside a per-task git +clone. The model iterates: read files, grep, run tests, edit, retry — +the environment-interaction loop that turns SWE-bench from "predict the +patch blind" (~0.30) into "actually fix the bug" (~0.77 for frontier +models). + +Two ways to use this module: + +1. **Standalone agent** — :class:`MiniSWEAgent` registered as + ``mini_swe_agent``. Use it directly as the agent for a cell. +2. **As a worker subroutine inside another paradigm** — call + :func:`run_swe_agent_loop(task, ...)`. Returns a dict with the final + patch, token totals, cost, etc. This is how Minions / Conductor / + Advisors / SkillOrchestra / ToolOrchestra / Archon swap their + one-shot worker call for a real agent loop when running SWE-bench. + +Differences vs. the upstream +(https://github.com/swe-agent/mini-swe-agent): + +- No Docker sandbox. We clone the SWE-bench repo into a tempdir and + exec bash there. Network is available (pip etc.). Treat outputs as + untrusted — model can run ``rm -rf`` against its own workdir, but the + workdir is disposable. Don't run this on a host with secrets in the + CWD. +- One tool, ``bash``. No separate ``submit`` — the loop ends when the + model produces a turn with no tool calls. We extract the patch from + ``git diff`` in the workdir at that point. +- Trace events captured via the LocalCloudAgent thread-local trace + buffer so every bash invocation + result lands in + ``experiments//logs/.json``. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import tempfile +import time +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._prices import ( + cost as estimate_cost, +) +from openjarvis.agents.hybrid._prices import ( + supports_temperature, +) +from openjarvis.core.registry import AgentRegistry + +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, \ +stderr, and the exit code. + +Your task: +1. Read the issue. +2. Use `bash` to explore the repo, read relevant files, and understand the bug. +3. Edit files to fix the bug. You can use `bash` for that too (sed, python -c '...', cat > file < None: + """Shallow-fetch the SWE-bench repo at the right commit into ``dest``.""" + url = f"https://github.com/{repo}.git" + subprocess.run( + ["git", "clone", "--quiet", url, str(dest)], + check=True, timeout=300, capture_output=True, + ) + subprocess.run( + ["git", "checkout", "--quiet", base_commit], + cwd=str(dest), check=True, timeout=120, capture_output=True, + ) + + +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.""" + t0 = time.time() + try: + proc = subprocess.run( + ["bash", "-lc", command], + cwd=str(workdir), + capture_output=True, text=True, timeout=timeout, + ) + stdout = proc.stdout + stderr = proc.stderr + 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 "" + exit_code = -1 + timed_out = True + truncated = False + if len(stdout) > output_cap: + stdout = stdout[:output_cap] + f"\n…[+{len(stdout) - output_cap} chars truncated]" + truncated = True + if len(stderr) > output_cap: + stderr = stderr[:output_cap] + f"\n…[+{len(stderr) - output_cap} chars truncated]" + truncated = True + return { + "stdout": stdout, + "stderr": stderr, + "exit_code": exit_code, + "timed_out": timed_out, + "truncated": truncated, + "latency_s": time.time() - t0, + } + + +def _format_observation(result: Dict[str, Any]) -> str: + parts = [f"exit_code: {result['exit_code']}"] + if result.get("timed_out"): + parts.append("[TIMED OUT]") + if result.get("truncated"): + parts.append("[output truncated]") + if result["stdout"]: + parts.append(f"--- stdout ---\n{result['stdout']}") + if result["stderr"]: + parts.append(f"--- stderr ---\n{result['stderr']}") + return "\n".join(parts) + + +def _extract_diff(workdir: Path) -> str: + """``git diff`` against the base commit — the final SWE-bench patch.""" + proc = subprocess.run( + ["git", "diff", "--no-color"], + cwd=str(workdir), capture_output=True, text=True, timeout=60, + ) + return proc.stdout + + +def _anthropic_assistant_block(block: Any) -> Dict[str, Any]: + """Convert an Anthropic content block back into the dict shape the API + expects for assistant-role messages.""" + btype = getattr(block, "type", None) + if btype == "tool_use": + return { + "type": "tool_use", + "id": block.id, + "name": block.name, + "input": block.input, + } + if hasattr(block, "text"): + return {"type": "text", "text": block.text} + return {"type": btype or "unknown"} + + +# ---------- Reusable agent-loop entry point ---------- + +def run_swe_agent_loop( + task: Dict[str, Any], + *, + backbone: str, # "cloud" or "local" + backbone_model: str, + cloud_endpoint: str = "anthropic", + local_endpoint: Optional[str] = None, + initial_prompt: Optional[str] = None, + max_turns: int = 30, + bash_timeout: int = 120, + output_cap: int = 10_000, + turn_max_tokens: int = 4096, + trace_prefix: str = "mini_swe", + workdir: Optional[Path] = None, +) -> Dict[str, Any]: + """Run a mini-SWE-agent loop for one SWE-bench task. Returns: + + .. code-block:: python + + { + "answer": str, # final framed answer with ```diff fence + "patch": str, # raw unified diff from git diff + "final_summary": str, # the no-tool-call assistant text (may be empty) + "tokens_in": int, + "tokens_out": int, + "tokens_local": int, # bookkeeping split for paradigms + "tokens_cloud": int, + "cost_usd": float, + "turns": int, + "max_turns_hit": bool, + "workdir": str, + } + + Captures every bash invocation + LLM turn into the active trace buffer + via :func:`_record_event` from the LocalCloudAgent base, so callers + don't have to do their own per-call instrumentation. + + Args: + task: SWE-bench-shaped dict with ``repo`` + ``base_commit`` + ``task_id`` + + (optional) ``problem_statement`` / ``hints_text``. + backbone: ``"cloud"`` to drive the loop with the cloud model + (Anthropic only today), ``"local"`` for vLLM. + backbone_model: model id for the loop's backbone. + cloud_endpoint / local_endpoint: SDK targets. + initial_prompt: if set, used as the first user message (paradigms + embed orchestrator context in here). If None, falls back to the + task's problem_statement. + workdir: pre-cloned repo path. If None, this function clones the + repo into a tempdir and cleans it up at the end. Paradigms that + want to chain multiple subloops over the same working tree can + manage their own workdir. + """ + repo = task.get("repo") or "" + base_commit = task.get("base_commit") or "" + if not repo or not base_commit: + raise ValueError( + f"run_swe_agent_loop needs task['repo'] + task['base_commit']; " + f"got repo={repo!r}, base_commit={base_commit!r}" + ) + + own_workdir = workdir is None + if own_workdir: + workdir = Path(tempfile.mkdtemp( + prefix=f"mini-swe-{task.get('task_id','x')}-" + )) + try: + _clone_repo(repo, base_commit, workdir) + except Exception: + shutil.rmtree(workdir, ignore_errors=True) + raise + + _record_event({ + "kind": f"{trace_prefix}_setup", + "repo": repo, + "base_commit": base_commit, + "workdir": str(workdir), + "owns_workdir": own_workdir, + "backbone": backbone, + "backbone_model": backbone_model, + "ts": time.time(), + }) + + user_prompt = initial_prompt or task.get("problem_statement") or "" + + try: + if backbone == "cloud": + result = _loop_cloud( + user_prompt, workdir, + model=backbone_model, + cloud_endpoint=cloud_endpoint, + max_turns=max_turns, + bash_timeout=bash_timeout, + output_cap=output_cap, + turn_max_tokens=turn_max_tokens, + trace_prefix=trace_prefix, + ) + elif backbone == "local": + if not local_endpoint: + raise ValueError("run_swe_agent_loop(backbone='local') needs local_endpoint") + result = _loop_local( + user_prompt, workdir, + model=backbone_model, + endpoint=local_endpoint, + max_turns=max_turns, + bash_timeout=bash_timeout, + output_cap=output_cap, + turn_max_tokens=turn_max_tokens, + trace_prefix=trace_prefix, + ) + else: + raise ValueError(f"unsupported backbone: {backbone!r}") + + patch = _extract_diff(workdir) + framed = (result["final_summary"] or "[mini-swe-agent produced no summary text]") + if patch.strip(): + framed = f"{framed}\n\n```diff\n{patch}```" + + return { + "answer": framed, + "patch": patch, + "final_summary": result["final_summary"], + "tokens_in": result["tokens_in"], + "tokens_out": result["tokens_out"], + "tokens_local": result["tokens_in"] + result["tokens_out"] if backbone == "local" else 0, + "tokens_cloud": result["tokens_in"] + result["tokens_out"] if backbone == "cloud" else 0, + "cost_usd": ( + estimate_cost(backbone_model, result["tokens_in"], result["tokens_out"]) + if backbone == "cloud" else 0.0 + ), + "turns": result["turns"], + "max_turns_hit": result["max_turns_hit"], + "workdir": str(workdir), + } + finally: + if own_workdir: + shutil.rmtree(workdir, ignore_errors=True) + + +# ---------- Cloud loop (Anthropic multi-turn with tools) ---------- + +def _loop_cloud( + problem: str, + workdir: Path, + *, + model: str, + cloud_endpoint: str, + max_turns: int, + bash_timeout: int, + output_cap: int, + 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}" + ) + import anthropic + client = anthropic.Anthropic(timeout=600.0, max_retries=5) + messages: List[Dict[str, Any]] = [{"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, + "system": SYSTEM_PROMPT, + "max_tokens": turn_max_tokens, + "tools": [BASH_TOOL_ANTHROPIC], + "messages": messages, + } + if supports_temperature(model): + kwargs["temperature"] = 0.0 + t0 = time.time() + msg = client.messages.create(**kwargs) + latency = time.time() - t0 + tokens_in += msg.usage.input_tokens + tokens_out += msg.usage.output_tokens + + content_blocks: List[Dict[str, Any]] = [] + tool_uses: List[Tuple[str, str, Dict[str, Any]]] = [] + text_parts: List[str] = [] + for block in msg.content: + btype = getattr(block, "type", None) + if btype == "tool_use": + tool_uses.append((block.id, block.name, dict(block.input or {}))) + content_blocks.append({ + "type": "tool_use", "id": block.id, "name": block.name, + "input": dict(block.input or {}), + }) + elif hasattr(block, "text"): + text_parts.append(block.text) + content_blocks.append({"type": "text", "text": block.text}) + else: + content_blocks.append({"type": btype or "unknown"}) + + _record_event({ + "kind": f"{trace_prefix}_turn", + "turn": turn, + "stop_reason": msg.stop_reason, + "tokens_in": msg.usage.input_tokens, + "tokens_out": msg.usage.output_tokens, + "latency_s": latency, + "content_blocks": content_blocks, + "ts": time.time(), + }) + + messages.append({"role": "assistant", "content": [ + _anthropic_assistant_block(b) for b in msg.content + ]}) + + if not tool_uses: + final_text = "\n".join(text_parts).strip() + break + + tool_result_blocks: List[Dict[str, Any]] = [] + for tu_id, tu_name, tu_input in tool_uses: + if tu_name != "bash": + obs = f"unknown tool: {tu_name!r}" + _record_event({ + "kind": f"{trace_prefix}_unknown_tool", + "turn": turn, "name": tu_name, "input": tu_input, + "ts": time.time(), + }) + else: + command = str(tu_input.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) + tool_result_blocks.append({ + "type": "tool_result", + "tool_use_id": tu_id, + "content": obs, + }) + messages.append({"role": "user", "content": tool_result_blocks}) + + 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) ---------- + +def _loop_local( + problem: str, + workdir: Path, + *, + model: str, + endpoint: str, + max_turns: int, + bash_timeout: int, + output_cap: int, + turn_max_tokens: int, + trace_prefix: str, +) -> Dict[str, Any]: + from openai import OpenAI + client = OpenAI(base_url=endpoint, api_key="EMPTY", 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 + 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}}, + ) + 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, + "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(), + }) + + messages.append({ + "role": "assistant", + "content": text or None, + "tool_calls": [ + { + "id": tc.id, "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in tool_calls + ] if tool_calls else None, + }) + + if not tool_calls: + 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}" + 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, + } + + +# ---------- Standalone agent ---------- + +@AgentRegistry.register("mini_swe_agent") +class MiniSWEAgent(LocalCloudAgent): + """Single-model bash-loop agent for SWE-bench-shaped tasks. + + Configurable knobs via ``cfg``: + + - ``backbone`` (str, default ``"cloud"``): ``"cloud"`` or ``"local"``. + - ``max_turns`` (int, default 30): hard cap on tool turns. + - ``bash_timeout_s`` (int, default 120): per-command timeout. + - ``output_cap`` (int, default 10_000): per-command stdout/stderr cap. + - ``turn_max_tokens`` (int, default 4096): max_tokens per LLM turn. + """ + + agent_id = "mini_swe_agent" + + 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 {} + + backbone = cfg.get("backbone", "cloud") + model = ( + self._cloud_model if backbone == "cloud" + else (self._local_model or "") + ) + + out = run_swe_agent_loop( + task, + backbone=backbone, + backbone_model=model, + cloud_endpoint=self._cloud_endpoint, + local_endpoint=self._local_endpoint, + initial_prompt=input, + max_turns=int(cfg.get("max_turns", 30)), + bash_timeout=int(cfg.get("bash_timeout_s", 120)), + output_cap=int(cfg.get("output_cap", 10_000)), + turn_max_tokens=int(cfg.get("turn_max_tokens", 4096)), + ) + meta = { + "tokens_local": out["tokens_local"], + "tokens_cloud": out["tokens_cloud"], + "cost_usd": out["cost_usd"], + "turns": out["turns"], + "traces": { + "backbone": backbone, + "max_turns_hit": out["max_turns_hit"], + "patch_chars": len(out["patch"]), + "final_summary": out["final_summary"], + }, + } + return out["answer"], meta + + +__all__ = ["MiniSWEAgent", "run_swe_agent_loop", "SYSTEM_PROMPT"] diff --git a/src/openjarvis/agents/hybrid/minions.py b/src/openjarvis/agents/hybrid/minions.py new file mode 100644 index 00000000..73d064f8 --- /dev/null +++ b/src/openjarvis/agents/hybrid/minions.py @@ -0,0 +1,554 @@ +"""MinionsAgent — port of HazyResearch Minions protocol. + +Cloud supervisor decomposes the task and reads back local-worker output; +local worker(s) do the bulk reading/extraction. Multi-turn loop until the +supervisor commits to a final answer. + +Two modes (``cfg["mode"]``): + +- ``"minion"`` — single local worker, one cloud supervisor (cheaper). + Default. +- ``"minions"`` — parallel local workers, cloud aggregator. + +Hybrid harness result: ``minions-swebenchverified-qwen27b-opus-500`` = +0.274 acc / $0.09 per task — beats baseline-cloud's 0.236 / $0.95 on +**both** accuracy and cost. GAIA at n=165 ties baseline-cloud at 0.576 +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()``. + +Compatibility patches applied at first ``run()`` (idempotent): + +- Strip ``temperature`` for Opus 4.7+ (rejected with 400). +- Inject server-side ``output_config`` JSON schema on supervisor turns + so Opus replies in the shape Minions's parser expects (per-turn schema + picked by sniffing the prompt for ``"decision": "provide_final_answer"``). +- Replace Minions's ``_extract_json`` with a wrapper that short-circuits + when the response is already valid JSON. +- Inject ``timeout=600``/``max_retries=5`` defaults into + ``anthropic.Anthropic()`` — Minions builds bare clients which 60s-timeout + under SWE-bench concurrency=8. + +Ported from ``hybrid-local-cloud-compute/adapters/minions_adapter.py``. +""" + +from __future__ import annotations + +import json as _json +import sys +import types +from typing import Any, Dict, List, Optional, Tuple + +from openjarvis.agents._stubs import AgentContext +from openjarvis.agents.hybrid._base import ( + ANTHROPIC_WEB_SEARCH_TOOL, + WEB_SEARCH_COST_PER_CALL, + LocalCloudAgent, +) +from openjarvis.agents.hybrid._prices import NO_TEMP_PREFIXES +from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop +from openjarvis.core.registry import AgentRegistry + +MINIONS_SWE_PLANNER_SYS = ( + "You are the cloud supervisor in a Minions setup. The small local model " + "is about to run an agent loop (with shell access) against a Python " + "repository to fix a bug. Read the issue and write a concise plan: " + "what files the local model should look at first, what tests are most " + "relevant, and 1-3 concrete approaches it should try. Be specific — " + "this is the local model's only briefing from you. Reply in 8 bullet " + "points or fewer." +) + + +# ---------- Per-turn JSON schemas (server-side enforcement) ---------- +# +# Minions's supervisor produces different JSON shapes per turn: +# turn 1 (decompose): {reasoning, message} +# turn 2+ (continue): {decision="request_additional_info", message} +# turn 2+ (final answer): {decision="provide_final_answer", answer} +# +# Anthropic strict mode requires additionalProperties:false + all-props +# required, so we pick the schema PER TURN by sniffing the prompt for +# Minions's turn-2 template marker. + +MINIONS_FIRST_TURN_SCHEMA = { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "reasoning": {"type": "string"}, + "message": {"type": "string"}, + }, + "required": ["reasoning", "message"], + "additionalProperties": False, + }, + } +} + +MINIONS_CONVERSATION_SCHEMA = { + "format": { + "type": "json_schema", + "schema": { + "anyOf": [ + { + "type": "object", + "properties": { + "decision": {"const": "request_additional_info"}, + "message": {"type": "string"}, + }, + "required": ["decision", "message"], + "additionalProperties": False, + }, + { + "type": "object", + "properties": { + "decision": {"const": "provide_final_answer"}, + "answer": {"type": "string"}, + }, + "required": ["decision", "answer"], + "additionalProperties": False, + }, + ], + }, + } +} + +# Markers from Minions's supervisor prompts (prompts/minion.py). Any one +# being present in the call's messages/system is a strong Minions signal. +MINIONS_PROMPT_MARKERS = ( + "small language model that has read", # SUPERVISOR_INITIAL_PROMPT + "provide_final_answer", # SUPERVISOR_CONVERSATION_PROMPT + "request_additional_info", +) + + +def _looks_like_minions_call(kwargs: Dict[str, Any]) -> bool: + blob_parts = [str(kwargs.get("system", ""))] + for msg in kwargs.get("messages", []) or []: + content = msg.get("content", "") if isinstance(msg, dict) else "" + blob_parts.append(str(content)) + blob = "\n".join(blob_parts) + return any(m in blob for m in MINIONS_PROMPT_MARKERS) + + +def _minions_turn_schema(kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Pick the schema by sniffing the prompt — see module docstring.""" + blob_parts = [str(kwargs.get("system", ""))] + for msg in kwargs.get("messages", []) or []: + content = msg.get("content", "") if isinstance(msg, dict) else "" + blob_parts.append(str(content)) + blob = "\n".join(blob_parts) + if ( + '"decision": "provide_final_answer"' in blob + or '"decision":"provide_final_answer"' in blob + ): + return MINIONS_CONVERSATION_SCHEMA + return MINIONS_FIRST_TURN_SCHEMA + + +# ---------- Compatibility patches (idempotent) ---------- + +_PATCHES_APPLIED = False + + +def _stub_missing_imports() -> None: + """Minions's clients/__init__.py eager-imports every provider client. + + Two annoyances: + 1. ``mistralai`` 2.x dropped top-level ``Mistral`` → ImportError. + 2. The secure-chat path imports ``nv_attestation_sdk`` which writes a log + file to CWD at import-time. We don't use secure chat. + """ + try: + import mistralai + if not hasattr(mistralai, "Mistral"): + mistralai.Mistral = type("Mistral", (), {}) # type: ignore[attr-defined] + except ImportError: + sys.modules["mistralai"] = types.ModuleType("mistralai") + sys.modules["mistralai"].Mistral = type("Mistral", (), {}) # type: ignore[attr-defined] + sys.modules.setdefault("nv_attestation_sdk", None) # type: ignore[arg-type] + + +def _patch_anthropic_globally() -> None: + import anthropic as _anth_mod + from anthropic.resources.beta.messages import messages as _beta_msgs_mod + from anthropic.resources.messages import messages as _msgs_mod + + # External Minions builds bare anthropic.Anthropic() clients (no timeout + # / max_retries). Under concurrency=8 SWE-bench load those default to + # ~60s and timeout in droves. Inject sane defaults at the constructor. + if not getattr(_anth_mod.Anthropic.__init__, "_hybrid_patched", False): + _orig_init = _anth_mod.Anthropic.__init__ + + def _patched_init(self, *args, **kwargs): # type: ignore[no-untyped-def] + kwargs.setdefault("timeout", 600.0) + kwargs.setdefault("max_retries", 5) + return _orig_init(self, *args, **kwargs) + + _patched_init._hybrid_patched = True # type: ignore[attr-defined] + _anth_mod.Anthropic.__init__ = _patched_init # type: ignore[assignment] + + for cls in (_msgs_mod.Messages, _beta_msgs_mod.Messages): + if getattr(cls.create, "_hybrid_patched", False): + continue + orig = cls.create + + def make_patched(orig): # type: ignore[no-untyped-def] + def patched(self, **kwargs): # type: ignore[no-untyped-def] + model = kwargs.get("model", "") + if model.startswith(NO_TEMP_PREFIXES): + kwargs.pop("temperature", None) + if ( + "output_config" not in kwargs + and _looks_like_minions_call(kwargs) + ): + kwargs["output_config"] = _minions_turn_schema(kwargs) + return orig(self, **kwargs) + patched._hybrid_patched = True # type: ignore[attr-defined] + return patched + + cls.create = make_patched(orig) # 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 + outputs the entire response IS valid JSON, so short-circuit on that. + """ + from minions import minion as _minion_mod # type: ignore[import-not-found] + + if getattr(_minion_mod._extract_json, "_hybrid_patched", False): + return + _orig = _minion_mod._extract_json + + def patched(text): # type: ignore[no-untyped-def] + s = (text or "").strip() + if s.startswith("{") and s.endswith("}"): + try: + return _json.loads(s) + except _json.JSONDecodeError: + pass + return _orig(text) + + patched._hybrid_patched = True # type: ignore[attr-defined] + _minion_mod._extract_json = patched # type: ignore[assignment] + + +def _apply_patches_once() -> None: + global _PATCHES_APPLIED + if _PATCHES_APPLIED: + return + _stub_missing_imports() + _patch_anthropic_globally() + _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]: + """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 + tasks ship with no doc, so we synthesize one by having Opus do an actual + web search first and dump the results back as the worker's context. + + Returns {text, tokens, cost_usd, n_searches}. On any failure: empty text + and zeros — the protocol still runs. + """ + out: Dict[str, Any] = { + "text": "", "tokens": 0, "cost_usd": 0.0, "n_searches": 0, + } + if cloud_endpoint != "anthropic" or not (question or "").strip(): + return out + try: + prompt = ( + "Research the following question using web_search. Do NOT answer it. " + "Instead, gather all relevant facts, numbers, names, dates, sources, " + "and direct quotes you find, and report them as a dense reference " + "document with URLs. The downstream reader is a small LLM that " + "needs raw material to reason over.\n\nQUESTION:\n" + question + ) + text, p, c, n_searches = LocalCloudAgent._call_anthropic( + cloud_model, + user=prompt, + max_tokens=8192, + tools=[ANTHROPIC_WEB_SEARCH_TOOL], + tool_choice={"type": "any"}, + ) + from openjarvis.agents.hybrid._prices import cost as _cost_usd + out.update( + text=text, + tokens=p + c, + cost_usd=_cost_usd(cloud_model, p, c) + n_searches * WEB_SEARCH_COST_PER_CALL, + n_searches=n_searches, + ) + except Exception as e: + out["error"] = f"{type(e).__name__}: {e}" + return out + + +def _context_for( + task: Optional[Dict[str, Any]], prefetched: str = "" +) -> List[str]: + """Minions wants a context list.""" + bits: List[str] = [] + task = task or {} + if task.get("hints_text"): + bits.append(task["hints_text"]) + if task.get("problem_statement") and not task.get("question"): + bits.append(task["problem_statement"]) + if prefetched: + bits.append(prefetched) + return bits or [""] + + +# ---------- Main agent ---------- + +@AgentRegistry.register("minions") +class MinionsAgent(LocalCloudAgent): + """HazyResearch Minions supervisor/worker protocol. See module docstring.""" + + agent_id = "minions" + + def _is_soft_failure(self, exc: BaseException) -> Optional[str]: + # Known soft-failure modes: Qwen worker JSON malformed, Anthropic + # 400/529, KeyError on missing schema fields. + try: + import anthropic + if isinstance(exc, anthropic.BadRequestError): + return f"{type(exc).__name__}: {str(exc)[:120]}" + except Exception: + pass + if isinstance(exc, (_json.JSONDecodeError, ValueError, KeyError)): + return f"{type(exc).__name__}: {str(exc)[:120]}" + if "JSONDecodeError" in type(exc).__name__: + return f"{type(exc).__name__}: {str(exc)[:120]}" + if "prompt is too long" in str(exc): + return f"{type(exc).__name__}: {str(exc)[:120]}" + return None + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + cfg = self._cfg + task_meta: Dict[str, Any] = {} + if context is not None: + task_meta = context.metadata.get("task", {}) or {} + + # SWE-bench branch: the upstream Minions library doesn't fit + # SWE-bench (it's "small model reads docs, big model summarizes"). + # Instead, mirror Minions's "cloud supervises, local does the + # work" pattern: cloud writes a high-level fix plan, local Qwen + # runs mini-SWE-agent with that plan as additional context. + 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: + return self._run_swe(input, task_meta, cfg) + + _apply_patches_once() + from minions.clients.anthropic import ( + AnthropicClient, # type: ignore[import-not-found] + ) + from minions.clients.openai import ( + OpenAIClient, # type: ignore[import-not-found] + ) + from minions.minion import Minion # type: ignore[import-not-found] + from minions.minions import Minions # type: ignore[import-not-found] + + mode = cfg.get("mode", "minion") + + if not self._local_endpoint or not self._local_model: + raise ValueError( + "MinionsAgent needs local_model + local_endpoint; got " + f"model={self._local_model!r} endpoint={self._local_endpoint!r}" + ) + + local_client = OpenAIClient( + model_name=self._local_model, + base_url=self._local_endpoint, + api_key="EMPTY", + temperature=cfg.get("local_temperature", 0.0), + max_tokens=cfg.get("worker_max_tokens", 4096), + local=True, + ) + if self._cloud_endpoint == "openai": + cloud_client = OpenAIClient( + model_name=self._cloud_model, + temperature=0.0, + max_tokens=4096, + ) + elif self._cloud_endpoint == "anthropic": + # Temperature stripping is handled by the global patch above for Opus 4.7+. + cloud_client = AnthropicClient( + model_name=self._cloud_model, + temperature=0.0, + max_tokens=4096, + ) + else: + raise ValueError(f"unsupported cloud endpoint: {self._cloud_endpoint!r}") + + cls = Minions if mode == "minions" else Minion + log_dir = cfg.get("log_dir") or "/tmp/minions_logs" + protocol = cls( + local_client=local_client, + remote_client=cloud_client, + max_rounds=cfg.get("max_rounds", 3), + log_dir=log_dir, + ) + + # 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. + prefetch: Dict[str, Any] = { + "text": "", "tokens": 0, "cost_usd": 0.0, "n_searches": 0, + } + if task_meta.get("question"): + prefetch = _prefetch_context( + task_meta["question"], self._cloud_endpoint, self._cloud_model + ) + + if prefetch.get("text"): + self.record_trace_event({ + "kind": "minions_prefetch", + "n_searches": prefetch["n_searches"], + "tokens": prefetch["tokens"], + "cost_usd": prefetch["cost_usd"], + "text": prefetch["text"], + "error": prefetch.get("error"), + }) + + out = protocol( + task=input, # full formatted prompt (with bench instruction) + context=_context_for(task_meta, prefetched=prefetch["text"]), + doc_metadata=cfg.get("doc_metadata", "task"), + max_rounds=cfg.get("max_rounds", 3), + ) + + # The Minions library doesn't go through our SDK helpers, so the + # auto-trace missed every turn. Record the protocol output directly — + # supervisor_messages + worker_messages contain the full conversation. + self.record_trace_event({ + "kind": "minions_protocol", + "mode": mode, + "supervisor_messages": out.get("supervisor_messages"), + "worker_messages": out.get("worker_messages"), + "timing": out.get("timing"), + "log_file": out.get("log_file"), + "final_answer": out.get("final_answer", ""), + }) + + local_usage = out.get("local_usage") + remote_usage = out.get("remote_usage") + lp = getattr(local_usage, "prompt_tokens", 0) + lc = getattr(local_usage, "completion_tokens", 0) + rp = getattr(remote_usage, "prompt_tokens", 0) + rc = getattr(remote_usage, "completion_tokens", 0) + + meta = { + "tokens_local": lp + lc, + "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), + "traces": { + "mode": mode, + "supervisor_messages": out.get("supervisor_messages"), + "worker_messages": out.get("worker_messages"), + "timing": out.get("timing"), + "log_file": out.get("log_file"), + "prefetch": { + "n_searches": prefetch["n_searches"], + "tokens": prefetch["tokens"], + "cost_usd": prefetch["cost_usd"], + "chars": len(prefetch["text"]), + "error": prefetch.get("error"), + }, + }, + } + return out.get("final_answer", ""), meta + + + # ------------------------------------------------------------------ + # SWE-bench variant + # ------------------------------------------------------------------ + + def _run_swe( + self, + input: str, + task: Dict[str, Any], + cfg: Dict[str, Any], + ) -> Tuple[str, Dict[str, Any]]: + if not self._local_endpoint or not self._local_model: + raise ValueError( + "MinionsAgent (swe mode) still needs local_model + local_endpoint" + ) + # 1. Cloud supervisor writes a high-level plan (no tools). + plan_text, p_in, p_out = self._call_cloud( + user=( + f"Issue:\n{task.get('problem_statement','')}\n\n" + f"Repo: {task.get('repo','')}\n" + f"Base commit: {task.get('base_commit','')}\n\n" + f"{task.get('hints_text','')}" + ), + system=MINIONS_SWE_PLANNER_SYS, + max_tokens=int(cfg.get("supervisor_max_tokens", 1024)), + temperature=0.0, + ) + self.record_trace_event({ + "kind": "minions_swe_plan", + "plan": plan_text, + "tokens_in": p_in, + "tokens_out": p_out, + }) + supervisor_cost = self.cost_usd(self._cloud_model, p_in, p_out) + + # 2. Local worker runs mini-SWE-agent with the plan as context. + worker_prompt = ( + f"{input}\n\n" + f"-----\n" + f"A cloud supervisor reviewed this issue and wrote a fix plan " + f"for you. Use it as guidance, but verify everything with the " + f"actual code via your bash tool:\n\n{plan_text}" + ) + out = run_swe_agent_loop( + task, + backbone="local", + backbone_model=self._local_model, + local_endpoint=self._local_endpoint, + initial_prompt=worker_prompt, + max_turns=int(cfg.get("swe_max_turns", 30)), + bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)), + output_cap=int(cfg.get("swe_output_cap", 10_000)), + turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)), + trace_prefix="minions_worker", + ) + + meta = { + "tokens_local": out["tokens_in"] + out["tokens_out"], + "tokens_cloud": p_in + p_out, + "cost_usd": supervisor_cost, + "turns": 1 + out["turns"], + "traces": { + "swe_mode": True, + "supervisor_plan": plan_text, + "worker_final_summary": out["final_summary"], + "worker_patch_chars": len(out["patch"]), + "max_turns_hit": out["max_turns_hit"], + }, + } + return out["answer"], meta + + +__all__ = ["MinionsAgent"] diff --git a/src/openjarvis/agents/hybrid/registry/advisors.toml b/src/openjarvis/agents/hybrid/registry/advisors.toml new file mode 100644 index 00000000..4f366360 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/advisors.toml @@ -0,0 +1,41 @@ +# Advisors cells (arXiv:2510.02453). +# +# Inference-only repro of the advisor paradigm: cloud drafts → local advisor +# critiques → cloud rewrites. The paper's RL-trained advisor isn't reproduced +# (no training in this harness) — these are the untrained-advisor lower bound. +# Adapter auto-resolves the local model name against /v1/models. + +[cells.advisors-gaia-qwen9b-opus-3] +method = "advisors" +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 = {} + +[cells.advisors-swebenchverified-qwen9b-opus-3] +method = "advisors" +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 = {} +concurrency = 3 + +[cells.advisors-gaia-qwen9b-opus-30] +method = "advisors" +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 = {} +concurrency = 4 + +[cells.advisors-swebenchverified-qwen9b-opus-30] +method = "advisors" +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 = {} +concurrency = 3 diff --git a/src/openjarvis/agents/hybrid/registry/archon.toml b/src/openjarvis/agents/hybrid/registry/archon.toml new file mode 100644 index 00000000..7a3c3096 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/archon.toml @@ -0,0 +1,94 @@ +# Archon cells. Inference-time architecture search (arXiv:2409.15254). +# +# Hybrid wiring: K local proposers (vLLM on 8001/8004) → 1 cloud ranker → +# 1 cloud fuser. method_cfg.architecture = "ensemble_rank_fuse" is the +# only non-debug preset; "single_local" exists for adapter smoke tests +# (no cloud calls). +# +# Naming: archon------ + +# ---- Stage 1 smokes (n=3) ---- + +[cells.archon-gaia-qwen27b-opus-ensemble-K3-3] +method = "archon" +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 = "ensemble_rank_fuse", n_samples = 3, max_tokens = 1024, temperature = 0.7 } + +# Cheap sanity smoke: local-only generator, zero cloud spend. Use to verify +# vLLM + Archon plumbing without paying for Opus. +[cells.archon-gaia-qwen27b-singlelocal-3] +method = "archon" +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 } + +[cells.archon-swebenchverified-qwen27b-opus-ensemble-K3-3] +method = "archon" +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 = { architecture = "ensemble_rank_fuse", n_samples = 3, max_tokens = 1024, temperature = 0.7 } +concurrency = 3 + + +# ---- Stage 1.5 mid (n=30) ---- + +[cells.archon-gaia-qwen27b-opus-ensemble-K5-30] +method = "archon" +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 = "ensemble_rank_fuse", n_samples = 5, max_tokens = 2048, temperature = 0.7 } +concurrency = 4 + +[cells.archon-gaia-qwen27b-singlelocal-30] +method = "archon" +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 } +concurrency = 4 + +[cells.archon-gaia-gemma31b-opus-ensemble-K5-30] +method = "archon" +bench = "gaia" +n = 30 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { architecture = "ensemble_rank_fuse", n_samples = 5, max_tokens = 2048, temperature = 0.7 } + +[cells.archon-swebenchverified-qwen27b-opus-ensemble-K5-30] +method = "archon" +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 = { architecture = "ensemble_rank_fuse", n_samples = 5, max_tokens = 2048, temperature = 0.7 } +concurrency = 3 + + +# ---- Stage 2 full ---- + +[cells.archon-gaia-qwen27b-opus-ensemble-K5-165] +method = "archon" +bench = "gaia" +n = 165 +local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { architecture = "ensemble_rank_fuse", n_samples = 5, max_tokens = 2048, temperature = 0.7 } + +[cells.archon-gaia-gemma31b-opus-ensemble-K5-165] +method = "archon" +bench = "gaia" +n = 165 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { architecture = "ensemble_rank_fuse", n_samples = 5, max_tokens = 2048, temperature = 0.7 } diff --git a/src/openjarvis/agents/hybrid/registry/conductor.toml b/src/openjarvis/agents/hybrid/registry/conductor.toml new file mode 100644 index 00000000..6e7eda06 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/conductor.toml @@ -0,0 +1,54 @@ +# 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. +# +# Naming: conductor--- + +[cells.conductor-gaia-opusplan-5] +method = "conductor" +bench = "gaia" +n = 5 +# vLLM is the default local worker; auto-skipped if /v1/models is unreachable. +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 = 4096, worker_temperature = 0.2, conductor_max_tokens = 2048 } + +[cells.conductor-gaia-gpt5miniplan-5] +method = "conductor" +bench = "gaia" +n = 5 +local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" } +cloud = { model = "gpt-5-mini", endpoint = "openai" } +method_cfg = { worker_max_tokens = 4096, worker_temperature = 0.2, conductor_max_tokens = 2048 } + +[cells.conductor-swebenchverified-opusplan-5] +method = "conductor" +bench = "swebench-verified" +n = 5 +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 } + + +# ---- n=30 (scaled-up cells for paradigms that worked at smoke tier) ---- + +[cells.conductor-gaia-opusplan-30] +method = "conductor" +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 = { worker_max_tokens = 4096, worker_temperature = 0.2, conductor_max_tokens = 2048 } +concurrency = 8 + +[cells.conductor-swebenchverified-opusplan-30] +method = "conductor" +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 = { worker_max_tokens = 8192, worker_temperature = 0.2, conductor_max_tokens = 4096 } +concurrency = 8 diff --git a/src/openjarvis/agents/hybrid/registry/mini_swe_agent.toml b/src/openjarvis/agents/hybrid/registry/mini_swe_agent.toml new file mode 100644 index 00000000..69bf7bc4 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/mini_swe_agent.toml @@ -0,0 +1,33 @@ +# mini-SWE-agent cells. Single-model bash-tool loop on SWE-bench-Verified. +# +# Vendored ~150-line port of mini-SWE-agent v2. No Docker; each task clones +# the repo into a tempdir and execs bash there. See +# `agents/hybrid/mini_swe_agent.py` docstring for the safety caveats. +# +# Backbone is `cloud` by default (drives the loop with the cloud model). +# Set `backbone = "local"` to drive with the vLLM-served local model +# (requires --enable-auto-tool-choice + a tool-call parser). + +[cells.mini-swe-agent-swebenchverified-opus-3] +method = "mini_swe_agent" +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 = { backbone = "cloud", max_turns = 30, bash_timeout_s = 120, turn_max_tokens = 4096 } + +[cells.mini-swe-agent-swebenchverified-opus-30] +method = "mini_swe_agent" +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 = { backbone = "cloud", max_turns = 50, bash_timeout_s = 120, turn_max_tokens = 4096 } + +[cells.mini-swe-agent-swebenchverified-qwen27b-3] +method = "mini_swe_agent" +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 = { backbone = "local", max_turns = 30, bash_timeout_s = 120, turn_max_tokens = 4096 } diff --git a/src/openjarvis/agents/hybrid/registry/minions.toml b/src/openjarvis/agents/hybrid/registry/minions.toml new file mode 100644 index 00000000..bf3eeff4 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/minions.toml @@ -0,0 +1,96 @@ +# Minions cells. One [cells.] block per (bench × local × cloud × N). +# +# Naming: minions---- +# HF model IDs verified Feb 2026 / Gemma 4 release: +# Qwen/Qwen3.5-27B (dense, 262K native context) +# google/gemma-4-31B-it (dense instruct, 256K context — 26B is the MoE variant) + +# ---- Stage 1 smokes (n=3) ---- + +[cells.minions-gaia-qwen27b-opus-3] +method = "minions" +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 = { mode = "minion", max_rounds = 3, worker_max_tokens = 4096 } + +[cells.minions-swebenchverified-qwen27b-opus-3] +method = "minions" +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 = { mode = "minion", max_rounds = 3, worker_max_tokens = 8192 } + + +# ---- Stage 1.5 mid (n=30) — catches cells that smoke-pass but full-fail ---- + +[cells.minions-gaia-qwen27b-opus-30] +method = "minions" +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 = { mode = "minion", max_rounds = 5, worker_max_tokens = 4096 } + +[cells.minions-gaia-gemma31b-opus-30] +method = "minions" +bench = "gaia" +n = 30 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 4096 } + +[cells.minions-swebenchverified-qwen27b-opus-30] +method = "minions" +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 = { mode = "minion", max_rounds = 5, worker_max_tokens = 8192 } + +[cells.minions-swebenchverified-gemma31b-opus-30] +method = "minions" +bench = "swebench-verified" +n = 30 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 8192 } + + +# ---- Stage 2 full ---- + +[cells.minions-gaia-qwen27b-opus-165] +method = "minions" +bench = "gaia" +n = 165 +local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 4096 } +concurrency = 4 + +[cells.minions-gaia-gemma31b-opus-165] +method = "minions" +bench = "gaia" +n = 165 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 4096 } + +[cells.minions-swebenchverified-qwen27b-opus-500] +method = "minions" +bench = "swebench-verified" +n = 500 +local = { model = "Qwen/Qwen3.5-27B-FP8", endpoint = "http://localhost:8001/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 8192 } +concurrency = 4 + +[cells.minions-swebenchverified-gemma31b-opus-500] +method = "minions" +bench = "swebench-verified" +n = 500 +local = { model = "google/gemma-4-31B-it", endpoint = "http://localhost:8004/v1" } +cloud = { model = "claude-opus-4-7", endpoint = "anthropic" } +method_cfg = { mode = "minion", max_rounds = 5, worker_max_tokens = 8192 } diff --git a/src/openjarvis/agents/hybrid/registry/skillorchestra.toml b/src/openjarvis/agents/hybrid/registry/skillorchestra.toml new file mode 100644 index 00000000..11db4a37 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/skillorchestra.toml @@ -0,0 +1,47 @@ +# 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. + +[cells.skillorchestra-gaia-qwen27b-opus-3] +method = "skillorchestra" +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 = {} + +[cells.skillorchestra-swebenchverified-qwen27b-opus-3] +method = "skillorchestra" +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 = {} +concurrency = 3 + +[cells.skillorchestra-gaia-qwen27b-opus-30] +method = "skillorchestra" +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 = {} +concurrency = 4 + +[cells.skillorchestra-swebenchverified-qwen27b-opus-30] +method = "skillorchestra" +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 = {} +concurrency = 3 diff --git a/src/openjarvis/agents/hybrid/registry/swe_agent_loop.toml b/src/openjarvis/agents/hybrid/registry/swe_agent_loop.toml new file mode 100644 index 00000000..96a36637 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/swe_agent_loop.toml @@ -0,0 +1,70 @@ +# SWE-bench cells that run each paradigm through mini-SWE-agent. +# +# Same paradigms as the headline cells, same models, but with +# `swe_use_agent_loop = true` so the worker step is a full agent loop +# (shell + repo browse + test runs) instead of one-shot patch prediction. +# Expected: 0.30 → ~0.5-0.77 jump per paradigm (depending on which model +# carries the loop). +# +# These cells are EXPENSIVE — each task is ~30-50 LLM turns with bash, so +# Opus-driven cells run $1-3/task and ~5-10 min/task. Smokes at n=3 first. + +# ------------- Minions: cloud planner + local mini-SWE-agent worker ------------- + +[cells.minions-swe-agent-swebenchverified-qwen27b-opus-3] +method = "minions" +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 = { swe_use_agent_loop = true, supervisor_max_tokens = 1024, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 } + +# ------------- Conductor: shared workdir across plan steps ------------- + +[cells.conductor-swe-agent-swebenchverified-opus-3] +method = "conductor" +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 = { swe_use_agent_loop = true, conductor_max_tokens = 2048, swe_max_turns = 25, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 } + +# ------------- Advisors: two full agent-loop attempts with local critique between ------------- + +[cells.advisors-swe-agent-swebenchverified-qwen9b-opus-3] +method = "advisors" +bench = "swebench-verified" +n = 3 +local = { model = "Qwen/Qwen3.5-9B", 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 } + +# ------------- SkillOrchestra: router routes to local or cloud, that one runs the loop ------------- + +[cells.skillorchestra-swe-agent-swebenchverified-qwen27b-opus-3] +method = "skillorchestra" +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 = { swe_use_agent_loop = true, router_max_tokens = 1024, swe_max_turns = 30, swe_bash_timeout_s = 120, swe_turn_max_tokens = 4096 } + +# ------------- ToolOrchestra: orchestrator dispatches solver workers through loop on shared workdir ------------- + +[cells.toolorchestra-swe-agent-swebenchverified-qwen27b-opus-3] +method = "toolorchestra" +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 = { 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 } + +# ------------- Archon: K=3 local mini-SWE-agent runs, cloud ranker picks best ------------- + +[cells.archon-swe-agent-swebenchverified-qwen27b-opus-3] +method = "archon" +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 = { 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 } diff --git a/src/openjarvis/agents/hybrid/registry/toolorchestra.toml b/src/openjarvis/agents/hybrid/registry/toolorchestra.toml new file mode 100644 index 00000000..c6893725 --- /dev/null +++ b/src/openjarvis/agents/hybrid/registry/toolorchestra.toml @@ -0,0 +1,30 @@ +# 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. +# +# Naming: toolorchestra---- + +[cells.toolorchestra-gaia-qwen27b-opus-3] +method = "toolorchestra" +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 = { max_turns = 6, orchestrator_max_tokens = 1024, worker_max_tokens = 4096 } + +[cells.toolorchestra-swebenchverified-qwen27b-opus-3] +method = "toolorchestra" +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 = { max_turns = 6, orchestrator_max_tokens = 1024, worker_max_tokens = 8192 } + +[cells.toolorchestra-gaia-qwen27b-opus-30] +method = "toolorchestra" +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 = { max_turns = 8, orchestrator_max_tokens = 1024, worker_max_tokens = 4096 } diff --git a/src/openjarvis/agents/hybrid/runner.py b/src/openjarvis/agents/hybrid/runner.py new file mode 100644 index 00000000..ad384f03 --- /dev/null +++ b/src/openjarvis/agents/hybrid/runner.py @@ -0,0 +1,501 @@ +"""CLI runner for hybrid paradigm experiments. + +:: + + python -m openjarvis.agents.hybrid.runner --cell minions-gaia-qwen27b-opus-3 + +Reads a cell definition from ``registry/.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 +``//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 +modification. +""" + +from __future__ import annotations + +import argparse +import fcntl +import json +import os +import sys +import threading +import time +import traceback +from concurrent.futures import ThreadPoolExecutor, as_completed +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Dict, List, Optional + +try: + import tomllib # py3.11+ +except ModuleNotFoundError: + import tomli as tomllib # type: ignore[import-not-found,no-redef] + +from openjarvis.agents._stubs import AgentContext, AgentResult +from openjarvis.agents.hybrid._prompts import format_prompt as _format_prompt + +PACKAGE_DIR = Path(__file__).parent +DEFAULT_REGISTRY_DIR = PACKAGE_DIR / "registry" +DEFAULT_EXPERIMENTS_DIR = Path( + os.environ.get( + "OPENJARVIS_HYBRID_EXPERIMENTS_DIR", + str(Path.home() / ".openjarvis-hybrid" / "experiments"), + ) +) + + +# ---------- Registry ---------- + +def load_registry(registry_dir: Optional[Path] = None) -> Dict[str, Dict[str, Any]]: + """Merge every ``/*.toml``. Cell names must be unique.""" + base = registry_dir or DEFAULT_REGISTRY_DIR + env_override = os.environ.get("OPENJARVIS_HYBRID_REGISTRY_DIR") + if env_override: + base = Path(env_override) + if not base.is_dir(): + return {} + cells: Dict[str, Dict[str, Any]] = {} + for p in sorted(base.glob("*.toml")): + data = tomllib.loads(p.read_text()) + for name, cell in data.get("cells", {}).items(): + if name in cells: + raise ValueError( + f"duplicate cell {name!r} (already defined before {p.name})" + ) + cells[name] = cell + return cells + + +# ---------- Bench dispatch ---------- + +def _load_gaia_tasks(n: Optional[int]) -> List[Dict[str, Any]]: + """GAIA validation. Each task is a dict with `task_id` + `question`.""" + from openjarvis.evals.datasets.gaia import GAIADataset + + ds = GAIADataset() + ds.load(max_samples=n) + 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. + out.append({ + "task_id": rec.record_id, + "question": rec.metadata.get("question", rec.problem), + "reference": rec.reference, + "metadata": dict(rec.metadata), + }) + return out + + +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.load(max_samples=n) + out: List[Dict[str, Any]] = [] + for rec in ds.iter_records(): + md = rec.metadata or {} + out.append({ + "task_id": md.get("instance_id", rec.record_id), + "repo": md.get("repo", ""), + "base_commit": md.get("base_commit", ""), + "problem_statement": md.get("problem_statement", rec.problem), + "hints_text": md.get("hints_text", ""), + "test_patch": md.get("test_patch", ""), + "FAIL_TO_PASS": md.get("FAIL_TO_PASS", []), + "PASS_TO_PASS": md.get("PASS_TO_PASS", []), + "version": md.get("version"), + "reference": rec.reference, + "metadata": dict(md), + }) + return out + + +def load_tasks(bench: str, n: Optional[int]) -> List[Dict[str, Any]]: + if bench == "gaia": + return _load_gaia_tasks(n) + if bench in ("swebench-verified", "swebench_verified", "swebench"): + return _load_swebench_tasks(n) + raise ValueError(f"unknown bench: {bench!r}") + + +# ---------- Scoring ---------- + +def _score_gaia(task: Dict[str, Any], answer: str) -> Dict[str, Any]: + """Exact-match-with-format-normalization GAIA scorer. + + Lightweight version: extracts the final-answer line and string-compares + against the reference. Use the OpenJarvis gaia_exact scorer for the + judge-tiebreaker path. + """ + import re + + 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, + ) + pred = (m.group(1).strip() if m else answer.strip()).rstrip(".").strip() + success = pred.lower() == ref.lower() + return { + "success": success, + "score": 1.0 if success else 0.0, + "details": {"prediction": pred, "reference": ref}, + } + + +def _score_swebench(task: Dict[str, Any], answer: str) -> Dict[str, Any]: + """Modal-backed SWE-bench Verified harness scorer.""" + from openjarvis.evals.core.types import EvalRecord + from openjarvis.evals.scorers.swebench_harness import ( + SWEBenchHarnessScorer, + extract_patch, + ) + + patch = extract_patch(answer) + if patch is None: + return {"success": False, "score": 0.0, "details": {"reason": "no_patch_extracted"}} + + record = EvalRecord( + record_id=task["task_id"], + problem=task.get("problem_statement", ""), + reference="", + category="agentic", + metadata={"instance_id": task["task_id"]}, + ) + scorer = SWEBenchHarnessScorer(timeout_s=int(os.environ.get("SWEBENCH_TIMEOUT_S", "1800"))) + is_correct, details = scorer.score(record, answer) + return { + "success": bool(is_correct), + "score": 1.0 if is_correct else 0.0, + "details": details, + } + + +def score(bench: str, task: Dict[str, Any], answer: str) -> Dict[str, Any]: + if bench == "gaia": + return _score_gaia(task, answer) + if bench in ("swebench-verified", "swebench_verified", "swebench"): + return _score_swebench(task, answer) + raise ValueError(f"unknown bench: {bench!r}") + + +# ---------- Cell run ---------- + +def _cell_dir(cell_name: str, root: Path) -> Path: + d = root / cell_name + d.mkdir(parents=True, exist_ok=True) + (d / "logs").mkdir(exist_ok=True) + return d + + +@contextmanager +def _cell_lock(out_dir: Path, cell_name: str): + """Exclusive flock on ``/.lock`` to prevent concurrent runner stomps.""" + lock_path = out_dir / ".lock" + f = lock_path.open("a+") + try: + fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + f.seek(0) + prev = (f.read() or "?").strip() or "?" + f.close() + raise SystemExit( + f"[lock] another runner is already running cell {cell_name!r} " + f"(holder pid: {prev}). refusing to start a second instance." + ) + f.seek(0) + f.truncate() + f.write(str(os.getpid())) + f.flush() + try: + yield + finally: + try: + fcntl.flock(f.fileno(), fcntl.LOCK_UN) + except Exception: + pass + f.close() + try: + lock_path.unlink() + except FileNotFoundError: + pass + + +def _build_agent(cell: Dict[str, Any]): + """Construct the registered agent for this cell.""" + import openjarvis.agents # noqa: F401 — populate registry + from openjarvis.core.registry import AgentRegistry + + method = cell["method"] + if not AgentRegistry.contains(method): + raise ValueError( + f"agent {method!r} not registered. Available: " + f"{', '.join(sorted(AgentRegistry.keys()))}" + ) + agent_cls = AgentRegistry.get(method) + local = cell.get("local") or {} + cloud = cell.get("cloud") or {} + method_cfg = dict(cell.get("method_cfg") or {}) + + return agent_cls( + engine=None, # raw SDK calls — engine unused + model=cloud.get("model", ""), + local_model=local.get("model"), + local_endpoint=local.get("endpoint"), + cloud_endpoint=(cloud.get("endpoint") or "anthropic").lower(), + cfg=method_cfg, + ) + + +def _run_one(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={ + "task": task, + "task_id": task["task_id"], + "log_dir": log_dir, + }) + t0 = time.time() + try: + result: AgentResult = agent.run(prompt, ctx) + meta = dict(result.metadata or {}) + out = { + "task_id": task["task_id"], + "answer": result.content or "", + "tokens_local": int(meta.get("tokens_local", 0)), + "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)), + "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()}", + } + + +def _heartbeat(done: int, total: int, row: Dict[str, Any], t_start: float) -> None: + elapsed = time.time() - t_start + eta = (total - done) * (elapsed / max(done, 1)) + ok = "OK" if not row.get("error") else "ERR" + s = row.get("score") or {} + sc = s.get("score") + sc_str = f"{sc:.2f}" if isinstance(sc, (int, float)) else "—" + print( + f"[{done}/{total}] {ok} task={row['task_id']} score={sc_str} " + f"local={row['tokens_local']} cloud={row['tokens_cloud']} " + f"${row['cost_usd']:.3f} {row['latency_s']:.1f}s eta={eta/60:.1f}m", + flush=True, + ) + + +def _write_summary( + out_dir: Path, + cell_name: str, + cell: Dict[str, Any], + tasks: List[Dict[str, Any]], + t_start: float, +) -> None: + results_path = out_dir / "results.jsonl" + rows = [ + json.loads(line) + for line in results_path.read_text().splitlines() + if line.strip() + ] + n_done = len(rows) + n_err = sum(1 for r in rows if r.get("error")) + successes = [r for r in rows if r.get("score") and r["score"].get("success")] + acc = (len(successes) / n_done) if n_done else 0.0 + 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) + elapsed = time.time() - t_start + + summary = { + "cell": cell_name, + "method": cell["method"], + "bench": cell["bench"], + "n_target": cell["n"], + "n_done": n_done, + "n_err": n_err, + "accuracy": acc, + "tokens_local_total": total_local, + "tokens_cloud_total": total_cloud, + "cost_usd_total": total_cost, + "wall_time_s": elapsed, + "task_count": len(tasks), + } + (out_dir / "summary.json").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", + flush=True, + ) + + +def run_cell( + cell_name: str, + cell: Dict[str, Any], + *, + do_score: bool = True, + resume: bool = True, + root: Optional[Path] = None, +) -> None: + out_root = root or DEFAULT_EXPERIMENTS_DIR + out_dir = _cell_dir(cell_name, out_root) + with _cell_lock(out_dir, cell_name): + _run_cell_locked( + cell_name, cell, out_dir, + do_score=do_score, resume=resume, + ) + + +def _run_cell_locked( + cell_name: str, + cell: Dict[str, Any], + out_dir: Path, + *, + do_score: bool, + resume: bool, +) -> None: + (out_dir / "config.json").write_text( + json.dumps({"name": cell_name, **cell}, indent=2) + ) + + results_path = out_dir / "results.jsonl" + done_ids: set = set() + if resume and results_path.exists(): + # Keep only successful rows; drop errored rows so they retry. + kept: List[str] = [] + for line in results_path.read_text().splitlines(): + try: + row = json.loads(line) + except Exception: + continue + if not row.get("error"): + kept.append(line) + done_ids.add(row["task_id"]) + results_path.write_text("\n".join(kept) + ("\n" if kept else "")) + print( + f"[resume] {len(done_ids)} tasks already done (errored rows dropped)", + flush=True, + ) + + 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) + + agent = _build_agent(cell) + + t_start = time.time() + write_lock = threading.Lock() + completed = [0] + log_dir = str(out_dir / "logs") + + def _process(task: Dict[str, Any]) -> None: + row = _run_one(agent, cell["bench"], task, log_dir) + scored: Optional[Dict[str, Any]] = None + if do_score and row.get("error") is None: + try: + scored = score(cell["bench"], task, row["answer"]) + 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() + 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() + + _write_summary(out_dir, cell_name, cell, tasks, t_start) + + +# ---------- CLI ---------- + +def main(argv: Optional[List[str]] = None) -> int: + p = argparse.ArgumentParser( + prog="python -m openjarvis.agents.hybrid.runner", + description="Run a hybrid paradigm experiment cell.", + ) + p.add_argument("--cell", required=True, help="Cell name from the registry TOMLs.") + p.add_argument( + "--registry-dir", + default=None, + help="Override registry dir (defaults to package registry/).", + ) + p.add_argument( + "--root", + default=None, + help="Override experiments output root.", + ) + p.add_argument("--no-score", action="store_true", help="Skip scoring.") + p.add_argument("--no-resume", action="store_true", help="Don't resume from results.jsonl.") + args = p.parse_args(argv) + + reg_dir = Path(args.registry_dir) if args.registry_dir else None + cells = load_registry(reg_dir) + if not cells: + print(f"[error] no cells found in {reg_dir or DEFAULT_REGISTRY_DIR}", file=sys.stderr) + return 2 + if args.cell not in cells: + print( + f"[error] unknown cell {args.cell!r}. Known: {', '.join(sorted(cells))}", + file=sys.stderr, + ) + return 2 + root = Path(args.root) if args.root else None + run_cell( + args.cell, cells[args.cell], + do_score=not args.no_score, + resume=not args.no_resume, + root=root, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +__all__ = [ + "DEFAULT_EXPERIMENTS_DIR", + "DEFAULT_REGISTRY_DIR", + "load_registry", + "load_tasks", + "main", + "run_cell", + "score", +] diff --git a/src/openjarvis/agents/hybrid/scripts/new_experiment.sh b/src/openjarvis/agents/hybrid/scripts/new_experiment.sh new file mode 100755 index 00000000..4d89ccd6 --- /dev/null +++ b/src/openjarvis/agents/hybrid/scripts/new_experiment.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# Scaffold a new hybrid paradigm experiment cell. +# +# Appends a [cells.] block to +# src/openjarvis/agents/hybrid/registry/.toml +# (registry is split by method — minions.toml, conductor.toml, etc.). +# +# Usage: +# src/openjarvis/agents/hybrid/scripts/new_experiment.sh \ +# --method minions --bench gaia \ +# --local qwen3.5-27b --cloud claude-opus-4-7 --n 50 \ +# [--mode minion|minions] [--max-rounds 3] +# +# Then run: +# python -m openjarvis.agents.hybrid.runner --cell + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +REGISTRY_DIR="$(cd "$SCRIPT_DIR/.." && pwd)/registry" + +method=""; bench=""; local_model=""; cloud_model=""; n="" +mode="minion"; max_rounds=3 + +while [[ $# -gt 0 ]]; do + case "$1" in + --method) method="$2"; shift 2;; + --bench) bench="$2"; shift 2;; + --local) local_model="$2"; shift 2;; + --cloud) cloud_model="$2"; shift 2;; + --n) n="$2"; shift 2;; + --mode) mode="$2"; shift 2;; + --max-rounds) max_rounds="$2"; shift 2;; + *) echo "unknown arg: $1" >&2; exit 2;; + esac +done + +for v in method bench local_model cloud_model n; do + if [[ -z "${!v}" ]]; then + echo "missing required --${v//_/-}" >&2 + exit 2 + fi +done + +# Map shortnames → full HF id + endpoint +case "$local_model" in + qwen3.5-27b|qwen-27b|qwen27b) local_full="Qwen/Qwen3.5-27B-FP8"; local_ep="http://localhost:8001/v1";; + qwen3.5-9b|qwen-9b|qwen9b) local_full="Qwen/Qwen3.5-9B"; local_ep="http://localhost:8001/v1";; + gemma-4-31b|gemma-31b|gemma31b) local_full="google/gemma-4-31B-it"; local_ep="http://localhost:8004/v1";; + gemma-4-26b|gemma-26b|gemma26b) local_full="google/gemma-4-26B-A4B-it"; local_ep="http://localhost:8004/v1";; + *) local_full="$local_model"; local_ep="${LOCAL_ENDPOINT:-http://localhost:8001/v1}";; +esac + +case "$cloud_model" in + gpt-5|gpt-5-mini|gpt-4o) cloud_ep="openai";; + claude-opus-4-7|claude-sonnet-4-6) cloud_ep="anthropic";; + *) cloud_ep="${CLOUD_ENDPOINT:-anthropic}";; +esac + +local_short="$(echo "$local_model" | tr '/' '_')" +cloud_short="$cloud_model" +name="${method}-${bench}-${local_short}-${cloud_short}-${n}" + +registry_file="$REGISTRY_DIR/${method}.toml" +if [[ ! -f "$registry_file" ]]; then + echo "no registry file for method '${method}' — expected $registry_file" >&2 + echo "known: $(ls "$REGISTRY_DIR" | tr '\n' ' ')" >&2 + exit 2 +fi + +# Method-specific method_cfg defaults — keep simple, override later in TOML. +case "$method" in + minions) + cfg_block="method_cfg = { mode = \"${mode}\", max_rounds = ${max_rounds}, worker_max_tokens = 4096 }" + ;; + conductor) + cfg_block="method_cfg = { conductor_max_tokens = 2048, worker_max_tokens = 4096 }" + ;; + advisors) + cfg_block="method_cfg = { executor_max_tokens = 4096, advisor_max_tokens = 1024 }" + ;; + archon) + cfg_block="method_cfg = { architecture = \"ensemble_rank_fuse\", n_samples = 5, max_tokens = 2048 }" + ;; + skillorchestra) + cfg_block="method_cfg = { router_max_tokens = 1024, local_max_tokens = 4096, cloud_max_tokens = 4096 }" + ;; + toolorchestra) + cfg_block="method_cfg = { max_turns = 6, orchestrator_max_tokens = 1024, worker_max_tokens = 4096 }" + ;; + *) + cfg_block="method_cfg = {}" + ;; +esac + +cat >> "$registry_file" <` rule, comma-separated lists with no units).", + "code_or_logic": "Write or trace code, or apply logical/symbolic constraints precisely.", +} + +DEFAULT_AGENT_COMPETENCE: Dict[str, Dict[str, float]] = { + "local-qwen-27b": { + "factual_recall": 0.25, + "multi_step_reasoning": 0.30, + "arithmetic": 0.55, + "web_grounding": 0.10, + "long_text_extraction": 0.55, + "format_compliance": 0.65, + "code_or_logic": 0.45, + }, + "cloud-opus-4-7": { + "factual_recall": 0.85, + "multi_step_reasoning": 0.88, + "arithmetic": 0.85, + "web_grounding": 0.70, + "long_text_extraction": 0.90, + "format_compliance": 0.92, + "code_or_logic": 0.90, + }, +} + +DEFAULT_AGENT_COST_USD: Dict[str, float] = { + "local-qwen-27b": 0.0, + "cloud-opus-4-7": 0.30, +} + +ROUTER_SYS_MARKER = "<>" + + +def _format_catalog() -> str: + return "\n".join(f"- {sid}: {desc}" for sid, desc in SKILL_CATALOG.items()) + + +def _format_agents( + competence: Dict[str, Dict[str, float]], + cost: Dict[str, float], +) -> str: + lines: List[str] = [] + for aid in competence: + comp = competence[aid] + comp_str = ", ".join(f"{k}={v:.2f}" for k, v in comp.items()) + lines.append(f"- **{aid}** (avg cost ${cost.get(aid, 0.0):.2f}/task)") + lines.append(f" Skill competences: {comp_str}") + return "\n".join(lines) + + +def _build_router_sys( + competence: Dict[str, Dict[str, float]], + cost: Dict[str, float], +) -> str: + return f"""{ROUTER_SYS_MARKER} +You are a skill-aware model router for a compound AI system (the SkillOrchestra deployment-time orchestrator). For each user question you must: + +1. Assign weights over the skill catalog (numbers in [0, 1], summing to ~1.0). \ +The weights reflect how much each skill matters for *this* question. +2. Score each candidate agent: score(agent) = sum_skill weight_skill * competence(agent, skill) - lambda_cost * avg_cost(agent), with lambda_cost = 0.5. +3. Pick the highest-scoring agent. Tie-break in favor of the cheaper agent. + +Skill catalog: +{_format_catalog()} + +Agent pool (with learned-prior competences and average costs): +{_format_agents(competence, cost)} + +Respond with ONLY a JSON object: {{"chosen_agent": ..., "skill_weights": {{...}}, "reasoning": "..."}}. No prose outside JSON. +""" + + +def _build_router_schema(agent_ids: List[str]) -> Dict[str, Any]: + return { + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "chosen_agent": {"type": "string", "enum": agent_ids}, + "skill_weights": { + "type": "object", + "properties": { + sid: {"type": "number"} for sid in SKILL_CATALOG + }, + "required": list(SKILL_CATALOG.keys()), + "additionalProperties": False, + }, + "reasoning": {"type": "string"}, + }, + "required": ["chosen_agent", "skill_weights", "reasoning"], + "additionalProperties": False, + }, + } + } + + +def _parse_router_json(text: str) -> Dict[str, Any]: + s = (text or "").strip() + try: + return json.loads(s) + except json.JSONDecodeError: + pass + start = s.find("{") + if start == -1: + raise ValueError(f"router emitted no JSON: {s[:200]!r}") + depth = 0 + for i in range(start, len(s)): + c = s[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + return json.loads(s[start : i + 1]) + raise ValueError(f"router JSON not balanced: {s[:200]!r}") + + +def _score_agents( + skill_weights: Dict[str, float], + competence: Dict[str, Dict[str, float]], + cost: Dict[str, float], +) -> Dict[str, Dict[str, float]]: + lam = 0.5 + scores: Dict[str, Dict[str, float]] = {} + for aid, comps in competence.items(): + comp = sum( + skill_weights.get(sid, 0.0) * comps[sid] for sid in SKILL_CATALOG + ) + cost_pen = lam * cost.get(aid, 0.0) + scores[aid] = { + "competence": comp, + "cost_penalty": cost_pen, + "final_score": comp - cost_pen, + } + return scores + + +@AgentRegistry.register("skillorchestra") +class SkillOrchestraAgent(LocalCloudAgent): + """Inference-time skill-aware router. See module docstring.""" + + agent_id = "skillorchestra" + + def _is_soft_failure(self, exc: BaseException) -> Optional[str]: + # Empty/unbalanced router JSON — treat as soft failure to match the + # hybrid adapter's behavior (matches `err=1` rows in the n=30 cell). + if isinstance(exc, (ValueError, json.JSONDecodeError)): + return f"{type(exc).__name__}: {str(exc)[:120]}" + return None + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + cfg = self._cfg + question = input + + competence: Dict[str, Dict[str, float]] = cfg.get( + "agent_competence", DEFAULT_AGENT_COMPETENCE + ) + cost: Dict[str, float] = cfg.get("agent_cost_usd", DEFAULT_AGENT_COST_USD) + agent_ids = list(competence.keys()) + router_sys = _build_router_sys(competence, cost) + router_schema = _build_router_schema(agent_ids) + + # 1. Route — Anthropic only (output_config schema is Anthropic-specific + # in the hybrid adapter). If you need OpenAI routing, swap the prompt + # to JSON-mode and bypass output_config. + if self._cloud_endpoint != "anthropic": + raise ValueError( + "SkillOrchestra router requires cloud_endpoint='anthropic'; " + f"got {self._cloud_endpoint!r}" + ) + router_max = int(cfg.get("router_max_tokens", 1024)) + # Strip temperature for Opus 4.7+; Anthropic's output_config does the schema. + if supports_temperature(self._cloud_model): + router_text, r_in, r_out, _ = self._call_anthropic( + self._cloud_model, + user=f"Question:\n{question}", + system=router_sys, + max_tokens=router_max, + temperature=0.0, + output_config=router_schema, + ) + else: + router_text, r_in, r_out, _ = self._call_anthropic( + self._cloud_model, + user=f"Question:\n{question}", + system=router_sys, + max_tokens=router_max, + output_config=router_schema, + ) + + decision = _parse_router_json(router_text) + skill_weights: Dict[str, float] = decision.get("skill_weights") or {} + for sid in SKILL_CATALOG: + skill_weights.setdefault(sid, 0.0) + chosen = decision.get("chosen_agent") or "" + + scored = _score_agents(skill_weights, competence, cost) + if chosen not in competence: + chosen = max(scored, key=lambda a: scored[a]["final_score"]) + + self.record_trace_event({ + "kind": "skillorchestra_route", + "chosen_agent": chosen, + "skill_weights": skill_weights, + "agent_scores": scored, + "reasoning": decision.get("reasoning", ""), + "router_raw": router_text, + }) + + tokens_local = 0 + tokens_cloud = r_in + r_out + run_cost = self.cost_usd(self._cloud_model, r_in, r_out) + + # 2. Execute via chosen agent + 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")) + ) + if chosen == "local-qwen-27b": + if not (self._local_model and self._local_endpoint): + raise ValueError( + "SkillOrchestra route hit local agent but local_model/" + f"local_endpoint missing: {self._local_model!r}/{self._local_endpoint!r}" + ) + if swe_mode: + out = run_swe_agent_loop( + task_meta, + backbone="local", + backbone_model=self._local_model, + local_endpoint=self._local_endpoint, + initial_prompt=question, + 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_local", + ) + ans = out["answer"] + tokens_local += out["tokens_in"] + out["tokens_out"] + else: + ans, w_in, w_out = self._call_vllm( + self._local_model, + self._local_endpoint, + user=question, + max_tokens=int(cfg.get("local_max_tokens", 4096)), + temperature=float(cfg.get("local_temperature", 0.2)), + enable_thinking=False, + ) + tokens_local += w_in + w_out + worker_model = self._local_model + else: + if swe_mode: + out = run_swe_agent_loop( + task_meta, + backbone="cloud", + backbone_model=self._cloud_model, + cloud_endpoint=self._cloud_endpoint, + initial_prompt=question, + 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_cloud", + ) + ans = out["answer"] + tokens_cloud += out["tokens_in"] + out["tokens_out"] + run_cost += out["cost_usd"] + else: + ans, w_in, w_out, _ = self._call_anthropic( + self._cloud_model, + user=question, + max_tokens=int(cfg.get("cloud_max_tokens", 4096)), + temperature=0.0, + ) + tokens_cloud += w_in + w_out + run_cost += self.cost_usd(self._cloud_model, w_in, w_out) + worker_model = self._cloud_model + + meta = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": run_cost, + "turns": 2, # router + worker + "traces": { + "chosen_agent": chosen, + "worker_model": worker_model, + "skill_weights": skill_weights, + "agent_scores": scored, + "reasoning": decision.get("reasoning", ""), + }, + } + return ans, meta + + +__all__ = [ + "DEFAULT_AGENT_COMPETENCE", + "DEFAULT_AGENT_COST_USD", + "SKILL_CATALOG", + "SkillOrchestraAgent", +] diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py new file mode 100644 index 00000000..b8195e61 --- /dev/null +++ b/src/openjarvis/agents/hybrid/toolorchestra.py @@ -0,0 +1,524 @@ +"""ToolOrchestraAgent — prompted port of NVlabs ToolOrchestra (arXiv:2511.21689). + +The paper RL-trains an 8B Orchestrator (``nvidia/Nemotron-Orchestrator-8B``) +to coordinate basic tools + specialist LLMs + generalist LLMs in +multi-turn agentic loops, ranked #1 on GAIA at release. + +The hybrid harness adapter for ToolOrchestra is a documented stub — +running the real thing needs a separate vLLM srun for the Orchestrator-8B +checkpoint, a FAISS wiki retriever, a Tavily API key, and a refactor of +the upstream eval scripts. None of that fits in our cluster allocation. + +This port keeps the same scope discipline: **inference-time only, +prompted, no RL**. A cloud model plays the role of the orchestrator, +dispatching to a pool of `(tool | specialist_llm | generalist_llm)` +workers in a reactive loop. The loop is the paradigm; the orchestrator +weights are not. + +Why ship this at all if it's not the "real" thing? Because the prompted +upper-bound is useful as a reference point alongside the other paradigms, +and because the OpenJarvis registry needs all six entries for the +distillation pipeline to slot ToolOrchestra in alongside the rest. + +Pipeline per task: + +1. Orchestrator (cloud) reads question + numbered worker pool. +2. Each turn it emits ``{"action": "call_worker", "worker_id": int, + "input": str}`` or ``{"action": "final_answer", "answer": str}``. +3. Up to ``max_turns`` (default 6) calls before forcing a final-answer + prompt; fallback to strongest worker on parse failure. + +Workers come from ``cfg["workers"]`` or a sensible default pool (local +Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7, +gpt-5-mini). + +Not yet validated end-to-end in the hybrid harness — the hybrid adapter +``raise NotImplementedError``s. Treat results from this paradigm as +preliminary until we have a real ToolOrchestra-8B deployment. +""" + +from __future__ import annotations + +import json +import re +import shutil +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from openjarvis.agents._stubs import AgentContext +from openjarvis.agents.hybrid._base import ( + ANTHROPIC_WEB_SEARCH_TOOL, + WEB_SEARCH_COST_PER_CALL, + LocalCloudAgent, +) +from openjarvis.agents.hybrid._prices import ( + is_gpt5_family, + supports_temperature, +) +from openjarvis.agents.hybrid.mini_swe_agent import ( + _clone_repo, + _extract_diff, + run_swe_agent_loop, +) +from openjarvis.core.registry import AgentRegistry + +ORCHESTRATOR_SYS = """\ +You are a tool-orchestrating agent. You coordinate a pool of workers to answer the user's question. Each turn you MUST emit exactly one JSON object — no prose, no markdown fences — taking one of two forms: + + {"action": "call_worker", "worker_id": , "input": ""} + + {"action": "final_answer", "answer": ""} + +Strategy: + +- Call cheap / specialized workers first (small local model for extraction or arithmetic on given data; web_search for unknowns; specialist LLMs for code/math). +- Call the frontier worker (Opus / GPT-5) sparingly, for hard reasoning or a final synthesis pass. +- Stop and emit `final_answer` as soon as the previous worker output is sufficient. Do NOT call a worker just to paraphrase. +- The user only sees the `answer` field of `final_answer`, so make sure it follows any answer-format rules in the question. +""" + +FORCE_FINAL_PROMPT = ( + "Worker-call budget exhausted. Emit `final_answer` now using everything " + "you've learned. Respect the question's answer-format rules." +) + + +def _build_pool_block(workers: List[Dict[str, Any]]) -> str: + return "\n".join( + f"Worker {w['id']} ({w['name']}): {w['description']}" for w in workers + ) + + +def _build_user_prompt( + question: str, + workers: List[Dict[str, Any]], + history: List[Dict[str, Any]], +) -> str: + pieces = [ + f"Worker pool:\n{_build_pool_block(workers)}", + f"User question:\n{question}", + ] + if history: + pieces.append("Conversation so far (orchestrator turns and worker outputs):") + for h in history: + if h["role"] == "orchestrator": + pieces.append(f"[Orchestrator turn {h['turn']}]\n{h['raw']}") + else: + pieces.append( + f"[Worker {h['worker_id']} ({h['worker_name']}) turn {h['turn']}]\n" + f"{h['output']}" + ) + pieces.append( + "Emit the next JSON action object now — exactly one object, no prose." + ) + return "\n\n".join(pieces) + + +def _strip_fences(s: str) -> str: + s = s.strip() + if s.startswith("```"): + first_nl = s.find("\n") + if first_nl != -1: + s = s[first_nl + 1:] + if s.endswith("```"): + s = s[:-3] + s = s.strip() + return s + + +def _parse_action(text: str) -> Optional[Dict[str, Any]]: + s = _strip_fences(text) + # First try direct parse, then balanced-brace extraction. + try: + obj = json.loads(s) + if isinstance(obj, dict) and "action" in obj: + return obj + except json.JSONDecodeError: + pass + start = s.find("{") + if start == -1: + return None + depth = 0 + for i in range(start, len(s)): + c = s[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + try: + obj = json.loads(s[start : i + 1]) + if isinstance(obj, dict) and "action" in obj: + return obj + except json.JSONDecodeError: + return None + return None + + +def _extract_final_answer_text(text: str) -> str: + """Best-effort: pull the answer string from a malformed action emission. + + Tries `"answer": "..."` regex, then the GAIA-style `FINAL ANSWER:` line. + """ + m = re.search(r'"answer"\s*:\s*"((?:\\.|[^"\\])*)"', text, re.DOTALL) + if m: + return m.group(1).encode("utf-8").decode("unicode_escape") + m = re.search(r"FINAL\s*ANSWER\s*:\s*(.+?)\s*$", text, re.IGNORECASE | re.MULTILINE) + if m: + return m.group(1).strip() + return text.strip() + + +# ---------- Worker pool ---------- + +def _default_pool(local_model: Optional[str], local_endpoint: Optional[str]) -> List[Dict[str, Any]]: + pool: List[Dict[str, Any]] = [] + if local_model and local_endpoint: + pool.append({ + "id": len(pool), + "name": "local-qwen", + "type": "vllm", + "model": local_model, + "base_url": local_endpoint, + "description": ( + "Open-weights Qwen3.5 served locally. Cheap and fast. Good at " + "concise extraction, formatting, arithmetic on given data." + ), + }) + pool.append({ + "id": len(pool), + "name": "web-search", + "type": "anthropic-web-search", + "model": "claude-haiku-4-5", + "description": ( + "Anthropic server-side web_search. Use for facts that need a lookup " + "(recent events, rare names/dates, niche sources). Returns a digest." + ), + }) + pool.append({ + "id": len(pool), + "name": "frontier-anthropic", + "type": "anthropic", + "model": "claude-opus-4-7", + "description": ( + "Frontier reasoning model. Use for hard multi-step reasoning, " + "code review, or a final synthesis pass. Expensive — use sparingly." + ), + }) + pool.append({ + "id": len(pool), + "name": "frontier-openai-mini", + "type": "openai", + "model": "gpt-5-mini", + "description": ( + "Mid-tier OpenAI model. Solid general knowledge and reasoning at a " + "fraction of frontier cost." + ), + }) + return pool + + +def _call_worker( + worker: Dict[str, Any], prompt: str, cfg: Dict[str, Any] +) -> Tuple[str, int, int, bool, float, int]: + """Returns (text, p_tok, c_tok, is_local, extra_cost, n_web_searches).""" + wtype = worker.get("type", "openai") + max_tok = int(cfg.get("worker_max_tokens", 4096)) + temp = float(cfg.get("worker_temperature", 0.2)) + + if wtype == "vllm": + text, p, c = LocalCloudAgent._call_vllm( + worker["model"], + worker["base_url"], + user=prompt, + max_tokens=max_tok, + temperature=temp, + enable_thinking=False, + ) + return text, p, c, True, 0.0, 0 + if wtype == "openai": + eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp + text, p, c = LocalCloudAgent._call_openai( + worker["model"], + user=prompt, + max_tokens=max_tok, + temperature=eff_temp, + ) + return text, p, c, False, 0.0, 0 + if wtype == "anthropic": + eff_temp = temp if supports_temperature(worker["model"]) else 0.0 + text, p, c, _ = LocalCloudAgent._call_anthropic( + worker["model"], + user=prompt, + max_tokens=max_tok, + temperature=eff_temp, + ) + return text, p, c, False, 0.0, 0 + if wtype == "anthropic-web-search": + eff_temp = temp if supports_temperature(worker["model"]) else 0.0 + text, p, c, n_searches = LocalCloudAgent._call_anthropic( + worker["model"], + user=prompt, + max_tokens=max_tok, + temperature=eff_temp, + tools=[ANTHROPIC_WEB_SEARCH_TOOL], + tool_choice={"type": "any"}, + ) + extra = n_searches * WEB_SEARCH_COST_PER_CALL + return text, p, c, False, extra, n_searches + raise ValueError(f"unsupported worker type: {wtype!r}") + + +def _swe_call_worker( + worker: Dict[str, Any], + prompt: str, + cfg: Dict[str, Any], + task: Dict[str, Any], + workdir: Path, + turn: int, +) -> Tuple[str, int, int, bool, float, int]: + """SWE-bench worker dispatch: route solver workers through + run_swe_agent_loop on a shared workdir. Web-search workers fall back + to the regular one-shot dispatch (search isn't an agent loop).""" + wtype = worker.get("type", "openai") + if wtype == "anthropic-web-search": + # Search workers stay one-shot. + return _call_worker(worker, prompt, cfg) + if wtype == "vllm": + backbone = "local" + endpoint = worker.get("base_url") + elif wtype == "anthropic": + backbone = "cloud" + endpoint = None + else: + # OpenAI workers fall back to one-shot. + return _call_worker(worker, prompt, cfg) + out = run_swe_agent_loop( + task, + backbone=backbone, + backbone_model=worker["model"], + cloud_endpoint="anthropic", + local_endpoint=endpoint, + initial_prompt=prompt, + max_turns=int(cfg.get("swe_max_turns", 30)), + bash_timeout=int(cfg.get("swe_bash_timeout_s", 120)), + output_cap=int(cfg.get("swe_output_cap", 10_000)), + turn_max_tokens=int(cfg.get("swe_turn_max_tokens", 4096)), + trace_prefix=f"toolorch_turn{turn}", + workdir=workdir, + ) + is_local = backbone == "local" + return ( + out["final_summary"] or out["answer"], + out["tokens_in"], out["tokens_out"], + is_local, 0.0, 0, + ) + + +@AgentRegistry.register("toolorchestra") +class ToolOrchestraAgent(LocalCloudAgent): + """Prompted multi-turn dispatcher over a mixed worker pool. + + Inference-only port — does NOT use the RL-trained Nemotron-Orchestrator-8B. + See module docstring for what's missing relative to the published paper. + """ + + agent_id = "toolorchestra" + + def _run_paradigm( + self, + input: str, + context: Optional[AgentContext], + **kwargs: Any, + ) -> Tuple[str, Dict[str, Any]]: + cfg = self._cfg + question = input + workers = cfg.get("workers") or _default_pool( + self._local_model, self._local_endpoint + ) + if not workers: + raise RuntimeError("toolorchestra: empty worker pool") + + max_turns = int(cfg.get("max_turns", 6)) + orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 1024)) + + task_meta = (context.metadata.get("task") if context is not None else {}) or {} + swe_mode = ( + bool(cfg.get("swe_use_agent_loop")) + and bool(task_meta.get("problem_statement")) + and bool(task_meta.get("repo")) + and bool(task_meta.get("base_commit")) + ) + shared_workdir: Optional[Path] = None + if swe_mode: + shared_workdir = Path(tempfile.mkdtemp( + prefix=f"toolorch-swe-{task_meta.get('task_id','x')}-" + )) + try: + _clone_repo(task_meta["repo"], task_meta["base_commit"], shared_workdir) + except Exception: + shutil.rmtree(shared_workdir, ignore_errors=True) + raise + self.record_trace_event({ + "kind": "toolorchestra_swe_workdir", + "workdir": str(shared_workdir), + "repo": task_meta["repo"], + "base_commit": task_meta["base_commit"], + }) + + # try/finally guards ``shared_workdir`` against exceptions raised + # anywhere in the turn loop, the worker calls, the fallback, or + # the diff-extraction step. Without this, at n=500 SWE-bench an + # exception leaves hundreds of MB of cloned repos in tempdir. + try: + history: List[Dict[str, Any]] = [] + tokens_local = 0 + tokens_cloud = 0 + cost = 0.0 + final_answer: Optional[str] = None + forced_final = False + parse_failures = 0 + + for turn in range(1, max_turns + 1): + sys_prompt = ORCHESTRATOR_SYS + if turn == max_turns and final_answer is None: + sys_prompt = ORCHESTRATOR_SYS + "\n\n" + FORCE_FINAL_PROMPT + forced_final = True + + user = _build_user_prompt(question, workers, history) + text, o_in, o_out = self._call_cloud( + user=user, + system=sys_prompt, + max_tokens=orch_max_tokens, + temperature=0.0, + ) + tokens_cloud += o_in + o_out + cost += self.cost_usd(self._cloud_model, o_in, o_out) + + action = _parse_action(text) + history.append({ + "role": "orchestrator", "turn": turn, "raw": text, "action": action, + }) + self.record_trace_event({ + "kind": "toolorchestra_action", + "turn": turn, + "action": action, + "raw": text, + }) + + if action is None: + parse_failures += 1 + if parse_failures >= 2 or forced_final: + final_answer = _extract_final_answer_text(text) + break + continue + + kind = action.get("action") + if kind == "final_answer": + final_answer = str(action.get("answer", "")).strip() + break + if kind == "call_worker": + wid = action.get("worker_id") + w_input = action.get("input", "") + if not isinstance(wid, int) or not (0 <= wid < len(workers)): + parse_failures += 1 + if parse_failures >= 2 or forced_final: + final_answer = _extract_final_answer_text(text) + break + continue + worker = workers[wid] + if swe_mode and shared_workdir is not None: + w_text, w_in, w_out, is_local, extra_cost, n_searches = ( + _swe_call_worker( + worker, str(w_input), cfg, task_meta, + shared_workdir, turn, + ) + ) + else: + w_text, w_in, w_out, is_local, extra_cost, n_searches = ( + _call_worker(worker, str(w_input), cfg) + ) + if is_local: + tokens_local += w_in + w_out + else: + tokens_cloud += w_in + w_out + cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost + history.append({ + "role": "worker", + "turn": turn, + "worker_id": wid, + "worker_name": worker["name"], + "worker_model": worker["model"], + "output": w_text, + "tokens_in": w_in, + "tokens_out": w_out, + "n_web_searches": n_searches, + }) + continue + # Unknown action kind — treat as parse failure. + parse_failures += 1 + + if final_answer is None: + # Hard fallback: call the strongest worker (last) directly. + worker = workers[-1] + if swe_mode and shared_workdir is not None: + ans, w_in, w_out, is_local, extra_cost, _ = _swe_call_worker( + worker, question, cfg, task_meta, + shared_workdir, max_turns + 1, + ) + else: + ans, w_in, w_out, is_local, extra_cost, _ = _call_worker( + worker, question, cfg + ) + if is_local: + tokens_local += w_in + w_out + else: + tokens_cloud += w_in + w_out + cost += self.cost_usd(worker["model"], w_in, w_out) + extra_cost + history.append({ + "role": "worker", + "turn": max_turns + 1, + "worker_id": worker["id"], + "worker_name": worker["name"], + "worker_model": worker["model"], + "output": ans, + "tokens_in": w_in, + "tokens_out": w_out, + "fallback": True, + }) + final_answer = ans + + # In SWE mode, the authoritative output is the working-tree diff — + # frame it (the runner extracts it via the scorer's ```diff fence). + if swe_mode and shared_workdir is not None: + patch = _extract_diff(shared_workdir) + if patch.strip(): + final_answer = ( + f"{final_answer}\n\n```diff\n{patch}```" + if final_answer else f"```diff\n{patch}```" + ) + + meta = { + "tokens_local": tokens_local, + "tokens_cloud": tokens_cloud, + "cost_usd": cost, + "turns": len([h for h in history if h["role"] == "orchestrator"]), + "traces": { + "history": history, + "forced_final": forced_final, + "parse_failures": parse_failures, + "workers": workers, + "note": ( + "inference-only port; the RL-trained Nemotron-Orchestrator-8B " + "is NOT in the loop. Results are preliminary." + ), + }, + } + return final_answer, meta + finally: + if shared_workdir is not None: + shutil.rmtree(shared_workdir, ignore_errors=True) + + +__all__ = ["ToolOrchestraAgent"] diff --git a/src/openjarvis/engine/_stubs.py b/src/openjarvis/engine/_stubs.py index 4a940274..c0629502 100644 --- a/src/openjarvis/engine/_stubs.py +++ b/src/openjarvis/engine/_stubs.py @@ -21,12 +21,18 @@ class StreamChunk: Used by ``stream_full()`` to yield rich streaming data including tool_calls fragments and finish_reason, unlike ``stream()`` which only yields plain content strings. + + ``content_blocks`` and ``tool_results`` are aggregate fields emitted + once at end-of-stream so streaming callers reach parity with the + non-streaming ``generate()`` return shape. """ content: Optional[str] = None tool_calls: Optional[List[Dict[str, Any]]] = None finish_reason: Optional[str] = None usage: Optional[Dict[str, Any]] = None + content_blocks: Optional[List[Dict[str, Any]]] = None + tool_results: Optional[List[Dict[str, Any]]] = None @dataclass(slots=True) diff --git a/src/openjarvis/engine/cloud.py b/src/openjarvis/engine/cloud.py index 7a6340c0..8a32fc84 100644 --- a/src/openjarvis/engine/cloud.py +++ b/src/openjarvis/engine/cloud.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import os import time from collections.abc import AsyncIterator, Sequence @@ -19,6 +20,8 @@ from openjarvis.engine._base import ( ) from openjarvis.engine._stubs import StreamChunk +logger = logging.getLogger(__name__) + # Pricing per million tokens (input, output) PRICING: Dict[str, tuple[float, float]] = { "gpt-4o": (2.50, 10.00), @@ -149,6 +152,37 @@ def estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> flo return input_cost + output_cost +def _serialize_anthropic_block(block: Any) -> Dict[str, Any]: + """Turn an Anthropic content block into a JSON-safe dict for tracing. + + Handles every block type Anthropic returns today (text, tool_use, + server_tool_use, web_search_tool_result, tool_result, thinking). Nested + content (e.g. ``web_search_tool_result.content`` is itself a list of + citation/result blocks) recurses so the trace has the full payload, not + a truncated summary. + """ + out: Dict[str, Any] = { + "type": getattr(block, "type", None) or type(block).__name__, + } + for attr in ( + "id", "name", "input", "text", "thinking", "signature", + "tool_use_id", "content", + ): + if not hasattr(block, attr): + continue + val = getattr(block, attr) + if attr == "content" and isinstance(val, list): + out[attr] = [_serialize_anthropic_block(b) for b in val] + elif hasattr(val, "model_dump"): + try: + out[attr] = val.model_dump() + except Exception: + out[attr] = str(val) + else: + out[attr] = val + return out + + def _annotate_anthropic_cache(messages: list[dict]) -> list[dict]: """Add cache_control to system message for Anthropic prompt caching.""" result = [] @@ -573,20 +607,48 @@ class CloudEngine(InferenceEngine): resp = self._anthropic_client.messages.create(**create_kwargs) elapsed = time.monotonic() - t0 - # Extract text and tool_use blocks from response content + # Walk every block in resp.content. Anthropic returns several kinds: + # - text (plain assistant text) + # - tool_use (model wants the caller to run a tool) + # - server_tool_use (model invoked a server-side tool, + # e.g. web_search; carries the actual + # query the model issued) + # - web_search_tool_result (server-tool result body) + # - tool_result (caller-side tool result echo) + # - thinking (Opus reasoning trace) + # ``content_blocks`` keeps the full serialized list for trace + # observability. ``tool_calls`` is the narrow caller-executable + # surface — only ``tool_use`` blocks (server_tool_use lives in + # content_blocks since Anthropic already ran it server-side). + # ``tool_results`` aggregates both result kinds. content_parts: list[str] = [] tool_calls: list[Dict[str, Any]] = [] + tool_results: list[Dict[str, Any]] = [] + content_blocks: list[Dict[str, Any]] = [] for block in resp.content: - if getattr(block, "type", None) == "tool_use": + btype = getattr(block, "type", None) or type(block).__name__ + serialized = _serialize_anthropic_block(block) + content_blocks.append(serialized) + if btype == "tool_use": + block_id = getattr(block, "id", None) + if not block_id: + logger.warning( + "Anthropic tool_use block without an id; skipping. " + "Round-trip into the next assistant turn would fail " + "Anthropic's tool_use_id matching." + ) + continue tool_calls.append( { - "id": block.id, - "name": block.name, - "arguments": json.dumps(block.input) - if isinstance(block.input, dict) - else str(block.input), + "id": block_id, + "name": getattr(block, "name", ""), + "arguments": json.dumps(getattr(block, "input", None)) + if isinstance(getattr(block, "input", None), dict) + else str(getattr(block, "input", "")), } ) + elif btype in ("web_search_tool_result", "tool_result"): + tool_results.append(serialized) elif hasattr(block, "text"): content_parts.append(block.text) @@ -605,10 +667,13 @@ class CloudEngine(InferenceEngine): "finish_reason": resp.stop_reason or "stop", "cost_usd": estimate_cost(model, prompt_tokens, completion_tokens), "ttft": elapsed, + "content_blocks": content_blocks, } if tool_calls: result["tool_calls"] = tool_calls + if tool_results: + result["tool_results"] = tool_results return result @@ -1281,6 +1346,31 @@ class CloudEngine(InferenceEngine): finish = "tool_calls" if stop_reason == "tool_use" else "stop" yield StreamChunk(finish_reason=finish) + # End-of-stream parity with ``_generate_anthropic``: emit + # ``content_blocks`` (every block kind, including server_tool_use + # and thinking) and ``tool_results`` (web_search_tool_result + + # tool_result) so streaming traces see what non-streaming traces + # see. + try: + final_msg = stream.get_final_message() + except Exception as exc: # noqa: BLE001 — SDK shape varies + logger.debug("Anthropic stream.get_final_message() failed: %s", exc) + final_msg = None + if final_msg is not None and getattr(final_msg, "content", None): + content_blocks: list[Dict[str, Any]] = [] + tool_results: list[Dict[str, Any]] = [] + for block in final_msg.content: + btype = getattr(block, "type", None) or type(block).__name__ + serialized = _serialize_anthropic_block(block) + content_blocks.append(serialized) + if btype in ("web_search_tool_result", "tool_result"): + tool_results.append(serialized) + if content_blocks or tool_results: + yield StreamChunk( + content_blocks=content_blocks or None, + tool_results=tool_results or None, + ) + async def stream_full( self, messages: Sequence[Message], diff --git a/src/openjarvis/evals/datasets/gaia.py b/src/openjarvis/evals/datasets/gaia.py index 6b519d18..18a89b3b 100644 --- a/src/openjarvis/evals/datasets/gaia.py +++ b/src/openjarvis/evals/datasets/gaia.py @@ -159,6 +159,7 @@ class GAIADataset(DatasetProvider): metadata = { "task_id": task_id, + "question": question, "level": level, "file_name": file_name, "file_path": str(file_path) if file_path else None, diff --git a/src/openjarvis/evals/datasets/swebench.py b/src/openjarvis/evals/datasets/swebench.py index 1117649f..d227dab6 100644 --- a/src/openjarvis/evals/datasets/swebench.py +++ b/src/openjarvis/evals/datasets/swebench.py @@ -142,6 +142,7 @@ class SWEBenchDataset(DatasetProvider): metadata: dict[str, Any] = { "instance_id": instance_id, + "problem_statement": problem_statement, "repo": repo, "base_commit": base_commit, "hints_text": hints_text, diff --git a/src/openjarvis/evals/scorers/swebench_harness.py b/src/openjarvis/evals/scorers/swebench_harness.py new file mode 100644 index 00000000..44d27f57 --- /dev/null +++ b/src/openjarvis/evals/scorers/swebench_harness.py @@ -0,0 +1,295 @@ +"""SWE-bench harness scorer — runs the official `swebench` test harness. + +This is the authoritative pass/fail scorer for SWE-bench-Verified. +The lightweight :class:`SWEBenchScorer` in ``swebench_structural.py`` +checks only that the model produced *something patch-shaped*; this one +actually applies the patch, runs the targeted tests, and reads the +harness's report JSON. + +Backends (selected by ``SWEBENCH_BACKEND`` env var): + +- ``modal`` (default) — runs on Modal in the cloud; needs ``swebench[modal]`` + installed and ``modal token new`` configured once. +- ``docker`` — runs locally; needs Docker daemon + user in ``docker`` group. + +Ported from ``hybrid-local-cloud-compute/benches/swebench_verified/{runner,parsing}.py``, +with two upstream-`swebench` patches applied at import time: + +1. **Modal cgroup-v2 fix**: + ``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. + +2. **Rescore `*_ids` fix**: older harness rescore code read + ``resolved_instances`` / ``unresolved_instances`` / ``error_instances`` + as lists. Current swebench writes counts there and puts IDs in + ``*_ids`` fields. Wherever we read these we use ``*_ids``. + +Both patches are idempotent and only fire when the harness modules are +imported via this scorer (we don't touch swebench until ``score()`` is +called for the first time). +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +from openjarvis.evals.core.scorer import Scorer +from openjarvis.evals.core.types import EvalRecord + +logger = logging.getLogger(__name__) + + +# ---------- Patch tracking ---------- + +_PATCHES_APPLIED = False + + +def _patch_modal_cgroup_v2() -> None: + """Wrap the cgroup-v1 write in run_evaluation_modal.py with try/except. + + The upstream line is ``Path("/sys/fs/cgroup/cpu/cpu.shares").write_text(...)``. + Modal v2 sandboxes are cgroup-v2 and that path doesn't exist; the write + raises FileNotFoundError and the sandbox dies. We replace the entire + ``set_cpu_quota`` function body (if present) with a try/except wrapper. + """ + try: + from swebench.harness.modal_eval import ( + run_evaluation_modal as _m, # type: ignore[import-not-found] + ) + except Exception: + return + if getattr(_m, "_hybrid_cgroup_patched", False): + return + orig = getattr(_m, "set_cpu_quota", None) + if orig is None: + # Upstream API changed: ``set_cpu_quota`` no longer exists. The + # cgroup-v1 write that used to live inside it is either gone (good) + # or now lives somewhere we can't patch (bad — every Modal-v2 + # sandbox will still die). Surface loudly so failures aren't + # silently scored as 0 via the ``no_report`` fallback in + # :func:`_run_harness`. + logger.warning( + "swebench.harness.modal_eval.run_evaluation_modal.set_cpu_quota " + "is missing — cgroup-v2 patch could not be applied. If your " + "Modal sandboxes are scoring 0 with `reason: no_report`, " + "verify upstream swebench's Modal cgroup handling." + ) + _m._hybrid_cgroup_patched = True # type: ignore[attr-defined] + return + + def patched(*args, **kwargs): # type: ignore[no-untyped-def] + try: + return orig(*args, **kwargs) + except FileNotFoundError: + # cgroup v2 sandbox — path missing is expected, skip. + return None + except PermissionError: + return None + + _m.set_cpu_quota = patched # type: ignore[assignment] + _m._hybrid_cgroup_patched = True # type: ignore[attr-defined] + + +def _apply_patches_once() -> None: + global _PATCHES_APPLIED + if _PATCHES_APPLIED: + return + _patch_modal_cgroup_v2() + _PATCHES_APPLIED = True + + +# ---------- Patch extraction ---------- + +_FENCE_PATTERNS = [ + re.compile(r"```(?:diff|patch)\n(.*?)```", re.DOTALL), + re.compile(r"```\n(diff --git .*?)```", re.DOTALL), +] + + +def extract_patch(text: str) -> Optional[str]: + """Pull a unified diff out of agent output. ``None`` if not found.""" + if not text: + return None + for pat in _FENCE_PATTERNS: + m = pat.search(text) + if m: + return m.group(1).strip() + "\n" + if "diff --git" in text: + start = text.index("diff --git") + return text[start:].strip() + "\n" + return None + + +# ---------- Harness invocation ---------- + +def _harness_cache_dir() -> Path: + """Where the swebench subprocess writes its report JSON + logs/ tree. + + Defaults to ``$OPENJARVIS_HOME/.swebench-cache`` if set, otherwise to a + process-shared tempdir. Pin both so we don't pollute the project root. + """ + home = os.environ.get("OPENJARVIS_HOME") + if home: + cache = Path(home) / ".swebench-cache" + else: + cache = Path(tempfile.gettempdir()) / "openjarvis-swebench-cache" + cache.mkdir(parents=True, exist_ok=True) + return cache + + +def _find_report(cache: Path, instance_id: str, run_id: str) -> Optional[Dict[str, Any]]: + """Find the harness's report JSON for one instance. + + swebench writes ``..json`` inside the + subprocess CWD. We use ``model_name_or_path="openjarvis-harness"``, + ``run_id=f"oj-{instance_id}"`` in :func:`_run_harness`. + """ + fname = f"openjarvis-harness.{run_id}.json" + p = cache / fname + if not p.exists(): + return None + try: + return json.loads(p.read_text()) + except Exception: + return None + + +def _run_harness(instance_id: str, patch: str, timeout_s: int) -> Dict[str, Any]: + """Hand one prediction to ``python -m swebench.harness.run_evaluation``. + + Returns ``{"success": bool, "score": float, "details": dict}``. + """ + _apply_patches_once() + backend = os.environ.get("SWEBENCH_BACKEND", "modal").lower() + cache = _harness_cache_dir() + run_id = f"oj-{instance_id}" + + # Defend against stale reports: ``run_id`` is deterministic per + # instance, the cache dir is shared across runs, and ``_find_report`` + # globs by filename. If a prior subprocess crashed mid-run (or was + # killed by timeout) and left a stale JSON, we'd silently read that + # old verdict as the current result. Delete it up front. + stale = cache / f"openjarvis-harness.{run_id}.json" + if stale.exists(): + try: + stale.unlink() + except OSError as exc: + logger.warning("Could not remove stale report %s: %s", stale, exc) + + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + preds_path = tmp_path / "predictions.jsonl" + preds_path.write_text(json.dumps({ + "instance_id": instance_id, + "model_name_or_path": "openjarvis-harness", + "model_patch": patch, + }) + "\n") + + cmd = [ + sys.executable, "-m", "swebench.harness.run_evaluation", + "--predictions_path", str(preds_path), + "--max_workers", "1", + "--run_id", run_id, + "--dataset_name", "SWE-bench/SWE-bench_Verified", + "--instance_ids", instance_id, + ] + if backend == "modal": + cmd += ["--modal", "true"] + + proc = subprocess.run( + cmd, capture_output=True, text=True, + timeout=timeout_s, cwd=str(cache), + ) + + report = _find_report(cache, instance_id, run_id) + if report is None: + return { + "success": False, + "score": 0.0, + "details": { + "reason": "no_report", + "stdout": proc.stdout[-2000:], + "stderr": proc.stderr[-2000:], + }, + } + + # The fix from `rescore.py`: read `resolved_ids` (current schema) + # not `resolved_instances` (older). Older harness wrote lists into + # `resolved_instances`; current swebench puts the count there and + # the actual instance ids in `resolved_ids`. + resolved_ids = report.get("resolved_ids") or [] + # Belt-and-suspenders: also accept the older list-typed field, in + # case the user is on an older swebench install. + if not resolved_ids: + legacy = report.get("resolved_instances") + if isinstance(legacy, list): + resolved_ids = legacy + resolved = instance_id in resolved_ids + return { + "success": resolved, + "score": 1.0 if resolved else 0.0, + "details": {"report": report}, + } + + +# ---------- Scorer ---------- + +class SWEBenchHarnessScorer(Scorer): + """SWE-bench Verified scorer that runs the official harness. + + ``score(record, model_answer)`` returns ``(is_correct, details)``: + + - ``is_correct = True`` if the harness marks the instance resolved. + - ``is_correct = False`` on harness failure or unresolved. + - ``details`` includes the raw harness report under ``["report"]`` plus + a ``"patch"`` field with the extracted patch text. + """ + + scorer_id = "swebench_harness" + + def __init__( + self, + *, + timeout_s: int = 1800, + judge_backend: object = None, # noqa: ARG002 — CLI factory compat + judge_model: str = "", # noqa: ARG002 — CLI factory compat + ) -> None: + self._timeout_s = int(timeout_s) + + def score( + self, + record: EvalRecord, + model_answer: str, + ) -> Tuple[Optional[bool], Dict[str, Any]]: + if not model_answer or not model_answer.strip(): + return False, {"reason": "empty_response"} + + patch = extract_patch(model_answer) + if patch is None: + return False, {"reason": "no_patch_extracted"} + + instance_id = ( + record.metadata.get("instance_id") + or record.record_id + or "" + ) + if not instance_id: + return False, {"reason": "missing_instance_id"} + + result = _run_harness(instance_id, patch, self._timeout_s) + details = dict(result.get("details", {})) + details["patch"] = patch + return bool(result["success"]), details + + +__all__ = ["SWEBenchHarnessScorer", "extract_patch"] diff --git a/src/openjarvis/telemetry/instrumented_engine.py b/src/openjarvis/telemetry/instrumented_engine.py index 1e6491d8..762849df 100644 --- a/src/openjarvis/telemetry/instrumented_engine.py +++ b/src/openjarvis/telemetry/instrumented_engine.py @@ -255,9 +255,11 @@ class InstrumentedEngine(InferenceEngine): "mean_itl_ms": mean_itl_ms, "energy_method": energy_method, "energy_vendor": energy_vendor, - # Rich trace data: model response content + # Rich trace data: model response content + structured blocks "content": result.get("content", ""), "tool_calls": result.get("tool_calls", []), + "tool_results": result.get("tool_results", []), + "content_blocks": result.get("content_blocks", []), "finish_reason": result.get("finish_reason", ""), } diff --git a/src/openjarvis/traces/collector.py b/src/openjarvis/traces/collector.py index 96f17311..0eb4186e 100644 --- a/src/openjarvis/traces/collector.py +++ b/src/openjarvis/traces/collector.py @@ -158,6 +158,8 @@ class TraceCollector: ), "content": data.get("content", ""), "tool_calls": data.get("tool_calls", []), + "tool_results": data.get("tool_results", []), + "content_blocks": data.get("content_blocks", []), "finish_reason": data.get("finish_reason", ""), }, metadata={