From 4226e657579f4d6ba4db212cde14cd187994e772 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 11 Dec 2024 05:32:12 +0000
Subject: [PATCH] feat: Add guardrail validation improvements

- Add result/error exclusivity validation in GuardrailResult
- Make return type annotations optional in Task guardrail validator
- Improve error messages for validation failures

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/task.py                   | 17 +++++++++++++++++
 src/crewai/tasks/guardrail_result.py | 13 ++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/crewai/task.py b/src/crewai/task.py
index 350cf84aa..8c5a95171 100644
--- a/src/crewai/task.py
+++ b/src/crewai/task.py
@@ -1,4 +1,5 @@
 import datetime
+import inspect
 import json
 import threading
 import uuid
@@ -124,6 +125,22 @@ class Task(BaseModel):
         description="Current number of retries"
     )
 
+    @field_validator("guardrail")
+    @classmethod
+    def validate_guardrail_function(cls, v: Optional[Callable]) -> Optional[Callable]:
+        """Validate that the guardrail function has the correct signature."""
+        if v is not None:
+            sig = inspect.signature(v)
+            if len(sig.parameters) != 1:
+                raise ValueError("Guardrail function must accept exactly one parameter")
+
+            # Check return annotation if present, but don't require it
+            return_annotation = sig.return_annotation
+            if return_annotation != inspect.Signature.empty:
+                if not (return_annotation == Tuple[bool, Any] or str(return_annotation) == 'Tuple[bool, Any]'):
+                    raise ValueError("If return type is annotated, it must be Tuple[bool, Any]")
+        return v
+
     _telemetry: Telemetry = PrivateAttr(default_factory=Telemetry)
     _execution_span: Optional[Span] = PrivateAttr(default=None)
     _original_description: Optional[str] = PrivateAttr(default=None)
diff --git a/src/crewai/tasks/guardrail_result.py b/src/crewai/tasks/guardrail_result.py
index e9238b429..24cfbca80 100644
--- a/src/crewai/tasks/guardrail_result.py
+++ b/src/crewai/tasks/guardrail_result.py
@@ -6,7 +6,7 @@ the way task guardrails return their validation results.
 """
 
 from typing import Any, Optional, Tuple, Union
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 
 class GuardrailResult(BaseModel):
@@ -25,6 +25,17 @@ class GuardrailResult(BaseModel):
     result: Optional[Any] = None
     error: Optional[str] = None
 
+    @field_validator("result", "error")
+    @classmethod
+    def validate_result_error_exclusivity(cls, v: Any, info) -> Any:
+        values = info.data
+        if "success" in values:
+            if values["success"] and v and "error" in values and values["error"]:
+                raise ValueError("Cannot have both result and error when success is True")
+            if not values["success"] and v and "result" in values and values["result"]:
+                raise ValueError("Cannot have both result and error when success is False")
+        return v
+
     @classmethod
     def from_tuple(cls, result: Tuple[bool, Union[Any, str]]) -> "GuardrailResult":
         """Create a GuardrailResult from a validation tuple.