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

Author: iam-dev

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,19 +44,26 @@ 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)
         self.client = anthropic.Anthropic(api_key=api_key, timeout=self.timeout)
         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),
         )

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()