mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-14 10:38:29 +00:00
feat: Add task guardrails feature
Add support for custom code guardrails in tasks that validate outputs before proceeding to the next task. Features include: - Optional task-level guardrail function - Pre-next-task execution timing - Tuple return format (success, data) - Automatic result/error routing - Configurable retry mechanism - Comprehensive documentation and tests Link to Devin run: https://app.devin.ai/sessions/39f6cfd6c5a24d25a7bd70ce070ed29a Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
134
tests/test_task_guardrails.py
Normal file
134
tests/test_task_guardrails.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""Tests for task guardrails functionality."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock
|
||||
|
||||
from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
from crewai.task import Task
|
||||
from crewai.tasks.task_output import TaskOutput
|
||||
|
||||
|
||||
def test_task_without_guardrail():
|
||||
"""Test that tasks work normally without guardrails (backward compatibility)."""
|
||||
agent = Mock()
|
||||
agent.role = "test_agent"
|
||||
agent.execute_task.return_value = "test result"
|
||||
agent.crew = None
|
||||
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Output"
|
||||
)
|
||||
|
||||
result = task.execute_sync(agent=agent)
|
||||
assert isinstance(result, TaskOutput)
|
||||
assert result.raw == "test result"
|
||||
|
||||
|
||||
def test_task_with_successful_guardrail():
|
||||
"""Test that successful guardrail validation passes transformed result."""
|
||||
def guardrail(result):
|
||||
return (True, result.upper())
|
||||
|
||||
agent = Mock()
|
||||
agent.role = "test_agent"
|
||||
agent.execute_task.return_value = "test result"
|
||||
agent.crew = None
|
||||
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Output",
|
||||
guardrail=guardrail
|
||||
)
|
||||
|
||||
result = task.execute_sync(agent=agent)
|
||||
assert isinstance(result, TaskOutput)
|
||||
assert result.raw == "TEST RESULT"
|
||||
|
||||
|
||||
def test_task_with_failing_guardrail():
|
||||
"""Test that failing guardrail triggers retry with error context."""
|
||||
def guardrail(result):
|
||||
return (False, "Invalid format")
|
||||
|
||||
agent = Mock()
|
||||
agent.role = "test_agent"
|
||||
agent.execute_task.side_effect = [
|
||||
"bad result",
|
||||
"good result"
|
||||
]
|
||||
agent.crew = None
|
||||
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Output",
|
||||
guardrail=guardrail,
|
||||
max_retries=1
|
||||
)
|
||||
|
||||
# First execution fails guardrail, second succeeds
|
||||
agent.execute_task.side_effect = ["bad result", "good result"]
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
task.execute_sync(agent=agent)
|
||||
|
||||
assert "Task failed guardrail validation" in str(exc_info.value)
|
||||
assert task.retry_count == 1
|
||||
|
||||
|
||||
def test_task_with_guardrail_retries():
|
||||
"""Test that guardrail respects max_retries configuration."""
|
||||
def guardrail(result):
|
||||
return (False, "Invalid format")
|
||||
|
||||
agent = Mock()
|
||||
agent.role = "test_agent"
|
||||
agent.execute_task.return_value = "bad result"
|
||||
agent.crew = None
|
||||
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Output",
|
||||
guardrail=guardrail,
|
||||
max_retries=2
|
||||
)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
task.execute_sync(agent=agent)
|
||||
|
||||
assert task.retry_count == 2
|
||||
assert "Task failed guardrail validation after 2 retries" in str(exc_info.value)
|
||||
assert "Invalid format" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_guardrail_error_in_context():
|
||||
"""Test that guardrail error is passed in context for retry."""
|
||||
def guardrail(result):
|
||||
return (False, "Expected JSON, got string")
|
||||
|
||||
agent = Mock()
|
||||
agent.role = "test_agent"
|
||||
agent.crew = None
|
||||
|
||||
task = Task(
|
||||
description="Test task",
|
||||
expected_output="Output",
|
||||
guardrail=guardrail,
|
||||
max_retries=1
|
||||
)
|
||||
|
||||
# Mock execute_task to succeed on second attempt
|
||||
first_call = True
|
||||
def execute_task(task, context, tools):
|
||||
nonlocal first_call
|
||||
if first_call:
|
||||
first_call = False
|
||||
return "invalid"
|
||||
return '{"valid": "json"}'
|
||||
|
||||
agent.execute_task.side_effect = execute_task
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
task.execute_sync(agent=agent)
|
||||
|
||||
assert "Task failed guardrail validation" in str(exc_info.value)
|
||||
assert "Expected JSON, got string" in str(exc_info.value)
|
||||
Reference in New Issue
Block a user