updated compaction logic

Author: alexzhang13

examples/compaction_example.py

@@ -1,11 +1,17 @@
 """
 Compaction example: run RLM with compaction enabled and a low threshold
-to trigger summarization on gpt-5-nano.
+to trigger summarization.
 
-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.
+The task uses random data so earlier results are impossible to recompute —
+the model *must* recover them from `history` after compaction fires.
+
+Phase 1: Generate random datasets and compute statistics (fills context).
+Phase 2: Compaction fires, summarizing the conversation.
+Phase 3: Combine the earlier statistics into a final answer. Since the data
+          was random, the model cannot just re-derive the numbers.
+
+Usage:
+    PORTKEY_API_KEY=... python examples/compaction_example.py
 """
 
 import os
@@ -17,9 +23,9 @@
 
 load_dotenv()
 
-# Low threshold so compaction triggers after a few iterations (~2% of context).
+# Low threshold so compaction triggers after a few iterations.
 # Use 0.85 in production.
-COMPACTION_THRESHOLD_PCT = 0.03
+COMPACTION_THRESHOLD_PCT = 0.02
 
 logger = RLMLogger()
 rlm = RLM(
@@ -31,35 +37,32 @@
     environment="local",
     environment_kwargs={},
     max_depth=1,
-    max_iterations=10,
+    max_iterations=12,
     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."
+    "Complete the following steps, each in its own REPL block (one block per message). "
+    "Do NOT combine steps.\n\n"
+    "Step 1: Use `import random; random.seed(42)` then generate a list of 50 random "
+    "integers between 1 and 1000. Print ALL of them. Store as `data_a`.\n\n"
+    "Step 2: Generate another 50 random integers (continuing the same RNG stream) "
+    "between 1 and 1000. Print ALL of them. Store as `data_b`.\n\n"
+    "Step 3: Compute and print the mean, median, min, max, and standard deviation "
+    "of `data_a`. Store the mean as `mean_a`.\n\n"
+    "Step 4: Compute and print the mean, median, min, max, and standard deviation "
+    "of `data_b`. Store the mean as `mean_b`.\n\n"
+    "Step 5: You previously computed `mean_a` and `mean_b`. "
+    "Compute final_answer = round(mean_a + mean_b, 2) and call FINAL_VAR(final_answer)."
 )
 
 result = rlm.completion(prompt, root_prompt=prompt)
 
-print("Response:", result.response)
-print("Execution time:", result.execution_time)
-print("Metadata:", result.metadata)
+print("\n" + "=" * 60)
+print("RESULT")
+print("=" * 60)
+print(f"Response: {result.response}")
+print(f"Execution time: {result.execution_time:.2f}s")

examples/compaction_history_retrieval_example.py

@@ -0,0 +1,65 @@
+"""
+Compaction + history retrieval example.
+
+Like compaction_example.py but with more intermediate results to recover.
+Uses random data across 4 groups so none of the statistics are re-derivable
+after compaction — the model must look in `history` to find them.
+
+Usage:
+    PORTKEY_API_KEY=... python examples/compaction_history_retrieval_example.py
+"""
+
+import os
+
+from dotenv import load_dotenv
+
+from rlm import RLM
+from rlm.logger import RLMLogger
+
+load_dotenv()
+
+# Very low threshold so compaction fires after the first few steps.
+COMPACTION_THRESHOLD_PCT = 0.02
+
+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=15,
+    compaction=True,
+    compaction_threshold_pct=COMPACTION_THRESHOLD_PCT,
+    verbose=True,
+    logger=logger,
+)
+
+prompt = (
+    "Complete the following steps, each in its own REPL block (one block per message). "
+    "Do NOT combine steps.\n\n"
+    "Step 1: `import random; random.seed(99)`. Generate 40 random floats in [0, 100). "
+    "Print all of them. Store as `group_1`.\n\n"
+    "Step 2: Generate another 40 random floats (same RNG stream). Print all. "
+    "Store as `group_2`.\n\n"
+    "Step 3: Generate another 40 random floats. Print all. Store as `group_3`.\n\n"
+    "Step 4: Generate another 40 random floats. Print all. Store as `group_4`.\n\n"
+    "Step 5: Compute the mean of each group. Print all four means. "
+    "Store as `m1`, `m2`, `m3`, `m4`.\n\n"
+    "Step 6: Compute the standard deviation of each group. Print all four. "
+    "Store as `s1`, `s2`, `s3`, `s4`.\n\n"
+    "Step 7: You previously computed four means and four standard deviations. "
+    "Compute final_answer = round((m1 + m2 + m3 + m4) / (s1 + s2 + s3 + s4), 4). "
+    "Print the full equation with values and call FINAL_VAR(final_answer)."
+)
+
+result = rlm.completion(prompt, root_prompt=prompt)
+
+print("\n" + "=" * 60)
+print("RESULT")
+print("=" * 60)
+print(f"Response: {result.response}")
+print(f"Execution time: {result.execution_time:.2f}s")

rlm/core/rlm.py

@@ -301,6 +301,7 @@ def completion(
         with self._spawn_completion_context(prompt) as (lm_handler, environment):
             message_history = self._setup_prompt(prompt)
 
+            compaction_count = 0
             try:
                 for i in range(self.max_iterations):
                     # Check timeout before each iteration
@@ -315,9 +316,10 @@ def completion(
                             current_tokens, threshold_tokens, max_tokens
                         )
                         if current_tokens >= threshold_tokens:
+                            compaction_count += 1
                             self.verbose.print_compaction()
                             message_history = self._compact_history(
-                                lm_handler, environment, message_history
+                                lm_handler, environment, message_history, compaction_count
                             )
 
                     # Current prompt = message history + additional prompt suffix
@@ -541,6 +543,7 @@ def _compact_history(
         lm_handler: LMHandler,
         environment: BaseEnv,
         message_history: list[dict[str, Any]],
+        compaction_count: int = 1,
     ) -> list[dict[str, Any]]:
         """
         Summarize current trajectory, append summary to REPL history, and return
@@ -549,7 +552,15 @@ def _compact_history(
         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.",
+                "content": (
+                    "Summarize your progress so far. Include:\n"
+                    "1. Which steps/sub-tasks you have completed and which remain.\n"
+                    "2. Any concrete intermediate results (numbers, values, variable names) "
+                    "you computed — preserve these exactly.\n"
+                    "3. What your next action should be.\n"
+                    "Be concise (1–3 paragraphs) but preserve all key results and your "
+                    "current position in the task."
+                ),
             }
         ]
         summary = lm_handler.completion(summary_prompt)
@@ -560,7 +571,13 @@ def _compact_history(
             {"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:",
+                "content": (
+                    f"Your conversation has been compacted {compaction_count} time(s). "
+                    "Continue from the above summary. Do NOT repeat work you have already "
+                    "completed. Use SHOW_VARS() to check which REPL variables exist, "
+                    "and check `history` for full context. "
+                    "Your next action:"
+                ),
             },
         ]
         return new_history

rlm/utils/prompts.py

@@ -21,8 +21,10 @@
 - Use `llm_query` for simple, one-shot tasks: extracting info from a chunk, summarizing text, answering a factual question, classifying content. These are fast single LLM calls.
 - Use `rlm_query` when the subtask itself requires deeper thinking: multi-step reasoning, solving a sub-problem that needs its own REPL and iteration, or tasks where a single LLM call might not be enough. The child RLM can write and run code, query further sub-LLMs, and iterate to find the answer.
 
+**Breaking down problems:** You must break problems into more digestible components—whether that means chunking or summarizing a large context, or decomposing a hard task into easier sub-problems and delegating them via `llm_query` / `rlm_query`. Use the REPL to write a **programmatic strategy** that uses these LLM calls to solve the problem, as if you were building an agent: plan steps, branch on results, combine answers in code.
+
 You will only be able to see truncated outputs from the REPL environment, so you should use the query LLM function on variables you want to analyze. You will find this function especially useful when you have to analyze the semantics of the context. Use these variables as buffers to build up your final answer.
-Make sure to explicitly look through the entire context in REPL before answering your query. An example strategy is to first look at the context and figure out a chunking strategy, then break up the context into smart chunks, and query an LLM per chunk with a particular question and save the answers to a buffer, then query an LLM with all the buffers to produce your final answer.
+Make sure to explicitly look through the entire context in REPL before answering your query. Break the context and the problem into digestible pieces: e.g. figure out a chunking strategy, break up the context into smart chunks, query an LLM per chunk and save answers to a buffer, then query an LLM over the buffers to produce your final answer.
 
 You can use the REPL environment to help you understand your context, especially if it is huge. Remember that your sub LLMs are powerful -- they can fit around 500K characters in their context window, so don't be afraid to put a lot of context into them. For example, a viable strategy is to feed 10 documents per sub-LLM query. Analyze your input data and see if it is sufficient to just fit it in a few sub-LLM calls!
 
@@ -66,25 +68,28 @@
 final_answer = llm_query(f"Aggregating all the answers per chunk, answer the original query about total number of jobs: {{query}}\\n\\nAnswers:\\n" + "\\n".join(answers))
 ```
 
-For subtasks that require deeper reasoning (e.g. solving a complex sub-problem), use `rlm_query` instead:
+For subtasks that require deeper reasoning (e.g. solving a complex sub-problem), use `rlm_query` instead. The child gets its own REPL to iterate; you can then use the result in parent logic:
 ```repl
-# The child RLM gets its own REPL and can iterate to solve this sub-problem
-hard_answer = rlm_query(f"Analyze this complex dataset and determine the trend: {{data}}")
-print(hard_answer)
+# Child RLM solves the sub-problem in its own REPL; we use the result in code
+trend = rlm_query(f"Analyze this dataset and conclude with one word: up, down, or stable: {{data}}")
+if "up" in trend.lower():
+    recommendation = "Consider increasing exposure."
+elif "down" in trend.lower():
+    recommendation = "Consider hedging."
+else:
+    recommendation = "Hold position."
+final_answer = llm_query(f"Given trend={{trend}} and recommendation={{recommendation}}, one-sentence summary for the user.")
 ```
 
-As a final example, after analyzing the context and realizing its separated by Markdown headers, we can maintain state through buffers by chunking the context by headers, and iteratively querying an LLM over it:
+As a final example, implement the solution as a **program**: try one approach via `rlm_query`; inspect the result and branch. If it suffices, use it. If not, break into one easier subproblem and delegate that only. More branches, one path runs—don't load the model. Example: prove √2 irrational.
 ```repl
-# After finding out the context is separated by Markdown headers, we can chunk, summarize, and answer
-import re
-sections = re.split(r'### (.+)', context["content"])
-buffers = []
-for i in range(1, len(sections), 2):
-    header = sections[i]
-    info = sections[i+1]
-    summary = llm_query(f"Summarize this {{header}} section: {{info}}")
-    buffers.append(f"{{header}}: {{summary}}")
-final_answer = llm_query(f"Based on these summaries, answer the original query: {{query}}\\n\\nSummaries:\\n" + "\\n".join(buffers))
+r = rlm_query("Prove √2 irrational. Give a 1-2 sentence proof, or reply only: USE_LEMMA or USE_CONTRADICTION.")
+if "USE_LEMMA" in r.upper():
+    final_answer = rlm_query("Prove 'n^2 even => n even' then use it to show √2 irrational. Two sentences.")
+elif "USE_CONTRADICTION" in r.upper():
+    final_answer = rlm_query("Prove √2 irrational by contradiction. Two sentences.")
+else:
+    final_answer = r
 ```
 In the next step, we can return FINAL_VAR(final_answer).