mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-12 09:38:31 +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:
@@ -608,6 +608,123 @@ While creating and executing tasks, certain validation mechanisms are in place t
|
||||
|
||||
These validations help in maintaining the consistency and reliability of task executions within the crewAI framework.
|
||||
|
||||
## Task Guardrails
|
||||
|
||||
Task guardrails provide a powerful way to validate, transform, or filter task outputs before they are passed to the next task. Guardrails are optional functions that execute before the next task starts, allowing you to ensure that task outputs meet specific requirements or formats.
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python Code
|
||||
from typing import Tuple, Union
|
||||
from crewai import Task
|
||||
|
||||
def validate_json_output(result: str) -> Tuple[bool, Union[dict, str]]:
|
||||
"""Validate that the output is valid JSON."""
|
||||
try:
|
||||
json_data = json.loads(result)
|
||||
return (True, json_data)
|
||||
except json.JSONDecodeError:
|
||||
return (False, "Output must be valid JSON")
|
||||
|
||||
task = Task(
|
||||
description="Generate JSON data",
|
||||
expected_output="Valid JSON object",
|
||||
guardrail=validate_json_output
|
||||
)
|
||||
```
|
||||
|
||||
### How Guardrails Work
|
||||
|
||||
1. **Optional Attribute**: Guardrails are an optional attribute at the task level, allowing you to add validation only where needed.
|
||||
2. **Execution Timing**: The guardrail function is executed before the next task starts, ensuring valid data flow between tasks.
|
||||
3. **Return Format**: Guardrails must return a tuple of `(success, data)`:
|
||||
- If `success` is `True`, `data` is the validated/transformed result
|
||||
- If `success` is `False`, `data` is the error message
|
||||
4. **Result Routing**:
|
||||
- On success (`True`), the result is automatically passed to the next task
|
||||
- On failure (`False`), the error is sent back to the agent to generate a new answer
|
||||
|
||||
### Common Use Cases
|
||||
|
||||
#### Data Format Validation
|
||||
```python Code
|
||||
def validate_email_format(result: str) -> Tuple[bool, Union[str, str]]:
|
||||
"""Ensure the output contains a valid email address."""
|
||||
import re
|
||||
email_pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$'
|
||||
if re.match(email_pattern, result.strip()):
|
||||
return (True, result.strip())
|
||||
return (False, "Output must be a valid email address")
|
||||
```
|
||||
|
||||
#### Content Filtering
|
||||
```python Code
|
||||
def filter_sensitive_info(result: str) -> Tuple[bool, Union[str, str]]:
|
||||
"""Remove or validate sensitive information."""
|
||||
sensitive_patterns = ['SSN:', 'password:', 'secret:']
|
||||
for pattern in sensitive_patterns:
|
||||
if pattern.lower() in result.lower():
|
||||
return (False, f"Output contains sensitive information ({pattern})")
|
||||
return (True, result)
|
||||
```
|
||||
|
||||
#### Data Transformation
|
||||
```python Code
|
||||
def normalize_phone_number(result: str) -> Tuple[bool, Union[str, str]]:
|
||||
"""Ensure phone numbers are in a consistent format."""
|
||||
import re
|
||||
digits = re.sub(r'\D', '', result)
|
||||
if len(digits) == 10:
|
||||
formatted = f"({digits[:3]}) {digits[3:6]}-{digits[6:]}"
|
||||
return (True, formatted)
|
||||
return (False, "Output must be a 10-digit phone number")
|
||||
```
|
||||
|
||||
### Advanced Features
|
||||
|
||||
#### Chaining Multiple Validations
|
||||
```python Code
|
||||
def chain_validations(*validators):
|
||||
"""Chain multiple validators together."""
|
||||
def combined_validator(result):
|
||||
for validator in validators:
|
||||
success, data = validator(result)
|
||||
if not success:
|
||||
return (False, data)
|
||||
result = data
|
||||
return (True, result)
|
||||
return combined_validator
|
||||
|
||||
# Usage
|
||||
task = Task(
|
||||
description="Get user contact info",
|
||||
expected_output="Email and phone",
|
||||
guardrail=chain_validations(
|
||||
validate_email_format,
|
||||
filter_sensitive_info
|
||||
)
|
||||
)
|
||||
```
|
||||
|
||||
#### Custom Retry Logic
|
||||
```python Code
|
||||
task = Task(
|
||||
description="Generate data",
|
||||
expected_output="Valid data",
|
||||
guardrail=validate_data,
|
||||
max_retries=5 # Override default retry limit
|
||||
)
|
||||
```
|
||||
|
||||
#### Async Guardrails
|
||||
```python Code
|
||||
async def validate_with_external_service(result: str) -> Tuple[bool, Union[str, str]]:
|
||||
"""Validate data using an external service."""
|
||||
# Example async validation
|
||||
validation_result = await external_service.validate(result)
|
||||
return (True, result) if validation_result.is_valid else (False, validation_result.error)
|
||||
```
|
||||
|
||||
## Creating Directories when Saving Files
|
||||
|
||||
You can now specify if a task should create directories when saving its output to a file. This is particularly useful for organizing outputs and ensuring that file paths are correctly structured.
|
||||
|
||||
@@ -6,7 +6,7 @@ from concurrent.futures import Future
|
||||
from copy import copy
|
||||
from hashlib import md5
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union
|
||||
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
|
||||
|
||||
from opentelemetry.trace import Span
|
||||
from pydantic import (
|
||||
@@ -22,6 +22,7 @@ from pydantic_core import PydanticCustomError
|
||||
from crewai.agents.agent_builder.base_agent import BaseAgent
|
||||
from crewai.tasks.output_format import OutputFormat
|
||||
from crewai.tasks.task_output import TaskOutput
|
||||
from crewai.tasks.guardrail_result import GuardrailResult
|
||||
from crewai.telemetry.telemetry import Telemetry
|
||||
from crewai.tools.base_tool import BaseTool
|
||||
from crewai.utilities.config import process_config
|
||||
@@ -110,6 +111,18 @@ class Task(BaseModel):
|
||||
default=None,
|
||||
)
|
||||
processed_by_agents: Set[str] = Field(default_factory=set)
|
||||
guardrail: Optional[Callable[[Any], Tuple[bool, Any]]] = Field(
|
||||
default=None,
|
||||
description="Function to validate task output before proceeding to next task"
|
||||
)
|
||||
max_retries: int = Field(
|
||||
default=3,
|
||||
description="Maximum number of retries when guardrail fails"
|
||||
)
|
||||
retry_count: int = Field(
|
||||
default=0,
|
||||
description="Current number of retries"
|
||||
)
|
||||
|
||||
_telemetry: Telemetry = PrivateAttr(default_factory=Telemetry)
|
||||
_execution_span: Optional[Span] = PrivateAttr(default=None)
|
||||
@@ -253,6 +266,22 @@ class Task(BaseModel):
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
# Add guardrail validation
|
||||
if self.guardrail:
|
||||
guardrail_result = GuardrailResult.from_tuple(self.guardrail(result))
|
||||
if not guardrail_result.success:
|
||||
if self.retry_count >= self.max_retries:
|
||||
raise Exception(
|
||||
f"Task failed guardrail validation after {self.max_retries} retries. "
|
||||
f"Last error: {guardrail_result.error}"
|
||||
)
|
||||
|
||||
self.retry_count += 1
|
||||
context = f"Previous attempt failed validation: {guardrail_result.error}\nPlease try again."
|
||||
return self._execute_core(agent, context, tools)
|
||||
|
||||
result = guardrail_result.result
|
||||
|
||||
pydantic_output, json_output = self._export_output(result)
|
||||
|
||||
task_output = TaskOutput(
|
||||
|
||||
44
src/crewai/tasks/guardrail_result.py
Normal file
44
src/crewai/tasks/guardrail_result.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
Module for handling task guardrail validation results.
|
||||
|
||||
This module provides the GuardrailResult class which standardizes
|
||||
the way task guardrails return their validation results.
|
||||
"""
|
||||
|
||||
from typing import Any, Optional, Tuple, Union
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class GuardrailResult(BaseModel):
|
||||
"""Result from a task guardrail execution.
|
||||
|
||||
This class standardizes the return format of task guardrails,
|
||||
converting tuple responses into a structured format that can
|
||||
be easily handled by the task execution system.
|
||||
|
||||
Attributes:
|
||||
success (bool): Whether the guardrail validation passed
|
||||
result (Any, optional): The validated/transformed result if successful
|
||||
error (str, optional): Error message if validation failed
|
||||
"""
|
||||
success: bool
|
||||
result: Optional[Any] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def from_tuple(cls, result: Tuple[bool, Union[Any, str]]) -> "GuardrailResult":
|
||||
"""Create a GuardrailResult from a validation tuple.
|
||||
|
||||
Args:
|
||||
result: A tuple of (success, data) where data is either
|
||||
the validated result or error message.
|
||||
|
||||
Returns:
|
||||
GuardrailResult: A new instance with the tuple data.
|
||||
"""
|
||||
success, data = result
|
||||
return cls(
|
||||
success=success,
|
||||
result=data if success else None,
|
||||
error=data if not success else None
|
||||
)
|
||||
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