evalrx 0.1.2__py3-none-any.whl

evalrx/models/__init__.py

@@ -0,0 +1,161 @@
+"""Model construction — one engine: ``compose(spec, backend)``.
+
+``evalrx.specs`` is the single source of truth for *which* models exist;
+``models/backends`` provides the runtimes.  Everything funnels through
+:func:`load` (friendly, key-based) and :func:`load_model` (``ModelConfig``-based),
+both of which build via ``compose`` — so there is exactly one way a model is
+constructed, and the capability set always reflects the chosen backend.
+"""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import Any
+
+from evalrx.config import ModelConfig
+from evalrx.core.capability import Capability
+from evalrx.core.model import Model
+
+# Backend layer + agent loop (torch-free at import; heavy deps are lazy in build/load).
+from evalrx.core.tool import ChatTurn, Tool, ToolCall
+from evalrx.models.agent import Agent, APIToolHandlerExecutor, ToolExecutor
+from evalrx.models.backends import BACKENDS, RuntimeConfig, call_vision_api_chat_fn
+from evalrx.models.base import BaseAgent
+from evalrx.models.compose import compose
+from evalrx.models.toolcodec import OpenAIToolCodec, QwenToolCodec, codec_for
+
+# Legacy concrete white-box model (now a thin deprecated alias) — optional on the
+# light, pure-API install (its module imports nothing heavy at top level).
+try:
+    from evalrx.models import whitebox as _whitebox  # noqa: F401
+except ImportError:  # pragma: no cover - light install
+    _whitebox = None
+
+__all__ = [
+    "Model",
+    "BaseAgent",
+    "load",
+    "load_model",
+    "wrap",
+    "compose",
+    "RuntimeConfig",
+    "BACKENDS",
+    "resolve_spec_key",
+    "Agent",
+    "ToolExecutor",
+    "APIToolHandlerExecutor",
+    "Tool",
+    "ToolCall",
+    "ChatTurn",
+    "OpenAIToolCodec",
+    "QwenToolCodec",
+    "codec_for",
+    "call_vision_api_chat_fn",
+]
+
+# Legacy name → canonical spec key.  Keeps old configs/code working.
+_SPEC_ALIASES: dict[str, str] = {
+    "qwen": "qwen2.5-7b-instruct",
+    "qwen2.5-7b": "qwen2.5-7b-instruct",
+}
+
+
+def resolve_spec_key(name: str) -> str:
+    """Resolve a spec key or legacy alias to a canonical :mod:`evalrx.specs` key."""
+    from evalrx.specs import REGISTRY
+
+    key = name.lower()
+    if key in REGISTRY:
+        return key
+    if key in _SPEC_ALIASES:
+        return _SPEC_ALIASES[key]
+    raise ValueError(
+        f"Unknown model {name!r}. Known specs: {sorted(REGISTRY)}; "
+        f"aliases: {sorted(_SPEC_ALIASES)}."
+    )
+
+
+def load(
+    key: str,
+    *,
+    backend: str = "hf_local",
+    want: "list[str] | set[Capability] | tuple" = (),
+    checkpoint: str | None = None,
+    **runtime: Any,
+) -> Model:
+    """Build a model by spec key — the friendly one-liner.
+
+    ``evalrx.load("qwen2.5-7b-instruct")`` is the modern replacement for the
+    old ``QwenLLM()``.  It wraps ``compose(get_spec(key), backend, RuntimeConfig(**runtime))``
+    and negotiates capabilities up front (``want``).
+
+    Args:
+        key:        Spec key or legacy alias (see :func:`resolve_spec_key`).
+        backend:    ``"hf_local"`` (internals) / ``"api"`` / ``"vllm_offline"``.
+                    Forced to ``"api"`` for ``api_only`` (closed-weight) specs.
+        want:       Capabilities the backend must provide (names or ``Capability``).
+        checkpoint: Optional override of the spec's ``hf_repo``.
+        **runtime:  Forwarded to :class:`RuntimeConfig` (``device``, ``dtype``, …).
+    """
+    from evalrx.specs import get_spec
+
+    spec = get_spec(resolve_spec_key(key))
+    if checkpoint:
+        spec = replace(spec, hf_repo=checkpoint)
+    chosen = "api" if spec.api_only else backend
+    caps = {w if isinstance(w, Capability) else Capability(w) for w in want}
+    return compose(spec, chosen, RuntimeConfig(**runtime), caps)
+
+
+def load_model(cfg: ModelConfig) -> Model:
+    """Instantiate a model from a :class:`~evalrx.config.ModelConfig` via :func:`load`."""
+    return load(
+        cfg.name,
+        backend=cfg.backend,
+        want=cfg.want,
+        checkpoint=cfg.checkpoint,
+        device=cfg.device,
+        dtype=cfg.dtype,
+    )
+
+
+def wrap(
+    model: Any,
+    tokenizer: Any = None,
+    *,
+    want: "list[str] | set[Capability] | tuple" = (),
+    **runtime: Any,
+) -> Model:
+    """Wrap an ALREADY-LOADED HF causal LM + tokenizer into an analyzable model.
+
+    The captum-style public on-ramp: the user brings their own model, no registry
+    key needed.  The result is the same :class:`~evalrx.models.backends.hf_local.HFLocalModel`
+    that ``evalrx.load("qwen...")`` produces, so every capability-compatible
+    analyzer works on it::
+
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        m = AutoModelForCausalLM.from_pretrained("my-org/my-llama")
+        tok = AutoTokenizer.from_pretrained("my-org/my-llama")
+
+        wrapped = evalrx.wrap(m, tok)
+        LogitLensAnalyzer().run(wrapped, "The capital of France is")
+
+    Capabilities are inferred from the live model; attention capture needs eager
+    attention, which ``wrap`` enables when it can (see ``HFLocalModel.from_loaded``).
+
+    Args:
+        model:     an instantiated ``transformers`` decoder-only model.
+        tokenizer: its tokenizer or processor (positional, recommended).
+        want:      capabilities to assert up front (names or ``Capability``); raises
+                   :class:`~evalrx.core.capability.CapabilityError` if unavailable.
+        **runtime: forwarded to :class:`RuntimeConfig` (``max_new_tokens``, ``device``…).
+
+    Raises:
+        NotImplementedError: if *model* looks like a VLM (Stage 2; text-only for now).
+    """
+    from evalrx.models.backends.hf_local import HFLocalModel
+    from evalrx.models.compose import negotiate
+
+    handle = HFLocalModel.from_loaded(model, tokenizer, runtime=RuntimeConfig(**runtime))
+    caps = {w if isinstance(w, Capability) else Capability(w) for w in want}
+    return negotiate(handle, caps, model_id=handle.spec.key, where="wrap")

evalrx/models/_discover.py

@@ -0,0 +1,101 @@
+"""Runtime module-path discovery for white-box hooking.
+
+Hardcoded dotted paths (``model.model.language_model.layers``) are the #1 source
+of silent white-box bugs: they differ per model AND per transformers release (the
+doubled-``.model.`` / Llama4-no-``.model`` / v5 fused-experts traps).  So instead
+of trusting ``ModelSpec.module_paths``, we DISCOVER the decoder-layer
+``ModuleList`` from the live module tree and derive everything relative to it.
+
+``torch`` is imported lazily inside the functions, so importing this module on the
+light (pure-API) install does not pull torch.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def find_decoder_layers(model: Any, num_hidden_layers: int | None = None) -> tuple[Any, str]:
+    """Find the decoder-layer ``ModuleList`` and its dotted path.
+
+    Strategy: among all ``nn.ModuleList`` modules, prefer the one whose length
+    equals ``num_hidden_layers`` (read from config when not given) and whose
+    children look like decoder layers (have a ``self_attn`` attr or a
+    ``*DecoderLayer`` class name).  Ties broken by the *shallowest* path so we
+    pick the language tower, not a vision encoder's block list.
+    """
+    import torch.nn as nn
+
+    if num_hidden_layers is None:
+        cfg = getattr(model, "config", None)
+        text_cfg = cfg.get_text_config() if hasattr(cfg, "get_text_config") else cfg
+        num_hidden_layers = getattr(text_cfg, "num_hidden_layers", None)
+
+    candidates: list[tuple[str, Any]] = []
+    for name, mod in model.named_modules():
+        if not isinstance(mod, nn.ModuleList) or len(mod) == 0:
+            continue
+        looks_like_layer = hasattr(mod[0], "self_attn") or "DecoderLayer" in type(mod[0]).__name__
+        if not looks_like_layer:
+            continue
+        if num_hidden_layers is not None and len(mod) != num_hidden_layers:
+            continue
+        candidates.append((name, mod))
+
+    if not candidates:
+        raise RuntimeError(
+            "Could not discover the decoder-layer ModuleList"
+            + (f" (expected length {num_hidden_layers})" if num_hidden_layers else "")
+            + f". Tree top-level: {[n for n, _ in model.named_children()]}"
+        )
+    name, layers = min(candidates, key=lambda kv: kv[0].count("."))
+    return layers, name
+
+
+def get_unembed(model: Any) -> Any:
+    """Return the unembedding (lm_head), handling tied embeddings.
+
+    Uses transformers' own ``get_output_embeddings()`` so we never hardcode the
+    ``lm_head`` path and the tied-weights case is handled by the model itself.
+    """
+    head = model.get_output_embeddings() if hasattr(model, "get_output_embeddings") else None
+    if head is None:  # tied / unusual — fall back to input embeddings' weight
+        emb = model.get_input_embeddings()
+        return emb
+    return head
+
+
+def get_final_norm(model: Any) -> Any:
+    """Return the final normalization module applied after the last decoder layer.
+
+    Found relative to the discovered decoder-layer list: its parent module's
+    norm child (``norm`` / ``final_layernorm`` / ``ln_f``, else any direct
+    child whose class name contains "Norm").  RMSNorm-family models (Llama,
+    Qwen, ...) need this applied before unembedding or logit-lens trajectories
+    distort.  Returns ``None`` when nothing can be discovered (lens callers
+    must degrade gracefully).
+    """
+    try:
+        _, layers_path = find_decoder_layers(model)
+    except Exception:
+        return None
+    parent = model if "." not in layers_path else resolve(model, layers_path.rsplit(".", 1)[0])
+    layers_attr = layers_path.rsplit(".", 1)[-1]
+    for attr in ("norm", "final_layernorm", "ln_f"):
+        mod = getattr(parent, attr, None)
+        if mod is not None and "norm" in type(mod).__name__.lower():
+            return mod
+    for name, mod in parent.named_children():
+        if name != layers_attr and "norm" in type(mod).__name__.lower():
+            return mod
+    return None
+
+
+def resolve(model: Any, dotted: str) -> Any:
+    """Resolve a dotted attribute path (``a.b.c``) against *model*, raising clearly."""
+    obj = model
+    for part in dotted.split("."):
+        if not hasattr(obj, part):
+            raise AttributeError(f"path {dotted!r} broke at {part!r}")
+        obj = getattr(obj, part)
+    return obj

evalrx/models/agent.py

@@ -0,0 +1,380 @@
+"""Agent — a backend-AGNOSTIC tool-calling loop over any handle.
+
+``Agent(wraps=handle)`` works identically on an API model and a local model: the
+loop only needs ``GENERATE`` + ``TOOL_CALLS`` (verified up front), never model
+internals.  The single backend/model-specific piece is the
+:class:`~evalrx.models.toolcodec.ToolCallCodec` (auto-selected).  Tool
+execution goes through a pluggable :class:`ToolExecutor` — swap in the existing
+``APIToolHandler`` for production (image handling, etc.).
+
+White-box backends additionally let an analyzer capture ONE step's internals via
+``handle.forward(...)``, but the trajectory production here is the same loop.
+Torch-free.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Iterable, Optional
+
+from evalrx.core.capability import Capability, CapabilityError
+from evalrx.core.case import (
+    FailureCase,
+    Inputs,
+    Label,
+    Step,
+    StepRole,
+    Trajectory,
+)
+from evalrx.core.tool import Tool, ToolCall, ToolResult
+from evalrx.models.toolcodec import ToolCallCodec, codec_for
+
+
+def _user_content(case: FailureCase) -> "str | list":
+    """Build the first user message's content: plain text, or content blocks
+    (transformers-style ``{"type": "image", "image": ...}``) when the case
+    carries an image — the multimodal chat convention every ``Model.chat``
+    implementation speaks (hf_local consumes it natively; API adapters convert
+    image blocks to their wire format).
+    """
+    image = case.inputs.image
+    if image is None:
+        return case.inputs.prompt
+    return [
+        {"type": "image", "image": image},
+        {"type": "text", "text": case.inputs.prompt},
+    ]
+
+
+def _observation_record(observation: Any) -> Any:
+    """What lands in ``Step.observation`` — structured for :class:`ToolResult`
+    (text + image count + host meta), the raw value otherwise."""
+    if isinstance(observation, ToolResult):
+        record: dict = {"text": observation.text, "n_images": len(observation.images)}
+        if observation.meta:
+            record["meta"] = observation.meta
+        return record
+    return observation
+
+
+class ToolExecutor:
+    """Run a decoded :class:`ToolCall` against the registered tools.
+
+    The simple default; production swaps in the existing ``APIToolHandler``
+    (which also handles tool-output images, retries, etc.).
+    """
+
+    def __init__(self, tools: Iterable[Tool]) -> None:
+        self._by_name = {t.name: t for t in tools}
+
+    def execute(self, call: ToolCall) -> Any:
+        tool = self._by_name.get(call.name)
+        if tool is None:
+            return f"[error: unknown tool {call.name!r}]"
+        if tool.fn is None:
+            return f"[error: tool {call.name!r} has no implementation]"
+        try:
+            return tool.fn(**(call.args or {}))
+        except Exception as exc:  # surface tool errors to the model, don't crash the loop
+            return f"[tool error in {call.name!r}: {exc}]"
+
+
+class APIToolHandlerExecutor:
+    """Bridge the Agent's tool execution to the XSkill engine's ``APIToolHandler``.
+
+    The handler is INJECTED (not imported here), keeping evalrx decoupled from
+    the engine repo.  It faithfully calls the real signature
+    ``execute_tool_call(tool_name, parameters, node, turn_idx, tool_call_id)`` and
+    returns the processed text; tool-output images accumulate in ``self.new_images``
+    for the caller to attach to the trajectory.
+
+    Usage::
+
+        from engine.api_tool_handler import APIToolHandler          # your repo
+        handler = APIToolHandler(args, save_dir)
+        agent = Agent(handle, tools, executor=APIToolHandlerExecutor(handler, node))
+
+    ``node`` is the engine's ``SearchNode`` (the handler mutates it, e.g. image_map);
+    pass the same node the engine drives so multi-turn image/state threading works.
+    """
+
+    def __init__(self, handler: Any, node: Any = None, *, result_key: str = "processed_result") -> None:
+        self.handler = handler
+        self.node = node
+        self.result_key = result_key
+        self._turn = 0
+        self.new_images: list = []
+        self.feedback_messages: list = []
+
+    def execute(self, call: ToolCall) -> Any:
+        self._turn += 1
+        out = self.handler.execute_tool_call(
+            call.name, call.args or {}, self.node, self._turn, getattr(call, "id", None)
+        )
+        if isinstance(out, dict):
+            self.new_images.extend(out.get("new_images") or [])
+            self.feedback_messages.extend(out.get("feedback_messages") or [])
+            return out.get(self.result_key) or out.get("tool_result") or ""
+        return out
+
+
+def _as_case(data: Any) -> FailureCase:
+    if isinstance(data, FailureCase):
+        return data
+    if isinstance(data, Inputs):
+        return FailureCase(inputs=data)
+    if isinstance(data, str):
+        return FailureCase.from_prompt(data)
+    raise TypeError(f"Agent.run expects str | Inputs | FailureCase, got {type(data).__name__}")
+
+
+# Loop-policy nudges (L2 scaffold messages, not stage prompts).
+_REPEAT_NUDGE = (
+    "[repeat blocked] You already made this exact tool call; the result would "
+    "be identical to what you already received. Choose DIFFERENT arguments or "
+    "give your final answer now."
+)
+_FINAL_NUDGE = (
+    "You have used your entire tool budget. Based on everything you have "
+    "already seen, give your final answer NOW as plain text."
+)
+
+
+def _call_signature(call: ToolCall) -> "tuple[str, str]":
+    import json
+
+    return call.name, json.dumps(call.args or {}, sort_keys=True, default=str)
+
+
+class Agent:
+    """A tool-calling agent composed over a model handle (any backend).
+
+    Loop-policy options (both default OFF — they are candidate L2 fixes, made
+    available as configuration so a validated repair is deployable):
+
+    * ``block_repeat_calls`` — a tool call identical to the immediately
+      preceding one is not executed; the model gets a nudge observation
+      instead (repeated identical calls return identical results and only
+      burn budget).
+    * ``force_final_answer`` — when ``max_turns`` runs out without a final
+      answer, ask once more WITHOUT tools so the run ends with an answer
+      instead of an empty trajectory (``terminated="forced_final"``).
+    """
+
+    requires = frozenset({Capability.GENERATE, Capability.TOOL_CALLS})
+
+    def __init__(
+        self,
+        handle,
+        tools: Iterable[Tool],
+        *,
+        codec: Optional[ToolCallCodec] = None,
+        executor: Optional[ToolExecutor] = None,
+        max_turns: int = 10,
+        system: Optional[str] = None,
+        block_repeat_calls: bool = False,
+        force_final_answer: bool = False,
+    ) -> None:
+        missing = self.requires - set(getattr(handle, "capabilities", frozenset()))
+        if missing:
+            raise CapabilityError(analyzer="Agent", model=repr(handle), missing=missing)
+        if not hasattr(handle, "chat"):
+            raise TypeError(
+                f"{type(handle).__name__} has no chat(); agent mode needs a tool-aware chat method."
+            )
+        self.handle = handle
+        self.tools = list(tools)
+        self.codec = codec or codec_for(handle)
+        self.executor = executor or ToolExecutor(self.tools)
+        self.max_turns = max_turns
+        self.system = system
+        self.block_repeat_calls = block_repeat_calls
+        self.force_final_answer = force_final_answer
+        # An Agent still exposes the underlying model's capabilities (pure-model
+        # analysis remains available on self.handle).
+        self.capabilities = handle.capabilities
+
+    def run(self, data: Any) -> Trajectory:
+        """Drive the tool loop to completion and return a :class:`Trajectory`.
+
+        Multimodal cases work end-to-end: ``case.inputs.image`` goes into the
+        first user message as a content block, and a :class:`ToolResult` whose
+        ``images`` are non-empty gets them re-injected as a follow-up user
+        message — so the model *sees* what its tool produced (zoom crops etc.).
+        """
+        case = _as_case(data)
+        goal = case.inputs.prompt
+        encoded = self.codec.encode(self.tools)
+
+        messages: list[dict] = []
+        if self.system:
+            messages.append({"role": "system", "content": self.system})
+        messages.append({"role": "user", "content": _user_content(case)})
+
+        steps: list[Step] = [
+            Step(
+                idx=0,
+                role=StepRole.USER,
+                content=goal,
+                span={"has_image": case.inputs.image is not None},
+            )
+        ]
+        final_answer: Optional[str] = None
+        terminated = "max_turns"
+        prev_sig: "Optional[tuple[str, str]]" = None
+
+        turn = 0
+        for turn in range(1, self.max_turns + 1):
+            t0 = time.perf_counter()
+            chat = self.handle.chat(messages, tools=encoded)
+            latency_ms = round((time.perf_counter() - t0) * 1000, 1)
+            call = self.codec.decode(chat)
+            span: dict = {"turn": turn, "latency_ms": latency_ms}
+            if chat.usage:
+                span.update(chat.usage)
+            steps.append(
+                Step(
+                    idx=len(steps),
+                    role=StepRole.ACTOR,
+                    content=chat.text,
+                    tool_call=call.to_dict() if call else None,
+                    span=span,
+                )
+            )
+            messages.append(self.codec.assistant_message(chat, call))
+
+            if call is None:
+                final_answer = self.codec.final_text(chat)  # strips <think> etc.
+                terminated = "final"
+                break
+
+            sig = _call_signature(call)
+            blocked = self.block_repeat_calls and sig == prev_sig
+            prev_sig = sig
+            observation: Any = _REPEAT_NUDGE if blocked else self.executor.execute(call)
+            tool_span: dict = {"turn": turn}
+            if blocked:
+                tool_span["repeat_blocked"] = True
+            steps.append(
+                Step(
+                    idx=len(steps),
+                    role=StepRole.TOOL,
+                    content=call.name,
+                    observation=_observation_record(observation),
+                    span=tool_span,
+                )
+            )
+            messages.append(self.codec.tool_message(call, str(observation)))
+            images = observation.images if isinstance(observation, ToolResult) else []
+            if images:
+                blocks: list = [{"type": "image", "image": im} for im in images]
+                blocks.append(
+                    {
+                        "type": "text",
+                        "text": f"The image{'s' if len(images) > 1 else ''} above "
+                        f"{'were' if len(images) > 1 else 'was'} returned by the "
+                        f"{call.name!r} tool call.",
+                    }
+                )
+                messages.append({"role": "user", "content": blocks})
+
+        if self.force_final_answer and final_answer is None:
+            messages.append({"role": "user", "content": _FINAL_NUDGE})
+            t0 = time.perf_counter()
+            chat = self.handle.chat(messages, tools=None)  # no tools: the reply IS the answer
+            span = {"forced_final": True, "latency_ms": round((time.perf_counter() - t0) * 1000, 1)}
+            if chat.usage:
+                span.update(chat.usage)
+            final_answer = self.codec.final_text(chat)
+            steps.append(Step(idx=len(steps), role=StepRole.ACTOR, content=chat.text, span=span))
+            terminated = "forced_final"
+
+        metrics: dict = {
+            "n_steps": len(steps),
+            "n_turns": turn,
+            "n_tool_calls": sum(1 for s in steps if s.tool_call),
+            "terminated": terminated,
+            "total_latency_ms": round(
+                sum(s.span.get("latency_ms", 0.0) for s in steps if s.role is StepRole.ACTOR), 1
+            ),
+        }
+        for key in ("prompt_tokens", "completion_tokens"):
+            total = sum(int(s.span.get(key) or 0) for s in steps if s.role is StepRole.ACTOR)
+            if total:
+                metrics[f"total_{key}"] = total
+        return Trajectory(
+            sample_id=case.id,
+            goal=goal,
+            steps=steps,
+            final_answer=final_answer,
+            ground_truth=case.expected,
+            outcome=Label.UNKNOWN,  # correctness is a separate analyzer's job
+            metrics=metrics,
+        )
+
+    def __repr__(self) -> str:
+        return f"Agent(handle={self.handle!r}, tools={[t.name for t in self.tools]}, codec={self.codec.name})"
+
+
+def run_batch(
+    handle,
+    cases: Iterable[Any],
+    *,
+    tools_factory: "Callable[[FailureCase], Iterable[Tool]]",
+    system: Optional[str] = None,
+    max_turns: int = 10,
+    codec: Optional[ToolCallCodec] = None,
+    concurrency: int = 4,
+    on_result: "Optional[Callable[[FailureCase, Trajectory], None]]" = None,
+    agent_kwargs: Optional[dict] = None,
+) -> list[Trajectory]:
+    """Run the tool loop over many cases — the M1-scale batch driver.
+
+    *tools_factory* builds the per-case tool set (visual tools bind to the
+    case's image, so each case needs its own instances).  Concurrency uses
+    threads and is only safe for API-backed handles (HTTP is reentrant); for
+    local backends the loop is forced sequential — one GPU forward at a time.
+
+    A case whose run raises (endpoint down, tool crash outside the executor
+    envelope) yields a stub trajectory with ``metrics["terminated"] == "error"``
+    instead of killing the batch; ``on_result`` fires per finished case (e.g.
+    incremental JSONL persistence).  Results keep the input order.
+    """
+    from concurrent.futures import ThreadPoolExecutor
+
+    from evalrx.models.backends.api import APIModel
+
+    normalized = [_as_case(c) for c in cases]
+    if concurrency > 1 and not isinstance(handle, APIModel):
+        import warnings
+
+        warnings.warn(
+            f"run_batch: {type(handle).__name__} is not an API handle; forcing "
+            "concurrency=1 (local backends are not thread-safe).",
+            stacklevel=2,
+        )
+        concurrency = 1
+
+    def _one(case: FailureCase) -> Trajectory:
+        try:
+            agent = Agent(
+                handle, tools_factory(case), codec=codec, max_turns=max_turns, system=system,
+                **(agent_kwargs or {}),
+            )
+            trajectory = agent.run(case)
+        except Exception as exc:  # keep the batch alive; the error IS the observation
+            trajectory = Trajectory(
+                sample_id=case.id,
+                goal=case.inputs.prompt,
+                steps=[Step(idx=0, role=StepRole.USER, content=case.inputs.prompt)],
+                ground_truth=case.expected,
+                metrics={"terminated": "error", "error": repr(exc)},
+            )
+        if on_result is not None:
+            on_result(case, trajectory)
+        return trajectory
+
+    if concurrency <= 1:
+        return [_one(c) for c in normalized]
+    with ThreadPoolExecutor(max_workers=concurrency) as pool:
+        return list(pool.map(_one, normalized))