Re-order extraction metadata union for better parsing (#865)

* Re-order args so that pydantic doesn't parse nested dict to a empty extraction result

* Use a citations array instead

Author: adrianlyjak

py/llama_cloud_services/beta/agent_data/schema.py

@@ -37,11 +37,10 @@ class Person(BaseModel):
 """
 
 from datetime import datetime
-import numbers
 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,
@@ -176,6 +175,16 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
     )
 
 
+class FieldCitation(BaseModel):
+    page: Optional[int] = Field(
+        None, description="The page number that the field occurred on"
+    )
+    matching_text: Optional[str] = Field(
+        None,
+        description="The original text this field's value was derived from",
+    )
+
+
 class ExtractedFieldMetadata(BaseModel):
     """
     Metadata for an extracted data field, such as confidence, and citation information.
@@ -193,14 +202,14 @@ class ExtractedFieldMetadata(BaseModel):
         None,
         description="The confidence score for the field based on the extracted text only",
     )
-    page_number: Optional[int] = Field(
-        None, description="The page number that the field occurred on"
-    )
-    matching_text: Optional[str] = Field(
+    citation: Optional[List[FieldCitation]] = Field(
         None,
-        description="The original text this field's value was derived from",
+        description="The citation for the field, including page number and matching text",
     )
 
+    # Forbid unknown keys to avoid swallowing nested dicts
+    model_config = ConfigDict(extra="forbid")
+
 
 ExtractedFieldMetaDataDict = Dict[
     str, Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]
@@ -238,19 +247,10 @@ 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.
-                if "citation" in field_value and len(field_value["citation"]) > 0:
-                    first_citation = field_value["citation"][0]
-                    if "page" in first_citation and isinstance(
-                        first_citation["page"], numbers.Number
-                    ):
-                        validated.page_number = int(first_citation["page"])  # type: ignore
-                    if "matching_text" in first_citation and isinstance(
-                        first_citation["matching_text"], str
-                    ):
-                        validated.matching_text = first_citation["matching_text"]
                 return validated
             except ValidationError:
                 pass
@@ -340,6 +340,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/llama_parse/pyproject.toml

@@ -11,13 +11,13 @@ dev = [
 
 [project]
 name = "llama-parse"
-version = "0.6.57"
+version = "0.6.58"
 description = "Parse files into RAG-Optimized formats."
 authors = [{name = "Logan Markewich", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
-dependencies = ["llama-cloud-services>=0.6.56"]
+dependencies = ["llama-cloud-services>=0.6.58"]
 
 [project.scripts]
 llama-parse = "llama_parse.cli.main:parse"

py/pyproject.toml

@@ -18,7 +18,7 @@ dev = [
 
 [project]
 name = "llama-cloud-services"
-version = "0.6.57"
+version = "0.6.58"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = [{name = "Logan Markewich", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"

py/unit_tests/beta/agent/test_agent_data_schema.py

@@ -1,15 +1,18 @@
 from datetime import datetime
-from typing import Any, Dict
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional
 
 import pytest
 from llama_cloud import ExtractRun, File
 from llama_cloud.types.agent_data import AgentData
 from llama_cloud.types.aggregate_group import AggregateGroup
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, Field, ValidationError
 
 from llama_cloud_services.beta.agent_data.schema import (
     ExtractedData,
     ExtractedFieldMetadata,
+    FieldCitation,
     InvalidExtractionData,
     TypedAgentData,
     TypedAggregateGroup,
@@ -81,8 +84,12 @@ def test_extracted_data_create_method():
 
     # Test with custom values using ExtractedFieldMetadata
     field_metadata = {
-        "name": ExtractedFieldMetadata(confidence=0.99, page_number=1),
-        "age": ExtractedFieldMetadata(confidence=0.85, page_number=1),
+        "name": ExtractedFieldMetadata(
+            confidence=0.99, citation=[FieldCitation(page=1)]
+        ),
+        "age": ExtractedFieldMetadata(
+            confidence=0.85, citation=[FieldCitation(page=1)]
+        ),
     }
     extracted_custom = ExtractedData.create(
         person, status="accepted", field_metadata=field_metadata
@@ -254,14 +261,16 @@ def test_parse_extracted_field_metadata():
     # name should have parsed citation data
     assert isinstance(result["name"], ExtractedFieldMetadata)
     assert result["name"].confidence == 0.95
-    assert result["name"].page_number == 1
-    assert result["name"].matching_text == "John Smith"
+    assert result["name"].citation == [
+        FieldCitation(page=1, matching_text="John Smith")
+    ]
 
     # age should handle float page number
     assert isinstance(result["age"], ExtractedFieldMetadata)
     assert result["age"].confidence == 0.87
-    assert result["age"].page_number == 2  # Should be converted to int
-    assert result["age"].matching_text == "25 years old"
+    assert result["age"].citation == [
+        FieldCitation(page=2, matching_text="25 years old")
+    ]
 
     # email should handle empty citations
     assert isinstance(result["email"], ExtractedFieldMetadata)
@@ -327,46 +336,52 @@ def test_parse_extracted_field_metadata_complex():
             reasoning="Combined key parametrics and construction from the datasheet for a structured title.",
             confidence=0.9470628580889779,
             extraction_confidence=0.9470628580889779,
-            page_number=1,
-            matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
+            citation=[
+                FieldCitation(
+                    page=1,
+                    matching_text="PHE844/F844, Film, Metallized Polypropylene, Safety, 0.47 uF",
+                )
+            ],
         ),
         "manufacturer": ExtractedFieldMetadata(
             reasoning="VERBATIM EXTRACTION",
             confidence=0.9997446550976602,
             extraction_confidence=0.9997446550976602,
-            page_number=1,
-            matching_text="YAGEO KEMET",
+            citation=[FieldCitation(page=1, matching_text="YAGEO KEMET")],
         ),
         "features": [
             ExtractedFieldMetadata(
                 reasoning="VERBATIM EXTRACTION",
                 confidence=0.9999308195540074,
                 extraction_confidence=0.9999308195540074,
-                page_number=1,
-                matching_text="Features</td><td>EMI Safety",
+                citation=[
+                    FieldCitation(
+                        page=1,
+                        matching_text="Features</td><td>EMI Safety",
+                    )
+                ],
             ),
             ExtractedFieldMetadata(
                 reasoning="VERBATIM EXTRACTION",
                 confidence=0.8642493886452225,
                 extraction_confidence=0.8642493886452225,
-                page_number=1,
-                matching_text="THB Performance</td><td>Yes",
+                citation=[
+                    FieldCitation(page=1, matching_text="THB Performance</td><td>Yes")
+                ],
             ),
         ],
         "dimensions": {
             "length": ExtractedFieldMetadata(
                 reasoning="VERBATIM EXTRACTION",
                 confidence=0.8986941382802304,
                 extraction_confidence=0.8986941382802304,
-                page_number=1,
-                matching_text="L</td><td>41mm MAX",
+                citation=[FieldCitation(page=1, matching_text="L</td><td>41mm MAX")],
             ),
             "width": ExtractedFieldMetadata(
                 reasoning="VERBATIM EXTRACTION",
                 confidence=0.9999377974447091,
                 extraction_confidence=0.9999377974447091,
-                page_number=1,
-                matching_text="T</td><td>13mm MAX",
+                citation=[FieldCitation(page=1, matching_text="T</td><td>13mm MAX")],
             ),
         },
     }
@@ -450,8 +465,9 @@ def test_extracted_data_from_extraction_result_success():
     # Verify field metadata was parsed
     assert isinstance(extracted.field_metadata["name"], ExtractedFieldMetadata)
     assert extracted.field_metadata["name"].confidence == 0.95
-    assert extracted.field_metadata["name"].page_number == 1
-    assert extracted.field_metadata["name"].matching_text == "John Doe"
+    assert extracted.field_metadata["name"].citation == [
+        FieldCitation(page=1, matching_text="John Doe")
+    ]
 
     # Verify overall confidence was calculated
     expected_confidence = (0.95 + 0.87 + 0.92) / 3
@@ -523,3 +539,54 @@ def test_extracted_data_from_extraction_result_invalid_data():
     assert isinstance(invalid_data.field_metadata["name"], ExtractedFieldMetadata)
     assert invalid_data.field_metadata["name"].confidence == 0.9
     assert invalid_data.overall_confidence == 0.9
+
+
+class Dimensions(BaseModel):
+    length: Optional[str] = Field(
+        None, description="Length in mm (Size, Longest Side, L)"
+    )
+    width: Optional[str] = Field(
+        None, description="Width in mm (Breadth, Side Width, W)"
+    )
+    height: Optional[str] = Field(
+        None, description="Height in mm (Thickness, Vertical Size, H)"
+    )
+    diameter: Optional[str] = Field(
+        None,
+        description="Diameter in mm (for radial or cylindrical types) (Outer Diameter, dt, OD, D, d<sub>t</sub>)",
+    )
+    lead_spacing: Optional[str] = Field(
+        None, description="Lead spacing in mm (Pin Pitch, Terminal Gap, LS)"
+    )
+
+
+class Capacitor(BaseModel):
+    dimensions: Optional[Dimensions] = None
+
+
+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)
+    expected = {
+        "dimensions": {
+            "diameter": ExtractedFieldMetadata(
+                reasoning="VERBATIM EXTRACTION",
+                confidence=1.0,
+                extraction_confidence=1.0,
+            ),
+            "lead_spacing": ExtractedFieldMetadata(
+                reasoning="VERBATIM EXTRACTION",
+                confidence=0.9999999031936799,
+                extraction_confidence=0.9999999031936799,
+            ),
+            "length": ExtractedFieldMetadata(
+                reasoning="VERBATIM EXTRACTION",
+                confidence=0.9999968039036192,
+                extraction_confidence=0.9999968039036192,
+            ),
+        }
+    }
+    assert result.field_metadata == expected
+    parsed = ExtractedData.model_validate_json(result.model_dump_json())
+    assert parsed.field_metadata == expected