Merge remote-tracking branch 'refs/remotes/origin/main'

Author: alexzhang13

rlm/clients/openai.py

@@ -14,12 +14,17 @@
 DEFAULT_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 DEFAULT_OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
 DEFAULT_VERCEL_API_KEY = os.getenv("AI_GATEWAY_API_KEY")
+DEFAULT_PRIME_API_KEY = os.getenv("PRIME_API_KEY")
 DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/"
 
 
 class OpenAIClient(BaseLM):
     """
     LM Client for running models with the OpenAI API. Works with vLLM as well.
+
+    Any additional keyword arguments (e.g. default_headers, default_query, max_retries)
+    are passed through to the underlying openai.OpenAI and openai.AsyncOpenAI constructors.
+    Only model_name is excluded, since it is not a client constructor argument.
     """
 
     def __init__(
@@ -38,12 +43,19 @@ def __init__(
                 api_key = DEFAULT_OPENROUTER_API_KEY
             elif base_url == "https://ai-gateway.vercel.sh/v1":
                 api_key = DEFAULT_VERCEL_API_KEY
-
-        # For vLLM, set base_url to local vLLM server address.
-        self.client = openai.OpenAI(api_key=api_key, base_url=base_url, timeout=self.timeout)
-        self.async_client = openai.AsyncOpenAI(
-            api_key=api_key, base_url=base_url, timeout=self.timeout
-        )
+            elif base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
+                api_key = DEFAULT_PRIME_API_KEY
+
+        # Pass through arbitrary kwargs to the OpenAI client (e.g. default_headers, default_query, max_retries).
+        # Exclude model_name since it is not an OpenAI client constructor argument.
+        client_kwargs = {
+            "api_key": api_key,
+            "base_url": base_url,
+            "timeout": self.timeout,
+            **{k: v for k, v in self.kwargs.items() if k != "model_name"},
+        }
+        self.client = openai.OpenAI(**client_kwargs)
+        self.async_client = openai.AsyncOpenAI(**client_kwargs)
         self.model_name = model_name
         self.base_url = base_url  # Track for cost extraction
 

rlm/core/lm_handler.py

@@ -85,8 +85,14 @@ def _handle_batched(self, request: LMRequest, handler: "LMHandler") -> LMRespons
 
         start_time = time.perf_counter()
 
+        sem = asyncio.Semaphore(handler.batch_max_concurrent)
+
+        async def run_one(prompt: str):
+            async with sem:
+                return await client.acompletion(prompt)
+
         async def run_all():
-            tasks = [client.acompletion(prompt) for prompt in request.prompts]
+            tasks = [run_one(prompt) for prompt in request.prompts]
             return await asyncio.gather(*tasks)
 
         results = asyncio.run(run_all())
@@ -132,6 +138,7 @@ def __init__(
         host: str = "127.0.0.1",
         port: int = 0,  # auto-assign available port
         other_backend_client: BaseLM | None = None,
+        batch_max_concurrent: int = 16,
     ):
         self.default_client = client
         self.other_backend_client = other_backend_client
@@ -140,6 +147,7 @@ def __init__(
         self._server: ThreadingLMServer | None = None
         self._thread: Thread | None = None
         self._port = port
+        self.batch_max_concurrent = batch_max_concurrent
 
         self.register_client(client.model_name, client)
 

rlm/core/rlm.py

@@ -70,6 +70,7 @@ def __init__(
         custom_sub_tools: dict[str, Any] | None = None,
         compaction: bool = False,
         compaction_threshold_pct: float = 0.85,
+        max_concurrent_subcalls: int = 4,
         on_subcall_start: Callable[[int, str, str], None] | None = None,
         on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None,
         on_iteration_start: Callable[[int, int], None] | None = None,
@@ -102,6 +103,8 @@ def __init__(
                 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).
+            max_concurrent_subcalls: Maximum number of parallel threads for rlm_query_batched subcalls.
+                Each child RLM runs in its own thread. Default 4.
             on_subcall_start: Callback fired when a child RLM starts. Args: (depth, model, prompt_preview).
             on_subcall_complete: Callback fired when a child RLM completes. Args: (depth, model, duration, error_or_none).
             on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num).
@@ -132,6 +135,7 @@ def __init__(
 
         self.compaction = compaction
         self.compaction_threshold_pct = compaction_threshold_pct
+        self.max_concurrent_subcalls = max_concurrent_subcalls
 
         self.depth = depth
         self.max_depth = max_depth
@@ -238,6 +242,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:
@@ -765,6 +770,8 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
             # Propagate custom tools to children (sub_tools become the child's tools)
             custom_tools=self.custom_sub_tools,
             custom_sub_tools=self.custom_sub_tools,
+            # Propagate concurrency settings to children
+            max_concurrent_subcalls=self.max_concurrent_subcalls,
             # Propagate callbacks to children for nested tracking
             on_subcall_start=self.on_subcall_start,
             on_subcall_complete=self.on_subcall_complete,

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

@@ -9,6 +9,7 @@
 import time
 import uuid
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from contextlib import contextmanager
 from typing import Any
 
@@ -135,9 +136,15 @@ def __init__(
         custom_tools: dict[str, Any] | None = None,
         custom_sub_tools: dict[str, Any] | None = None,
         compaction: bool = False,
+        max_concurrent_subcalls: int = 4,
         **kwargs,
     ):
-        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)
@@ -316,9 +323,13 @@ 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. When multiple
+        prompts are provided, subcalls run concurrently using a thread pool
+        (bounded by max_concurrent_subcalls) since they are independent and
+        I/O-bound. Results are returned in the same order as input prompts.
 
-        Each prompt gets its own child RLM for deeper thinking.
         Falls back to llm_query_batched if no recursive capability is configured.
 
         Args:
@@ -329,14 +340,47 @@ def _rlm_query_batched(self, prompts: list[str], model: str | None = None) -> li
             List of responses in the same order as input prompts.
         """
         if self.subcall_fn is not None:
-            results = []
-            for prompt in prompts:
+            # For 0 or 1 prompts, no need for thread pool overhead
+            if len(prompts) <= 1:
+                results = []
+                for prompt in prompts:
+                    try:
+                        completion = self.subcall_fn(prompt, model)
+                        self._pending_llm_calls.append(completion)
+                        results.append(completion.response)
+                    except Exception as e:
+                        results.append(f"Error: RLM query failed - {e}")
+                return results
+
+            # Parallel execution for multiple prompts
+            max_workers = min(self.max_concurrent_subcalls, len(prompts))
+            # Pre-allocate result slots to preserve ordering
+            results: list[str] = [""] * len(prompts)
+            completions: list[tuple[int, RLMChatCompletion]] = []
+            lock = threading.Lock()
+
+            def _run_subcall(index: int, prompt: str) -> None:
                 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

tests/mock_lm.py

@@ -1,27 +1,61 @@
+"""Mock LM for testing without a real model."""
+
+from collections.abc import Callable
+from typing import Any
+
 from rlm.clients.base_lm import BaseLM
 from rlm.core.types import ModelUsageSummary, UsageSummary
 
 
 class MockLM(BaseLM):
-    """Simple mock LM that echoes prompts."""
+    """
+    In-memory mock LM that implements BaseLM for testing.
 
-    def __init__(self):
-        super().__init__(model_name="mock-model")
+    - completion() / acompletion(): echo-style response by default, or use
+      optional responses list / response_fn if provided.
+    - get_usage_summary() / get_last_usage(): return usage suitable for tests.
+    """
 
-    def completion(self, prompt):
-        return f"Mock response to: {prompt[:50]}"
+    def __init__(
+        self,
+        model_name: str = "mock-model",
+        responses: list[str] | None = None,
+        response_fn: Callable[[str | dict[str, Any]], str] | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name=model_name, **kwargs)
+        self._responses = list(responses) if responses is not None else None
+        self._response_fn = response_fn
+        self._call_count = 0
 
-    async def acompletion(self, prompt):
+    def completion(self, prompt: str | dict[str, Any]) -> str:
+        self._call_count += 1
+        if self._responses is not None:
+            if not self._responses:
+                raise IndexError("MockLM: no more responses in list")
+            return self._responses.pop(0)
+        if self._response_fn is not None:
+            return self._response_fn(prompt)
+        prompt_str = prompt if isinstance(prompt, str) else str(prompt)[:80]
+        return f"Mock response to: {prompt_str}"
+
+    async def acompletion(self, prompt: str | dict[str, Any]) -> str:
         return self.completion(prompt)
 
-    def get_usage_summary(self):
+    def get_usage_summary(self) -> UsageSummary:
         return UsageSummary(
             model_usage_summaries={
-                "mock-model": ModelUsageSummary(
-                    total_calls=1, total_input_tokens=10, total_output_tokens=10
+                self.model_name: ModelUsageSummary(
+                    total_calls=self._call_count,
+                    total_input_tokens=self._call_count * 10,
+                    total_output_tokens=self._call_count * 10,
                 )
             }
         )
 
-    def get_last_usage(self):
-        return self.get_usage_summary()
+    def get_last_usage(self) -> ModelUsageSummary:
+        return ModelUsageSummary(
+            total_calls=1,
+            total_input_tokens=10,
+            total_output_tokens=10,
+        )

tests/test_lm_handler.py

@@ -0,0 +1,45 @@
+"""Tests for LMHandler using MockLM (no real LM required)."""
+
+from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
+from rlm.core.lm_handler import LMHandler
+from tests.mock_lm import MockLM
+
+
+def test_lm_handler_single_request():
+    """Single prompt request returns success and echo-style content."""
+    mock = MockLM(responses=["hello back"])
+    with LMHandler(client=mock) as handler:
+        request = LMRequest(prompt="hello")
+        response = send_lm_request(handler.address, request)
+    assert response.success
+    assert response.chat_completion is not None
+    assert response.chat_completion.response == "hello back"
+
+
+def test_lm_handler_batched_request():
+    """Batched prompts return one response per prompt in order."""
+    responses = [f"r{i}" for i in range(5)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=3) as handler:
+        prompts = [f"prompt-{i}" for i in range(5)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == 5
+    for i, resp in enumerate(result):
+        assert resp.success, resp.error
+        assert resp.chat_completion is not None
+        assert resp.chat_completion.response == f"r{i}"
+
+
+def test_lm_handler_batched_many_prompts_semaphore_cap():
+    """Many prompts complete successfully with semaphore limiting concurrency."""
+    # 50 prompts, max 4 concurrent: should still all complete
+    count = 50
+    responses = [f"resp-{i}" for i in range(count)]
+    mock = MockLM(responses=responses)
+    with LMHandler(client=mock, batch_max_concurrent=4) as handler:
+        prompts = [f"p-{i}" for i in range(count)]
+        result = send_lm_request_batched(handler.address, prompts)
+    assert len(result) == count
+    for i, resp in enumerate(result):
+        assert resp.success, (i, resp.error)
+        assert resp.chat_completion.response == f"resp-{i}"