add more permission capabilities to PythonInterpreter (#8296)

* add more permission capabilities to PythonInterpreter

* update tests

* update file access/write tests

* resolve comments

* resolve comments

* update flags and interpreter usage

* resolve comments

Author: arnavsinghvi11

dspy/predict/code_act.py

@@ -1,7 +1,7 @@
 import inspect
+from typing import Callable, Union, Type, Optional
 import logging
 from inspect import Signature
-from typing import Callable, Type, Union
 
 import dspy
 from dspy.adapters.types.tool import Tool
@@ -17,15 +17,15 @@ class CodeAct(ReAct, ProgramOfThought):
     CodeAct is a module that utilizes the Code Interpreter and predefined tools to solve the problem.
     """
 
-    def __init__(self, signature: Union[str, Type[Signature]], tools: list[Callable], max_iters: int = 5):
+    def __init__(self, signature: Union[str, Type[Signature]], tools: list[Callable], max_iters: int = 5, interpreter: Optional[PythonInterpreter] = None):
         """
         Initializes the CodeAct class with the specified model, temperature, and max tokens.
 
         Args:
             signature (Union[str, Type[Signature]]): The signature of the module.
             tools (list[Callable]): The tool callables to be used. CodeAct only accepts functions and not callable objects.
             max_iters (int): The maximum number of iterations to generate the answer.
-        
+            interpreter: PythonInterpreter instance to use. If None, a new one is instantiated.
         Example:
             ```python
             from dspy.predict import CodeAct
@@ -67,8 +67,8 @@ def factorial(n):
         self.codeact = dspy.Predict(codeact_signature)
         self.extractor = dspy.ChainOfThought(extract_signature)
         # It will raises exception when dspy cannot find available deno instance by now.
-        self.interpreter = PythonInterpreter()
-
+        self.interpreter = interpreter or PythonInterpreter()
+    
     def _build_instructions(self, signature, tools):
         instructions = [f"{signature.instructions}\n"] if signature.instructions else []
         inputs = ", ".join([f"`{k}`" for k in signature.input_fields.keys()])

dspy/predict/program_of_thought.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import re
-from typing import Type, Union
+from typing import Type, Union, Optional
 
 import dspy
 from dspy.primitives.program import Module
@@ -27,11 +27,12 @@ class ProgramOfThought(Module):
     ```
     """
 
-    def __init__(self, signature: Union[str, Type[Signature]], max_iters=3):
+    def __init__(self, signature: Union[str, Type[Signature]], max_iters=3, interpreter: Optional[PythonInterpreter] = None):
         """
         Args:
             signature: The signature of the module.
             max_iters: The maximum number of iterations to retry code generation and execution.
+            interpreter: PythonInterpreter instance to use. If None, a new one is instantiated.
         """
         super().__init__()
         self.signature = signature = ensure_signature(signature)
@@ -59,7 +60,7 @@ def __init__(self, signature: Union[str, Type[Signature]], max_iters=3):
             ),
         )
         # It will raises exception when dspy cannot find available deno instance by now.
-        self.interpreter = PythonInterpreter()
+        self.interpreter = interpreter or PythonInterpreter()
 
     def _generate_signature(self, mode):
         signature_dict = dict(self.input_fields)

dspy/primitives/python_interpreter.py

@@ -2,7 +2,8 @@
 import os
 import subprocess
 from types import TracebackType
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
+from os import PathLike
 
 
 class InterpreterError(RuntimeError):
@@ -24,16 +25,99 @@ class PythonInterpreter:
     ```
     """
 
-    def __init__(self, deno_command: Optional[List[str]] = None) -> None:
+    def __init__(
+        self,
+        deno_command: Optional[List[str]] = None,
+        enable_read_paths: Optional[List[Union[PathLike, str]]] = None,
+        enable_write_paths: Optional[List[Union[PathLike, str]]] = None,
+        enable_env_vars: Optional[List[str]] = None,
+        enable_network_access: Optional[List[str]] = None,
+        sync_files: bool = True,
+    ) -> None:
+        """
+        Args:
+            deno_command: command list to launch Deno.
+            enable_read_paths: Files or directories to allow reading from in the sandbox.
+            enable_write_paths: Files or directories to allow writing to in the sandbox.
+            enable_env_vars: Environment variable names to allow in the sandbox.
+            enable_network_access: Domains or IPs to allow network access in the sandbox.
+            sync_files: If set, syncs changes within the sandbox back to original files after execution.
+        """
         if isinstance(deno_command, dict):
             deno_command = None  # no-op, just a guard in case someone passes a dict
-        self.deno_command = deno_command or ["deno", "run", "--allow-read", self._get_runner_path()]
+
+        self.enable_read_paths = enable_read_paths or []
+        self.enable_write_paths = enable_write_paths or []
+        self.enable_env_vars = enable_env_vars or []
+        self.enable_network_access = enable_network_access or []
+        self.sync_files = sync_files
+        # TODO later on add enable_run (--allow-run) by proxying subprocess.run through Deno.run() to fix 'emscripten does not support processes' error
+
+        if deno_command:
+            self.deno_command = list(deno_command)
+        else:
+            args = ['deno', 'run', '--allow-read']
+            self._env_arg  = ""
+            if self.enable_env_vars:
+                user_vars = [str(v).strip() for v in self.enable_env_vars]
+                args.append("--allow-env=" + ",".join(user_vars))
+                self._env_arg = ",".join(user_vars)
+            if self.enable_network_access:
+                args.append(f"--allow-net={','.join(str(x) for x in self.enable_network_access)}")
+            if self.enable_write_paths:
+                args.append(f"--allow-write={','.join(str(x) for x in self.enable_write_paths)}")
+                
+            args.append(self._get_runner_path())
+
+            # For runner.js to load in env vars
+            if self._env_arg:
+                args.append(self._env_arg)
+            self.deno_command = args
+
         self.deno_process = None
+        self._mounted_files = False
 
     def _get_runner_path(self) -> str:
         current_dir = os.path.dirname(os.path.abspath(__file__))
         return os.path.join(current_dir, "runner.js")
 
+    def _mount_files(self):
+        if self._mounted_files:
+            return
+        paths_to_mount = []
+        if self.enable_read_paths:
+            paths_to_mount.extend(self.enable_read_paths)
+        if self.enable_write_paths:
+            paths_to_mount.extend(self.enable_write_paths)
+        if not paths_to_mount:
+            return
+        for path in paths_to_mount:
+            if not path:
+                continue
+            if not os.path.exists(path):
+                if self.enable_write_paths and path in self.enable_write_paths:
+                    open(path, "a").close()
+                else:
+                    raise FileNotFoundError(f"Cannot mount non-existent file: {path}")
+            virtual_path = f"/sandbox/{os.path.basename(path)}"
+            mount_msg = json.dumps({"mount_file": str(path), "virtual_path": virtual_path})
+            self.deno_process.stdin.write(mount_msg + "\n")
+            self.deno_process.stdin.flush()
+        self._mounted_files = True
+
+    def _sync_files(self):
+        if not self.enable_write_paths or not self.sync_files:
+            return
+        for path in self.enable_write_paths:
+            virtual_path = f"/sandbox/{os.path.basename(path)}"
+            sync_msg = json.dumps({
+                "sync_file": virtual_path,
+                "host_file": str(path)
+            })
+            self.deno_process.stdin.write(sync_msg + "\n")
+            self.deno_process.stdin.flush()
+
+
     def _ensure_deno_process(self) -> None:
         if self.deno_process is None or self.deno_process.poll() is not None:
             try:
@@ -43,6 +127,7 @@ def _ensure_deno_process(self) -> None:
                     stdout=subprocess.PIPE,
                     stderr=subprocess.PIPE,
                     text=True,
+                    env=os.environ.copy()
                 )
             except FileNotFoundError as e:
                 install_instructions = (
@@ -87,6 +172,7 @@ def execute(
         variables = variables or {}
         code = self._inject_variables(code, variables)
         self._ensure_deno_process()
+        self._mount_files()
 
         # Send the code as JSON
         input_data = json.dumps({"code": code})
@@ -127,6 +213,7 @@ def execute(
                 raise InterpreterError(f"{error_type}: {result.get('errorArgs') or error_msg}")
 
         # If there's no error or got `FinalAnswer`, return the "output" field
+        self._sync_files()
         return result.get("output", None)
 
     def __enter__(self):

dspy/primitives/runner.js

@@ -5,6 +5,21 @@ import { readLines } from "https://deno.land/[email protected]/io/mod.ts";
 
 const pyodide = await pyodideModule.loadPyodide();
 
+try {
+  const env_vars = (Deno.args[0] ?? "").split(",").filter(Boolean);
+  for (const key of env_vars) {
+    const val = Deno.env.get(key);
+    if (val !== undefined) {
+      pyodide.runPython(`
+import os
+os.environ[${JSON.stringify(key)}] = ${JSON.stringify(val)}
+      `);
+    }
+  }
+} catch (e) {
+  console.error("Error setting environment variables in Pyodide:", e);
+}
+
 for await (const line of readLines(Deno.stdin)) {
   let input;
   try {
@@ -17,6 +32,43 @@ for await (const line of readLines(Deno.stdin)) {
     continue;
   }
 
+  if (input.mount_file) {
+      const hostPath = input.mount_file;
+      const virtualPath = input.virtual_path || hostPath;
+      try {
+          const contents = await Deno.readFile(hostPath);
+          const dirs = virtualPath.split('/').slice(1, -1);
+          let cur = '';
+          for (const d of dirs) {
+              cur += '/' + d;
+              try {
+                  pyodide.FS.mkdir(cur);
+              } catch (e) {
+                  if (!(e && e.message && e.message.includes('File exists'))) {
+                      console.log("[DEBUG] Error creating directory in Pyodide file system:", cur, "|", e.message);
+                  }
+              }
+          }
+          pyodide.FS.writeFile(virtualPath, contents);
+      } catch (e) {
+          console.log(JSON.stringify({error: "Failed to mount file: " + e.message}));
+      }
+      continue;      
+  }
+
+  if (input.sync_file) {
+      const virtualPath = input.sync_file;
+      const hostPath = input.host_file || virtualPath;
+      try {
+          const contents = pyodide.FS.readFile(virtualPath);
+          await Deno.writeFile(hostPath, contents);
+      } catch (e) {
+          console.log("[DEBUG] Failed to sync file:", hostPath, "|", e.message);
+      }
+      continue;
+  }
+
+
   // Expecting an object like { "code": "...", ... }
   if (typeof input !== 'object' || input === null) {
     console.log(JSON.stringify({

tests/primitives/test_python_interpreter.py

@@ -1,8 +1,7 @@
 import random
 import shutil
-
+import os
 import pytest
-
 from dspy.primitives.python_interpreter import InterpreterError, PythonInterpreter
 
 # This test suite requires deno to be installed. Please install deno following https://docs.deno.com/runtime/getting_started/installation/
@@ -59,5 +58,112 @@ def test_final_answer_trick():
         code = f"final_answer('The result is', {token})"
         result = interpreter(code)
 
-        # They should matain the same order
+        # They should maintain the same order
         assert result == ["The result is", token], "The returned results are differ, `final_answer` trick doesn't work"
+    
+def test_enable_env_vars_flag():
+    os.environ["FOO_TEST_ENV"] = "test_value"
+
+    with PythonInterpreter(enable_env_vars=None) as interpreter:
+        code = "import os\nresult = os.getenv('FOO_TEST_ENV')\nresult"
+        result = interpreter.execute(code)
+        assert result == "", "Environment variables should be inaccessible without allow-env"
+
+    with PythonInterpreter(enable_env_vars=["FOO_TEST_ENV"]) as interpreter:
+        code = "import os\nresult = os.getenv('FOO_TEST_ENV')\nresult"
+        result = interpreter.execute(code)
+        assert result == "test_value", "Environment variables should be accessible with allow-env"
+
+
+
+def test_read_file_access_control(tmp_path):
+    testfile_path = tmp_path / "test_temp_file.txt"
+    virtual_path = f"/sandbox/{testfile_path.name}"
+    with open(testfile_path, "w") as f:
+        f.write("test content")
+
+    with PythonInterpreter(enable_read_paths=[str(testfile_path)]) as interpreter:
+        code = (
+            f"with open({repr(virtual_path)}, 'r') as f:\n"
+            f"    data = f.read()\n"
+            f"data"
+        )
+        result = interpreter.execute(code)
+        assert result == "test content", "Test file should be accessible with enable_read_paths and specified file"
+
+    with PythonInterpreter(enable_read_paths=None) as interpreter:
+        code = (
+            f"try:\n"
+            f"    with open({repr(virtual_path)}, 'r') as f:\n"
+            f"        data = f.read()\n"
+            f"except Exception as e:\n"
+            f"    data = str(e)\n"
+            f"data"
+        )
+        result = interpreter.execute(code)
+        assert ("PermissionDenied" in result or "denied" in result.lower() or "no such file" in result.lower()), "Test file should not be accessible without enable_read_paths"
+
+def test_enable_write_flag(tmp_path):
+    testfile_path = tmp_path / "test_temp_output.txt"
+    virtual_path = f"/sandbox/{testfile_path.name}"
+
+    with PythonInterpreter(enable_write_paths=None) as interpreter:
+        code = (
+            f"try:\n"
+            f"    with open({repr(virtual_path)}, 'w') as f:\n"
+            f"        f.write('blocked')\n"
+            f"    result = 'wrote'\n"
+            f"except Exception as e:\n"
+            f"    result = str(e)\n"
+            f"result"
+        )
+        result = interpreter.execute(code)
+        assert ("PermissionDenied" in result or "denied" in result.lower() or "no such file" in result.lower()), "Test file should not be writable without enable_write_paths"
+
+    with PythonInterpreter(enable_write_paths=[str(testfile_path)]) as interpreter:
+        code = (
+            f"with open({repr(virtual_path)}, 'w') as f:\n"
+            f"    f.write('allowed')\n"
+            f"'ok'"
+        )
+        result = interpreter.execute(code)
+        assert result == "ok", "Test file should be writable with enable_write_paths"
+    assert testfile_path.exists()
+    with open(testfile_path, "r") as f:
+        assert f.read() == "allowed", "Test file outputs should match content written during execution"
+    
+    with open(testfile_path, "w") as f:
+        f.write("original_content")
+    with PythonInterpreter(enable_write_paths=[str(testfile_path)], sync_files=False) as interpreter:
+        code = (
+            f"with open({repr(virtual_path)}, 'w') as f:\n"
+            f"    f.write('should_not_sync')\n"
+            f"'done_no_sync'"
+        )
+        result = interpreter.execute(code)
+        assert result == "done_no_sync"
+    with open(testfile_path, "r") as f:
+        assert f.read() == "original_content", "File should not be changed when sync_files is False"
+
+
+
+def test_enable_net_flag():
+    test_url = "https://example.com"
+
+    with PythonInterpreter(enable_network_access=None) as interpreter:
+        code = (
+            "import js\n"
+            f"resp = await js.fetch({repr(test_url)})\n"
+            "resp.status"
+        )
+        with pytest.raises(InterpreterError, match="PythonError"):
+            interpreter.execute(code)
+
+    with PythonInterpreter(enable_network_access=["example.com"]) as interpreter:
+        code = (
+            "import js\n"
+            f"resp = await js.fetch({repr(test_url)})\n"
+            "resp.status"
+        )
+        result = interpreter.execute(code)
+        assert int(result) == 200, "Network access is permitted with enable_network_access"