diff --git a/examples/compaction_example.py b/examples/compaction_example.py new file mode 100644 index 000000000..69e3cecfd --- /dev/null +++ b/examples/compaction_example.py @@ -0,0 +1,65 @@ +""" +Compaction example: run RLM with compaction enabled and a low threshold +to trigger summarization on gpt-5-nano. + +Uses a low compaction_threshold_pct so compaction runs after several iterations. +The task forces many separate REPL turns with substantial output so the root +context grows and compaction definitely triggers. The REPL variable `history` +holds trajectory segments and summaries. +""" + +import os + +from dotenv import load_dotenv + +from rlm import RLM +from rlm.logger import RLMLogger + +load_dotenv() + +# Low threshold so compaction triggers after a few iterations (~2% of context). +# Use 0.85 in production. +COMPACTION_THRESHOLD_PCT = 0.03 + +logger = RLMLogger() +rlm = RLM( + backend="portkey", + backend_kwargs={ + "model_name": "@openai/gpt-5-nano", + "api_key": os.getenv("PORTKEY_API_KEY"), + }, + environment="local", + environment_kwargs={}, + max_depth=1, + max_iterations=10, + compaction=True, + compaction_threshold_pct=COMPACTION_THRESHOLD_PCT, + verbose=True, + logger=logger, +) + +# Hard task that forces many iterations: find the 50th prime using at least 8 +# separate REPL blocks, one per turn. Each block must produce visible output. +# This grows message history (long reasoning + code + execution results) so +# compaction triggers. +prompt = ( + "Find the 50th prime number. You MUST use at least 8 separate REPL blocks, " + "each in its own response (one block per message). Do NOT combine steps.\n\n" + "Required structure: " + "Block 1: Define a function is_prime(n) and test it on a few numbers; print the results. " + "Block 2: Write a loop that counts primes and print the first 10 primes. " + "Block 3: Extend to count up to 20 primes and print them. " + "Block 4: Count up to 30 primes and print the 25th. " + "Block 5: Count up to 40 primes and print the 35th. " + "Block 6: Count up to 50 primes and print the 45th. " + "Block 7: Print the full list of the first 50 primes (so we see all 50). " + "Block 8: Set answer = (the 50th prime) and call FINAL_VAR(answer).\n\n" + "Each block must run alone. Show your reasoning briefly before each block. " + "After each block, wait for the execution result before writing the next block." +) + +result = rlm.completion(prompt, root_prompt=prompt) + +print("Response:", result.response) +print("Execution time:", result.execution_time) +print("Metadata:", result.metadata) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 6c0646ec3..f236f45b8 100644 --- a/rlm/core/rlm.py +++ b/rlm/core/rlm.py @@ -27,6 +27,7 @@ build_user_prompt, ) from rlm.utils.rlm_utils import filter_sensitive_keys +from rlm.utils.token_utils import count_tokens, get_context_limit class RLM: @@ -54,6 +55,8 @@ def __init__( persistent: bool = False, custom_tools: dict[str, Any] | None = None, custom_sub_tools: dict[str, Any] | None = None, + compaction: bool = False, + compaction_threshold_pct: float = 0.85, ): """ Args: @@ -74,6 +77,10 @@ def __init__( values are callable functions. These are injected into the REPL globals. custom_sub_tools: Dict of custom tools for sub-agents (llm_query calls). If None, inherits from custom_tools. Pass an empty dict {} to disable tools for sub-agents. + compaction: If True, keep full root model history in REPL variable `history` and compact + when root context reaches compaction_threshold_pct of the model's context limit. + compaction_threshold_pct: When compaction is on, trigger summarization when root + message token count reaches this fraction of the model context limit (default 0.85). """ # Store config for spawning per-completion self.backend = backend @@ -98,6 +105,9 @@ def __init__( # Sub-tools: if None, inherit from custom_tools; if {}, no tools for sub-agents self.custom_sub_tools = custom_sub_tools if custom_sub_tools is not None else custom_tools + self.compaction = compaction + self.compaction_threshold_pct = compaction_threshold_pct + self.depth = depth self.max_depth = max_depth self.max_iterations = max_iterations @@ -181,6 +191,8 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]): env_kwargs["custom_tools"] = self.custom_tools if self.custom_sub_tools is not None: env_kwargs["custom_sub_tools"] = self.custom_sub_tools + if self.compaction and self.environment_type == "local": + env_kwargs["compaction"] = True environment: BaseEnv = get_environment(self.environment_type, env_kwargs) if self.persistent: @@ -204,7 +216,11 @@ def _setup_prompt(self, prompt: str | dict[str, Any]) -> list[dict[str, Any]]: query_metadata=metadata, custom_tools=self.custom_tools, ) - + if self.compaction: + message_history[0]["content"] += ( + "\n\nThe full conversation history (trajectory segments and any summaries) " + "is available in the REPL variable `history` as a list." + ) return message_history def completion( @@ -236,6 +252,18 @@ def completion( message_history = self._setup_prompt(prompt) for i in range(self.max_iterations): + if self.compaction and hasattr(environment, "append_compaction_entry"): + current_tokens, threshold_tokens, max_tokens = self._get_compaction_status( + message_history + ) + self.verbose.print_compaction_status( + current_tokens, threshold_tokens, max_tokens + ) + if current_tokens >= threshold_tokens: + self.verbose.print_compaction() + message_history = self._compact_history( + lm_handler, environment, message_history + ) # Current prompt = message history + additional prompt suffix context_count = ( environment.get_context_count() @@ -257,8 +285,14 @@ def completion( environment=environment, ) - # Check if RLM is done and has a final answer. - final_answer = find_final_answer(iteration.response, environment=environment) + # Check if RLM is done and has a final answer. Prefer FINAL_VAR result from REPL execution. + final_answer = None + for block in iteration.code_blocks: + if getattr(block.result, "final_answer", None): + final_answer = block.result.final_answer + break + if final_answer is None: + final_answer = find_final_answer(iteration.response, environment=environment) iteration.final_answer = final_answer # If logger is used, log the iteration. @@ -294,6 +328,8 @@ def completion( # Update message history with the new messages. message_history.extend(new_messages) + if self.compaction and hasattr(environment, "append_compaction_entry"): + environment.append_compaction_entry(new_messages) # Default behavior: we run out of iterations, provide one final answer time_end = time.perf_counter() @@ -317,6 +353,50 @@ def completion( metadata=self.logger.get_trajectory() if self.logger else None, ) + def _get_compaction_status(self, message_history: list[dict[str, Any]]) -> tuple[int, int, int]: + """Return (current_tokens, threshold_tokens, max_tokens) for compaction.""" + model_name = ( + self.backend_kwargs.get("model_name", "unknown") if self.backend_kwargs else "unknown" + ) + max_tokens = get_context_limit(model_name) + current_tokens = count_tokens(message_history, model_name) + threshold_tokens = int(self.compaction_threshold_pct * max_tokens) + return current_tokens, threshold_tokens, max_tokens + + def _should_compact(self, message_history: list[dict[str, Any]]) -> bool: + """True when root message history is at or over the compaction threshold.""" + current_tokens, threshold_tokens, _ = self._get_compaction_status(message_history) + return current_tokens >= threshold_tokens + + def _compact_history( + self, + lm_handler: LMHandler, + environment: BaseEnv, + message_history: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """ + Summarize current trajectory, append summary to REPL history, and return + a short message_history with the summary as the new starting point. + """ + summary_prompt = message_history + [ + { + "role": "user", + "content": "Very concisely summarize what you have been doing so far in 1–3 short paragraphs. Be extremely brief. This summary will be used to continue the conversation.", + } + ] + summary = lm_handler.completion(summary_prompt) + if hasattr(environment, "append_compaction_entry"): + environment.append_compaction_entry({"type": "summary", "content": summary}) + # Keep system + initial assistant (metadata), then summary + continue + new_history = message_history[:2] + [ + {"role": "assistant", "content": summary}, + { + "role": "user", + "content": "Continue from the above summary. The full history (including this summary) is in the REPL variable `history`. Your next action:", + }, + ] + return new_history + def _completion_turn( self, prompt: str | dict[str, Any], diff --git a/rlm/core/types.py b/rlm/core/types.py index 631051ea3..239e092ca 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -131,6 +131,7 @@ class REPLResult: locals: dict execution_time: float llm_calls: list["RLMChatCompletion"] + final_answer: str | None = None def __init__( self, @@ -139,12 +140,14 @@ def __init__( locals: dict, execution_time: float = None, rlm_calls: list["RLMChatCompletion"] = None, + final_answer: str | None = None, ): self.stdout = stdout self.stderr = stderr self.locals = locals self.execution_time = execution_time self.rlm_calls = rlm_calls or [] + self.final_answer = final_answer def __str__(self): return f"REPLResult(stdout={self.stdout}, stderr={self.stderr}, locals={self.locals}, execution_time={self.execution_time}, rlm_calls={len(self.rlm_calls)})" @@ -156,6 +159,7 @@ def to_dict(self): "locals": {k: _serialize_value(v) for k, v in self.locals.items()}, "execution_time": self.execution_time, "rlm_calls": [call.to_dict() for call in self.rlm_calls], + "final_answer": self.final_answer, } diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index 65a49697c..d2ba13fd9 100644 --- a/rlm/environments/local_repl.py +++ b/rlm/environments/local_repl.py @@ -132,6 +132,7 @@ def __init__( depth: int = 1, custom_tools: dict[str, Any] | None = None, custom_sub_tools: dict[str, Any] | None = None, + compaction: bool = False, **kwargs, ): super().__init__(persistent=persistent, depth=depth, **kwargs) @@ -142,6 +143,7 @@ def __init__( self._lock = threading.Lock() self._context_count: int = 0 self._history_count: int = 0 + self.compaction = compaction # Custom tools: functions available in the REPL self.custom_tools = custom_tools or {} @@ -156,6 +158,10 @@ def __init__( # Setup globals, locals, and modules in environment. self.setup() + if compaction: + self._compaction_history: list[Any] = [] + self.locals["history"] = self._compaction_history + # Load context if provided if context_payload is not None: self.load_context(context_payload) @@ -175,6 +181,8 @@ def setup(self): # Track LLM calls made during code execution self._pending_llm_calls: list[RLMChatCompletion] = [] + # When FINAL_VAR is called inside a REPL block, we store the value here for the main loop + self._last_final_answer: str | None = None # Add helper functions self.globals["FINAL_VAR"] = self._final_var @@ -192,13 +200,19 @@ def setup(self): # For non-callable values (constants, data), add to locals self.locals[name] = value - def _final_var(self, variable_name: str) -> str: - """Return the value of a variable as a final answer.""" + def _final_var(self, variable_name: str | Any) -> str: + """Return the value of a variable as a final answer for the main model, or stringify a direct value.""" + if not isinstance(variable_name, str): + answer = str(variable_name) + self._last_final_answer = answer + return answer variable_name = variable_name.strip().strip("\"'") if variable_name in self.locals: - return str(self.locals[variable_name]) + answer = str(self.locals[variable_name]) + self._last_final_answer = answer + return answer - # Provide helpful error message with available variables + # Provide helpful error message with available variables (do not set _last_final_answer) available = [k for k in self.locals.keys() if not k.startswith("_")] if available: return ( @@ -358,6 +372,17 @@ def get_history_count(self) -> int: """Return the number of conversation histories stored.""" return self._history_count + def append_compaction_entry(self, entry: list[dict[str, Any]] | dict[str, Any]) -> None: + """ + Append a trajectory segment or a summary to the compaction history. + + Entry is either a list of message dicts (trajectory segment) or + a dict with "type": "summary" and "content": str. + """ + if not self.compaction: + return + self._compaction_history.append(copy.deepcopy(entry)) + @contextmanager def _capture_output(self): """Thread-safe context manager to capture stdout/stderr.""" @@ -393,8 +418,10 @@ def _restore_scaffold(self) -> None: self.globals["SHOW_VARS"] = self._show_vars elif name == "context" and "context_0" in self.locals: self.locals["context"] = self.locals["context_0"] - elif name == "history" and "history_0" in self.locals: + elif name == "history" and "history_0" in self.locals and not self.compaction: self.locals["history"] = self.locals["history_0"] + elif name == "history" and self.compaction: + self.locals["history"] = self._compaction_history def execute_code(self, code: str) -> REPLResult: """Execute code in the persistent namespace and return result.""" @@ -422,12 +449,16 @@ def execute_code(self, code: str) -> REPLResult: stdout = stdout_buf.getvalue() stderr = stderr_buf.getvalue() + f"\n{type(e).__name__}: {e}" + final_answer = self._last_final_answer + self._last_final_answer = None + return REPLResult( stdout=stdout, stderr=stderr, locals=self.locals.copy(), execution_time=time.perf_counter() - start_time, rlm_calls=self._pending_llm_calls.copy(), + final_answer=final_answer, ) def __enter__(self): diff --git a/rlm/logger/verbose.py b/rlm/logger/verbose.py index 2f612a389..9f88b0a3e 100644 --- a/rlm/logger/verbose.py +++ b/rlm/logger/verbose.py @@ -324,6 +324,49 @@ def print_iteration(self, iteration: RLMIteration, iteration_num: int) -> None: execution_time=call.execution_time, ) + def print_compaction_status( + self, + current_tokens: int, + threshold_tokens: int, + max_tokens: int, + ) -> None: + """Print how close root context is to compaction threshold (before next turn).""" + if not self.enabled: + return + pct = (current_tokens / threshold_tokens * 100) if threshold_tokens else 0 + line = Text() + line.append("Context: ", style=STYLE_MUTED) + line.append(f"{current_tokens:,}", style=STYLE_TEXT) + line.append(" / ", style=STYLE_MUTED) + line.append(f"{threshold_tokens:,}", style=STYLE_SECONDARY) + line.append(" tokens ", style=STYLE_MUTED) + line.append(f"({pct:.0f}% of compaction threshold)", style=STYLE_MUTED) + if current_tokens >= threshold_tokens: + line.append(" — compacting", style=STYLE_WARNING) + self.console.print(line) + + def print_compaction(self) -> None: + """Print that context compaction (summarization) is running.""" + if not self.enabled: + return + + title = Text() + title.append("◐ ", style=STYLE_ACCENT) + title.append("Compaction", style=STYLE_SECONDARY) + title.append(" — summarizing context, continuing from summary", style=STYLE_MUTED) + + self.console.print() + self.console.print( + Panel( + Text("Root context reached threshold. Summarizing and continuing."), + title=title, + title_align="left", + border_style=COLORS["muted"], + padding=(0, 1), + ) + ) + self.console.print() + def print_final_answer(self, answer: Any) -> None: """Print the final answer.""" if not self.enabled: diff --git a/rlm/utils/parsing.py b/rlm/utils/parsing.py index f0bdd40dc..904aa29db 100644 --- a/rlm/utils/parsing.py +++ b/rlm/utils/parsing.py @@ -49,7 +49,14 @@ def find_final_answer(text: str, environment: "BaseEnv | None" = None) -> str | result = environment.execute_code(f"print(FINAL_VAR({variable_name!r}))") final_answer = result.stdout.strip() if final_answer == "": - final_answer = result.stderr.strip() or "" + return None + # Don't treat FINAL_VAR "variable not found" as final answer (so RLM continues) + if ( + "Variable '" in final_answer + and "' not found" in final_answer + and "FINAL_VAR" in final_answer + ): + return None return final_answer return None diff --git a/rlm/utils/token_utils.py b/rlm/utils/token_utils.py new file mode 100644 index 000000000..79febbe10 --- /dev/null +++ b/rlm/utils/token_utils.py @@ -0,0 +1,143 @@ +""" +Token counting and model context limits for compaction and context sizing. + +Uses tiktoken for OpenAI-style models when available; otherwise estimates +with ~4 characters per token. +""" + +from typing import Any + +# Default context limit when model is unknown (tokens) +DEFAULT_CONTEXT_LIMIT = 128_000 + +# Characters per token when tokenizer is unavailable (conservative estimate) +CHARS_PER_TOKEN_ESTIMATE = 4 + +# Model context limits (max input context in tokens). +# Match: key contained in model_name (e.g. "gpt-4o" matches "@openai/gpt-4o"). +# Longest matching key wins. +MODEL_CONTEXT_LIMITS: dict[str, int] = { + # OpenAI (GPT-5: 272k input, 128k reasoning+output) + "gpt-5-nano": 128_000, + "gpt-5": 272_000, + "gpt-4o-mini": 128_000, + "gpt-4o-2024": 128_000, + "gpt-4o": 128_000, + "gpt-4-turbo-preview": 128_000, + "gpt-4-turbo": 128_000, + "gpt-4-32k": 32_768, + "gpt-4": 8_192, + "gpt-3.5-turbo-16k": 16_385, + "gpt-3.5-turbo": 16_385, + "o1-mini": 128_000, + "o1-preview": 128_000, + "o1": 200_000, + # Anthropic + "claude-3-5-sonnet": 200_000, + "claude-3-5-haiku": 200_000, + "claude-3-opus": 200_000, + "claude-3-sonnet": 200_000, + "claude-3-haiku": 200_000, + "claude-2.1": 200_000, + "claude-2": 100_000, + # Gemini + "gemini-2.5-flash": 1_000_000, + "gemini-2.5-pro": 1_000_000, + "gemini-2.0-flash": 1_000_000, + "gemini-1.5-pro": 1_000_000, + "gemini-1.5-flash": 1_000_000, + "gemini-1.0-pro": 30_720, + # Qwen (Alibaba) + "qwen3-max": 256_000, + "qwen3-72b": 128_000, + "qwen3-32b": 128_000, + "qwen3-8b": 32_768, + "qwen3": 128_000, + # Kimi (Moonshot) + "kimi-k2.5": 262_000, + "kimi-k2-0905": 256_000, + "kimi-k2-thinking": 256_000, + "kimi-k2": 128_000, + "kimi": 128_000, + # GLM (Zhipu) + "glm-4.6": 200_000, + "glm-4-9b": 1_000_000, + "glm-4": 128_000, + "glm": 128_000, +} + + +def get_context_limit(model_name: str) -> int: + """ + Return max context size in tokens for a model. + + Matches when the dict key is contained in model_name (e.g. "gpt-4o" matches + "@openai/gpt-4o"). Longest matching key wins. Falls back to DEFAULT_CONTEXT_LIMIT + for unknown models. + """ + if not model_name or model_name == "unknown": + return DEFAULT_CONTEXT_LIMIT + exact = MODEL_CONTEXT_LIMITS.get(model_name) + if exact is not None: + return exact + best = 0 + best_limit = DEFAULT_CONTEXT_LIMIT + for key, limit in MODEL_CONTEXT_LIMITS.items(): + if key in model_name and len(key) > best: + best = len(key) + best_limit = limit + return best_limit + + +def _count_tokens_tiktoken(messages: list[dict[str, Any]], model_name: str) -> int | None: + """Count tokens with tiktoken if available. Returns None on failure.""" + try: + import tiktoken + except ImportError: + return None + try: + enc = tiktoken.encoding_for_model(model_name) + except Exception: + try: + enc = tiktoken.get_encoding("cl100k_base") + except Exception: + return None + total = 0 + # Approximate OpenAI message format overhead per message + tokens_per_message = 3 + tokens_per_name = 1 + for m in messages: + total += tokens_per_message + content = m.get("content") + if isinstance(content, str): + total += len(enc.encode(content)) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + total += len(enc.encode(part.get("text", "") or "")) + elif content is not None and content != "": + total += len(enc.encode(str(content))) + if m.get("name"): + total += tokens_per_name + return total + + +def count_tokens(messages: list[dict[str, Any]], model_name: str) -> int: + """ + Count tokens in a list of message dicts (role, content). + + Uses tiktoken for OpenAI-style models when the package is available; + otherwise estimates with character length / CHARS_PER_TOKEN_ESTIMATE. + """ + if not messages: + return 0 + if model_name and model_name != "unknown": + n = _count_tokens_tiktoken(messages, model_name) + if n is not None: + return n + # Fallback: count chars (stringify in case content is not str, e.g. list) + total_chars = 0 + for m in messages: + raw = m.get("content", "") or "" + total_chars += len(raw) if isinstance(raw, str) else len(str(raw)) + return (total_chars + CHARS_PER_TOKEN_ESTIMATE - 1) // CHARS_PER_TOKEN_ESTIMATE diff --git a/tests/test_parsing.py b/tests/test_parsing.py index 0ff910faa..ddb8d436f 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -259,13 +259,25 @@ def test_final_var_retrieves_actual_variables_from_environment(self): result = find_final_answer(text, environment=env) assert result == "100", f"Expected '100', got '{result}'" - # Test that non-existent variable returns error message + # Non-existent variable: find_final_answer must return None (not the error string) + # so the RLM loop continues and the model can fix it text = "FINAL_VAR(nonexistent)" result = find_final_answer(text, environment=env) - assert "Error" in result or "not found" in result.lower() + assert result is None, "must return None for variable-not-found, not the error message" finally: env.cleanup() + def test_final_var_variable_not_found_returns_none(self): + """When env returns FINAL_VAR 'variable not found' error, find_final_answer must return None.""" + mock_env = Mock() + mock_env.execute_code.return_value = REPLResult( + stdout="Error: Variable 'missing' not found. Available variables: []. You must create and assign a variable BEFORE calling FINAL_VAR on it.", + stderr="", + locals={}, + ) + result = find_final_answer("FINAL_VAR(missing)", environment=mock_env) + assert result is None + class TestFormatExecutionResult: """Tests for format_execution_result function."""