fix tests

llama_cloud_services/beta/agent_data/__init__.py

@@ -6,6 +6,11 @@ from .schema import (
     ExtractedT,
     AgentDataT,
     ComparisonOperator,
+    parse_extracted_field_metadata,
+    calculate_overall_confidence,
+    InvalidExtractionData,
+    ExtractedFieldMetadata,
+    ExtractedFieldMetaDataDict,
 )
 from .client import AsyncAgentDataClient
 
@@ -18,4 +23,9 @@ __all__ = [
     "ExtractedT",
     "AgentDataT",
     "ComparisonOperator",
+    "parse_extracted_field_metadata",
+    "calculate_overall_confidence",
+    "InvalidExtractionData",
+    "ExtractedFieldMetadata",
+    "ExtractedFieldMetaDataDict",
 ]

llama_cloud_services/beta/agent_data/schema.py

@@ -37,9 +37,11 @@ Example Usage:
 """
 
 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
+from pydantic import BaseModel, Field, ValidationError
 from typing import (
     Generic,
     List,
@@ -174,6 +176,74 @@ class TypedAgentDataItems(BaseModel, Generic[AgentDataT]):
     )
 
 
+class ExtractedFieldMetadata(BaseModel):
+    """
+    Metadata for an extracted data field, such as confidence, and citation information.
+    """
+
+    confidence: Optional[float] = Field(
+        None,
+        description="The confidence score for the field, combined with parsing confidence if applicable",
+    )
+    extracted_confidence: Optional[float] = Field(
+        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(
+        None,
+        description="The original text this field's value was derived from",
+    )
+
+
+ExtractedFieldMetaDataDict = Dict[
+    str, Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]]
+]
+
+
+def parse_extracted_field_metadata(
+    field_metadata: dict[str, Any],
+) -> ExtractedFieldMetaDataDict:
+    """
+    Parse the extracted field metadata into a dictionary of field names to field metadata.
+    """
+    result: ExtractedFieldMetaDataDict = {}
+    for field_name, field_value in field_metadata.items():
+        if isinstance(field_value, ExtractedFieldMetadata):
+            # support running this multiple times
+            result[field_name] = field_value
+        elif isinstance(field_value, dict):
+            if "confidence" in field_value or "citations" in field_value:
+                try:
+                    validated = ExtractedFieldMetadata.model_validate(field_value)
+
+                    # grab the citation from the array. This is just an array for backwards compatibility.
+                    if "citations" in field_value and len(field_value["citations"]) > 0:
+                        first_citation = field_value["citations"][0]
+                        if "page_number" in first_citation and isinstance(
+                            first_citation["page_number"], numbers.Number
+                        ):
+                            validated.page_number = int(first_citation["page_number"])  # type: ignore
+                        if "matching_text" in first_citation and isinstance(
+                            first_citation["matching_text"], str
+                        ):
+                            validated.matching_text = first_citation["matching_text"]
+                    result[field_name] = validated
+                    continue
+                except ValidationError:
+                    pass
+            result[field_name] = parse_extracted_field_metadata(field_value)
+        elif isinstance(field_value, list):
+            result[field_name] = [
+                parse_extracted_field_metadata(item) for item in field_value
+            ]
+        else:
+            result[field_name] = field_value
+    return result
+
+
 class ExtractedData(BaseModel, Generic[ExtractedT]):
     """
     Wrapper for extracted data with workflow status tracking.
@@ -220,9 +290,13 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         description="The latest state of the data. Will differ if data has been updated"
     )
     status: StatusType = Field(description="The status of the extracted data")
-    confidence: Dict[str, Any] = Field(
+    overall_confidence: Optional[float] = Field(
+        None,
+        description="The overall confidence score for the extracted data",
+    )
+    field_metadata: ExtractedFieldMetaDataDict = Field(
         default_factory=dict,
-        description="Confidence scores, if any, for each primitive field in the original_data data",
+        description="Page links, and perhaps eventually bounding boxes, for individual fields in the extracted data. Structure is expected to have a ",
     )
     file_id: Optional[str] = Field(
         None, description="The ID of the file that was used to extract the data"
@@ -243,7 +317,7 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         cls,
         data: ExtractedT,
         status: StatusType = "pending_review",
-        confidence: Optional[Dict[str, Any]] = None,
+        field_metadata: ExtractedFieldMetaDataDict = {},
         file_id: Optional[str] = None,
         file_name: Optional[str] = None,
         file_hash: Optional[str] = None,
@@ -255,25 +329,128 @@ class ExtractedData(BaseModel, Generic[ExtractedT]):
         Args:
             extracted_data: The extracted data payload
             status: Initial workflow status
-            confidence: Optional confidence scores for fields
+            field_metadata: Optional confidence scores, citations, and other metadata for fields
             file_id: The llamacloud file ID of the file that was used to extract the data
             file_name: The name of the file that was used to extract the data
             file_hash: A content hash of the file that was used to extract the data, for de-duplication
+            metadata: Arbitrary additional application-specific data about the extracted data
 
         Returns:
             New ExtractedData instance ready for storage
         """
+        normalized_field_metadata = parse_extracted_field_metadata(field_metadata)
         return cls(
             original_data=data,
             data=data,
             status=status,
-            confidence=confidence or {},
+            field_metadata=normalized_field_metadata,
+            overall_confidence=calculate_overall_confidence(normalized_field_metadata),
             file_id=file_id,
             file_name=file_name,
             file_hash=file_hash,
             metadata=metadata or {},
         )
 
+    @classmethod
+    def from_extraction_result(
+        cls,
+        result: ExtractRun,
+        schema: Type[ExtractedT],
+        file_hash: Optional[str] = None,
+        file_name: Optional[str] = None,
+        file_id: Optional[str] = None,
+        status: StatusType = "pending_review",
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> "ExtractedData[ExtractedT]":
+        """
+        Create an ExtractedData instance from an extraction result.
+        """
+        file_id = file_id or result.file.id
+        file_name = file_name or result.file.name
+
+        try:
+            field_metadata = parse_extracted_field_metadata(
+                result.extraction_metadata.get("field_metadata", {})
+            )
+        except ValidationError:
+            field_metadata = {}
+
+        try:
+            data = schema.model_validate(result.data)  # type: ignore
+            return cls.create(
+                data=data,
+                status=status,
+                field_metadata=field_metadata,
+                file_id=file_id,
+                file_name=file_name,
+                file_hash=file_hash,
+                metadata=metadata or {},
+            )
+        except ValidationError as e:
+            invalid_item = ExtractedData[Dict[str, Any]].create(
+                data=result.data or {},
+                status="error",
+                field_metadata=field_metadata,
+                metadata={"extraction_error": str(e), **(metadata or {})},
+                file_id=file_id,
+                file_name=file_name,
+                file_hash=file_hash,
+            )
+            raise InvalidExtractionData(invalid_item) from e
+
+
+class InvalidExtractionData(Exception):
+    """
+    Exception raised when the extracted data does not conform to the schema.
+    """
+
+    def __init__(self, invalid_item: ExtractedData[Dict[str, Any]]):
+        self.invalid_item = invalid_item
+        super().__init__("Not able to parse the extracted data, parsed invalid format")
+
+
+def calculate_overall_confidence(
+    metadata: ExtractedFieldMetaDataDict,
+) -> Optional[float]:
+    """
+    Calculate the overall confidence score for the extracted data.
+    """
+    numerator, denominator = _calculate_overall_confidence_recursive(metadata)
+    if denominator == 0:
+        return None
+    return numerator / denominator
+
+
+def _calculate_overall_confidence_recursive(
+    confidence: Union[ExtractedFieldMetadata, Dict[str, Any], list[Any]],
+) -> tuple[float, int]:
+    """
+    Calculate the overall confidence score for the extracted data.
+    """
+    if isinstance(confidence, ExtractedFieldMetadata):
+        if confidence.confidence is not None:
+            return confidence.confidence, 1
+        else:
+            return 0, 0
+    if isinstance(confidence, dict):
+        numerator: float = 0
+        denominator: int = 0
+        for value in confidence.values():
+            num, den = _calculate_overall_confidence_recursive(value)
+            numerator += num
+            denominator += den
+        return numerator, denominator
+    elif isinstance(confidence, list):
+        numerator = 0
+        denominator = 0
+        for value in confidence:
+            num, den = _calculate_overall_confidence_recursive(value)
+            numerator += num
+            denominator += den
+        return numerator, denominator
+    else:
+        return 0, 0
+
 
 class TypedAggregateGroup(BaseModel, Generic[AgentDataT]):
     """

llama_parse/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "llama-parse"
-version = "0.6.52"
+version = "0.6.53"
 description = "Parse files into RAG-Optimized formats."
 authors = ["Logan Markewich <logan@llamaindex.ai>"]
 license = "MIT"
@@ -13,7 +13,7 @@ packages = [{include = "llama_parse"}]
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
-llama-cloud-services = ">=0.6.52"
+llama-cloud-services = ">=0.6.53"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

pyproject.toml

@@ -8,7 +8,7 @@ python_version = "3.10"
 
 [tool.poetry]
 name = "llama-cloud-services"
-version = "0.6.52"
+version = "0.6.53"
 description = "Tailored SDK clients for LlamaCloud services."
 authors = ["Logan Markewich <logan@runllama.ai>"]
 license = "MIT"

tests/beta/agent_data/test_agent_data_schema.py

@@ -2,14 +2,19 @@ from datetime import datetime
 from typing import Any, Dict
 
 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 llama_cloud_services.beta.agent_data.schema import (
     ExtractedData,
+    ExtractedFieldMetadata,
+    InvalidExtractionData,
     TypedAgentData,
     TypedAggregateGroup,
+    calculate_overall_confidence,
+    parse_extracted_field_metadata,
 )
 
 
@@ -71,14 +76,21 @@ def test_extracted_data_create_method():
     assert extracted.original_data == person
     assert extracted.data == person
     assert extracted.status == "pending_review"
-    assert extracted.confidence == {}
+    assert extracted.field_metadata == {}
+    assert extracted.overall_confidence is None
 
-    # Test with custom values
+    # Test with custom values using ExtractedFieldMetadata
+    field_metadata = {
+        "name": ExtractedFieldMetadata(confidence=0.99, page_number=1),
+        "age": ExtractedFieldMetadata(confidence=0.85, page_number=1),
+    }
     extracted_custom = ExtractedData.create(
-        person, status="accepted", confidence={"name": 0.99}
+        person, status="accepted", field_metadata=field_metadata
     )
     assert extracted_custom.status == "accepted"
-    assert extracted_custom.confidence["name"] == 0.99
+    assert extracted_custom.field_metadata["name"].confidence == 0.99
+    assert extracted_custom.field_metadata["age"].confidence == 0.85
+    assert extracted_custom.overall_confidence == pytest.approx((0.99 + 0.85) / 2)
 
 
 def test_extracted_data_with_dict():
@@ -107,3 +119,303 @@ def test_typed_aggregate_group_from_raw():
     assert typed_group.count == 25
     assert typed_group.first_item.name == "Tech Corp"
     assert typed_group.first_item.employees == 500
+
+
+def test_calculate_overall_confidence_simple_flat():
+    """Test calculate_overall_confidence with simple flat dictionary of ExtractedFieldMetadata."""
+    field_metadata = {
+        "name": ExtractedFieldMetadata(confidence=0.9),
+        "age": ExtractedFieldMetadata(confidence=0.8),
+        "email": ExtractedFieldMetadata(confidence=0.95),
+    }
+    result = calculate_overall_confidence(field_metadata)
+    expected = (0.9 + 0.8 + 0.95) / 3
+    assert result == pytest.approx(expected, rel=1e-9)
+
+
+def test_calculate_overall_confidence_nested():
+    """Test calculate_overall_confidence with nested dictionary structure."""
+    field_metadata = {
+        "person": {
+            "name": ExtractedFieldMetadata(confidence=0.9),
+            "age": ExtractedFieldMetadata(confidence=0.8),
+        },
+        "contact": {
+            "email": ExtractedFieldMetadata(confidence=0.95),
+            "phone": ExtractedFieldMetadata(confidence=0.85),
+        },
+        "score": ExtractedFieldMetadata(confidence=0.7),
+    }
+    result = calculate_overall_confidence(field_metadata)
+    # Should average all leaf values: (0.9 + 0.8 + 0.95 + 0.85 + 0.7) / 5
+    expected = (0.9 + 0.8 + 0.95 + 0.85 + 0.7) / 5
+    assert result == pytest.approx(expected, rel=1e-9)
+
+
+def test_calculate_overall_confidence_with_lists():
+    """Test calculate_overall_confidence with lists of ExtractedFieldMetadata and nested structures."""
+    field_metadata = {
+        "scores": [
+            ExtractedFieldMetadata(confidence=0.9),
+            ExtractedFieldMetadata(confidence=0.8),
+            ExtractedFieldMetadata(confidence=0.95),
+        ],
+        "nested_data": [
+            {
+                "field1": ExtractedFieldMetadata(confidence=0.7),
+                "field2": ExtractedFieldMetadata(confidence=0.6),
+            },
+            {
+                "field1": ExtractedFieldMetadata(confidence=0.8),
+                "field2": ExtractedFieldMetadata(confidence=0.9),
+            },
+        ],
+        "single_value": ExtractedFieldMetadata(confidence=0.85),
+    }
+    result = calculate_overall_confidence(field_metadata)
+    # Should count: [0.9, 0.8, 0.95] + [0.7, 0.6, 0.8, 0.9] + [0.85] = 8 values
+    expected = (0.9 + 0.8 + 0.95 + 0.7 + 0.6 + 0.8 + 0.9 + 0.85) / 8
+    assert result == pytest.approx(expected, rel=1e-9)
+
+
+def test_calculate_overall_confidence_invalid_types():
+    """Test calculate_overall_confidence with invalid/mixed types alongside valid ExtractedFieldMetadata."""
+    field_metadata = {
+        "valid_metadata": ExtractedFieldMetadata(confidence=0.8),
+        "valid_metadata_no_confidence": ExtractedFieldMetadata(),  # No confidence
+        "valid_list": [
+            ExtractedFieldMetadata(confidence=0.5),
+            ExtractedFieldMetadata(confidence=0.6),
+        ],
+        "invalid_string": "not_a_number",
+        "invalid_list_mixed": [
+            ExtractedFieldMetadata(confidence=0.7),
+            "invalid",
+            ExtractedFieldMetadata(confidence=0.8),
+        ],
+        "invalid_none": None,
+        "random_dict": {"a": 1, "b": 2},
+    }
+    result = calculate_overall_confidence(field_metadata)
+    expected = (0.8 + 0.5 + 0.6 + 0.7 + 0.8) / 5
+    assert result == pytest.approx(expected, rel=1e-9)
+
+
+def test_calculate_overall_confidence_empty():
+    """Test calculate_overall_confidence with empty inputs."""
+    # Empty dict
+    assert calculate_overall_confidence({}) is None
+
+    # Empty list
+    assert calculate_overall_confidence([]) is None
+
+    # Dict with only invalid values
+    field_metadata_invalid = {"invalid": "not_a_number", "also_invalid": None}
+    assert calculate_overall_confidence(field_metadata_invalid) is None
+
+    # List with only invalid values
+    field_metadata_invalid_list = ["invalid", None, {}]
+    assert calculate_overall_confidence(field_metadata_invalid_list) is None
+
+    # Dict with ExtractedFieldMetadata but no confidence values
+    field_metadata_no_confidence = {
+        "field1": ExtractedFieldMetadata(),
+        "field2": ExtractedFieldMetadata(),
+    }
+    assert calculate_overall_confidence(field_metadata_no_confidence) is None
+
+
+def test_parse_extracted_field_metadata():
+    """Test parse_extracted_field_metadata with legacy citation format."""
+    raw_metadata = {
+        "name": {
+            "confidence": 0.95,
+            "citations": [{"page_number": 1, "matching_text": "John Smith"}],
+        },
+        "age": {
+            "confidence": 0.87,
+            "citations": [
+                {
+                    "page_number": 2.0,  # Float page number
+                    "matching_text": "25 years old",
+                }
+            ],
+        },
+        "email": {
+            "confidence": 0.92,
+            "citations": [],  # Empty citations
+        },
+    }
+
+    result = parse_extracted_field_metadata(raw_metadata)
+    result2 = parse_extracted_field_metadata(result)
+    assert result2 == result
+
+    # 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"
+
+    # 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"
+
+    # email should handle empty citations
+    assert isinstance(result["email"], ExtractedFieldMetadata)
+    assert result["email"].confidence == 0.92
+
+
+def create_file(
+    id: str = "file-456",
+    name: str = "resume.pdf",
+    external_file_id: str = "external-file-id",
+    project_id: str = "project-123",
+) -> File:
+    return File.parse_obj(
+        {
+            "id": id,
+            "name": name,
+            "external_file_id": external_file_id,
+            "project_id": project_id,
+        }
+    )
+
+
+def create_extract_run(
+    id: str = "extract-123",
+    data: Dict[str, Any] = {"name": "John Doe", "age": 30, "email": "john@example.com"},
+    extraction_metadata: Dict[str, Any] = {
+        "name": {
+            "confidence": 0.95,
+            "citations": [{"page_number": 1, "matching_text": "John Doe"}],
+        },
+        "age": {"confidence": 0.87},
+        "email": {
+            "confidence": 0.92,
+            "citations": [{"page_number": 1, "matching_text": "john@example.com"}],
+        },
+    },
+    data_schema: Dict[str, Any] = {},
+    file: File = create_file(),
+) -> ExtractRun:
+    return ExtractRun.parse_obj(
+        {
+            "id": id,
+            "data": data,
+            "extraction_metadata": {
+                "field_metadata": extraction_metadata,
+            },
+            "data_schema": data_schema,
+            "file": file,
+            "extraction_agent_id": "extraction-agent-123",
+            "config": {},
+            "status": "SUCCESS",
+            "from_ui": False,
+        }
+    )
+
+
+def test_extracted_data_from_extraction_result_success():
+    """Test ExtractedData.from_extraction_result with valid data."""
+    # Create mock ExtractRun with valid data
+    extract_run = create_extract_run(
+        file=create_file(id="file-456", name="resume.pdf"),
+    )
+
+    # Create with file object
+    extracted: ExtractedData[Person] = ExtractedData.from_extraction_result(
+        extract_run,
+        Person,
+        file_hash="abc123",
+        status="accepted",
+    )
+
+    # Verify the extracted data
+    assert isinstance(extracted.data, Person)
+    assert extracted.data.name == "John Doe"
+    assert extracted.data.age == 30
+    assert extracted.data.email == "john@example.com"
+    assert extracted.status == "accepted"
+    assert extracted.file_id == "file-456"
+    assert extracted.file_name == "resume.pdf"
+    assert extracted.file_hash == "abc123"
+
+    # 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"
+
+    # Verify overall confidence was calculated
+    expected_confidence = (0.95 + 0.87 + 0.92) / 3
+    assert extracted.overall_confidence == pytest.approx(expected_confidence)
+
+
+def test_extracted_data_from_extraction_result_with_file_params():
+    """Test ExtractedData.from_extraction_result with explicit file parameters."""
+    extract_run = create_extract_run(
+        file=create_file(id="original-file", name="original.pdf"),
+    )
+
+    # Override file parameters
+    extracted: ExtractedData[Person] = ExtractedData.from_extraction_result(
+        extract_run,
+        Person,
+        file_id="custom-file-id",  # Should override file.id
+        file_name="custom-name.pdf",  # Should override file.name
+        file_hash="custom-hash",
+        metadata={"source": "api_test"},
+    )
+
+    assert extracted.file_id == "custom-file-id"  # Overridden
+    assert extracted.file_name == "custom-name.pdf"  # Overridden
+    assert extracted.file_hash == "custom-hash"
+    assert extracted.metadata["source"] == "api_test"
+
+
+def test_extracted_data_from_extraction_result_invalid_data():
+    """Test ExtractedData.from_extraction_result with invalid data raises custom exception."""
+    # Create ExtractRun with data that doesn't match Person schema
+    extract_run = create_extract_run(
+        data={
+            "name": "Valid Name",
+            "age": "not_a_number",
+            "missing_email": True,
+        },  # Invalid age, missing email
+        extraction_metadata={
+            "name": {"confidence": 0.9},
+        },
+        data_schema={},
+        file=create_file(id="error-file", name="bad_data.pdf"),
+    )
+
+    # Should raise InvalidExtractionData with ExtractedData containing error info
+    with pytest.raises(InvalidExtractionData) as exc_info:
+        ExtractedData.from_extraction_result(
+            extract_run, Person, metadata={"test": "metadata"}
+        )
+
+    # Verify the exception contains the invalid ExtractedData
+    invalid_data = exc_info.value.invalid_item
+    assert isinstance(invalid_data, ExtractedData)
+    assert invalid_data.status == "error"
+    assert invalid_data.data == {
+        "name": "Valid Name",
+        "age": "not_a_number",
+        "missing_email": True,
+    }
+    assert invalid_data.file_id == "error-file"
+    assert invalid_data.file_name == "bad_data.pdf"
+
+    # Check error metadata was added
+    assert "extraction_error" in invalid_data.metadata
+    assert "test" in invalid_data.metadata  # Original metadata preserved
+    assert "2 validation errors" in invalid_data.metadata["extraction_error"]
+
+    # Verify field metadata was still parsed (before validation failed)
+    assert isinstance(invalid_data.field_metadata["name"], ExtractedFieldMetadata)
+    assert invalid_data.field_metadata["name"].confidence == 0.9
+    assert invalid_data.overall_confidence == 0.9