Skip to content

Modules

Modules are composable units that encapsulate LLM calls. They provide a standard interface for building complex LLM-powered applications.

Overview

All modules inherit from the base Module class and implement a unified execution pattern. This allows modules to be composed, nested, and combined to create sophisticated behaviors.

Core Concepts

Unified Interface

Every module implements: - aexecute(*, stream: bool, **inputs) - Core async execution - aforward(**inputs) - Async convenience (no streaming) - __call__(**inputs) - Synchronous wrapper

Composition

Modules can contain other modules:

from udspy import Module, Predict, ChainOfThought, Prediction

class Pipeline(Module):
    def __init__(self):
        self.analyze = Predict("text -> analysis")
        self.summarize = ChainOfThought("text, analysis -> summary")

    async def aexecute(self, *, stream: bool = False, **inputs):
        # First module: get analysis (stream=False since we need full result)
        analysis = None
        async for event in self.analyze.aexecute(stream=False, text=inputs["text"]):
            if isinstance(event, Prediction):
                analysis = event

        if not analysis:
            raise ValueError("First module did not produce a result")

        # Second module: pass down stream parameter
        async for event in self.summarize.aexecute(
            stream=stream,
            text=inputs["text"],
            analysis=analysis.analysis
        ):
            yield event

Streaming Support

All modules support streaming for real-time output:

async for event in module.aexecute(stream=True, **inputs):
    if isinstance(event, OutputStreamChunk):
        print(event.delta, end="", flush=True)
    elif isinstance(event, Prediction):
        print(f"\nFinal: {event}")

Dynamic Tool Management

Modules support runtime modification of their toolset via the init_module() method. This enables: - Loading tools on demand (reduce initial context size) - Progressive tool discovery (agent figures out what it needs) - Adaptive behavior (add tools based on task complexity)

Key method: init_module(tools=None)

This method is essential for dynamic tools because when the toolset changes, the module needs to fully reconfigure itself: - Regenerate tool schemas (so the LLM knows how to call new tools) - Rebuild signatures (so tool descriptions appear in the prompt) - Update the tool registry (so the module can execute the tools)

It's typically called from within a module callback:

from udspy import tool, module_callback, ReAct

@tool(name="load_calculator")
def load_calculator() -> callable:
    """Load calculator tool when needed."""

    @module_callback
    def add_calculator(context):
        # Get current tools
        current = list(context.module.tools.values())

        # Add new tool
        context.module.init_module(tools=current + [calculator])

        return "Calculator loaded"

    return add_calculator

# Agent starts with only the loader
agent = ReAct(signature, tools=[load_calculator])

# Agent loads calculator when needed
result = agent(question="What is 157 * 834?")

See: Dynamic Tool Management Guide for detailed examples and patterns.

Built-in Modules

udspy provides three core modules:

Base Module

The foundation for all modules. Provides: - Unified execution interface - Streaming infrastructure - Async-first design - Composition support

When to use: When creating custom modules

Predict

The core module for LLM predictions. Features: - Maps signature inputs to outputs - Native tool calling support - Conversation history management - Streaming and async execution

When to use: For basic LLM calls, tool usage, and as a building block for other modules

predictor = Predict("question -> answer")
result = predictor(question="What is AI?")

ChainOfThought

Adds step-by-step reasoning before outputs. Features: - Automatic reasoning field injection - Improves answer quality on complex tasks - Transparent reasoning process - Works with any signature

When to use: For tasks requiring reasoning (math, analysis, decision-making)

cot = ChainOfThought("question -> answer")
result = cot(question="What is 157 * 234?")
print(result.reasoning)  # Shows step-by-step work
print(result.answer)     # "36738"

ReAct

Agent that reasons and acts with tools. Features: - Iterative reasoning and tool usage - Human-in-the-loop support - Built-in user_clarification and finish tools - Full trajectory tracking

When to use: For tasks requiring multiple steps, tool usage, or agent-like behavior

@tool(name="search")
def search(query: str = Field(...)) -> str:
    return search_web(query)

agent = ReAct("question -> answer", tools=[search])
result = agent(question="What's the weather in Tokyo?")

Module Comparison

Feature Predict ChainOfThought ReAct
Basic LLM calls
Step-by-step reasoning
Tool usage
Multi-step iteration
Human-in-the-loop
Trajectory tracking
Complexity Low Low Medium
Token usage Low Medium High
Latency Low Medium High

Creating Custom Modules

To create a custom module:

  1. Subclass Module
  2. Implement aexecute() method
  3. Yield StreamEvent objects during execution
  4. Yield final Prediction at the end
from udspy import Module, Prediction, OutputStreamChunk

class CustomModule(Module):
    def __init__(self, signature):
        self.predictor = Predict(signature)

    async def aexecute(self, *, stream: bool = False, **inputs):
        # Custom logic before prediction
        processed_inputs = preprocess(inputs)

        # Call nested module's aexecute, passing down stream parameter
        async for event in self.predictor.aexecute(stream=stream, **processed_inputs):
            if isinstance(event, Prediction):
                # Custom logic after prediction
                final_result = postprocess(event)
                # Yield final prediction
                yield Prediction(**final_result)
            else:
                # Pass through other events (OutputStreamChunks, etc.)
                yield event

See Base Module for detailed guidance.

Best Practices

Choose the Right Module

  • Simple tasks: Use Predict
  • Need reasoning: Use ChainOfThought
  • Need tools/agents: Use ReAct
  • Custom logic: Create custom module

Composition over Inheritance

Build complex behaviors by composing modules rather than deep inheritance:

# Good: Composition
class Pipeline(Module):
    def __init__(self):
        self.step1 = Predict(sig1)
        self.step2 = ChainOfThought(sig2)

# Avoid: Deep inheritance
class MyComplexModule(ChainOfThought):
    # Complex overrides

Async Best Practices

  • Use aforward() when you don't need streaming
  • Use aexecute(stream=True) for real-time output
  • Use __call__() only in sync contexts
  • Always await async operations

Error Handling

try:
    result = await module.aforward(**inputs)
except ConfirmationRequired as e:
    # Handle confirmations
    result = await module.aresume(user_input, e)
except Exception as e:
    # Handle other errors
    logger.error(f"Module failed: {e}")

See Also