Add recursive RLM client factory injection

Author: vishaltandale00

README.md

@@ -129,6 +129,23 @@ export PRIME_API_KEY=...
 ### Model Providers
 We currently support most major clients (OpenAI, Anthropic), as well as the router platforms (OpenRouter, Portkey). For local models, we recommend using vLLM (which interfaces with the [OpenAI client](https://github.com/alexzhang13/rlm/blob/main/rlm/clients/openai.py)). To view or add support for more clients, start by looking at [`rlm/clients/`](https://github.com/alexzhang13/rlm/tree/main/rlm/clients).
 
+Applications can also inject a client factory without modifying the global backend
+registry. The factory receives the backend identifier and backend arguments, returns a
+[`BaseLM`](https://github.com/alexzhang13/rlm/blob/main/rlm/clients/base_lm.py), and is
+used for the root model, recursive children, explicit model overrides, and max-depth
+fallback calls:
+
+```python
+from rlm import RLM
+
+rlm = RLM(
+    backend="my-runtime",
+    backend_kwargs={"model_name": "orchestrator"},
+    client_factory=lambda backend, kwargs: MyClient(**kwargs),
+    max_depth=3,
+)
+```
+
 ## Training
 We provide a simple RL training harness for training RLMs used in this repo (specifically the `local` REPL). The implementation uses no sandboxes for simplicity and slots easily your use case, but an ideal setup would use sandboxes for safety. Training logic is isolated to the [`training/`](https://github.com/alexzhang13/rlm/tree/main/training) folder, which exposes `rlm.RLM` as a [`verifiers`](https://github.com/willccbb/verifiers) `Environment` and plugs straight into [`prime-rl`](https://github.com/PrimeIntellect-ai/prime-rl). See the [training README](https://github.com/alexzhang13/rlm/tree/main/training#readme) for the launch command. The harness uses subprocess-isolated local REPL execution (no cloud sandboxes), matching the `local` environment above.
 

rlm/clients/__init__.py

@@ -1,15 +1,18 @@
+from collections.abc import Callable
 from typing import Any
 
 from dotenv import load_dotenv
 
 from rlm.clients.base_lm import BaseLM
 from rlm.core.types import ClientBackend
 
+ClientFactory = Callable[[str, dict[str, Any]], BaseLM]
+
 load_dotenv()
 
 
 def get_client(
-    backend: ClientBackend,
+    backend: ClientBackend | str,
     backend_kwargs: dict[str, Any],
 ) -> BaseLM:
     """

rlm/core/rlm.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from typing import Any
 
-from rlm.clients import BaseLM, get_client
+from rlm.clients import BaseLM, ClientFactory, get_client
 from rlm.core.lm_handler import LMHandler
 from rlm.core.types import (
     ClientBackend,
@@ -48,7 +48,7 @@ class RLM:
 
     def __init__(
         self,
-        backend: ClientBackend = "openai",
+        backend: ClientBackend | str = "openai",
         backend_kwargs: dict[str, Any] | None = None,
         environment: EnvironmentType = "local",
         environment_kwargs: dict[str, Any] | None = None,
@@ -60,7 +60,7 @@ def __init__(
         max_tokens: int | None = None,
         max_errors: int | None = None,
         custom_system_prompt: str | None = None,
-        other_backends: list[ClientBackend] | None = None,
+        other_backends: list[ClientBackend | str] | None = None,
         other_backend_kwargs: list[dict[str, Any]] | None = None,
         logger: RLMLogger | None = None,
         verbose: bool = False,
@@ -78,10 +78,12 @@ def __init__(
         sub_sampling_args: dict[str, Any] | None = None,
         orchestrator: bool = True,
         user_prologue: str | None = None,
+        client_factory: ClientFactory | None = None,
     ):
         """
         Args:
-            backend: The backend to use for the RLM.
+            backend: The backend identifier to use for the RLM. Custom identifiers are
+                allowed when ``client_factory`` is supplied.
             backend_kwargs: The kwargs to pass to the backend.
             environment: The environment to use for the RLM.
             environment_kwargs: The kwargs to pass to the environment.
@@ -112,6 +114,9 @@ def __init__(
             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).
             on_iteration_complete: Callback fired when an iteration completes. Args: (depth, iteration_num, duration).
+            client_factory: Optional backend client constructor. It receives the selected
+                backend and backend kwargs and must return a ``BaseLM``. The same factory
+                is propagated to recursive children and max-depth fallback calls.
         """
         # Sampling args plumbed into backend_kwargs / other_backend_kwargs
         # before the clients are constructed, so they reach the chat-completions
@@ -154,6 +159,7 @@ def __init__(
 
         self.other_backends = other_backends
         self.other_backend_kwargs = other_backend_kwargs
+        self.client_factory = client_factory or get_client
 
         # Custom tools: functions available in the REPL environment
         self.custom_tools = custom_tools
@@ -231,12 +237,14 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
         When persistent=False (default), creates fresh environment each call.
         """
         # Create client and wrap in handler
-        client: BaseLM = get_client(self.backend, self.backend_kwargs)
+        client: BaseLM = self.client_factory(self.backend, self.backend_kwargs)
 
         # Create other_backend_client if provided (for depth=1 routing)
         other_backend_client: BaseLM | None = None
         if self.other_backends and self.other_backend_kwargs:
-            other_backend_client = get_client(self.other_backends[0], self.other_backend_kwargs[0])
+            other_backend_client = self.client_factory(
+                self.other_backends[0], self.other_backend_kwargs[0]
+            )
 
         lm_handler = LMHandler(client, other_backend_client=other_backend_client)
 
@@ -250,7 +258,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
                 self.other_backend_kwargs[1:],
                 strict=True,
             ):
-                other_client: BaseLM = get_client(backend, kwargs)
+                other_client: BaseLM = self.client_factory(backend, kwargs)
                 lm_handler.register_client(other_client.model_name, other_client)
 
         lm_handler.start()
@@ -699,7 +707,7 @@ def _fallback_answer(self, message: str | dict[str, Any]) -> str:
         """
         Fallback behavior if the RLM is actually at max depth, and should be treated as an LM.
         """
-        client: BaseLM = get_client(self.backend, self.backend_kwargs)
+        client: BaseLM = self.client_factory(self.backend, self.backend_kwargs)
         response = client.completion(message)
         return response
 
@@ -734,9 +742,9 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
         if next_depth >= self.max_depth:
             # Use other_backend if available, otherwise use main backend
             if self.other_backends and self.other_backend_kwargs:
-                client = get_client(self.other_backends[0], self.other_backend_kwargs[0])
+                client = self.client_factory(self.other_backends[0], self.other_backend_kwargs[0])
             else:
-                client = get_client(self.backend, child_backend_kwargs or {})
+                client = self.client_factory(self.backend, child_backend_kwargs or {})
             root_model = model or client.model_name
             start_time = time.perf_counter()
             try:
@@ -831,6 +839,7 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion:
             # Propagate callbacks to children for nested tracking
             on_subcall_start=self.on_subcall_start,
             on_subcall_complete=self.on_subcall_complete,
+            client_factory=self.client_factory,
         )
         try:
             result = child.completion(prompt, root_prompt=None)

tests/test_subcall.py

@@ -13,6 +13,7 @@
 import rlm.core.rlm as rlm_module
 from rlm import RLM
 from rlm.core.types import ModelUsageSummary, UsageSummary
+from tests.mock_lm import MockLM
 
 
 def create_mock_lm(responses: list[str], model_name: str = "mock-model") -> Mock:
@@ -36,6 +37,109 @@ def final(content: str) -> str:
     return f"```repl\nanswer['content'] = {content!r}\nanswer['ready'] = True\n```"
 
 
+class TestClientFactoryInjection:
+    def test_factory_builds_root_client(self):
+        calls = []
+
+        def client_factory(backend, backend_kwargs):
+            calls.append((backend, dict(backend_kwargs)))
+            return MockLM(
+                model_name=backend_kwargs["model_name"],
+                responses=[final("root answer")],
+            )
+
+        rlm = RLM(
+            backend="custom-runtime",
+            backend_kwargs={"model_name": "root-model"},
+            client_factory=client_factory,
+        )
+
+        result = rlm.completion("root prompt")
+
+        assert result.response == "root answer"
+        assert calls == [("custom-runtime", {"model_name": "root-model"})]
+
+    def test_factory_propagates_to_recursive_child_with_model_override(self):
+        calls = []
+
+        def client_factory(backend, backend_kwargs):
+            calls.append((backend, dict(backend_kwargs)))
+            return MockLM(
+                model_name=backend_kwargs["model_name"],
+                responses=[final("child answer")],
+            )
+
+        parent = RLM(
+            backend="openai",
+            backend_kwargs={"model_name": "root-model"},
+            client_factory=client_factory,
+            max_depth=3,
+        )
+
+        result = parent._subcall("review this", model="review-model")
+
+        assert result.response == "child answer"
+        assert calls == [("openai", {"model_name": "review-model"})]
+
+    def test_factory_is_used_by_max_depth_plain_lm_fallback(self):
+        calls = []
+
+        def client_factory(backend, backend_kwargs):
+            calls.append((backend, dict(backend_kwargs)))
+            return MockLM(
+                model_name=backend_kwargs["model_name"],
+                responses=["plain answer"],
+            )
+
+        parent = RLM(
+            backend="openai",
+            backend_kwargs={"model_name": "root-model"},
+            client_factory=client_factory,
+            max_depth=1,
+        )
+
+        result = parent._subcall("review this", model="review-model")
+
+        assert result.response == "plain answer"
+        assert result.root_model == "review-model"
+        assert calls == [("openai", {"model_name": "review-model"})]
+
+    def test_rlm_query_uses_factory_for_a_distinct_recursive_child(self):
+        calls = []
+
+        def client_factory(backend, backend_kwargs):
+            model = backend_kwargs["model_name"]
+            calls.append((backend, dict(backend_kwargs)))
+            if model == "root-model":
+                return MockLM(
+                    model_name=model,
+                    responses=[
+                        "```repl\nreview = rlm_query('review this', model='review-model')\nprint(review)\n```",
+                        final("root accepted child review"),
+                    ],
+                )
+            return MockLM(
+                model_name=model,
+                responses=[final("child review complete")],
+            )
+
+        rlm = RLM(
+            backend="openai",
+            backend_kwargs={"model_name": "root-model"},
+            client_factory=client_factory,
+            max_depth=3,
+            max_iterations=3,
+        )
+
+        result = rlm.completion("root prompt")
+
+        assert result.response == "root accepted child review"
+        assert calls == [
+            ("openai", {"model_name": "root-model"}),
+            ("openai", {"model_name": "review-model"}),
+        ]
+
+
 class TestSubcallTimeoutPropagation:
     """Tests for max_timeout propagation to child RLM."""