diff --git a/verifiers/v1/harnesses/rlm/harness.py b/verifiers/v1/harnesses/rlm/harness.py index 6be4fb0412..83712dbaf1 100644 --- a/verifiers/v1/harnesses/rlm/harness.py +++ b/verifiers/v1/harnesses/rlm/harness.py @@ -1,20 +1,18 @@ """RLM over ACP, with MCP tools exposed as pre-imported IPython skills.""" -import json import logging import random import shlex from typing import Literal -from pydantic import Field, PositiveInt, model_validator +from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator -from verifiers.v1.acp import ACPConfig, ACPHarness +from verifiers.v1.acp import ACPConfig, ACPHarness, ACPTurnResult, JsonObject from verifiers.v1.clients import ModelContext from verifiers.v1.configs.harness import HarnessConfig from verifiers.v1.runtimes import Runtime from verifiers.v1.task import TaskData from verifiers.v1.trace import Trace -from verifiers.v1.utils.decorators import metric logger = logging.getLogger(__name__) @@ -25,15 +23,24 @@ RLM_BIN = f"{RLM_DIR}/bin/rlm" SKILLS_DIR = "/task/rlm-skills" RLM_STATE_DIR = ".vf-rlm" +RLM_RUNTIME_METADATA_KEY = "ai.prime.rlm/runtime-v1" +RLM_SESSION_METADATA_KEY = "ai.prime.rlm/session-v1" + + +class _SessionSnapshot(BaseModel): + model_config = ConfigDict(extra="ignore", strict=True) + + session_id: str = Field(pattern=r"^[A-Za-z0-9._:-]{1,128}$") + metrics: dict[str, int | float] class RLMHarnessConfig(HarnessConfig): version: str = Field( - default="83ef01f7a6c97328919387343bd30cf4edaac20d", min_length=1 + default="5ee1c34024a183bbbd3a38a6129995f5b982631d", min_length=1 ) """Git ref (branch, tag, or commit) of nano-rlm to install.""" max_depth: int = 0 - """Recursion depth rlm may spawn sub-harnesses to (RLM_MAX_DEPTH).""" + """Recursion depth RLM may spawn sub-harnesses to.""" builtin_skills: list[BuiltinSkill] = Field(default_factory=list) """Built-in rlm skills to enable (RLM_SKILLS), e.g. `["edit"]`; empty enables none. The tool set is fixed (ipython); the base `skills` field takes SKILL.md paths.""" @@ -54,7 +61,6 @@ def validate_range(self) -> "RLMHarnessConfig": @model_validator(mode="after") def reject_disabled_tools(self) -> "RLMHarnessConfig": - # rlm's only tool is ipython, which must stay enabled, so there's nothing to disable. if self.disabled_tools: raise ValueError( "the rlm harness has a fixed tool set (ipython) and does not support " @@ -91,41 +97,45 @@ async def setup(self, runtime: Runtime) -> None: raise RuntimeError(f"rlm install failed: {result.stderr.strip()[-500:]}") await super().setup(runtime) - def summarize_threshold(self, task_idx: int | None) -> str: - """The `RLM_SUMMARIZE_AT_TOKENS` value: a range draws per-group (seeded by task index — - 0 when unset — so a task's rollouts share one threshold). Always set — "" when disabled — - so the typed field, not a host var the subprocess runtime would inherit, wins.""" + def summarize_threshold(self, task_idx: int | None) -> int | None: + """Resolve a fixed or per-task compaction threshold.""" value = self.config.summarize_at_tokens if value is None: - return "" + return None if isinstance(value, tuple): lo, hi = value - return str(random.Random(task_idx or 0).randint(lo, hi)) - return str(value) + return random.Random(task_idx or 0).randint(lo, hi) + return value - def _env( + def _runtime_metadata( self, ctx: ModelContext, trace: Trace, + runtime: Runtime, endpoint: str, secret: str, data: TaskData, system_prompt: str | None, - ) -> dict[str, str]: - env = { - **self.config.resolved_env, - "RLM_BASE_URL": endpoint, - "RLM_API_KEY": secret, - "RLM_MODEL": ctx.model, - "RLM_MAX_DEPTH": str(self.config.max_depth), - "RLM_HOME": self._home(trace), - "RLM_SUMMARIZE_AT_TOKENS": self.summarize_threshold(data.idx), + ) -> JsonObject: + payload = { + "session_id": trace.id, + "model": ctx.model, + "provider": { + "base_url": endpoint, + "api_key": secret, + }, + "policy": { + "max_depth": self.config.max_depth, + "summarize_at_tokens": self.summarize_threshold(data.idx), + "max_concurrent_subagents": max(4, self.config.max_depth), + }, + "system_prompt_path": None, + "append_to_system_prompt": system_prompt, + "skills": list(self.config.builtin_skills), + "kernel_env": runtime.env, + "search_api_key": self.config.resolved_env.get("SERPER_API_KEY"), } - if system_prompt is not None: - env["RLM_APPEND_TO_SYSTEM_PROMPT"] = system_prompt - if self.config.builtin_skills: - env["RLM_SKILLS"] = ",".join(self.config.builtin_skills) - return env + return {RLM_RUNTIME_METADATA_KEY: payload} async def prepare_acp( self, @@ -139,29 +149,21 @@ async def prepare_acp( ) -> ACPConfig: system_prompt, prompt = self.resolve_prompt(data) return ACPConfig( - env=self._env(ctx, trace, endpoint, secret, data, system_prompt), + env={**self.config.resolved_env, "RLM_HOME": self._home(trace)}, command=[RLM_BIN, "--acp"], prompt=prompt, + session_meta=self._runtime_metadata( + ctx, trace, runtime, endpoint, secret, data, system_prompt + ), ) - @metric - async def rlm(self, trace: Trace, runtime: Runtime) -> dict[str, float]: - # RolloutRun closes the harness session before metrics, which finalizes - # RLM's meta.json while leaving the harness-owned state available here. - home = shlex.quote(self._home(trace)) - latest = f'cat "$(ls -t {home}/sessions/*/meta.json | head -1)"' - result = await runtime.run(["sh", "-c", latest], {}) - if result.exit_code != 0 or not result.stdout.strip(): - return {} - try: - meta = json.loads(result.stdout) - except json.JSONDecodeError: - return {} - return { - key: float(value) - for key, value in meta.get("metrics", {}).items() - if isinstance(value, (int, float)) and not isinstance(value, bool) - } + def acp_turn_result(self, trace: Trace, result: ACPTurnResult) -> None: + snapshot = _SessionSnapshot.model_validate( + result.response_metadata.get(RLM_SESSION_METADATA_KEY) + ) + if snapshot.session_id != trace.id: + raise ValueError("RLM session snapshot does not match the rollout") + trace.record_metrics(snapshot.metrics) async def cleanup(self, trace: Trace, runtime: Runtime) -> None: await runtime.run(["rm", "-rf", f"{RLM_STATE_DIR}/{trace.id}"], {})