Adjust tests for pydantic2

Author: irgolic

tests/integration_tests/mock_llm_outputs.py

@@ -78,6 +78,10 @@ def _invoke_llm(
                 python_rail.COMPILED_PROMPT_1_WITHOUT_INSTRUCTIONS,
                 python_rail.COMPILED_INSTRUCTIONS,
             ): python_rail.LLM_OUTPUT_1_FAIL_GUARDRAILS_VALIDATION,
+            (
+                python_rail.COMPILED_PROMPT_1_PYDANTIC_2_WITHOUT_INSTRUCTIONS,
+                python_rail.COMPILED_INSTRUCTIONS,
+            ): python_rail.LLM_OUTPUT_1_FAIL_GUARDRAILS_VALIDATION,
             (
                 python_rail.COMPILED_PROMPT_2_WITHOUT_INSTRUCTIONS,
                 python_rail.COMPILED_INSTRUCTIONS,

tests/integration_tests/test_assets/pydantic/validated_response_reask.py

@@ -3,6 +3,7 @@
 
 from pydantic import BaseModel, Field
 
+from guardrails.utils.pydantic_utils import PYDANTIC_VERSION
 from guardrails.utils.reask_utils import FieldReAsk
 from guardrails.validators import (
     FailResult,
@@ -59,14 +60,31 @@ class Person(BaseModel):
     """
 
     name: str
-    age: int = Field(..., validators=[AgeMustBeBetween0And150(on_fail="reask")])
-    zip_code: str = Field(
-        ...,
-        validators=[
-            ZipCodeMustBeNumeric(on_fail="reask"),
-            ZipCodeInCalifornia(on_fail="reask"),
-        ],
-    )
+    if PYDANTIC_VERSION.startswith("1"):
+        age: int = Field(..., validators=[AgeMustBeBetween0And150(on_fail="reask")])
+        zip_code: str = Field(
+            ...,
+            validators=[
+                ZipCodeMustBeNumeric(on_fail="reask"),
+                ZipCodeInCalifornia(on_fail="reask"),
+            ],
+        )
+    else:
+        age: int = Field(
+            ...,
+            json_schema_extra={
+                "validators": [AgeMustBeBetween0And150(on_fail="reask")]
+            },
+        )
+        zip_code: str = Field(
+            ...,
+            json_schema_extra={
+                "validators": [
+                    ZipCodeMustBeNumeric(on_fail="reask"),
+                    ZipCodeInCalifornia(on_fail="reask"),
+                ],
+            },
+        )
 
 
 class ListOfPeople(BaseModel):

tests/integration_tests/test_assets/python_rail/__init__.py

@@ -7,6 +7,9 @@
 )
 
 COMPILED_PROMPT_1_WITHOUT_INSTRUCTIONS = reader("compiled_prompt_1.txt")
+COMPILED_PROMPT_1_PYDANTIC_2_WITHOUT_INSTRUCTIONS = reader(
+    "compiled_prompt_1_pydantic_2.txt"
+)
 COMPILED_PROMPT_2_WITHOUT_INSTRUCTIONS = reader("compiled_prompt_2.txt")
 COMPILED_INSTRUCTIONS = reader("compiled_instructions.txt")
 

tests/integration_tests/test_assets/python_rail/compiled_prompt_1_pydantic_2.txt

@@ -0,0 +1,34 @@
+Provide detailed information about the top 5 grossing movies from Christopher Nolan including release date, duration, budget, whether it's a sequel, website, and contact email.
+
+Given below is XML that describes the information to extract from this document and the tags to extract it into.
+
+<output>
+    <string name="name" format="is-valid-director"/>
+    <list name="movies">
+        <object>
+            <integer name="rank"/>
+            <string name="title"/>
+            <object name="details">
+                <date name="release_date"/>
+                <time name="duration"/>
+                <float name="budget"/>
+                <bool name="is_sequel" required="false"/>
+                <string name="website" format="length: min=9 max=100"/>
+                <string name="contact_email"/>
+                <choice name="revenue" discriminator="revenue_type">
+                    <case name="box_office">
+                        <float name="gross"/>
+                        <float name="opening_weekend"/>
+                    </case>
+                    <case name="streaming">
+                        <integer name="subscriptions"/>
+                        <float name="subscription_fee"/>
+                    </case>
+                </choice>
+            </object>
+        </object>
+    </list>
+</output>
+
+
+ONLY return a valid JSON object (no other text is necessary), where the key of the field in JSON is the `name` attribute of the corresponding XML, and the value is of the type specified by the corresponding XML's tag. The JSON MUST conform to the XML format, including any types and format requests e.g. requests for lists, objects and specific types. Be correct and concise. If you are unsure anywhere, enter `null`.

tests/integration_tests/test_python_rail.py

@@ -4,10 +4,10 @@
 
 import openai
 import pytest
-from pydantic import BaseModel, Field, root_validator, validator
+from pydantic import BaseModel, Field
 
 import guardrails as gd
-from guardrails.utils.pydantic_utils import add_validator
+from guardrails.utils.pydantic_utils import PYDANTIC_VERSION, add_validator
 from guardrails.validators import (
     FailResult,
     PassResult,
@@ -52,7 +52,16 @@ class BoxOfficeRevenue(BaseModel):
         opening_weekend: float
 
         # Field-level validation using Pydantic (not Guardrails)
-        @validator("gross")
+        if PYDANTIC_VERSION.startswith("1"):
+            from pydantic import validator
+
+            decorator = validator("gross")
+        else:
+            from pydantic import field_validator
+
+            decorator = field_validator("gross")
+
+        @decorator
         def validate_gross(cls, gross):
             if gross <= 0:
                 raise ValueError("Gross revenue must be a positive value")
@@ -68,23 +77,47 @@ class Details(BaseModel):
         duration: time
         budget: float
         is_sequel: bool = Field(default=False)
-        website: str = Field(validators=[ValidLength(min=9, max=100, on_fail="reask")])
+
+        # Root-level validation using Pydantic (Not in Guardrails)
+        if PYDANTIC_VERSION.startswith("1"):
+            website: str = Field(
+                validators=[ValidLength(min=9, max=100, on_fail="reask")]
+            )
+            from pydantic import root_validator
+
+            @root_validator
+            def validate_budget_and_gross(cls, values):
+                budget = values.get("budget")
+                revenue = values.get("revenue")
+                if isinstance(revenue, BoxOfficeRevenue):
+                    gross = revenue.gross
+                    if budget >= gross:
+                        raise ValueError("Budget must be less than gross revenue")
+                return values
+
+        else:
+            website: str = Field(
+                json_schema_extra={
+                    "validators": [ValidLength(min=9, max=100, on_fail="reask")]
+                }
+            )
+            from pydantic import model_validator
+
+            @model_validator(mode="before")
+            def validate_budget_and_gross(cls, values):
+                budget = values.get("budget")
+                revenue = values.get("revenue")
+                if revenue["revenue_type"] == "box_office":
+                    gross = revenue["gross"]
+                    if budget >= gross:
+                        raise ValueError("Budget must be less than gross revenue")
+                return values
+
         contact_email: str
         revenue: Union[BoxOfficeRevenue, StreamingRevenue] = Field(
             ..., discriminator="revenue_type"
         )
 
-        # Root-level validation using Pydantic (Not in Guardrails)
-        @root_validator
-        def validate_budget_and_gross(cls, values):
-            budget = values.get("budget")
-            revenue = values.get("revenue")
-            if isinstance(revenue, BoxOfficeRevenue):
-                gross = revenue.gross
-                if budget >= gross:
-                    raise ValueError("Budget must be less than gross revenue")
-            return values
-
     class Movie(BaseModel):
         rank: int
         title: str
@@ -125,9 +158,15 @@ class Director(BaseModel):
     # Check that the guard state object has the correct number of re-asks.
     assert len(guard_history) == 2
 
-    assert guard_history[0].prompt == gd.Prompt(
-        python_rail.COMPILED_PROMPT_1_WITHOUT_INSTRUCTIONS
-    )
+    if PYDANTIC_VERSION.startswith("1"):
+        assert guard_history[0].prompt == gd.Prompt(
+            python_rail.COMPILED_PROMPT_1_WITHOUT_INSTRUCTIONS
+        )
+    else:
+        assert guard_history[0].prompt == gd.Prompt(
+            python_rail.COMPILED_PROMPT_1_PYDANTIC_2_WITHOUT_INSTRUCTIONS
+        )
+
     assert (
         guard_history[0].output == python_rail.LLM_OUTPUT_1_FAIL_GUARDRAILS_VALIDATION
     )
@@ -140,20 +179,32 @@ class Director(BaseModel):
         == python_rail.LLM_OUTPUT_2_SUCCEED_GUARDRAILS_BUT_FAIL_PYDANTIC_VALIDATION
     )
 
-    with pytest.raises(ValueError):
-        Director.parse_raw(
-            python_rail.LLM_OUTPUT_2_SUCCEED_GUARDRAILS_BUT_FAIL_PYDANTIC_VALIDATION
-        )
+    if PYDANTIC_VERSION.startswith("1"):
+        with pytest.raises(ValueError):
+            Director.parse_raw(
+                python_rail.LLM_OUTPUT_2_SUCCEED_GUARDRAILS_BUT_FAIL_PYDANTIC_VALIDATION
+            )
 
-    # The user can take corrective action based on the failed validation.
-    # Either manipulating the output themselves, taking corrective action
-    # in their application, or upstreaming their validations into Guardrails.
+        # The user can take corrective action based on the failed validation.
+        # Either manipulating the output themselves, taking corrective action
+        # in their application, or upstreaming their validations into Guardrails.
 
-    # The fixed output should pass validation using Pydantic
-    Director.parse_raw(python_rail.LLM_OUTPUT_3_SUCCEED_GUARDRAILS_AND_PYDANTIC)
+        # The fixed output should pass validation using Pydantic
+        Director.parse_raw(python_rail.LLM_OUTPUT_3_SUCCEED_GUARDRAILS_AND_PYDANTIC)
+    else:
+        with pytest.raises(ValueError):
+            Director.model_validate_json(
+                python_rail.LLM_OUTPUT_2_SUCCEED_GUARDRAILS_BUT_FAIL_PYDANTIC_VALIDATION
+            )
+        Director.model_validate_json(
+            python_rail.LLM_OUTPUT_3_SUCCEED_GUARDRAILS_AND_PYDANTIC
+        )
 
 
+@pytest.mark.skipif(not PYDANTIC_VERSION.startswith("1"), reason="Pydantic 1.x only")
 def test_python_rail_add_validator(mocker):
+    from pydantic import root_validator, validator
+
     mocker.patch(
         "guardrails.llm_providers.OpenAIChatCallable",
         new=MockOpenAIChatCallable,

tests/unit_tests/test_llm_providers.py

@@ -283,7 +283,7 @@ async def return_async():
 
 class ReturnTempCallable(Callable):
     def __call__(*args, **kwargs) -> Any:
-        return kwargs.get("temperature")
+        return ""
 
 
 @pytest.mark.parametrize(
@@ -295,7 +295,8 @@ def __call__(*args, **kwargs) -> Any:
 )
 def test_get_llm_ask_temperature(llm_api, args, kwargs, expected_temperature):
     result = get_llm_ask(llm_api, *args, **kwargs)
-    assert result().output == str(expected_temperature)
+    assert "temperature" in result.init_kwargs
+    assert result.init_kwargs["temperature"] == expected_temperature
 
 
 def test_chat_prompt():

tests/unit_tests/utils/test_pydantic_utils.py

@@ -1,12 +1,18 @@
+import pytest
 from pydantic import BaseModel, Field
 
 from guardrails.utils.pydantic_utils import (
+    PYDANTIC_VERSION,
     add_pydantic_validators_as_guardrails_validators,
     add_validator,
 )
 from guardrails.validators import FailResult, PassResult, ValidChoices, ValidLength
 
 
+@pytest.mark.skipif(
+    not PYDANTIC_VERSION.startswith("1"),
+    reason="Tests validators syntax for Pydantic v1",
+)
 def test_add_pydantic_validators_as_guardrails_validators():
     # TODO(shreya): Uncomment when custom validators are supported
     # def dummy_validator(name: str):
@@ -75,3 +81,33 @@ class DummyModel(BaseModel):
     # validators[3].validate(None, "Bob", None)
     # with pytest.raises(EventDetail):
     #     validators[3].validate(None, "Alex", None)
+
+
+@pytest.mark.skipif(
+    not PYDANTIC_VERSION.startswith("2"),
+    reason="Tests validators syntax for Pydantic v2",
+)
+def test_add_pydantic_validators_as_guardrails_validators_v2():
+    class DummyModel(BaseModel):
+        name: str = Field(..., validators=[ValidLength(min=1, max=10)])
+
+    model_fields = add_pydantic_validators_as_guardrails_validators(DummyModel)
+    name_field = model_fields["name"]
+
+    # Should have 1 field
+    assert len(model_fields) == 1, "Should only have one field"
+
+    # Should have 4 validators: 1 from the field, 2 from the add_validator method,
+    # and 1 from the validator decorator.
+    validators = name_field.json_schema_extra["validators"]
+    assert len(validators) == 1, "Should have 1 validator"
+
+    # The BaseModel field should not be modified
+    assert len(DummyModel.model_fields["name"].json_schema_extra["validators"]) == 1
+
+    # The first validator should be the ValidLength validator
+    assert isinstance(
+        validators[0], ValidLength
+    ), "First validator should be ValidLength"
+    assert isinstance(validators[0].validate("Beatrice", {}), PassResult)
+    assert isinstance(validators[0].validate("MrAlexander", {}), FailResult)