first commit
This commit is contained in:
DigiJ
517
agent_dhal/agentdhal_extensions/code_executors/local/__init__.py
Normal file
517
agent_dhal/agentdhal_extensions/code_executors/local/__init__.py
Normal file
@@ -0,0 +1,517 @@
|
||||
# File based from: https://github.com/microsoft/autogen/blob/main/autogen/coding/local_commandline_code_executor.py
|
||||
# Credit to original authors
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import warnings
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from string import Template
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union
|
||||
|
||||
from agentdhal_core import CancellationToken, Component
|
||||
from agentdhal_core.code_executor import CodeBlock, CodeExecutor, FunctionWithRequirements, FunctionWithRequirementsStr
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ParamSpec, Self
|
||||
|
||||
from .._common import (
|
||||
PYTHON_VARIANTS,
|
||||
CommandLineCodeResult,
|
||||
build_python_functions_file,
|
||||
get_file_name_from_content,
|
||||
lang_to_cmd,
|
||||
silence_pip,
|
||||
to_stub,
|
||||
)
|
||||
|
||||
__all__ = ("LocalCommandLineCodeExecutor",)
|
||||
|
||||
A = ParamSpec("A")
|
||||
|
||||
|
||||
class LocalCommandLineCodeExecutorConfig(BaseModel):
|
||||
"""Configuration for LocalCommandLineCodeExecutor"""
|
||||
|
||||
timeout: int = 60
|
||||
work_dir: Optional[str] = None
|
||||
functions_module: str = "functions"
|
||||
cleanup_temp_files: bool = True
|
||||
|
||||
|
||||
class LocalCommandLineCodeExecutor(CodeExecutor, Component[LocalCommandLineCodeExecutorConfig]):
|
||||
"""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", "sh", "pwsh", "powershell", or "ps1" for the code
|
||||
block.
|
||||
|
||||
.. note::
|
||||
|
||||
On Windows, the event loop policy must be set to `WindowsProactorEventLoopPolicy` to avoid issues with subprocesses.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import sys
|
||||
import asyncio
|
||||
|
||||
if sys.platform == "win32":
|
||||
asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
|
||||
|
||||
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 a temporary 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".
|
||||
cleanup_temp_files (bool, optional): Whether to automatically clean up temporary files after execution. Defaults to True.
|
||||
virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None.
|
||||
|
||||
.. note::
|
||||
Using the current directory (".") as working directory is deprecated. Using it will raise a deprecation warning.
|
||||
|
||||
|
||||
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 agentdhal_core import CancellationToken
|
||||
from agentdhal_core.code_executor import CodeBlock
|
||||
from agentdhal_extensions.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())
|
||||
|
||||
"""
|
||||
|
||||
component_config_schema = LocalCommandLineCodeExecutorConfig
|
||||
component_provider_override = "agentdhal_extensions.code_executors.local.LocalCommandLineCodeExecutor"
|
||||
|
||||
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: Optional[Union[Path, str]] = None,
|
||||
functions: Sequence[
|
||||
Union[
|
||||
FunctionWithRequirements[Any, A],
|
||||
Callable[..., Any],
|
||||
FunctionWithRequirementsStr,
|
||||
]
|
||||
] = [],
|
||||
functions_module: str = "functions",
|
||||
cleanup_temp_files: bool = True,
|
||||
virtual_env_context: Optional[SimpleNamespace] = None,
|
||||
):
|
||||
if timeout < 1:
|
||||
raise ValueError("Timeout must be greater than or equal to 1.")
|
||||
self._timeout = timeout
|
||||
|
||||
self._work_dir: Optional[Path] = None
|
||||
if work_dir is not None:
|
||||
# Check if user provided work_dir is the current directory and warn if so.
|
||||
if Path(work_dir).resolve() == Path.cwd().resolve():
|
||||
warnings.warn(
|
||||
"Using the current directory as work_dir is deprecated.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
if isinstance(work_dir, str):
|
||||
self._work_dir = Path(work_dir)
|
||||
else:
|
||||
self._work_dir = work_dir
|
||||
self._work_dir.mkdir(exist_ok=True)
|
||||
|
||||
self._functions = functions
|
||||
# Setup could take some time so we intentionally wait for the first code block to do it.
|
||||
if len(functions) > 0:
|
||||
self._setup_functions_complete = False
|
||||
else:
|
||||
self._setup_functions_complete = True
|
||||
|
||||
if not functions_module.isidentifier():
|
||||
raise ValueError("Module name must be a valid Python identifier")
|
||||
self._functions_module = functions_module
|
||||
|
||||
self._cleanup_temp_files = cleanup_temp_files
|
||||
self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context
|
||||
|
||||
self._temp_dir: Optional[tempfile.TemporaryDirectory[str]] = None
|
||||
self._started = False
|
||||
|
||||
# Check the current event loop policy if on windows.
|
||||
if sys.platform == "win32":
|
||||
current_policy = asyncio.get_event_loop_policy()
|
||||
if hasattr(asyncio, "WindowsProactorEventLoopPolicy") and not isinstance(
|
||||
current_policy, asyncio.WindowsProactorEventLoopPolicy
|
||||
):
|
||||
warnings.warn(
|
||||
"The current event loop policy is not WindowsProactorEventLoopPolicy. "
|
||||
"This may cause issues with subprocesses. "
|
||||
"Try setting the event loop policy to WindowsProactorEventLoopPolicy. "
|
||||
"For example: `asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())`. "
|
||||
"See https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.ProactorEventLoop.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
def format_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)
|
||||
return template.substitute(
|
||||
module_name=self._functions_module,
|
||||
functions="\n\n".join([to_stub(func) for func in self._functions]),
|
||||
)
|
||||
|
||||
@property
|
||||
def timeout(self) -> int:
|
||||
"""(Experimental) The timeout for code execution."""
|
||||
return self._timeout
|
||||
|
||||
@property
|
||||
def work_dir(self) -> Path:
|
||||
"""(Experimental) The working directory for the code execution."""
|
||||
if self._work_dir is not None:
|
||||
return self._work_dir
|
||||
else:
|
||||
# Automatically create temp directory if not exists
|
||||
if self._temp_dir is None:
|
||||
self._temp_dir = tempfile.TemporaryDirectory()
|
||||
self._started = True
|
||||
return Path(self._temp_dir.name)
|
||||
|
||||
@property
|
||||
def functions(self) -> List[str]:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def functions_module(self) -> str:
|
||||
"""(Experimental) The module name for the functions."""
|
||||
return self._functions_module
|
||||
|
||||
@property
|
||||
def cleanup_temp_files(self) -> bool:
|
||||
"""(Experimental) Whether to automatically clean up temporary files after execution."""
|
||||
return self._cleanup_temp_files
|
||||
|
||||
async def _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 requirements
|
||||
lists_of_packages = [x.python_packages for x in self._functions if isinstance(x, FunctionWithRequirements)]
|
||||
flattened_packages = [item for sublist in lists_of_packages for item in sublist]
|
||||
required_packages = list(set(flattened_packages))
|
||||
if len(required_packages) > 0:
|
||||
logging.info("Ensuring packages are installed in executor.")
|
||||
|
||||
cmd_args = ["-m", "pip", "install"]
|
||||
cmd_args.extend(required_packages)
|
||||
|
||||
if self._virtual_env_context:
|
||||
py_executable = self._virtual_env_context.env_exe
|
||||
else:
|
||||
py_executable = sys.executable
|
||||
|
||||
task = 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 = await task
|
||||
stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
|
||||
except asyncio.TimeoutError as e:
|
||||
raise ValueError("Pip install timed out") from e
|
||||
except asyncio.CancelledError as e:
|
||||
raise ValueError("Pip install was cancelled") from e
|
||||
|
||||
if proc.returncode is not None and proc.returncode != 0:
|
||||
raise ValueError(f"Pip install failed. {stdout.decode()}, {stderr.decode()}")
|
||||
|
||||
# Attempt to load the function file to check for syntax errors, imports etc.
|
||||
exec_result = await self._execute_code_dont_check_setup(
|
||||
[CodeBlock(code=func_file_content, language="python")], cancellation_token
|
||||
)
|
||||
|
||||
if exec_result.exit_code != 0:
|
||||
raise ValueError(f"Functions failed to load: {exec_result.output}")
|
||||
|
||||
self._setup_functions_complete = True
|
||||
|
||||
async def execute_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."""
|
||||
|
||||
if not self._setup_functions_complete:
|
||||
await self._setup_functions(cancellation_token)
|
||||
|
||||
return await self._execute_code_dont_check_setup(code_blocks, cancellation_token)
|
||||
|
||||
async def _execute_code_dont_check_setup(
|
||||
self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
|
||||
) -> CommandLineCodeResult:
|
||||
"""
|
||||
Execute the provided code blocks in the local command line without re-checking setup.
|
||||
Returns a CommandLineCodeResult indicating success or failure.
|
||||
"""
|
||||
logs_all: str = ""
|
||||
file_names: List[Path] = []
|
||||
exitcode = 0
|
||||
|
||||
for code_block in code_blocks:
|
||||
lang, code = code_block.language, code_block.code
|
||||
lang = lang.lower()
|
||||
|
||||
# Remove pip output where possible
|
||||
code = silence_pip(code, lang)
|
||||
|
||||
# Normalize python variants to "python"
|
||||
if lang in PYTHON_VARIANTS:
|
||||
lang = "python"
|
||||
|
||||
# Abort if not supported
|
||||
if lang not in self.SUPPORTED_LANGUAGES:
|
||||
exitcode = 1
|
||||
logs_all += "\n" + f"unknown language {lang}"
|
||||
break
|
||||
|
||||
# Try extracting a filename (if present)
|
||||
try:
|
||||
filename = get_file_name_from_content(code, self.work_dir)
|
||||
except ValueError:
|
||||
return CommandLineCodeResult(
|
||||
exit_code=1,
|
||||
output="Filename is not in the workspace",
|
||||
code_file=None,
|
||||
)
|
||||
|
||||
# If no filename is found, create one
|
||||
if filename is None:
|
||||
code_hash = sha256(code.encode()).hexdigest()
|
||||
if lang.startswith("python"):
|
||||
ext = "py"
|
||||
elif lang in ["pwsh", "powershell", "ps1"]:
|
||||
ext = "ps1"
|
||||
else:
|
||||
ext = lang
|
||||
|
||||
filename = f"tmp_code_{code_hash}.{ext}"
|
||||
|
||||
written_file = (self.work_dir / filename).resolve()
|
||||
with written_file.open("w", encoding="utf-8") as f:
|
||||
f.write(code)
|
||||
file_names.append(written_file)
|
||||
|
||||
# Build environment
|
||||
env = os.environ.copy()
|
||||
if self._virtual_env_context:
|
||||
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']}"
|
||||
|
||||
# Decide how to invoke the script
|
||||
if lang == "python":
|
||||
program = (
|
||||
os.path.abspath(self._virtual_env_context.env_exe) if self._virtual_env_context else sys.executable
|
||||
)
|
||||
extra_args = [str(written_file.absolute())]
|
||||
else:
|
||||
# Get the appropriate command for the language
|
||||
program = lang_to_cmd(lang)
|
||||
|
||||
# Special handling for PowerShell
|
||||
if program == "pwsh":
|
||||
extra_args = [
|
||||
"-NoProfile",
|
||||
"-ExecutionPolicy",
|
||||
"Bypass",
|
||||
"-File",
|
||||
str(written_file.absolute()),
|
||||
]
|
||||
else:
|
||||
# Shell commands (bash, sh, etc.)
|
||||
extra_args = [str(written_file.absolute())]
|
||||
|
||||
# Create a subprocess and run
|
||||
task = asyncio.create_task(
|
||||
asyncio.create_subprocess_exec(
|
||||
program,
|
||||
*extra_args,
|
||||
cwd=self.work_dir,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
)
|
||||
cancellation_token.link_future(task)
|
||||
|
||||
proc = None # Track the process
|
||||
try:
|
||||
proc = await task
|
||||
stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
|
||||
exitcode = proc.returncode or 0
|
||||
except asyncio.TimeoutError:
|
||||
logs_all += "\nTimeout"
|
||||
exitcode = 124
|
||||
if proc:
|
||||
proc.terminate()
|
||||
await proc.wait() # Ensure process is fully dead
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
logs_all += "\nCancelled"
|
||||
exitcode = 125
|
||||
if proc:
|
||||
proc.terminate()
|
||||
await proc.wait()
|
||||
break
|
||||
|
||||
logs_all += stderr.decode()
|
||||
logs_all += stdout.decode()
|
||||
|
||||
if exitcode != 0:
|
||||
break
|
||||
|
||||
code_file = str(file_names[0]) if file_names else None
|
||||
code_result = CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file)
|
||||
|
||||
if self._cleanup_temp_files:
|
||||
for file in file_names:
|
||||
try:
|
||||
file.unlink(missing_ok=True)
|
||||
except OSError as error:
|
||||
logging.error(f"Failed to delete temporary file {file}: {error}")
|
||||
|
||||
return code_result
|
||||
|
||||
async def restart(self) -> None:
|
||||
"""(Experimental) Restart the code executor."""
|
||||
warnings.warn(
|
||||
"Restarting local command line code executor is not supported. No action is taken.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
async def start(self) -> None:
|
||||
"""(Experimental) Start the code executor.
|
||||
|
||||
Initializes the local code executor and should be called before executing any code blocks.
|
||||
It marks the executor internal state as started.
|
||||
If no working directory is provided, the method creates a temporary directory for the executor to use.
|
||||
"""
|
||||
if self._work_dir is None and self._temp_dir is None:
|
||||
self._temp_dir = tempfile.TemporaryDirectory()
|
||||
self._started = True
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""(Experimental) Stop the code executor.
|
||||
|
||||
Stops the local code executor and performs the cleanup of the temporary working directory (if it was created).
|
||||
The executor's internal state is markes as no longer started.
|
||||
"""
|
||||
if self._temp_dir is not None:
|
||||
self._temp_dir.cleanup()
|
||||
self._temp_dir = None
|
||||
self._started = False
|
||||
pass
|
||||
|
||||
def _to_config(self) -> LocalCommandLineCodeExecutorConfig:
|
||||
if self._functions:
|
||||
logging.info("Functions will not be included in serialized configuration")
|
||||
if self._virtual_env_context:
|
||||
logging.info("Virtual environment context will not be included in serialized configuration")
|
||||
|
||||
return LocalCommandLineCodeExecutorConfig(
|
||||
timeout=self._timeout,
|
||||
work_dir=str(self.work_dir),
|
||||
functions_module=self._functions_module,
|
||||
cleanup_temp_files=self._cleanup_temp_files,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _from_config(cls, config: LocalCommandLineCodeExecutorConfig) -> Self:
|
||||
return cls(
|
||||
timeout=config.timeout,
|
||||
work_dir=Path(config.work_dir) if config.work_dir is not None else None,
|
||||
functions_module=config.functions_module,
|
||||
cleanup_temp_files=config.cleanup_temp_files,
|
||||
)
|
||||
Reference in New Issue
Block a user