pydantic: Allow Dict/List field types (fix #319)

Author: irgolic

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[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)

guardrails/utils/pydantic_utils/v2.py

@@ -332,25 +332,25 @@ 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,
                     name=field_name,
                 )
-            children[field_name] = datatype_to_pydantic_field(
+            children[field_name] = pydantic_field_to_datatype(
                 ListDataType,
                 field,
                 children={"item": child},
@@ -374,7 +374,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,
@@ -390,33 +390,30 @@ def convert_pydantic_model_to_datatype(
                 name=field_name,
             )
         else:
-            children[field_name] = datatype_to_pydantic_field(
+            children[field_name] = pydantic_field_to_datatype(
                 target_datatype,
                 field,
                 strict=strict,
                 name=field_name,
             )
 
     if isinstance(model_field, FieldInfo):
-        return datatype_to_pydantic_field(
+        return pydantic_field_to_datatype(
             datatype,
             model_field,
             children=children,
             strict=strict,
             name=name,
         )
     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[DataTypeT],
     field: FieldInfo,
     children: Optional[Dict[str, "DataType"]] = None,
@@ -431,15 +428,39 @@ def datatype_to_pydantic_field(
         validators = []
     else:
         validators = field.json_schema_extra.get("validators", [])
-    format_attr = FormatAttr.from_validators(validators, datatype.tag, strict)
 
     is_optional = not field.is_required()
 
     if field.title is not None:
         name = field.title
     description = field.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[DataTypeT],
+    children: Optional[Dict[str, Any]] = None,
+    validators: Optional[list[Validator]] = None,
+    optional: bool = False,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    strict: bool = False,
+    **kwargs,
+) -> DataTypeT:
+    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)

tests/integration_tests/test_pydantic.py

@@ -1,4 +1,9 @@
+import json
+from typing import Dict, List
+
 import openai
+import pytest
+from pydantic import BaseModel
 
 import guardrails as gd
 
@@ -91,3 +96,48 @@ def test_pydantic_with_full_schema_reask(mocker):
     )
     assert guard_history[2].output == pydantic.LLM_OUTPUT_FULL_REASK_2
     assert guard_history[2].validated_output == pydantic.VALIDATED_OUTPUT_REASK_3
+
+
+class ContainerModel(BaseModel):
+    annotated_dict: Dict[str, str] = {}
+    annotated_dict_in_list: List[Dict[str, str]] = []
+    annotated_list: List[str] = []
+    annotated_list_in_dict: Dict[str, List[str]] = {}
+
+
+class ContainerModel2(BaseModel):
+    dict_: Dict = {}
+    dict_in_list: List[Dict] = []
+    list_: List = []
+    list_in_dict: Dict[str, List] = {}
+
+
+@pytest.mark.parametrize(
+    "model, output",
+    [
+        (
+            ContainerModel,
+            {
+                "annotated_dict": {"a": "b"},
+                "annotated_dict_in_list": [{"a": "b"}],
+                "annotated_list": ["a"],
+                "annotated_list_in_dict": {"a": ["b"]},
+            },
+        ),
+        (
+            ContainerModel2,
+            {
+                "dict_": {"a": "b"},
+                "dict_in_list": [{"a": "b"}],
+                "list_": ["a"],
+                "list_in_dict": {"a": ["b"]},
+            },
+        ),
+    ],
+)
+def test_container_types(model, output):
+    output_str = json.dumps(output)
+
+    guard = gd.Guard.from_pydantic(model)
+    out = guard.parse(output_str)
+    assert out == output