Merge branch 'main' into parallel-rlm-subcalls

Author: alexzhang13

examples/quickstart.py

@@ -1,4 +1,6 @@
 import os
+import random
+import string
 
 from dotenv import load_dotenv
 
@@ -7,21 +9,31 @@
 
 load_dotenv()
 
-logger = RLMLogger(log_dir="./logs")
+# Generate a large text file with a hidden secret number
+secret_number = random.randint(100_000_000, 999_999_999)
+filler_lines = ["".join(random.choices(string.ascii_lowercase + " ", k=120)) for _ in range(50_000)]
+insert_at = random.randint(len(filler_lines) // 3, 2 * len(filler_lines) // 3)
+filler_lines.insert(insert_at, f"SECRET_NUMBER={secret_number}")
+haystack = "\n".join(filler_lines)
 
 rlm = RLM(
-    backend="portkey",  # or "openai", etc.
+    backend="openai",
     backend_kwargs={
-        "model_name": "@openai/gpt-5-nano",
-        "api_key": os.getenv("PORTKEY_API_KEY"),
+        "model_name": "gpt-5-nano",
+        "api_key": os.getenv("OPENAI_API_KEY"),
     },
     environment="local",
-    environment_kwargs={},
-    max_depth=1,
-    logger=logger,
-    verbose=True,  # For printing to console with rich, disabled by default.
+    max_iterations=10,
+    logger=RLMLogger(log_dir="./logs"),
+    verbose=True,
 )
 
-result = rlm.completion("Print me the first 5 powers of two, each on a newline.")
+result = rlm.completion(
+    "The context contains ~50k lines of random text with a single line "
+    "matching the pattern SECRET_NUMBER=<digits>. Find and return ONLY the "
+    f"numeric value.\n\n{haystack}"
+)
 
-print(result)
+print(f"\nModel found: {result.response}")
+print(f"Actual number: {secret_number}")
+print(f"Correct: {str(secret_number) in result.response}")

rlm/core/rlm.py

@@ -248,6 +248,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
                 env_kwargs["custom_sub_tools"] = self.custom_sub_tools
             if self.compaction and self.environment_type == "local":
                 env_kwargs["compaction"] = True
+            env_kwargs["max_concurrent_subcalls"] = self.max_concurrent_subcalls
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
 
             if self.persistent:

rlm/environments/base_env.py

@@ -212,9 +212,12 @@ class BaseEnv(ABC):
             }
     """
 
-    def __init__(self, persistent: bool = False, depth: int = 1, **kwargs):
+    def __init__(
+        self, persistent: bool = False, depth: int = 1, max_concurrent_subcalls: int = 4, **kwargs
+    ):
         self.persistent = persistent
         self.depth = depth
+        self.max_concurrent_subcalls = max_concurrent_subcalls
         self.kwargs = kwargs
 
     @abstractmethod

rlm/environments/local_repl.py

@@ -172,16 +172,12 @@ def __init__(
         subcall_semaphore: threading.Semaphore | None = None,
         **kwargs,
     ):
-        """
-        Args:
-            max_concurrent_subcalls: Max children to dispatch in parallel from
-                THIS REPL instance (controls local ThreadPoolExecutor size).
-                Default 1 = sequential (identical to previous behavior).
-            subcall_semaphore: Global semaphore bounding total concurrent child
-                RLMs across ALL depths. If None, uses the module-level default
-                (16 slots). Pass a custom one for testing or tighter control.
-        """
-        super().__init__(persistent=persistent, depth=depth, **kwargs)
+        super().__init__(
+            persistent=persistent,
+            depth=depth,
+            max_concurrent_subcalls=max_concurrent_subcalls,
+            **kwargs,
+        )
 
         self.lm_handler_address = lm_handler_address
         self.subcall_fn = subcall_fn  # Callback for recursive RLM calls (depth > 1 support)
@@ -363,8 +359,7 @@ def _rlm_query(self, prompt: str, model: str | None = None) -> str:
         return self._llm_query(prompt, model)
 
     def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]:
-        """Spawn recursive RLM sub-calls for multiple prompts.
-
+        """Spawn recursive RLM sub-calls for multiple prompts in parallel.
         Each prompt gets its own child RLM for deeper thinking.
 
         Execution modes (controlled by max_concurrent_subcalls):
@@ -401,10 +396,25 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li
             for prompt in prompts:
                 try:
                     completion = self.subcall_fn(prompt, model)
-                    self._pending_llm_calls.append(completion)
-                    results.append(completion.response)
+                    with lock:
+                        completions.append((index, completion))
+                    results[index] = completion.response
                 except Exception as e:
-                    results.append(f"Error: RLM query failed - {e}")
+                    results[index] = f"Error: RLM query failed - {e}"
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = [
+                    executor.submit(_run_subcall, i, prompt) for i, prompt in enumerate(prompts)
+                ]
+                # Wait for all futures to complete; exceptions are captured inside _run_subcall
+                for future in as_completed(futures):
+                    future.result()  # Re-raises unexpected executor errors
+
+            # Append completions in original prompt order for deterministic metadata
+            completions.sort(key=lambda x: x[0])
+            for _, completion in completions:
+                self._pending_llm_calls.append(completion)
+
             return results
 
         # Fall back to plain batched LM call if no recursive capability