Merge pull request alexzhang13#152 from alexzhang13/ipython_repl

Add iPython REPL environment support

Author: alexzhang13

pyproject.toml

@@ -39,6 +39,7 @@ modal = ["modal>=0.73.0", "dill>=0.3.7"]
 e2b = ["e2b-code-interpreter>=0.0.11", "dill>=0.3.7"]
 daytona = ["daytona>=0.128.1", "dill>=0.3.7"]
 prime = ["prime-sandboxes>=0.2.0", "dill>=0.3.7"]
+ipython = ["ipython>=8.0.0", "jupyter_client>=8.0.0", "ipykernel>=6.0.0", "dill>=0.3.7"]
 
 [build-system]
 requires = ["setuptools>=61.0"]

rlm/core/rlm.py

@@ -232,8 +232,8 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
             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
-            # For local environment with max_depth > 1, pass subcall callback for recursive RLM calls
-            if self.environment_type == "local" and self.max_depth > 1:
+            # For local/ipython environments with max_depth > 1, pass subcall callback for recursive RLM calls
+            if self.environment_type in ("local", "ipython") and self.max_depth > 1:
                 env_kwargs["subcall_fn"] = self._subcall
             # Pass custom tools to the environment
             if self.custom_tools is not None:
@@ -828,7 +828,7 @@ def _validate_persistent_environment_support(self) -> None:
             ValueError: If the environment type does not support persistent mode.
         """
         # Known environments that support persistence
-        persistent_supported_environments = {"local"}
+        persistent_supported_environments = {"local", "ipython"}
 
         if self.environment_type not in persistent_supported_environments:
             raise ValueError(

rlm/core/types.py

@@ -12,7 +12,7 @@
     "azure_openai",
     "gemini",
 ]
-EnvironmentType = Literal["local", "docker", "modal", "prime", "daytona", "e2b"]
+EnvironmentType = Literal["local", "ipython", "docker", "modal", "prime", "daytona", "e2b"]
 
 
 def _serialize_value(value: Any) -> Any:

rlm/environments/__init__.py

@@ -14,8 +14,20 @@
 )
 from rlm.environments.local_repl import LocalREPL
 
+
+def __getattr__(name: str) -> Any:
+    # Lazy-load environments with optional dependencies so that the package
+    # imports cleanly even when those extras aren't installed.
+    if name == "IPythonREPL":
+        from rlm.environments.ipython_repl import IPythonREPL
+
+        return IPythonREPL
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "BaseEnv",
+    "IPythonREPL",
     "LocalREPL",
     "RESERVED_TOOL_NAMES",
     "SupportsCustomTools",
@@ -31,15 +43,19 @@
 
 
 def get_environment(
-    environment: Literal["local", "modal", "docker", "daytona", "prime", "e2b"],
+    environment: Literal["local", "ipython", "modal", "docker", "daytona", "prime", "e2b"],
     environment_kwargs: dict[str, Any],
 ) -> BaseEnv:
     """
     Routes a specific environment and the args (as a dict) to the appropriate environment if supported.
-    Currently supported environments: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b']
+    Currently supported environments: ['local', 'ipython', 'modal', 'docker', 'daytona', 'prime', 'e2b']
     """
     if environment == "local":
         return LocalREPL(**environment_kwargs)
+    elif environment == "ipython":
+        from rlm.environments.ipython_repl import IPythonREPL
+
+        return IPythonREPL(**environment_kwargs)
     elif environment == "modal":
         from rlm.environments.modal_repl import ModalREPL
 
@@ -62,5 +78,5 @@ def get_environment(
         return E2BREPL(**environment_kwargs)
     else:
         raise ValueError(
-            f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b']"
+            f"Unknown environment: {environment}. Supported: ['local', 'ipython', 'modal', 'docker', 'daytona', 'prime', 'e2b']"
         )