# File based from: https://github.com/microsoft/autogen/blob/main/autogen/coding/local_commandline_code_executor.py# Credit to original authorsimportasyncioimportloggingimportosimportsysimportwarningsfromhashlibimportsha256frompathlibimportPathfromstringimportTemplatefromtypesimportSimpleNamespacefromtypingimportAny,Callable,ClassVar,List,Optional,Sequence,Unionfromautogen_coreimportCancellationTokenfromautogen_core.code_executorimportCodeBlock,CodeExecutor,FunctionWithRequirements,FunctionWithRequirementsStrfromtyping_extensionsimportParamSpecfrom.._commonimport(PYTHON_VARIANTS,CommandLineCodeResult,build_python_functions_file,get_file_name_from_content,lang_to_cmd,silence_pip,to_stub,)__all__=("LocalCommandLineCodeExecutor",)A=ParamSpec("A")
[docs]classLocalCommandLineCodeExecutor(CodeExecutor):"""A code executor class that executes code through a local command line environment. .. danger:: This will execute code on the local machine. If being used with LLM generated code, caution should be used. Each code block is saved as a file and executed in a separate process in the working directory, and a unique file is generated and saved in the working directory for each code block. The code blocks are executed in the order they are received. Command line code is sanitized using regular expression match against a list of dangerous commands in order to prevent self-destructive commands from being executed which may potentially affect the users environment. Currently the only supported languages is Python and shell scripts. For Python code, use the language "python" for the code block. For shell scripts, use the language "bash", "shell", or "sh" for the code block. Args: timeout (int): The timeout for the execution of any single code block. Default is 60. work_dir (str): The working directory for the code execution. If None, a default working directory will be used. The default working directory is the current directory ".". functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions". virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None. Example: How to use `LocalCommandLineCodeExecutor` with a virtual environment different from the one used to run the autogen application: Set up a virtual environment using the `venv` module, and pass its context to the initializer of `LocalCommandLineCodeExecutor`. This way, the executor will run code within the new environment. .. code-block:: python import venv from pathlib import Path import asyncio from autogen_core import CancellationToken from autogen_core.code_executor import CodeBlock from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor async def example(): work_dir = Path("coding") work_dir.mkdir(exist_ok=True) venv_dir = work_dir / ".venv" venv_builder = venv.EnvBuilder(with_pip=True) venv_builder.create(venv_dir) venv_context = venv_builder.ensure_directories(venv_dir) local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context) await local_executor.execute_code_blocks( code_blocks=[ CodeBlock(language="bash", code="pip install matplotlib"), ], cancellation_token=CancellationToken(), ) asyncio.run(example()) """SUPPORTED_LANGUAGES:ClassVar[List[str]]=["bash","shell","sh","pwsh","powershell","ps1","python",]FUNCTION_PROMPT_TEMPLATE:ClassVar[str]="""You have access to the following user defined functions. They can be accessed from the module called `$module_name` by their function names.For example, if there was a function called `foo` you could import it by writing `from $module_name import foo`$functions"""def__init__(self,timeout:int=60,work_dir:Union[Path,str]=Path("."),functions:Sequence[Union[FunctionWithRequirements[Any,A],Callable[...,Any],FunctionWithRequirementsStr,]]=[],functions_module:str="functions",virtual_env_context:Optional[SimpleNamespace]=None,):iftimeout<1:raiseValueError("Timeout must be greater than or equal to 1.")ifisinstance(work_dir,str):work_dir=Path(work_dir)ifnotfunctions_module.isidentifier():raiseValueError("Module name must be a valid Python identifier")self._functions_module=functions_modulework_dir.mkdir(exist_ok=True)self._timeout=timeoutself._work_dir:Path=work_dirself._functions=functions# Setup could take some time so we intentionally wait for the first code block to do it.iflen(functions)>0:self._setup_functions_complete=Falseelse:self._setup_functions_complete=Trueself._virtual_env_context:Optional[SimpleNamespace]=virtual_env_context
[docs]defformat_functions_for_prompt(self,prompt_template:str=FUNCTION_PROMPT_TEMPLATE)->str:"""(Experimental) Format the functions for a prompt. The template includes two variables: - `$module_name`: The module name. - `$functions`: The functions formatted as stubs with two newlines between each function. Args: prompt_template (str): The prompt template. Default is the class default. Returns: str: The formatted prompt. """template=Template(prompt_template)returntemplate.substitute(module_name=self._functions_module,functions="\n\n".join([to_stub(func)forfuncinself._functions]),)
@propertydeffunctions_module(self)->str:"""(Experimental) The module name for the functions."""returnself._functions_module@propertydeffunctions(self)->List[str]:raiseNotImplementedError@propertydeftimeout(self)->int:"""(Experimental) The timeout for code execution."""returnself._timeout@propertydefwork_dir(self)->Path:"""(Experimental) The working directory for the code execution."""returnself._work_dirasyncdef_setup_functions(self,cancellation_token:CancellationToken)->None:func_file_content=build_python_functions_file(self._functions)func_file=self._work_dir/f"{self._functions_module}.py"func_file.write_text(func_file_content)# Collect requirementslists_of_packages=[x.python_packagesforxinself._functionsifisinstance(x,FunctionWithRequirements)]flattened_packages=[itemforsublistinlists_of_packagesforiteminsublist]required_packages=list(set(flattened_packages))iflen(required_packages)>0:logging.info("Ensuring packages are installed in executor.")cmd_args=["-m","pip","install"]cmd_args.extend(required_packages)ifself._virtual_env_context:py_executable=self._virtual_env_context.env_exeelse:py_executable=sys.executabletask=asyncio.create_task(asyncio.create_subprocess_exec(py_executable,*cmd_args,cwd=self._work_dir,stdout=asyncio.subprocess.PIPE,stderr=asyncio.subprocess.PIPE,))cancellation_token.link_future(task)try:proc=awaittaskstdout,stderr=awaitasyncio.wait_for(proc.communicate(),self._timeout)exceptasyncio.TimeoutErrorase:raiseValueError("Pip install timed out")fromeexceptasyncio.CancelledErrorase:raiseValueError("Pip install was cancelled")fromeifproc.returncodeisnotNoneandproc.returncode!=0:raiseValueError(f"Pip install failed. {stdout.decode()}, {stderr.decode()}")# Attempt to load the function file to check for syntax errors, imports etc.exec_result=awaitself._execute_code_dont_check_setup([CodeBlock(code=func_file_content,language="python")],cancellation_token)ifexec_result.exit_code!=0:raiseValueError(f"Functions failed to load: {exec_result.output}")self._setup_functions_complete=True
[docs]asyncdefexecute_code_blocks(self,code_blocks:List[CodeBlock],cancellation_token:CancellationToken)->CommandLineCodeResult:"""(Experimental) Execute the code blocks and return the result. Args: code_blocks (List[CodeBlock]): The code blocks to execute. cancellation_token (CancellationToken): a token to cancel the operation Returns: CommandLineCodeResult: The result of the code execution."""ifnotself._setup_functions_complete:awaitself._setup_functions(cancellation_token)returnawaitself._execute_code_dont_check_setup(code_blocks,cancellation_token)
asyncdef_execute_code_dont_check_setup(self,code_blocks:List[CodeBlock],cancellation_token:CancellationToken)->CommandLineCodeResult:logs_all:str=""file_names:List[Path]=[]exitcode=0forcode_blockincode_blocks:lang,code=code_block.language,code_block.codelang=lang.lower()code=silence_pip(code,lang)iflanginPYTHON_VARIANTS:lang="python"iflangnotinself.SUPPORTED_LANGUAGES:# In case the language is not supported, we return an error message.exitcode=1logs_all+="\n"+f"unknown language {lang}"breaktry:# Check if there is a filename commentfilename=get_file_name_from_content(code,self._work_dir)exceptValueError:returnCommandLineCodeResult(exit_code=1,output="Filename is not in the workspace",code_file=None,)iffilenameisNone:# create a file with an automatically generated namecode_hash=sha256(code.encode()).hexdigest()filename=f"tmp_code_{code_hash}.{'py'iflang.startswith('python')elselang}"written_file=(self._work_dir/filename).resolve()withwritten_file.open("w",encoding="utf-8")asf:f.write(code)file_names.append(written_file)env=os.environ.copy()ifself._virtual_env_context:virtual_env_exe_abs_path=os.path.abspath(self._virtual_env_context.env_exe)virtual_env_bin_abs_path=os.path.abspath(self._virtual_env_context.bin_path)env["PATH"]=f"{virtual_env_bin_abs_path}{os.pathsep}{env['PATH']}"program=virtual_env_exe_abs_pathiflang.startswith("python")elselang_to_cmd(lang)else:program=sys.executableiflang.startswith("python")elselang_to_cmd(lang)# Wrap in a task to make it cancellabletask=asyncio.create_task(asyncio.create_subprocess_exec(program,str(written_file.absolute()),cwd=self._work_dir,stdout=asyncio.subprocess.PIPE,stderr=asyncio.subprocess.PIPE,env=env,))cancellation_token.link_future(task)try:proc=awaittaskstdout,stderr=awaitasyncio.wait_for(proc.communicate(),self._timeout)exitcode=proc.returncodeor0exceptasyncio.TimeoutError:logs_all+="\n Timeout"# Same exit code as the timeout command on linux.exitcode=124breakexceptasyncio.CancelledError:logs_all+="\n Cancelled"# TODO: which exit code? 125 is Operation Canceledexitcode=125breakself._running_cmd_task=Nonelogs_all+=stderr.decode()logs_all+=stdout.decode()ifexitcode!=0:breakcode_file=str(file_names[0])iflen(file_names)>0elseNonereturnCommandLineCodeResult(exit_code=exitcode,output=logs_all,code_file=code_file)
[docs]asyncdefrestart(self)->None:"""(Experimental) Restart the code executor."""warnings.warn("Restarting local command line code executor is not supported. No action is taken.",stacklevel=2,)