The ADE Storybook: Part 2

# Define your sorting rules
split_classes = [
    {
        "name": "Bank Statement",
        "description": "A monthly financial summary."
    },
    {
        "name": "Pay Stub",
        "description": "A document detailing an employee's earnings.",
        "identifier": "Payment Date" # Optional. Helps if you have multiple stubs
    }
]

import json
from pathlib import Path
from landingai_ade import LandingAIADE

client = LandingAIADE()

# Revisit Part 1 for document parsing basics. We need the parsed output for splitting.
response = client.split(
    markdown=Path("/path/to/parsed_mixed_document.md"),    split_class= json.dumps(split_classes),
        save_to="output_folder"  # optional: saves as {input_file}_split_output.json
)

import json
from pydantic import BaseModel, Field
from landingai_ade.lib import pydantic_to_json_schema

# 1. Define the Reference List (Schema) for the document we care about
class PaystubSchema(BaseModel):
    employee_name: str = Field(..., description="The name of the employee receiving the pay")
    net_pay: float = Field(..., description="The final amount paid after taxes")
    pay_period_end: str = Field(..., description="The date the pay period ended")

# 2. Group the documents by identifier first
# This handles the edge case where one document is fragmented across multiple splits
grouped_paystubs = {}

for split in split_response.get("splits", []):
    # We only care about Paystubs right now
    if split.get("classification") == "Pay Stub":
        doc_id = split.get("identifier")
        
        # Only group if an identifier exists
        if doc_id:
            if doc_id not in grouped_paystubs:
                grouped_paystubs[doc_id] = []
            
            # Collect the markdowns for this specific identifier
            grouped_paystubs[doc_id].extend(split.get("markdowns", []))

# 3. Loop through the grouped documents
for doc_id, markdowns in grouped_paystubs.items():
    
    # We merge all pages for this identifier to provide the full context for extraction
    merged_markdown = "\n".join(markdowns)
    
    # Run the extraction using a unique filename per identifier to prevent overwriting
    extract_response = client.extract(
        schema=pydantic_to_json_schema(PaystubSchema),
        markdown=merged_markdown,
        save_to=f"output_folder/{doc_id}_extract_output.json" 
    )
    
    print(f"Results for Paystub ID {doc_id}:")
    print(extract_response.extraction)

# Now we have taken out all the relevant information we want for all pay stubs in this messy mixed document.