fix: Improve guardrail validation and add missing type hints

- Fixed critical bug: validation now properly checks for GuardrailResult type instead of accepting any class with the same name
- Added missing TaskOutput type hints in example file
- Refactored conditional logic for better readability
- Maintains full backward compatibility

Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>

Author: claude[bot]

examples/guardrail_example_fixed.py

@@ -4,12 +4,12 @@
 This addresses the issues reported in issue #875.
 """
 
-from praisonaiagents import Agent, Task, GuardrailResult, PraisonAIAgents
+from praisonaiagents import Agent, Task, GuardrailResult, PraisonAIAgents, TaskOutput
 from typing import Tuple, Any
 import trafilatura
 
 # Example 1: Using GuardrailResult return type (now supported!)
-def validate_length_guardrailresult(output) -> GuardrailResult:
+def validate_length_guardrailresult(output: TaskOutput) -> GuardrailResult:
     """Ensure output is between 100-500 characters using GuardrailResult"""
     # Extract the raw text from the TaskOutput object
     text = output.raw if hasattr(output, 'raw') else str(output)
@@ -29,7 +29,7 @@ def validate_length_guardrailresult(output) -> GuardrailResult:
         )
 
 # Example 2: Using Tuple[bool, Any] return type (original method)
-def validate_length_tuple(output) -> Tuple[bool, Any]:
+def validate_length_tuple(output: TaskOutput) -> Tuple[bool, Any]:
     """Ensure output is between 100-500 characters using tuple"""
     text = output.raw if hasattr(output, 'raw') else str(output)
     length = len(text)

src/praisonai-agents/praisonaiagents/task/task.py

@@ -176,29 +176,26 @@ def _setup_guardrail(self):
                 from ..guardrails import GuardrailResult
 
                 # Check if it's a GuardrailResult type
-                if return_annotation is GuardrailResult or (
-                    hasattr(return_annotation, '__name__') and 
-                    return_annotation.__name__ == 'GuardrailResult'
-                ):
-                    # Valid GuardrailResult return type
-                    pass
-                else:
-                    # Check for tuple return type
-                    return_annotation_args = get_args(return_annotation)
-                    if not (
-                        get_origin(return_annotation) is tuple
-                        and len(return_annotation_args) == 2
-                        and return_annotation_args[0] is bool
-                        and (
-                            return_annotation_args[1] is Any
-                            or return_annotation_args[1] is str
-                            or return_annotation_args[1] is TaskOutput
-                            or return_annotation_args[1] == Union[str, TaskOutput]
-                        )
-                    ):
-                        raise ValueError(
-                            "If return type is annotated, it must be GuardrailResult or Tuple[bool, Any]"
-                        )
+                is_guardrail_result = return_annotation is GuardrailResult
+                
+                # Check for tuple return type
+                return_annotation_args = get_args(return_annotation)
+                is_tuple = (
+                    get_origin(return_annotation) is tuple
+                    and len(return_annotation_args) == 2
+                    and return_annotation_args[0] is bool
+                    and (
+                        return_annotation_args[1] is Any
+                        or return_annotation_args[1] is str
+                        or return_annotation_args[1] is TaskOutput
+                        or return_annotation_args[1] == Union[str, TaskOutput]
+                    )
+                )
+                
+                if not (is_guardrail_result or is_tuple):
+                    raise ValueError(
+                        "If return type is annotated, it must be GuardrailResult or Tuple[bool, Any]"
+                    )
 
             self._guardrail_fn = self.guardrail
         elif isinstance(self.guardrail, str):