From fcd036d48e0f24b63759a95f89c8ecfb537572b9 Mon Sep 17 00:00:00 2001 From: iam-dev Date: Thu, 26 Feb 2026 00:04:01 +0100 Subject: [PATCH 1/2] add claude code CLI client and unit tests Add ClaudeCodeCLI backend that shells out to `claude --print` for API-key-free usage. Includes 27 unit tests covering init, command building, prompt preparation, sync/async completion, usage tracking, and get_client integration. --- rlm/clients/__init__.py | 6 +- rlm/clients/claude_cli.py | 192 +++++++++++++++++++++ rlm/core/types.py | 1 + rlm/utils/token_utils.py | 4 + tests/clients/test_claude_cli.py | 288 +++++++++++++++++++++++++++++++ 5 files changed, 490 insertions(+), 1 deletion(-) create mode 100644 rlm/clients/claude_cli.py create mode 100644 tests/clients/test_claude_cli.py diff --git a/rlm/clients/__init__.py b/rlm/clients/__init__.py index 09d3008b6..53c9d7626 100644 --- a/rlm/clients/__init__.py +++ b/rlm/clients/__init__.py @@ -57,7 +57,11 @@ def get_client( from rlm.clients.azure_openai import AzureOpenAIClient return AzureOpenAIClient(**backend_kwargs) + elif backend == "claude_cli": + from rlm.clients.claude_cli import ClaudeCodeCLI + + return ClaudeCodeCLI(**backend_kwargs) else: raise ValueError( - f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai', 'gemini', 'vercel']" + f"Unknown backend: {backend}. Supported backends: ['openai', 'vllm', 'portkey', 'openrouter', 'litellm', 'anthropic', 'azure_openai', 'gemini', 'vercel', 'claude_cli']" ) diff --git a/rlm/clients/claude_cli.py b/rlm/clients/claude_cli.py new file mode 100644 index 000000000..650ca9a04 --- /dev/null +++ b/rlm/clients/claude_cli.py @@ -0,0 +1,192 @@ +import asyncio +import os +import shutil +import subprocess +from collections import defaultdict +from typing import Any + +from rlm.clients.base_lm import BaseLM +from rlm.core.types import ModelUsageSummary, UsageSummary + +# Characters-per-token estimate (CLI doesn't return usage data) +_CHARS_PER_TOKEN = 4 + + +class ClaudeCodeCLI(BaseLM): + """ + LM Client that shells out to the ``claude`` CLI (``claude --print``) + instead of calling an HTTP API. + + Designed for use inside cc-dirigent's RLM integration so that the full + RLM REPL loop can run without any API keys — Claude Code handles auth. + + Model names use CLI aliases: ``sonnet``, ``opus``, ``haiku``. + """ + + def __init__( + self, + model_name: str = "sonnet", + max_budget_usd: float | None = None, + permission_mode: str = "bypassPermissions", + **kwargs, + ): + super().__init__(model_name=model_name, **kwargs) + self.model_name = model_name + self.max_budget_usd = max_budget_usd + self.permission_mode = permission_mode + + # Verify ``claude`` is on PATH + if not shutil.which("claude"): + raise FileNotFoundError( + "claude CLI not found on PATH. Install Claude Code first: " + "https://docs.anthropic.com/en/docs/claude-code" + ) + + # Per-model usage tracking (estimated) + self.model_call_counts: dict[str, int] = defaultdict(int) + self.model_input_tokens: dict[str, int] = defaultdict(int) + self.model_output_tokens: dict[str, int] = defaultdict(int) + + self.last_prompt_tokens: int = 0 + self.last_completion_tokens: int = 0 + + # Build a clean env that allows nested claude invocations. + # CLAUDECODE env var blocks nested sessions — strip it. + self._env = {k: v for k, v in os.environ.items() if k != "CLAUDECODE"} + + # ── helpers ──────────────────────────────────────────── + + def _build_cmd(self, system: str | None, model: str | None) -> list[str]: + """Build the ``claude`` CLI argument list.""" + model = model or self.model_name + cmd = [ + "claude", + "--print", + "--model", model, + "--no-session-persistence", + "--permission-mode", self.permission_mode, + ] + if system: + cmd.extend(["--system-prompt", system]) + if self.max_budget_usd is not None: + cmd.extend(["--max-budget-usd", str(self.max_budget_usd)]) + return cmd + + def _prepare_prompt( + self, prompt: str | list[dict[str, Any]] + ) -> tuple[str, str | None]: + """ + Convert a prompt (string or message list) into ``(user_text, system)`` + for piping to ``claude --print``. + + Mirrors ``AnthropicClient._prepare_messages`` — system messages are + extracted and passed via ``--system-prompt``. + """ + system: str | None = None + + if isinstance(prompt, str): + return prompt, None + + if isinstance(prompt, list) and all(isinstance(m, dict) for m in prompt): + parts: list[str] = [] + for msg in prompt: + role = msg.get("role", "") + content = msg.get("content", "") + if role == "system": + system = content + else: + parts.append(str(content)) + return "\n\n".join(parts), system + + raise ValueError(f"Invalid prompt type: {type(prompt)}") + + @staticmethod + def _estimate_tokens(text: str) -> int: + return (len(text) + _CHARS_PER_TOKEN - 1) // _CHARS_PER_TOKEN + + def _track_usage(self, prompt_text: str, response_text: str, model: str) -> None: + input_tokens = self._estimate_tokens(prompt_text) + output_tokens = self._estimate_tokens(response_text) + + self.model_call_counts[model] += 1 + self.model_input_tokens[model] += input_tokens + self.model_output_tokens[model] += output_tokens + + self.last_prompt_tokens = input_tokens + self.last_completion_tokens = output_tokens + + # ── BaseLM interface ────────────────────────────────── + + def completion( + self, prompt: str | list[dict[str, Any]], model: str | None = None + ) -> str: + user_text, system = self._prepare_prompt(prompt) + model = model or self.model_name + cmd = self._build_cmd(system, model) + + # Uses subprocess.run with explicit argument list (no shell=True) + # to prevent command injection. User input is passed via stdin only. + result = subprocess.run( + cmd, + input=user_text, + capture_output=True, + text=True, + timeout=self.timeout, + env=self._env, + ) + + if result.returncode != 0: + raise RuntimeError( + f"claude CLI exited with code {result.returncode}: {result.stderr.strip()}" + ) + + response = result.stdout.strip() + self._track_usage(user_text, response, model) + return response + + async def acompletion( + self, prompt: str | list[dict[str, Any]], model: str | None = None + ) -> str: + user_text, system = self._prepare_prompt(prompt) + model = model or self.model_name + cmd = self._build_cmd(system, model) + + # Uses asyncio.create_subprocess_exec (not _shell) with explicit + # argument list to prevent command injection. Input via stdin only. + proc = await asyncio.create_subprocess_exec( + *cmd, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=self._env, + ) + stdout, stderr = await asyncio.wait_for( + proc.communicate(input=user_text.encode()), + timeout=self.timeout, + ) + + if proc.returncode != 0: + raise RuntimeError( + f"claude CLI exited with code {proc.returncode}: {stderr.decode().strip()}" + ) + + response = stdout.decode().strip() + self._track_usage(user_text, response, model) + return response + + def get_usage_summary(self) -> UsageSummary: + model_summaries = {} + for model in self.model_call_counts: + model_summaries[model] = ModelUsageSummary( + total_calls=self.model_call_counts[model], + total_input_tokens=self.model_input_tokens[model], + total_output_tokens=self.model_output_tokens[model], + ) + return UsageSummary(model_usage_summaries=model_summaries) + + def get_last_usage(self) -> ModelUsageSummary: + return ModelUsageSummary( + total_calls=1, + total_input_tokens=self.last_prompt_tokens, + total_output_tokens=self.last_completion_tokens, + ) diff --git a/rlm/core/types.py b/rlm/core/types.py index bd9e0d0a7..5f5501a12 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -12,6 +12,7 @@ "anthropic", "azure_openai", "gemini", + "claude_cli", ] EnvironmentType = Literal["local", "docker", "modal", "prime", "daytona", "e2b"] diff --git a/rlm/utils/token_utils.py b/rlm/utils/token_utils.py index 44772ce93..4963a5a81 100644 --- a/rlm/utils/token_utils.py +++ b/rlm/utils/token_utils.py @@ -64,6 +64,10 @@ "glm-4-9b": 1_000_000, "glm-4": 128_000, "glm": 128_000, + # Claude CLI aliases (used by claude_cli backend) + "sonnet": 200_000, + "opus": 200_000, + "haiku": 200_000, } diff --git a/tests/clients/test_claude_cli.py b/tests/clients/test_claude_cli.py new file mode 100644 index 000000000..f5f25df83 --- /dev/null +++ b/tests/clients/test_claude_cli.py @@ -0,0 +1,288 @@ +"""Tests for ClaudeCodeCLI client.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from rlm.core.types import ModelUsageSummary, UsageSummary + + +@pytest.fixture +def cli_client(): + """Create a ClaudeCodeCLI instance with mocked shutil.which.""" + with patch("rlm.clients.claude_cli.shutil.which", return_value="/usr/local/bin/claude"): + from rlm.clients.claude_cli import ClaudeCodeCLI + + return ClaudeCodeCLI(model_name="sonnet") + + +@pytest.fixture +def cli_client_custom(): + """Create a ClaudeCodeCLI with custom options.""" + with patch("rlm.clients.claude_cli.shutil.which", return_value="/usr/local/bin/claude"): + from rlm.clients.claude_cli import ClaudeCodeCLI + + return ClaudeCodeCLI( + model_name="opus", + max_budget_usd=5.0, + permission_mode="plan", + ) + + +class TestInit: + def test_defaults(self, cli_client): + assert cli_client.model_name == "sonnet" + assert cli_client.max_budget_usd is None + assert cli_client.permission_mode == "bypassPermissions" + + def test_custom_options(self, cli_client_custom): + assert cli_client_custom.model_name == "opus" + assert cli_client_custom.max_budget_usd == 5.0 + assert cli_client_custom.permission_mode == "plan" + + def test_missing_cli_raises(self): + with patch("rlm.clients.claude_cli.shutil.which", return_value=None): + from rlm.clients.claude_cli import ClaudeCodeCLI + + with pytest.raises(FileNotFoundError, match="claude CLI not found"): + ClaudeCodeCLI() + + def test_env_strips_claudecode(self, cli_client): + assert "CLAUDECODE" not in cli_client._env + + +class TestBuildCmd: + def test_basic_cmd(self, cli_client): + cmd = cli_client._build_cmd(system=None, model="sonnet") + assert cmd == [ + "claude", "--print", + "--model", "sonnet", + "--no-session-persistence", + "--permission-mode", "bypassPermissions", + ] + + def test_with_system_prompt(self, cli_client): + cmd = cli_client._build_cmd(system="You are helpful.", model="sonnet") + assert "--system-prompt" in cmd + idx = cmd.index("--system-prompt") + assert cmd[idx + 1] == "You are helpful." + + def test_with_budget(self, cli_client_custom): + cmd = cli_client_custom._build_cmd(system=None, model="opus") + assert "--max-budget-usd" in cmd + idx = cmd.index("--max-budget-usd") + assert cmd[idx + 1] == "5.0" + + def test_model_override(self, cli_client): + cmd = cli_client._build_cmd(system=None, model="haiku") + assert cmd[cmd.index("--model") + 1] == "haiku" + + +class TestPreparePrompt: + def test_string_prompt(self, cli_client): + text, system = cli_client._prepare_prompt("Hello") + assert text == "Hello" + assert system is None + + def test_message_list_no_system(self, cli_client): + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + text, system = cli_client._prepare_prompt(messages) + assert text == "Hi\n\nHello" + assert system is None + + def test_message_list_with_system(self, cli_client): + messages = [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Hi"}, + ] + text, system = cli_client._prepare_prompt(messages) + assert text == "Hi" + assert system == "Be concise." + + def test_invalid_prompt_type(self, cli_client): + with pytest.raises(ValueError, match="Invalid prompt type"): + cli_client._prepare_prompt(12345) + + +class TestEstimateTokens: + def test_basic_estimate(self): + from rlm.clients.claude_cli import ClaudeCodeCLI + + assert ClaudeCodeCLI._estimate_tokens("abcd") == 1 # 4 chars / 4 + assert ClaudeCodeCLI._estimate_tokens("abcde") == 2 # 5 chars -> ceil(5/4) + assert ClaudeCodeCLI._estimate_tokens("") == 0 + + def test_longer_text(self): + from rlm.clients.claude_cli import ClaudeCodeCLI + + text = "a" * 100 + assert ClaudeCodeCLI._estimate_tokens(text) == 25 + + +class TestCompletion: + def test_successful_completion(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "Hello, world!\n" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result) as mock_run: + result = cli_client.completion("Say hello") + + assert result == "Hello, world!" + mock_run.assert_called_once() + call_kwargs = mock_run.call_args[1] + assert call_kwargs["input"] == "Say hello" + assert call_kwargs["capture_output"] is True + assert call_kwargs["text"] is True + assert call_kwargs["timeout"] == cli_client.timeout + + def test_completion_with_model_override(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "response" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result) as mock_run: + cli_client.completion("test", model="opus") + + cmd = mock_run.call_args[0][0] + assert cmd[cmd.index("--model") + 1] == "opus" + + def test_completion_nonzero_exit(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 1 + mock_result.stderr = "Error: something went wrong" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result): + with pytest.raises(RuntimeError, match="claude CLI exited with code 1"): + cli_client.completion("test") + + def test_completion_tracks_usage(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "response text" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result): + cli_client.completion("input text") + + assert cli_client.model_call_counts["sonnet"] == 1 + assert cli_client.model_input_tokens["sonnet"] > 0 + assert cli_client.model_output_tokens["sonnet"] > 0 + + def test_completion_with_message_list(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "ok" + mock_result.stderr = "" + + messages = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Hi"}, + ] + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result) as mock_run: + cli_client.completion(messages) + + cmd = mock_run.call_args[0][0] + assert "--system-prompt" in cmd + assert mock_run.call_args[1]["input"] == "Hi" + + +class TestAcompletion: + def test_successful_acompletion(self, cli_client): + mock_proc = MagicMock() + mock_proc.returncode = 0 + mock_proc.communicate = AsyncMock(return_value=(b"Hello async!", b"")) + + with patch("rlm.clients.claude_cli.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=mock_proc): + result = asyncio.run(cli_client.acompletion("Say hello")) + + assert result == "Hello async!" + + def test_acompletion_nonzero_exit(self, cli_client): + mock_proc = MagicMock() + mock_proc.returncode = 1 + mock_proc.communicate = AsyncMock(return_value=(b"", b"async error")) + + with patch("rlm.clients.claude_cli.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=mock_proc): + with pytest.raises(RuntimeError, match="claude CLI exited with code 1"): + asyncio.run(cli_client.acompletion("test")) + + def test_acompletion_tracks_usage(self, cli_client): + mock_proc = MagicMock() + mock_proc.returncode = 0 + mock_proc.communicate = AsyncMock(return_value=(b"async response", b"")) + + with patch("rlm.clients.claude_cli.asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=mock_proc): + asyncio.run(cli_client.acompletion("async input")) + + assert cli_client.model_call_counts["sonnet"] == 1 + + +class TestUsageSummary: + def test_get_usage_summary_empty(self, cli_client): + summary = cli_client.get_usage_summary() + assert isinstance(summary, UsageSummary) + assert summary.model_usage_summaries == {} + + def test_get_usage_summary_after_calls(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "resp" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result): + cli_client.completion("a]") + cli_client.completion("b") + + summary = cli_client.get_usage_summary() + assert "sonnet" in summary.model_usage_summaries + assert summary.model_usage_summaries["sonnet"].total_calls == 2 + + def test_get_last_usage(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "response" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result): + cli_client.completion("input") + + last = cli_client.get_last_usage() + assert isinstance(last, ModelUsageSummary) + assert last.total_calls == 1 + assert last.total_input_tokens > 0 + assert last.total_output_tokens > 0 + + def test_multi_model_tracking(self, cli_client): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "ok" + mock_result.stderr = "" + + with patch("rlm.clients.claude_cli.subprocess.run", return_value=mock_result): + cli_client.completion("a", model="sonnet") + cli_client.completion("b", model="opus") + + summary = cli_client.get_usage_summary() + assert "sonnet" in summary.model_usage_summaries + assert "opus" in summary.model_usage_summaries + assert summary.model_usage_summaries["sonnet"].total_calls == 1 + assert summary.model_usage_summaries["opus"].total_calls == 1 + + +class TestGetClient: + def test_get_client_claude_cli(self): + with patch("rlm.clients.claude_cli.shutil.which", return_value="/usr/local/bin/claude"): + from rlm.clients import get_client + + client = get_client("claude_cli", {}) + from rlm.clients.claude_cli import ClaudeCodeCLI + + assert isinstance(client, ClaudeCodeCLI) From 3ce2dc81111db329a4750cf8bd97ace0c15b0a81 Mon Sep 17 00:00:00 2001 From: iam-dev Date: Thu, 26 Feb 2026 00:53:11 +0100 Subject: [PATCH 2/2] add anthropic prompt caching, cost tracking, and shared handler optimizations - Add prefix caching support with cache_control annotations for Anthropic API - Add per-model USD cost tracking using built-in pricing table - Track cache creation/read tokens in usage summaries - Add shared LM handler for child RLMs to reduce per-subcall overhead - Update local_repl and test_rlm_query for new features - Sync uv.lock version to 0.1.1 --- rlm/clients/anthropic.py | 114 +++++++++++- rlm/core/rlm.py | 126 ++++++++++--- rlm/environments/local_repl.py | 36 ++-- tests/test_anthropic_optimizations.py | 247 ++++++++++++++++++++++++++ tests/test_rlm_query.py | 94 ++++++++-- uv.lock | 2 +- 6 files changed, 557 insertions(+), 62 deletions(-) create mode 100644 tests/test_anthropic_optimizations.py diff --git a/rlm/clients/anthropic.py b/rlm/clients/anthropic.py index 5c747de4b..f4f6121bf 100644 --- a/rlm/clients/anthropic.py +++ b/rlm/clients/anthropic.py @@ -6,6 +6,33 @@ from rlm.clients.base_lm import BaseLM from rlm.core.types import ModelUsageSummary, UsageSummary +# Per-1M-token pricing (USD) as of 2025-06. Update as needed or use pricing_override. +_ANTHROPIC_PRICING: dict[str, tuple[float, float]] = { + # (input_per_1M, output_per_1M) + # Claude 4 family + "claude-opus-4-20250514": (15.0, 75.0), + "claude-opus-4-latest": (15.0, 75.0), + "claude-sonnet-4-20250514": (3.0, 15.0), + "claude-sonnet-4-latest": (3.0, 15.0), + # Claude 3.7 + "claude-3-7-sonnet-20250219": (3.0, 15.0), + "claude-3-7-sonnet-latest": (3.0, 15.0), + # Claude 3.5 + "claude-3-5-sonnet-20241022": (3.0, 15.0), + "claude-3-5-sonnet-20240620": (3.0, 15.0), + "claude-3-5-sonnet-latest": (3.0, 15.0), + "claude-3-5-haiku-20241022": (0.80, 4.0), + "claude-3-5-haiku-latest": (0.80, 4.0), + # Claude 3 + "claude-3-opus-20240229": (15.0, 75.0), + "claude-3-opus-latest": (15.0, 75.0), + "claude-3-sonnet-20240229": (3.0, 15.0), + "claude-3-haiku-20240307": (0.25, 1.25), +} + +_CACHE_CREATION_MULTIPLIER = 1.25 +_CACHE_READ_MULTIPLIER = 0.10 + class AnthropicClient(BaseLM): """ @@ -17,6 +44,8 @@ def __init__( api_key: str, model_name: str | None = None, max_tokens: int = 32768, + enable_prefix_cache: bool = False, + pricing_override: dict[str, tuple[float, float]] | None = None, **kwargs, ): super().__init__(model_name=model_name, **kwargs) @@ -24,12 +53,17 @@ def __init__( self.async_client = anthropic.AsyncAnthropic(api_key=api_key, timeout=self.timeout) self.model_name = model_name self.max_tokens = max_tokens + self.enable_prefix_cache = enable_prefix_cache + self._pricing = {**_ANTHROPIC_PRICING, **(pricing_override or {})} # Per-model usage tracking self.model_call_counts: dict[str, int] = defaultdict(int) self.model_input_tokens: dict[str, int] = defaultdict(int) self.model_output_tokens: dict[str, int] = defaultdict(int) self.model_total_tokens: dict[str, int] = defaultdict(int) + self.model_costs: dict[str, float] = defaultdict(float) + self.model_cache_creation_tokens: dict[str, int] = defaultdict(int) + self.model_cache_read_tokens: dict[str, int] = defaultdict(int) def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str: messages, system = self._prepare_messages(prompt) @@ -65,9 +99,13 @@ async def acompletion( def _prepare_messages( self, prompt: str | list[dict[str, Any]] - ) -> tuple[list[dict[str, Any]], str | None]: - """Prepare messages and extract system prompt for Anthropic API.""" - system = None + ) -> tuple[list[dict[str, Any]], str | list[dict[str, Any]] | None]: + """Prepare messages and extract system prompt for Anthropic API. + + When enable_prefix_cache is True, annotates the system prompt and the + last user message with cache_control for Anthropic prompt caching. + """ + system: str | list[dict[str, Any]] | None = None if isinstance(prompt, str): messages = [{"role": "user", "content": prompt}] @@ -82,25 +120,84 @@ def _prepare_messages( else: raise ValueError(f"Invalid prompt type: {type(prompt)}") + if self.enable_prefix_cache: + # Annotate system prompt with cache_control + if system is not None: + if isinstance(system, str): + system = [ + {"type": "text", "text": system, "cache_control": {"type": "ephemeral"}} + ] + elif isinstance(system, list) and system: + # System is already a list of content blocks; annotate the last block + last = system[-1] + if isinstance(last, dict): + system = system[:-1] + [{**last, "cache_control": {"type": "ephemeral"}}] + + # Annotate last user message content with cache_control + if messages: + for i in range(len(messages) - 1, -1, -1): + if messages[i].get("role") == "user": + content = messages[i].get("content", "") + if isinstance(content, str): + messages[i] = { + **messages[i], + "content": [ + { + "type": "text", + "text": content, + "cache_control": {"type": "ephemeral"}, + } + ], + } + break + return messages, system def _track_cost(self, response: anthropic.types.Message, model: str): self.model_call_counts[model] += 1 - self.model_input_tokens[model] += response.usage.input_tokens - self.model_output_tokens[model] += response.usage.output_tokens - self.model_total_tokens[model] += response.usage.input_tokens + response.usage.output_tokens + input_tokens = response.usage.input_tokens + output_tokens = response.usage.output_tokens + self.model_input_tokens[model] += input_tokens + self.model_output_tokens[model] += output_tokens + self.model_total_tokens[model] += input_tokens + output_tokens + + # Cache token tracking (Anthropic extended usage fields) + cache_creation = getattr(response.usage, "cache_creation_input_tokens", 0) or 0 + cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0 + self.model_cache_creation_tokens[model] += cache_creation + self.model_cache_read_tokens[model] += cache_read # Track last call for handler to read - self.last_prompt_tokens = response.usage.input_tokens - self.last_completion_tokens = response.usage.output_tokens + self.last_prompt_tokens = input_tokens + self.last_completion_tokens = output_tokens + + # Compute USD cost if pricing is known. + # Anthropic usage fields decompose input tokens into three disjoint groups: + # input_tokens: non-cached input tokens (charged at base rate) + # cache_creation_input_tokens: tokens written to cache (1.25x base rate) + # cache_read_input_tokens: tokens read from cache (0.10x base rate) + self.last_cost: float | None = None + pricing = self._pricing.get(model) + if pricing is not None: + input_price, output_price = pricing + cost = ( + input_tokens * input_price + + output_tokens * output_price + + cache_creation * input_price * _CACHE_CREATION_MULTIPLIER + + cache_read * input_price * _CACHE_READ_MULTIPLIER + ) / 1_000_000 + self.last_cost = cost + self.model_costs[model] += cost def get_usage_summary(self) -> UsageSummary: model_summaries = {} for model in self.model_call_counts: + cost = self.model_costs.get(model) model_summaries[model] = ModelUsageSummary( total_calls=self.model_call_counts[model], total_input_tokens=self.model_input_tokens[model], total_output_tokens=self.model_output_tokens[model], + total_cost=cost if cost else None, ) return UsageSummary(model_usage_summaries=model_summaries) @@ -109,4 +206,5 @@ def get_last_usage(self) -> ModelUsageSummary: total_calls=1, total_input_tokens=self.last_prompt_tokens, total_output_tokens=self.last_completion_tokens, + total_cost=getattr(self, "last_cost", None), ) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index ce247ecdc..b97bf3c55 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -1,3 +1,4 @@ +import threading import time from collections.abc import Callable from contextlib import contextmanager @@ -74,6 +75,9 @@ def __init__( on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None, on_iteration_start: Callable[[int, int], None] | None = None, on_iteration_complete: Callable[[int, int, float], None] | None = None, + # Internal: shared handler/client for child RLMs (reduces per-subcall overhead) + _shared_lm_handler_address: tuple[str, int] | None = None, + _shared_client: "BaseLM | None" = None, ): """ Args: @@ -152,11 +156,18 @@ def __init__( # Tracking (cumulative across all calls including children) self._cumulative_cost: float = 0.0 + self._cumulative_cost_lock = threading.Lock() self._consecutive_errors: int = 0 self._last_error: str | None = None self._best_partial_answer: str | None = None self._completion_start_time: float | None = None # Set when completion() starts + # Shared handler/client for child RLMs (Phase D optimization) + self._shared_lm_handler_address = _shared_lm_handler_address + self._shared_client = _shared_client + self._active_handler_address: tuple[str, int] | None = None + self._active_client: BaseLM | None = None + # Persistence support self.persistent = persistent self._persistent_env: SupportsPersistence | None = None @@ -192,24 +203,44 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): When persistent=True, the environment is reused across calls. When persistent=False (default), creates fresh environment each call. - """ - # Create client and wrap in handler - client: BaseLM = get_client(self.backend, self.backend_kwargs) - # Create other_backend_client if provided (for depth=1 routing) - other_backend_client: BaseLM | None = None - if self.other_backends and self.other_backend_kwargs: - other_backend_client = get_client(self.other_backends[0], self.other_backend_kwargs[0]) + When _shared_lm_handler_address is set (child RLM reusing parent's handler), + skips creating a new LMHandler/TCP server and creates a proxy instead. + """ + owns_handler = False + + if self._shared_lm_handler_address is not None and self._shared_client is not None: + # Reuse parent's handler — create a lightweight proxy + lm_handler = _SharedLMHandlerProxy(self._shared_client, self._shared_lm_handler_address) + # Active handler info for children to reuse + self._active_handler_address = self._shared_lm_handler_address + self._active_client = self._shared_client + else: + # Create client and wrap in handler + client: BaseLM = get_client(self.backend, self.backend_kwargs) - lm_handler = LMHandler(client, other_backend_client=other_backend_client) + # Create other_backend_client if provided (for depth=1 routing) + other_backend_client: BaseLM | None = None + if self.other_backends and self.other_backend_kwargs: + other_backend_client = get_client( + self.other_backends[0], self.other_backend_kwargs[0] + ) - # Register other clients to be available as sub-call options (by model name) - if self.other_backends and self.other_backend_kwargs: - for backend, kwargs in zip(self.other_backends, self.other_backend_kwargs, strict=True): - other_client: BaseLM = get_client(backend, kwargs) - lm_handler.register_client(other_client.model_name, other_client) + lm_handler = LMHandler(client, other_backend_client=other_backend_client) - lm_handler.start() + # Register other clients to be available as sub-call options (by model name) + if self.other_backends and self.other_backend_kwargs: + for backend, kwargs in zip( + self.other_backends, self.other_backend_kwargs, strict=True + ): + other_client: BaseLM = get_client(backend, kwargs) + lm_handler.register_client(other_client.model_name, other_client) + + lm_handler.start() + owns_handler = True + # Active handler info for children to reuse + self._active_handler_address = (lm_handler.host, lm_handler.port) + self._active_client = client # Environment: reuse if persistent, otherwise create fresh if self.persistent and self._persistent_env is not None: @@ -221,11 +252,12 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): f"implement required methods (update_handler_address, add_context, get_context_count). " f"This should have been caught at initialization." ) - environment.update_handler_address((lm_handler.host, lm_handler.port)) + handler_addr = self._active_handler_address + environment.update_handler_address(handler_addr) environment.add_context(prompt) else: env_kwargs = self.environment_kwargs.copy() - env_kwargs["lm_handler_address"] = (lm_handler.host, lm_handler.port) + env_kwargs["lm_handler_address"] = self._active_handler_address env_kwargs["context_payload"] = prompt env_kwargs["depth"] = self.depth + 1 # Environment depth is RLM depth + 1 # For local environment with max_depth > 1, pass subcall callback for recursive RLM calls @@ -246,7 +278,10 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): try: yield lm_handler, environment finally: - lm_handler.stop() + self._active_handler_address = None + self._active_client = None + if owns_handler: + lm_handler.stop() if not self.persistent and hasattr(environment, "cleanup"): environment.cleanup() @@ -491,7 +526,8 @@ def _check_iteration_limits( if self.max_budget is not None: current_usage = lm_handler.get_usage_summary() current_cost = current_usage.total_cost or 0.0 - self._cumulative_cost = current_cost + with self._cumulative_cost_lock: + self._cumulative_cost = current_cost if self._cumulative_cost > self.max_budget: self.verbose.print_budget_exceeded(self._cumulative_cost, self.max_budget) raise BudgetExceededError( @@ -703,7 +739,8 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: # Calculate remaining budget for child (if budget tracking enabled) remaining_budget = None if self.max_budget is not None: - remaining_budget = self.max_budget - self._cumulative_cost + with self._cumulative_cost_lock: + remaining_budget = self.max_budget - self._cumulative_cost if remaining_budget <= 0: return RLMChatCompletion( root_model=resolved_model, @@ -743,7 +780,7 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: subcall_start = time.perf_counter() error_msg: str | None = None - # Spawn a child RLM with its own LocalREPL + # Spawn a child RLM, sharing parent's handler when possible child = RLM( backend=self.backend, backend_kwargs=child_backend_kwargs, @@ -768,16 +805,21 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: # Propagate callbacks to children for nested tracking on_subcall_start=self.on_subcall_start, on_subcall_complete=self.on_subcall_complete, + # Share parent's handler/client to avoid per-child TCP server overhead + _shared_lm_handler_address=self._active_handler_address, + _shared_client=self._active_client, ) try: result = child.completion(prompt, root_prompt=None) # Track child's cost in parent's cumulative cost if result.usage_summary and result.usage_summary.total_cost: - self._cumulative_cost += result.usage_summary.total_cost + with self._cumulative_cost_lock: + self._cumulative_cost += result.usage_summary.total_cost return result except BudgetExceededError as e: # Propagate child's spending to parent - self._cumulative_cost += e.spent + with self._cumulative_cost_lock: + self._cumulative_cost += e.spent error_msg = f"Budget exceeded - {e}" return RLMChatCompletion( root_model=resolved_model, @@ -849,3 +891,43 @@ def __enter__(self) -> "RLM": def __exit__(self, exc_type, exc_val, exc_tb) -> bool: self.close() return False + + +class _SharedLMHandlerProxy: + """Lightweight proxy that reuses a parent's BaseLM client without spawning a TCP server. + + Exposes the same interface as LMHandler for completion() and get_usage_summary(), + but start()/stop() are no-ops since the parent owns the actual handler lifecycle. + """ + + def __init__(self, client: BaseLM, address: tuple[str, int]): + self._client = client + self._host, self._port = address + + @property + def host(self) -> str: + return self._host + + @property + def port(self) -> int: + return self._port + + @property + def address(self) -> tuple[str, int]: + return (self._host, self._port) + + def start(self) -> tuple[str, int]: + return self.address + + def stop(self) -> None: + pass # Parent owns the handler + + def completion(self, prompt: str, model: str | None = None) -> str: + # BaseLM.completion() does not accept model, but concrete implementations + # (AnthropicClient, OpenAIClient) do. Forward when possible. + if model is not None: + return self._client.completion(prompt, model=model) + return self._client.completion(prompt) + + def get_usage_summary(self) -> UsageSummary: + return self._client.get_usage_summary() diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index afdde3b3d..687528f6c 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -9,6 +9,7 @@ import time import uuid from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed from contextlib import contextmanager from typing import Any @@ -316,9 +317,10 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str: return self._llm_query(prompt, model) def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]: - """Spawn recursive RLM sub-calls for multiple prompts. + """Spawn recursive RLM sub-calls for multiple prompts in parallel. - Each prompt gets its own child RLM for deeper thinking. + Each prompt gets its own child RLM for deeper thinking. When a subcall_fn + is configured, all prompts are dispatched concurrently via ThreadPoolExecutor. Falls back to llm_query_batched if no recursive capability is configured. Args: @@ -329,14 +331,23 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li List of responses in the same order as input prompts. """ if self.subcall_fn is not None: - results = [] - for prompt in prompts: - try: - completion = self.subcall_fn(prompt, model) - self._pending_llm_calls.append(completion) - results.append(completion.response) - except Exception as e: - results.append(f"Error: RLM query failed - {e}") + if not prompts: + return [] + results: list[str | None] = [None] * len(prompts) + with ThreadPoolExecutor(max_workers=min(len(prompts), 16)) as executor: + future_to_idx = { + executor.submit(self.subcall_fn, prompt, model): i + for i, prompt in enumerate(prompts) + } + for future in as_completed(future_to_idx): + idx = future_to_idx[future] + try: + completion = future.result() + # list.append is GIL-protected in CPython — safe for concurrent use + self._pending_llm_calls.append(completion) + results[idx] = completion.response + except Exception as e: + results[idx] = f"Error: RLM query failed - {e}" return results # Fall back to plain batched LM call if no recursive capability @@ -483,8 +494,9 @@ def execute_code(self, code: str) -> REPLResult: """Execute code in the persistent namespace and return result.""" start_time = time.perf_counter() - # Clear pending LLM calls from previous execution - self._pending_llm_calls = [] + # Clear pending LLM calls from previous execution (use .clear() to keep + # the same list reference, preventing stale references from other threads) + self._pending_llm_calls.clear() with self._capture_output() as (stdout_buf, stderr_buf), self._temp_cwd(): try: diff --git a/tests/test_anthropic_optimizations.py b/tests/test_anthropic_optimizations.py new file mode 100644 index 000000000..c97240b03 --- /dev/null +++ b/tests/test_anthropic_optimizations.py @@ -0,0 +1,247 @@ +"""Tests for Anthropic client cost estimation and prefix caching.""" + +from unittest.mock import MagicMock, patch + +from rlm.clients.anthropic import ( + _ANTHROPIC_PRICING, + _CACHE_CREATION_MULTIPLIER, + _CACHE_READ_MULTIPLIER, + AnthropicClient, +) + + +def _mock_usage(input_tokens=100, output_tokens=50, cache_creation=0, cache_read=0): + """Create a mock Anthropic usage object.""" + usage = MagicMock() + usage.input_tokens = input_tokens + usage.output_tokens = output_tokens + usage.cache_creation_input_tokens = cache_creation + usage.cache_read_input_tokens = cache_read + return usage + + +def _mock_response(input_tokens=100, output_tokens=50, cache_creation=0, cache_read=0): + """Create a mock Anthropic Message response.""" + response = MagicMock() + response.usage = _mock_usage(input_tokens, output_tokens, cache_creation, cache_read) + response.content = [MagicMock(text="test response")] + return response + + +class TestAnthropicPricingTable: + """Tests for the pricing table.""" + + def test_known_models_present(self): + assert "claude-sonnet-4-20250514" in _ANTHROPIC_PRICING + assert "claude-opus-4-20250514" in _ANTHROPIC_PRICING + assert "claude-3-5-sonnet-latest" in _ANTHROPIC_PRICING + assert "claude-3-haiku-20240307" in _ANTHROPIC_PRICING + + def test_pricing_format(self): + for model, (input_price, output_price) in _ANTHROPIC_PRICING.items(): + assert isinstance(input_price, (int, float)), f"Bad input price for {model}" + assert isinstance(output_price, (int, float)), f"Bad output price for {model}" + assert input_price > 0, f"Non-positive input price for {model}" + assert output_price > 0, f"Non-positive output price for {model}" + + +class TestAnthropicCostEstimation: + """Tests for cost tracking in _track_cost.""" + + @patch("rlm.clients.anthropic.anthropic") + def test_known_model_cost_positive(self, mock_anthropic): + """Cost should be > 0 for a known model.""" + client = AnthropicClient(api_key="test-key", model_name="claude-sonnet-4-20250514") + response = _mock_response(input_tokens=1000, output_tokens=500) + + client._track_cost(response, "claude-sonnet-4-20250514") + + assert client.last_cost is not None + assert client.last_cost > 0 + # Verify: (1000 * 3.0 + 500 * 15.0) / 1_000_000 = 0.0105 + expected = (1000 * 3.0 + 500 * 15.0) / 1_000_000 + assert abs(client.last_cost - expected) < 1e-10 + + @patch("rlm.clients.anthropic.anthropic") + def test_unknown_model_cost_none(self, mock_anthropic): + """Cost should be None for an unknown model.""" + client = AnthropicClient(api_key="test-key", model_name="unknown-model-xyz") + response = _mock_response(input_tokens=100, output_tokens=50) + + client._track_cost(response, "unknown-model-xyz") + + assert client.last_cost is None + + @patch("rlm.clients.anthropic.anthropic") + def test_cache_token_cost(self, mock_anthropic): + """Cache creation and read tokens should affect cost.""" + client = AnthropicClient(api_key="test-key", model_name="claude-sonnet-4-20250514") + response = _mock_response( + input_tokens=1000, + output_tokens=500, + cache_creation=200, + cache_read=300, + ) + + client._track_cost(response, "claude-sonnet-4-20250514") + + assert client.last_cost is not None + input_price, output_price = _ANTHROPIC_PRICING["claude-sonnet-4-20250514"] + expected = ( + 1000 * input_price + + 500 * output_price + + 200 * input_price * _CACHE_CREATION_MULTIPLIER + + 300 * input_price * _CACHE_READ_MULTIPLIER + ) / 1_000_000 + assert abs(client.last_cost - expected) < 1e-10 + + @patch("rlm.clients.anthropic.anthropic") + def test_cache_tokens_tracked(self, mock_anthropic): + """Cache tokens should be accumulated in model_cache_* dicts.""" + client = AnthropicClient(api_key="test-key", model_name="claude-sonnet-4-20250514") + model = "claude-sonnet-4-20250514" + + client._track_cost(_mock_response(cache_creation=100, cache_read=200), model) + client._track_cost(_mock_response(cache_creation=50, cache_read=80), model) + + assert client.model_cache_creation_tokens[model] == 150 + assert client.model_cache_read_tokens[model] == 280 + + @patch("rlm.clients.anthropic.anthropic") + def test_usage_summary_includes_cost(self, mock_anthropic): + """get_usage_summary should include total_cost for known models.""" + client = AnthropicClient(api_key="test-key", model_name="claude-sonnet-4-20250514") + client._track_cost( + _mock_response(input_tokens=1000, output_tokens=500), + "claude-sonnet-4-20250514", + ) + + summary = client.get_usage_summary() + model_summary = summary.model_usage_summaries["claude-sonnet-4-20250514"] + assert model_summary.total_cost is not None + assert model_summary.total_cost > 0 + + @patch("rlm.clients.anthropic.anthropic") + def test_last_usage_includes_cost(self, mock_anthropic): + """get_last_usage should include total_cost.""" + client = AnthropicClient(api_key="test-key", model_name="claude-sonnet-4-20250514") + client._track_cost( + _mock_response(input_tokens=1000, output_tokens=500), + "claude-sonnet-4-20250514", + ) + + last = client.get_last_usage() + assert last.total_cost is not None + assert last.total_cost > 0 + + @patch("rlm.clients.anthropic.anthropic") + def test_pricing_override(self, mock_anthropic): + """pricing_override should allow custom per-model pricing.""" + custom_pricing = {"my-custom-model": (1.0, 2.0)} + client = AnthropicClient( + api_key="test-key", + model_name="my-custom-model", + pricing_override=custom_pricing, + ) + client._track_cost( + _mock_response(input_tokens=1000, output_tokens=500), + "my-custom-model", + ) + expected = (1000 * 1.0 + 500 * 2.0) / 1_000_000 + assert abs(client.last_cost - expected) < 1e-10 + + +class TestAnthropicPrefixCaching: + """Tests for enable_prefix_cache message structure.""" + + @patch("rlm.clients.anthropic.anthropic") + def test_prefix_cache_disabled_plain_strings(self, mock_anthropic): + """With prefix caching disabled, system and content should be plain strings.""" + client = AnthropicClient( + api_key="test-key", + model_name="test-model", + enable_prefix_cache=False, + ) + prompt = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + messages, system = client._prepare_messages(prompt) + + assert system == "You are helpful." + assert messages == [{"role": "user", "content": "Hello"}] + + @patch("rlm.clients.anthropic.anthropic") + def test_prefix_cache_system_as_list_of_blocks(self, mock_anthropic): + """With prefix caching enabled, system should be a list with cache_control.""" + client = AnthropicClient( + api_key="test-key", + model_name="test-model", + enable_prefix_cache=True, + ) + prompt = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + messages, system = client._prepare_messages(prompt) + + assert isinstance(system, list) + assert len(system) == 1 + assert system[0]["type"] == "text" + assert system[0]["text"] == "You are helpful." + assert system[0]["cache_control"] == {"type": "ephemeral"} + + @patch("rlm.clients.anthropic.anthropic") + def test_prefix_cache_last_user_message(self, mock_anthropic): + """With prefix caching, last user message content should have cache_control.""" + client = AnthropicClient( + api_key="test-key", + model_name="test-model", + enable_prefix_cache=True, + ) + prompt = [ + {"role": "user", "content": "First message"}, + {"role": "assistant", "content": "Response"}, + {"role": "user", "content": "Second message"}, + ] + messages, system = client._prepare_messages(prompt) + + # First user message should be unchanged + assert messages[0]["content"] == "First message" + # Last user message should be annotated + last_user = messages[2] + assert isinstance(last_user["content"], list) + assert last_user["content"][0]["type"] == "text" + assert last_user["content"][0]["text"] == "Second message" + assert last_user["content"][0]["cache_control"] == {"type": "ephemeral"} + + @patch("rlm.clients.anthropic.anthropic") + def test_prefix_cache_no_system_prompt(self, mock_anthropic): + """Prefix caching should work even without a system prompt.""" + client = AnthropicClient( + api_key="test-key", + model_name="test-model", + enable_prefix_cache=True, + ) + prompt = [{"role": "user", "content": "Hello"}] + messages, system = client._prepare_messages(prompt) + + assert system is None + assert isinstance(messages[0]["content"], list) + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral"} + + @patch("rlm.clients.anthropic.anthropic") + def test_prefix_cache_string_prompt(self, mock_anthropic): + """Prefix caching with a simple string prompt.""" + client = AnthropicClient( + api_key="test-key", + model_name="test-model", + enable_prefix_cache=True, + ) + messages, system = client._prepare_messages("Hello world") + + assert system is None + assert len(messages) == 1 + assert isinstance(messages[0]["content"], list) + assert messages[0]["content"][0]["text"] == "Hello world" + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral"} diff --git a/tests/test_rlm_query.py b/tests/test_rlm_query.py index 850a87fff..b72ad7a78 100644 --- a/tests/test_rlm_query.py +++ b/tests/test_rlm_query.py @@ -1,5 +1,6 @@ """Tests for rlm_query and rlm_query_batched in LocalREPL.""" +import time from unittest.mock import MagicMock from rlm.core.types import RLMChatCompletion, UsageSummary @@ -75,13 +76,14 @@ class TestRlmQueryBatchedWithSubcallFn: """Tests for rlm_query_batched when subcall_fn is provided.""" def test_batched_calls_subcall_fn_per_prompt(self): - """rlm_query_batched should call subcall_fn once per prompt.""" - completions = [ - _make_completion("answer 1"), - _make_completion("answer 2"), - _make_completion("answer 3"), - ] - subcall_fn = MagicMock(side_effect=completions) + """rlm_query_batched should call subcall_fn once per prompt (thread-safe).""" + # Use prompt-keyed dict instead of ordered side_effect for thread safety + prompt_responses = { + "q1": _make_completion("answer 1"), + "q2": _make_completion("answer 2"), + "q3": _make_completion("answer 3"), + } + subcall_fn = MagicMock(side_effect=lambda prompt, model: prompt_responses[prompt]) repl = LocalREPL(subcall_fn=subcall_fn) result = repl.execute_code( "answers = rlm_query_batched(['q1', 'q2', 'q3'])\nprint(len(answers))" @@ -94,12 +96,18 @@ def test_batched_calls_subcall_fn_per_prompt(self): def test_batched_tracks_all_pending_calls(self): """rlm_query_batched should track all completions in rlm_calls.""" - completions = [_make_completion(f"resp {i}") for i in range(3)] - subcall_fn = MagicMock(side_effect=completions) + prompt_responses = { + "a": _make_completion("resp 0"), + "b": _make_completion("resp 1"), + "c": _make_completion("resp 2"), + } + subcall_fn = MagicMock(side_effect=lambda prompt, model: prompt_responses[prompt]) repl = LocalREPL(subcall_fn=subcall_fn) result = repl.execute_code("rlm_query_batched(['a', 'b', 'c'])") assert len(result.rlm_calls) == 3 - assert [c.response for c in result.rlm_calls] == ["resp 0", "resp 1", "resp 2"] + # Order may vary due to parallel execution; check all responses present + responses = sorted(c.response for c in result.rlm_calls) + assert responses == ["resp 0", "resp 1", "resp 2"] repl.cleanup() def test_batched_with_model_override(self): @@ -114,21 +122,21 @@ def test_batched_with_model_override(self): def test_batched_partial_failure(self): """If one subcall_fn call fails, others should still succeed.""" - subcall_fn = MagicMock( - side_effect=[ - _make_completion("ok 1"), - RuntimeError("boom"), - _make_completion("ok 3"), - ] - ) + + def _side_effect(prompt, model): + if prompt == "b": + raise RuntimeError("boom") + return _make_completion(f"ok {prompt}") + + subcall_fn = MagicMock(side_effect=_side_effect) repl = LocalREPL(subcall_fn=subcall_fn) result = repl.execute_code("answers = rlm_query_batched(['a', 'b', 'c'])") assert result.stderr == "" answers = repl.locals["answers"] - assert answers[0] == "ok 1" + assert answers[0] == "ok a" assert "Error" in answers[1] assert "boom" in answers[1] - assert answers[2] == "ok 3" + assert answers[2] == "ok c" repl.cleanup() def test_batched_empty_prompts(self): @@ -208,3 +216,51 @@ def test_rlm_query_batched_restored_after_overwrite(self): repl.execute_code("answers = rlm_query_batched(['q1'])") assert repl.locals["answers"] == ["real"] repl.cleanup() + + +class TestRlmQueryBatchedParallel: + """Tests that rlm_query_batched runs concurrently, not sequentially.""" + + def test_batched_runs_concurrently(self): + """N subcalls with 0.1s delay each should complete in ~0.1s, not N*0.1s.""" + n = 5 + delay = 0.1 + + def _slow_subcall(prompt, model): + time.sleep(delay) + return _make_completion(f"done-{prompt}") + + subcall_fn = MagicMock(side_effect=_slow_subcall) + repl = LocalREPL(subcall_fn=subcall_fn) + prompts = [f"p{i}" for i in range(n)] + code = f"answers = rlm_query_batched({prompts!r})" + + start = time.perf_counter() + result = repl.execute_code(code) + elapsed = time.perf_counter() - start + + assert result.stderr == "" + assert subcall_fn.call_count == n + answers = repl.locals["answers"] + assert len(answers) == n + assert all(a.startswith("done-") for a in answers) + # Parallel: should take ~delay, not n*delay. Use 2*delay as generous bound. + assert elapsed < 2 * delay, f"Took {elapsed:.2f}s, expected < {2 * delay:.2f}s (parallel)" + repl.cleanup() + + def test_batched_preserves_order(self): + """Results should be in the same order as input prompts despite parallel execution.""" + + def _varied_delay(prompt, model): + # Vary delays to encourage out-of-order completion + idx = int(prompt.replace("p", "")) + time.sleep(0.05 * (5 - idx)) # p0 slowest, p4 fastest + return _make_completion(f"result-{prompt}") + + subcall_fn = MagicMock(side_effect=_varied_delay) + repl = LocalREPL(subcall_fn=subcall_fn) + result = repl.execute_code("answers = rlm_query_batched(['p0','p1','p2','p3','p4'])") + assert result.stderr == "" + answers = repl.locals["answers"] + assert answers == [f"result-p{i}" for i in range(5)] + repl.cleanup() diff --git a/uv.lock b/uv.lock index 6abd703d0..5e1be9de9 100644 --- a/uv.lock +++ b/uv.lock @@ -1760,7 +1760,7 @@ wheels = [ [[package]] name = "rlms" -version = "0.1.0" +version = "0.1.1" source = { editable = "." } dependencies = [ { name = "anthropic" },