mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
feat: add auto-discovery for Guardrail code execution mode
This commit is contained in:
@@ -11,7 +11,8 @@ class GuardrailTask:
|
|||||||
|
|
||||||
This class generates and executes Python code to validate task outputs based on
|
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
|
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:
|
Args:
|
||||||
description (str): The description of the validation criteria.
|
description (str): The description of the validation criteria.
|
||||||
@@ -28,11 +29,9 @@ class GuardrailTask:
|
|||||||
description: str,
|
description: str,
|
||||||
task: Task | None = None,
|
task: Task | None = None,
|
||||||
llm: LLM | None = None,
|
llm: LLM | None = None,
|
||||||
unsafe_mode: bool = False,
|
|
||||||
additional_instructions: str = "",
|
additional_instructions: str = "",
|
||||||
):
|
):
|
||||||
self.description = description
|
self.description = description
|
||||||
self.unsafe_mode: bool = unsafe_mode
|
|
||||||
|
|
||||||
fallback_llm: LLM | None = (
|
fallback_llm: LLM | None = (
|
||||||
task.agent.llm
|
task.agent.llm
|
||||||
@@ -138,7 +137,10 @@ class GuardrailTask:
|
|||||||
from crewai_tools import CodeInterpreterTool
|
from crewai_tools import CodeInterpreterTool
|
||||||
|
|
||||||
code = self.generate_code(task_output)
|
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 = [
|
error_messages = [
|
||||||
"Something went wrong while running the code",
|
"Something went wrong while running the code",
|
||||||
@@ -152,3 +154,12 @@ class GuardrailTask:
|
|||||||
result = ast.literal_eval(result)
|
result = ast.literal_eval(result)
|
||||||
|
|
||||||
return 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
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
from unittest.mock import Mock, patch
|
from unittest.mock import ANY, Mock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
@@ -319,3 +319,45 @@ def test_guardrail_emits_events(sample_agent):
|
|||||||
]
|
]
|
||||||
assert started_guardrail == expected_started_events
|
assert started_guardrail == expected_started_events
|
||||||
assert completed_guardrail == expected_completed_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)
|
||||||
|
|||||||
Reference in New Issue
Block a user