diff --git a/README.md b/README.md index e4a8f33ce..536a332ca 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,27 @@ 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, +) +``` + +When `other_backends` is configured, its first client is the default for an +`rlm_query(prompt)` recursive child as well as a max-depth plain sub-call. Passing +`model=` to `rlm_query` explicitly overrides that default on the parent backend. + ## 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. diff --git a/rlm/clients/__init__.py b/rlm/clients/__init__.py index a57c8e8df..ed9fedb60 100644 --- a/rlm/clients/__init__.py +++ b/rlm/clients/__init__.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from typing import Any from dotenv import load_dotenv @@ -5,11 +6,13 @@ 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: """ diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 43a5819d3..f0b3ff596 100644 --- a/rlm/core/rlm.py +++ b/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 @@ -722,21 +730,24 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: """ next_depth = self.depth + 1 - # Determine which backend/kwargs to use (model override or parent's default) + # Determine which backend/kwargs to use. An explicit model stays on the + # parent's backend. Without one, the configured recursive child backend is + # the default for both child RLMs and max-depth plain calls. if model is not None: + child_backend = self.backend child_backend_kwargs = (self.backend_kwargs or {}).copy() child_backend_kwargs["model_name"] = model + elif self.other_backends and self.other_backend_kwargs: + child_backend = self.other_backends[0] + child_backend_kwargs = self.other_backend_kwargs[0] else: + child_backend = self.backend child_backend_kwargs = self.backend_kwargs resolved_model = model or (child_backend_kwargs or {}).get("model_name", "unknown") # If we'd hit/exceed the cap, do a normal LM completion (no REPL) 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]) - else: - client = get_client(self.backend, child_backend_kwargs or {}) + client = self.client_factory(child_backend, child_backend_kwargs or {}) root_model = model or client.model_name start_time = time.perf_counter() try: @@ -806,7 +817,7 @@ def _subcall(self, prompt: str, model: str | None = None) -> RLMChatCompletion: # Spawn a child RLM with its own LocalREPL child = RLM( - backend=self.backend, + backend=child_backend, backend_kwargs=child_backend_kwargs, environment=self.environment_type, environment_kwargs=self.environment_kwargs, @@ -831,6 +842,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) diff --git a/tests/test_subcall.py b/tests/test_subcall.py index f6c10a28e..cb92abb93 100644 --- a/tests/test_subcall.py +++ b/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,133 @@ 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_recursive_child_defaults_to_configured_other_backend(self): + calls = [] + + def client_factory(backend, backend_kwargs): + calls.append((backend, dict(backend_kwargs))) + return MockLM( + model_name=backend_kwargs["model_name"], + responses=[final("independent child answer")], + ) + + parent = RLM( + backend="owner-runtime", + backend_kwargs={"model_name": "content-owner"}, + other_backends=["review-runtime"], + other_backend_kwargs=[{"model_name": "self-assess"}], + client_factory=client_factory, + max_depth=3, + ) + + result = parent._subcall("review this") + + assert result.response == "independent child answer" + assert calls[0] == ("review-runtime", {"model_name": "self-assess"}) + + 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."""