stricter typing for custom on_fail methods

Author: CalebCourier

guardrails/validator_base.py

@@ -12,6 +12,7 @@
 from string import Template
 from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
 from warnings import warn
+import warnings
 
 import nltk
 import requests
@@ -26,6 +27,7 @@
 from guardrails.logger import logger
 from guardrails.remote_inference import remote_inference
 from guardrails.types.on_fail import OnFailAction
+from guardrails.utils.safe_get import safe_get
 from guardrails.utils.hub_telemetry_utils import HubTelemetry
 
 #   See: https://github.com/guardrails-ai/guardrails/issues/829
@@ -78,7 +80,7 @@ class Validator:
 
     def __init__(
         self,
-        on_fail: Optional[Union[Callable, OnFailAction]] = None,
+        on_fail: Optional[Union[Callable[[Any, FailResult], Any], OnFailAction]] = None,
         **kwargs,
     ):
         self.creds = Credentials.from_rc_file()
@@ -127,7 +129,7 @@ def __init__(
             self.on_fail_method = None
         else:
             self.on_fail_descriptor = OnFailAction.CUSTOM
-            self.on_fail_method = on_fail
+            self._set_on_fail_method(on_fail)
 
         # Store the kwargs for the validator.
         self._kwargs = kwargs
@@ -136,6 +138,31 @@ def __init__(
             self.rail_alias in validators_registry
         ), f"Validator {self.__class__.__name__} is not registered. "
 
+    def _set_on_fail_method(self, on_fail: Callable[[Any, FailResult], Any]):
+        """Set the on_fail method for the validator."""
+        on_fail_args = inspect.getfullargspec(on_fail)
+        second_arg = safe_get(on_fail_args.args, 1)
+        if second_arg is None:
+            raise ValueError(
+                "The on_fail method must take two arguments: "
+                "the value being validated and the FailResult."
+            )
+        second_arg_type = on_fail_args.annotations.get(second_arg)
+        if second_arg_type == List[FailResult]:
+            warnings.warn(
+                "Specifying a List[FailResult] as the second argument"
+                " for a custom on_fail handler is deprecated. "
+                "Please use FailResult instead.",
+                DeprecationWarning,
+            )
+
+            def on_fail_wrapper(value: Any, fail_result: FailResult) -> Any:
+                return on_fail(value, [fail_result])  # type: ignore
+
+            self.on_fail_method = on_fail_wrapper
+        else:
+            self.on_fail_method = on_fail
+
     def _validate(self, value: Any, metadata: Dict[str, Any]) -> ValidationResult:
         """User implementable function.
 

guardrails/validator_service/validator_service_base.py

@@ -95,7 +95,7 @@ def perform_correction(
         if on_fail_descriptor == OnFailAction.CUSTOM:
             if validator.on_fail_method is None:
                 raise ValueError("on_fail is 'custom' but on_fail_method is None")
-            return validator.on_fail_method(value, [result])
+            return validator.on_fail_method(value, result)
         if on_fail_descriptor == OnFailAction.REASK:
             return FieldReAsk(
                 incorrect_value=value,

tests/unit_tests/test_validator_base.py

@@ -1,4 +1,5 @@
 import json
+import re
 from typing import Any, Dict, List
 
 import pytest
@@ -209,106 +210,66 @@ def test_to_xml_attrib(min, max, expected_xml):
     assert xml_validator == expected_xml
 
 
-def custom_fix_on_fail_handler(value: Any, fail_results: List[FailResult]):
+def custom_deprecated_on_fail_handler(value: Any, fail_results: List[FailResult]):
+    return value + " deprecated"
+
+
+def custom_fix_on_fail_handler(value: Any, fail_result: FailResult):
     return value + " " + value
 
 
-def custom_reask_on_fail_handler(value: Any, fail_results: List[FailResult]):
-    return FieldReAsk(incorrect_value=value, fail_results=fail_results)
+def custom_reask_on_fail_handler(value: Any, fail_result: FailResult):
+    return FieldReAsk(incorrect_value=value, fail_results=[fail_result])
 
 
-def custom_exception_on_fail_handler(value: Any, fail_results: List[FailResult]):
+def custom_exception_on_fail_handler(value: Any, fail_result: FailResult):
     raise ValidationError("Something went wrong!")
 
 
-def custom_filter_on_fail_handler(value: Any, fail_results: List[FailResult]):
+def custom_filter_on_fail_handler(value: Any, fail_result: FailResult):
     return Filter()
 
 
-def custom_refrain_on_fail_handler(value: Any, fail_results: List[FailResult]):
+def custom_refrain_on_fail_handler(value: Any, fail_result: FailResult):
     return Refrain()
 
 
-@pytest.mark.parametrize(
-    "custom_reask_func, expected_result",
-    [
-        (
-            custom_fix_on_fail_handler,
-            {"pet_type": "dog dog", "name": "Fido"},
-        ),
-        (
-            custom_reask_on_fail_handler,
-            FieldReAsk(
-                incorrect_value="dog",
-                path=["pet_type"],
-                fail_results=[
-                    FailResult(
-                        error_message="must be exactly two words",
-                        fix_value="dog dog",
-                    )
-                ],
-            ),
-        ),
-        (
-            custom_exception_on_fail_handler,
-            ValidationError,
-        ),
-        (
-            custom_filter_on_fail_handler,
-            None,
-        ),
-        (
-            custom_refrain_on_fail_handler,
-            None,
-        ),
-    ],
-)
-# @pytest.mark.parametrize(
-#     "validator_spec",
-#     [
-#         lambda val_func: TwoWords(on_fail=val_func),
-#         # This was never supported even pre-0.5.x.
-#         # Trying this with function calling will throw.
-#         lambda val_func: ("two-words", val_func),
-#     ],
-# )
-def test_custom_on_fail_handler(
-    custom_reask_func,
-    expected_result,
-):
-    prompt = """
-        What kind of pet should I get and what should I name it?
+class TestCustomOnFailHandler:
+    def test_deprecated_on_fail_handler(self):
+        prompt = """
+            What kind of pet should I get and what should I name it?
 
-        ${gr.complete_json_suffix_v2}
-    """
+            ${gr.complete_json_suffix_v2}
+        """
 
-    output = """
-    {
-       "pet_type": "dog",
-       "name": "Fido"
-    }
-    """
+        output = """
+        {
+        "pet_type": "dog",
+        "name": "Fido"
+        }
+        """
+        expected_result = {"pet_type": "dog deprecated", "name": "Fido"}
+
+        with pytest.warns(
+            DeprecationWarning,
+            match=re.escape(  # Becuase of square brackets in the message
+                "Specifying a List[FailResult] as the second argument"
+                " for a custom on_fail handler is deprecated. "
+                "Please use FailResult instead."
+            ),
+        ):
+            validator: Validator = TwoWords(on_fail=custom_deprecated_on_fail_handler)  # type: ignore
 
-    validator: Validator = TwoWords(on_fail=custom_reask_func)
+        class Pet(BaseModel):
+            pet_type: str = Field(description="Species of pet", validators=[validator])
+            name: str = Field(description="a unique pet name")
 
-    class Pet(BaseModel):
-        pet_type: str = Field(description="Species of pet", validators=[validator])
-        name: str = Field(description="a unique pet name")
+        guard = Guard.from_pydantic(output_class=Pet, prompt=prompt)
 
-    guard = Guard.from_pydantic(output_class=Pet, prompt=prompt)
-    if isinstance(expected_result, type) and issubclass(expected_result, Exception):
-        with pytest.raises(ValidationError) as excinfo:
-            guard.parse(output, num_reasks=0)
-        assert str(excinfo.value) == "Something went wrong!"
-    else:
         response = guard.parse(output, num_reasks=0)
-        if isinstance(expected_result, FieldReAsk):
-            assert guard.history.first.iterations.first.reasks[0] == expected_result
-        else:
-            assert response.validated_output == expected_result
-
+        assert response.validation_passed is True
+        assert response.validated_output == expected_result
 
-class TestCustomOnFailHandler:
     def test_custom_fix(self):
         prompt = """
             What kind of pet should I get and what should I name it?