diff --git a/tests/m1/test_opto_integration.py b/tests/m1/test_opto_integration.py new file mode 100644 index 0000000..9eb67b6 --- /dev/null +++ b/tests/m1/test_opto_integration.py @@ -0,0 +1,87 @@ +"""Tests for the generic opto/OpenTrace bundle-evaluation integration.""" +from __future__ import annotations + +import pytest + +from trace_bench.integrations.opto import BundleEval, evaluate_bundle +from trace_bench.registry import normalize_task_id + + +class _Node: + """Node-like param: _extract_response reads .data (production shape).""" + def __init__(self, data="GOOD"): self.data = data + + +class _Guide: + # Rewards the known-good response; mirrors a real guide scoring a fixed + # artifact response against each example. + def __call__(self, x, response, info): + return (0.5 if response == "GOOD" else 0.0), f"resp={response}" + + def get_score_dict(self, x, response, info): + return {"accuracy": 1.0 if response == "GOOD" else 0.0, "cost": 0.1} + + +def _bundle(**over): + base = {"param": _Node("GOOD"), "guide": _Guide(), + "train_dataset": {"inputs": ["a", "b"], "infos": [{}, {}]}} + base.update(over) + return base + + +def test_evaluate_bundle_returns_mean_reward_and_score_dicts(): + mean_reward, evals = evaluate_bundle(_bundle(), max_examples=2) + assert mean_reward == 0.5 + assert [type(e) for e in evals] == [BundleEval, BundleEval] + assert evals[0].score_dict == {"accuracy": 1.0, "cost": 0.1} + + +def test_evaluate_bundle_prefers_validation_split_and_caps_examples(): + bundle = _bundle(validate_dataset={"inputs": ["a"], "infos": [{}]}) + mean_reward, evals = evaluate_bundle(bundle, max_examples=5) + assert len(evals) == 1 and mean_reward == 0.5 # validate split used, capped + + +def test_evaluate_bundle_unwraps_trainable_param(): + mean_reward, evals = evaluate_bundle(_bundle(param=_Node("GOOD")), max_examples=1) + assert mean_reward == 0.5 and evals[0].feedback == "resp=GOOD" + + +def test_evaluate_bundle_raises_on_empty_dataset(): + with pytest.raises(ValueError, match="no evaluable examples"): + evaluate_bundle(_bundle(train_dataset={"inputs": [], "infos": []})) + + +def test_normalize_task_id_is_public_and_stable(): + assert normalize_task_id("example:foo") == "trace_examples:foo" + assert normalize_task_id("hf:GSM8K") == "hf:GSM8K" + assert normalize_task_id("online_bin_packing") == "llm4ad:online_bin_packing" + + +def test_evaluate_bundle_strict_score_dict_reraises(): + class _BadDict(_Guide): + def get_score_dict(self, x, response, info): + raise RuntimeError("scorer unavailable") + + bundle = _bundle(guide=_BadDict()) + # default: swallow -> score_dict None, reward still present + mean_reward, evals = evaluate_bundle(bundle, max_examples=1) + assert evals[0].score_dict is None and mean_reward == 0.5 + # strict: re-raise so a silent None can't masquerade as a real vector + with pytest.raises(RuntimeError, match="scorer unavailable"): + evaluate_bundle(bundle, max_examples=1, strict_score_dict=True) + + +def test_evaluate_bundle_mean_reward_matches_runner_mean(): + # parity guard: aggregate reward equals the runner's own mean over the split + from trace_bench.runner import _score_dataset, _evaluation_dataset + bundle = _bundle() + name, dataset = _evaluation_dataset(bundle) + runner_summary = _score_dataset(bundle, dataset, name) + mean_reward, _ = evaluate_bundle(bundle, max_examples=len(dataset["inputs"])) + assert mean_reward == pytest.approx(runner_summary["score"]) + + +def test_integrations_package_exports_public_api(): + import trace_bench.integrations as integ + assert set(integ.__all__) == {"BundleEval", "evaluate_bundle", "load_and_evaluate"} diff --git a/trace_bench/integrations/__init__.py b/trace_bench/integrations/__init__.py index a04cd9d..f0f6bac 100644 --- a/trace_bench/integrations/__init__.py +++ b/trace_bench/integrations/__init__.py @@ -1,4 +1,6 @@ from __future__ import annotations -__all__ = [] +# Public integration surface for external OpenTrace / opto consumers. +from .opto import BundleEval, evaluate_bundle, load_and_evaluate +__all__ = ["BundleEval", "evaluate_bundle", "load_and_evaluate"] diff --git a/trace_bench/integrations/opto.py b/trace_bench/integrations/opto.py new file mode 100644 index 0000000..10c12cd --- /dev/null +++ b/trace_bench/integrations/opto.py @@ -0,0 +1,112 @@ +"""Bundle evaluation helper for external OpenTrace / opto ``recursive_opt`` consumers. + +WHY THIS EXISTS +--------------- +NewTrace's ``opto.features.recursive_opt`` adapter (``TraceBenchTaskAdapter``) +needs to evaluate a loaded Trace-Bench bundle on its own dataset and read back +BOTH the scalar reward and the optional multi-objective ``score_dict`` for its +recursive scoring transforms (relative deltas, clipping, objective +scalarization). The benchmark repo is the right home for "load a bundle and run +it on its data"; the consumer keeps only the recursion-specific transforms. + +This module is therefore a thin, stable, PUBLIC API surface for that external +consumer. It is intentionally not wired into ``runner.py``'s job pipeline (the +runner has its own batched/staged evaluation with stub-LLM handling). To avoid +the semantic drift the reviewer rightly flagged, the per-example evaluation here +REUSES the runner's own ``_extract_response`` and ``_evaluation_dataset`` rather +than re-deriving them — so "what counts as the response / which split is used" +has exactly one definition in this repo. + +The only deliberate behavioral difference from ``runner._score_dataset`` is +shape, by design for the consumer: this returns per-example ``BundleEval`` +records (reward + feedback + optional score_dict) instead of a single +``{score, feedback}`` summary, because the recursive optimizer scores each +example's objective vector. Aggregate reward matches the runner's mean. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +from trace_bench.registry import load_task_bundle, normalize_task_id + +__all__ = ["BundleEval", "evaluate_bundle", "load_and_evaluate"] + + +@dataclass +class BundleEval: + """One example's evaluation: official reward, feedback, optional score dict.""" + + reward: float + feedback: str + score_dict: Optional[Dict[str, float]] = None + + +def evaluate_bundle( + bundle: Dict[str, Any], + *, + max_examples: int = 1, + strict_score_dict: bool = False, +) -> Tuple[float, List[BundleEval]]: + """Evaluate a loaded bundle on up to ``max_examples`` of its own dataset. + + Returns ``(mean_reward, per_example_evals)``. Split selection and response + extraction are delegated to the runner's helpers (single source of truth). + + - ``strict_score_dict=False`` (default): a guide whose ``get_score_dict`` + raises yields ``score_dict=None`` for that example (the scalar reward is + still recorded). This matches "score_dict is an optional enrichment". + - ``strict_score_dict=True``: re-raise instead of swallowing, for consumers + whose multi-objective scoring MUST have the dict (so a silent None can't + masquerade as a degenerate objective vector). + + Raises ``ValueError`` when the chosen dataset is empty — a silent + zero-example "success" would hide a real loading/dataset failure. + """ + # Reuse the runner's canonical helpers so this integration cannot drift from + # the repo's own response-extraction / split-selection logic. Imported lazily + # to avoid a circular import (runner -> integrations package -> this module). + from trace_bench.runner import _evaluation_dataset, _extract_response + + dataset_name, dataset = _evaluation_dataset(bundle) + inputs = list(dataset.get("inputs") or []) + infos = list(dataset.get("infos") or dataset.get("info") or []) + limit = min(len(inputs), len(infos) or len(inputs), + max_examples if max_examples and max_examples > 0 else len(inputs)) + if limit <= 0: + raise ValueError(f"bundle dataset {dataset_name!r} has no evaluable examples") + + guide = bundle["guide"] + param = bundle["param"] + out: List[BundleEval] = [] + for i in range(limit): + info = infos[i] if i < len(infos) else None + response = _extract_response(param, inputs[i]) + reward, feedback = guide(inputs[i], response, info) + score_dict = None + if hasattr(guide, "get_score_dict"): + try: + score_dict = guide.get_score_dict(inputs[i], response, info) + except Exception: + if strict_score_dict: + raise + score_dict = None + out.append(BundleEval(reward=float(reward), feedback=str(feedback), + score_dict=score_dict)) + return sum(ev.reward for ev in out) / len(out), out + + +def load_and_evaluate( + task_id: str, + tasks_root: str, + *, + eval_kwargs: Optional[Dict[str, Any]] = None, + max_examples: int = 1, + strict_score_dict: bool = False, +) -> Tuple[Dict[str, Any], float, List[BundleEval]]: + """One-call helper: load the bundle for ``task_id``, then evaluate it.""" + bundle = load_task_bundle(normalize_task_id(task_id), tasks_root, + eval_kwargs=eval_kwargs) + mean_reward, evals = evaluate_bundle( + bundle, max_examples=max_examples, strict_score_dict=strict_score_dict) + return bundle, mean_reward, evals diff --git a/trace_bench/registry.py b/trace_bench/registry.py index 54f2f94..1e62e43 100644 --- a/trace_bench/registry.py +++ b/trace_bench/registry.py @@ -401,6 +401,11 @@ def discover_tasks(tasks_root: str | Path, bench: Optional[str] = None) -> List[ return specs +def normalize_task_id(task_id: str) -> str: + """Public task-id normalization for external integrations (stable API).""" + return _normalize_task_id(task_id) + + def _normalize_task_id(task_id: str) -> str: if task_id.startswith("example:"): return task_id.replace("example:", "trace_examples:", 1) @@ -516,5 +521,6 @@ def load_task_bundle(task_id: str, tasks_root: str | Path, eval_kwargs: Optional "expand_special_tasks", "load_task_bundle", "load_task_module", + "normalize_task_id", "priority_search_example_trainers_supported", ]