add depth component (for depth=1) to allow alternative backends (#44)

Author: alexzhang13

rlm/core/comms_utils.py

@@ -28,6 +28,7 @@ class LMRequest:
     prompt: str | dict[str, Any] | None = None
     prompts: list[str | dict[str, Any]] | None = None
     model: str | None = None
+    depth: int = 0
 
     @property
     def is_batched(self) -> bool:
@@ -43,6 +44,7 @@ def to_dict(self) -> dict:
             d["prompts"] = self.prompts
         if self.model is not None:
             d["model"] = self.model
+        d["depth"] = self.depth
         return d
 
     @classmethod
@@ -52,6 +54,7 @@ def from_dict(cls, data: dict) -> "LMRequest":
             prompt=data.get("prompt"),
             prompts=data.get("prompts"),
             model=data.get("model"),
+            depth=data.get("depth", -1),  # TODO: Default should throw an error
         )
 
 
@@ -198,18 +201,23 @@ def socket_request(address: tuple[str, int], data: dict, timeout: int = 300) ->
 # =============================================================================
 
 
-def send_lm_request(address: tuple[str, int], request: LMRequest, timeout: int = 300) -> LMResponse:
+def send_lm_request(
+    address: tuple[str, int], request: LMRequest, timeout: int = 300, depth: int | None = None
+) -> LMResponse:
     """Send an LM request and return typed response.
 
     Args:
         address: (host, port) tuple of LM Handler server.
         request: LMRequest to send.
         timeout: Socket timeout in seconds.
+        depth: Optional depth to override request depth.
 
     Returns:
         LMResponse with content or error.
     """
     try:
+        if depth is not None:
+            request.depth = depth
         response_data = socket_request(address, request.to_dict(), timeout)
         return LMResponse.from_dict(response_data)
     except Exception as e:
@@ -221,6 +229,7 @@ def send_lm_request_batched(
     prompts: list[str | dict[str, Any]],
     model: str | None = None,
     timeout: int = 300,
+    depth: int = 0,
 ) -> list[LMResponse]:
     """Send a batched LM request and return a list of typed responses.
 
@@ -229,12 +238,13 @@ def send_lm_request_batched(
         prompts: List of prompts to send.
         model: Optional model name to use.
         timeout: Socket timeout in seconds.
+        depth: Depth for routing (default 0).
 
     Returns:
         List of LMResponse objects, one per prompt, in the same order.
     """
     try:
-        request = LMRequest(prompts=prompts, model=model)
+        request = LMRequest(prompts=prompts, model=model, depth=depth)
         response_data = socket_request(address, request.to_dict(), timeout)
         response = LMResponse.from_dict(response_data)
 

rlm/core/lm_handler.py

@@ -45,7 +45,7 @@ def handle(self):
 
     def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse:
         """Handle a single prompt request."""
-        client = handler.get_client(request.model)
+        client = handler.get_client(request.model, request.depth)
 
         start_time = time.perf_counter()
         content = client.completion(request.prompt)
@@ -64,7 +64,7 @@ def _handle_single(self, request: LMRequest, handler: "LMHandler") -> LMResponse
 
     def _handle_batched(self, request: LMRequest, handler: "LMHandler") -> LMResponse:
         """Handle a batched prompts request using async for concurrency."""
-        client = handler.get_client(request.model)
+        client = handler.get_client(request.model, request.depth)
 
         start_time = time.perf_counter()
 
@@ -112,8 +112,10 @@ def __init__(
         client: BaseLM,
         host: str = "127.0.0.1",
         port: int = 0,  # auto-assign available port
+        other_backend_client: BaseLM | None = None,
     ):
         self.default_client = client
+        self.other_backend_client = other_backend_client
         self.clients: dict[str, BaseLM] = {}
         self.host = host
         self._server: ThreadingLMServer | None = None
@@ -126,10 +128,21 @@ def register_client(self, model_name: str, client: BaseLM) -> None:
         """Register a client for a specific model name."""
         self.clients[model_name] = client
 
-    def get_client(self, model: str | None = None) -> BaseLM:
-        """Get client by model name, or return default."""
+    def get_client(self, model: str | None = None, depth: int = 0) -> BaseLM:
+        """Get client by model name or depth, or return default.
+
+        Routing logic:
+        - depth=0: use default_client (main backend)
+        - depth=1: use other_backend_client if it exists, otherwise default_client
+        - If model is specified and exists in clients, use that (overrides depth routing)
+        """
         if model and model in self.clients:
             return self.clients[model]
+
+        # Route based on depth
+        if depth == 1 and self.other_backend_client is not None:
+            return self.other_backend_client
+
         return self.default_client
 
     @property
@@ -179,6 +192,14 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     def get_usage_summary(self) -> UsageSummary:
         """Get the usage summary for all clients, merged into a single dict."""
         merged = {}
+        # Include default client
+        default_summary = self.default_client.get_usage_summary()
+        merged.update(default_summary.model_usage_summaries)
+        # Include other backend client if it exists
+        if self.other_backend_client is not None:
+            other_summary = self.other_backend_client.get_usage_summary()
+            merged.update(other_summary.model_usage_summaries)
+        # Include all registered clients
         for client in self.clients.values():
             client_summary = client.get_usage_summary()
             merged.update(client_summary.model_usage_summaries)

rlm/core/rlm.py

@@ -76,6 +76,14 @@ def __init__(
         self.environment_kwargs = (
             environment_kwargs.copy() if environment_kwargs is not None else {}
         )
+        # Validate other_backends: currently only support one additional backend
+        if other_backends is not None:
+            if len(other_backends) != 1:
+                raise ValueError(
+                    "We currently only support one additional backend for the recursive sub-calls! "
+                    "This model will be the model used for recursive sub-calls, but this will change in the future"
+                )
+
         self.other_backends = other_backends
         self.other_backend_kwargs = other_backend_kwargs
 
@@ -124,9 +132,15 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
         """
         # Create client and wrap in handler
         client: BaseLM = get_client(self.backend, self.backend_kwargs)
-        lm_handler = LMHandler(client)
 
-        # Register other clients to be available as sub-call options
+        # 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])
+
+        lm_handler = LMHandler(client, other_backend_client=other_backend_client)
+
+        # Register other clients to be available as sub-call options (by model name)
         if self.other_backends and self.other_backend_kwargs:
             for backend, kwargs in zip(self.other_backends, self.other_backend_kwargs, strict=True):
                 other_client: BaseLM = get_client(backend, kwargs)
@@ -150,6 +164,7 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
             env_kwargs = self.environment_kwargs.copy()
             env_kwargs["lm_handler_address"] = (lm_handler.host, lm_handler.port)
             env_kwargs["context_payload"] = prompt
+            env_kwargs["depth"] = self.depth + 1  # Environment depth is RLM depth + 1
             environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
 
             if self.persistent:

rlm/environments/base_env.py

@@ -10,8 +10,9 @@ class BaseEnv(ABC):
     where isolated environments are on a separate machine from the LM.
     """
 
-    def __init__(self, persistent: bool = False, **kwargs):
+    def __init__(self, persistent: bool = False, depth: int = 1, **kwargs):
         self.persistent = persistent
+        self.depth = depth
         self.kwargs = kwargs
 
     @abstractmethod

rlm/environments/docker_repl.py

@@ -28,6 +28,7 @@ class LLMProxyHandler(BaseHTTPRequestHandler):
     lm_handler_address: tuple[str, int] | None = None
     pending_calls: list[RLMChatCompletion] = []
     lock: threading.Lock = threading.Lock()
+    depth: int = 1
 
     def log_message(self, *args):
         pass
@@ -55,7 +56,7 @@ def _handle_single(self, body: dict) -> dict:
         if not self.lm_handler_address:
             return {"error": "No LM handler configured"}
 
-        request = LMRequest(prompt=body.get("prompt"), model=body.get("model"))
+        request = LMRequest(prompt=body.get("prompt"), model=body.get("model"), depth=self.depth)
         response = send_lm_request(self.lm_handler_address, request)
 
         if not response.success:
@@ -72,7 +73,7 @@ def _handle_batched(self, body: dict) -> dict:
 
         prompts = body.get("prompts", [])
         responses = send_lm_request_batched(
-            self.lm_handler_address, prompts, model=body.get("model")
+            self.lm_handler_address, prompts, model=body.get("model"), depth=self.depth
         )
 
         results = []
@@ -87,7 +88,7 @@ def _handle_batched(self, body: dict) -> dict:
         return {"responses": results}
 
 
-def _build_exec_script(code: str, proxy_port: int) -> str:
+def _build_exec_script(code: str, proxy_port: int, depth: int = 1) -> str:
     """Build execution script for the container."""
     code_b64 = base64.b64encode(code.encode()).decode()
 
@@ -104,15 +105,15 @@ def _build_exec_script(code: str, proxy_port: int) -> str:
 
 def llm_query(prompt, model=None):
     try:
-        r = requests.post(f"{{PROXY}}/llm_query", json={{"prompt": prompt, "model": model}}, timeout=300)
+        r = requests.post(f"{{PROXY}}/llm_query", json={{"prompt": prompt, "model": model, "depth": {depth}}}, timeout=300)
         d = r.json()
         return d.get("response") or f"Error: {{d.get('error')}}"
     except Exception as e:
         return f"Error: {{e}}"
 
 def llm_query_batched(prompts, model=None):
     try:
-        r = requests.post(f"{{PROXY}}/llm_query_batched", json={{"prompts": prompts, "model": model}}, timeout=300)
+        r = requests.post(f"{{PROXY}}/llm_query_batched", json={{"prompts": prompts, "model": model, "depth": {depth}}}, timeout=300)
         d = r.json()
         return d.get("responses") or [f"Error: {{d.get('error')}}"] * len(prompts)
     except Exception as e:
@@ -181,13 +182,14 @@ def __init__(
         context_payload: dict | list | str | None = None,
         setup_code: str | None = None,
         persistent: bool = False,
+        depth: int = 1,
         **kwargs,
     ):
         if persistent:
             raise NotImplementedError(
                 "Persistent REPLs are currently not supported for environment: DockerREPL"
             )
-        super().__init__(persistent=persistent, **kwargs)
+        super().__init__(persistent=persistent, depth=depth, **kwargs)
 
         self.image = image
         self.lm_handler_address = lm_handler_address
@@ -216,6 +218,7 @@ def setup(self):
                 "lm_handler_address": self.lm_handler_address,
                 "pending_calls": self.pending_calls,
                 "lock": self._calls_lock,
+                "depth": self.depth,
             },
         )
         self.proxy_server = HTTPServer(("127.0.0.1", 0), handler)
@@ -266,7 +269,7 @@ def execute_code(self, code: str) -> REPLResult:
         with self._calls_lock:
             self.pending_calls.clear()
 
-        script = _build_exec_script(code, self.proxy_port)
+        script = _build_exec_script(code, self.proxy_port, self.depth)
         result = subprocess.run(
             ["docker", "exec", self.container_id, "python", "-c", script],
             capture_output=True,

rlm/environments/local_repl.py

@@ -124,9 +124,10 @@ def __init__(
         context_payload: dict | list | str | None = None,
         setup_code: str | None = None,
         persistent: bool = False,
+        depth: int = 1,
         **kwargs,
     ):
-        super().__init__(persistent=persistent, **kwargs)
+        super().__init__(persistent=persistent, depth=depth, **kwargs)
 
         self.lm_handler_address = lm_handler_address
         self.original_cwd = os.getcwd()
@@ -181,7 +182,7 @@ def _llm_query(self, prompt: str, model: str | None = None) -> str:
             return "Error: No LM handler configured"
 
         try:
-            request = LMRequest(prompt=prompt, model=model)
+            request = LMRequest(prompt=prompt, model=model, depth=self.depth)
             response = send_lm_request(self.lm_handler_address, request)
 
             if not response.success:
@@ -210,7 +211,9 @@ def _llm_query_batched(self, prompts: list[str], model: str | None = None) -> li
             return ["Error: No LM handler configured"] * len(prompts)
 
         try:
-            responses = send_lm_request_batched(self.lm_handler_address, prompts, model=model)
+            responses = send_lm_request_batched(
+                self.lm_handler_address, prompts, model=model, depth=self.depth
+            )
 
             results = []
             for response in responses:

rlm/environments/modal_repl.py

@@ -112,7 +112,7 @@ def respond():
 # =============================================================================
 
 
-def _build_exec_script(code: str, broker_port: int = 8080) -> str:
+def _build_exec_script(code: str, broker_port: int = 8080, depth: int = 1) -> str:
     """
     Build a script that executes code with state persistence.
     LLM queries go through the local broker server.
@@ -145,7 +145,7 @@ def llm_query(prompt, model=None):
     try:
         response = requests.post(
             f"{{BROKER_URL}}/enqueue",
-            json={{"type": "single", "prompt": prompt, "model": model}},
+            json={{"type": "single", "prompt": prompt, "model": model, "depth": {depth}}},
             timeout=300,
         )
         data = response.json()
@@ -161,7 +161,7 @@ def llm_query_batched(prompts, model=None):
     try:
         response = requests.post(
             f"{{BROKER_URL}}/enqueue",
-            json={{"type": "batched", "prompts": prompts, "model": model}},
+            json={{"type": "batched", "prompts": prompts, "model": model, "depth": {depth}}},
             timeout=300,
         )
         data = response.json()
@@ -284,13 +284,14 @@ def __init__(
         context_payload: dict | list | str | None = None,
         setup_code: str | None = None,
         persistent: bool = False,
+        depth: int = 1,
         **kwargs,
     ):
         if persistent:
             raise NotImplementedError(
                 "Persistent REPLs are currently not supported for environment: ModalREPL"
             )
-        super().__init__(persistent=persistent, **kwargs)
+        super().__init__(persistent=persistent, depth=depth, **kwargs)
 
         self.app_name = app_name
         self.timeout = timeout
@@ -387,7 +388,7 @@ def _handle_llm_request(self, req_data: dict) -> dict:
 
         if req_type == "single":
             prompt = req_data.get("prompt")
-            request = LMRequest(prompt=prompt, model=model)
+            request = LMRequest(prompt=prompt, model=model, depth=self.depth)
             response = send_lm_request(self.lm_handler_address, request)
 
             if not response.success:
@@ -401,7 +402,9 @@ def _handle_llm_request(self, req_data: dict) -> dict:
 
         elif req_type == "batched":
             prompts = req_data.get("prompts", [])
-            responses = send_lm_request_batched(self.lm_handler_address, prompts, model=model)
+            responses = send_lm_request_batched(
+                self.lm_handler_address, prompts, model=model, depth=self.depth
+            )
 
             results = []
             for resp in responses:
@@ -437,7 +440,7 @@ def execute_code(self, code: str) -> REPLResult:
             self.pending_llm_calls.clear()
 
         # Build and execute the script
-        script = _build_exec_script(code, self.BROKER_PORT)
+        script = _build_exec_script(code, self.BROKER_PORT, self.depth)
         process = self.sandbox.exec("python", "-c", script)
 
         # Read output