From acb0d7e05613484c1420c777046835647ede8d61 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Thu, 9 Jan 2025 14:30:39 -0500 Subject: [PATCH] add back in guardrails --- tests/test_task_guardrails.py | 184 +++++++++++++++++----------------- 1 file changed, 92 insertions(+), 92 deletions(-) diff --git a/tests/test_task_guardrails.py b/tests/test_task_guardrails.py index 271a5a600..e22e76234 100644 --- a/tests/test_task_guardrails.py +++ b/tests/test_task_guardrails.py @@ -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)