Free lesson · AI Engineering

Type Hints and Pydantic

Easy10 min readaillmrag

Type Hints and Pydantic

The Bug That Type Hints Would Have Caught in 30 Seconds

At Meta, our AI pipeline had a function called score_results that took a list of search results and returned a list of scored results. A junior engineer called it with a single result (a dictionary) instead of a list of results (a list containing dictionaries). Python happily accepted the dictionary, iterated over its keys (strings), and produced garbage output without raising any error. The pipeline ran for six hours before anyone noticed that the scoring was wrong, processing 14,000 queries with incorrect scores.

The fix was one line - wrapping the single result in a list. The diagnosis took 45 minutes because the function signature was def score_results(results) with no type information. The engineer had to read the function body to understand that results should be a list of dictionaries, not a single dictionary. With a type hint - def score_results(results: list[dict]) -> list[dict] - the engineer's IDE would have underlined the call in red, the CI type checker (mypy) would have blocked the merge, and the bug would have been caught in 30 seconds instead of 6 hours.

Python's type hint system adds optional static type checking to a dynamically typed language. It does not change how Python runs - type hints are ignored at runtime. But it changes how Python is developed: IDEs provide better autocompletion, type checkers catch bugs before deployment, and function signatures document themselves. For AI applications where data flows through complex pipelines with many stages, type hints are the difference between "read the function body to understand the input format" and "read the function signature."

Python Type Hints: The Essentials

Type hints annotate function arguments and return values with their expected types. They are syntactic annotations - Python ignores them at runtime, but tools like mypy, pyright, and IDE type checkers use them to catch type errors before the code runs.

# Basic type hints
def count_tokens(text: str) -> int:
    """Count approximate tokens in text."""
    return len(text.split()) * 4 // 3  # Rough token estimate

def compute_cost(input_tokens: int, output_tokens: int, 
                 price_per_million: float) -> float:
    """Compute inference cost in dollars."""
    total_tokens = input_tokens + output_tokens
    return (total_tokens / 1_000_000) * price_per_million

# The type hints document the function's contract
result = count_tokens("What is retrieval augmented generation?")
cost = compute_cost(input_tokens=150, output_tokens=300, price_per_million=10.0)
print(f"Tokens: {result}, Cost: ${cost:.6f}")

Collection Types

from typing import Optional

# List of specific types
def filter_by_score(results: list[dict], threshold: float) -> list[dict]:
    """Filter results above a score threshold."""
    return [r for r in results if r.get("score", 0) >= threshold]

# Dictionary with typed keys and values
def build_metadata(model: str, tokens: int) -> dict[str, str | int | float]:
    """Build a metadata dictionary."""
    return {"model": model, "tokens": tokens, "timestamp": 1234567890}

# Optional - value that might be None
def find_best_result(results: list[dict]) -> Optional[dict]:
    """Find the highest-scoring result, or None if empty."""
    if not results:
        return None
    return max(results, key=lambda r: r.get("score", 0))

# Union types (Python 3.10+ syntax with |)
def process_input(data: str | list[str]) -> list[str]:
    """Accept either a single string or a list of strings."""
    if isinstance(data, str):
        return [data]
    return data

# Test
results = [{"text": "A", "score": 0.9}, {"text": "B", "score": 0.6}]
filtered = filter_by_score(results, 0.7)
best = find_best_result(results)
print(f"Filtered: {[r['text'] for r in filtered]}")
print(f"Best: {best['text'] if best else 'None'}")

TypedDict - Type Hints for Dictionaries

When dictionaries have a known structure (specific keys with specific value types), TypedDict adds type checking to dictionary access - catching the typo bugs that regular dict type hints miss. TypedDict is a static typing construct - it provides IDE support and mypy checking but does not validate data at runtime. For runtime validation, use Pydantic (covered below). The rule of thumb: use TypedDict for dictionaries you create internally (where you control the data), use Pydantic for dictionaries that arrive from external sources (where the data might be malformed).

Running type checks with mypy: Install mypy (pip install mypy) and run mypy your_file.py to check for type errors. In CI/CD pipelines, mypy runs on every pull request, catching type-related bugs before they reach production. At Meta, adding mypy to our AI pipeline's CI pipeline caught an average of 3-4 type errors per pull request during the first month - errors that would otherwise have reached production and caused runtime failures.

from typing import TypedDict

class SearchResult(TypedDict):
    text: str
    score: float
    document_id: str

class APIResponse(TypedDict):
    results: list[SearchResult]
    total_count: int
    latency_ms: float

# Usage - IDE knows exactly what keys exist
def format_results(response: APIResponse) -> str:
    """Format API response for display."""
    lines = []
    for result in response["results"]:
        # IDE autocompletion works here - knows result has "text", "score", "document_id"
        lines.append(f"  [{result['score']:.2f}] {result['text']}")
    return f"Found {response['total_count']} results:\n" + "\n".join(lines)

# Test
response: APIResponse = {
    "results": [
        {"text": "RAG combines retrieval...", "score": 0.94, "document_id": "doc_1"},
        {"text": "Vector databases store...", "score": 0.87, "document_id": "doc_2"},
    ],
    "total_count": 2,
    "latency_ms": 23.4
}
print(format_results(response))

Pydantic: Runtime Data Validation

Type hints are checked by tools but ignored at runtime. Pydantic bridges this gap - it validates data at runtime using the same type annotation syntax. When data does not match the declared types, Pydantic raises a clear validation error rather than allowing invalid data to flow through the pipeline.

Why Pydantic matters for AI engineering: AI applications receive data from many sources - user input, API responses, database queries, file uploads, LLM outputs. Each source can produce malformed data: missing fields, wrong types, values outside expected ranges, unexpected nulls. Without validation, this malformed data flows silently through the pipeline, producing incorrect results or crashing at an unrelated point downstream where the error is hard to diagnose. Pydantic validates the data at the boundary where it enters your application, catching errors immediately with clear messages that name the exact field and the exact violation.

Pydantic v1 vs v2: The AI engineering ecosystem went through a major migration from Pydantic v1 to v2 in 2023-2024. The key changes: validator became field_validator, Config class became model_config, and .dict() became .model_dump(). Most current tutorials and libraries use Pydantic v2. If you encounter v1 syntax in older codebases, the Pydantic migration guide covers every change. This course uses Pydantic v2 syntax exclusively.

At Meta, adopting Pydantic for all external data boundaries reduced pipeline-related bugs by roughly 40% in the first quarter after adoption.

Pydantic Models

A Pydantic model is a class that inherits from BaseModel. Fields are defined with type annotations. When you create an instance, Pydantic validates every field against its declared type and raises ValidationError if validation fails.

from pydantic import BaseModel, Field
from typing import Optional
from enum import Enum

class ModelName(str, Enum):
    GPT4_TURBO = "gpt-4-turbo"
    GPT4O = "gpt-4o"
    GPT4O_MINI = "gpt-4o-mini"
    CLAUDE_SONNET = "claude-3-5-sonnet"

class LLMRequest(BaseModel):
    """Validated LLM API request."""
    model: ModelName
    prompt: str = Field(min_length=1, max_length=100000)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=500, ge=1, le=4096)
    system_prompt: Optional[str] = None

# Valid request - passes validation
request = LLMRequest(
    model="gpt-4-turbo",
    prompt="What is RAG?",
    temperature=0.3,
    max_tokens=300
)
print(f"Valid request: model={request.model.value}, temp={request.temperature}")

# Pydantic automatically converts "gpt-4-turbo" string to ModelName enum
print(f"Model type: {type(request.model)}")

# Invalid request - caught at creation time
try:
    bad_request = LLMRequest(
        model="gpt-4-turbo",
        prompt="",  # Empty string violates min_length=1
        temperature=3.0  # Exceeds le=2.0
    )
except Exception as e:
    print(f"Validation error: {e}")

Field Validation and Constraints

Pydantic's Field function adds constraints beyond basic type checking: minimum and maximum values for numbers, minimum and maximum lengths for strings, regex patterns for format validation, and custom validators for complex business rules.

from pydantic import BaseModel, Field, field_validator

class SearchQuery(BaseModel):
    """Validated search query with business rules."""
    query: str = Field(min_length=3, max_length=500)
    top_k: int = Field(default=5, ge=1, le=100)
    score_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
    filters: dict[str, str] = Field(default_factory=dict)
    
    @field_validator("query")
    @classmethod
    def query_must_not_be_empty_whitespace(cls, v: str) -> str:
        """Custom validator: query must contain non-whitespace characters."""
        if not v.strip():
            raise ValueError("Query must contain non-whitespace characters")
        return v.strip()  # Also normalise whitespace

# Valid query
query = SearchQuery(query="  What is RAG?  ", top_k=10, score_threshold=0.8)
print(f"Query: '{query.query}', top_k={query.top_k}")
# Note: whitespace was stripped by the validator

# Invalid query
try:
    bad_query = SearchQuery(query="   ", top_k=10)  # Only whitespace
except Exception as e:
    print(f"Validation error: {e}")

Pydantic for LLM Output Parsing

One of the most valuable Pydantic patterns in AI engineering is validating LLM output. When you ask an LLM to return JSON, the response might be malformed, missing fields, or have wrong types. Pydantic validates the parsed JSON and produces clear errors when the output does not match expectations.

from pydantic import BaseModel
from typing import Optional
import json

class ExtractedEntity(BaseModel):
    name: str
    type: str  # "person", "company", "location"
    confidence: float

class ExtractionResult(BaseModel):
    entities: list[ExtractedEntity]
    summary: str
    language: str = "en"

def parse_llm_extraction(llm_output: str) -> Optional[ExtractionResult]:
    """Parse and validate LLM extraction output."""
    try:
        data = json.loads(llm_output)
        result = ExtractionResult(**data)
        return result
    except json.JSONDecodeError:
        print("LLM returned invalid JSON")
        return None
    except Exception as e:
        print(f"LLM output failed validation: {e}")
        return None

# Test with valid output
valid_output = json.dumps({
    "entities": [
        {"name": "Meta", "type": "company", "confidence": 0.98},
        {"name": "Menlo Park", "type": "location", "confidence": 0.85}
    ],
    "summary": "The document discusses Meta's AI infrastructure."
})

result = parse_llm_extraction(valid_output)
if result:
    print(f"Found {len(result.entities)} entities")
    for entity in result.entities:
        print(f"  {entity.name} ({entity.type}): {entity.confidence:.0%}")

# Test with invalid output - missing required field
invalid_output = json.dumps({
    "entities": [{"name": "Meta", "type": "company"}],  # Missing "confidence"
    "summary": "Test"
})
parse_llm_extraction(invalid_output)

Pydantic Settings - Configuration Management

Pydantic's BaseSettings class loads configuration from environment variables with type validation. This is the standard pattern for managing API keys, model configurations, and deployment settings in AI applications.

from pydantic import BaseModel, Field

# Simulating BaseSettings (actual BaseSettings reads from env vars)
class AISettings(BaseModel):
    """Application settings with defaults and validation."""
    openai_api_key: str = "sk-test-key"
    default_model: str = "gpt-4-turbo"
    default_temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_retries: int = Field(default=3, ge=1, le=10)
    cache_ttl_seconds: int = Field(default=3600, ge=60)
    debug_mode: bool = False

# Load settings (in production, values come from environment variables)
settings = AISettings()
print(f"Model: {settings.default_model}")
print(f"Temperature: {settings.default_temperature}")
print(f"Debug: {settings.debug_mode}")

# Settings are immutable by default in Pydantic v2
# settings.default_model = "gpt-3.5"  # Would raise error with frozen=True

graph TD
    subgraph Without_Pydantic["Without Pydantic"]
        W1["Data arrives\n(JSON from API, user input)"]
        W2["Parse with json.loads()\n→ raw dict"]
        W3["Access fields directly\n→ KeyError at runtime"]
        W4["No validation\n→ invalid data propagates"]
        W5["Bug found in production\n→ 6 hours to diagnose"]
    end
    
    subgraph With_Pydantic["With Pydantic"]
        P1["Data arrives\n(JSON from API, user input)"]
        P2["Parse into Pydantic model\n→ validated, typed object"]
        P3["Access fields with autocomplete\n→ IDE catches typos"]
        P4["Validation at boundary\n→ invalid data rejected immediately"]
        P5["Bug prevented\n→ 30 seconds to fix"]
    end
    
    style Without_Pydantic fill:#991b1b,stroke:#fca5a5,stroke-width:2px
    style With_Pydantic fill:#166534,stroke:#86efac,stroke-width:2px

This diagram contrasts data handling with and without Pydantic. Without Pydantic, data arrives, is parsed with json.loads into a raw dictionary, fields are accessed directly risking KeyError at runtime, no validation is performed so invalid data propagates, and bugs are found in production taking 6 hours to diagnose. With Pydantic, data arrives, is parsed into a validated typed Pydantic model, fields are accessed with IDE autocomplete catching typos, validation at the boundary rejects invalid data immediately, and bugs are prevented in 30 seconds.

graph LR
    subgraph Boundaries["Where Pydantic Guards Your Pipeline"]
        B1["User input\n→ validate query format"]
        B2["LLM response\n→ validate JSON structure"]
        B3["API response\n→ validate expected fields"]
        B4["Config loading\n→ validate settings"]
        B5["Database query\n→ validate row structure"]
    end
    
    subgraph Pipeline["Clean, Validated Data Flows Through"]
        P1["embed()"] --> P2["search()"]
        P2 --> P3["generate()"]
        P3 --> P4["respond()"]
    end
    
    Boundaries --> Pipeline
    
    style Boundaries fill:#312e81,stroke:#a5b4fc,stroke-width:2px
    style Pipeline fill:#166534,stroke:#86efac,stroke-width:2px

This diagram shows where Pydantic validates data at pipeline boundaries. Five boundary points are guarded: user input with query format validation, LLM response with JSON structure validation, API response with expected field validation, configuration loading with settings validation, and database query results with row structure validation. Once data passes through these boundary guards, it flows through the pipeline stages - embed, search, generate, respond - as clean, validated data.

Interview: AI engineering interview connection: Type hints and Pydantic are not directly tested in interviews, but they are expected in your code. When you write a function in an interview, adding type hints signals professionalism and production awareness. When the interviewer asks "how would you validate LLM output?", describing Pydantic model parsing is the production-grade answer.

What to demonstrate: Add type hints to every function you write in the interview. If asked about data validation, describe the Pydantic boundary pattern - validate data when it enters your system (from APIs, user input, LLM output) and trust it once validated. This shows you think about data integrity at the system level, not just at the function level.

Exercise: Starter: Add type hints to the following function signature: a function that takes a query string, a list of document dictionaries (each with "text" and "score" keys), and an optional maximum number of results, and returns a list of the document dictionaries that match.

Intermediate: Create a Pydantic model for an LLM chat completion response. It should validate: a list of choices (each with an index and a message containing role and content), a usage object (prompt_tokens, completion_tokens, total_tokens), and a model name. Write a parser function that takes a raw JSON string and returns the validated model or None.

Advanced: Design a Pydantic-based configuration system for an AI pipeline that uses three different models (embedding, classification, generation). Each model has its own settings (model name, API key, temperature, max tokens, timeout). The configuration should load from environment variables, provide sensible defaults, validate all constraints (temperature between 0 and 2, max tokens between 1 and 128000), and be immutable after creation. Write the complete Pydantic models and a factory function that creates the configuration.

Key Takeaways

Type hints transform Python from "read the function body to understand the types" to "read the function signature." For AI pipelines where data flows through many stages, this self-documentation prevents the category of bugs where a function receives data in an unexpected format - the most common production bug in dynamically typed AI codebases.

Pydantic extends type hints from static analysis (IDE and type checker) to runtime validation. Data entering your system from external sources - user input, LLM responses, API calls, file loads - is validated against Pydantic models at the boundary. Invalid data is rejected with clear error messages before it can propagate through the pipeline. This boundary validation pattern is the single most impactful code quality improvement for AI applications.

The Pydantic-for-LLM-output pattern is particularly valuable: ask the LLM to return JSON, parse it with json.loads, validate it with a Pydantic model, and handle validation failures gracefully (retry with a clarification prompt, or fall back to a default). This three-step pattern - parse, validate, handle errors - appears in every production AI application that consumes structured LLM output. When combined with the structured output features that model providers now offer (OpenAI's response_format, Anthropic's tool use for structured output), Pydantic validation becomes the safety net that catches the cases where the model's structured output is not quite right - missing an optional field, producing a value outside the expected range, or using a string where a number was expected.

Back to AI Engineering