openenv/repl

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Controller composition for recursive backends.9 10This keeps server-side recursion assembly outside `REPLEnvironment`:11- backend selection based on max_depth12- limits configuration13- uniform callable interface for the environment14"""15 16from __future__ import annotations17 18from dataclasses import dataclass19from typing import Callable20 21from .prompts import RLM_SYSTEM_PROMPT22from .recursive_backends import BackendLimits, DirectLMBackend, LocalChildRLMBackend23 24 25ChatFn = Callable[..., str]26 27 28@dataclass29class RecursiveController:30    llm_query_fn: Callable[[str, str | None], str]31    llm_batch_fn: Callable[[list[str], str | None], list[str]]32    rlm_query_fn: Callable[[str, str | None], str] | None33    rlm_batch_fn: Callable[[list[str], str | None], list[str]] | None34    backend: object35 36    def close(self) -> None:37        pass38 39 40def create_server_recursive_controller(41    chat_fn: ChatFn,42    *,43    max_depth: int,44    max_iterations: int,45    system_prompt: str = RLM_SYSTEM_PROMPT,46    max_batch_workers: int = 8,47    max_children_total: int | None = None,48    max_children_per_batch: int | None = None,49    result_truncation_limit: int | None = None,50    per_child_timeout_s: float | None = None,51    env_max_iterations_multiplier: int = 5,52) -> RecursiveController:53    limits = BackendLimits(54        max_depth=max_depth,55        max_batch_workers=max_batch_workers,56        max_children_total=max_children_total,57        max_children_per_batch=max_children_per_batch,58        result_truncation_limit=result_truncation_limit,59        per_child_timeout_s=per_child_timeout_s,60    )61    if max_depth > 1:62        from .runner import LocalRLMRunner63 64        backend = LocalChildRLMBackend(65            chat_fn,66            runner_factory=LocalRLMRunner,67            system_prompt=system_prompt,68            max_iterations=max_iterations,69            env_max_iterations_multiplier=env_max_iterations_multiplier,70            depth=0,71            limits=limits,72        )73        return RecursiveController(74            llm_query_fn=backend.query,75            llm_batch_fn=backend.query_batched,76            rlm_query_fn=backend.recursive_query,77            rlm_batch_fn=backend.recursive_query_batched,78            backend=backend,79        )80 81    backend = DirectLMBackend(chat_fn, depth=0, limits=limits)82    return RecursiveController(83        llm_query_fn=backend.query,84        llm_batch_fn=backend.query_batched,85        rlm_query_fn=None,86        rlm_batch_fn=None,87        backend=backend,88    )89