Authorize heapq, bisect, operator and cmath; make the env constructor accept them

envs/coding_env/server/app.py

@@ -21,6 +21,7 @@ Usage:
     python -m envs.coding_env.server.app
 """
 
+import functools
 import os
 from contextlib import suppress
 
@@ -30,7 +31,16 @@ from openenv.core.env_server import create_app
 
 # Create the app with web interface and README integration
 # Pass the class (factory) instead of an instance for WebSocket session support
-app = create_app(PythonCodeActEnv, CodeAction, CodeObservation, env_name="coding_env", max_concurrent_envs=512)
+# MBPP reference solutions need these four, and smolagents' default allowlist has none of
+# them: without this, 8 of the 170 sanitized training problems are unsolvable whatever the
+# model writes. `sys` is deliberately left out of a sandbox that runs model-written code.
+app = create_app(
+    functools.partial(PythonCodeActEnv, additional_imports=["heapq", "bisect", "operator", "cmath"]),
+    CodeAction,
+    CodeObservation,
+    env_name="coding_env",
+    max_concurrent_envs=512,
+)
 
 
 def main():

envs/coding_env/server/python_codeact_env.py

@@ -13,7 +13,7 @@ Python code actions using PyExecutor.
 
 import uuid
 
-from openenv.core.env_server.interfaces import Action, Environment, Observation
+from openenv.core.env_server.interfaces import Action, Environment, Observation, Transform
 
 from ..models import CodeAction, CodeObservation, CodeState
 from .python_executor import PyExecutor
@@ -21,8 +21,6 @@ from .transforms import create_safe_coding_transform
 
 
 class PythonCodeActEnv(Environment):
-    SUPPORTS_CONCURRENT_SESSIONS = True
-
     """
     Python Code Action Environment for executing code and tracking state.
 
@@ -45,11 +43,17 @@ class PythonCodeActEnv(Environment):
         >>> print(env.state.last_exit_code)  # 0
     """
 
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
     def __init__(
         self,
+        transform: Transform | None = None,
+        additional_imports: list[str] | None = None,
     ):
-        self.transform = create_safe_coding_transform()
-        self._executor = PyExecutor()
+        self._transform = transform
+        self._additional_imports = additional_imports
+        self.transform = transform or create_safe_coding_transform()
+        self._executor = PyExecutor(additional_imports=additional_imports)
         self._state = CodeState()
 
     def reset(self) -> Observation:
@@ -65,10 +69,10 @@ class PythonCodeActEnv(Environment):
         self._state.last_exit_code = 0
 
         # Reset executor to clear any previously defined variables/functions
-        self._executor = PyExecutor()
+        self._executor = PyExecutor(additional_imports=self._additional_imports)
 
         # Reset transform to clear any accumulated state
-        self.transform = create_safe_coding_transform()
+        self.transform = self._transform or create_safe_coding_transform()
 
         # Return initial observation
         observation = CodeObservation(

server/app.py

@@ -21,6 +21,7 @@ Usage:
     python -m envs.coding_env.server.app
 """
 
+import functools
 import os
 from contextlib import suppress
 
@@ -30,7 +31,16 @@ from openenv.core.env_server import create_app
 
 # Create the app with web interface and README integration
 # Pass the class (factory) instead of an instance for WebSocket session support
-app = create_app(PythonCodeActEnv, CodeAction, CodeObservation, env_name="coding_env", max_concurrent_envs=512)
+# MBPP reference solutions need these four, and smolagents' default allowlist has none of
+# them: without this, 8 of the 170 sanitized training problems are unsolvable whatever the
+# model writes. `sys` is deliberately left out of a sandbox that runs model-written code.
+app = create_app(
+    functools.partial(PythonCodeActEnv, additional_imports=["heapq", "bisect", "operator", "cmath"]),
+    CodeAction,
+    CodeObservation,
+    env_name="coding_env",
+    max_concurrent_envs=512,
+)
 
 
 def main():

server/python_codeact_env.py

@@ -13,7 +13,7 @@ Python code actions using PyExecutor.
 
 import uuid
 
-from openenv.core.env_server.interfaces import Action, Environment, Observation
+from openenv.core.env_server.interfaces import Action, Environment, Observation, Transform
 
 from ..models import CodeAction, CodeObservation, CodeState
 from .python_executor import PyExecutor
@@ -21,8 +21,6 @@ from .transforms import create_safe_coding_transform
 
 
 class PythonCodeActEnv(Environment):
-    SUPPORTS_CONCURRENT_SESSIONS = True
-
     """
     Python Code Action Environment for executing code and tracking state.
 
@@ -45,11 +43,17 @@ class PythonCodeActEnv(Environment):
         >>> print(env.state.last_exit_code)  # 0
     """
 
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
     def __init__(
         self,
+        transform: Transform | None = None,
+        additional_imports: list[str] | None = None,
     ):
-        self.transform = create_safe_coding_transform()
-        self._executor = PyExecutor()
+        self._transform = transform
+        self._additional_imports = additional_imports
+        self.transform = transform or create_safe_coding_transform()
+        self._executor = PyExecutor(additional_imports=additional_imports)
         self._state = CodeState()
 
     def reset(self) -> Observation:
@@ -65,10 +69,10 @@ class PythonCodeActEnv(Environment):
         self._state.last_exit_code = 0
 
         # Reset executor to clear any previously defined variables/functions
-        self._executor = PyExecutor()
+        self._executor = PyExecutor(additional_imports=self._additional_imports)
 
         # Reset transform to clear any accumulated state
-        self.transform = create_safe_coding_transform()
+        self.transform = self._transform or create_safe_coding_transform()
 
         # Return initial observation
         observation = CodeObservation(