88from pathlib import Path
99from time import sleep
1010from types import TracebackType
11- from typing import Any , List , Optional , Type , Union
11+ from typing import Any , ClassVar , Dict , List , Optional , Type , Union
1212
1313import docker
1414from docker .errors import ImageNotFound
@@ -39,14 +39,30 @@ def _wait_for_ready(container: Any, timeout: int = 60, stop_time: float = 0.1) -
3939
4040
4141class DockerCommandLineCodeExecutor (CodeExecutor ):
42+ DEFAULT_EXECUTION_POLICY : ClassVar [Dict [str , bool ]] = {
43+ "bash" : True ,
44+ "shell" : True ,
45+ "sh" : True ,
46+ "pwsh" : True ,
47+ "powershell" : True ,
48+ "ps1" : True ,
49+ "python" : True ,
50+ "javascript" : False ,
51+ "html" : False ,
52+ "css" : False ,
53+ }
54+ LANGUAGE_ALIASES : ClassVar [Dict [str , str ]] = {"py" : "python" , "js" : "javascript" }
55+
4256 def __init__ (
4357 self ,
4458 image : str = "python:3-slim" ,
4559 container_name : Optional [str ] = None ,
4660 timeout : int = 60 ,
4761 work_dir : Union [Path , str ] = Path ("." ),
62+ bind_dir : Optional [Union [Path , str ]] = None ,
4863 auto_remove : bool = True ,
4964 stop_container : bool = True ,
65+ execution_policies : Optional [Dict [str , bool ]] = None ,
5066 ):
5167 """(Experimental) A code executor class that executes code through
5268 a command line environment in a Docker container.
@@ -67,6 +83,9 @@ def __init__(
6783 timeout (int, optional): The timeout for code execution. Defaults to 60.
6884 work_dir (Union[Path, str], optional): The working directory for the code
6985 execution. Defaults to Path(".").
86+ bind_dir (Union[Path, str], optional): The directory that will be bound
87+ to the code executor container. Useful for cases where you want to spawn
88+ the container from within a container. Defaults to work_dir.
7089 auto_remove (bool, optional): If true, will automatically remove the Docker
7190 container when it is stopped. Defaults to True.
7291 stop_container (bool, optional): If true, will automatically stop the
@@ -76,17 +95,19 @@ def __init__(
7695 Raises:
7796 ValueError: On argument error, or if the container fails to start.
7897 """
79-
8098 if timeout < 1 :
8199 raise ValueError ("Timeout must be greater than or equal to 1." )
82100
83101 if isinstance (work_dir , str ):
84102 work_dir = Path (work_dir )
85-
86103 work_dir .mkdir (exist_ok = True )
87104
88- client = docker .from_env ()
105+ if bind_dir is None :
106+ bind_dir = work_dir
107+ elif isinstance (bind_dir , str ):
108+ bind_dir = Path (bind_dir )
89109
110+ client = docker .from_env ()
90111 # Check if the image exists
91112 try :
92113 client .images .get (image )
@@ -105,7 +126,7 @@ def __init__(
105126 entrypoint = "/bin/sh" ,
106127 tty = True ,
107128 auto_remove = auto_remove ,
108- volumes = {str (work_dir .resolve ()): {"bind" : "/workspace" , "mode" : "rw" }},
129+ volumes = {str (bind_dir .resolve ()): {"bind" : "/workspace" , "mode" : "rw" }},
109130 working_dir = "/workspace" ,
110131 )
111132 self ._container .start ()
@@ -118,7 +139,6 @@ def cleanup() -> None:
118139 container .stop ()
119140 except docker .errors .NotFound :
120141 pass
121-
122142 atexit .unregister (cleanup )
123143
124144 if stop_container :
@@ -132,6 +152,10 @@ def cleanup() -> None:
132152
133153 self ._timeout = timeout
134154 self ._work_dir : Path = work_dir
155+ self ._bind_dir : Path = bind_dir
156+ self .execution_policies = self .DEFAULT_EXECUTION_POLICY .copy ()
157+ if execution_policies is not None :
158+ self .execution_policies .update (execution_policies )
135159
136160 @property
137161 def timeout (self ) -> int :
@@ -143,6 +167,11 @@ def work_dir(self) -> Path:
143167 """(Experimental) The working directory for the code execution."""
144168 return self ._work_dir
145169
170+ @property
171+ def bind_dir (self ) -> Path :
172+ """(Experimental) The binding directory for the code execution container."""
173+ return self ._bind_dir
174+
146175 @property
147176 def code_extractor (self ) -> CodeExtractor :
148177 """(Experimental) Export a code extractor that can be used by an agent."""
@@ -164,35 +193,42 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeRe
164193 files = []
165194 last_exit_code = 0
166195 for code_block in code_blocks :
167- lang = code_block .language
196+ lang = self .LANGUAGE_ALIASES .get (code_block .language .lower (), code_block .language .lower ())
197+ if lang not in self .DEFAULT_EXECUTION_POLICY :
198+ outputs .append (f"Unsupported language { lang } \n " )
199+ last_exit_code = 1
200+ break
201+
202+ execute_code = self .execution_policies .get (lang , False )
168203 code = silence_pip (code_block .code , lang )
169204
205+ # Check if there is a filename comment
170206 try :
171- # Check if there is a filename comment
172- filename = _get_file_name_from_content (code , Path ("/workspace" ))
207+ filename = _get_file_name_from_content (code , self ._work_dir )
173208 except ValueError :
174- return CommandLineCodeResult (exit_code = 1 , output = "Filename is not in the workspace" )
209+ outputs .append ("Filename is not in the workspace" )
210+ last_exit_code = 1
211+ break
175212
176- if filename is None :
177- # create a file with an automatically generated name
178- code_hash = md5 (code .encode ()).hexdigest ()
179- filename = f"tmp_code_{ code_hash } { 'py' if lang .startswith ('python' ) else lang }
213+ if not filename :
214+ filename = f"tmp_code_{ md5 (code .encode ()).hexdigest ()} { lang }
180215
181216 code_path = self ._work_dir / filename
182217 with code_path .open ("w" , encoding = "utf-8" ) as fout :
183218 fout .write (code )
219+ files .append (code_path )
184220
185- command = ["timeout" , str (self ._timeout ), _cmd (lang ), filename ]
221+ if not execute_code :
222+ outputs .append (f"Code saved to { str (code_path )} \n " )
223+ continue
186224
225+ command = ["timeout" , str (self ._timeout ), _cmd (lang ), filename ]
187226 result = self ._container .exec_run (command )
188227 exit_code = result .exit_code
189228 output = result .output .decode ("utf-8" )
190229 if exit_code == 124 :
191- output += "\n "
192- output += TIMEOUT_MSG
193-
230+ output += "\n " + TIMEOUT_MSG
194231 outputs .append (output )
195- files .append (code_path )
196232
197233 last_exit_code = exit_code
198234 if exit_code != 0 :
0 commit comments