schemas.py

"""
Pydantic schemas for structured data extraction.

These models define the shape of data we want to extract from unstructured text.
Each model includes validation rules to ensure data quality.
"""

from datetime import date
from enum import Enum
from typing import Optional, List, Union
from pydantic import BaseModel, Field, field_validator, ConfigDict


class Priority(str, Enum):
    """Support ticket priority levels."""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    URGENT = "urgent"


class InvoiceItem(BaseModel):
    """Individual line item in an invoice."""
    description: str = Field(..., min_length=1, description="Item description")
    quantity: int = Field(..., gt=0, description="Quantity ordered")
    unit_price: float = Field(..., gt=0, description="Price per unit")
    total: float = Field(..., gt=0, description="Line item total")
    
    @field_validator('total')
    @classmethod
    def validate_total(cls, v, info):
        """Ensure total matches quantity * unit_price."""
        # Note: We can't access other fields in Pydantic v2 validators easily during validation
        # This is a simplified check - in production, use model_validator
        if v <= 0:
            raise ValueError("Total must be positive")
        return v


class InvoiceData(BaseModel):
    """Structured data extracted from invoices."""
    
    model_config = ConfigDict(str_strip_whitespace=True)
    
    invoice_number: str = Field(
        ..., 
        min_length=1,
        description="Unique invoice identifier"
    )
    invoice_date: str = Field(
        ..., 
        pattern=r"^\d{4}-\d{2}-\d{2}$",
        description="Invoice date in YYYY-MM-DD format"
    )
    due_date: str = Field(
        ..., 
        pattern=r"^\d{4}-\d{2}-\d{2}$",
        description="Payment due date in YYYY-MM-DD format"
    )
    vendor_name: str = Field(..., min_length=1, description="Vendor/seller name")
    vendor_email: Optional[str] = Field(None, description="Vendor email address")
    customer_name: str = Field(..., min_length=1, description="Customer/buyer name")
    customer_email: Optional[str] = Field(None, description="Customer email address")
    
    items: List[InvoiceItem] = Field(
        ..., 
        min_length=1,
        description="List of items/services"
    )
    
    subtotal: float = Field(..., gt=0, description="Subtotal before tax")
    tax_amount: float = Field(..., ge=0, description="Tax amount")
    total_amount: float = Field(..., gt=0, description="Total amount due")
    currency: str = Field(
        default="USD", 
        pattern=r"^[A-Z]{3}$",
        description="Currency code (ISO 4217)"
    )
    
    @field_validator('invoice_date', 'due_date')
    @classmethod
    def validate_date_format(cls, v):
        """Validate date format and ensure it's a valid date."""
        try:
            year, month, day = v.split('-')
            date(int(year), int(month), int(day))
        except (ValueError, AttributeError):
            raise ValueError(f"Invalid date format: {v}. Expected YYYY-MM-DD")
        return v


class EmailData(BaseModel):
    """Structured data extracted from emails."""
    
    model_config = ConfigDict(str_strip_whitespace=True)
    
    sender_email: str = Field(
        ..., 
        pattern=r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$",
        description="Sender's email address"
    )
    sender_name: Optional[str] = Field(None, description="Sender's name")
    recipient_email: str = Field(
        ..., 
        pattern=r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$",
        description="Primary recipient's email address"
    )
    recipient_name: Optional[str] = Field(None, description="Recipient's name")
    
    subject: str = Field(..., min_length=1, description="Email subject line")
    date_sent: str = Field(
        ..., 
        pattern=r"^\d{4}-\d{2}-\d{2}$",
        description="Date sent in YYYY-MM-DD format"
    )
    
    intent: str = Field(
        ..., 
        description="Primary intent/purpose of the email"
    )
    action_items: List[str] = Field(
        default_factory=list,
        description="List of action items or tasks mentioned"
    )
    mentioned_dates: List[str] = Field(
        default_factory=list,
        description="Important dates mentioned (YYYY-MM-DD format)"
    )
    key_entities: List[str] = Field(
        default_factory=list,
        description="Important entities mentioned (people, companies, products)"
    )


class SupportTicketData(BaseModel):
    """Structured data extracted from support tickets."""
    
    model_config = ConfigDict(str_strip_whitespace=True)
    
    ticket_id: Optional[str] = Field(
        None, 
        description="Ticket ID if present in the text"
    )
    customer_name: str = Field(..., min_length=1, description="Customer name")
    customer_email: str = Field(
        ..., 
        pattern=r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$",
        description="Customer email"
    )
    
    issue_category: str = Field(
        ..., 
        description="Category of the issue (e.g., 'billing', 'technical', 'feature request')"
    )
    priority: Priority = Field(
        default=Priority.MEDIUM,
        description="Ticket priority level"
    )
    
    summary: str = Field(
        ..., 
        min_length=10,
        max_length=200,
        description="Brief summary of the issue"
    )
    description: str = Field(
        ..., 
        min_length=1,
        description="Detailed description of the issue"
    )
    
    affected_product: Optional[str] = Field(
        None,
        description="Product or service affected"
    )
    error_codes: List[str] = Field(
        default_factory=list,
        description="Any error codes mentioned"
    )
    requested_action: Optional[str] = Field(
        None,
        description="What the customer wants us to do"
    )


# Type alias for all supported extraction types
ExtractionSchema = Union[InvoiceData, EmailData, SupportTicketData]


# Schema metadata for function calling
EXTRACTION_SCHEMAS = {
    "invoice": {
        "model": InvoiceData,
        "description": "Extract structured data from an invoice or receipt",
        "name": "extract_invoice_data"
    },
    "email": {
        "model": EmailData,
        "description": "Extract structured data from an email",
        "name": "extract_email_data"
    },
    "support_ticket": {
        "model": SupportTicketData,
        "description": "Extract structured data from a support ticket",
        "name": "extract_support_ticket_data"
    }
}