diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index d2d8b7d34..ff7470082 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -505,6 +505,76 @@ jobs: pip install torch --index-url https://download.pytorch.org/whl/cpu pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow + + - name: Install + shell: bash + run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps + + + - name: Execute + shell: bash + run: | + + TEST_PATH="${{ matrix.info.test_file }}" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + TEST_ARGS="${{ matrix.info.test_args || '' }}" + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "${{ matrix.info.num_gpus }}" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + + + agent-adapter-test: + needs: pre-commit + + if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' + + + runs-on: ubuntu-latest + + strategy: + fail-fast: false + matrix: + info: [{"num_gpus": 0, "test_file": "test_agent_trajectory.py"}, {"num_gpus": 0, "test_file": "test_agent_adapters.py"}, {"num_gpus": 0, "test_file": "test_agent_sdk_adapters.py"}] + defaults: + run: + working-directory: ${{ github.workspace }} + env: + GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} + SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} + SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} + SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + cache: 'pip' + + - name: Install dependencies + shell: bash + run: | + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow + + pip install openai openai-agents anthropic + + - name: Install shell: bash run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 54a37d279..37d4b79cc 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -69,6 +69,18 @@ ], }, + 'agent-adapter-test': { + 'label': 'run-ci-agent-adapter', + 'always': True, + 'cpu': True, + 'extra_pip_deps': 'openai openai-agents anthropic', + 'tests': [ + {'test_file': 'test_agent_trajectory.py', 'num_gpus': 0}, + {'test_file': 'test_agent_adapters.py', 'num_gpus': 0}, + {'test_file': 'test_agent_sdk_adapters.py', 'num_gpus': 0}, + ], + }, + 'e2e-test-image': { 'label': 'run-ci-image', 'image': 'inferactinc/public:vime-vllm-cu129-latest', @@ -175,6 +187,9 @@ jobs: run: | pip install torch --index-url https://download.pytorch.org/whl/cpu pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow +<% if config.get('extra_pip_deps') %> + pip install << config.extra_pip_deps >> +<% endif %> - name: Install shell: bash diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md new file mode 100644 index 000000000..61c81e036 --- /dev/null +++ b/examples/coding_agent_rl/README.md @@ -0,0 +1,179 @@ +# Coding-Agent RL + +This directory provides an example of running end-to-end **SWE (Software-Engineering) coding-agent RL** with vime: a real coding agent (claude-code CLI) drives `Read/Edit/Grep/Bash/Agent` tools inside a fresh sandbox per sample, the model produces a `git diff`, and the diff is graded against the dataset's test harness in a second clean sandbox (no test-cheating). + +Two example files and one shared adapter implement the loop: + +- `generate.py` — per-sample `generate()` registered via `--custom-generate-function-path`. Boots the sandbox, runs claude-code, captures the diff, scores it, and emits one or more `Sample`s back to vime. +- `vime.agent.adapters.AnthropicAdapter` — the shared Anthropic Messages adapter. claude-code talks to it as if it were Anthropic; the adapter tokenizes the current message history each turn, records prompt/output token snapshots, preserves model-generated tokens (`loss_mask=1`) only while later prompts stitch onto them, masks template/observation tokens (`0`), and emits **three kinds of segments** per trajectory: `subagent` (completed `Task/Agent` dispatch), `wipe` (chain frozen by auto-compact), `final` (tail of the main chain). +- `sandbox.py` — coding-agent/SWE helpers built on `vime.agent.sandbox`: install bootstraps, spawn claude-code, capture patches, and run the fresh-sandbox evaluator. The shared sandbox contract lives in `vime.agent.sandbox.Sandbox`. + +`generate.py` owns one `AnthropicAdapter` instance. For each sample it calls +`adapter.open_session(...)` before starting claude-code, serves `adapter.app` as +the Anthropic-compatible endpoint, and drains trainable `TokenSegment`s with +`await adapter.finish_session(...)` when the trajectory ends. + +## Environment Setup + +The vime training stack itself follows the standard setup. On top of that you need: + +1. **An E2B-compatible sandbox cluster** (or any provider that speaks the E2B SDK). Configure via `E2B_API_KEY` (e.g. the standard `e2b_xxx` key from https://e2b.dev, or any internal endpoint that accepts the same SDK). The official SDK validates this value locally, so internal gateways that ignore auth still need a syntactically valid `e2b_` + 40 hex-character placeholder. +2. **Host-side tarballs** that get uploaded into each sandbox at boot: + - Node 22 (`node-v22.x-linux-x64.tar.xz`) — exported as `SWE_HOST_NODE_TARBALL`. + - Claude Code CLI npm tarball (`anthropic-ai-claude-code-local-linux-x64.tgz`) — exported as `SWE_HOST_CC_TARBALL`. +3. **A sandbox metadata file** (`SWE_SANDBOX_METADATA_FILE`, or the generic `VIME_AGENT_SANDBOX_METADATA_FILE`) — JSON dict whose keys are passed as routing tags when booting an E2B sandbox. Must contain the image key referenced by `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` (e.g. `image`). +4. **Network reachability**: each sandbox dials back to the vime head node's Anthropic adapter over `http://${VIME_HEAD_HOST}:${SHIM_PORT}`. The head host must be reachable from inside the sandboxes (set `VIME_HEAD_HOST` to a routable IP, not `127.0.0.1`). + +## Dataset Format + +Standard vime JSONL with three keys: + +```jsonc +{ + "prompt": "", + "label": "", + "metadata": { + "image": "swedev/scaleswe.oh.34:", // sandbox image reference + "workdir": "/workspace/", // repo path inside the sandbox + "problem_statement": "", + // exactly one of the following two graders: + "swepro": { /* SWE-bench Pro test harness — preferred */ }, + "eval_cmd": "pytest -x tests/..." // last-resort: exit 0 = solved + // sweb-style rows: metadata.remote_env_info.f2p_script (Python file + // ending in `sys.exit(pytest.main(...))`) is auto-wrapped into eval_cmd. + } +} +``` + +Wire it up with `--input-key prompt --label-key label --metadata-key metadata`. + +## Running the Script + +Override the paths at the top of the launcher, then run from a long-lived shell on the Ray head node (do **not** wrap in `nohup` — Ray child processes get cleaned up with it): + +```bash +cd vime/ + +export HF_CHECKPOINT=/path/to/Qwen3.6-35B-A3B +export REF_MODEL_PATH=/path/to/Qwen3.6-35B-A3B_torch_dist +export PROMPT_DATA=/path/to/swe_train.jsonl +export SANDBOX_METADATA_FILE=/path/to/sandbox_metadata.json +export SWE_HOST_NODE_TARBALL=/path/to/node-v22.20.0-linux-x64.tar.xz +export SWE_HOST_CC_TARBALL=/path/to/anthropic-ai-claude-code-local-linux-x64.tgz + +bash examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh +``` + +The launcher brings up Ray across all hosts in `/root/mpi_rack_hostfile`, dumps every rollout to `runs/${EXP_TAG}_${STAMP}/rollout_dumps/`, and tees stdout into `runs/${EXP_TAG}_${STAMP}/run.log`. + +## New Arguments + +`generate.py` is wired in through vime's standard custom-generate hook: + +```bash +ROLLOUT_ARGS=( + --custom-generate-function-path examples.coding_agent_rl.generate.generate + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --metadata-key metadata + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-context-len 96000 + --rollout-max-response-len 32768 + --rollout-stop-token-ids 248046 248044 + --save-debug-rollout-data "${RUN_ROOT}/rollout_dumps/rollout_{rollout_id}.pt" +) +``` + +claude-code's tool invocations are parsed from the model output. By default the +adapter uses the built-in XML tool-call fallback (`parse_xml_tool_uses`), which +handles the `` text Qwen3-Coder emits, so no extra +engine-side parser configuration is required. (The adapter can optionally +delegate to SGLang's reasoning/function-call parsers when +`vllm_tool_call_parser` / `vllm_reasoning_parser` are set on `args`, but this is +not used by this example and is not wired into vime's vLLM arguments.) + +## SWE-specific Environment Knobs + +All set in the launcher; tune per cluster. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `VIME_HEAD_HOST` | `${MASTER_ADDR}` | Public IP the sandbox uses to reach the Anthropic adapter. **Must be routable from inside the sandbox.** | +| `SHIM_BIND_HOST` / `SHIM_PORT` | `0.0.0.0` / `18001` | Bind address of the adapter shim on the head node. | +| `E2B_API_KEY` | — | E2B (or compatible) API key. | +| `SWE_SANDBOX_METADATA_FILE` / `VIME_AGENT_SANDBOX_METADATA_FILE` | — | JSON dict of routing metadata passed at sandbox boot. | +| `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` | — | Which key in the metadata file holds the image reference (e.g. `image`). | +| `SWE_HOST_NODE_TARBALL` | — | Host path to Node 22 tarball uploaded into each sandbox. | +| `SWE_HOST_CC_TARBALL` | — | Host path to the Claude Code CLI npm tarball. | +| `SWE_TIME_BUDGET_SEC` | `1800` | Wallclock budget for one agent run. | +| `SWE_EVAL_TIMEOUT_SEC` | `600` | Wallclock cap on the evaluator sandbox. | +| `SWE_BOOT_CONCURRENCY` | `6` | Cap on simultaneous sandbox boots (eases h2/SSL long-tail). | +| `SWE_CLAUDE_EXTRA_ARGS` | (see launcher) | Extra flags appended to the `claude` CLI invocation — registers the read-only `investigator` sub-agent, disables `WebFetch`/`WebSearch`, disables slash commands. | +| `SWE_CC_PROMPT` | unset | Optional override for the user-turn prompt. Setting this to require sub-agent dispatch is the most reliable way to maximize fan-out. | + +`--rollout-max-response-len` is the per-turn generation cap passed to each vLLM +`/inference/v1/generate` call as the sampling-params `max_tokens`. +`--rollout-max-context-len` is the multi-turn prompt+response budget enforced +only during generation: each turn clamps the generation length to the remaining +context. Trajectory merge/export keeps the emitted segments and does not drop +them for length. + +## String-in, Token-out Trajectories + +The coding-agent environment is string/message based: claude-code sends +Anthropic Messages requests, receives streamed text/thinking/tool-use blocks, +and later sends back rendered tool observations. Training, however, must stay +token based. A trajectory is only a valid RL target when the optimized tokens +are the same tokens the rollout model actually sampled. + +The Anthropic adapter therefore follows a **string in, token out** contract: + +- Each incoming message history is rendered with the served model's chat + template and sent to vLLM as `token_ids`. +- vLLM's `/inference/v1/generate` is called with `logprobs` set; the adapter + records the exact `prompt_ids`, sampled `output_ids` (`choices[0].token_ids`), + and per-token rollout logprobs (`choices[0].logprobs.content[i].logprob`) for + that turn. +- At training export time, samples are assembled from those saved token ids. + The decoded `response` field is only a readable sidecar; it is not + re-tokenized to recover the training sequence. + +Multi-turn agents still force the adapter to tokenize later message +histories, because tool observations and claude-code's own compacted messages +arrive as strings. `vime.agent.trajectory.merge_turns` stitches those later +prompts against the saved token stream: + +- New prompt suffixes that are tool/user/environment context are appended with + `loss_mask=0`. +- Fresh model outputs from vLLM are appended with `loss_mask=1`. +- If a later prompt no longer token-matches an earlier sampled output, the + unmatched suffix is dropped. If the drift cuts through the middle of a + previous model output, the retained prefix of that whole output turn is also + assigned `loss_mask=0`. + +That last case is the important correctness guard. A re-tokenization mismatch +can make a string-level conversation look continuous while token-level +provenance is broken. vime keeps the context needed to continue the agent, but +does not backprop through tokens whose sampled origin can no longer be proven. + +## Fan-out Semantics + +- `generate()` returns `list[Sample]` — one Sample per trajectory **segment** (`subagent` / `wipe` / `final`). +- Per-trajectory reward is split as `reward / K` across segments; `rollout_id` is shared so the per-rollout-mean loss reducer still counts the trajectory once. +- Sub-agent dispatch increases `K` (each completed `Agent` turn block becomes its own segment), so the effective batch after flatten can be much larger than `rollout_batch_size * n_samples_per_prompt`. + +## Porting to a New Sandbox Backend + +`vime.agent.sandbox.Sandbox` exposes the shared sandbox contract, and +`vime.agent.sandbox.E2BSandbox` is the E2B implementation: + +```python +await sb.exec(cmd, user=..., check=..., timeout=...) +await sb.write_file(sandbox_path, content_or_host_path, user=...) +await sb.read_file(sandbox_path, user=...) +async with E2BSandbox(...) as sb: ... +``` + +Reimplement those on Docker / Modal / a local VM and everything in `generate.py` keeps working unchanged. diff --git a/examples/coding_agent_rl/aiohttp_threaded.py b/examples/coding_agent_rl/aiohttp_threaded.py new file mode 100644 index 000000000..a5a17652d --- /dev/null +++ b/examples/coding_agent_rl/aiohttp_threaded.py @@ -0,0 +1,90 @@ +"""Run an ``aiohttp.web.Application`` in a background daemon thread.""" + +from __future__ import annotations + +import asyncio +import threading +from dataclasses import dataclass +from typing import Any + +from aiohttp import web + + +@dataclass +class AppHandle: + host: str + port: int + thread: threading.Thread + loop: asyncio.AbstractEventLoop + runner: web.AppRunner + + @property + def url(self) -> str: + return f"http://{self.host}:{self.port}" + + def stop(self) -> None: + async def _shutdown() -> None: + await self.runner.cleanup() + + try: + fut = asyncio.run_coroutine_threadsafe(_shutdown(), self.loop) + fut.result(timeout=10) + except Exception: + pass + self.loop.call_soon_threadsafe(self.loop.stop) + self.thread.join(timeout=5) + + +def run_app_in_thread( + app: web.Application, + *, + host: str = "0.0.0.0", + port: int = 0, + thread_name: str = "aiohttp-app", + start_timeout_sec: float = 15.0, + runner_kwargs: dict[str, Any] | None = None, +) -> AppHandle: + """Spin up ``app`` on a daemon thread; block until it is listening. + + ``runner_kwargs`` is forwarded to ``web.AppRunner`` (e.g. pass + ``{"handler_cancellation": True}`` to make a client disconnect cancel + the in-flight handler coroutine). + """ + started = threading.Event() + err_box: list[BaseException] = [] + box: dict[str, Any] = {} + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + runner = web.AppRunner(app, **(runner_kwargs or {})) + loop.run_until_complete(runner.setup()) + site = web.TCPSite(runner, host, port) + loop.run_until_complete(site.start()) + actual_port = port + for sock in site._server.sockets: # type: ignore[attr-defined] + actual_port = sock.getsockname()[1] + break + box["loop"] = loop + box["runner"] = runner + box["port"] = actual_port + started.set() + loop.run_forever() + except BaseException as e: # pragma: no cover + err_box.append(e) + started.set() + raise + + thread = threading.Thread(target=_run, name=thread_name, daemon=True) + thread.start() + started.wait(timeout=start_timeout_sec) + if err_box: + raise err_box[0] + return AppHandle( + host=host, + port=int(box["port"]), + thread=thread, + loop=box["loop"], + runner=box["runner"], + ) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py new file mode 100644 index 000000000..b048cff52 --- /dev/null +++ b/examples/coding_agent_rl/generate.py @@ -0,0 +1,364 @@ +"""Coding-Agent RL: per-sample generate() function for vime. + +Wire-up: + + --custom-generate-function-path examples.coding_agent_rl.generate.generate + +``generate()`` is intentionally a small four-stage orchestrator: + + 1. ``sandbox.run_claude_code`` prepares the agent sandbox and runs claude-code. + 2. ``sandbox.git_diff`` captures the model-produced patch. + 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. + 4. ``_merge_samples`` combines reward + adapter ``TokenSegment``s, + delegating segment-to-``Sample`` fan-out to ``vime.agent.trajectory``. + +All sandbox-side details live in ``sandbox.py``; the LLM plumbing +(Anthropic <-> vLLM ``/inference/v1/generate``, token capture, 3-kind segment +split) uses ``vime.agent.adapters.AnthropicAdapter``. + +Dataset row ``metadata`` schema:: + + image: str # sandbox image + workdir: str # repo path inside the sandbox + problem_statement: str # issue body (falls back to sample.prompt) + swepro: dict|None # SWE-bench Pro test harness (preferred) + eval_cmd: str|None # last-resort: shell command (exit 0 = solved) + +Also accepted (sweb-style rows): ``metadata.remote_env_info.f2p_script`` — +a self-contained Python test file ending in ``sys.exit(pytest.main(...))``. +When ``eval_cmd`` is absent, ``_metadata`` wraps this script into a base64 +materialize-and-run shell command so the existing eval path stays unchanged. + +Env knobs (set in run.sh): + + SWE_HOST_NODE_TARBALL host path to a Node 22 tarball (REQUIRED) + SWE_HOST_CC_TARBALL host path to the Claude Code npm tarball (REQUIRED) + SWE_TIME_BUDGET_SEC 1800 per agent run, wallclock + SWE_EVAL_TIMEOUT_SEC 600 per eval test execution + SHIM_BIND_HOST 0.0.0.0 + SHIM_PORT 18001 + VIME_HEAD_HOST public host the sandboxes use to reach the adapter (REQUIRED) +""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +import os +import secrets +import time +import traceback +from dataclasses import dataclass +from typing import Any + +from vime.agent.adapters import AnthropicAdapter +from vime.agent.trajectory import TokenSegment, fan_out_sample_segments +from vime.utils.misc import SingletonMeta +from vime.utils.processing_utils import load_tokenizer +from vime.utils.types import Sample + +from . import sandbox +from .aiohttp_threaded import run_app_in_thread + +logger = logging.getLogger(__name__) + + +SWE_TIME_BUDGET_SEC = int(os.environ.get("SWE_TIME_BUDGET_SEC", "1800")) +SWE_EVAL_TIMEOUT_SEC = int(os.environ.get("SWE_EVAL_TIMEOUT_SEC", "600")) +# Wall-clock guard for the entire generate() call. Defaults to +# SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 (buffer for sandbox boot, +# diff capture, etc). When exceeded, the in-flight sample is aborted with +# reason `wall_clock_timeout` and the rest of the rollout continues -- this +# isolates a single hung trajectory (e.g. stuck in sandbox.evaluate) so it +# does not kill the whole training step. +SWE_GENERATE_GUARD_SEC = int(os.environ.get("SWE_GENERATE_GUARD_SEC", "0") or 0) or ( + SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 +) +SHIM_BIND_HOST = os.environ.get("SHIM_BIND_HOST", "0.0.0.0") +SHIM_PORT = int(os.environ.get("SHIM_PORT", "18001")) + + +# --------------------------------------------------------------------------- +# Singleton: tokenizer + in-process Anthropic adapter + reducer +# --------------------------------------------------------------------------- +class _State(metaclass=SingletonMeta): + def __init__(self, args) -> None: + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + self.max_context_len = int(getattr(args, "rollout_max_context_len", 0) or 0) + self.tool_parser = getattr(args, "vllm_tool_call_parser", None) or None + self.reasoning_parser = getattr(args, "vllm_reasoning_parser", None) or None + vllm_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + public_host = os.environ.get("VIME_HEAD_HOST") + if not public_host: + raise RuntimeError( + "VIME_HEAD_HOST is not set. Export it to the host IP that " + "sandboxes can reach for reverse-connection to the Anthropic adapter. " + "Without it the sandbox cannot dial back and the rollout will " + "silently abort." + ) + self.adapter = AnthropicAdapter( + tokenizer=self.tokenizer, + vllm_url=vllm_url, + model=args.hf_checkpoint, + tool_parser=self.tool_parser, + reasoning_parser=self.reasoning_parser, + ) + # handler_cancellation=True so a client disconnect cancels the handler + # coroutine, tearing down the in-flight engine ``/inference/v1/generate`` + # request. Without it a cancelled client leaves an inflight generate that + # races with the next release_memory_occupation. + self.app_handle = run_app_in_thread( + self.adapter.app, + host=SHIM_BIND_HOST, + port=SHIM_PORT, + thread_name="anthropic-adapter", + runner_kwargs={"handler_cancellation": True}, + ) + self.adapter_url = f"http://{public_host}:{self.app_handle.port}" + logger.info( + "[coding_agent_rl] tokenizer=%s adapter=%s max_context_len=%s tool_parser=%s reasoning_parser=%s", + args.hf_checkpoint, + self.adapter_url, + self.max_context_len, + self.tool_parser, + self.reasoning_parser, + ) + + +# --------------------------------------------------------------------------- +# Trajectory -> Sample conversion +# adapter.finish_session() returns TokenSegments. One trajectory yields >=1 +# segments because the agent may compact + reset mid-run; trajectory.py handles +# the mechanical segment -> Sample fan-out. +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class RewardResult: + reward: float + is_solved: bool + applied_cleanly: bool + + +def _start_session( + state: _State, + sample: Sample, + md: dict[str, Any], + sampling_params: dict[str, Any], +) -> str: + # claude-code inside the sandbox dials back to the adapter with this + # session_id (passed as the Bearer token) so its turns are grouped under + # one chain history. Build from (instance_id, index, group_index) when + # possible; fall back to random hex if either index is missing. + if sample.session_id: + session_id = sample.session_id + elif sample.index is not None and sample.group_index is not None: + session_id = f"cagent-{md['instance_id']}-{sample.index}-{sample.group_index}" + else: + session_id = f"cagent-{md['instance_id']}-{secrets.token_hex(8)}" + sample.session_id = session_id + state.adapter.open_session( + session_id, + sampling_defaults=sampling_params, + max_context_tokens=state.max_context_len, + ) + return session_id + + +def _merge_samples( + *, + sample: Sample, + state: _State, + segments: list[TokenSegment], + reward_result: RewardResult, + elapsed_sec: float, + instance_id: str, +): + if not segments: + return _abort_result(sample, "adapter_session_empty") + + trajectory_metadata = { + **(sample.metadata or {}), + "instance_id": instance_id, + "is_solved": reward_result.is_solved, + "applied_cleanly": reward_result.applied_cleanly, + "elapsed_sec": elapsed_sec, + } + + # All K samples share rollout_id so the loss reducer counts this + # trajectory once. + fanned = fan_out_sample_segments( + sample, + segments, + reward_result.reward, + state.tokenizer, + metadata=trajectory_metadata, + ) + if not fanned: + raise ValueError("fan-out produced no samples") + + logger.info( + "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", + instance_id, + reward_result.reward, + reward_result.is_solved, + reward_result.applied_cleanly, + elapsed_sec, + len(fanned), + ) + return fanned + + +# --------------------------------------------------------------------------- +# Main per-sample agent function +# +# The four calls inside the timeout are the high-level rollout recipe: +# run_claude_code -> git_diff -> sandbox.evaluate -> merge_samples. +# --------------------------------------------------------------------------- +async def generate(args, sample: Sample, sampling_params: dict[str, Any]): + """Per-sample agent function with wall-clock guard. See + SWE_GENERATE_GUARD_SEC docstring above.""" + state = _State(args) + md = _metadata(sample) + if not md["image"] or not md["workdir"]: + return _abort_result(sample, "missing_image_or_workdir") + + instance_id = md["instance_id"] + session_id = _start_session(state, sample, md, sampling_params) + t0 = time.time() + try: + async with asyncio.timeout(SWE_GENERATE_GUARD_SEC): + async with sandbox.boot_agent_sandbox(md["image"]) as sb: + await sandbox.run_claude_code( + sb, + workdir=md["workdir"], + session_id=session_id, + adapter_url=state.adapter_url, + time_budget_sec=SWE_TIME_BUDGET_SEC, + problem_statement=md["problem_statement"], + swepro=md["swepro"], + pre_commands=md["pre_commands"], + ) + diff_text = await sandbox.git_diff(sb, md["workdir"]) + + reward, is_solved, applied_cleanly = await sandbox.evaluate( + image=md["image"], + workdir=md["workdir"], + diff_text=diff_text, + swepro=md["swepro"], + eval_cmd=md["eval_cmd"], + pre_commands=md["pre_commands"], + timeout_sec=SWE_EVAL_TIMEOUT_SEC, + ) + reward_result = RewardResult( + reward=float(reward), + is_solved=bool(is_solved), + applied_cleanly=bool(applied_cleanly), + ) + segments = await state.adapter.finish_session(session_id) + return _merge_samples( + sample=sample, + state=state, + segments=segments, + reward_result=reward_result, + elapsed_sec=time.time() - t0, + instance_id=instance_id, + ) + + except asyncio.TimeoutError: + _log_timeout_diagnostic(t0) + return _abort_result(sample, "wall_clock_timeout") + except Exception as e: + logger.error( + "[coding_agent_rl] %s: rollout failed: %s\n%s", + instance_id, + e, + traceback.format_exc(), + ) + return _abort_result(sample, f"exception:{type(e).__name__}") + finally: + # Close the sid before next train step's release_memory_occupation; + # stragglers from this trajectory would otherwise race its idle assert. + await state.adapter.finish_session(session_id) # idempotent + + +def _log_timeout_diagnostic(t0: float) -> None: + """Dump pending-task names when the wall-clock guard fires so future + debugging can see which await was stuck. Must never crash.""" + try: + elapsed = time.time() - t0 + pending = [t for t in asyncio.all_tasks() if not t.done()] + stuck = [] + for t in pending[:5]: # cap to avoid log spam + coro = getattr(t, "_coro", None) + stuck.append(getattr(coro, "__qualname__", repr(coro))) + logger.warning( + "[coding_agent_rl] generate() wall_clock_timeout after %.1fs " + "(guard=%ds); %d tasks pending; sample of stuck: %s", + elapsed, + SWE_GENERATE_GUARD_SEC, + len(pending), + stuck, + ) + except Exception: # pragma: no cover - diag must never crash + pass + + +# --------------------------------------------------------------------------- +# Metadata helpers +# --------------------------------------------------------------------------- +def _wrap_f2p_script(script: str | None) -> str | None: + # Materialize a self-contained pytest script (typical sweb f2p_script: + # ends with `sys.exit(pytest.main([...]))`) into the sandbox via base64 + # so we sidestep all shell quoting; python's exit code carries the + # pytest pass/fail signal that `_run_eval_cmd` turns into reward. + if not script: + return None + b64 = base64.b64encode(script.encode("utf-8")).decode("ascii") + return f"echo {b64} | base64 -d > /tmp/vime_f2p.py && python /tmp/vime_f2p.py" + + +def _metadata(sample: Sample) -> dict[str, Any]: + """Normalize the two dataset schemas (flat vs ``remote_env_info``).""" + m = sample.metadata or {} + rem = m.get("remote_env_info") or {} + label = sample.label if (isinstance(sample.label, str) and len(sample.label) < 256) else None + return { + "instance_id": m.get("instance_id") or rem.get("instance_id") or label or "unknown", + "image": m.get("image") or rem.get("image_url"), + "workdir": m.get("workdir") or rem.get("workdir"), + "problem_statement": m.get("problem_statement") or _coerce_prompt(sample.prompt), + "swepro": m.get("swepro"), + "eval_cmd": m.get("eval_cmd") or _wrap_f2p_script(rem.get("f2p_script")), + "pre_commands": m.get("pre_commands") or rem.get("pre_commands"), + } + + +def _coerce_prompt(prompt) -> str: + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list): + for m in prompt: + if isinstance(m, dict) and m.get("role") == "user": + c = m.get("content") + if isinstance(c, str): + return c + if isinstance(c, list): + return "\n".join(p.get("text", "") for p in c if isinstance(p, dict) and p.get("type") == "text") + return "" + + +def _abort(sample: Sample, reason: str) -> Sample: + sample.tokens = [0, 0] + sample.response = "" + sample.response_length = 1 + sample.loss_mask = [0] + sample.reward = 0.0 + sample.status = Sample.Status.ABORTED + sample.metadata = {**(sample.metadata or {}), "abort_reason": reason} + logger.warning("[coding_agent_rl] aborted: %s", reason) + return sample + + +def _abort_result(sample: Sample, reason: str): + """Return a uniform list shape for this fan-out generate function.""" + return [_abort(sample, reason)] diff --git a/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh new file mode 100644 index 000000000..ef14087c0 --- /dev/null +++ b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh @@ -0,0 +1,374 @@ +#!/usr/bin/env bash +# End-to-end SWE coding-agent RL on 8 nodes. +# +# Same model and training loop as run_qwen36_35b_a3b_swe_8node.sh, with three +# extra layers that actively encourage the rollout to dispatch sub-agents. +# Trajectory trees produced by this script show real `sibling` branches: +# +# (1) An `investigator` sub-agent is registered via claude-code's --agents +# flag (Grep/Read/Glob only, no edits) — a concrete, narrowly-scoped +# dispatch target. +# (2) SWE_CC_PROMPT requires the model to dispatch the investigator before +# any edit, naming the exact call form (Agent tool with +# subagent_type=investigator). +# (3) Agent/Task tools stay in the allowed set; WebFetch/WebSearch are +# disabled (sandbox has no outbound internet); --disable-slash-commands +# removes /compact as a competing branching pathway. +# +# Fan-out semantics: +# * generate() returns list[Sample] (one Sample per trajectory segment); +# the per-trajectory reward is split as reward/K across segments. +# * Sub-agent dispatch increases K (each sub-agent turn block becomes its +# own segment), so the effective batch after flatten can be much larger +# than rollout_batch_size * n_samples_per_prompt. If pinned-memory or +# GPU wake_up OOM appears, lower rollout_batch_size or n_samples_per_prompt +# first — not max-tokens-per-gpu. +# Run from a long-lived shell / tmux session on the Ray head node; do not wrap +# in a short-lived nohup launcher or Ray child processes get cleaned up with it. + +# Best-effort cleanup so a rerun does not collide with stale workers. +pkill -9 -f "vllm serve" || true +sleep 3 +ray stop --force || true +pkill -9 ray || true +sleep 3 +pkill -9 ray || true + +set -ex + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_DIR="${VIME_DIR:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" + +# ============ model parallelism ============ +# CP=8 (higher than the baseline script's CP=2) gives more rank-local context +# room for the longer per-segment payloads typical under sub-agent dispatch. +export TP_SIZE="${TP_SIZE:-2}" +export PP_SIZE="${PP_SIZE:-1}" +export CP_SIZE="${CP_SIZE:-8}" +export EP_SIZE="${EP_SIZE:-8}" +export ETP_SIZE="${ETP_SIZE:-1}" + +# ============ rollout engine ============ +ROLLOUT_TP_SIZE="${ROLLOUT_TP_SIZE:-8}" +ROLLOUT_DP_SIZE="${ROLLOUT_DP_SIZE:-8}" +ROLLOUT_MEM_UTILIZATION="${ROLLOUT_MEM_UTILIZATION:-0.75}" + +# ============ Qwen3.5-35B-A3B architecture ============ +NLAYERS=40 +FIRST_K_DENSE_REPLACE=0 + +arr=() +for ((i=0; ivLLM arg map (see translation_guide.md "Arg map"): +# --sglang-mem-fraction-static -> --vllm-gpu-memory-utilization +# --sglang-dp-size N -> --vllm-data-parallel-size N +# --sglang-ep-size N -> --vllm-enable-expert-parallel (boolean) +# --sglang-{enable-dp-attention,enable-dp-lm-head,moe-dense-tp-size, +# mamba-scheduler-strategy} -> sglang-only, no vLLM equivalent (dropped) +# --sglang-{tool-call,reasoning}-parser -> no vLLM arg; coding_agent_rl uses the +# XML tool-call fallback (parsers=None), so these are dropped. +# --sglang-speculative-* per-field flags -> a single --vllm-speculative-config +# JSON dict (commented below; supply your EAGLE draft model to enable). +VLLM_ARGS=( + --rollout-num-gpus 64 + --rollout-num-gpus-per-engine ${ROLLOUT_TP_SIZE} + --vllm-gpu-memory-utilization ${ROLLOUT_MEM_UTILIZATION} + --vllm-data-parallel-size ${ROLLOUT_DP_SIZE} + --vllm-enable-expert-parallel + --prefill-num-servers 1 +) +# Speculative decoding (EAGLE): vLLM takes one JSON dict rather than the +# sglang per-field flags. Uncomment and point "model" at your draft checkpoint. +# VLLM_ARGS+=( +# --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":4,"model":"/path/to/eagle_draft"}' +# ) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --colocate +) + +# ============ ray cluster network ============ +# Set MASTER_ADDR before the SWE block: VIME_HEAD_HOST below falls back to it. +export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" +export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" + +# ============ SWE / claude-code rollout knobs ============ + +# --- sandbox provisioning (E2B) --- +# The E2B SDK validates the API key format locally before it reaches +# E2B-compatible gateways. If your internal gateway ignores auth, use any +# syntactically valid e2b_... hex placeholder. +E2B_DUMMY_API_KEY="e2b_0000000000000000000000000000000000000000" +E2B_API_KEY="${E2B_API_KEY:-${E2B_DUMMY_API_KEY}}" +if [[ ! "${E2B_API_KEY}" =~ ^e2b_[0-9a-fA-F]{40}$ ]]; then + echo "WARN: E2B_API_KEY does not pass local E2B SDK format validation; using dummy key." >&2 + E2B_API_KEY="${E2B_DUMMY_API_KEY}" +fi +export E2B_API_KEY +export SWE_SANDBOX_METADATA_FILE="${SANDBOX_METADATA_FILE}" +export SWE_SANDBOX_IMAGE_METADATA_KEY="${SWE_SANDBOX_IMAGE_METADATA_KEY:-glm-platform/image}" +# Host-side tarballs injected into each sandbox at boot. +export SWE_HOST_NODE_TARBALL="${SWE_HOST_NODE_TARBALL:-/path/to/node-v22.x-linux-x64.tar.xz}" +export SWE_HOST_CC_TARBALL="${SWE_HOST_CC_TARBALL:-/path/to/anthropic-ai-claude-code-local-linux-x64.tgz}" + +# --- reply path (sandbox -> host shim) --- +export VIME_HEAD_HOST="${VIME_HEAD_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export SHIM_BIND_HOST="${SHIM_BIND_HOST:-0.0.0.0}" +export SHIM_PORT="${SHIM_PORT:-18001}" + +# --- per-trajectory time / concurrency budgets --- +# Time budget 1800s (vs baseline 1200): sub-agent dispatch on large repos blows +# past a tighter budget — investigator passes are the long tail. +# Boot concurrency 6 (vs baseline 8) eases h2/SSL long-tail stalls under +# heavier sub-agent dispatch. +export SWE_TIME_BUDGET_SEC="${SWE_TIME_BUDGET_SEC:-1800}" +export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-6}" + +# --- trajectory fan-out --- +# generate() emits one Sample per segment (reducer splits reward/K); +# group_id is shared so the per-rollout-mean loss reducer still counts +# the trajectory once. +# --rollout-max-response-len caps one model turn. The custom generate function +# uses --rollout-max-context-len as the multi-turn prompt+response budget. + +# --- claude-code CLI extras --- +# SETTINGS_JSON: autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI +# compacts before any segment crosses the training-side cap. +# AGENTS_JSON: register a read-only `investigator` sub-agent (Grep/Read/Glob) +# as a concrete, narrowly-scoped dispatch target. +# SWE_CLAUDE_EXTRA_ARGS: WebFetch/WebSearch are off (sandbox has no outbound +# internet); --disable-slash-commands keeps the model from emitting /compact +# as a competing branching pathway. +SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +AGENTS_JSON='{"investigator":{"description":"Searches the repo for relevant files before any edit","prompt":"You are an investigator sub-agent. Use Grep/Read/Glob to find every file relevant to the user task, then return a short bulleted summary. Do NOT edit anything.","tools":["Grep","Read","Glob"]}}' +export SWE_CLAUDE_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" + +# Optional: bias the model to dispatch the investigator before any edit. +# Uncomment to maximize sub-agent dispatch — naming the exact call form +# (Agent tool with subagent_type=investigator) is what reliably triggers it. +# export SWE_CC_PROMPT="Read PROBLEM_STATEMENT.md. BEFORE editing any file, dispatch the 'investigator' sub-agent (via the Agent tool with subagent_type=investigator) to locate every file relevant to the issue. Then fix the issue and run the tests." + +# ============ proxy bypass for in-cluster traffic ============ +export no_proxy="127.0.0.1,${MASTER_ADDR},${VIME_HEAD_HOST}" +export NO_PROXY="${no_proxy}" + +cd "${VIME_DIR}" + +# ============ bring up ray cluster ============ +HOSTFILE="${HOSTFILE:-/root/mpi_rack_hostfile}" +ACTOR_NUM_NODES="${ACTOR_NUM_NODES:-${MLP_WORKER_NUM:-8}}" +ACTOR_NUM_GPUS_PER_NODE="${ACTOR_NUM_GPUS_PER_NODE:-8}" + +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +if [[ -f "${HOSTFILE}" ]]; then + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + [[ -z "${WORKER_IP}" ]] && continue + [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue + echo "Starting Ray worker on ${WORKER_IP}" + ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ + "pkill -9 -f 'vllm serve' ; ray stop --force ; pkill -9 python ; \ + ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ + --node-ip-address ${WORKER_IP} --disable-usage-stats" & + done + wait +fi + +echo "Waiting for Ray cluster to stabilize..." +sleep 30 +ray status + +# ============ runtime env propagated to ray workers ============ +export VIME_DIR +RUNTIME_ENV_JSON=$(python3 - <&1 | tee "${LOG_FILE}" + +echo "RUN_ROOT=${RUN_ROOT}" diff --git a/examples/coding_agent_rl/sandbox.py b/examples/coding_agent_rl/sandbox.py new file mode 100644 index 000000000..cd9fbe48b --- /dev/null +++ b/examples/coding_agent_rl/sandbox.py @@ -0,0 +1,400 @@ +"""Coding-agent sandbox helpers. + +The provider-agnostic sandbox contract and E2B backend live in +``vime.agent.sandbox``. This module keeps the coding-agent/SWE-specific +bootstrap, Claude Code runner, diff capture, and fresh-sandbox evaluator. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import lzma +import os +import shlex +import shutil +import tempfile +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path + +from vime.agent.sandbox import E2BSandbox, Sandbox + + +logger = logging.getLogger(__name__) + +# Paths inside the sandbox (avoid clashes with image-shipped paths). +_PATCH = "/workspace/__cagent_patch__.diff" +_PRE = "/workspace/__cagent_pre__.sh" +_SWEPRO_DIR = "/workspace/swepro_eval" + +SWE_HOST_NODE_TARBALL = Path( + os.environ.get( + "SWE_HOST_NODE_TARBALL", + "/path/to/node-v22.20.0-linux-x64.tar.xz", + ) +) +SWE_HOST_CC_TARBALL = Path( + os.environ.get( + "SWE_HOST_CC_TARBALL", + "/path/to/anthropic-ai-claude-code.tgz", + ) +) +SWE_BOOT_CONCURRENCY = int(os.environ.get("SWE_BOOT_CONCURRENCY", "16")) +SWE_BOOT_RETRIES = int(os.environ.get("SWE_BOOT_RETRIES", "2")) +CC_PROMPT = os.environ.get( + "SWE_CC_PROMPT", + "Read PROBLEM_STATEMENT.md in the current directory and resolve the issue. " + "Edit source files only (do NOT touch tests). After editing, run the relevant " + "tests to verify your fix passes. Do NOT modify PROBLEM_STATEMENT.md and do " + "NOT commit. When finished, print a one-line summary and exit.", +) + +_BOOT_SEM: asyncio.Semaphore | None = None + + +# --------------------------------------------------------------------------- +# Sandbox bootstrap (Node + Claude Code + agent user) +# --------------------------------------------------------------------------- +@asynccontextmanager +async def boot_agent_sandbox(image: str) -> AsyncIterator[E2BSandbox]: + """Boot a fresh E2B sandbox and install the Claude Code toolchain. + + This is the provisioning wrapper for the work sandbox: create the sandbox + from the dataset image, install Node 22 + Claude Code CLI from host + tarballs, retry transient boot/install failures, and close the sandbox when + the caller leaves the context. + """ + global _BOOT_SEM + if _BOOT_SEM is None: + _BOOT_SEM = asyncio.Semaphore(SWE_BOOT_CONCURRENCY) + + sb = None + last_err: Exception | None = None + for attempt in range(SWE_BOOT_RETRIES): + cand = E2BSandbox(image) + try: + async with _BOOT_SEM: + await cand.__aenter__() + try: + await install_node22(cand, SWE_HOST_NODE_TARBALL) + await install_claude_code(cand, SWE_HOST_CC_TARBALL) + except BaseException: + await cand.__aexit__(None, None, None) + raise + sb = cand + break + except Exception as e: + last_err = e + logger.warning( + "[coding_agent_rl] provision attempt %d/%d failed: %s: %s", + attempt + 1, + SWE_BOOT_RETRIES, + type(e).__name__, + str(e)[:200], + ) + await asyncio.sleep(1 + attempt) + if sb is None: + assert last_err is not None + raise last_err + try: + yield sb + finally: + await sb.__aexit__(None, None, None) + + +async def install_node22(sb: Sandbox, host_tarball: Path) -> None: + """Node 22 over the base image (Debian 12 ships 16; cli.js needs >= 20). + Decompresses .xz on the host (cached) so sandboxes without xz-utils can + still run plain `tar xf`. npm prefix=/usr/local required for sweap-images.""" + host_tarball = Path(host_tarball) + if host_tarball.suffix == ".xz": + plain = Path(tempfile.gettempdir()) / f"coding_agent_rl.{host_tarball.stem}.tar" + if not plain.exists(): + tmp = plain.with_suffix(".tar.partial") + with lzma.open(host_tarball, "rb") as src, open(tmp, "wb") as dst: + shutil.copyfileobj(src, dst) + os.replace(tmp, plain) + host_tarball = plain + await sb.write_file("/tmp/node22.tar", host_tarball) + await sb.exec( + "set -e && mkdir -p /opt/node22 && " + "tar xf /tmp/node22.tar -C /opt/node22 --strip-components=1 && " + "ln -sf /opt/node22/bin/node /usr/local/bin/node && " + "ln -sf /opt/node22/bin/npm /usr/local/bin/npm && " + "ln -sf /opt/node22/bin/npx /usr/local/bin/npx && " + "hash -r 2>/dev/null || true && node --version && npm --version", + user="root", + timeout=180, + check=True, + ) + + +async def install_claude_code(sb: Sandbox, host_tarball: Path) -> None: + await sb.write_file("/tmp/claude-code.tgz", host_tarball) + await sb.exec( + "npm install -g --prefix=/usr/local --no-audit --no-fund /tmp/claude-code.tgz " + "&& ls -la /usr/local/bin/claude && /usr/local/bin/claude --version", + user="root", + timeout=300, + check=True, + ) + + +async def ensure_agent_user(sb: Sandbox, workdir: str) -> None: + """Create the unprivileged 'agent' user that owns workdir + can git diff. + Settings file pre-acks bypass-permissions so claude-code starts headless.""" + await sb.exec( + f"id agent >/dev/null 2>&1 || useradd -m -s /bin/bash agent && " + f"chown -R agent:agent /home/agent {workdir} && " + f"git config --system --add safe.directory '*' && id agent && " + f"mkdir -p /home/agent/.claude && " + f'echo \'{{"hasCompletedOnboarding": true, "bypassPermissionsModeAccepted": true}}\' ' + f"| tee /home/agent/.claude.json /home/agent/.claude/settings.json > /dev/null && " + f"chown -R agent:agent /home/agent/.claude /home/agent/.claude.json", + user="root", + check=True, + timeout=60, + ) + + +async def apply_before_repo_set_cmd(sb: Sandbox, workdir: str, swepro: dict) -> None: + """Run swepro['before_repo_set_cmd'] in the sandbox if present (no-op if not).""" + before = swepro.get("before_repo_set_cmd") if swepro else None + if not before: + return + payload = f"set -e\ncd {workdir}\n{before}\n" + await sb.exec( + "mkdir -p /workspace/swepro_setup && chown agent:agent /workspace/swepro_setup", user="root", check=True + ) + await sb.write_file("/workspace/swepro_setup/before.sh", payload, user="agent") + await sb.exec("bash /workspace/swepro_setup/before.sh", user="agent", check=False, timeout=600) + + +# --------------------------------------------------------------------------- +# Agent run (workspace prep + claude-code spawn + done-marker poll) +# --------------------------------------------------------------------------- +async def run_claude_code( + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + time_budget_sec: int, + problem_statement: str = "", + swepro: dict | None = None, + pre_commands: list[str] | str | None = None, + prompt: str | None = None, +) -> int: + """Prepare the SWE workspace, write PROBLEM_STATEMENT.md, then run CC.""" + await ensure_agent_user(sb, workdir) + if swepro: + await apply_before_repo_set_cmd(sb, workdir, swepro) + if pre_commands: + await apply_pre_commands(sb, workdir, pre_commands) + await sb.write_file( + f"{workdir}/PROBLEM_STATEMENT.md", + problem_statement or "", + user="agent", + ) + return await _spawn_claude_code( + sb, + workdir=workdir, + session_id=session_id, + adapter_url=adapter_url, + prompt=prompt or CC_PROMPT, + time_budget_sec=time_budget_sec, + ) + + +async def _spawn_claude_code( + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + prompt: str, + time_budget_sec: int, +) -> int: + """Spawn claude-code detached + poll a done-marker file. + + E2B's gateway resets HTTP/2 around 6.5 min, so we can't keep a long-lived + foreground exec. The launcher writes the exit code into a marker file + and we poll it every 5s via short RPCs (which also keeps the sandbox + alive against idle GC).""" + done = f"{workdir}/.cagent_done" + launcher = f"{workdir}/.cagent_run.sh" + traj = f"{workdir}/claude_code_trajectory.jsonl" + + launcher_body = ( + "#!/bin/bash\n" + f"cd {workdir}\n" + "export HOME=/home/agent\n" + f"/usr/local/bin/claude -p {json.dumps(prompt)} " + f"--permission-mode bypassPermissions " + f"--output-format stream-json --include-partial-messages " + f"--include-hook-events --verbose " + f"{os.environ.get('SWE_CLAUDE_EXTRA_ARGS', '').strip()} " + f"2>&1 | tee {shlex.quote(traj)}\n" + f"echo $? > {done}\n" + ) + await sb.write_file(launcher, launcher_body, user="agent") + await sb.exec(f"chmod +x {launcher}", user="agent", timeout=30) + + env = { + "ANTHROPIC_BASE_URL": adapter_url, + "ANTHROPIC_AUTH_TOKEN": session_id, + "ANTHROPIC_MODEL": "vime-actor", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + } + env_keys = ",".join(env.keys()) + await sb.exec( + f"runuser -u agent --whitelist-environment={env_keys}" + f" -- bash -c 'setsid {launcher} < /dev/null > /dev/null 2>&1 &'", + user="root", + env=env, + timeout=30, + check=True, + ) + + deadline = time.time() + time_budget_sec + exit_code = -2 # convention: -2 = budget exceeded + while time.time() < deadline: + await asyncio.sleep(5) + ec, out, _ = await sb.exec( + f"test -f {done} && cat {done}", + user="agent", + timeout=15, + check=False, + ) + if ec == 0: + try: + exit_code = int((out or "").strip() or "-1") + except ValueError: + exit_code = -1 + break + return exit_code + + +async def git_diff(sb: Sandbox, workdir: str) -> str: + cmd = ( + f"cd {workdir} && git add -N . && " + f"git diff -- . ':(exclude)PROBLEM_STATEMENT.md' " + f"':(exclude)claude_code_trajectory.jsonl' " + f"':(exclude).cagent_done' ':(exclude).cagent_run.sh'" + ) + _, out, _ = await sb.exec(cmd, user="agent", timeout=120) + return out + + +# --------------------------------------------------------------------------- +# Eval (fresh sandbox, apply diff, run dataset tests) +# --------------------------------------------------------------------------- +async def evaluate( + *, + image: str, + workdir: str, + diff_text: str, + swepro: dict | None = None, + eval_cmd: str | None = None, + pre_commands: list[str] | str | None = None, + timeout_sec: int = 600, +) -> tuple[float, bool, bool]: + """Returns (reward, solved, applied_cleanly). + + No-test-cheating guarantee: the eval sandbox is built from the same image + but starts CLEAN, so only the model-produced diff affects reward.""" + if not (swepro or eval_cmd): + logger.warning("[e2b.evaluate] no swepro/eval_cmd; reward=0") + return 0.0, False, True + + async with E2BSandbox(image) as ev: + await ensure_agent_user(ev, workdir) + if swepro: + await _setup_swepro_assets(ev, swepro) + await apply_before_repo_set_cmd(ev, workdir, swepro) + if pre_commands: + await apply_pre_commands(ev, workdir, pre_commands) + + applied = await _apply_diff(ev, workdir, diff_text) + if not applied: + return 0.0, False, False + + if swepro: + r, s = await _run_swepro(ev, workdir, swepro, timeout_sec) + return r, s, True + r, s = await _run_eval_cmd(ev, workdir, eval_cmd, timeout_sec) + return r, s, True + + +async def _setup_swepro_assets(ev: Sandbox, swepro: dict) -> None: + await ev.exec(f"mkdir -p {_SWEPRO_DIR} && chmod 777 {_SWEPRO_DIR}", user="root", check=True) + for k, dst in [("run_script_path", "run_script.sh"), ("parser_script_path", "parser.py")]: + host_p = swepro.get(k) + if host_p: + text = Path(host_p).read_text() + await ev.write_file(f"{_SWEPRO_DIR}/{dst}", text, user="root") + await ev.exec(f"chmod 755 {_SWEPRO_DIR}/* && chown -R agent:agent {_SWEPRO_DIR}", user="root", check=True) + + +async def apply_pre_commands(ev: Sandbox, workdir: str, pre: list[str] | str) -> None: + # Public: also called by generate.py to keep the work sandbox baseline + # aligned with eval (sweb-style pre_commands typically `git checkout + # -f`, so skipping in work sandbox makes the model's diff + # context mismatch the eval base -> 100% apply failure). + if isinstance(pre, str): + body = pre.replace("\\n", "\n") + else: + body = "\n".join(c for c in (pre or []) if c) + await ev.write_file(_PRE, "set -e\n" + body, user="agent") + await ev.exec(f"chmod 755 {_PRE} && cd {workdir} && bash {_PRE}", user="agent", check=False, timeout=600) + + +async def _apply_diff(ev: Sandbox, workdir: str, diff_text: str) -> bool: + if not diff_text.strip(): + return True + await ev.write_file(_PATCH, diff_text, user="agent") + for cmd in [ + f"cd {workdir} && git apply --3way --whitespace=nowarn {_PATCH}", + f"cd {workdir} && git apply --whitespace=nowarn {_PATCH}", + f"cd {workdir} && patch -p1 --no-backup-if-mismatch < {_PATCH}", + ]: + ec, _, _ = await ev.exec(cmd, user="agent", check=False, timeout=120) + if ec == 0: + return True + return False + + +async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict, timeout: int) -> tuple[float, bool]: + test_arg = ",".join(swepro.get("selected_test_files") or []) + stdout_f = f"{_SWEPRO_DIR}/stdout.log" + stderr_f = f"{_SWEPRO_DIR}/stderr.log" + result_f = f"{_SWEPRO_DIR}/result.json" + await ev.exec( + f"cd {workdir} && bash {_SWEPRO_DIR}/run_script.sh " + f"{json.dumps(test_arg)} > {stdout_f} 2> {stderr_f} || true", + user="agent", + check=False, + timeout=timeout, + ) + await ev.exec( + f"python3 {_SWEPRO_DIR}/parser.py {stdout_f} {stderr_f} {result_f}", + user="agent", + check=False, + timeout=120, + ) + raw = await ev.read_file(result_f, user="agent") + parsed = json.loads(raw) if raw else {"tests": []} + passed = {t["name"] for t in parsed.get("tests", []) if t.get("status") == "PASSED"} + required = set(swepro.get("fail_to_pass") or []) | set(swepro.get("pass_to_pass") or []) + solved = bool(required) and required.issubset(passed) + return (1.0 if solved else 0.0), solved + + +async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> tuple[float, bool]: + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) + return (1.0 if ec == 0 else 0.0), ec == 0 diff --git a/tests/test_agent_adapters.py b/tests/test_agent_adapters.py new file mode 100644 index 000000000..94e055f67 --- /dev/null +++ b/tests/test_agent_adapters.py @@ -0,0 +1,853 @@ +import asyncio +import json +import sys +from pathlib import Path + +import pytest +from aiohttp import web +from aiohttp.test_utils import TestClient, TestServer + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.adapters import anthropic, openai +from vime.agent.adapters.common import VLLM_URL_KEY +from vime.agent.trajectory import TurnRecord + + +NUM_GPUS = 0 + + +class ToyTokenizer: + def __init__(self, outputs: dict[tuple[int, ...], str] | None = None) -> None: + self.outputs = outputs or {} + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + return list(range(1, len(messages) + 2)) + + def decode(self, ids, skip_special_tokens=False): + return self.outputs.get(tuple(ids), "") + + +class ScriptedTokenizer(ToyTokenizer): + def __init__(self, prompts: list[list[int]], outputs: dict[tuple[int, ...], str]) -> None: + super().__init__(outputs) + self.prompts = [list(prompt) for prompt in prompts] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + assert self.prompts, "unexpected chat-template render" + return self.prompts.pop(0) + + +class FakeVLLM: + def __init__(self, turns: list[list[tuple[float, int]]]) -> None: + self.turns = [list(turn) for turn in turns] + self.requests: list[dict] = [] + self.routing_keys: list[str | None] = [] + + async def handle_generate(self, request): + self.routing_keys.append(request.headers.get("x-session-id")) + self.requests.append(await request.json()) + assert self.turns, "unexpected /inference/v1/generate call" + turn = self.turns.pop(0) + token_ids = [token_id for _logprob, token_id in turn] + content = [{"logprob": logprob} for logprob, _token_id in turn] + return web.json_response( + { + "choices": [ + { + "token_ids": token_ids, + "logprobs": {"content": content}, + "finish_reason": "stop", + } + ] + } + ) + + +class FakeRequest: + def __init__(self, headers: dict[str, str]) -> None: + self.headers = headers + + +def _parse_sse(raw: str) -> list[tuple[str, object]]: + events: list[tuple[str, object]] = [] + event_name = "message" + data_lines: list[str] = [] + + def flush() -> None: + nonlocal event_name, data_lines + if not data_lines: + event_name = "message" + return + data = "\n".join(data_lines) + payload: object + if data == "[DONE]": + payload = data + else: + payload = json.loads(data) + events.append((event_name, payload)) + event_name = "message" + data_lines = [] + + for line in raw.splitlines(): + if not line: + flush() + elif line.startswith("event:"): + event_name = line.removeprefix("event:").strip() + elif line.startswith("data:"): + data_lines.append(line.removeprefix("data:").strip()) + flush() + return events + + +@pytest.mark.unit +def test_session_id_comes_from_protocol_fields_not_custom_header(): + assert ( + openai._request_session_id( + FakeRequest({"X-Slime-Session-Id": "custom"}), + {"metadata": {"session_id": "meta-session"}, "user": "body-user"}, + ) + == "meta-session" + ) + assert ( + openai._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom"}), {"user": "body-user"}) == "body-user" + ) + assert ( + anthropic._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom", "X-Api-Key": "anthropic-key"})) + == "anthropic-key" + ) + assert ( + anthropic._request_session_id( + FakeRequest({"Authorization": "Bearer bearer-session", "X-Api-Key": "anthropic-key"}) + ) + == "bearer-session" + ) + + +@pytest.mark.unit +def test_anthropic_translation_keeps_tool_results_and_tool_schema(): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan"}, + {"type": "text", "text": "ok"}, + {"type": "tool_use", "name": "lookup", "input": {"q": "slime"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "u1", "content": "result"}]}, + ] + + translated = anthropic._translate_anthropic(messages, system="sys") + tools = anthropic._anthropic_tools_to_chat_tools( + [{"name": "lookup", "description": "search", "input_schema": {"type": "object"}}] + ) + + assert translated == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "ok", + "reasoning_content": "plan", + "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "slime"}}}], + }, + {"role": "tool", "content": "result"}, + ] + assert tools == [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object"}, + }, + } + ] + + +@pytest.mark.unit +def test_openai_translation_and_responses_input_shapes(): + chat_messages = openai._translate_chat_messages( + [ + {"role": "developer", "content": "rules"}, + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "slime"}}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + ) + response_messages = openai._responses_input_to_messages( + [ + {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, + {"type": "function_call_output", "call_id": "call_1", "output": "answer"}, + ], + instructions="be brief", + ) + + assert chat_messages == [ + {"role": "system", "content": "rules"}, + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "slime"}'}, + } + ], + }, + {"role": "tool", "content": "found", "tool_call_id": "call_1"}, + ] + assert response_messages == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "answer"}, + ] + + +@pytest.mark.unit +def test_openai_chat_completion_endpoint_records_token_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[101], finish_reason="stop", output_log_probs=[-0.1]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(101,): "hello"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat"}, + json={ + "model": "actor", + "messages": [{"role": "user", "content": "hello?"}], + "max_tokens": 4, + }, + ) + data = await resp.json() + finally: + await client.close() + + segments = await adapter.finish_session("sid-chat") + assert resp.status == 200 + assert data["object"] == "chat.completion" + assert data["choices"][0]["message"] == {"role": "assistant", "content": "hello"} + assert data["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + assert segments[0].prompt_ids == [1, 2] + assert segments[0].response_ids == [101] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_chat_completion_streaming_returns_sse_chunks_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[401], finish_reason="stop", output_log_probs=[-0.4]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(401,): "streamed text"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat-stream"}, + json={ + "model": "actor", + "stream": True, + "messages": [{"role": "user", "content": "hello?"}], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + chunks = [payload for _, payload in events if isinstance(payload, dict)] + segments = await adapter.finish_session("sid-chat-stream") + assert resp.status == 200 + assert chunks[0]["object"] == "chat.completion.chunk" + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant"} + assert any(c["choices"][0]["delta"] == {"content": "streamed text"} for c in chunks) + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + assert events[-1] == ("message", "[DONE]") + assert segments[0].prompt_ids == [1, 2] + assert segments[0].response_ids == [401] + assert segments[0].rollout_log_probs == [-0.4] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_chat_completion_streaming_returns_tool_call_delta(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[451], finish_reason="stop", output_log_probs=[-0.45] + ) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "use it slime" + tokenizer = ToyTokenizer({(451,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat-tool-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat-tool-stream"}, + json={ + "model": "actor", + "stream": True, + "messages": [{"role": "user", "content": "call lookup"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + } + ], + }, + ) + raw_sse = await resp.text() + finally: + await client.close() + + chunks = [payload for _, payload in _parse_sse(raw_sse) if isinstance(payload, dict)] + tool_delta = next(c["choices"][0]["delta"] for c in chunks if "tool_calls" in c["choices"][0]["delta"]) + segments = await adapter.finish_session("sid-chat-tool-stream") + assert resp.status == 200 + assert any(c["choices"][0]["delta"] == {"content": "use it"} for c in chunks) + assert tool_delta["tool_calls"][0]["index"] == 0 + assert tool_delta["tool_calls"][0]["function"]["name"] == "lookup" + assert tool_delta["tool_calls"][0]["function"]["arguments"] == '{"query": "slime"}' + assert chunks[-1]["choices"][0]["finish_reason"] == "tool_calls" + assert segments[0].response_ids == [451] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_endpoint_returns_function_calls(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[301], finish_reason="stop", output_log_probs=[-0.3]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "look slime" + tokenizer = ToyTokenizer({(301,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses"}, + json={ + "model": "actor", + "input": "find it", + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + data = await resp.json() + finally: + await client.close() + + output_types = [item["type"] for item in data["output"]] + function_call = next(item for item in data["output"] if item["type"] == "function_call") + segments = await adapter.finish_session("sid-responses") + assert resp.status == 200 + assert data["object"] == "response" + assert output_types == ["message", "function_call"] + assert data["output"][0]["content"][0]["text"] == "look" + assert function_call["name"] == "lookup" + assert function_call["arguments"] == '{"query": "slime"}' + assert segments[0].response_ids == [301] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_streaming_preserves_function_call_output(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[551], finish_reason="stop", output_log_probs=[-0.55] + ) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "slime" + tokenizer = ToyTokenizer({(551,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses-tool-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses-tool-stream"}, + json={ + "model": "actor", + "stream": True, + "input": "call lookup", + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + raw_sse = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw_sse) + created = next(payload for name, payload in events if name == "response.created") + completed = next(payload for name, payload in events if name == "response.completed") + completed_call = next(item for item in completed["response"]["output"] if item["type"] == "function_call") + segments = await adapter.finish_session("sid-responses-tool-stream") + assert resp.status == 200 + assert created["type"] == "response.created" + assert [item["type"] for item in created["response"]["output"]] == ["function_call"] + assert completed_call["name"] == "lookup" + assert completed_call["arguments"] == '{"query": "slime"}' + assert segments[0].response_ids == [551] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_streaming_returns_sse_events_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[501], finish_reason="stop", output_log_probs=[-0.5]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(501,): "response text"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses-stream"}, + json={ + "model": "actor", + "stream": True, + "instructions": "be brief", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello?"}]}], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + event_names = [name for name, _ in events] + text_delta = next(payload for name, payload in events if name == "response.output_text.delta") + completed = next(payload for name, payload in events if name == "response.completed") + segments = await adapter.finish_session("sid-responses-stream") + assert resp.status == 200 + assert event_names == ["response.created", "response.output_text.delta", "response.completed"] + assert text_delta == {"type": "response.output_text.delta", "delta": "response text"} + assert completed["response"]["status"] == "completed" + assert completed["response"]["usage"] == {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4} + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [501] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_endpoint_returns_non_stream_json_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[581], finish_reason="stop", output_log_probs=[-0.58] + ) + + async def run_case(): + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = ToyTokenizer({(581,): "plain response"}) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-anthropic-json", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-json"}, + json={ + "model": "actor", + "system": "be useful", + "max_tokens": 4, + "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], + }, + ) + data = await resp.json() + finally: + await client.close() + + segments = await adapter.finish_session("sid-anthropic-json") + assert resp.status == 200 + assert data["type"] == "message" + assert data["model"] == "actor" + assert data["content"] == [{"type": "text", "text": "plain response"}] + assert data["stop_reason"] == "end_turn" + assert data["usage"] == {"input_tokens": 3, "output_tokens": 1} + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [581] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_endpoint_streams_blocks_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[601], finish_reason="stop", output_log_probs=[-0.6]) + + async def run_case(): + monkeypatch.setattr(anthropic, "_generate", fake_generate) + raw_output = ( + "delegate inspect" + ) + tokenizer = ToyTokenizer({(601,): raw_output}) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-anthropic-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-stream"}, + json={ + "model": "actor", + "system": "be useful", + "stream": True, + "max_tokens": 4, + "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], + "tools": [ + { + "name": "Task", + "description": "spawn subagent", + "input_schema": {"type": "object", "properties": {"description": {"type": "string"}}}, + } + ], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + names = [name for name, _ in events] + starts = [payload for name, payload in events if name == "content_block_start"] + deltas = [payload for name, payload in events if name == "content_block_delta"] + message_delta = next(payload for name, payload in events if name == "message_delta") + segments = await adapter.finish_session("sid-anthropic-stream") + assert resp.status == 200 + assert names[0] == "message_start" + assert names[-1] == "message_stop" + assert any(s["content_block"]["type"] == "text" for s in starts) + assert any(s["content_block"]["type"] == "tool_use" and s["content_block"]["name"] == "Task" for s in starts) + assert any(d["delta"].get("text") == "delegate" for d in deltas) + assert any(json.loads(d["delta"].get("partial_json", "{}")) == {"description": "inspect"} for d in deltas) + assert message_delta["delta"]["stop_reason"] == "tool_use" + assert message_delta["usage"] == {"input_tokens": 3, "output_tokens": 1} + assert segments[0].metadata["segment_kind"] == "final" + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [601] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_multiturn_uses_vllm_tokens_for_training_segment(): + async def run_case(): + upstream = FakeVLLM( + [ + [(-0.20, 20), (-0.21, 21)], + [(-0.40, 40)], + ] + ) + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) + upstream_server = TestServer(upstream_app) + await upstream_server.start_server() + + tool_raw = "slime" + tokenizer = ScriptedTokenizer( + prompts=[ + [10, 11], + [10, 11, 20, 21, 30, 31], + ], + outputs={ + (20, 21): tool_raw, + (40,): "done", + }, + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url=str(upstream_server.make_url("")).rstrip("/")) + adapter.open_session("sid-openai-token", sampling_defaults={"max_new_tokens": 99}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + first = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-openai-token"}, + json={ + "model": "actor", + "input": "find slime", + "max_output_tokens": 5, + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + first_data = await first.json() + function_call = next(item for item in first_data["output"] if item["type"] == "function_call") + + second = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-openai-token"}, + json={ + "model": "actor", + "input": [ + {"role": "user", "content": "find slime"}, + function_call, + { + "type": "function_call_output", + "call_id": function_call["call_id"], + "output": "found slime", + }, + ], + "max_output_tokens": 7, + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + second_data = await second.json() + finally: + await client.close() + await upstream_server.close() + + segments = await adapter.finish_session("sid-openai-token") + assert first.status == 200 + assert second.status == 200 + assert function_call["name"] == "lookup" + assert function_call["arguments"] == '{"query": "slime"}' + assert second_data["output"][0]["content"][0]["text"] == "done" + assert [req["token_ids"] for req in upstream.requests] == [[10, 11], [10, 11, 20, 21, 30, 31]] + assert upstream.routing_keys == ["sid-openai-token", "sid-openai-token"] + assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 + assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 + assert segments[0].prompt_ids == [10, 11] + assert segments[0].response_ids == [20, 21, 30, 31, 40] + assert segments[0].loss_mask == [1, 1, 0, 0, 1] + assert segments[0].rollout_log_probs == [-0.20, -0.21, 0.0, 0.0, -0.40] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_multiturn_uses_vllm_tokens_for_training_segment(): + async def run_case(): + upstream = FakeVLLM( + [ + [(-1.20, 120), (-1.21, 121)], + [(-1.40, 140), (-1.41, 141)], + ] + ) + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) + upstream_server = TestServer(upstream_app) + await upstream_server.start_server() + + tool_raw = "slime" + tokenizer = ScriptedTokenizer( + prompts=[ + [110, 111], + [110, 111, 120, 121, 130], + ], + outputs={ + (120, 121): tool_raw, + (140, 141): "anthropic done", + }, + ) + adapter = anthropic.AnthropicAdapter( + tokenizer=tokenizer, + vllm_url=str(upstream_server.make_url("")).rstrip("/"), + ) + adapter.open_session("sid-anthropic-token", sampling_defaults={"max_new_tokens": 99}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + first = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-token"}, + json={ + "model": "actor", + "max_tokens": 5, + "messages": [{"role": "user", "content": [{"type": "text", "text": "find slime"}]}], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + first_data = await first.json() + tool_use = next(block for block in first_data["content"] if block["type"] == "tool_use") + + second = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-token"}, + json={ + "model": "actor", + "max_tokens": 7, + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "find slime"}]}, + {"role": "assistant", "content": first_data["content"]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_use["id"], + "content": "found slime", + } + ], + }, + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + second_data = await second.json() + finally: + await client.close() + await upstream_server.close() + + segments = await adapter.finish_session("sid-anthropic-token") + assert first.status == 200 + assert second.status == 200 + assert tool_use["name"] == "lookup" + assert tool_use["input"] == {"query": "slime"} + assert second_data["content"] == [{"type": "text", "text": "anthropic done"}] + assert [req["token_ids"] for req in upstream.requests] == [[110, 111], [110, 111, 120, 121, 130]] + assert upstream.routing_keys == ["sid-anthropic-token", "sid-anthropic-token"] + assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 + assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 + assert segments[0].prompt_ids == [110, 111] + assert segments[0].response_ids == [120, 121, 130, 140, 141] + assert segments[0].loss_mask == [1, 1, 0, 1, 1] + assert segments[0].rollout_log_probs == [-1.20, -1.21, 0.0, -1.40, -1.41] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_generate_posts_token_ids_and_extracts_logprobs(): + async def run_case(): + captured = {} + captured_headers = {} + + async def handle_generate(request): + captured_headers.update(request.headers) + captured.update(await request.json()) + return web.json_response( + { + "choices": [ + { + "token_ids": [701, 702], + "logprobs": {"content": [{"logprob": -0.7}, {"logprob": -0.8}]}, + "finish_reason": "stop", + } + ] + } + ) + + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", handle_generate) + server = TestServer(upstream_app) + await server.start_server() + try: + session = openai.Session(sampling_defaults={"max_new_tokens": 9}) + turn = await openai._generate( + [11, 12], + session, + {"max_tokens": 3, "temperature": 0.25, "stop": [""]}, + {VLLM_URL_KEY: str(server.make_url("")).rstrip("/")}, + ) + finally: + await server.close() + + assert captured["token_ids"] == [11, 12] + assert captured["sampling_params"]["logprobs"] == 1 + assert captured_headers.get("x-session-id") is None + assert captured["sampling_params"]["max_tokens"] == 3 + assert captured["sampling_params"]["temperature"] == 0.25 + assert captured["sampling_params"]["stop"] == [""] + assert turn.prompt_ids == [11, 12] + assert turn.output_ids == [701, 702] + assert turn.output_log_probs == [-0.7, -0.8] + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_sdk_adapters.py b/tests/test_agent_sdk_adapters.py new file mode 100644 index 000000000..a4ad3d4ce --- /dev/null +++ b/tests/test_agent_sdk_adapters.py @@ -0,0 +1,418 @@ +import asyncio +import sys +from pathlib import Path + +import httpx +import pytest +from aiohttp.test_utils import TestClient, TestServer + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.adapters import anthropic, openai +from vime.agent.trajectory import TurnRecord + + +NUM_GPUS = 0 + + +agents = pytest.importorskip("agents") +anthropic_sdk = pytest.importorskip("anthropic") +openai_sdk = pytest.importorskip("openai") + + +class SDKTokenizer: + def __init__(self, outputs: list[str]) -> None: + self.outputs = outputs + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + return list(range(1, len(messages) + 2)) + + def decode(self, ids, skip_special_tokens=False): + return self.outputs[ids[0] - 1] + + +@pytest.mark.integration +def test_openai_agents_sdk_responses_runs_tool_loop_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=[len(calls)], + finish_reason="stop", + output_log_probs=[-0.1 * len(calls)], + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer( + [ + "slime", + "final after tool", + ] + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + @agents.function_tool + def lookup(query: str) -> str: + return f"found {query}" + + agents.set_tracing_disabled(True) + model = agents.OpenAIResponsesModel(model="actor", openai_client=oai) + agent = agents.Agent( + name="sdk-responses", + instructions="Use lookup.", + model=model, + tools=[lookup], + model_settings=agents.ModelSettings(max_tokens=4), + ) + try: + result = await agents.Runner.run(agent, "find slime") + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai") + assert result.final_output == "final after tool" + assert len(calls) == 2 + assert calls[0]["body"]["max_output_tokens"] == 4 + assert calls[0]["body"]["tools"][0]["name"] == "lookup" + assert calls[1]["body"]["input"][-1] == { + "call_id": calls[1]["body"]["input"][-2]["call_id"], + "output": "found slime", + "type": "function_call_output", + } + assert tokenizer.rendered[0][0] == [ + {"role": "system", "content": "Use lookup."}, + {"role": "user", "content": "find slime"}, + ] + assert tokenizer.rendered[1][0][-1] == { + "role": "tool", + "content": "found slime", + "tool_call_id": calls[1]["body"]["input"][-2]["call_id"], + } + assert segments[0].metadata["segment_kind"] == "final" + assert segments[0].response_ids[-1] == 2 + assert segments[0].loss_mask[-1] == 1 + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_agents_sdk_chat_completions_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.2] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer(["chat final"]) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-chat", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + agents.set_tracing_disabled(True) + model = agents.OpenAIChatCompletionsModel(model="actor", openai_client=oai) + agent = agents.Agent( + name="sdk-chat", + instructions="Be short.", + model=model, + model_settings=agents.ModelSettings(max_tokens=5), + ) + try: + result = await agents.Runner.run(agent, "say hi") + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-chat") + assert result.final_output == "chat final" + assert calls[0]["body"]["max_tokens"] == 5 + assert calls[0]["body"]["messages"] == [ + {"content": "Be short.", "role": "system"}, + {"role": "user", "content": "say hi"}, + ] + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [1] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_sdk_chat_completion_streaming_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.25] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer( + ["streamed via sdk slime"] + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-chat-stream", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await oai.chat.completions.create( + model="actor", + messages=[{"role": "user", "content": "call lookup"}], + tools=[ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + } + ], + stream=True, + ) + text_parts = [] + tool_names = [] + tool_arguments = [] + finish_reasons = [] + usages = [] + async for chunk in stream: + choice = chunk.choices[0] + if choice.delta.content: + text_parts.append(choice.delta.content) + if choice.delta.tool_calls: + for tool_call in choice.delta.tool_calls: + tool_names.append(tool_call.function.name) + tool_arguments.append(tool_call.function.arguments) + if choice.finish_reason: + finish_reasons.append(choice.finish_reason) + if chunk.usage: + usages.append(chunk.usage) + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-chat-stream") + assert "".join(text_parts) == "streamed via sdk" + assert tool_names == ["lookup"] + assert tool_arguments == ['{"query": "slime"}'] + assert finish_reasons == ["tool_calls"] + assert usages[-1].prompt_tokens == 2 + assert usages[-1].completion_tokens == 1 + assert calls[0]["body"]["stream"] is True + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_sdk_responses_streaming_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.35] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer(["response stream via sdk"]) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-responses-stream", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await oai.responses.create( + model="actor", + instructions="Be brief.", + input="say hi", + stream=True, + ) + event_types = [] + deltas = [] + completed_response = None + async for event in stream: + event_types.append(event.type) + if event.type == "response.output_text.delta": + deltas.append(event.delta) + if event.type == "response.completed": + completed_response = event.response + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-responses-stream") + assert event_types == ["response.created", "response.output_text.delta", "response.completed"] + assert "".join(deltas) == "response stream via sdk" + assert completed_response.status == "completed" + assert completed_response.usage.input_tokens == 3 + assert completed_response.usage.output_tokens == 1 + assert calls[0]["body"]["stream"] is True + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_anthropic_sdk_non_streaming_messages_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.28] + ) + + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = SDKTokenizer(["anthropic json"]) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/")) + http_client = httpx.AsyncClient(trust_env=False) + anth = anthropic_sdk.AsyncAnthropic( + api_key="sdk-anthropic-json", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + message = await anth.messages.create( + model="actor", + max_tokens=6, + system="Be direct.", + messages=[{"role": "user", "content": "say hi"}], + ) + finally: + await client.close() + await anth.close() + + segments = await adapter.finish_session("sdk-anthropic-json") + assert message.type == "message" + assert message.content[0].type == "text" + assert message.content[0].text == "anthropic json" + assert message.stop_reason == "end_turn" + assert message.usage.input_tokens == 3 + assert message.usage.output_tokens == 1 + assert calls[0]["body"]["max_tokens"] == 6 + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_anthropic_sdk_streaming_messages_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.3] + ) + + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = SDKTokenizer(["anthropic final"]) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/")) + http_client = httpx.AsyncClient(trust_env=False) + anth = anthropic_sdk.AsyncAnthropic( + api_key="sdk-anthropic", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await anth.messages.create( + model="actor", + max_tokens=6, + system="Be direct.", + messages=[{"role": "user", "content": "say hi"}], + stream=True, + ) + text_parts = [] + event_types = [] + async for event in stream: + event_types.append(event.type) + if event.type == "content_block_delta" and getattr(event.delta, "text", None): + text_parts.append(event.delta.text) + finally: + await client.close() + await anth.close() + + segments = await adapter.finish_session("sdk-anthropic") + assert "".join(text_parts) == "anthropic final" + assert event_types == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert calls[0]["body"]["max_tokens"] == 6 + assert tokenizer.rendered[0][0] == [ + {"role": "system", "content": "Be direct."}, + {"role": "user", "content": "say hi"}, + ] + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [1] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_trajectory.py b/tests/test_agent_trajectory.py new file mode 100644 index 000000000..3abf88610 --- /dev/null +++ b/tests/test_agent_trajectory.py @@ -0,0 +1,143 @@ +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.trajectory import TurnRecord, TurnSegment, merge_turn_segments, merge_turns + + +NUM_GPUS = 0 + + +def _turn(prompt_ids: list[int], output_ids: list[int], output_log_probs: list[float] | None = None) -> TurnRecord: + return TurnRecord( + prompt_ids=prompt_ids, + output_ids=output_ids, + finish_reason="stop", + output_log_probs=( + output_log_probs if output_log_probs is not None else [-token_id / 100 for token_id in output_ids] + ), + ) + + +@pytest.mark.unit +def test_merge_turns_preserves_matched_prefix_on_prompt_drift(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 21, 12, 31], [13]), + _turn([10, 11, 21, 12, 22], [14]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 21, 12, 22, 14] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.12, 0.0, -0.14] + + +@pytest.mark.unit +def test_merge_turns_drops_middle_turn_when_next_prompt_skips_it(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 22], [13]), + _turn([10, 11, 22, 13, 31], [14]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 22, 13, 31, 14] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.13, 0.0, -0.14] + + +@pytest.mark.unit +def test_merge_turns_handles_consecutive_prompt_drifts(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 22], [13]), + _turn([10, 11, 23], [14]), + _turn([10, 11, 23, 14, 31], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 23, 14, 31, 15] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.14, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_masks_whole_output_when_prompt_drift_splits_it(): + segment = merge_turns( + [ + _turn([10], [11, 12, 13, 14]), + _turn([10, 11, 12, 99, 14], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 12, 99, 14, 15] + assert segment.loss_mask == [0, 0, 0, 0, 1] + assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_masks_whole_output_when_prompt_drift_changes_token_count(): + segment = merge_turns( + [ + _turn([10], [11, 12, 13, 14]), + _turn([10, 11, 12, 99, 100, 14], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 12, 99, 100, 14, 15] + assert segment.loss_mask == [0, 0, 0, 0, 0, 1] + assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_restarts_when_prompt_base_changes(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([20, 21], [22]), + _turn([20, 21, 22, 23], [24]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [20, 21] + assert segment.response_ids == [22, 23, 24] + assert segment.loss_mask == [1, 0, 1] + assert segment.rollout_log_probs == [-0.22, 0.0, -0.24] + + +@pytest.mark.unit +def test_merge_turn_segments_keeps_oversized_segments(): + segments = [TurnSegment(turns=[_turn([10, 11, 12], [13, 14])])] + + merged = merge_turn_segments(segments) + + assert len(merged) == 1 + assert merged[0].prompt_ids == [10, 11, 12] + assert merged[0].response_ids == [13, 14] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/vime/agent/__init__.py b/vime/agent/__init__.py new file mode 100644 index 000000000..1048d52b8 --- /dev/null +++ b/vime/agent/__init__.py @@ -0,0 +1 @@ +"""Agent rollout building blocks.""" diff --git a/vime/agent/adapters/__init__.py b/vime/agent/adapters/__init__.py new file mode 100644 index 000000000..4f19a829f --- /dev/null +++ b/vime/agent/adapters/__init__.py @@ -0,0 +1,7 @@ +"""HTTP adapters for agent rollouts.""" + +from vime.agent.adapters.anthropic import AnthropicAdapter +from vime.agent.adapters.common import BaseAdapter +from vime.agent.adapters.openai import OpenAIAdapter + +__all__ = ["AnthropicAdapter", "BaseAdapter", "OpenAIAdapter"] diff --git a/vime/agent/adapters/anthropic.py b/vime/agent/adapters/anthropic.py new file mode 100644 index 000000000..f0a85f5c0 --- /dev/null +++ b/vime/agent/adapters/anthropic.py @@ -0,0 +1,481 @@ +"""Anthropic Messages adapter for agent rollouts. + +The adapter exposes ``/v1/messages`` and ``/v1/messages/count_tokens``. It +renders each Anthropic message history with the served model's chat template, +calls vLLM's ``/inference/v1/generate`` with ``token_ids``, and +records the exact sampled token ids/logprobs as ``TurnRecord`` objects. New +code should use ``AnthropicAdapter`` and call ``finish_session()`` at trajectory +end to drain trainable ``TokenSegment`` objects. + +It also handles Claude Code sub-agent and compaction patterns by splitting one +session into ``subagent``, ``wipe``, and ``final`` segments. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import secrets +from typing import Any + +from aiohttp import web + +from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY +from vime.agent.adapters.common import AdapterChain as Chain +from vime.agent.adapters.common import ( + BaseAdapter, + call_vllm_generate, + ok_response, + render_token_ids, + request_session_id, +) +from vime.agent.adapters.common import stable_hash as _hash +from vime.agent.parsing import parse_model_output +from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments + +logger = logging.getLogger(__name__) + + +# Tool names claude-code uses to dispatch a sub-agent. +_SUBAGENT_TOOLS = {"Task", "Agent"} + + +@dataclasses.dataclass +class Session: + main: Chain = dataclasses.field(default_factory=Chain) + active_sub: Chain | None = None # at most one sub-agent at a time + pending_dispatch_id: str = "" # tool_use_id we're waiting to close + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + segments: list[TurnSegment] = dataclasses.field(default_factory=list) # frozen output + + +class AnthropicAdapter(BaseAdapter): + """Anthropic Messages-compatible HTTP adapter with session lifecycle helpers.""" + + session_cls = Session + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + super().__init__( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + ) + self.app.router.add_post("/v1/messages", _handle_request) + self.app.router.add_post("/v1/messages/count_tokens", _count_tokens) + self.app.router.add_get("/healthz", _ok) + self.app.router.add_get("/v1/models", _ok) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + await self.shutdown_session(sid, wait_timeout=wait_timeout) + s = self.store.pop(sid, None) + if s is None: + return [] + if s.active_sub is not None and s.active_sub.turns: + s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) + if s.main.turns: + s.segments.append(make_turn_segment(s.main.turns, kind="final")) + + return merge_turn_segments(s.segments) + + +# ============================================================================= +# 2. Per-turn stages +# ============================================================================= + + +def _select_chain(s: Session, body: dict) -> tuple[Chain, bool, str]: + """Decide which chain this turn operates on. + + 1. fingerprint body.messages and body.system into hashes + 2. if main now contains the tool_result for a pending sub dispatch, + snapshot the sub chain into s.segments and clear s.active_sub + 3. pick main vs s.active_sub based on whether request continues main's prefix + 4. classify as 'new' | 'append' | 'wipe' against the chosen target; + a wipe also snapshots the target's current state into s.segments + + Returns (target_chain, is_sub, kind). + """ + all_msgs = body.get("messages") or [] + msg_hashes = [_hash(m) for m in all_msgs] + req_system_hash = _hash(body.get("system")) if "system" in body else s.main.system_hash + + # Close active sub-agent if its dispatch tool_result has landed on main. + if s.pending_dispatch_id and s.active_sub is not None: + tu_id = s.pending_dispatch_id + for m in all_msgs: + if not isinstance(m, dict) or m.get("role") != "user": + continue + content = m.get("content") + if not isinstance(content, list): + continue + done = any( + isinstance(b, dict) and b.get("type") == "tool_result" and b.get("tool_use_id") == tu_id + for b in content + ) + if done: + if s.active_sub.turns: + s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) + s.active_sub = None + s.pending_dispatch_id = "" + break + + # Route: main iff request continues main's prefix. Sub system_hash can be + # "" (armed before sub dialled in), so never route by sub equality alone. + if s.active_sub is None: + target, is_sub = s.main, False + else: + main_continues = ( + req_system_hash == s.main.system_hash + and len(msg_hashes) >= s.main.seen_msgs + and msg_hashes[: s.main.seen_msgs] == s.main.msg_hashes[: s.main.seen_msgs] + ) + target, is_sub = (s.main, False) if main_continues else (s.active_sub, True) + + # Classify; snapshot a "wipe" segment first if we're discarding work. + if target.seen_msgs == 0: + kind = "new" + else: + is_append = ( + req_system_hash == target.system_hash + and len(msg_hashes) >= target.seen_msgs + and msg_hashes[: target.seen_msgs] == target.msg_hashes[: target.seen_msgs] + ) + if is_append: + kind = "append" + else: + if target.turns: + s.segments.append(make_turn_segment(target.turns, kind="wipe")) + kind = "wipe" + + return target, is_sub, kind + + +def _flatten(c: Any) -> str: + """Recursive Anthropic content flattener: text/tool_result(content) joined + by newline, images replaced with a placeholder.""" + if c is None: + return "" + if isinstance(c, str): + return c + if not isinstance(c, list): + return str(c) + parts: list[str] = [] + for b in c: + if isinstance(b, dict): + t = b.get("type") + if t == "text": + parts.append(b.get("text", "")) + elif t == "tool_result": + parts.append(_flatten(b.get("content"))) + elif t == "image": + parts.append("[image omitted]") + elif isinstance(b, str): + parts.append(b) + return "\n".join(p for p in parts if p) + + +def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: + """Anthropic messages + system -> chat-template messages. Pure function.""" + translated: list[dict] = [] + if system: + translated.append({"role": "system", "content": _flatten(system)}) + for m in msgs: + if not isinstance(m, dict): + continue + role, content = m.get("role"), m.get("content") + if role == "user": + blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + for b in blocks: + if isinstance(b, dict) and b.get("type") == "tool_result": + translated.append({"role": "tool", "content": _flatten(b.get("content"))}) + elif isinstance(b, dict) and b.get("type") == "text": + translated.append({"role": "user", "content": b.get("text", "")}) + else: + translated.append({"role": "user", "content": _flatten(b)}) + elif role == "assistant": + texts, thinkings, tcs = [], [], [] + blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + for b in blocks: + if not isinstance(b, dict): + continue + if b.get("type") == "text": + texts.append(b.get("text", "")) + elif b.get("type") == "thinking": + thinkings.append(b.get("thinking", "")) + elif b.get("type") == "tool_use": + tcs.append({"function": {"name": b.get("name", "tool"), "arguments": b.get("input") or {}}}) + mo: dict[str, Any] = {"role": "assistant", "content": "".join(texts)} + if thinkings: + mo["reasoning_content"] = "".join(thinkings) + if tcs: + mo["tool_calls"] = tcs + translated.append(mo) + elif role == "system": + translated.append({"role": "system", "content": _flatten(content)}) + return translated + + +def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] | None: + """Convert Anthropic tools to tokenizer chat-template tool schema.""" + if not anth_tools: + return None + ts: list[dict] = [] + for t in anth_tools: + if not isinstance(t, dict) or "name" not in t: + continue + ts.append( + { + "type": "function", + "function": { + "name": t["name"], + "description": t.get("description", ""), + "parameters": t.get("input_schema") or t.get("parameters") or {"type": "object", "properties": {}}, + }, + } + ) + return ts or None + + +def _replace_chat_messages(target: Chain, body: dict) -> None: + """new/wipe: full reset of chat state and turn log.""" + all_msgs = body.get("messages") or [] + target.chat_messages = _translate_anthropic(all_msgs, body.get("system")) + if "system" in body: + target.system_hash = _hash(body.get("system")) + target.turns.clear() + target.seen_msgs = len(all_msgs) + target.msg_hashes = [_hash(m) for m in all_msgs] + if target.tools_schema is None: + target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) + + +def _extend_chat_messages(target: Chain, body: dict) -> None: + """append: translate only the new tail.""" + all_msgs = body.get("messages") or [] + translated = _translate_anthropic(all_msgs[target.seen_msgs :], None) + target.chat_messages.extend(translated) + + target.seen_msgs = len(all_msgs) + target.msg_hashes = [_hash(m) for m in all_msgs] + if target.tools_schema is None: + target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) + + +def _build_prompt(target: Chain, body: dict, kind: str, tok) -> list[int]: + """Replace/extend chat_messages and render input ids for vLLM.""" + (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, body) + return render_token_ids(target, tok) + + +async def _generate( + prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None +) -> TurnRecord: + """Call vLLM and return a TurnRecord. + + 1. build sampling_params (session defaults overlaid with body overrides) + 2. POST vLLM ``/inference/v1/generate``; on cancel/error tear down + the request (vLLM has no per-request HTTP abort endpoint) + 3. keep the exact prompt/output token ids; trajectory merge later compares + later prompt tokens with earlier outputs to build the loss mask + """ + return await call_vllm_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_tokens",), + stop_keys=("stop_sequences",), + log_prefix="anthropic_adapter", + logger=logger, + session_id=session_id, + ) + + +def _build_reply(target: Chain, output_ids: list[int], finish: str, app) -> tuple[list[dict], str, str]: + """Turn the model's raw output ids into the reply we send back to claude-code. + + 1. parse decoded text -> (thinking, visible, tool_uses) via parsers + 2. pack into Anthropic content blocks; tag dispatch_id when a tool_use + names Task/Agent (sub-agent trigger) + 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' + + Returns (blocks, stop_reason, dispatch_id). + """ + tok = app[TOKENIZER_KEY] + + raw_output = tok.decode(output_ids, skip_special_tokens=False) if output_ids else "" + parsed = parse_model_output( + raw_output, + tools_schema=target.tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + blocks, dispatch_id = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) + return blocks, _stop_reason(parsed.tool_uses, finish), dispatch_id + + +def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> tuple[list[dict], str]: + """Pack parsed model output into Anthropic content blocks.""" + blocks: list[dict] = [] + if thinking: + blocks.append({"type": "thinking", "thinking": thinking}) + if visible: + blocks.append({"type": "text", "text": visible}) + dispatch_id = "" + for tu in tool_uses: + tu_id = f"toolu_{secrets.token_hex(8)}" + blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) + if tu["name"] in _SUBAGENT_TOOLS: + dispatch_id = tu_id + if not blocks: + blocks.append({"type": "text", "text": ""}) + return blocks, dispatch_id + + +def _stop_reason(tool_uses: list[dict], finish: str) -> str: + if tool_uses: + return "tool_use" + if finish == "length": + return "max_tokens" + return "end_turn" + + +def _start_sub_chain(s: Session, dispatch_id: str) -> None: + """Start a fresh sub chain on this session and remember the tool_use_id + we'll watch for on main to know when this sub is done. The matching + 'sub done' step lives inside _select_chain.""" + s.pending_dispatch_id = dispatch_id + if s.active_sub is None: + s.active_sub = Chain() + + +# ============================================================================= +# 3. Request handling -- one full turn + SSE wrap +# ============================================================================= + + +def _request_session_id(request: web.Request) -> str: + return request_session_id(request, include_x_api_key=True) + + +async def _handle_request(request: web.Request) -> web.StreamResponse: + body = await request.json() + sid = _request_session_id(request) + adapter = request.app[ADAPTER_KEY] + if sid in adapter.closed: # session drained; refuse stragglers + return web.Response(status=503, text="session closed") + app = request.app + s = adapter.store.setdefault(sid, Session()) + task = asyncio.current_task() + adapter.inflight.setdefault(sid, set()).add(task) + try: + async with s.lock: # same sid -> serialized + target, is_sub, kind = _select_chain(s, body) + ideal_ids = _build_prompt(target, body, kind, app[TOKENIZER_KEY]) + turn = await _generate(ideal_ids, s, body, app, session_id=sid) + blocks, stop, did = _build_reply(target, turn.output_ids, turn.finish_reason, app) + target.turns.append(turn) + if did and not is_sub: # sub doesn't nest + _start_sub_chain(s, did) + in_tok, out_tok = len(ideal_ids), len(turn.output_ids) + if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): + return await _stream_response(request, blocks, stop, in_tok, out_tok) + return web.json_response(_message_response(body, blocks, stop, in_tok, out_tok)) + finally: + adapter.inflight.get(sid, set()).discard(task) + + +def _message_response(body: dict, blocks: list[dict], stop_reason: str, in_tok: int, out_tok: int) -> dict: + return { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "model": body.get("model", "vime-actor"), + "content": blocks, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + + +async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web.StreamResponse: + """Stream blocks back to claude-code as an Anthropic Messages SSE + response: message_start, (content_block_start, content_block_delta, + content_block_stop)*N, message_delta, message_stop.""" + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + + # message_start + ms_data = { + "type": "message_start", + "message": { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "model": "vime-actor", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": in_tok, "output_tokens": 0}, + }, + } + await out.write(f"event: message_start\ndata: {json.dumps(ms_data, ensure_ascii=False)}\n\n".encode()) + + for idx, block in enumerate(blocks): + bt = block["type"] + if bt == "thinking": + start = {"type": "thinking", "thinking": ""} + delta = {"type": "thinking_delta", "thinking": block["thinking"]} + elif bt == "text": + start = {"type": "text", "text": ""} + delta = {"type": "text_delta", "text": block["text"]} + else: # tool_use + start = {"type": "tool_use", "id": block["id"], "name": block["name"], "input": {}} + delta = { + "type": "input_json_delta", + "partial_json": json.dumps(block["input"], ensure_ascii=False), + } + + cbs_data = {"type": "content_block_start", "index": idx, "content_block": start} + await out.write(f"event: content_block_start\ndata: {json.dumps(cbs_data, ensure_ascii=False)}\n\n".encode()) + + cbd_data = {"type": "content_block_delta", "index": idx, "delta": delta} + await out.write(f"event: content_block_delta\ndata: {json.dumps(cbd_data, ensure_ascii=False)}\n\n".encode()) + + cbe_data = {"type": "content_block_stop", "index": idx} + await out.write(f"event: content_block_stop\ndata: {json.dumps(cbe_data, ensure_ascii=False)}\n\n".encode()) + + md_data = { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + await out.write(f"event: message_delta\ndata: {json.dumps(md_data, ensure_ascii=False)}\n\n".encode()) + + mst_data = {"type": "message_stop"} + await out.write(f"event: message_stop\ndata: {json.dumps(mst_data, ensure_ascii=False)}\n\n".encode()) + + return out + + +# Trivial endpoints claude-code probes during a session: count_tokens runs +# every turn (return 0 -- client uses it as a hint, not a hard budget), +# healthz/v1/models are startup readiness checks. +async def _count_tokens(request: web.Request) -> web.Response: + return web.json_response({"input_tokens": 0}) + + +async def _ok(request: web.Request) -> web.Response: + return await ok_response(request) diff --git a/vime/agent/adapters/common.py b/vime/agent/adapters/common.py new file mode 100644 index 000000000..14fb2e850 --- /dev/null +++ b/vime/agent/adapters/common.py @@ -0,0 +1,308 @@ +"""Shared adapter primitives for token-capturing agent rollouts.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import hashlib +import json +import logging +import uuid +from collections.abc import Callable +from typing import Any + +import aiohttp +from aiohttp import web + +from vime.agent.trajectory import TokenSegment, TurnRecord + + +ADAPTER_KEY = web.AppKey("adapter", object) +TOKENIZER_KEY = web.AppKey("tokenizer", object) +VLLM_URL_KEY = web.AppKey("vllm_url", object) +TOOL_PARSER_KEY = web.AppKey("tool_parser", object) +REASONING_PARSER_KEY = web.AppKey("reasoning_parser", object) + + +@dataclasses.dataclass +class AdapterChain: + """Protocol-neutral chat chain state used by HTTP adapters.""" + + system_hash: str = "" + chat_messages: list[dict] = dataclasses.field(default_factory=list) + tools_schema: list[dict] | None = None + seen_msgs: int = 0 + msg_hashes: list[str] = dataclasses.field(default_factory=list) + turns: list[TurnRecord] = dataclasses.field(default_factory=list) + + +class BaseAdapter: + """Base HTTP adapter with per-instance session lifecycle state.""" + + session_cls: type + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + self.store: dict[str, Any] = {} + self.inflight: dict[str, set[asyncio.Task]] = {} + self.closed: set[str] = set() + self.app = web.Application(client_max_size=64 * 1024 * 1024) + self.app[ADAPTER_KEY] = self + self.app[TOKENIZER_KEY] = tokenizer + self.app[VLLM_URL_KEY] = vllm_url.rstrip("/") if isinstance(vllm_url, str) else vllm_url + self.app[TOOL_PARSER_KEY] = tool_parser + self.app[REASONING_PARSER_KEY] = reasoning_parser + + def open_session( + self, + sid: str, + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, + ) -> None: + register_session( + self.store, + sid, + self.session_cls, + sampling_defaults=sampling_defaults, + max_context_tokens=max_context_tokens, + ) + + async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: + await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + raise NotImplementedError + + +def strip_cache_control(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: strip_cache_control(v) for k, v in obj.items() if k != "cache_control"} + if isinstance(obj, list): + return [strip_cache_control(x) for x in obj] + return obj + + +def stable_hash(obj: Any) -> str: + payload = json.dumps(strip_cache_control(obj), sort_keys=True, ensure_ascii=False, default=str).encode("utf-8") + return hashlib.sha1(payload).hexdigest()[:12] + + +def json_arguments(value: Any) -> str: + if value is None: + return "{}" + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def render_token_ids(chain: AdapterChain, tokenizer) -> list[int]: + enc = tokenizer.apply_chat_template( + chain.chat_messages, + tools=chain.tools_schema, + tokenize=True, + add_generation_prompt=True, + ) + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) + + +def request_session_id( + request: web.Request, + *, + body: dict | None = None, + include_x_api_key: bool = False, +) -> str: + auth = request.headers.get("Authorization", "") + if auth.lower().startswith("bearer "): + sid = auth[7:].strip() + if sid: + return sid + + if body is not None: + metadata = body.get("metadata") + if isinstance(metadata, dict) and metadata.get("session_id"): + return str(metadata["session_id"]) + if body.get("user"): + return str(body["user"]) + + if include_x_api_key: + api_key = request.headers.get("X-Api-Key") + if api_key: + return api_key.strip() + + return "default" + + +def register_session( + store: dict[str, Any], + sid: str, + session_factory: Callable[[], Any], + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, +) -> None: + if sid in store: + raise ValueError(f"session_id {sid!r} already exists; sids must be unique per agent run") + session = store[sid] = session_factory() + session.sampling_defaults = dict(sampling_defaults or {}) + session.max_context_tokens = int(max_context_tokens or 0) + + +def _sampling_params(session: Any, body: dict, *, max_token_keys: tuple[str, ...], stop_keys: tuple[str, ...]) -> dict: + sp: dict[str, Any] = { + "skip_special_tokens": False, + "spaces_between_special_tokens": False, + "no_stop_trim": True, + "max_new_tokens": 4096, + **(session.sampling_defaults or {}), + } + + for key in max_token_keys: + if body.get(key) is not None: + sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", body[key])), int(body[key])) + break + + for src_k, dst_k in (("temperature", "temperature"), ("top_p", "top_p"), ("top_k", "top_k")): + if src_k in body: + sp[dst_k] = body[src_k] + + for key in stop_keys: + if body.get(key): + sp["stop"] = body[key] + break + + return sp + + +def _vllm_sampling_body(sp: dict) -> dict: + """Map the canonical (sglang-shaped) sampling dict to a vLLM ``/inference/v1/generate`` + ``sampling_params`` body. vLLM uses ``max_tokens`` (not ``max_new_tokens``) and returns + per-token logprobs when ``logprobs`` is set.""" + body: dict[str, Any] = { + "max_tokens": int(sp.get("max_new_tokens", 4096)), + "logprobs": 1, + } + if "temperature" in sp: + body["temperature"] = sp["temperature"] + if "top_p" in sp: + body["top_p"] = sp["top_p"] + tk = sp.get("top_k") + if tk is not None and (tk > 0 or tk == -1): + body["top_k"] = tk + if sp.get("stop"): + body["stop"] = sp["stop"] + if sp.get("stop_token_ids"): + body["stop_token_ids"] = sp["stop_token_ids"] + if sp.get("skip_special_tokens") is not None: + body["skip_special_tokens"] = bool(sp["skip_special_tokens"]) + return body + + +def _tokens_and_logprobs_from_choice(choice: dict) -> tuple[list[int], list[float]]: + """Parse ``token_ids`` + ``logprobs.content[i].logprob`` from a vLLM + ``/inference/v1/generate`` choice. Mirrors vime ``_inference_generate_tokens_and_logprobs``.""" + tids_raw = choice.get("token_ids") + if not (isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw)): + return [], [] + tids = [int(x) for x in tids_raw] + lp = choice.get("logprobs") + if not isinstance(lp, dict): + return tids, [0.0] * len(tids) + content = lp.get("content") + if isinstance(content, list) and content: + lps: list[float] = [] + for i in range(len(tids)): + if i < len(content) and isinstance(content[i], dict): + lps.append(float(content[i].get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + return tids, [0.0] * len(tids) + + +async def call_vllm_generate( + prompt_ids: list[int], + session: Any, + body: dict, + app, + *, + max_token_keys: tuple[str, ...], + stop_keys: tuple[str, ...], + log_prefix: str, + logger: logging.Logger, + session_id: str | None = None, +) -> TurnRecord: + sp = _sampling_params(session, body, max_token_keys=max_token_keys, stop_keys=stop_keys) + + if session.max_context_tokens > 0: + remaining_context = session.max_context_tokens - len(prompt_ids) + if remaining_context <= 0: + logger.warning( + "[%s] prompt exceeds max_context_tokens (%d >= %d)", + log_prefix, + len(prompt_ids), + session.max_context_tokens, + ) + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[], finish_reason="length") + sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", remaining_context)), remaining_context) + + vllm_url = app[VLLM_URL_KEY] + payload: dict[str, Any] = { + "token_ids": list(prompt_ids), + "sampling_params": _vllm_sampling_body(sp), + } + # session_id routes via vllm-router's consistent_hash policy (x-session-id header); + # see vime ``vllm_rollout.py`` headers handling. + headers = {"x-session-id": session_id} if session_id and session_id != "default" else None + timeout = aiohttp.ClientTimeout(total=None, sock_read=900) + task = asyncio.current_task() + try: + async with aiohttp.ClientSession(timeout=timeout) as sess, sess.post( + f"{vllm_url}/inference/v1/generate", + json=payload, + headers=headers, + ) as r: + if r.status >= 400: + text = await r.text() + raise RuntimeError(f"vllm upstream {r.status}: {text[:400]}") + data = await r.json(content_type=None) + choice = (data.get("choices") or [{}])[0] + output_ids, output_log_probs = _tokens_and_logprobs_from_choice(choice) + fr = choice.get("finish_reason") + finish = fr if isinstance(fr, str) and fr else "stop" + except (asyncio.CancelledError, aiohttp.ClientError, asyncio.TimeoutError): + # vLLM ``/inference/v1/generate`` has no per-request HTTP abort endpoint (unlike + # sglang ``/abort_request``). Cancelling the in-flight task tears down the aiohttp + # request, which drops the streaming connection so vLLM stops generating. + if task is not None: + task.cancel() + raise + + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=output_ids, + finish_reason=finish, + output_log_probs=output_log_probs, + ) + + +async def shutdown_session_tasks( + sid: str, + closed: set[str], + inflight: dict[str, set[asyncio.Task]], + *, + wait_timeout: float = 5.0, +) -> None: + closed.add(sid) + tasks = [t for t in inflight.pop(sid, ()) if not t.done()] + if not tasks: + return + _, pending = await asyncio.wait(tasks, timeout=wait_timeout) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + +async def ok_response(request: web.Request) -> web.Response: + return web.json_response({"ok": True}) diff --git a/vime/agent/adapters/openai.py b/vime/agent/adapters/openai.py new file mode 100644 index 000000000..5042d1dd6 --- /dev/null +++ b/vime/agent/adapters/openai.py @@ -0,0 +1,572 @@ +"""OpenAI-compatible adapters for agent rollouts. + +The adapter exposes ``/v1/chat/completions`` and ``/v1/responses``. Both +endpoints render incoming messages with the served model's chat template, call +vLLM's ``/inference/v1/generate`` with ``token_ids``, and record +the exact sampled token ids/logprobs as ``TurnRecord`` objects. New code should +use ``OpenAIAdapter`` and call ``finish_session()`` at trajectory end to drain +trainable ``TokenSegment`` objects. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import secrets +import time +from typing import Any + +from aiohttp import web + +from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY +from vime.agent.adapters.common import AdapterChain as Chain +from vime.agent.adapters.common import BaseAdapter, call_vllm_generate +from vime.agent.adapters.common import json_arguments as _json_arguments +from vime.agent.adapters.common import ok_response, render_token_ids, request_session_id +from vime.agent.adapters.common import stable_hash as _hash +from vime.agent.parsing import ParsedModelOutput, parse_model_output +from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Session: + main: Chain = dataclasses.field(default_factory=Chain) + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + segments: list[TurnSegment] = dataclasses.field(default_factory=list) + + +class OpenAIAdapter(BaseAdapter): + """OpenAI-compatible HTTP adapter with session lifecycle helpers.""" + + session_cls = Session + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + super().__init__( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + ) + self.app.router.add_post("/v1/chat/completions", _handle_chat_completions) + self.app.router.add_post("/v1/responses", _handle_responses) + self.app.router.add_get("/healthz", _ok) + self.app.router.add_get("/v1/models", _ok) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + await self.shutdown_session(sid, wait_timeout=wait_timeout) + s = self.store.pop(sid, None) + if s is None: + return [] + if s.main.turns: + s.segments.append(make_turn_segment(s.main.turns, kind="final")) + return merge_turn_segments(s.segments) + + +def _flatten_content(content: Any) -> str: + """Flatten OpenAI text/content parts into a chat-template string.""" + if content is None: + return "" + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) + + parts: list[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + continue + if not isinstance(item, dict): + parts.append(str(item)) + continue + typ = item.get("type") + if typ in {"text", "input_text", "output_text"}: + parts.append(item.get("text", "")) + elif typ in {"image_url", "input_image"}: + parts.append("[image omitted]") + elif "content" in item: + parts.append(_flatten_content(item.get("content"))) + elif "text" in item: + parts.append(str(item.get("text") or "")) + return "\n".join(p for p in parts if p) + + +def _normalize_tool_call(call: dict[str, Any]) -> dict[str, Any]: + function = call.get("function") or {} + name = function.get("name") or call.get("name") or "tool" + arguments = function.get("arguments", call.get("arguments", {})) + out = { + "type": "function", + "function": { + "name": name, + "arguments": _json_arguments(arguments), + }, + } + if call.get("id"): + out["id"] = call["id"] + return out + + +def _translate_chat_messages(messages: list[dict]) -> list[dict]: + """OpenAI chat messages -> tokenizer chat-template messages.""" + translated: list[dict] = [] + for msg in messages: + if not isinstance(msg, dict): + continue + role = msg.get("role") + content = msg.get("content") + if role == "developer": + role = "system" + + if role in {"system", "user"}: + translated.append({"role": role, "content": _flatten_content(content)}) + elif role == "tool": + tool_msg = {"role": "tool", "content": _flatten_content(content)} + if msg.get("tool_call_id"): + tool_msg["tool_call_id"] = msg["tool_call_id"] + translated.append(tool_msg) + elif role == "assistant": + assistant: dict[str, Any] = {"role": "assistant", "content": _flatten_content(content)} + if msg.get("reasoning_content"): + assistant["reasoning_content"] = msg["reasoning_content"] + tool_calls = msg.get("tool_calls") or [] + if tool_calls: + assistant["tool_calls"] = [_normalize_tool_call(c) for c in tool_calls if isinstance(c, dict)] + translated.append(assistant) + return translated + + +def _normalize_tool(tool: dict[str, Any]) -> dict[str, Any] | None: + if not isinstance(tool, dict): + return None + if tool.get("type") != "function": + return None + if isinstance(tool.get("function"), dict): + function = tool["function"] + name = function.get("name") + if not name: + return None + return { + "type": "function", + "function": { + "name": name, + "description": function.get("description", ""), + "parameters": function.get("parameters") or {"type": "object", "properties": {}}, + }, + } + name = tool.get("name") + if not name: + return None + return { + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, + }, + } + + +def _normalize_tools(tools: list[dict] | None) -> list[dict] | None: + normalized = [_normalize_tool(t) for t in tools or []] + return [t for t in normalized if t is not None] or None + + +def _responses_input_to_messages(input_value: Any, instructions: Any = None) -> list[dict]: + """Responses API input -> OpenAI chat message list. + + This intentionally covers the common message/function-call shapes used by + agent SDKs. Unknown input items are preserved as user text where possible. + """ + messages: list[dict] = [] + if instructions: + messages.append({"role": "system", "content": _flatten_content(instructions)}) + + if isinstance(input_value, str): + messages.append({"role": "user", "content": input_value}) + return messages + + if not isinstance(input_value, list): + messages.append({"role": "user", "content": _flatten_content(input_value)}) + return messages + + for item in input_value: + if isinstance(item, str): + messages.append({"role": "user", "content": item}) + continue + if not isinstance(item, dict): + messages.append({"role": "user", "content": str(item)}) + continue + + typ = item.get("type") + if typ == "function_call_output": + messages.append( + { + "role": "tool", + "tool_call_id": item.get("call_id") or item.get("id") or "", + "content": item.get("output", ""), + } + ) + elif typ == "function_call": + messages.append( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": item.get("call_id") or item.get("id") or f"call_{secrets.token_hex(8)}", + "type": "function", + "function": { + "name": item.get("name", "tool"), + "arguments": item.get("arguments", "{}"), + }, + } + ], + } + ) + elif item.get("role"): + messages.append({"role": item.get("role"), "content": item.get("content", "")}) + elif typ == "message": + messages.append({"role": item.get("role", "user"), "content": item.get("content", "")}) + else: + messages.append({"role": "user", "content": _flatten_content(item)}) + return messages + + +def _select_kind(s: Session, messages: list[dict]) -> str: + target = s.main + msg_hashes = [_hash(m) for m in messages] + if target.seen_msgs == 0: + kind = "new" + else: + is_append = len(msg_hashes) >= target.seen_msgs and msg_hashes[: target.seen_msgs] == target.msg_hashes + if is_append: + kind = "append" + else: + if target.turns: + s.segments.append(make_turn_segment(target.turns, kind="wipe")) + kind = "wipe" + return kind + + +def _replace_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: + target.chat_messages = _translate_chat_messages(messages) + target.turns.clear() + target.seen_msgs = len(messages) + target.msg_hashes = [_hash(m) for m in messages] + if tools_schema is not None: + target.tools_schema = tools_schema + + +def _extend_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: + translated = _translate_chat_messages(messages[target.seen_msgs :]) + target.chat_messages.extend(translated) + target.seen_msgs = len(messages) + target.msg_hashes = [_hash(m) for m in messages] + if tools_schema is not None: + target.tools_schema = tools_schema + + +def _build_prompt(target: Chain, messages: list[dict], tools_schema: list[dict] | None, kind: str, tok) -> list[int]: + (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, messages, tools_schema) + return render_token_ids(target, tok) + + +async def _generate( + prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None +) -> TurnRecord: + return await call_vllm_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_output_tokens", "max_completion_tokens", "max_tokens"), + stop_keys=("stop",), + log_prefix="openai_adapter", + logger=logger, + session_id=session_id, + ) + + +def _parse_turn(target: Chain, turn: TurnRecord, app) -> ParsedModelOutput: + tok = app[TOKENIZER_KEY] + raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" + return parse_model_output( + raw_output, + tools_schema=target.tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + + +def _openai_tool_calls(tool_uses: list[dict[str, Any]]) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + for tool_use in tool_uses: + call_id = f"call_{secrets.token_hex(12)}" + calls.append( + { + "id": call_id, + "type": "function", + "function": { + "name": tool_use.get("name", "tool"), + "arguments": _json_arguments(tool_use.get("input") or {}), + }, + } + ) + return calls + + +def _finish_reason(parsed: ParsedModelOutput, finish: str) -> str: + if parsed.tool_uses: + return "tool_calls" + if finish == "length": + return "length" + return "stop" + + +def _chat_message(parsed: ParsedModelOutput) -> dict[str, Any]: + tool_calls = _openai_tool_calls(parsed.tool_uses) + message: dict[str, Any] = { + "role": "assistant", + "content": parsed.text if parsed.text else None, + } + if parsed.reasoning: + message["reasoning_content"] = parsed.reasoning + if tool_calls: + message["tool_calls"] = tool_calls + return message + + +def _usage(in_tok: int, out_tok: int) -> dict[str, int]: + return { + "prompt_tokens": in_tok, + "completion_tokens": out_tok, + "total_tokens": in_tok + out_tok, + } + + +def _responses_usage(in_tok: int, out_tok: int) -> dict[str, int]: + return { + "input_tokens": in_tok, + "output_tokens": out_tok, + "total_tokens": in_tok + out_tok, + } + + +def _request_session_id(request: web.Request, body: dict) -> str: + return request_session_id(request, body=body) + + +async def _run_turn( + request: web.Request, body: dict, messages: list[dict] +) -> tuple[TurnRecord, ParsedModelOutput, int, int]: + sid = _request_session_id(request, body) + adapter = request.app[ADAPTER_KEY] + if sid in adapter.closed: + raise web.HTTPServiceUnavailable(text="session closed") + app = request.app + s = adapter.store.setdefault(sid, Session()) + task = asyncio.current_task() + adapter.inflight.setdefault(sid, set()).add(task) + try: + async with s.lock: + target = s.main + tools_schema = _normalize_tools(body.get("tools")) + kind = _select_kind(s, messages) + prompt_ids = _build_prompt(target, messages, tools_schema, kind, app[TOKENIZER_KEY]) + turn = await _generate(prompt_ids, s, body, app, session_id=sid) + parsed = _parse_turn(target, turn, app) + target.turns.append(turn) + return turn, parsed, len(prompt_ids), len(turn.output_ids) + finally: + adapter.inflight.get(sid, set()).discard(task) + + +async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: + body = await request.json() + messages = body.get("messages") or [] + if not isinstance(messages, list): + raise web.HTTPBadRequest(text="messages must be a list") + turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + if body.get("stream"): + return await _stream_chat_completion(request, body, parsed, turn.finish_reason, in_tok, out_tok) + return web.json_response(_chat_completion_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + + +def _chat_completion_response( + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> dict[str, Any]: + return { + "id": f"chatcmpl_{secrets.token_hex(12)}", + "object": "chat.completion", + "created": int(time.time()), + "model": body.get("model", "vime-actor"), + "choices": [ + { + "index": 0, + "message": _chat_message(parsed), + "finish_reason": _finish_reason(parsed, finish), + } + ], + "usage": _usage(in_tok, out_tok), + } + + +async def _stream_chat_completion( + request: web.Request, + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> web.StreamResponse: + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + completion_id = f"chatcmpl_{secrets.token_hex(12)}" + created = int(time.time()) + + async def emit(choice_delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": body.get("model", "vime-actor"), + "choices": [{"index": 0, "delta": choice_delta, "finish_reason": finish_reason}], + } + if usage is not None: + chunk["usage"] = usage + await out.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()) + + await emit({"role": "assistant"}) + if parsed.reasoning: + await emit({"reasoning_content": parsed.reasoning}) + if parsed.text: + await emit({"content": parsed.text}) + for idx, call in enumerate(_openai_tool_calls(parsed.tool_uses)): + await emit({"tool_calls": [{**call, "index": idx}]}) + await emit({}, finish_reason=_finish_reason(parsed, finish), usage=_usage(in_tok, out_tok)) + await out.write(b"data: [DONE]\n\n") + return out + + +async def _handle_responses(request: web.Request) -> web.StreamResponse: + body = await request.json() + messages = _responses_input_to_messages(body.get("input", ""), body.get("instructions")) + turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + if body.get("stream"): + return await _stream_response(request, body, parsed, turn.finish_reason, in_tok, out_tok) + return web.json_response(_response_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + + +def _response_output(parsed: ParsedModelOutput) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + if parsed.reasoning: + output.append( + { + "id": f"rs_{secrets.token_hex(12)}", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": parsed.reasoning}], + } + ) + if parsed.text: + output.append( + { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": parsed.text, "annotations": []}], + } + ) + for call in _openai_tool_calls(parsed.tool_uses): + output.append( + { + "id": f"fc_{secrets.token_hex(12)}", + "type": "function_call", + "status": "completed", + "call_id": call["id"], + "name": call["function"]["name"], + "arguments": call["function"]["arguments"], + } + ) + if not output: + output.append( + { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "", "annotations": []}], + } + ) + return output + + +def _response_response( + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> dict[str, Any]: + status = "incomplete" if finish == "length" else "completed" + return { + "id": f"resp_{secrets.token_hex(12)}", + "object": "response", + "created_at": int(time.time()), + "status": status, + "model": body.get("model", "vime-actor"), + "output": _response_output(parsed), + "usage": _responses_usage(in_tok, out_tok), + } + + +async def _stream_response( + request: web.Request, + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> web.StreamResponse: + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + response = _response_response(body, parsed, finish, in_tok, out_tok) + created = {"type": "response.created", "response": response} + await out.write(f"event: response.created\ndata: {json.dumps(created, ensure_ascii=False)}\n\n".encode()) + if parsed.text: + delta = {"type": "response.output_text.delta", "delta": parsed.text} + await out.write( + f"event: response.output_text.delta\ndata: {json.dumps(delta, ensure_ascii=False)}\n\n".encode() + ) + completed = {"type": "response.completed", "response": response} + await out.write(f"event: response.completed\ndata: {json.dumps(completed, ensure_ascii=False)}\n\n".encode()) + return out + + +async def _ok(request: web.Request) -> web.Response: + return await ok_response(request) diff --git a/vime/agent/parsing.py b/vime/agent/parsing.py new file mode 100644 index 000000000..db1913608 --- /dev/null +++ b/vime/agent/parsing.py @@ -0,0 +1,113 @@ +"""Model-output parsing helpers for agent harnesses.""" + +from __future__ import annotations + +import dataclasses +import json +import logging +import re +from typing import Any + + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class ParsedModelOutput: + """Structured view of one decoded model output.""" + + reasoning: str + text: str + tool_uses: list[dict[str, Any]] + + +def parse_model_output( + raw_output: str, + *, + tools_schema: list[dict] | None, + tool_parser_name: str | None, + reasoning_parser_name: str | None, +) -> ParsedModelOutput: + """Parse raw model text into reasoning, visible text, and tool uses. + + When ``reasoning_parser_name`` / ``tool_parser_name`` are set the heavy + format-specific work is delegated to SGLang's reasoning and function-call + parsers (imported lazily, so they are only required when explicitly + enabled). The coding-agent example leaves both unset and relies on the XML + fallback, which covers the Anthropic-style tool-call text that some + coding-agent models still emit occasionally. + """ + reasoning, body_text = "", raw_output + if reasoning_parser_name: + from sglang.srt.parser.reasoning_parser import ReasoningParser + + r, b = ReasoningParser( + model_type=reasoning_parser_name, + stream_reasoning=False, + ).parse_non_stream(raw_output) + reasoning, body_text = r or "", b or "" + if not reasoning and "" in body_text: + reasoning, body_text = body_text.split("", 1) + + body_text, tool_uses = parse_tool_uses(body_text, tools_schema, tool_parser_name) + return ParsedModelOutput( + reasoning=reasoning, + text=(body_text or "").strip(), + tool_uses=tool_uses, + ) + + +def parse_tool_uses( + body_text: str, + tools_schema: list[dict] | None, + tool_parser_name: str | None, +) -> tuple[str, list[dict[str, Any]]]: + """Parse tool calls from body text and return visible text plus tool uses.""" + tool_uses: list[dict[str, Any]] = [] + if tool_parser_name and tools_schema: + from sglang.srt.entrypoints.openai.protocol import Function, Tool + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + sg_tools = [Tool(type="function", function=Function(**d["function"])) for d in tools_schema] + parser = FunctionCallParser(tools=sg_tools, tool_call_parser=tool_parser_name) + calls = [] + if parser.has_tool_call(body_text): + try: + body_text, calls = parser.parse_non_stream(body_text) + except Exception: + logger.exception("[agent.parsing] sglang tool-call parsing failed; falling back") + for c in calls: + try: + args = json.loads(c.parameters or "{}") + except json.JSONDecodeError: + args = {"_raw_arguments": c.parameters} + tool_uses.append({"name": c.name or "tool", "input": args}) + + if not tool_uses and tools_schema: + body_text, tool_uses = parse_xml_tool_uses(body_text, tools_schema) + + return body_text, tool_uses + + +def parse_xml_tool_uses(body_text: str, tools_schema: list[dict]) -> tuple[str, list[dict[str, Any]]]: + """Fallback parser for Anthropic-style XML tool calls.""" + valid_tools = {t.get("function", {}).get("name") for t in tools_schema} + tool_uses: list[dict[str, Any]] = [] + cleaned_parts: list[str] = [] + last = 0 + for m in re.finditer( + r"\s*]+)>(.*?)\s*", + body_text, + flags=re.DOTALL, + ): + name, inner = m.group(1), m.group(2) + if name in valid_tools: + args = { + p.group(1): p.group(2).strip() + for p in re.finditer(r"]+)>(.*?)", inner, flags=re.DOTALL) + } + tool_uses.append({"name": name, "input": args}) + cleaned_parts.append(body_text[last : m.start()]) + last = m.end() + cleaned_parts.append(body_text[last:]) + return "".join(cleaned_parts), tool_uses diff --git a/vime/agent/sandbox.py b/vime/agent/sandbox.py new file mode 100644 index 000000000..6447bd2ee --- /dev/null +++ b/vime/agent/sandbox.py @@ -0,0 +1,281 @@ +"""Sandbox backends for agent rollouts. + +The public sandbox contract is intentionally small: async context management, +command execution, and file read/write. Agent examples can build task-specific +setup, runner, and evaluator logic on top of this without depending directly on +one sandbox provider. +""" + +from __future__ import annotations + +import asyncio +import io +import json +import logging +import os +from pathlib import Path +from typing import Protocol, runtime_checkable + +logger = logging.getLogger(__name__) + + +ExecResult = tuple[int, str, str] +FileContent = str | bytes | Path + + +@runtime_checkable +class Sandbox(Protocol): + """Minimal async sandbox interface used by agent rollouts. + + ``write_file`` accepts either in-memory content (``str``/``bytes``) or a + host ``Path`` to stream into the sandbox. + """ + + sandbox_id: str + + async def __aenter__(self) -> Sandbox: ... + + async def __aexit__(self, exc_type, exc, tb) -> None: ... + + async def exec( + self, + cmd: str, + *, + user: str = "root", + env: dict[str, str] | None = None, + timeout: int = 120, + check: bool = False, + ) -> ExecResult: ... + + async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: ... + + async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: ... + + +def _getenv(*names: str, default: str = "") -> str: + for name in names: + value = os.environ.get(name) + if value is not None and value.strip(): + return value + return default + + +class E2BSandbox: + """Async context manager around e2b.AsyncSandbox.""" + + metadata_file_env = ("VIME_AGENT_SANDBOX_METADATA_FILE", "SWE_SANDBOX_METADATA_FILE") + metadata_json_env = ("VIME_AGENT_SANDBOX_METADATA_JSON", "SWE_SANDBOX_METADATA_JSON") + image_metadata_key_env = ("VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY", "SWE_SANDBOX_IMAGE_METADATA_KEY") + lifetime_sec_env = ("VIME_AGENT_SANDBOX_LIFETIME_SEC", "SWE_SANDBOX_LIFETIME_SEC") + rpc_retries_env = ("VIME_AGENT_SANDBOX_RPC_RETRIES", "SWE_RPC_RETRIES") + + default_lifetime_sec = 3600 + default_rpc_retries = 3 + # With retries=3 the sleep budget is 3s, which handles common E2B h2 reset + # / SSL / pool-timeout flaps without stalling rollout steps for too long. + rpc_backoff_base_sec = 1.0 + + def __init__( + self, + image: str, + *, + timeout: int | None = None, + metadata: dict[str, str] | None = None, + image_metadata_key: str | None = None, + rpc_retries: int | None = None, + ) -> None: + self.image = image + self.timeout = timeout if timeout is not None else self._lifetime_sec_from_env() + self.metadata = dict(metadata) if metadata is not None else self._metadata_from_env() + self.image_metadata_key = image_metadata_key or self._image_metadata_key_from_env() + self.rpc_retries = rpc_retries if rpc_retries is not None else self._rpc_retries_from_env() + self._sb = None + self.sandbox_id = "" + + @classmethod + def _metadata_from_env(cls) -> dict[str, str]: + """Read E2B routing metadata from file or JSON environment values.""" + file_path = _getenv(*cls.metadata_file_env) + raw = "" + if file_path: + try: + raw = Path(file_path).read_text() + except OSError as e: + logger.warning("[agent.sandbox] metadata file %s unreadable: %s", file_path, e) + raw = "" + if not raw: + raw = _getenv(*cls.metadata_json_env) + if not raw: + return {} + try: + md = json.loads(raw) + except json.JSONDecodeError as e: + logger.warning("[agent.sandbox] metadata not valid JSON, ignoring: %s", e) + return {} + if not isinstance(md, dict): + logger.warning("[agent.sandbox] metadata must be a JSON object, got %s", type(md).__name__) + return {} + return {str(k): str(v) for k, v in md.items()} + + @classmethod + def _image_metadata_key_from_env(cls) -> str | None: + return _getenv(*cls.image_metadata_key_env) or None + + @classmethod + def _lifetime_sec_from_env(cls) -> int: + return int(_getenv(*cls.lifetime_sec_env, default=str(cls.default_lifetime_sec))) + + @classmethod + def _rpc_retries_from_env(cls) -> int: + return int(_getenv(*cls.rpc_retries_env, default=str(cls.default_rpc_retries))) + + @staticmethod + def _is_transient_rpc_error(e: BaseException) -> bool: + """True if e is a transient E2B client-side failure safe to retry.""" + name = type(e).__name__ + if name in { + "ProtocolError", + "LocalProtocolError", + "WriteError", + "ReadError", + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "WriteTimeout", + "PoolTimeout", + "RemoteProtocolError", + "SSLError", + }: + return True + msg = str(e) + if name == "SandboxException": + if "does not exist" in msg or "STOPPED state" in msg: + return False + return True + return False + + async def _rpc_retry(self, op_name: str, coro_factory): + """Run coro_factory() with retries for transient E2B RPC failures.""" + last_err = None + for attempt in range(self.rpc_retries): + try: + return await coro_factory() + except Exception as e: + if not self._is_transient_rpc_error(e): + raise + last_err = e + if attempt + 1 < self.rpc_retries: + backoff = self.rpc_backoff_base_sec * (2**attempt) + logger.debug( + "[agent.sandbox] %s transient %s, retry %d/%d in %.1fs: %s", + op_name, + type(e).__name__, + attempt + 1, + self.rpc_retries, + backoff, + str(e)[:120], + ) + await asyncio.sleep(backoff) + assert last_err is not None + raise last_err + + async def __aenter__(self) -> E2BSandbox: + if self.image_metadata_key is None: + raise RuntimeError( + "VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY is not set. Export it " + "to the metadata key your E2B gateway uses for image routing. " + "The legacy SWE_SANDBOX_IMAGE_METADATA_KEY name is also " + "accepted for coding-agent examples." + ) + from e2b import AsyncSandbox # type: ignore + + md = dict(self.metadata) + md.setdefault(self.image_metadata_key, self.image) + self._sb = await AsyncSandbox.create(timeout=self.timeout, metadata=md) + self.sandbox_id = self._sb.sandbox_id + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + try: + if self._sb is not None: + await self._sb.kill() + except Exception as e: + logger.warning("[agent.sandbox] kill %s failed: %s", self.sandbox_id[:8], e) + + async def exec( + self, + cmd: str, + *, + user: str = "root", + env: dict[str, str] | None = None, + timeout: int = 120, + check: bool = False, + ) -> ExecResult: + from e2b.sandbox.commands.command_handle import CommandExitException + + try: + res = await self._rpc_retry( + f"exec({cmd[:60]!r})", + lambda: self._sb.commands.run( + cmd, + user=user, + envs=env, + timeout=timeout, + on_stdout=lambda s: None, + on_stderr=lambda s: None, + ), + ) + return res.exit_code, res.stdout or "", res.stderr or "" + except CommandExitException as e: + if check: + raise RuntimeError( + f"e2b exec failed (exit={e.exit_code}): {cmd[:120]}\n{(e.stderr or '')[:400]}" + ) from None + return e.exit_code, e.stdout or "", e.stderr or "" + + async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: + if isinstance(content, Path): + host_path = content + + async def _do_path(): + with open(host_path, "rb") as fp: + await self._sb.files.write( + sandbox_path, + fp, + user=user, + gzip=False, + use_octet_stream=True, + request_timeout=600, + ) + + await self._rpc_retry(f"write_file({sandbox_path} <- {host_path.name})", _do_path) + return + + if isinstance(content, bytes): + + async def _do_bytes(): + await self._sb.files.write( + sandbox_path, + io.BytesIO(content), + user=user, + gzip=False, + use_octet_stream=True, + request_timeout=600, + ) + + await self._rpc_retry(f"write_file({sandbox_path}, bytes={len(content)})", _do_bytes) + return + + await self._rpc_retry( + f"write_file({sandbox_path})", + lambda: self._sb.files.write(sandbox_path, content, user=user), + ) + + async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: + try: + return await self._rpc_retry( + f"read_file({sandbox_path})", + lambda: self._sb.files.read(sandbox_path, user=user), + ) + except Exception: + return "" diff --git a/vime/agent/trajectory.py b/vime/agent/trajectory.py new file mode 100644 index 000000000..b30b69080 --- /dev/null +++ b/vime/agent/trajectory.py @@ -0,0 +1,208 @@ +"""Token-level trajectory helpers for agent rollouts.""" + +from __future__ import annotations + +import copy +import dataclasses +import logging +from typing import Any + +from vime.utils.types import Sample + + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class TurnRecord: + """Exact token snapshot for one assistant generation. + + ``prompt_ids`` is the full tokenized prompt sent to the generator for that + turn. ``output_ids`` is the raw generated output, and + ``output_log_probs`` is aligned with it when the rollout engine returns + per-token log probabilities. + """ + + prompt_ids: list[int] + output_ids: list[int] + finish_reason: str + output_log_probs: list[float] = dataclasses.field(default_factory=list) + + +@dataclasses.dataclass(frozen=True) +class TokenSegment: + """One training segment assembled from an agent trajectory.""" + + prompt_ids: list[int] + response_ids: list[int] + loss_mask: list[int] + rollout_log_probs: list[float] = dataclasses.field(default_factory=list) + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +@dataclasses.dataclass(frozen=True) +class TurnSegment: + """A frozen group of turns before token-level merge.""" + + turns: list[TurnRecord] + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +def make_turn_segment( + turns: list[TurnRecord], + *, + kind: str = "", + metadata: dict[str, Any] | None = None, +) -> TurnSegment: + """Freeze turns and attach conventional segment metadata.""" + frozen_turns = list(turns) + segment_metadata = dict(metadata or {}) + if kind: + segment_metadata.setdefault("segment_kind", kind) + segment_metadata.setdefault("finish_reason", frozen_turns[-1].finish_reason if frozen_turns else "") + return TurnSegment(turns=frozen_turns, metadata=segment_metadata) + + +def _common_prefix_len(a: list[int], b: list[int]) -> int: + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + +def _output_log_probs(turn: TurnRecord) -> list[float]: + if len(turn.output_log_probs) == len(turn.output_ids): + return list(turn.output_log_probs) + logger.warning( + "[trajectory] turn logprob length mismatch; zeroing output logprobs (%d ids, %d logprobs)", + len(turn.output_ids), + len(turn.output_log_probs), + ) + return [0.0] * len(turn.output_ids) + + +def merge_turns(turns: list[TurnRecord], *, metadata: dict[str, Any] | None = None) -> TokenSegment | None: + """Replay turn records into one linear training segment. + + The first turn's prompt becomes the segment prompt. Later turn prompts are + stitched against ``prompt + response_so_far``. Any new prompt suffix is + non-model context and receives loss mask 0. If a later prompt diverges + inside a previous model output, the retained prefix of that whole output + turn is also masked out, because partial token matches are not a faithful + training target for that turn. + """ + if not turns: + return None + + prompt_ids = list(turns[0].prompt_ids) + response_ids: list[int] = [] + loss_mask: list[int] = [] + rollout_log_probs: list[float] = [] + output_spans: list[tuple[int, int]] = [] + + for i, turn in enumerate(turns): + if i > 0: + if turn.prompt_ids[: len(prompt_ids)] != prompt_ids: + logger.warning("[trajectory] merge prompt base changed; starting segment from drifted prompt") + prompt_ids = list(turn.prompt_ids) + response_ids = [] + loss_mask = [] + rollout_log_probs = [] + output_spans = [] + else: + prompt_suffix = turn.prompt_ids[len(prompt_ids) :] + matched_len = _common_prefix_len(response_ids, prompt_suffix) + if matched_len < len(response_ids): + logger.warning( + "[trajectory] merge prefix drift; truncating %d unstitched response tokens", + len(response_ids) - matched_len, + ) + for start, end in output_spans: + if start < matched_len < end: + loss_mask[start:matched_len] = [0] * (matched_len - start) + rollout_log_probs[start:matched_len] = [0.0] * (matched_len - start) + response_ids = response_ids[:matched_len] + loss_mask = loss_mask[:matched_len] + rollout_log_probs = rollout_log_probs[:matched_len] + output_spans = [ + (start, min(end, matched_len)) for start, end in output_spans if start < matched_len + ] + + context_tail = prompt_suffix[matched_len:] + response_ids.extend(context_tail) + loss_mask.extend([0] * len(context_tail)) + rollout_log_probs.extend([0.0] * len(context_tail)) + + output_start = len(response_ids) + response_ids.extend(turn.output_ids) + loss_mask.extend([1] * len(turn.output_ids)) + rollout_log_probs.extend(_output_log_probs(turn)) + output_spans.append((output_start, len(response_ids))) + + rollout_log_probs = [logprob if mask else 0.0 for logprob, mask in zip(rollout_log_probs, loss_mask, strict=True)] + + return TokenSegment( + prompt_ids=prompt_ids, + response_ids=response_ids, + loss_mask=loss_mask, + rollout_log_probs=rollout_log_probs, + metadata=dict(metadata or {}), + ) + + +def merge_turn_segments(segments: list[TurnSegment]) -> list[TokenSegment]: + """Merge frozen turn segments and keep every non-empty output.""" + out: list[TokenSegment] = [] + for turn_segment in segments: + token_segment = merge_turns(turn_segment.turns, metadata=turn_segment.metadata) + if token_segment is None: + continue + if token_segment.response_ids: + out.append(token_segment) + return out + + +def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: + """Populate token, mask, response, reward, and status fields from a segment.""" + sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) + sample.response_length = len(segment.response_ids) + sample.loss_mask = list(segment.loss_mask) + sample.rollout_log_probs = list(segment.rollout_log_probs) + sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) + sample.reward = float(reward) + sample.status = Sample.Status.COMPLETED + + +def fan_out_sample_segments( + sample: Sample, + segments: list[TokenSegment], + reward: float, + tokenizer, + *, + metadata: dict[str, Any] | None = None, + rollout_id: int | None = None, +) -> list[Sample]: + """Emit one Sample per segment, splitting reward uniformly across them. + + Sibling samples share ``rollout_id`` so reducers that average by rollout do + not over-count trajectories split by compaction or sub-agent dispatch. + """ + k = len(segments) + per_segment_reward = float(reward) / max(1, k) + shared_rollout_id = getattr(sample, "index", None) if rollout_id is None else rollout_id + base_metadata = {**(sample.metadata or {}), **(metadata or {})} + + out: list[Sample] = [] + for i, segment in enumerate(segments): + sub = sample if i == 0 else copy.copy(sample) + write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) + sub.rollout_id = shared_rollout_id + sub.metadata = { + **base_metadata, + **(segment.metadata or {}), + "segment_idx": i, + "num_segments": k, + } + out.append(sub) + return out