From a8c6f9254d0e2ad67d2c44992daf6a840760a18d Mon Sep 17 00:00:00 2001 From: Alex Zhang Date: Sat, 10 Jan 2026 01:19:41 -0500 Subject: [PATCH] add depth component (for depth=1) to allow alternative backends --- rlm/core/comms_utils.py | 14 ++++++++++++-- rlm/core/lm_handler.py | 29 +++++++++++++++++++++++++---- rlm/core/rlm.py | 19 +++++++++++++++++-- rlm/environments/base_env.py | 3 ++- rlm/environments/docker_repl.py | 17 ++++++++++------- rlm/environments/local_repl.py | 9 ++++++--- rlm/environments/modal_repl.py | 17 ++++++++++------- rlm/environments/prime_repl.py | 17 ++++++++++------- 8 files changed, 92 insertions(+), 33 deletions(-) diff --git a/rlm/core/comms_utils.py b/rlm/core/comms_utils.py index 07e92aeff..4dbc4ffc1 100644 --- a/rlm/core/comms_utils.py +++ b/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) diff --git a/rlm/core/lm_handler.py b/rlm/core/lm_handler.py index 36d27b3a1..46d20339b 100644 --- a/rlm/core/lm_handler.py +++ b/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) diff --git a/rlm/core/rlm.py b/rlm/core/rlm.py index 1de297769..5f303928e 100644 --- a/rlm/core/rlm.py +++ b/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: diff --git a/rlm/environments/base_env.py b/rlm/environments/base_env.py index 963018b4e..e7b1be8fd 100644 --- a/rlm/environments/base_env.py +++ b/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 diff --git a/rlm/environments/docker_repl.py b/rlm/environments/docker_repl.py index 197148573..f79d70609 100644 --- a/rlm/environments/docker_repl.py +++ b/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,7 +105,7 @@ 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: @@ -112,7 +113,7 @@ def llm_query(prompt, model=None): 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, diff --git a/rlm/environments/local_repl.py b/rlm/environments/local_repl.py index 05de1d3dc..5aa1b696c 100644 --- a/rlm/environments/local_repl.py +++ b/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: diff --git a/rlm/environments/modal_repl.py b/rlm/environments/modal_repl.py index 2a61be1ee..dc6801871 100644 --- a/rlm/environments/modal_repl.py +++ b/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 diff --git a/rlm/environments/prime_repl.py b/rlm/environments/prime_repl.py index c2b33f885..aad4069a8 100644 --- a/rlm/environments/prime_repl.py +++ b/rlm/environments/prime_repl.py @@ -111,7 +111,7 @@ def respond(): # ============================================================================= -def _build_exec_script(code: str, broker_port: int = 8888) -> str: +def _build_exec_script(code: str, broker_port: int = 8888, depth: int = 1) -> str: """ Build a script that executes code with state persistence. LLM queries go through the local broker server. @@ -144,7 +144,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() @@ -160,7 +160,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,9 +284,10 @@ def __init__( setup_code: str | None = None, network_access: bool = True, persistent: bool = False, + depth: int = 1, **kwargs: Any, ): - super().__init__(persistent=persistent, **kwargs) + super().__init__(persistent=persistent, depth=depth, **kwargs) if persistent: raise NotImplementedError( @@ -454,7 +455,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: @@ -468,7 +469,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: @@ -504,7 +507,7 @@ def execute_code(self, code: str) -> REPLResult: self.pending_llm_calls.clear() # Build and write the script - script = _build_exec_script(code, self.BROKER_PORT) + script = _build_exec_script(code, self.BROKER_PORT, self.depth) script_b64 = base64.b64encode(script.encode()).decode() self.client.execute_command( self.sandbox_id,