Free lesson · AI Engineering
Object-Oriented Python and Dataclasses
Object-Oriented Python and Dataclasses
When to Use OOP in AI Applications
At Meta, I have watched two types of AI codebases evolve. The first type uses classes for everything - a PromptManager class, a ResponseParser class, a ChunkProcessor class, each with five methods and three instance variables. The code is well-organised but over-engineered - most of these classes have a single method that does the actual work, and the rest is boilerplate. The second type uses functions for everything - a flat collection of 200 functions in a single file, with data passed as dictionaries. The code is hard to navigate and data structures are implicit - you have to read the function body to know what keys a dictionary should have.
The right approach is in the middle. Use classes when you need to encapsulate state that persists across multiple method calls (an LLM client with connection pooling, a pipeline with configuration). Use functions when you need stateless data transformation (parsing a response, formatting a prompt). Use dataclasses when you need structured data containers (search results, API responses, configuration objects).
This topic covers the Python OOP features you need for AI engineering - not everything Python OOP offers, but the subset that produces clean, maintainable AI application code. Specifically: basic classes for stateful components, dataclasses for structured data, frozen dataclasses for immutable configuration, enums for named constants, abstract base classes for swappable interfaces, and protocols for lightweight duck-typed interfaces. If you are coming from Java, you will find Python OOP familiar but lighter - fewer access modifiers, no mandatory getters/setters, and dataclasses that eliminate the boilerplate that Java records were introduced to solve.
Classes: The Basics
A Python class bundles data (attributes) and behaviour (methods) into a single object. The __init__ method initialises the object when it is created. The self parameter refers to the object itself - it is explicit in Python, unlike this in Java or C++ which is implicit.
class LLMClient:
"""A simple LLM client that maintains state across calls."""
def __init__(self, model: str, temperature: float = 0.7):
self.model = model
self.temperature = temperature
self.call_count = 0
self.total_tokens = 0
def complete(self, prompt: str) -> str:
"""Simulate an LLM completion call."""
self.call_count += 1
# Simulate token counting
input_tokens = len(prompt.split()) * 2
output_tokens = 50 # Simulated
self.total_tokens += input_tokens + output_tokens
return f"[{self.model}] Response to: {prompt[:30]}..."
def get_stats(self) -> dict:
"""Return usage statistics."""
return {
"model": self.model,
"calls": self.call_count,
"total_tokens": self.total_tokens,
"avg_tokens_per_call": self.total_tokens / max(self.call_count, 1)
}
# Usage
client = LLMClient(model="gpt-4-turbo", temperature=0.3)
response = client.complete("What is retrieval augmented generation?")
print(response)
response = client.complete("How do vector databases work?")
print(response)
print(f"Stats: {client.get_stats()}")
When classes make sense in AI applications: The LLMClient above maintains state - call count and token usage - that persists across method calls. This state tracking would be awkward with standalone functions (you would need global variables or explicit state-passing). The class encapsulates the state naturally. Use classes for: API clients (maintaining connection pools and auth tokens), pipeline orchestrators (managing stage execution order and shared context), configuration managers (loading and providing settings across the application), and any component that maintains state across multiple interactions.
Properties - computed attributes that look like data access: The @property decorator (shown in the SearchResponse dataclass above) creates methods that are accessed like attributes. results.avg_score looks like attribute access but actually computes the average on each call. Properties are useful for derived values that should not be stored separately because they would go stale when the underlying data changes.
Dunder methods for Pythonic behaviour: Python uses "dunder" (double underscore) methods to make your classes work with built-in operations. __len__ enables len(obj), __str__ enables print(obj), __getitem__ enables obj[key], and __iter__ enables for item in obj. In AI applications, implementing __len__ on a Conversation class to return the message count, or __iter__ to iterate over messages, makes your classes feel like native Python objects rather than Java-style objects transplanted into Python.
When classes are overkill: A function that parses a JSON response does not need state. A function that formats a prompt template does not need state. A function that computes cosine similarity does not need state. Making these into classes adds boilerplate without adding value. Keep them as functions.
Dataclasses: Structured Data Without Boilerplate
Dataclasses are Python's solution for data containers - objects that primarily hold data rather than implementing complex behaviour. They automatically generate __init__, __repr__, __eq__, and optionally __hash__ methods based on the fields you define. This eliminates the boilerplate of writing self.field = field for every attribute.
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class SearchResult:
"""A single result from a vector search."""
text: str
score: float
document_id: str
metadata: dict = field(default_factory=dict)
@dataclass
class SearchResponse:
"""Response from a vector search operation."""
query: str
results: list[SearchResult]
latency_ms: float
total_candidates: int
@property
def top_result(self) -> Optional[SearchResult]:
"""The highest-scoring result, or None if no results."""
return self.results[0] if self.results else None
@property
def avg_score(self) -> float:
"""Average score across all results."""
if not self.results:
return 0.0
return sum(r.score for r in self.results) / len(self.results)
# Usage - clean, readable, type-safe
results = SearchResponse(
query="What is RAG?",
results=[
SearchResult(text="RAG combines retrieval with generation...",
score=0.94, document_id="doc_001"),
SearchResult(text="Retrieval systems fetch relevant context...",
score=0.87, document_id="doc_002"),
],
latency_ms=23.4,
total_candidates=1500
)
print(f"Query: {results.query}")
print(f"Top result: {results.top_result.text[:50]}...")
print(f"Average score: {results.avg_score:.2f}")
print(f"Latency: {results.latency_ms}ms")
Why dataclasses matter for AI engineering: AI applications pass structured data between pipeline stages. Without dataclasses, this data is usually dictionaries - which have no structure enforcement, no autocompletion in IDEs, and no documentation of what fields they contain. With dataclasses, the data structure is explicit, enforced at creation time, and documented by the class definition. When a function receives a SearchResponse instead of a dict, the developer knows exactly what fields are available without reading the upstream code.
The dictionary problem in practice: Consider a pipeline where the search function returns {"text": "...", "score": 0.94, "doc_id": "001"}. Downstream code accesses result["scroe"] - a typo that Python does not catch until runtime, causing a KeyError in production. With a dataclass, the attribute access result.scroe is caught by the IDE immediately and by type checkers (mypy) before deployment. This matters at scale: a pipeline with 10 stages passing dictionaries between them has 10 opportunities for silent key name mismatches. Dataclasses eliminate all of them.
Dataclass conversion methods: Dataclasses integrate cleanly with the JSON-based world of AI APIs:
from dataclasses import dataclass, asdict, field
@dataclass
class APIRequest:
model: str
messages: list[dict]
temperature: float = 0.7
max_tokens: int = 500
# Create from structured data
request = APIRequest(
model="gpt-4-turbo",
messages=[{"role": "user", "content": "What is RAG?"}]
)
# Convert to dict for JSON serialisation (API call body)
import json
request_dict = asdict(request)
print(json.dumps(request_dict, indent=2))
Enums - Named Constants for AI Configuration
Enums define a fixed set of named values. In AI applications, they replace magic strings for model names, task types, and status values - preventing bugs where a typo in a string goes undetected.
from enum import Enum
class ModelTier(Enum):
FRONTIER = "gpt-4-turbo"
MID_RANGE = "gpt-4o"
EFFICIENT = "gpt-4o-mini"
class TaskType(Enum):
CLASSIFICATION = "classification"
GENERATION = "generation"
EMBEDDING = "embedding"
# Usage - type-safe, no magic strings
def select_model(task: TaskType) -> ModelTier:
routing = {
TaskType.CLASSIFICATION: ModelTier.EFFICIENT,
TaskType.GENERATION: ModelTier.FRONTIER,
TaskType.EMBEDDING: ModelTier.EFFICIENT,
}
return routing[task]
selected = select_model(TaskType.CLASSIFICATION)
print(f"Selected: {selected.value} for classification")
Frozen Dataclasses - Immutable Data
Adding frozen=True to a dataclass makes it immutable - instances cannot be modified after creation. This is valuable for data that flows through a pipeline and should not be accidentally modified by a downstream stage.
from dataclasses import dataclass
@dataclass(frozen=True)
class ModelConfig:
"""Immutable model configuration - cannot be changed after creation."""
model: str
temperature: float
max_tokens: int
top_p: float = 1.0
config = ModelConfig(model="gpt-4-turbo", temperature=0.7, max_tokens=500)
print(f"Config: {config}")
# This would raise FrozenInstanceError:
# config.temperature = 0.3 # Error! Cannot modify frozen dataclass
Inheritance and Abstract Base Classes
Inheritance in AI applications is used sparingly but precisely - primarily to define interfaces that multiple implementations must conform to. Python's abc module provides abstract base classes (ABCs) that define required methods without implementing them.
graph TD
subgraph Dict_Approach["Using Dictionaries (fragile)"]
DA1["result = {'text': '...', 'score': 0.94}"]
DA2["Access: result['scroe']\n→ KeyError at RUNTIME"]
DA3["No autocompletion\nNo documentation\nNo type checking"]
end
subgraph DC_Approach["Using Dataclasses (robust)"]
DC1["@dataclass\nclass SearchResult:\n text: str\n score: float"]
DC2["Access: result.scroe\n→ Caught by IDE + mypy"]
DC3["Autocompletion ✓\nSelf-documenting ✓\nType-safe ✓"]
end
style Dict_Approach fill:#991b1b,stroke:#fca5a5,stroke-width:2px
style DC_Approach fill:#166534,stroke:#86efac,stroke-width:2px
This diagram compares using dictionaries versus dataclasses for structured data in AI applications. With dictionaries, a result is accessed via string keys like result bracket scroe, which produces a KeyError at runtime due to the typo, with no autocompletion, documentation, or type checking. With dataclasses, the same typo result dot scroe is caught by the IDE and mypy before runtime, providing autocompletion, self-documentation, and type safety.
from abc import ABC, abstractmethod
class EmbeddingModel(ABC):
"""Abstract base class - defines the interface for all embedding models."""
@abstractmethod
def embed(self, text: str) -> list[float]:
"""Generate embedding vector for text. All subclasses must implement."""
pass
@abstractmethod
def embed_batch(self, texts: list[str]) -> list[list[float]]:
"""Generate embeddings for a batch of texts."""
pass
@property
@abstractmethod
def dimension(self) -> int:
"""The dimensionality of the embedding vector."""
pass
class OpenAIEmbeddings(EmbeddingModel):
"""OpenAI implementation of the embedding interface."""
def __init__(self, model: str = "text-embedding-3-small"):
self.model = model
self._dimension = 1536
def embed(self, text: str) -> list[float]:
# In production: call OpenAI API
return [0.1] * self._dimension # Simulated
def embed_batch(self, texts: list[str]) -> list[list[float]]:
return [self.embed(t) for t in texts]
@property
def dimension(self) -> int:
return self._dimension
class LocalEmbeddings(EmbeddingModel):
"""Local model implementation of the embedding interface."""
def __init__(self, model_path: str):
self.model_path = model_path
self._dimension = 384
def embed(self, text: str) -> list[float]:
return [0.2] * self._dimension # Simulated
def embed_batch(self, texts: list[str]) -> list[list[float]]:
return [self.embed(t) for t in texts]
@property
def dimension(self) -> int:
return self._dimension
# Both implementations work with any code that expects an EmbeddingModel
def build_index(model: EmbeddingModel, documents: list[str]):
"""Works with any embedding model implementation."""
print(f"Building index with {model.__class__.__name__} (dim={model.dimension})")
embeddings = model.embed_batch(documents)
print(f" Generated {len(embeddings)} embeddings")
# Swap implementations without changing pipeline code
docs = ["doc1", "doc2", "doc3"]
build_index(OpenAIEmbeddings(), docs)
build_index(LocalEmbeddings("/models/all-MiniLM-L6"), docs)
When to use ABCs in AI applications: When you have multiple implementations of the same interface that need to be swappable - different embedding models, different vector databases, different LLM providers. The ABC defines the contract. Each implementation fulfils it. Pipeline code depends on the abstract interface, not on any specific implementation. This is the dependency inversion principle applied to AI components.
When ABCs are overkill: If you only have one implementation and no plans for a second, an ABC is unnecessary ceremony. A simple class or even a module of functions is cleaner. Add the ABC when the second implementation appears, not in anticipation of it.
Alternative to ABCs: Protocol typing. Python 3.8+ offers typing.Protocol as a lighter-weight alternative to ABCs for defining interfaces. Protocols use structural typing (duck typing with type checking) rather than nominal typing (explicit inheritance). A class satisfies a Protocol if it has the required methods, even if it does not explicitly inherit from the Protocol. This is often more Pythonic than ABCs for simple interfaces:
from typing import Protocol
class Embeddable(Protocol):
"""Any class with an embed method satisfies this protocol."""
def embed(self, text: str) -> list[float]: ...
# This class satisfies Embeddable without inheriting from it
class SimpleEmbedder:
def embed(self, text: str) -> list[float]:
return [0.1] * 384 # Simulated
def use_embedder(model: Embeddable, text: str) -> list[float]:
"""Works with any object that has an embed method."""
return model.embed(text)
embedder = SimpleEmbedder()
result = use_embedder(embedder, "test")
print(f"Embedding dimension: {len(result)}")
Protocols are preferred for simple interfaces (1-2 methods). ABCs are preferred for complex interfaces (3+ methods) or when you need to share implementation between subclasses (via non-abstract methods on the base class).
graph TD
subgraph OOP_Guide["When to Use What"]
subgraph Functions["Use Functions"]
F1["Stateless transformations\n(parse response, format prompt)"]
F2["Pure computations\n(cosine similarity, token counting)"]
F3["One-off operations\n(validate schema, clean text)"]
end
subgraph Classes["Use Classes"]
C1["Stateful components\n(API clients, pipeline orchestrators)"]
C2["Resource management\n(connection pools, model loaders)"]
C3["Complex behaviour\n(multi-step processing)"]
end
subgraph Dataclasses["Use Dataclasses"]
D1["Data containers\n(search results, API responses)"]
D2["Configuration objects\n(model config, pipeline settings)"]
D3["Pipeline stage outputs\n(structured intermediate results)"]
end
end
style Functions fill:#1e3a5f,stroke:#7dd3fc,stroke-width:2px
style Classes fill:#312e81,stroke:#a5b4fc,stroke-width:2px
style Dataclasses fill:#166534,stroke:#86efac,stroke-width:2px
This diagram provides guidance on when to use functions, classes, and dataclasses in AI applications. Use functions for stateless transformations like parsing responses and formatting prompts, pure computations like cosine similarity and token counting, and one-off operations like schema validation and text cleaning. Use classes for stateful components like API clients and pipeline orchestrators, resource management like connection pools and model loaders, and complex multi-step processing behaviour. Use dataclasses for data containers like search results and API responses, configuration objects like model config and pipeline settings, and pipeline stage outputs that need structured intermediate results.
Interview: AI engineering interview connection: AI engineering interviews test OOP through system design questions - "design a pluggable embedding pipeline" or "design an LLM client that supports multiple providers." The expected answer uses abstract base classes for the interface, concrete classes for each provider, and dataclasses for request/response objects. Demonstrating that you know when not to use classes (stateless transformations should be functions) is as important as knowing when to use them.
What to practice: Implement an abstract
LLMProviderclass withcomplete()andstream()methods, then implement it for two providers (OpenAI and Anthropic). Use dataclasses for the request and response objects. This exercise covers the OOP patterns you will use in production AI code.
Exercise: Starter: Create a
Messagedataclass with fields: role (str), content (str), and timestamp (float, defaulting to current time). Create aConversationdataclass containing a list of Messages and a model name. Add a method to compute the total word count across all messages.Intermediate: Design an abstract
VectorStoreclass with methods:add(id, embedding, metadata),search(query_embedding, top_k) → list[SearchResult], anddelete(id). Implement it for an in-memory store (using a Python dict). Use dataclasses forSearchResult.Advanced: Build a
PipelineBuilderclass that lets you compose an AI pipeline from stages:builder.add_stage("embed", embed_fn).add_stage("search", search_fn).add_stage("generate", generate_fn).build(). The built pipeline should execute stages in order, passing each stage's output as the next stage's input. Use dataclasses for the pipeline configuration and stage results.
Key Takeaways
OOP in AI applications follows a "right tool for the right job" principle. Functions handle stateless transformations (parsing, formatting, computing). Classes handle stateful components (API clients, pipeline orchestrators, model managers). Dataclasses handle structured data (search results, configurations, pipeline stage outputs). Enums handle fixed sets of named constants (model tiers, task types, status values). Using all four appropriately produces clean, maintainable code. Over-using any one pattern produces either boilerplate-heavy code (too many classes) or unstructured code (too many raw dictionaries and magic strings).
Dataclasses are the most immediately useful Python OOP feature for AI engineering. They replace the dictionaries that most AI code uses for structured data, adding type safety, IDE autocompletion, and self-documentation. The transition from result["score"] to result.score is small in syntax but large in code quality - the dataclass definition documents the data structure explicitly, the type checker catches access to nonexistent fields before runtime, and the asdict() function converts dataclasses to dictionaries for JSON serialisation when you need to send them to an API. Frozen dataclasses (@dataclass(frozen=True)) add immutability guarantees that prevent downstream pipeline stages from accidentally modifying data that upstream stages depend on.
Abstract base classes and Protocols define swappable interfaces for AI components. Different embedding models, different vector databases, and different LLM providers all conform to the same interface, enabling pipeline code that works with any implementation. Use ABCs for complex interfaces with shared implementation. Use Protocols for simple interfaces where structural typing (duck typing with type safety) is sufficient. Skip both when you have only one implementation and no concrete plans for a second - premature abstraction is as harmful as no abstraction.