mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-25 16:18:13 +00:00
feat: allow to set unsafe_mode from Guardrail task
This commit is contained in:
@@ -19,7 +19,7 @@ class GuardrailTask:
|
|||||||
task (Task, optional): The task whose output needs validation.
|
task (Task, optional): The task whose output needs validation.
|
||||||
llm (LLM, optional): The language model to use for code generation.
|
llm (LLM, optional): The language model to use for code generation.
|
||||||
additional_instructions (str, optional): Additional instructions for the guardrail task.
|
additional_instructions (str, optional): Additional instructions for the guardrail task.
|
||||||
|
unsafe_mode (bool, optional): Whether to run the code in unsafe mode.
|
||||||
Raises:
|
Raises:
|
||||||
ValueError: If no valid LLM is provided.
|
ValueError: If no valid LLM is provided.
|
||||||
"""
|
"""
|
||||||
@@ -30,6 +30,7 @@ class GuardrailTask:
|
|||||||
task: Task | None = None,
|
task: Task | None = None,
|
||||||
llm: LLM | None = None,
|
llm: LLM | None = None,
|
||||||
additional_instructions: str = "",
|
additional_instructions: str = "",
|
||||||
|
unsafe_mode: bool | None = None,
|
||||||
):
|
):
|
||||||
self.description = description
|
self.description = description
|
||||||
|
|
||||||
@@ -44,6 +45,7 @@ class GuardrailTask:
|
|||||||
self.llm: LLM | None = llm or fallback_llm
|
self.llm: LLM | None = llm or fallback_llm
|
||||||
|
|
||||||
self.additional_instructions = additional_instructions
|
self.additional_instructions = additional_instructions
|
||||||
|
self.unsafe_mode = unsafe_mode
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def system_instructions(self) -> str:
|
def system_instructions(self) -> str:
|
||||||
@@ -138,7 +140,11 @@ class GuardrailTask:
|
|||||||
|
|
||||||
code = self.generate_code(task_output)
|
code = self.generate_code(task_output)
|
||||||
|
|
||||||
unsafe_mode = not self.check_docker_available()
|
unsafe_mode = (
|
||||||
|
self.unsafe_mode
|
||||||
|
if self.unsafe_mode is not None
|
||||||
|
else not self.check_docker_available()
|
||||||
|
)
|
||||||
|
|
||||||
result = CodeInterpreterTool(code=code, unsafe_mode=unsafe_mode).run()
|
result = CodeInterpreterTool(code=code, unsafe_mode=unsafe_mode).run()
|
||||||
|
|
||||||
|
|||||||
@@ -334,14 +334,12 @@ def test_guardrail_task_when_docker_is_not_available(mock_llm, task_output):
|
|||||||
) as mock_init,
|
) as mock_init,
|
||||||
patch(
|
patch(
|
||||||
"crewai_tools.CodeInterpreterTool.run", return_value=(True, "Valid output")
|
"crewai_tools.CodeInterpreterTool.run", return_value=(True, "Valid output")
|
||||||
) as mock_run,
|
),
|
||||||
patch(
|
patch(
|
||||||
"subprocess.run",
|
"subprocess.run",
|
||||||
side_effect=FileNotFoundError,
|
side_effect=FileNotFoundError,
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
mock_init.return_value = None
|
|
||||||
mock_run.return_value = (True, "Valid output")
|
|
||||||
guardrail(task_output)
|
guardrail(task_output)
|
||||||
|
|
||||||
mock_init.assert_called_once_with(code=ANY, unsafe_mode=True)
|
mock_init.assert_called_once_with(code=ANY, unsafe_mode=True)
|
||||||
@@ -361,8 +359,6 @@ def test_guardrail_task_when_docker_is_available(mock_llm, task_output):
|
|||||||
return_value=True,
|
return_value=True,
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
mock_init.return_value = None
|
|
||||||
mock_run.return_value = (True, "Valid output")
|
|
||||||
guardrail(task_output)
|
guardrail(task_output)
|
||||||
|
|
||||||
mock_init.assert_called_once_with(code=ANY, unsafe_mode=False)
|
mock_init.assert_called_once_with(code=ANY, unsafe_mode=False)
|
||||||
@@ -380,10 +376,32 @@ def test_guardrail_task_when_tool_output_is_not_valid(mock_llm, task_output):
|
|||||||
patch(
|
patch(
|
||||||
"subprocess.run",
|
"subprocess.run",
|
||||||
return_value=True,
|
return_value=True,
|
||||||
),
|
) as docker_check,
|
||||||
):
|
):
|
||||||
mock_init.return_value = None
|
|
||||||
mock_run.return_value = (True, "Valid output")
|
|
||||||
guardrail(task_output)
|
guardrail(task_output)
|
||||||
|
|
||||||
mock_init.assert_called_once_with(code=ANY, unsafe_mode=False)
|
mock_init.assert_called_once_with(code=ANY, unsafe_mode=False)
|
||||||
|
docker_check.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("unsafe_mode", [True, False])
|
||||||
|
def test_guardrail_task_force_code_tool_unsafe_mode(mock_llm, task_output, unsafe_mode):
|
||||||
|
guardrail = GuardrailTask(
|
||||||
|
description="Test validation", llm=mock_llm, unsafe_mode=unsafe_mode
|
||||||
|
)
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
"crewai_tools.CodeInterpreterTool.__init__", return_value=None
|
||||||
|
) as mock_init,
|
||||||
|
patch(
|
||||||
|
"crewai_tools.CodeInterpreterTool.run", return_value=(True, "Valid output")
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"subprocess.run",
|
||||||
|
side_effect=FileNotFoundError,
|
||||||
|
) as docker_check,
|
||||||
|
):
|
||||||
|
guardrail(task_output)
|
||||||
|
|
||||||
|
docker_check.assert_not_called()
|
||||||
|
mock_init.assert_called_once_with(code=ANY, unsafe_mode=unsafe_mode)
|
||||||
|
|||||||
Reference in New Issue
Block a user