Solve this a better way

Author: adrianlyjak

py/llama_cloud_services/beta/agent_data/schema.py

@@ -41,7 +41,7 @@ class Person(BaseModel):
 from llama_cloud import ExtractRun
 from llama_cloud.types.agent_data import AgentData
 from llama_cloud.types.aggregate_group import AggregateGroup
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel, Field, ValidationError, model_validator, ConfigDict
 from typing import (
     Generic,
     List,
@@ -201,9 +201,12 @@ class ExtractedFieldMetadata(BaseModel):
         description="The original text this field's value was derived from",
     )
 
+    # Forbid unknown keys to avoid swallowing nested dicts
+    model_config = ConfigDict(extra="forbid")
+
 
 ExtractedFieldMetaDataDict = Dict[
-    str, Union[Dict[str, Any], ExtractedFieldMetadata, list[Any]]
+    str, Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]
 ]
 
 
@@ -223,7 +226,7 @@ def parse_extracted_field_metadata(
 def _parse_extracted_field_metadata_recursive(
     field_value: Any,
     additional_fields: dict[str, Any] = {},
-) -> Union[Dict[str, Any], ExtractedFieldMetadata, list[Any]]:
+) -> Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]:
     """
     Parse the extracted field metadata into a dictionary of field names to field metadata.
     """
@@ -238,6 +241,8 @@ def _parse_extracted_field_metadata_recursive(
         if len(indicator_fields.intersection(field_value.keys())) > 0:
             try:
                 merged = {**field_value, **additional_fields}
+                allowed_fields = ExtractedFieldMetadata.model_fields.keys()
+                merged = {k: v for k, v in merged.items() if k in allowed_fields}
                 validated = ExtractedFieldMetadata.model_validate(merged)
 
                 # grab the citation from the array. This is just an array for backwards compatibility.
@@ -340,6 +345,28 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         description="Additional metadata about the extracted data, such as errors, tokens, etc.",
     )
 
+    @model_validator(mode="before")
+    @classmethod
+    def _normalize_field_metadata_on_input(cls, value: Any) -> Any:
+        # Ensure any inbound representation (including JSON round-trips)
+        # gets normalized so nested dicts become ExtractedFieldMetadata where appropriate.
+        if (
+            isinstance(value, dict)
+            and "field_metadata" in value
+            and isinstance(value["field_metadata"], dict)
+        ):
+            try:
+                value = {
+                    **value,
+                    "field_metadata": parse_extracted_field_metadata(
+                        value["field_metadata"]
+                    ),
+                }
+            except Exception:
+                # Let pydantic surface detailed errors later rather than swallowing completely
+                pass
+        return value
+
     @classmethod
     def create(
         cls,

py/unit_tests/beta/agent/test_agent_data_schema.py

@@ -554,7 +554,7 @@ def test_full_parse_nested_dimensions():
     with open(Path(__file__).parent.parent.parent / "data" / "capacitor.json") as f:
         data = json.load(f)
     result = ExtractedData.from_extraction_result(ExtractRun.parse_obj(data), Capacitor)
-    assert result.field_metadata == {
+    expected = {
         "dimensions": {
             "diameter": ExtractedFieldMetadata(
                 reasoning="VERBATIM EXTRACTION",
@@ -577,3 +577,6 @@ def test_full_parse_nested_dimensions():
             ),
         }
     }
+    assert result.field_metadata == expected
+    parsed = ExtractedData.model_validate_json(result.model_dump_json())
+    assert parsed.field_metadata == expected