Merge pull request #420 from irgolic:pydantic-v2-impl

Pydantic v2 Implementation

Author: zsimjee

.github/workflows/ci.yml

@@ -20,7 +20,6 @@ jobs:
         strategy:
             matrix:
                 python-version: ['3.8', '3.9', '3.10', '3.11']
-        
         steps:
             -   uses: actions/checkout@v2
             -   name: Set up Python ${{ matrix.python-version }}
@@ -51,7 +50,7 @@ jobs:
         strategy:
             matrix:
                 python-version: ['3.8', '3.9', '3.10', '3.11']
-
+                pydantic-version: ['1.10.9', '2.4.2']
         steps:
             -   uses: actions/checkout@v2
             -   name: Set up Python ${{ matrix.python-version }}
@@ -63,7 +62,7 @@ jobs:
                 uses: actions/cache@v3
                 with:
                     path: ~/.cache/pypoetry
-                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}
+                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ matrix.pydantic-version }}
 
             -   name: Install Poetry
                 uses: snok/install-poetry@v1
@@ -72,10 +71,17 @@ jobs:
                 # TODO: fix errors so that we can run `make dev` instead
                 run: |
                     make full
+                    poetry run pip install pydantic==${{ matrix.pydantic-version }}
+
+            -   if: matrix.pydantic-version == '2.4.2'
+                name: Static analysis with pyright (ignoring pydantic v1)
+                run: |
+                    make type-pydantic-v2
 
-            -   name: Static analysis with pyright
+            -   if: matrix.pydantic-version == '1.10.9'
+                name: Static analysis with mypy (ignoring pydantic v2)
                 run: |
-                    make type
+                    make type-pydantic-v1
 
     Pytests:
         runs-on: ubuntu-latest
@@ -85,7 +91,7 @@ jobs:
                 # TODO: fix errors so that we can run both `make dev` and `make full`
                 # dependencies: ['dev', 'full']
                 dependencies: ['full']
-
+                pydantic-version: ['1.10.9', '2.4.2']
         steps:
             -   uses: actions/checkout@v2
             -   name: Set up Python ${{ matrix.python-version }}
@@ -97,15 +103,15 @@ jobs:
                 uses: actions/cache@v3
                 with:
                     path: ~/.cache/pypoetry
-                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}
+                    key: poetry-cache-${{ runner.os }}-${{ steps.setup_python.outputs.python-version }}-${{ env.POETRY_VERSION }}-${{ matrix.pydantic-version }}
 
             -   name: Install Poetry
                 uses: snok/install-poetry@v1
 
             -   name: Install Dependencies
                 run: |
-                    python -m pip install --upgrade pip
                     make ${{ matrix.dependencies }}
+                    python -m pip install pydantic==${{ matrix.pydantic-version }}
 
             -   name: Run Pytests
                 run: |

Makefile

@@ -8,6 +8,16 @@ autoformat:
 type:
 	poetry run pyright guardrails/
 
+type-pydantic-v1:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v2.py"]}' > pyrightconfig.json
+	poetry run pyright guardrails/
+	rm pyrightconfig.json
+
+type-pydantic-v2:
+	echo '{"exclude": ["guardrails/utils/pydantic_utils/v1.py"]}' > pyrightconfig.json
+	poetry run pyright guardrails/
+	rm pyrightconfig.json
+
 lint:
 	poetry run isort -c guardrails/ tests/
 	poetry run black guardrails/ tests/ --check
@@ -44,4 +54,4 @@ all: autoformat type lint docs test
 precommit:
 	# pytest -x -q --no-summary
 	pyright guardrails/
-	make lint
+	make lint

guardrails/utils/parsing_utils.py

@@ -65,7 +65,7 @@ def get_code_block(
 
 def get_template_variables(template: str) -> List[str]:
     if hasattr(Template, "get_identifiers"):
-        return Template(template).get_identifiers()
+        return Template(template).get_identifiers()  # type: ignore
     else:
         d = collections.defaultdict(str)
         Template(template).safe_substitute(d)

guardrails/utils/pydantic_utils/__init__.py

@@ -11,7 +11,13 @@
         convert_pydantic_model_to_openai_fn,
     )
 else:
-    raise NotImplementedError(f"Pydantic version {PYDANTIC_VERSION} is not supported.")
+    from .v2 import (
+        ArbitraryModel,
+        add_pydantic_validators_as_guardrails_validators,
+        add_validator,
+        convert_pydantic_model_to_datatype,
+        convert_pydantic_model_to_openai_fn,
+    )
 
 
 __all__ = [

guardrails/utils/pydantic_utils/v1.py

@@ -311,22 +311,22 @@ def convert_pydantic_model_to_datatype(
             inner_type = get_args(type_annotation)
             if len(inner_type) == 0:
                 # If the list is empty, we cannot infer the type of the elements
-                children[field_name] = datatype_to_pydantic_field(
+                children[field_name] = pydantic_field_to_datatype(
                     ListDataType,
                     field,
                     strict=strict,
                 )
+                continue
             inner_type = inner_type[0]
             if is_pydantic_base_model(inner_type):
                 child = convert_pydantic_model_to_datatype(inner_type)
             else:
                 inner_target_datatype = field_to_datatype(inner_type)
-                child = datatype_to_pydantic_field(
+                child = construct_datatype(
                     inner_target_datatype,
-                    inner_type,
                     strict=strict,
                 )
-            children[field_name] = datatype_to_pydantic_field(
+            children[field_name] = pydantic_field_to_datatype(
                 ListDataType,
                 field,
                 children={"item": child},
@@ -349,7 +349,7 @@ def convert_pydantic_model_to_datatype(
                     strict=strict,
                     excluded_fields=[discriminator],
                 )
-            children[field_name] = datatype_to_pydantic_field(
+            children[field_name] = pydantic_field_to_datatype(
                 Choice,
                 field,
                 children=choice_children,
@@ -361,31 +361,28 @@ def convert_pydantic_model_to_datatype(
                 field, datatype=target_datatype, strict=strict
             )
         else:
-            children[field_name] = datatype_to_pydantic_field(
+            children[field_name] = pydantic_field_to_datatype(
                 target_datatype,
                 field,
                 strict=strict,
             )
 
     if isinstance(model_field, ModelField):
-        return datatype_to_pydantic_field(
+        return pydantic_field_to_datatype(
             datatype,
             model_field,
             children=children,
             strict=strict,
         )
     else:
-        format_attr = FormatAttr.from_validators([], ObjectDataType.tag, strict)
-        return datatype(
+        return construct_datatype(
+            datatype,
             children=children,
-            format_attr=format_attr,
-            optional=False,
             name=name,
-            description=None,
         )
 
 
-def datatype_to_pydantic_field(
+def pydantic_field_to_datatype(
     datatype: Type[T],
     field: ModelField,
     children: Optional[Dict[str, "DataType"]] = None,
@@ -396,14 +393,38 @@ def datatype_to_pydantic_field(
         children = {}
 
     validators = field.field_info.extra.get("validators", [])
-    format_attr = FormatAttr.from_validators(validators, datatype.tag, strict)
 
     is_optional = field.required is False
 
     name = field.name
     description = field.field_info.description
 
-    data_type = datatype(
-        children, format_attr, is_optional, name, description, **kwargs
+    return construct_datatype(
+        datatype,
+        children,
+        validators,
+        is_optional,
+        name,
+        description,
+        strict=strict,
+        **kwargs,
     )
-    return data_type
+
+
+def construct_datatype(
+    datatype: Type[T],
+    children: Optional[Dict[str, Any]] = None,
+    validators: Optional[typing.List[Validator]] = None,
+    optional: bool = False,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    strict: bool = False,
+    **kwargs,
+) -> T:
+    if children is None:
+        children = {}
+    if validators is None:
+        validators = []
+
+    format_attr = FormatAttr.from_validators(validators, datatype.tag, strict)
+    return datatype(children, format_attr, optional, name, description, **kwargs)