Free lesson · AI Engineering

Functions, Decorators, and Generators

Easy11 min readaillmrag

Functions, Decorators, and Generators

Why This Topic Matters for AI Engineering

AI applications are built from functions that call other functions. A RAG pipeline is a chain: embed_query() → search_vectors() → assemble_context() → call_llm() → parse_response(). Each step is a function. The quality of your functions - how you design their signatures, how you handle their errors, how you compose them - determines the quality of your pipeline.

This topic covers three Python features that elevate functions from basic to production-grade: first-class functions (functions as values you can pass around), decorators (functions that wrap other functions to add behaviour), and generators (functions that produce values lazily for memory-efficient processing). These are not exotic features - they appear in every AI application codebase at Meta.

First-Class Functions

In Python, functions are objects. You can assign them to variables, pass them as arguments to other functions, return them from other functions, and store them in data structures. This is what "first-class" means - functions are treated the same as any other value.

Why this matters for AI engineering: AI pipelines are configurable. The same pipeline might use different embedding functions (OpenAI embeddings vs Cohere embeddings vs local model), different scoring functions (cosine similarity vs dot product vs euclidean distance), or different post-processing functions (truncate vs summarise vs filter) depending on the configuration. First-class functions enable this configurability without inheritance hierarchies or strategy pattern boilerplate. You pass the function as an argument, and the pipeline uses whatever function it receives.

The callback pattern: In asynchronous AI applications, callbacks are functions passed as arguments that are called when an operation completes. When an LLM streaming response finishes, the callback processes the complete response. When a batch embedding job completes, the callback stores the embeddings. Callbacks are first-class functions applied to event-driven architectures.

# Functions as values
def cosine_similarity(a, b):
    return sum(x * y for x, y in zip(a, b))

def euclidean_distance(a, b):
    return sum((x - y) ** 2 for x, y in zip(a, b)) ** 0.5

# Choose the similarity function at runtime
def find_most_similar(query_embedding, candidates, similarity_fn):
    """Find the most similar candidate using a configurable similarity function."""
    best_score = -float('inf')
    best_candidate = None
    
    for candidate in candidates:
        score = similarity_fn(query_embedding, candidate["embedding"])
        if score > best_score:
            best_score = score
            best_candidate = candidate
    
    return best_candidate, best_score

# Usage - same function, different similarity metrics
query = [0.1, 0.8, 0.3]
candidates = [
    {"text": "What is RAG?", "embedding": [0.2, 0.7, 0.4]},
    {"text": "How do embeddings work?", "embedding": [0.1, 0.9, 0.2]},
]

result, score = find_most_similar(query, candidates, cosine_similarity)
print(f"Best match (cosine): '{result['text']}' with score {score:.4f}")

Closures - Functions That Remember Their Environment

A closure is a function that captures variables from its enclosing scope. When the enclosing function returns, the inner function retains access to those variables. Closures are how you create configurable functions without classes.

def make_threshold_filter(threshold):
    """Create a filter function with a baked-in threshold."""
    def filter_fn(results):
        return [r for r in results if r["score"] >= threshold]
    return filter_fn

# Create filters with different thresholds
strict_filter = make_threshold_filter(0.90)
lenient_filter = make_threshold_filter(0.70)

results = [
    {"text": "Result A", "score": 0.95},
    {"text": "Result B", "score": 0.82},
    {"text": "Result C", "score": 0.67},
]

print(f"Strict: {[r['text'] for r in strict_filter(results)]}")
print(f"Lenient: {[r['text'] for r in lenient_filter(results)]}")

Lambda Functions - Anonymous One-Liners

Lambda functions are anonymous functions for simple, one-expression operations. They are most commonly used with sorted(), map(), and filter() - functions that take a function as an argument.

# Sorting with a custom key
results = [
    {"text": "Result A", "score": 0.82},
    {"text": "Result B", "score": 0.95},
    {"text": "Result C", "score": 0.67},
]

# Sort by score descending
ranked = sorted(results, key=lambda r: r["score"], reverse=True)
for r in ranked:
    print(f"  {r['score']:.2f} - {r['text']}")

Decorators - Functions That Wrap Functions

A decorator is a function that takes a function as input and returns a new function that adds behaviour before, after, or around the original function. Decorators are Python's mechanism for cross-cutting concerns: logging, caching, timing, retry logic, and authentication - behaviours that apply to many functions but should not be duplicated inside each one.

graph LR
    subgraph Without["Without Decorators"]
        W1["call_llm()"] --> W2["logging code\n+ timing code\n+ retry code\n+ cost tracking\n+ actual LLM call\n(all mixed together)"]
    end
    
    subgraph With["With Decorators"]
        D1["@retry"] --> D2["@log_and_time"]
        D2 --> D3["@cost_track"]
        D3 --> D4["call_llm()\n(only business logic)"]
    end
    
    style Without fill:#991b1b,stroke:#fca5a5,stroke-width:2px
    style With fill:#166534,stroke:#86efac,stroke-width:2px

This diagram contrasts two approaches to adding cross-cutting concerns to an LLM call function. Without decorators, the logging code, timing code, retry code, cost tracking, and actual LLM call are all mixed together in a single function. With decorators, three separate decorators - retry, log_and_time, and cost_track - wrap the call_llm function, which contains only the business logic. Each concern is separated and reusable.

The AI engineering use case: Every LLM API call needs logging (what prompt was sent, what response was received, which model was used), timing (how long the call took, what was time-to-first-token), retry logic (what to do when the provider returns a 429 or 503), and cost tracking (how many tokens were consumed, what was the cost). Writing this logic inside every API call function is repetitive, error-prone, and makes the core logic harder to read. Decorators apply these concerns once as reusable wrappers and attach them with a single @decorator line. The function stays clean - it contains only the business logic. The decorator handles the infrastructure concerns.

import time
import functools

def log_and_time(func):
    """Decorator that logs function calls and measures execution time."""
    @functools.wraps(func)  # Preserve the original function's name and docstring
    def wrapper(*args, **kwargs):
        start = time.time()
        print(f"→ Calling {func.__name__}...")
        
        try:
            result = func(*args, **kwargs)
            elapsed = (time.time() - start) * 1000
            print(f"← {func.__name__} completed in {elapsed:.1f}ms")
            return result
        except Exception as e:
            elapsed = (time.time() - start) * 1000
            print(f"✗ {func.__name__} failed after {elapsed:.1f}ms: {e}")
            raise
    
    return wrapper

@log_and_time
def embed_text(text):
    """Simulate embedding a text string."""
    time.sleep(0.1)  # Simulate API call
    return [0.1, 0.8, 0.3, 0.5]  # Fake embedding

@log_and_time
def search_similar(embedding, top_k=3):
    """Simulate vector search."""
    time.sleep(0.05)
    return [{"text": "result", "score": 0.92}]

# Both functions now have logging and timing automatically
result = embed_text("What is RAG?")
results = search_similar(result)

Practical Decorators for AI Engineering

Retry decorator - automatically retries failed API calls with exponential backoff. This is the single most important decorator in AI engineering because every external API call can fail transiently.

import time
import functools

def retry(max_attempts=3, backoff_factor=2):
    """Retry decorator with exponential backoff."""
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(1, max_attempts + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_attempts:
                        wait = backoff_factor ** (attempt - 1)
                        print(f"  Attempt {attempt} failed: {e}. Retrying in {wait}s...")
                        time.sleep(wait)
                    else:
                        print(f"  All {max_attempts} attempts failed.")
            raise last_exception
        return wrapper
    return decorator

@retry(max_attempts=3, backoff_factor=1)
def flaky_api_call():
    """Simulates a flaky API that fails sometimes."""
    import random
    if random.random() < 0.6:
        raise ConnectionError("API timeout")
    return {"status": "success"}

# This will retry up to 3 times
try:
    result = flaky_api_call()
    print(f"Result: {result}")
except ConnectionError:
    print("API unavailable after all retries")

Caching decorator - Python's built-in @functools.lru_cache caches function results based on arguments. For AI applications, this is useful for caching embeddings (the same text always produces the same embedding), caching API responses for identical prompts, and caching expensive computations that are called repeatedly with the same inputs.

Important limitation: @lru_cache requires hashable arguments. Lists and dicts are not hashable, so you cannot cache a function that takes a list as input without converting it to a tuple first. For AI applications where inputs are often complex objects, consider using a custom caching decorator with a hash of the serialised input as the cache key.

import functools

@functools.lru_cache(maxsize=1000)
def expensive_computation(text):
    """Simulate an expensive embedding computation."""
    print(f"  Computing embedding for: '{text[:30]}...'")
    return hash(text)  # Fake embedding as hash

# First call - computes
result1 = expensive_computation("What is retrieval augmented generation?")
# Second call - cache hit, no computation
result2 = expensive_computation("What is retrieval augmented generation?")
# Third call - different input, computes
result3 = expensive_computation("How do vector databases work?")

print(f"Cache info: {expensive_computation.cache_info()}")

Stacking Decorators

Decorators can be stacked - applying multiple decorators to the same function. The order matters: decorators are applied bottom-up (the decorator closest to the function is applied first).

import time
import functools

def log_call(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"  [LOG] Calling {func.__name__}")
        return func(*args, **kwargs)
    return wrapper

def time_call(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        elapsed = (time.time() - start) * 1000
        print(f"  [TIME] {func.__name__}: {elapsed:.1f}ms")
        return result
    return wrapper

@log_call      # Applied second (outer)
@time_call     # Applied first (inner)
def process_document(text):
    """Process a document with logging and timing."""
    time.sleep(0.05)
    return f"Processed: {text[:20]}..."

result = process_document("This is a sample document for processing")
print(result)

In AI pipelines, stacking is common: @retry @log_and_time @cost_track on an LLM API call gives you resilience, observability, and cost tracking - three cross-cutting concerns applied without touching the function's core logic.

Generators - Lazy Evaluation for Large Data

Generators produce values one at a time, on demand, rather than computing all values upfront and storing them in memory. This is critical for AI applications that process large datasets - embedding 100,000 documents, streaming LLM responses token by token, or iterating over search results from a paginated API.

The memory problem generators solve: If you have 100,000 documents to embed and you load them all into a list, they consume memory simultaneously - potentially gigabytes for large document collections. If the embedding function crashes on document 50,000, you have wasted the memory for all 100,000 documents while only processing half. With a generator, only one document is in memory at a time - the generator produces the next document when asked and the previous one is garbage collected. This means your pipeline can process arbitrarily large datasets regardless of available memory, and a failure on document 50,000 does not waste memory for documents 50,001 through 100,000.

When to use generators vs lists: Use a list when you need random access (accessing element [47] directly), when the dataset fits comfortably in memory, or when you need to iterate over the data multiple times. Use a generator when you process elements sequentially (one by one, in order), when the dataset might not fit in memory, or when you only need to iterate once. In AI pipelines, sequential processing of large datasets is the dominant pattern - which is why generators appear throughout this course.

# List - loads everything into memory at once
def get_all_chunks_list(documents):
    chunks = []
    for doc in documents:
        for paragraph in doc.split("\n\n"):
            chunks.append(paragraph)
    return chunks  # All chunks in memory simultaneously

# Generator - produces one chunk at a time
def get_all_chunks_generator(documents):
    for doc in documents:
        for paragraph in doc.split("\n\n"):
            yield paragraph  # Produce one, then pause until next is requested

# Usage
docs = ["First paragraph\n\nSecond paragraph", "Third paragraph\n\nFourth paragraph"]

# Generator version - memory efficient
for chunk in get_all_chunks_generator(docs):
    print(f"  Processing: '{chunk[:30]}...'")
    # Process one chunk, then the generator produces the next

Generator Expressions - One-Line Generators

Just as list comprehensions create lists, generator expressions create generators. The syntax is identical except with parentheses instead of brackets. This one-character difference changes the memory profile fundamentally - a list comprehension allocates memory for all values at once, while a generator expression allocates memory for only one value at a time.

# List comprehension - creates all values in memory
scores_list = [x * 0.1 for x in range(10)]

# Generator expression - creates values on demand
scores_gen = (x * 0.1 for x in range(10))

# Use with functions that consume iterables
total = sum(x * 0.1 for x in range(10))  # No extra memory for the list
print(f"Total: {total}")

# Practical: count tokens across many documents without loading all into memory
documents = ["short doc", "a slightly longer document here", "another one"]
total_words = sum(len(doc.split()) for doc in documents)
print(f"Total words: {total_words}")

The Streaming Pattern: Generators for LLM Token Streaming

The most important generator pattern in AI engineering is streaming LLM responses. When a model generates a response token by token, a generator function yields each token as it arrives, enabling the application to display tokens to the user in real time rather than waiting for the full response.

import time

def simulate_llm_stream(prompt):
    """Simulate an LLM streaming response, yielding one word at a time."""
    response_words = "Retrieval Augmented Generation combines a retrieval system with a language model".split()
    for word in response_words:
        time.sleep(0.05)  # Simulate token generation delay
        yield word + " "

# Consume the stream - display tokens as they arrive
print("Response: ", end="")
for token in simulate_llm_stream("What is RAG?"):
    print(token, end="", flush=True)
print()  # Newline after response

This pattern - yield inside a loop that receives data from an external source - is how every streaming LLM integration works. The OpenAI, Anthropic, and Google SDKs all return generator-like objects that yield tokens. Your application consumes them with a for loop. Understanding generators is understanding streaming.

graph TD
    subgraph Concepts["Python Function Concepts for AI"]
        subgraph FirstClass["First-Class Functions"]
            FC1["Functions as arguments\n→ configurable pipelines"]
            FC2["Closures\n→ parameterised filters"]
            FC3["Lambda\n→ inline sort keys"]
        end
        
        subgraph Decorators["Decorators"]
            D1["@log_and_time\n→ observability"]
            D2["@retry\n→ resilience"]
            D3["@lru_cache\n→ performance"]
        end
        
        subgraph Generators["Generators"]
            G1["yield keyword\n→ memory efficiency"]
            G2["Generator expressions\n→ one-line lazy iteration"]
            G3["Streaming patterns\n→ token-by-token processing"]
        end
    end
    
    style FirstClass fill:#1e1b4b,stroke:#818cf8,stroke-width:2px
    style Decorators fill:#312e81,stroke:#a5b4fc,stroke-width:2px
    style Generators fill:#166534,stroke:#86efac,stroke-width:2px

This diagram summarises three Python function concepts essential for AI engineering. First-class functions include functions as arguments for configurable pipelines, closures for parameterised filters, and lambda for inline sort keys. Decorators include log_and_time for observability, retry for resilience, and lru_cache for performance. Generators include the yield keyword for memory efficiency, generator expressions for one-line lazy iteration, and streaming patterns for token-by-token processing.

Interview: AI engineering interview connection: Decorators and generators appear in AI engineering interviews in two forms. First, you may be asked to implement a retry decorator or a caching decorator - these test your understanding of higher-order functions. Second, you may be asked to process a large dataset efficiently - the expected answer uses generators, not lists, to demonstrate memory awareness.

What to practice: Implement a retry decorator with configurable max attempts and backoff. Implement a generator that reads a large file line by line and yields processed results. These are the two patterns that come up most frequently.

Exercise: Starter: Write a decorator called @count_calls that tracks how many times a function has been called and prints the count on each invocation. Test it on a function that simulates an API call.

Intermediate: Write a generator function stream_tokens(text, chunk_size=5) that takes a string and yields it in chunks of chunk_size words, simulating token streaming. Then write a consumer function that prints each chunk with a 0.1-second delay, simulating real-time display.

Advanced: Write a @cost_tracker decorator for LLM API calls. The decorator should accept a price_per_1k_tokens parameter, extract the token count from the function's return value (assume it returns a dict with a "usage" → "total_tokens" key), compute the cost, and maintain a running total accessible via function_name.total_cost. Test it with a mock LLM call function.

Key Takeaways

First-class functions enable configurable AI pipelines. Instead of hardcoding a similarity function, pass it as a parameter. Instead of writing a different filter for each threshold, create a closure that captures the threshold. This flexibility is what makes Python AI code composable and reusable.

Decorators apply cross-cutting concerns - logging, timing, retrying, caching - without modifying the functions they wrap. The three decorators every AI engineer needs are: @retry for resilient API calls, @log_and_time for observability, and @functools.lru_cache for caching repeated computations. Writing these once and applying them with a single line (@retry(max_attempts=3)) is Python's alternative to the boilerplate that other languages require.

Generators produce values lazily, enabling memory-efficient processing of large datasets. When you are embedding 100,000 documents, streaming LLM responses, or iterating over paginated API results, generators ensure that only one item is in memory at a time. The yield keyword is the signal: a function that yields is a generator, and it pauses after each yield until the next value is requested.

Back to AI Engineering