Add Anthropic prompt caching, cost tracking, and shared handler optimizations

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']"
         )

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/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,
+        )