Why this matters
AI systems deal with a uniquely uncomfortable combination: structured types in your code, semi-structured data at provider boundaries, and model outputs that may or may not match the schema you specified. Weak typing means bugs that surface as silent data corruption rather than explicit errors. Strong typing — done at the right layer — gives you fast feedback, self-documenting code, and refactoring safety.
This lesson covers the Pydantic and typing patterns that AI engineers reach for in real production systems.
Core concepts
Discriminated unions for provider responses
Different LLM providers return fundamentally different response shapes. A naive approach uses optional fields and conditional logic everywhere. A discriminated union gives you a clean type per provider and exhaustive handling at the use site:
from __future__ import annotations
from typing import Literal, Union
from pydantic import BaseModel
class OpenAIResponse(BaseModel):
provider: Literal["openai"] = "openai"
id: str
choices: list[dict]
usage: dict
@property
def text(self) -> str:
return self.choices[0]["message"]["content"]
class AnthropicResponse(BaseModel):
provider: Literal["anthropic"] = "anthropic"
id: str
content: list[dict]
stop_reason: str
@property
def text(self) -> str:
return self.content[0]["text"]
class LocalResponse(BaseModel):
provider: Literal["local"] = "local"
text: str
model_path: str
ProviderResponse = Union[OpenAIResponse, AnthropicResponse, LocalResponse]
Pydantic's discriminated union uses the provider literal field to route validation:
from pydantic import TypeAdapter
adapter = TypeAdapter(ProviderResponse)
def parse_response(raw: dict, provider: str) -> ProviderResponse:
return adapter.validate_python({**raw, "provider": provider})
Now when you add a new provider, you add a new model class and Python's type checker will tell you every place that needs to handle the new case.
Generic types for AI abstractions
When building reusable AI infrastructure, generics let you preserve type information through abstractions without losing it to Any:
from typing import Generic, TypeVar
from pydantic import BaseModel
T = TypeVar("T", bound=BaseModel)
class ProviderRequest(BaseModel, Generic[T]):
model: str
payload: T
request_id: str
timeout_s: float = 30.0
class ProviderResult(BaseModel, Generic[T]):
request_id: str
data: T
latency_ms: int
cached: bool = False
class BatchResult(BaseModel, Generic[T]):
successes: list[ProviderResult[T]]
failures: list[dict]
total_latency_ms: int
A function typed as async def call(request: ProviderRequest[T]) -> ProviderResult[T] tells callers that what goes in comes out, typed. No Any in the middle.
Runtime validation for LLM structured outputs
When you ask an LLM to return JSON, it often returns almost-JSON. Robust validation handles the common failure modes:
import json
import re
from pydantic import BaseModel, ValidationError
class ExtractedEntities(BaseModel):
names: list[str]
dates: list[str]
locations: list[str]
def parse_llm_json(raw_text: str, model_class: type[BaseModel]) -> BaseModel | None:
# Strip markdown code fences if present
text = raw_text.strip()
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\n?", "", text)
text = re.sub(r"\n?```$", "", text)
# Find the outermost JSON object or array
match = re.search(r"(\{.*\}|\[.*\])", text, re.DOTALL)
if not match:
return None
try:
data = json.loads(match.group(1))
return model_class.model_validate(data)
except (json.JSONDecodeError, ValidationError):
return None
This handles the three most common LLM JSON failures: markdown code fences, leading/trailing prose, and minor structural errors.
Nested models with computed fields
Complex AI payloads often have relationships between fields that need validation:
from pydantic import BaseModel, Field, model_validator
from typing import Optional
class RetrievedContext(BaseModel):
chunks: list[str]
scores: list[float]
total_tokens: int
@model_validator(mode="after")
def validate_parallel_lengths(self) -> "RetrievedContext":
if len(self.chunks) != len(self.scores):
raise ValueError(
f"chunks and scores must have equal length, "
f"got {len(self.chunks)} and {len(self.scores)}"
)
return self
class GenerationRequest(BaseModel):
user_query: str = Field(min_length=1)
context: Optional[RetrievedContext] = None
system_prompt: str = Field(default="You are a helpful assistant.")
max_tokens: int = Field(default=1024, ge=1, le=8192)
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
The @model_validator runs after all individual field validators, letting you check cross-field invariants.
Working example
A type-safe provider abstraction that handles multiple providers with full type preservation:
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Generic, TypeVar
from pydantic import BaseModel
ResponseT = TypeVar("ResponseT", bound=BaseModel)
class NormalizedResponse(BaseModel):
request_id: str
text: str
model: str
input_tokens: int
output_tokens: int
latency_ms: int
class ProviderAdapter(ABC):
@abstractmethod
async def generate(self, prompt: str, model: str) -> NormalizedResponse:
...
@abstractmethod
def normalize(self, raw: dict) -> NormalizedResponse:
...
class OpenAIAdapter(ProviderAdapter):
async def generate(self, prompt: str, model: str = "gpt-4o-mini") -> NormalizedResponse:
import time
start = time.monotonic()
raw = await self._client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
)
return self.normalize(raw.model_dump(), latency_ms=int((time.monotonic() - start) * 1000))
def normalize(self, raw: dict, latency_ms: int = 0) -> NormalizedResponse:
return NormalizedResponse(
request_id=raw["id"],
text=raw["choices"][0]["message"]["content"],
model=raw["model"],
input_tokens=raw["usage"]["prompt_tokens"],
output_tokens=raw["usage"]["completion_tokens"],
latency_ms=latency_ms,
)
Every provider adapter produces the same NormalizedResponse. Application code only ever sees the normalized shape. Provider-specific details never leak past the adapter.
Common mistakes
Using Optional[str] where you mean str. Optional[str] means the field can be None. If None is not a valid value in your domain, use a non-optional type and let validation fail explicitly when the field is missing.
Skipping validation on LLM outputs because "it usually works". LLM outputs that are sometimes invalid are a reliability bug. Validate every structured output. Handle the failure case explicitly.
Putting all providers in one model with optional fields. class Response(BaseModel): openai_id: Optional[str]; anthropic_id: Optional[str] is a maintenance nightmare. Use discriminated unions.
Not using model_validator for cross-field constraints. Individual field validators cannot see other fields. Cross-field invariants belong in a model_validator.
Try it yourself
- Build a discriminated union for two LLM providers you have worked with. Write a
parse_responsefunction that takes raw provider output and returns the correct typed model. - Add a
@model_validatorto a model that has two parallel lists (like chunks and scores). Verify it raises a clear error when the lists have different lengths. - Write a
parse_llm_jsonfunction that handles markdown code fences. Test it against: raw JSON, JSON with triple backticks, JSON wrapped in prose like "Here is the result: {...}".