support unsafe code execution. add in docker install and running checks. (#1496)

* support unsafe code execution. add in docker install and running checks.

* Update return type
This commit is contained in:
Brandon Hancock (bhancock_ai)
2024-10-23 11:01:00 -04:00
committed by GitHub
parent 7d68e287cc
commit e1fd83e6a7
2 changed files with 37 additions and 5 deletions

View File

@@ -1,6 +1,8 @@
import os
import shutil
import subprocess
from inspect import signature
from typing import Any, List, Optional, Union
from typing import Any, List, Literal, Optional, Union
from pydantic import Field, InstanceOf, PrivateAttr, model_validator
@@ -112,6 +114,10 @@ class Agent(BaseAgent):
default=2,
description="Maximum number of retries for an agent to execute a task when an error occurs.",
)
code_execution_mode: Literal["safe", "unsafe"] = Field(
default="safe",
description="Mode for code execution: 'safe' (using Docker) or 'unsafe' (direct execution).",
)
@model_validator(mode="after")
def post_init_setup(self):
@@ -173,6 +179,9 @@ class Agent(BaseAgent):
if not self.agent_executor:
self._setup_agent_executor()
if self.allow_code_execution:
self._validate_docker_installation()
return self
def _setup_agent_executor(self):
@@ -308,7 +317,9 @@ class Agent(BaseAgent):
try:
from crewai_tools import CodeInterpreterTool
return [CodeInterpreterTool()]
# Set the unsafe_mode based on the code_execution_mode attribute
unsafe_mode = self.code_execution_mode == "unsafe"
return [CodeInterpreterTool(unsafe_mode=unsafe_mode)]
except ModuleNotFoundError:
self._logger.log(
"info", "Coding tools not available. Install crewai_tools. "
@@ -408,6 +419,25 @@ class Agent(BaseAgent):
return "\n".join(tool_strings)
def _validate_docker_installation(self) -> None:
"""Check if Docker is installed and running."""
if not shutil.which("docker"):
raise RuntimeError(
f"Docker is not installed. Please install Docker to use code execution with agent: {self.role}"
)
try:
subprocess.run(
["docker", "info"],
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
except subprocess.CalledProcessError:
raise RuntimeError(
f"Docker is not running. Please start Docker to use code execution with agent: {self.role}"
)
@staticmethod
def __tools_names(tools) -> str:
return ", ".join([t.name for t in tools])