from openai import OpenAI
from pydantic import BaseModel
from typing import Optional

client = OpenAI()

class ExtractedCompany(BaseModel):
    name: str
    industry: str
    founded_year: Optional[int]
    headquarters: str
    employee_count_range: str  # e.g., "100-500", "1000+"
    is_public: bool
    notable_products: list[str]

def extract_company_info(text: str) -> ExtractedCompany:
    response = client.beta.chat.completions.parse(
        model="gpt-4o-2024-11-20",
        messages=[
            {
                "role": "system",
                "content": "Extract company information from the provided text. "
                          "Be precise; use null for fields where information is not present."
            },
            {
                "role": "user",
                "content": text
            }
        ],
        response_format=ExtractedCompany
    )
    
    # If structured outputs worked, this is guaranteed valid
    return response.choices[0].message.parsed

# Usage
company = extract_company_info("""
Stripe was founded in 2010 by Patrick and John Collison. 
The fintech company is headquartered in South San Francisco 
and has over 7,000 employees. Products include Stripe Payments, 
Stripe Connect, and Stripe Billing.
""")

print(company.name)              # "Stripe"
print(company.founded_year)      # 2010
print(company.is_public)         # False
print(company.notable_products)  # ["Stripe Payments", "Stripe Connect", "Stripe Billing"]

import anthropic
from pydantic import BaseModel
import json

client = anthropic.Anthropic()

class SentimentAnalysis(BaseModel):
    sentiment: str  # positive, negative, neutral
    confidence: float  # 0.0 to 1.0
    key_phrases: list[str]
    summary: str

def analyze_sentiment(text: str) -> SentimentAnalysis:
    tool_schema = SentimentAnalysis.model_json_schema()
    
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=1024,
        tools=[{
            "name": "report_sentiment",
            "description": "Report the sentiment analysis results",
            "input_schema": tool_schema
        }],
        tool_choice={"type": "tool", "name": "report_sentiment"},
        messages=[{
            "role": "user",
            "content": f"Analyze the sentiment of this text:\n\n{text}"
        }]
    )
    
    # Extract tool use result
    tool_use = next(
        block for block in response.content 
        if block.type == "tool_use"
    )
    
    return SentimentAnalysis(**tool_use.input)

from abc import ABC, abstractmethod
from typing import TypeVar, Type
from pydantic import BaseModel
import anthropic
import openai

T = TypeVar('T', bound=BaseModel)

class StructuredLLM(ABC):
    @abstractmethod
    def extract(self, prompt: str, schema: Type[T], system: str = "") -> T:
        pass

class OpenAIStructured(StructuredLLM):
    def __init__(self, model: str = "gpt-4o-2024-11-20"):
        self.client = openai.OpenAI()
        self.model = model
    
    def extract(self, prompt: str, schema: Type[T], system: str = "") -> T:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        response = self.client.beta.chat.completions.parse(
            model=self.model,
            messages=messages,
            response_format=schema
        )
        return response.choices[0].message.parsed

class AnthropicStructured(StructuredLLM):
    def __init__(self, model: str = "claude-sonnet-4-5"):
        self.client = anthropic.Anthropic()
        self.model = model
    
    def extract(self, prompt: str, schema: Type[T], system: str = "") -> T:
        tool_name = f"extract_{schema.__name__.lower()}"
        
        kwargs = {
            "model": self.model,
            "max_tokens": 2048,
            "tools": [{
                "name": tool_name,
                "description": f"Extract and return {schema.__name__} data",
                "input_schema": schema.model_json_schema()
            }],
            "tool_choice": {"type": "tool", "name": tool_name},
            "messages": [{"role": "user", "content": prompt}]
        }
        if system:
            kwargs["system"] = system
        
        response = self.client.messages.create(**kwargs)
        tool_use = next(b for b in response.content if b.type == "tool_use")
        return schema(**tool_use.input)

# Usage — same code regardless of provider
def process_document(text: str, llm: StructuredLLM) -> DocumentSummary:
    return llm.extract(
        prompt=f"Summarize this document:\n\n{text}",
        schema=DocumentSummary,
        system="You are a precise document analyzer."
    )

# Swap providers trivially
llm = OpenAIStructured()  # or AnthropicStructured()
summary = process_document(document_text, llm)

from enum import Enum

class Priority(str, Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

class IssueClassification(BaseModel):
    priority: Priority  # LLM can only return valid enum values
    category: str
    affected_users: int

from typing import Optional

class PersonInfo(BaseModel):
    name: str                      # Required — fail if not present
    email: Optional[str] = None    # Optional — null if not found
    age: Optional[int] = None      # Optional — null if not mentioned
    company: Optional[str] = None

class Address(BaseModel):
    street: str
    city: str
    country: str
    postal_code: str

class ContactInfo(BaseModel):
    name: str
    email: str
    phone: Optional[str] = None
    address: Optional[Address] = None  # Nested object

class BusinessCard(BaseModel):
    contacts: list[ContactInfo]    # Array of nested objects
    notes: str

from pydantic import Field

class FinancialMetrics(BaseModel):
    revenue: float = Field(description="Annual revenue in USD millions")
    growth_rate: float = Field(description="Year-over-year growth rate as percentage, e.g., 23.5 for 23.5%")
    profit_margin: float = Field(description="Net profit margin as percentage")
    fiscal_year: int = Field(description="Fiscal year this data refers to, e.g., 2025")

import time
from typing import Optional

class StructuredExtractionError(Exception):
    pass

def extract_with_retry(
    llm: StructuredLLM,
    prompt: str,
    schema: Type[T],
    max_attempts: int = 3,
    fallback_value: Optional[T] = None
) -> T:
    last_error = None
    
    for attempt in range(max_attempts):
        try:
            result = llm.extract(prompt, schema)
            
            # Additional validation beyond schema
            validate_business_rules(result)
            
            return result
            
        except ValidationError as e:
            # Pydantic validation failed — schema conformant but invalid values
            last_error = e
            if attempt < max_attempts - 1:
                # Add error context to prompt and retry
                error_context = f"\n\nPrevious attempt failed validation: {str(e)}\nPlease ensure values meet the constraints."
                prompt = prompt + error_context
                time.sleep(0.5 * (attempt + 1))  # exponential backoff
        
        except Exception as e:
            last_error = e
            if attempt < max_attempts - 1:
                time.sleep(1.0 * (attempt + 1))
    
    if fallback_value is not None:
        return fallback_value
    
    raise StructuredExtractionError(
        f"Failed after {max_attempts} attempts"
    ) from last_error

def validate_business_rules(result: BaseModel):
    """Additional validation beyond Pydantic schema."""
    if hasattr(result, 'confidence') and not (0 <= result.confidence <= 1):
        raise ValidationError("confidence must be between 0 and 1")
    if hasattr(result, 'email') and result.email and '@' not in result.email:
        raise ValidationError(f"Invalid email format: {result.email}")

# 1. Keep schemas lean — only extract what you need
class MinimalExtraction(BaseModel):
    key_fact: str       # Don't add 10 optional fields "just in case"
    confidence: float

# 2. Batch when possible
def batch_extract(texts: list[str], schema: Type[T]) -> list[T]:
    # Some providers support batch APIs
    # Or use asyncio for parallel requests
    import asyncio
    
    async def extract_one(text: str) -> T:
        return await async_llm.extract(text, schema)
    
    return asyncio.run(
        asyncio.gather(*[extract_one(t) for t in texts])
    )

# 3. Cache schema compilation
from functools import lru_cache

@lru_cache(maxsize=128)
def get_compiled_schema(model_class: Type[BaseModel]) -> dict:
    return model_class.model_json_schema()