Skip to content

Module Callbacks API

API reference for module callbacks - the mechanism that enables dynamic tool management during execution.

Overview

Module callbacks are special callables that tools can return to modify the module during execution. They receive context about the current execution state and can add/remove tools, modify configuration, or perform other operations on the module.


@module_callback Decorator

@module_callback
def callback(context: ModuleContext) -> str:
    # Modify module
    context.module.init_module(tools=[...])
    return "Observation string"

Decorator that marks a function as a module callback. The decorated function: - Receives a ModuleContext (or subclass) as its only parameter - Must return a string that becomes the observation in the trajectory - Can modify the module by calling context.module.init_module()

Parameters:

  • func (Callable[[ModuleContext], str]): Function that takes context and returns observation string

Returns:

  • ModuleCallback: Wrapped callback object

Example:

from udspy import module_callback

@module_callback
def add_tools(context):
    # Access the module
    current_tools = list(context.module.tools.values())

    # Add new tools
    context.module.init_module(tools=current_tools + [new_tool])

    # Return observation
    return "Added new tool"

ModuleCallback Class

class ModuleCallback:
    """Wrapper for module callback functions."""

    func: Callable[[ModuleContext], str]

    def __call__(self, context: ModuleContext) -> str: ...

The ModuleCallback class wraps a callback function. You typically don't instantiate this directly; use the @module_callback decorator instead.

Attributes

func

func: Callable[[ModuleContext], str]

The underlying callback function.

Methods

__call__(context: ModuleContext) -> str

Execute the callback with the provided context.

Parameters:

  • context (ModuleContext): Execution context with module and state

Returns:

  • str: Observation string to add to trajectory/history

Context Classes

ModuleContext

class ModuleContext:
    """Base context for module callbacks."""

    module: Module

Base context class that provides access to the module instance.

Attributes:

  • module (Module): The module instance (Predict or ReAct)

Example:

@module_callback
def callback(context: ModuleContext):
    # Access module
    module = context.module

    # Access current tools
    tools = context.module.tools

    # Modify module
    context.module.init_module(tools=[...])

    return "Modified module"

PredictContext

class PredictContext(ModuleContext):
    """Context for Predict module callbacks."""

    module: Predict
    history: Optional[History]

Context for callbacks in Predict modules. Provides access to conversation history in addition to the module.

Attributes:

  • module (Predict): The Predict module instance
  • history (Optional[History]): Conversation history (if provided)

Example:

from udspy.module.callbacks import PredictContext

@module_callback
def callback(context: PredictContext):
    # Access Predict-specific features
    module = context.module  # Type: Predict

    # Access conversation history
    if context.history:
        messages = context.history.messages
        print(f"Conversation has {len(messages)} messages")

    # Modify tools
    context.module.init_module(tools=[...])

    return "Updated tools based on conversation history"

ReactContext

class ReactContext(ModuleContext):
    """Context for ReAct module callbacks."""

    module: ReAct
    trajectory: dict[str, Any]

Context for callbacks in ReAct modules. Provides access to the agent's trajectory in addition to the module.

Attributes:

  • module (ReAct): The ReAct module instance
  • trajectory (dict[str, Any]): Current trajectory with thought/action/observation history

Example:

from udspy.module.callbacks import ReactContext

@module_callback
def callback(context: ReactContext):
    # Access ReAct-specific features
    module = context.module  # Type: ReAct

    # Access trajectory
    trajectory = context.trajectory

    # Analyze agent's thoughts
    thoughts = [v for k, v in trajectory.items() if k.startswith("thought_")]
    print(f"Agent has had {len(thoughts)} thoughts")

    # Modify tools based on trajectory
    if len(thoughts) > 3:
        # Agent is struggling, add more tools
        context.module.init_module(tools=[...])

    return "Adapted tools based on trajectory"

Trajectory Structure:

The trajectory dictionary contains entries like:

{
    "thought_0": "I need to search for information",
    "tool_calls_0": [{"name": "search", "args": {...}}],
    "observation_0": "Search results: ...",
    "thought_1": "Now I need to calculate",
    "tool_calls_1": [{"name": "calculator", "args": {...}}],
    "observation_1": "Result: 42",
    # ...
}

Helper Functions

is_module_callback(obj: Any) -> bool

def is_module_callback(obj: Any) -> bool:
    """Check if an object is a module callback."""

Check if an object is a ModuleCallback instance.

Parameters:

  • obj (Any): Object to check

Returns:

  • bool: True if obj is a ModuleCallback instance

Example:

from udspy import module_callback, is_module_callback

@module_callback
def my_callback(context):
    return "Done"

def regular_function():
    return "Done"

print(is_module_callback(my_callback))      # True
print(is_module_callback(regular_function))  # False
print(is_module_callback("not a function"))  # False

Usage with Tools

Module callbacks are typically returned by tools to enable dynamic tool loading:

from udspy import tool, module_callback
from pydantic import Field

@tool(name="load_calculator", description="Load calculator tool")
def load_calculator() -> callable:
    """Tool that returns a module callback."""

    @module_callback
    def add_calculator(context):
        # Get current tools
        current = [
            t for t in context.module.tools.values()
            if t.name not in ("finish", "user_clarification")
        ]

        # Add calculator
        from my_tools import calculator
        context.module.init_module(tools=current + [calculator])

        return "Calculator loaded and ready to use"

    return add_calculator

Execution Flow:

  1. LLM calls load_calculator() tool
  2. Tool returns a ModuleCallback instance
  3. Module detects the callback with is_module_callback()
  4. Module creates appropriate context (PredictContext or ReactContext)
  5. Module executes callback: callback(context)
  6. Callback modifies module via context.module.init_module()
  7. Callback returns observation string
  8. Observation is added to history/trajectory
  9. Execution continues with modified tool set

Common Patterns

Loading Tools on Demand

@tool(name="load_tools", description="Load specialized tools")
def load_tools(category: str = Field(...)) -> callable:
    @module_callback
    def add_tools(context):
        current = list(context.module.tools.values())

        if category == "math":
            new_tools = [calculator, statistics]
        elif category == "web":
            new_tools = [search, scrape]

        context.module.init_module(tools=current + new_tools)
        return f"Loaded {len(new_tools)} {category} tools"

    return add_tools

Conditional Loading Based on History

@tool(name="smart_load", description="Intelligently load tools")
def smart_load() -> callable:
    @module_callback
    def analyze_and_load(context: PredictContext):
        # Analyze conversation history
        if context.history:
            messages = context.history.messages
            # Determine what tools are needed
            # ...

        context.module.init_module(tools=[...])
        return "Loaded appropriate tools"

    return analyze_and_load

Progressive Tool Discovery

@tool(name="discover_tools", description="Discover needed tools")
def discover_tools(task: str = Field(...)) -> callable:
    @module_callback
    def discover(context: ReactContext):
        # Analyze trajectory to see what agent has tried
        thoughts = [v for k, v in context.trajectory.items()
                   if k.startswith("thought_")]

        # Load tools the agent seems to need
        # ...

        context.module.init_module(tools=[...])
        return "Discovered and loaded needed tools"

    return discover

Type Annotations

from typing import Callable, Any
from udspy import Module, History

# Decorator
def module_callback(
    func: Callable[[ModuleContext], str]
) -> ModuleCallback: ...

# Callback class
class ModuleCallback:
    func: Callable[[ModuleContext], str]
    def __call__(self, context: ModuleContext) -> str: ...

# Context classes
class ModuleContext:
    module: Module

class PredictContext(ModuleContext):
    module: Predict
    history: Optional[History]

class ReactContext(ModuleContext):
    module: ReAct
    trajectory: dict[str, Any]

# Helper
def is_module_callback(obj: Any) -> bool: ...

Important Notes

  1. Return String Required: Callbacks MUST return a string - this becomes the observation

  2. Thread Safety: Callbacks execute synchronously during tool execution

  3. Built-in Preservation: ReAct's finish and user clarification tools are automatically preserved

  4. Tool Persistence: Tools added via callbacks remain available for the entire execution

  5. Error Handling: Wrap callback logic in try/except to handle errors gracefully:

    @module_callback
    def safe_callback(context):
        try:
            context.module.init_module(tools=[...])
            return "Success"
        except Exception as e:
            return f"Failed: {e}"
    


See Also

Example Code: See examples/dynamic_calculator.py and examples/dynamic_tools.py in the GitHub repository