方寸 Portal

AI Engineer Portal

Your personal operating system for career transition.

Private mode
Back to Learning Paths
Lesson

Type safety and validation patterns for AI systems

Advanced Pydantic patterns, discriminated unions for provider responses, generic types for AI abstractions, and runtime validation strategies.

70 min
python-for-ai-engineersphase-1portfolio

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

  1. Build a discriminated union for two LLM providers you have worked with. Write a parse_response function that takes raw provider output and returns the correct typed model.
  2. Add a @model_validator to a model that has two parallel lists (like chunks and scores). Verify it raises a clear error when the lists have different lengths.
  3. Write a parse_llm_json function that handles markdown code fences. Test it against: raw JSON, JSON with triple backticks, JSON wrapped in prose like "Here is the result: {...}".
Lesson Deep Dive

Ask a follow-up question about this lesson and get an AI-powered explanation.

Loading previous questions...