add back in guardrails

This commit is contained in:
Brandon Hancock
2025-01-09 14:30:39 -05:00
parent 9fb98dbf01
commit acb0d7e056

View File

@@ -1,129 +1,129 @@
# """Tests for task guardrails functionality."""
"""Tests for task guardrails functionality."""
# from unittest.mock import Mock
from unittest.mock import Mock
# import pytest
import pytest
# from crewai.task import Task
# from crewai.tasks.task_output import TaskOutput
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
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")
task = Task(description="Test task", expected_output="Output")
# result = task.execute_sync(agent=agent)
# assert isinstance(result, TaskOutput)
# assert result.raw == "test result"
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 test_task_with_successful_guardrail():
"""Test that successful guardrail validation passes transformed result."""
# def guardrail(result: TaskOutput):
# return (True, result.raw.upper())
def guardrail(result: TaskOutput):
return (True, result.raw.upper())
# agent = Mock()
# agent.role = "test_agent"
# agent.execute_task.return_value = "test result"
# agent.crew = None
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)
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"
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 test_task_with_failing_guardrail():
"""Test that failing guardrail triggers retry with error context."""
# def guardrail(result: TaskOutput):
# return (False, "Invalid format")
def guardrail(result: TaskOutput):
return (False, "Invalid format")
# agent = Mock()
# agent.role = "test_agent"
# agent.execute_task.side_effect = ["bad result", "good result"]
# agent.crew = None
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,
# )
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)
# 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
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 test_task_with_guardrail_retries():
"""Test that guardrail respects max_retries configuration."""
# def guardrail(result: TaskOutput):
# return (False, "Invalid format")
def guardrail(result: TaskOutput):
return (False, "Invalid format")
# agent = Mock()
# agent.role = "test_agent"
# agent.execute_task.return_value = "bad result"
# agent.crew = None
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,
# )
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)
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)
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 test_guardrail_error_in_context():
"""Test that guardrail error is passed in context for retry."""
# def guardrail(result: TaskOutput):
# return (False, "Expected JSON, got string")
def guardrail(result: TaskOutput):
return (False, "Expected JSON, got string")
# agent = Mock()
# agent.role = "test_agent"
# agent.crew = None
agent = Mock()
agent.role = "test_agent"
agent.crew = None
# task = Task(
# description="Test task",
# expected_output="Output",
# guardrail=guardrail,
# max_retries=1,
# )
task = Task(
description="Test task",
expected_output="Output",
guardrail=guardrail,
max_retries=1,
)
# # Mock execute_task to succeed on second attempt
# first_call = True
# 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"}'
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
agent.execute_task.side_effect = execute_task
# with pytest.raises(Exception) as exc_info:
# task.execute_sync(agent=agent)
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)
assert "Task failed guardrail validation" in str(exc_info.value)
assert "Expected JSON, got string" in str(exc_info.value)