feat: add auto-discovery for Guardrail code execution mode

This commit is contained in:
Lucas Gomide
2025-04-21 19:19:37 -03:00
parent 91b618b4e0
commit 50453c6984
2 changed files with 58 additions and 5 deletions

View File

@@ -11,7 +11,8 @@ class GuardrailTask:
This class generates and executes Python code to validate task outputs based on
specified criteria. It uses an LLM to generate the validation code and provides
safety guardrails for code execution.
safety guardrails for code execution. The code is executed in a Docker container
if available, otherwise it is executed in the current environment.
Args:
description (str): The description of the validation criteria.
@@ -28,11 +29,9 @@ class GuardrailTask:
description: str,
task: Task | None = None,
llm: LLM | None = None,
unsafe_mode: bool = False,
additional_instructions: str = "",
):
self.description = description
self.unsafe_mode: bool = unsafe_mode
fallback_llm: LLM | None = (
task.agent.llm
@@ -138,7 +137,10 @@ class GuardrailTask:
from crewai_tools import CodeInterpreterTool
code = self.generate_code(task_output)
result = CodeInterpreterTool(code=code, unsafe_mode=self.unsafe_mode).run()
unsafe_mode = not self.check_docker_available()
result = CodeInterpreterTool(code=code, unsafe_mode=unsafe_mode).run()
error_messages = [
"Something went wrong while running the code",
@@ -152,3 +154,12 @@ class GuardrailTask:
result = ast.literal_eval(result)
return result
def check_docker_available(self) -> bool:
import subprocess
try:
subprocess.run(["docker", "--version"], check=True)
return True
except (subprocess.CalledProcessError, FileNotFoundError):
return False

View File

@@ -1,4 +1,4 @@
from unittest.mock import Mock, patch
from unittest.mock import ANY, Mock, patch
import pytest
@@ -319,3 +319,45 @@ def test_guardrail_emits_events(sample_agent):
]
assert started_guardrail == expected_started_events
assert completed_guardrail == expected_completed_events
def test_guardrail_task_when_docker_is_not_available(mock_llm, task_output):
guardrail = GuardrailTask(description="Test validation", llm=mock_llm)
with (
patch(
"crewai_tools.CodeInterpreterTool.__init__", return_value=None
) as mock_init,
patch(
"crewai_tools.CodeInterpreterTool.run", return_value=(True, "Valid output")
) as mock_run,
patch(
"subprocess.run",
side_effect=FileNotFoundError,
),
):
mock_init.return_value = None
mock_run.return_value = (True, "Valid output")
guardrail(task_output)
mock_init.assert_called_once_with(code=ANY, unsafe_mode=True)
def test_guardrail_task_when_docker_is_available(mock_llm, task_output):
guardrail = GuardrailTask(description="Test validation", llm=mock_llm)
with (
patch(
"crewai_tools.CodeInterpreterTool.__init__", return_value=None
) as mock_init,
patch(
"crewai_tools.CodeInterpreterTool.run", return_value=(True, "Valid output")
) as mock_run,
patch(
"subprocess.run",
return_value=True,
),
):
mock_init.return_value = None
mock_run.return_value = (True, "Valid output")
guardrail(task_output)
mock_init.assert_called_once_with(code=ANY, unsafe_mode=False)