diff --git a/fern/versions/latest/pages/reference/cli-commands.mdx b/fern/versions/latest/pages/reference/cli-commands.mdx index fc9518f5d5..9237ccfc9e 100644 --- a/fern/versions/latest/pages/reference/cli-commands.mdx +++ b/fern/versions/latest/pages/reference/cli-commands.mdx @@ -36,6 +36,7 @@ gym dataset collate # validate and collate a dataset # Environments gym env init # scaffold a new resources server gym env resolve # resolve and print the final merged config +gym env compose # swap a benchmark's agent/dataset params and dump the merged config gym env validate # validate a config (no Ray, no servers) — fast pre-flight check gym env packages # list packages in a server's virtual environment gym env test # test resources server(s); all of them if none is given @@ -361,6 +362,29 @@ gym env resolve \ ++responses_create_params.temperature=0.6 ``` +### `gym env compose` + +Rewrite a merged benchmark config along the **agent** and **dataset** axes, then dump the result as YAML. No servers are started and no secrets are read. The agent axis swaps the composable harness, carrying over the environment's `resources_server` / `model_server` / `datasets` wiring; the dataset axis edits the benchmark dataset's run params. Model selection stays with the `--model*` flags (applied before composition). + +| Option | Description | +| --- | --- | +| `--config PATH` | Config file to load. Repeatable. | +| `--benchmark NAME` | Compose the named benchmark config. | +| `--agent NAME` | Agent harness to compose in (`name` or `name/variant`). Must be composable (Pattern A); a self-contained Pattern B agent is rejected. | +| `--num-repeats N` | Override `num_repeats` on the benchmark dataset. | +| `--prompt-config PATH` | Prompt-config YAML to set on the benchmark dataset. | +| `--search-dir DIR` | Extra root directory to search for named components. Repeatable. | + +```bash +# Swap in a different agent harness and bump rollouts per task +gym env compose --benchmark aime24 --agent simple_agent --num-repeats 4 + +# Compose an explicit config with an agent variant +gym env compose \ + --config benchmarks/aime24/config.yaml \ + --agent langgraph_agent/rewoo_agent +``` + ### `gym env validate` Validate a config without starting Ray or any server subprocess — a fast pre-flight check that catches config mistakes (missing/malformed `config_paths`, unknown server cross-references, unset mandatory `???` values, schema errors) in well under a second instead of after a Ray bootstrap. Exits `0` when valid, or `1` with a clean message (no traceback) when not. A model config is **not** required — model interpolations resolve against a dummy model; pass one (or `--model-type`) if you want it validated too. diff --git a/nemo_gym/cli/env.py b/nemo_gym/cli/env.py index b018947bb5..f11861445f 100644 --- a/nemo_gym/cli/env.py +++ b/nemo_gym/cli/env.py @@ -905,6 +905,66 @@ def dump_config(): # pragma: no cover print(OmegaConf.to_yaml(global_config_dict, resolve=True)) +def compose_config(): # pragma: no cover + """Compose a benchmark config: swap the agent harness and/or edit benchmark dataset params, then dump. + + The agent axis swaps the composable harness (carrying the environment's resources_server / + model_server / datasets wiring); the dataset axis edits the benchmark dataset's run params. Model + selection stays with the `--model*` flags. The result is dumped as YAML (no servers are started). + + Examples: + + ```bash + gym env compose --benchmark aime24-x --agent simple_agent --num-repeats 4 + gym env compose --config --agent langgraph_agent/rewoo_agent + ``` + """ + from nemo_gym.agent_registry import discover_agents + from nemo_gym.cli.main import _asset_config_path + from nemo_gym.config_composer import AgentNotComposableError, ComposeRequest, ConfigComposerError, compose + + # NO_MODEL injects a dummy `policy_model` so a benchmark's model cross-refs resolve without real + # model creds — composition is about structure; model selection stays with the `--model*` flags. + global_config_dict = get_global_config_dict( + global_config_dict_parser_config=GlobalConfigDictParserConfig( + initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT, + ), + ) + BaseNeMoGymCLIConfig.model_validate(global_config_dict) + + # Pull the compose directives out so they don't leak into the dumped config (MISSING `???` preserved). + raw = OmegaConf.to_container(global_config_dict, resolve=False, throw_on_missing=False) + agent = raw.pop("compose_agent", None) + num_repeats = raw.pop("compose_num_repeats", None) + prompt_config = raw.pop("compose_prompt_config", None) + merged = OmegaConf.create(raw) + + def _resolve_agent_config_path(name, require_composable=False): + # Reuse the unified CLI's asset resolver for name->path; add only the composability guard. + agent_dir = name.split("/", 1)[0] + if require_composable: + entry = discover_agents().get(agent_dir) + if entry is not None and entry.self_contained: + raise AgentNotComposableError( + f"Agent '{agent_dir}' is self-contained (Pattern B) and cannot be composed into an " + "environment; run it with its own config instead (see `gym list agents`)." + ) + return _asset_config_path("agent", name) + + request = ComposeRequest( + agent=agent, + num_repeats=int(num_repeats) if num_repeats is not None else None, + prompt_config=prompt_config, + ) + try: + composed = compose(merged, request, resolve_agent_config_path=_resolve_agent_config_path) + except (ConfigComposerError, ValueError) as e: + rich.print(f"[red]Error:[/red] {escape(str(e))}") + raise SystemExit(1) + + print(OmegaConf.to_yaml(composed, resolve=False)) + + @exit_cleanly_on_config_error def validate(): """Validate a config without starting Ray or any server subprocess. diff --git a/nemo_gym/cli/main.py b/nemo_gym/cli/main.py index 30cc261f15..6e5e3d5693 100644 --- a/nemo_gym/cli/main.py +++ b/nemo_gym/cli/main.py @@ -130,6 +130,18 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: MODEL_URL = _value_flag("model-url", "policy_base_url", "Model server base URL.") MODEL_API_KEY = _value_flag("model-api-key", "policy_api_key", "Model server API key.") +# Compose flags for `gym env compose`: rewrite a merged benchmark config along the agent/dataset axes. +# Namespaced (`compose_*`) so they don't collide with server-config keys or `eval run`'s --agent/agent_name. +COMPOSE_AGENT = _value_flag( + "agent", "compose_agent", "Agent harness to compose into the environment (name or name/variant)." +) +COMPOSE_NUM_REPEATS = _value_flag( + "num-repeats", "compose_num_repeats", "Override num_repeats on the benchmark dataset." +) +COMPOSE_PROMPT_CONFIG = _value_flag( + "prompt-config", "compose_prompt_config", "Prompt-config path to set on the benchmark dataset." +) + # Shared flag: select a single resources server by name. Reused by `env test`, `env init`, and `env packages`. RESOURCES_SERVER = Flag( register=lambda p: p.add_argument("--resources-server", metavar="NAME", help="Name of the resources server."), @@ -157,6 +169,7 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag: "environment": ("environments", "", "config"), "resources-server": ("resources_servers", "configs", None), "model-type": ("responses_api_models", "configs", None), + "agent": ("responses_api_agents", "configs", None), } @@ -412,6 +425,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None: summary="Resolve the final config from configs, flags, and overrides.", flags=(CONFIG,), ), + "env compose": Command( + target="nemo_gym.cli.env:compose_config", + summary="Compose a benchmark config with a swapped agent and/or dataset params; dump the result.", + flags=(CONFIG, BENCHMARK, SEARCH_DIR, COMPOSE_AGENT, COMPOSE_NUM_REPEATS, COMPOSE_PROMPT_CONFIG), + ), "env validate": Command( target="nemo_gym.cli.env:validate", summary="Validate a config (paths, cross-refs, ??? values, servers) fast — no Ray, no servers.", diff --git a/nemo_gym/config_composer.py b/nemo_gym/config_composer.py new file mode 100644 index 0000000000..e4e61b02d9 --- /dev/null +++ b/nemo_gym/config_composer.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Compose an already-merged NeMo Gym config along the agent / dataset axes. + +This module takes a *merged* OmegaConf :class:`~omegaconf.DictConfig` — the kind produced by Gym's +Hydra config stack from a benchmark's ``config.yaml`` — and rewrites it in place along two axes: + +- **Agent axis:** swap the composable agent harness (``--agent``) carrying the benchmark dataset for + a different one, re-using the environment's own wiring (resources server, model server, datasets). +- **Dataset axis:** edit the benchmark dataset entry's run parameters (``--num-repeats``, + ``--prompt-config``). + +The *model axis* is intentionally out of scope: model selection is handled by the CLI's +``--model*`` flags long before composition runs, so :func:`substitute_model` is a documented no-op. + +Like :mod:`nemo_gym.agent_registry` and :mod:`nemo_gym.benchmarks`, composition is *pure and +resolution-safe*: it never resolves interpolations (``resolve=False``), never starts servers, and +never reads secrets, so it is safe to call when API keys referenced by a config are unset. It also +never emits Hydra override tokens — it returns a new :class:`~omegaconf.DictConfig`. +""" + +from dataclasses import dataclass +from typing import Callable, List, Optional + +from omegaconf import DictConfig, ListConfig, OmegaConf + + +# Keys carried over from the environment's existing agent block onto a freshly substituted agent, so +# the agent stays wired to the same resources server, policy model, and benchmark dataset. +_AGENT_TYPE_KEY = "responses_api_agents" +_WIRING_KEYS = ("resources_server", "model_server", "datasets") +_BENCHMARK_DATASET_TYPE = "benchmark" + + +class ConfigComposerError(ValueError): + """Base error for failures while composing a merged config.""" + + +class NoComposableAgentBlockError(ConfigComposerError): + """The merged config has zero (or several ambiguous) agent blocks to compose against.""" + + +class MandatoryPlaceholderError(ConfigComposerError): + """A mandatory ``???`` (OmegaConf MISSING) field remained after composition.""" + + +class AgentNotComposableError(ConfigComposerError): + """The requested agent is self-contained (Pattern B) and cannot be wired into an environment. + + Raised by the agent-resolution callable passed to :func:`substitute_agent` / :func:`compose` + when ``require_composable=True`` and the agent ships its own environment/framework. Composability + is classified by :func:`nemo_gym.agent_registry.discover_agents` (``AgentEntry.self_contained``). + """ + + +@dataclass(frozen=True) +class ComposeRequest: + """A composition request along the agent and dataset axes (model axis handled by the CLI).""" + + agent: Optional[str] = None + num_repeats: Optional[int] = None + prompt_config: Optional[str] = None + + @property + def is_empty(self) -> bool: + return self.agent is None and self.num_repeats is None and self.prompt_config is None + + +def _deepcopy(merged: DictConfig) -> DictConfig: + """Return an independent, resolution-safe copy of ``merged`` (MISSING fields preserved).""" + return OmegaConf.create(OmegaConf.to_container(merged, resolve=False, throw_on_missing=False)) + + +def _agent_block_type(block: DictConfig) -> Optional[str]: + """Return the single ``responses_api_agents.`` inner key of a top-level block, if any.""" + agents = block.get(_AGENT_TYPE_KEY) if isinstance(block, DictConfig) else None + if not isinstance(agents, DictConfig): + return None + keys = list(agents.keys()) + if len(keys) != 1: + return None + return keys[0] + + +def _dataset_is_benchmark(dataset) -> bool: + return isinstance(dataset, DictConfig) and dataset.get("type") == _BENCHMARK_DATASET_TYPE + + +def find_agent_block_key(merged: DictConfig) -> str: + """Return the top-level key whose ``responses_api_agents`` block is the composition target. + + Mirroring :mod:`nemo_gym.benchmarks`, the target is the agent block whose ``datasets`` contain a + ``type: benchmark`` entry. If exactly one agent block carries such a dataset, it wins even when + other (auxiliary) agent blocks exist. Otherwise the choice must be unambiguous: a single agent + block overall is accepted. Raises :class:`NoComposableAgentBlockError` for zero candidates or an + ambiguous selection, listing the candidate keys. + """ + agent_block_keys: List[str] = [] + benchmark_block_keys: List[str] = [] + for top_level_key in merged: + block = merged[top_level_key] + agent_type = _agent_block_type(block) + if agent_type is None: + continue + agent_block_keys.append(top_level_key) + inner = block[_AGENT_TYPE_KEY][agent_type] + datasets = inner.get("datasets") if isinstance(inner, DictConfig) else None + if isinstance(datasets, (list, ListConfig)) and any(_dataset_is_benchmark(d) for d in datasets): + benchmark_block_keys.append(top_level_key) + + if len(benchmark_block_keys) == 1: + return benchmark_block_keys[0] + if len(benchmark_block_keys) > 1: + raise NoComposableAgentBlockError( + "Multiple agent blocks carry a `type: benchmark` dataset; cannot pick a composition " + f"target unambiguously. Candidates: {sorted(benchmark_block_keys)}." + ) + if len(agent_block_keys) == 1: + return agent_block_keys[0] + if len(agent_block_keys) == 0: + raise NoComposableAgentBlockError( + "No `responses_api_agents` block found in the merged config to compose against." + ) + raise NoComposableAgentBlockError( + "Ambiguous composition target: several agent blocks exist and none carries a " + f"`type: benchmark` dataset. Candidates: {sorted(agent_block_keys)}." + ) + + +def substitute_agent( + merged: DictConfig, + agent_block_key: str, + new_agent: str, + *, + resolve_agent_config_path: Callable[..., str], +) -> DictConfig: + """Replace the agent harness in ``agent_block_key`` with ``new_agent``, keeping env wiring. + + ``new_agent`` is resolved via ``resolve_agent_config_path(new_agent, require_composable=True)``; + a Pattern B (self-contained) agent raises :class:`AgentNotComposableError`, which is left to + propagate. The new agent config's inner block (shape + ``.responses_api_agents.``) replaces the existing inner agent block, but the + environment's ``resources_server`` / ``model_server`` / ``datasets`` win over the agent config's + own values, and the inner key is re-keyed to ````. + """ + merged = _deepcopy(merged) + + new_agent_path = resolve_agent_config_path(new_agent, require_composable=True) + new_agent_config = OmegaConf.load(new_agent_path) + + new_inner_block: Optional[DictConfig] = None + for top_level_value in new_agent_config.values(): + agent_type = _agent_block_type(top_level_value) + if agent_type is not None: + new_inner_block = top_level_value[_AGENT_TYPE_KEY][agent_type] + break + if new_inner_block is None: + raise NoComposableAgentBlockError( + f"Agent config for '{new_agent}' at {new_agent_path} has no `responses_api_agents` block." + ) + + old_agents = merged[agent_block_key][_AGENT_TYPE_KEY] + old_agent_type = next(iter(old_agents.keys())) + old_inner_block = old_agents[old_agent_type] + + new_inner_block = _deepcopy(new_inner_block) + for wiring_key in _WIRING_KEYS: + if wiring_key in old_inner_block: + new_inner_block[wiring_key] = old_inner_block[wiring_key] + + merged[agent_block_key][_AGENT_TYPE_KEY] = OmegaConf.create({new_agent: new_inner_block}) + return merged + + +def substitute_dataset_params( + merged: DictConfig, + agent_block_key: str, + *, + num_repeats: Optional[int] = None, + prompt_config: Optional[str] = None, +) -> DictConfig: + """Edit the single ``type: benchmark`` dataset in the agent block; no-op for ``None`` fields.""" + if num_repeats is None and prompt_config is None: + return _deepcopy(merged) + + merged = _deepcopy(merged) + agents = merged[agent_block_key][_AGENT_TYPE_KEY] + inner = agents[next(iter(agents.keys()))] + datasets = inner.get("datasets") or [] + benchmark_datasets = [d for d in datasets if _dataset_is_benchmark(d)] + if len(benchmark_datasets) != 1: + raise NoComposableAgentBlockError( + f"Expected exactly one `type: benchmark` dataset in agent block '{agent_block_key}', " + f"found {len(benchmark_datasets)}." + ) + + dataset = benchmark_datasets[0] + if num_repeats is not None: + dataset["num_repeats"] = num_repeats + if prompt_config is not None: + dataset["prompt_config"] = prompt_config + return merged + + +def substitute_model(merged: DictConfig) -> DictConfig: + """No-op shim: the model axis is owned by the CLI's ``--model*`` flags, not the composer.""" + return merged + + +def _validate_no_mandatory_placeholders(merged: DictConfig, agent_block_key: str) -> None: + """Raise :class:`MandatoryPlaceholderError` if any ``???`` remains in the agent or its resources. + + Walks the composed agent block and, if it references one, the resources-server block it points + at (``resources_server.name`` -> top-level key). Interpolations are *not* resolved. + """ + missing: List[str] = [] + + def _walk(node, path: str) -> None: + if isinstance(node, DictConfig): + for key in node: + if OmegaConf.is_missing(node, key): + missing.append(f"{path}.{key}" if path else str(key)) + else: + # Use `_get_node` so interpolations (e.g. a composable agent's `${nvidia_api_key}`) + # are NOT resolved while we only scan for remaining `???`; `node[key]` would + # resolve and raise InterpolationKeyError for keys NO_MODEL doesn't inject. + _walk(node._get_node(key), f"{path}.{key}" if path else str(key)) + elif isinstance(node, (list, ListConfig)): + for index in range(len(node)): + item = node._get_node(index) if isinstance(node, ListConfig) else node[index] + _walk(item, f"{path}[{index}]") + + agent_block = merged[agent_block_key] + _walk(agent_block, agent_block_key) + + agents = agent_block[_AGENT_TYPE_KEY] + inner = agents[next(iter(agents.keys()))] + resources_server = inner.get("resources_server") if isinstance(inner, DictConfig) else None + if isinstance(resources_server, DictConfig) and not OmegaConf.is_missing(resources_server, "name"): + resources_key = resources_server.get("name") + if isinstance(resources_key, str) and resources_key in merged: + _walk(merged[resources_key], resources_key) + + if missing: + raise MandatoryPlaceholderError( + "Mandatory `???` field(s) remain after composition; provide them via CLI/config overrides: " + + ", ".join(sorted(missing)) + ) + + +def compose( + merged: DictConfig, + request: ComposeRequest, + *, + resolve_agent_config_path: Callable[..., str], +) -> DictConfig: + """Compose ``merged`` per ``request`` and return a new validated config. + + Pipeline: locate the agent block, optionally swap the agent (carrying env wiring), edit the + benchmark dataset params, then assert no mandatory ``???`` field remains. An empty request is a + pure passthrough (still deep-copied and validated). + """ + result = _deepcopy(merged) + agent_block_key = find_agent_block_key(result) + + if request.agent is not None: + result = substitute_agent( + result, + agent_block_key, + request.agent, + resolve_agent_config_path=resolve_agent_config_path, + ) + + result = substitute_dataset_params( + result, + agent_block_key, + num_repeats=request.num_repeats, + prompt_config=request.prompt_config, + ) + + result = substitute_model(result) + + _validate_no_mandatory_placeholders(result, agent_block_key) + return result diff --git a/tests/unit_tests/test_cli_main.py b/tests/unit_tests/test_cli_main.py index 20d6ae205c..a8259b6aa5 100644 --- a/tests/unit_tests/test_cli_main.py +++ b/tests/unit_tests/test_cli_main.py @@ -262,6 +262,27 @@ def test_direct_entrypoint_override_also_runs_single(self, monkeypatch: MonkeyPa assert overrides == ["+entrypoint=resources_servers/gpqa"] +class TestEnvComposeFlags: + def test_compose_dispatches_to_compose_config(self, monkeypatch: MonkeyPatch) -> None: + target, _ = _dispatch_for(monkeypatch, ["env", "compose", "--benchmark", "gsm8k"]) + assert target == "nemo_gym.cli.env:compose_config" + + def test_compose_flags_map_to_namespaced_hydra_keys(self, monkeypatch: MonkeyPatch) -> None: + # --agent/--num-repeats/--prompt-config are namespaced (compose_*) so they don't collide with + # server-config keys or eval run's --agent/agent_name. + _, overrides = _dispatch_for( + monkeypatch, + ["env", "compose", "--agent", "simple_agent", "--num-repeats", "4", "--prompt-config", "p.yaml"], + ) + assert "+compose_agent=simple_agent" in overrides + assert "+compose_num_repeats=4" in overrides + assert "+compose_prompt_config=p.yaml" in overrides + + def test_compose_benchmark_selector_resolves_to_config_paths(self, monkeypatch: MonkeyPatch) -> None: + _, overrides = _dispatch_for(monkeypatch, ["env", "compose", "--benchmark", "gsm8k"]) + assert any(o.startswith("+config_paths=[") and "gsm8k" in o for o in overrides) + + class TestDatasetFlags: def test_upload_hf_default(self, monkeypatch: MonkeyPatch) -> None: target, overrides = _dispatch_for( diff --git a/tests/unit_tests/test_config_composer.py b/tests/unit_tests/test_config_composer.py new file mode 100644 index 0000000000..1a2181a4fb --- /dev/null +++ b/tests/unit_tests/test_config_composer.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from pathlib import Path + +from omegaconf import OmegaConf +from pytest import raises + +from nemo_gym.config_composer import ( + AgentNotComposableError, + ComposeRequest, + MandatoryPlaceholderError, + NoComposableAgentBlockError, + _validate_no_mandatory_placeholders, + compose, + find_agent_block_key, + substitute_agent, + substitute_dataset_params, + substitute_model, +) + + +def _merged(*, agent_type: str = "simple_agent", resources_name: str = "bench_resources", num_repeats: int = 16): + """A merged benchmark config: one resources block + one agent block carrying the benchmark.""" + return OmegaConf.create( + { + "policy_model": {"responses_api_models": {"vllm_model": {"entrypoint": "app.py"}}}, + resources_name: { + "resources_servers": { + "bench": { + "entrypoint": "app.py", + "domain": "other", + "judge_model_server": {"type": "responses_api_models", "name": "policy_model"}, + } + } + }, + "bench_agent": { + "responses_api_agents": { + agent_type: { + "entrypoint": "app.py", + "resources_server": {"type": "resources_servers", "name": resources_name}, + "model_server": {"type": "responses_api_models", "name": "policy_model"}, + "datasets": [ + { + "name": "bench", + "type": "benchmark", + "jsonl_fpath": "benchmarks/bench/data/bench.jsonl", + "prompt_config": None, + "prepare_script": "benchmarks/bench/prepare.py", + "num_repeats": num_repeats, + } + ], + } + } + }, + } + ) + + +def _fake_resolver(tmp_path: Path, agent_name: str, *, body: str): + """Return a resolve_agent_config_path stub that writes ``body`` and returns its path.""" + config_path = tmp_path / f"{agent_name}.yaml" + config_path.write_text(body) + + def _resolve(name, require_composable=False): + if name == "swe_agents" and require_composable: + raise AgentNotComposableError("Agent 'swe_agents' is self-contained") + return str(config_path) + + return _resolve + + +def _new_agent_body(agent_name: str = "react_agent") -> str: + # A composable agent config with its own (to-be-overridden) wiring. + return ( + f"some_key:\n responses_api_agents:\n {agent_name}:\n entrypoint: app.py\n" + " max_turns: 8\n" + " resources_server:\n type: resources_servers\n name: ???\n" + " model_server:\n type: responses_api_models\n name: policy_model\n" + ) + + +class TestFindAgentBlockKey: + def test_finds_benchmark_carrying_block(self) -> None: + assert find_agent_block_key(_merged()) == "bench_agent" + + def test_prefers_benchmark_block_over_auxiliary_agent(self) -> None: + merged = _merged() + merged["aux_agent"] = OmegaConf.create({"responses_api_agents": {"helper_agent": {"entrypoint": "app.py"}}}) + assert find_agent_block_key(merged) == "bench_agent" + + def test_single_agent_block_without_benchmark_is_accepted(self) -> None: + merged = OmegaConf.create({"only_agent": {"responses_api_agents": {"simple_agent": {"entrypoint": "app.py"}}}}) + assert find_agent_block_key(merged) == "only_agent" + + def test_zero_agent_blocks_raises(self) -> None: + merged = OmegaConf.create({"policy_model": {"responses_api_models": {"vllm_model": {}}}}) + with raises(NoComposableAgentBlockError, match="No `responses_api_agents` block"): + find_agent_block_key(merged) + + def test_ambiguous_agent_blocks_without_benchmark_raises(self) -> None: + merged = OmegaConf.create( + { + "a": {"responses_api_agents": {"x": {"entrypoint": "app.py"}}}, + "b": {"responses_api_agents": {"y": {"entrypoint": "app.py"}}}, + } + ) + with raises(NoComposableAgentBlockError, match="Ambiguous composition target"): + find_agent_block_key(merged) + + def test_multiple_benchmark_blocks_raises(self) -> None: + merged = _merged() + merged["bench_agent_2"] = OmegaConf.create(merged["bench_agent"]) + with raises(NoComposableAgentBlockError, match="Multiple agent blocks"): + find_agent_block_key(merged) + + def test_block_with_multiple_agent_types_is_ignored(self) -> None: + # A `responses_api_agents` mapping with two inner types is not a valid single-agent block. + merged = OmegaConf.create( + {"weird": {"responses_api_agents": {"a": {"entrypoint": "app.py"}, "b": {"entrypoint": "app.py"}}}} + ) + with raises(NoComposableAgentBlockError, match="No `responses_api_agents` block"): + find_agent_block_key(merged) + + +class TestSubstituteDatasetParams: + def test_sets_num_repeats_and_prompt_config(self) -> None: + merged = _merged(num_repeats=16) + out = substitute_dataset_params(merged, "bench_agent", num_repeats=4, prompt_config="prompts/p.yaml") + + dataset = out["bench_agent"]["responses_api_agents"]["simple_agent"]["datasets"][0] + assert dataset["num_repeats"] == 4 + assert dataset["prompt_config"] == "prompts/p.yaml" + # Original is untouched (deep copy). + assert merged["bench_agent"]["responses_api_agents"]["simple_agent"]["datasets"][0]["num_repeats"] == 16 + + def test_none_fields_are_noop(self) -> None: + merged = _merged(num_repeats=16) + out = substitute_dataset_params(merged, "bench_agent") + assert out["bench_agent"]["responses_api_agents"]["simple_agent"]["datasets"][0]["num_repeats"] == 16 + + def test_missing_benchmark_dataset_raises(self) -> None: + merged = _merged() + merged["bench_agent"]["responses_api_agents"]["simple_agent"]["datasets"] = [] + with raises(NoComposableAgentBlockError, match="Expected exactly one"): + substitute_dataset_params(merged, "bench_agent", num_repeats=2) + + +class TestSubstituteAgent: + def test_swaps_agent_and_carries_over_wiring(self, tmp_path: Path) -> None: + merged = _merged(agent_type="simple_agent") + resolver = _fake_resolver(tmp_path, "react_agent", body=_new_agent_body("react_agent")) + + out = substitute_agent(merged, "bench_agent", "react_agent", resolve_agent_config_path=resolver) + + agents = out["bench_agent"]["responses_api_agents"] + # Re-keyed to the new agent name; old type gone. + assert list(agents.keys()) == ["react_agent"] + block = agents["react_agent"] + # New agent's own field is preserved. + assert block["max_turns"] == 8 + # Env wiring wins over the agent config's values. + assert block["resources_server"]["name"] == "bench_resources" + assert block["model_server"]["name"] == "policy_model" + assert block["datasets"][0]["name"] == "bench" + + def test_pattern_b_agent_propagates_not_composable(self, tmp_path: Path) -> None: + merged = _merged() + resolver = _fake_resolver(tmp_path, "swe_agents", body=_new_agent_body()) + with raises(AgentNotComposableError): + substitute_agent(merged, "bench_agent", "swe_agents", resolve_agent_config_path=resolver) + + def test_agent_config_without_agent_block_raises(self, tmp_path: Path) -> None: + merged = _merged() + resolver = _fake_resolver(tmp_path, "broken", body="k:\n not_an_agent: true\n") + with raises(NoComposableAgentBlockError, match="has no `responses_api_agents` block"): + substitute_agent(merged, "bench_agent", "broken", resolve_agent_config_path=resolver) + + +class TestValidateNoMandatoryPlaceholders: + def test_passes_when_filled(self) -> None: + _validate_no_mandatory_placeholders(_merged(), "bench_agent") + + def test_raises_on_missing_in_agent_block(self) -> None: + merged = _merged() + merged["bench_agent"]["responses_api_agents"]["simple_agent"]["resources_server"]["name"] = "???" + with raises(MandatoryPlaceholderError, match="resources_server.name"): + _validate_no_mandatory_placeholders(merged, "bench_agent") + + def test_unresolvable_interpolation_does_not_crash(self) -> None: + # A composable (Pattern A) agent may interpolate a key NO_MODEL doesn't inject (e.g. a + # provider api key like ${nvidia_api_key}). Scanning for `???` must NOT resolve it — doing so + # would raise InterpolationKeyError and block an otherwise-valid composition. + merged = _merged() + merged["bench_agent"]["responses_api_agents"]["simple_agent"]["api_key"] = "${nvidia_api_key}" + _validate_no_mandatory_placeholders(merged, "bench_agent") # must not raise + + def test_raises_on_missing_in_referenced_resources_block(self) -> None: + merged = _merged() + merged["bench_resources"]["resources_servers"]["bench"]["judge_model_server"]["name"] = "???" + with raises(MandatoryPlaceholderError, match="judge_model_server.name"): + _validate_no_mandatory_placeholders(merged, "bench_agent") + + def test_missing_resources_server_name_does_not_crash(self) -> None: + merged = _merged() + # If the resources_server reference itself is missing, validation skips it (the agent-block + # walk still flags it as a placeholder). + merged["bench_agent"]["responses_api_agents"]["simple_agent"]["resources_server"]["name"] = "???" + with raises(MandatoryPlaceholderError): + _validate_no_mandatory_placeholders(merged, "bench_agent") + + def test_dangling_resources_reference_is_ignored(self) -> None: + merged = _merged() + merged["bench_agent"]["responses_api_agents"]["simple_agent"]["resources_server"]["name"] = "nonexistent" + # No such top-level block; validation simply does not descend into it. + _validate_no_mandatory_placeholders(merged, "bench_agent") + + +class TestSubstituteModel: + def test_is_noop(self) -> None: + merged = _merged() + assert substitute_model(merged) is merged + + +class TestCompose: + def test_empty_request_is_passthrough(self) -> None: + merged = _merged() + out = compose(merged, ComposeRequest(), resolve_agent_config_path=lambda *a, **k: "") + assert OmegaConf.to_container(out, resolve=False) == OmegaConf.to_container(merged, resolve=False) + assert out is not merged + + def test_compose_request_is_empty_property(self) -> None: + assert ComposeRequest().is_empty is True + assert ComposeRequest(num_repeats=2).is_empty is False + + def test_full_compose_swaps_agent_and_edits_dataset(self, tmp_path: Path) -> None: + merged = _merged(agent_type="simple_agent", num_repeats=16) + resolver = _fake_resolver(tmp_path, "react_agent", body=_new_agent_body("react_agent")) + + out = compose( + merged, + ComposeRequest(agent="react_agent", num_repeats=2, prompt_config="prompts/p.yaml"), + resolve_agent_config_path=resolver, + ) + + block = out["bench_agent"]["responses_api_agents"]["react_agent"] + assert block["resources_server"]["name"] == "bench_resources" + assert block["datasets"][0]["num_repeats"] == 2 + assert block["datasets"][0]["prompt_config"] == "prompts/p.yaml" + + def test_compose_validates_placeholders(self, tmp_path: Path) -> None: + merged = _merged() + merged["bench_resources"]["resources_servers"]["bench"]["judge_model_server"]["name"] = "???" + with raises(MandatoryPlaceholderError): + compose(merged, ComposeRequest(num_repeats=2), resolve_agent_config_path=lambda *a, **k: "") + + def test_compose_propagates_not_composable(self, tmp_path: Path) -> None: + merged = _merged() + resolver = _fake_resolver(tmp_path, "swe_agents", body=_new_agent_body()) + with raises(AgentNotComposableError): + compose(merged, ComposeRequest(agent="swe_agents"), resolve_agent_config_path=resolver)