Merge pull request #511 from guardrails-ai/two-word-fix

zsimjee · web-flow · commit 67529d1fb851 · 2023-12-13T14:59:54.000-08:00
Better Fix Logic for TwoWord Validator
diff --git a/docs/examples/extracting_entities.ipynb b/docs/examples/extracting_entities.ipynb
diff --git a/guardrails/classes/history/call.py b/guardrails/classes/history/call.py
@@ -2,6 +2,7 @@
 
 from pydantic import Field
 from rich.panel import Panel
+from rich.pretty import pretty_repr
 from rich.tree import Tree
 
 from guardrails.classes.generic.stack import Stack
@@ -319,4 +320,23 @@ def tree(self) -> Tree:
         tree = Tree("Logs")
         for i, iteration in enumerate(self.iterations):
             tree.add(Panel(iteration.rich_group, title=f"Step {i}"))
+
+        # Replace the last Validated Output panel if we applied fixes
+        if self.failed_validations.length > 0 and self.status == pass_status:
+            previous_panels = tree.children[  # type: ignore
+                -1
+            ].label.renderable._renderables[  # type: ignore
+                :-1
+            ]
+            validated_outcome_panel = Panel(
+                pretty_repr(self.validated_output),
+                title="Validated Output",
+                style="on #F0FFF0",
+            )
+            tree.children[
+                -1
+            ].label.renderable._renderables = previous_panels + (  # type: ignore
+                validated_outcome_panel,
+            )
+
         return tree
diff --git a/guardrails/utils/reask_utils.py b/guardrails/utils/reask_utils.py
@@ -213,14 +213,15 @@ def sub_reasks_with_fixed_values(value: Any) -> Any:
     Returns:
         The value with ReAsk objects replaced with their fixed values.
     """
-    if isinstance(value, list):
-        for index, item in enumerate(value):
-            value[index] = sub_reasks_with_fixed_values(item)
-    elif isinstance(value, dict):
+    copy = deepcopy(value)
+    if isinstance(copy, list):
+        for index, item in enumerate(copy):
+            copy[index] = sub_reasks_with_fixed_values(item)
+    elif isinstance(copy, dict):
         for dict_key, dict_value in value.items():
-            value[dict_key] = sub_reasks_with_fixed_values(dict_value)
-    elif isinstance(value, FieldReAsk):
+            copy[dict_key] = sub_reasks_with_fixed_values(dict_value)
+    elif isinstance(copy, FieldReAsk):
         # TODO handle multiple fail results
-        value = value.fail_results[0].fix_value
+        copy = copy.fail_results[0].fix_value
 
-    return value
+    return copy
diff --git a/guardrails/validators/two_words.py b/guardrails/validators/two_words.py
@@ -1,5 +1,7 @@
 from typing import Any, Dict
 
+from pydash.strings import words as _words
+
 from guardrails.logger import logger
 from guardrails.validator_base import (
     FailResult,
@@ -23,13 +25,24 @@ class TwoWords(Validator):
     | Programmatic fix              | Pick the first two words.         |
     """
 
+    def _get_fix_value(self, value: str) -> str:
+        words = value.split()
+        if len(words) == 1:
+            words = _words(value)
+
+        if len(words) == 1:
+            value = f"{value} {value}"
+            words = value.split()
+
+        return " ".join(words[:2])
+
     def validate(self, value: Any, metadata: Dict) -> ValidationResult:
         logger.debug(f"Validating {value} is two words...")
 
         if len(value.split()) != 2:
             return FailResult(
                 error_message="must be exactly two words",
-                fix_value=" ".join(value.split()[:2]),
+                fix_value=self._get_fix_value(str(value)),
             )
 
         return PassResult()
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -47,6 +47,7 @@ anthropic = {version = "^0.7.2", optional = true}
 torch = {version = "^2.1.1", optional = true}
 guardrails-ai-unbabel-comet = {version = "^2.2.1", optional = true}
 huggingface_hub = {version = "^0.16.4", optional = true}
+pydash = "^7.0.6"
 
 
 [tool.poetry.extras]
diff --git a/tests/integration_tests/test_assets/entity_extraction/validated_output_fix.py b/tests/integration_tests/test_assets/entity_extraction/validated_output_fix.py
@@ -39,7 +39,7 @@
         },
         {
             "index": 7,
-            "name": "over-the-credit-limit",
+            "name": "over the",
             "explanation": "Over-the-Credit-Limit None",
             "value": 0,
         },
diff --git a/tests/integration_tests/test_assets/entity_extraction/validated_output_reask_1.py b/tests/integration_tests/test_assets/entity_extraction/validated_output_reask_1.py
@@ -56,7 +56,7 @@
                 fail_results=[
                     FailResult(
                         error_message="must be exactly two words",
-                        fix_value="over-the-credit-limit",
+                        fix_value="over the",
                     )
                 ],
                 path=["fees", 6, "name"],
diff --git a/tests/integration_tests/test_assets/entity_extraction/validated_output_reask_2.py b/tests/integration_tests/test_assets/entity_extraction/validated_output_reask_2.py
@@ -39,7 +39,7 @@
         },
         {
             "index": 7,
-            "name": "over-the-credit-limit",
+            "name": "over the",
             "explanation": "Over-the-Credit-Limit None",
             "value": 0,
         },
diff --git a/tests/integration_tests/test_assets/entity_extraction/validated_output_skeleton_reask_2.py b/tests/integration_tests/test_assets/entity_extraction/validated_output_skeleton_reask_2.py
@@ -1,15 +1,15 @@
 # flake8: noqa: E501
 VALIDATED_OUTPUT_SKELETON_REASK_2 = {
     "fees": [
-        {"name": "annual_membership_fee", "explanation": "", "value": 0.0},
-        {"name": "my_chase_plan_fee", "explanation": "", "value": 1.72},
-        {"name": "balance_transfers", "explanation": "", "value": 5.0},
-        {"name": "cash_advances", "explanation": "", "value": 5.0},
-        {"name": "foreign_transactions", "explanation": "", "value": 3.0},
-        {"name": "late_payment", "explanation": "", "value": 0.0},
-        {"name": "over-the-credit-limit", "explanation": "", "value": 0.0},
-        {"name": "return_payment", "explanation": "", "value": 0.0},
-        {"name": "return_check", "explanation": "", "value": 0.0},
+        {"name": "annual membership", "explanation": "", "value": 0.0},
+        {"name": "my chase", "explanation": "", "value": 1.72},
+        {"name": "balance transfers", "explanation": "", "value": 5.0},
+        {"name": "cash advances", "explanation": "", "value": 5.0},
+        {"name": "foreign transactions", "explanation": "", "value": 3.0},
+        {"name": "late payment", "explanation": "", "value": 0.0},
+        {"name": "over the", "explanation": "", "value": 0.0},
+        {"name": "return payment", "explanation": "", "value": 0.0},
+        {"name": "return check", "explanation": "", "value": 0.0},
     ],
     "interest_rates": {
         "purchase": {
diff --git a/tests/unit_tests/test_validators.py b/tests/unit_tests/test_validators.py
@@ -624,7 +624,7 @@ def custom_refrain_on_fail_handler(value: Any, fail_results: List[FailResult]):
                 fail_results=[
                     FailResult(
                         error_message="must be exactly two words",
-                        fix_value="dog",
+                        fix_value="dog dog",
                     )
                 ],
             ),
diff --git a/tests/unit_tests/validators/test_two_words.py b/tests/unit_tests/validators/test_two_words.py
@@ -0,0 +1,32 @@
+import pytest
+
+from guardrails.validators import FailResult, PassResult, TwoWords
+
+
+def test_two_words_happy_path():
+    validator = TwoWords()
+
+    result: PassResult = validator.validate("Hello there", {})
+
+    assert result.outcome == "pass"
+
+
+@pytest.mark.parametrize(
+    "input, expected_output",
+    [
+        ("Hello there general", "Hello there"),
+        ("hello-there-general", "hello there"),
+        ("hello_there_general", "hello there"),
+        ("helloThereGeneral", "hello There"),
+        ("HelloThereGeneral", "Hello There"),
+        ("hello.there.general", "hello there"),
+        ("hello", "hello hello"),
+    ],
+)
+def test_two_words_failures(input, expected_output):
+    validator = TwoWords()
+
+    result: FailResult = validator.validate(input, {})
+
+    assert result.outcome == "fail"
+    assert result.fix_value == expected_output

Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@`
`39`	`39`	`},`
`40`	`40`	`{`
`41`	`41`	`"index": 7,`
`42`		`- "name": "over-the-credit-limit",`
	`42`	`+ "name": "over the",`
`43`	`43`	`"explanation": "Over-the-Credit-Limit None",`
`44`	`44`	`"value": 0,`
`45`	`45`	`},`
Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@`
`56`	`56`	`fail_results=[`
`57`	`57`	`FailResult(`
`58`	`58`	`error_message="must be exactly two words",`
`59`		`- fix_value="over-the-credit-limit",`
	`59`	`+ fix_value="over the",`
`60`	`60`	`)`
`61`	`61`	`],`
`62`	`62`	`path=["fees", 6, "name"],`