Implement token limit feature with cost protection and logging

Author: pairing4good

README.md

@@ -41,6 +41,19 @@ You can try out RLMs quickly by installing from PyPi:
 pip install rlms
 ```
 
+> [!IMPORTANT]
+> **Token Limit & Cost Protection**: RLM has a **default limit of 1,000,000 tokens per session** to prevent unexpectedly large bills. This costs approximately:
+> - GPT-3.5 Turbo: ~$1 per session
+> - GPT-4o: ~$6 per session  
+> - GPT-4 Turbo: ~$20 per session
+> - Claude Opus: ~$45 per session
+>
+> To change this limit:
+> ```python
+> rlm = RLM(max_tokens=500_000)  # Reduce to 500k tokens
+> rlm = RLM(max_tokens=None)      # Remove limit (use with caution!)
+> ```
+
 The default RLM client uses a REPL environment that runs on the host process through Python `exec` calls. It uses the same virtual environment as the host process (i.e. it will have access to the same dependencies), but with some limitations in its available global modules. As an example, we can call RLM completions using GPT-5-nano:
 ```python
 from rlm import RLM
@@ -54,6 +67,11 @@ rlm = RLM(
 print(rlm.completion("Print me the first 100 powers of two, each on a newline.").response)
 ```
 
+**Understanding Your Token Budget:**
+- Each session (one `completion()` call) is limited to 1M tokens by default
+- You can check token usage: `result.usage_summary.to_dict()`
+- Adjust the limit with `max_tokens` parameter (see [Token Limits](#important-token-limit--cost-protection) above)
+
 <details>
 <summary><b>Manual Setup</b></summary>
 

docs/api/rlm.md

@@ -45,6 +45,7 @@ RLM(
     depth: int = 0,
     max_depth: int = 1,
     max_iterations: int = 30,
+    max_tokens: int | None = None,
     custom_system_prompt: str | None = None,
     other_backends: list[str] | None = None,
     other_backend_kwargs: list[dict] | None = None,
@@ -187,6 +188,83 @@ rlm = RLM(
 
 ---
 
+#### `max_tokens`
+{: .no_toc }
+
+**Type:** `int | None`  
+**Default:** `1_000_000` (1 million tokens)
+
+Maximum total tokens (input + output combined) allowed per `completion()` session. When this limit is reached or exceeded, the session ends early to stay within the token budget.
+
+{: .warning }
+> **This is a cost protection feature!** The default limit prevents unexpectedly large API bills.
+
+**Cost implications of 1M tokens:**
+
+| Model | Cost for 1M Tokens |
+|:------|:-------------------|
+| GPT-3.5 Turbo | ~$1 |
+| GPT-4o | ~$6 |
+| GPT-4o-mini | ~$0.30 |
+| GPT-4 Turbo | ~$20 |
+| Claude Sonnet 4 | ~$9 |
+| Claude Opus | ~$45 |
+
+**Token counting:**
+- Includes all LM calls (main model + sub-LM calls via `llm_query()`)
+- Tracks cumulative usage across all iterations in a single session
+- Resets for each new `completion()` call (even in persistent mode)
+
+**Examples:**
+
+```python
+# Use default 1M token limit (recommended for most use cases)
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    # max_tokens=1_000_000 is implicit
+)
+
+# Reduce limit for stricter budget control
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=100_000,  # Limit to ~$0.60 per session
+)
+
+# Increase limit for complex, expensive tasks
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=5_000_000,  # Allow up to ~$30 per session
+)
+
+# Remove limit entirely (not recommended - use with caution!)
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=None,  # Unlimited - can result in large bills!
+)
+
+result = rlm.completion("Analyze this dataset...")
+# If limit is hit:
+# result.response == "Session ended: Token limit exceeded. Used 1,000,234 tokens..."
+```
+
+**When the limit is exceeded:**
+- The session stops before the next iteration
+- A clear message is returned in `result.response`
+- The event is logged (if a logger is configured)
+- Usage summary and execution time are still populated
+
+**Use cases:**
+- **Budget control**: Prevent runaway costs on complex tasks (default behavior)
+- **Cost management**: Set organization-wide limits in production
+- **Development**: Use lower limits during testing/debugging
+- **Production**: Increase limits for known expensive workflows
+
+---
+
 #### `custom_system_prompt`
 {: .no_toc }
 
@@ -446,6 +524,7 @@ rlm = RLM(
 
     # Behavior
     max_iterations=40,
+    max_tokens=100_000,  # Budget control
     max_depth=1,
 
     # Debugging

docs/getting-started.md

@@ -62,6 +62,60 @@ docker --version
 
 ---
 
+## Understanding Token Limits & Costs
+
+{: .warning }
+> **IMPORTANT: Default Token Limit**  
+> RLM has a built-in safety limit of **1,000,000 tokens per session** to prevent unexpectedly large API bills.
+
+### What does 1M tokens cost?
+
+The cost varies by model:
+
+| Model | Approximate Cost for 1M Tokens |
+|:------|:-------------------------------|
+| GPT-3.5 Turbo | ~$1 |
+| GPT-4o | ~$6 |
+| GPT-4o-mini | ~$0.30 |
+| GPT-4 Turbo | ~$20 |
+| Claude Sonnet 4 | ~$9 |
+| Claude Opus | ~$45 |
+
+### How to adjust the limit
+
+```python
+# Reduce limit for tighter budget control
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=100_000,  # Limit to ~$0.60 per session
+)
+
+# Increase limit for complex tasks
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=5_000_000,  # Allow up to ~$30 per session
+)
+
+# Remove limit entirely (use with caution!)
+rlm = RLM(
+    backend="openai",
+    backend_kwargs={"model_name": "gpt-4o"},
+    max_tokens=None,  # Unlimited - monitor your usage!
+)
+```
+
+### What happens when the limit is reached?
+
+When your session hits the token limit:
+1. RLM stops before the next iteration
+2. Returns a clear message: `"Session ended: Token limit exceeded. Used X tokens (limit: Y)..."`
+3. Logs the event (if logging is enabled)
+4. You still receive usage statistics and timing information
+
+---
+
 ## Your First RLM Call
 
 ### Step 1: Set Up API Keys
@@ -133,6 +187,7 @@ This will display:
 | `environment_kwargs` | `dict` | `None` | Environment configuration |
 | `max_depth` | `int` | `1` | Maximum recursion depth |
 | `max_iterations` | `int` | `30` | Max REPL iterations per call |
+| `max_tokens` | `int` | `1_000_000` | **Max total tokens per session** (see [Token Limits](#understanding-token-limits--costs)) |
 | `custom_system_prompt` | `str` | `None` | Override default system prompt |
 | `other_backends` | `list` | `None` | Additional backends for sub-calls |
 | `other_backend_kwargs` | `list` | `None` | Configs for additional backends |

rlm/core/rlm.py

@@ -12,6 +12,7 @@
     RLMChatCompletion,
     RLMIteration,
     RLMMetadata,
+    UsageSummary,
 )
 from rlm.environments import BaseEnv, SupportsPersistence, get_environment
 from rlm.logger import RLMLogger, VerbosePrinter
@@ -46,6 +47,7 @@ def __init__(
         depth: int = 0,
         max_depth: int = 1,
         max_iterations: int = 30,
+        max_tokens: int | None = 1_000_000,
         custom_system_prompt: str | None = None,
         other_backends: list[ClientBackend] | None = None,
         other_backend_kwargs: list[dict[str, Any]] | None = None,
@@ -62,6 +64,12 @@ def __init__(
             depth: The current depth of the RLM (0-indexed).
             max_depth: The maximum depth of the RLM. Currently, only depth 1 is supported.
             max_iterations: The maximum number of iterations of the RLM.
+            max_tokens: The maximum total tokens (input + output) allowed per completion() session.
+                DEFAULT: 1,000,000 tokens to prevent unexpectedly large API bills.
+                COST ESTIMATE: 1M tokens costs ~$1-45 depending on your model:
+                  - GPT-3.5 Turbo: ~$1  |  GPT-4o: ~$6  |  GPT-4 Turbo: ~$20  |  Claude Opus: ~$45
+                Set to None for unlimited tokens (use with caution - can result in expensive bills).
+                Set to a smaller value (e.g., 100_000) for tighter budget control.
             custom_system_prompt: The custom system prompt to use for the RLM.
             other_backends: A list of other client backends that the environments can use to make sub-calls.
             other_backend_kwargs: The kwargs to pass to the other client backends (ordered to match other_backends).
@@ -90,6 +98,7 @@ def __init__(
         self.depth = depth
         self.max_depth = max_depth
         self.max_iterations = max_iterations
+        self.max_tokens = max_tokens
         self.system_prompt = custom_system_prompt if custom_system_prompt else RLM_SYSTEM_PROMPT
         self.logger = logger
         self.verbose = VerbosePrinter(enabled=verbose)
@@ -110,6 +119,7 @@ def __init__(
                 else "unknown",
                 max_depth=max_depth,
                 max_iterations=max_iterations,
+                max_tokens=max_tokens,
                 backend=backend,
                 backend_kwargs=filter_sensitive_keys(backend_kwargs) if backend_kwargs else {},
                 environment_type=environment,
@@ -215,6 +225,16 @@ def completion(
             message_history = self._setup_prompt(prompt)
 
             for i in range(self.max_iterations):
+                # Check token limit before starting next iteration
+                if self.max_tokens is not None:
+                    current_usage = lm_handler.get_usage_summary()
+                    tokens_used = self._calculate_total_tokens(current_usage)
+                    if tokens_used >= self.max_tokens:
+                        # Store message history in persistent environment
+                        if self.persistent and isinstance(environment, SupportsPersistence):
+                            environment.add_history(message_history)
+                        return self._token_limit_exceeded_answer(i, tokens_used, time_start, lm_handler)
+
                 # Current prompt = message history + additional prompt suffix
                 context_count = (
                     environment.get_context_count()
@@ -354,6 +374,62 @@ def _fallback_answer(self, message: str | dict[str, Any]) -> str:
         response = client.completion(message)
         return response
 
+    def _calculate_total_tokens(self, usage: UsageSummary) -> int:
+        """Calculate total tokens (input + output) from usage summary."""
+        total = 0
+        for model_usage in usage.model_usage_summaries.values():
+            total += model_usage.total_input_tokens + model_usage.total_output_tokens
+        return total
+
+    def _token_limit_exceeded_answer(
+        self,
+        iteration: int,
+        tokens_used: int,
+        time_start: float,
+        lm_handler: LMHandler,
+    ) -> RLMChatCompletion:
+        """
+        Handle early termination when token limit is exceeded.
+        Returns a completion indicating the session ended due to token limit.
+        """
+        time_end = time.perf_counter()
+        usage = lm_handler.get_usage_summary()
+
+        # Log the limit hit event
+        if self.logger:
+            self.logger.log_limit_hit(
+                limit_type="token_limit",
+                max_value=self.max_tokens,
+                current_value=tokens_used,
+                iteration=iteration,
+            )
+
+        # Print verbose output
+        self.verbose.print_token_limit_hit(
+            max_tokens=self.max_tokens,
+            tokens_used=tokens_used,
+            iteration=iteration,
+        )
+        self.verbose.print_summary(
+            iteration, time_end - time_start, usage.to_dict(), stopped_reason="token_limit"
+        )
+
+        response_message = (
+            f"Session ended: Token limit exceeded. "
+            f"Used {tokens_used:,} tokens (limit: {self.max_tokens:,}). "
+            f"Completed {iteration} iteration(s) before reaching limit."
+        )
+
+        return RLMChatCompletion(
+            root_model=self.backend_kwargs.get("model_name", "unknown")
+            if self.backend_kwargs
+            else "unknown",
+            prompt="",
+            response=response_message,
+            usage_summary=usage,
+            execution_time=time_end - time_start,
+        )
+
     def _validate_persistent_environment_support(self) -> None:
         """
         Validate that the configured environment type supports persistent mode.

rlm/core/types.py

@@ -192,6 +192,7 @@ class RLMMetadata:
     root_model: str
     max_depth: int
     max_iterations: int
+    max_tokens: int | None
     backend: str
     backend_kwargs: dict[str, Any]
     environment_type: str
@@ -203,6 +204,7 @@ def to_dict(self):
             "root_model": self.root_model,
             "max_depth": self.max_depth,
             "max_iterations": self.max_iterations,
+            "max_tokens": self.max_tokens,
             "backend": self.backend,
             "backend_kwargs": {k: _serialize_value(v) for k, v in self.backend_kwargs.items()},
             "environment_type": self.environment_type,

rlm/logger/rlm_logger.py

@@ -58,6 +58,23 @@ def log(self, iteration: RLMIteration):
             json.dump(entry, f)
             f.write("\n")
 
+    def log_limit_hit(
+        self, limit_type: str, max_value: int, current_value: int, iteration: int
+    ):
+        """Log a limit hit event (e.g., token limit exceeded)."""
+        entry = {
+            "type": "limit_hit",
+            "limit_type": limit_type,
+            "max_value": max_value,
+            "current_value": current_value,
+            "iteration": iteration,
+            "timestamp": datetime.now().isoformat(),
+        }
+
+        with open(self.log_file_path, "a") as f:
+            json.dump(entry, f)
+            f.write("\n")
+
     @property
     def iteration_count(self) -> int:
         return self._iteration_count