Creating Custom Amplifier Modules

This guide shows how to create custom modules for Amplifier, based on real experience building the email assistant example. We'll cover the patterns, pitfalls, and best practices learned from actual implementation.

When to Create a Module vs Inline Code

Module Structure (Standard Pattern)

your-module/
├── pyproject.toml                      # Package metadata
├── amplifier_module_tool_yourname/
│   ├── __init__.py                     # mount() function
│   ├── tool.py                         # Tool implementation
│   └── manager.py                      # State/business logic (if needed)
└── README.md

Example: pyproject.toml

[project]
name = "amplifier-module-tool-email"
version = "1.0.0"
description = "Email management for Amplifier"
requires-python = ">=3.11"
dependencies = []

[project.entry-points."amplifier.modules"]
tool-email = "amplifier_module_tool_email:mount"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

# In bundle.md
tools:
  - module: tool-email
    source: ./modules/tool-email
    config:
      db_path: ../data/emails.db  # Relative to where module runs

The mount() Function (Entry Point)

Every module must have a mount() function:

async def mount(coordinator, config: dict | None = None):
    """Mount the module.

    Args:
        coordinator: ModuleCoordinator from amplifier-core
        config: Configuration from bundle

    Returns:
        Optional cleanup function
    """
    config = config or {}

    # Create your tool/hook/provider
    tool = EmailTool(config)

    # Register with coordinator
    await coordinator.mount("tools", tool, name=tool.name)

    return  # Or return cleanup function

Pattern: Tool with Embedded State

For tools that need to maintain state (like database connections, API clients), use the embedded manager pattern. This is how tool-issue works.

Manager Class (Holds State)

# manager.py
import sqlite3

class EmailManager:
    """Manages email state: IMAP connection, database, operations."""

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self.imap = None
        self._init_db()

    def _init_db(self):
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS emails (
                id TEXT PRIMARY KEY,
                sender TEXT,
                subject TEXT,
                body TEXT
            )
        """)

    def connect_imap(self, server, email, password):
        import imaplib
        self.imap = imaplib.IMAP4_SSL(server)
        self.imap.login(email, password)

    def fetch_emails(self, limit=20):
        # IMAP operations...
        pass

    def get_email(self, email_id):
        cursor = self.conn.execute(
            "SELECT * FROM emails WHERE id = ?", (email_id,)
        )
        return cursor.fetchone()

Tool Class (Exposes Operations)

# tool.py
from amplifier_core import ToolResult
from .manager import EmailManager

class EmailTool:
    """Tool that wraps EmailManager and exposes operations to AI."""

    name = "email"
    description = "Manage emails: fetch, reply, archive"

    def __init__(self, email_manager: EmailManager):
        self.email_manager = email_manager

    @property
    def input_schema(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["fetch", "get", "reply", "archive"],
                    "description": "Operation to perform"
                },
                "params": {
                    "type": "object",
                    "description": "Parameters for operation"
                }
            },
            "required": ["operation"]
        }

    async def execute(self, input: dict) -> ToolResult:
        operation = input.get("operation")
        params = input.get("params", {})

        if operation == "fetch":
            emails = self.email_manager.fetch_emails(
                limit=params.get("limit", 20)
            )
            return ToolResult(
                output=f"Fetched {len(emails)} emails",
                data={"emails": emails},
                success=True
            )
        elif operation == "get":
            email = self.email_manager.get_email(params.get("email_id"))
            if not email:
                return ToolResult(
                    success=False,
                    error={"message": "Email not found"}
                )
            return ToolResult(
                output="Email retrieved",
                data=email,
                success=True
            )
        # ... more operations

Mount Function (Wires It Together)

# __init__.py
from .manager import EmailManager
from .tool import EmailTool

async def mount(coordinator, config: dict | None = None):
    """Mount email tool with embedded state."""
    config = config or {}

    # Create manager (holds state)
    db_path = config.get("db_path", "emails.db")
    email_manager = EmailManager(db_path)

    # Create tool (exposes to AI)
    tool = EmailTool(email_manager)

    # Register with coordinator
    await coordinator.mount("tools", tool, name=tool.name)

    return

Critical Learning: Accessing Tools from Your App

# Your web app
from my_module.manager import EmailManager

email_mgr = EmailManager()  # Creates NEW instance
# This is separate from the one in Amplifier sessions!

# Your web app
from amplifier_foundation.registry import load_bundle

# Load and prepare bundle
bundle = await load_bundle("./bundle.md")
prepared_bundle = await bundle.prepare()  # This loads and validates modules

# Create session using PreparedBundle (not manual AmplifierSession)
session = await prepared_bundle.create_session(
    session_id="temp-for-tool-access"
)

# Get the mounted tool from coordinator
email_tool = session.coordinator.get("tools", "email")

# Now use it!
result = await email_tool.execute({
    "operation": "fetch",
    "params": {"limit": 50}
})

Real Example: Email Assistant

The examples/email-assistant/ directory contains a complete, working implementation showing all these patterns in practice.

What It Demonstrates

• Custom module - modules/tool-email/ with EmailManager + EmailTool
• Web app integration - FastAPI backend accesses tool through coordinator
• Shared state - Web API and AI sessions use same EmailManager instance
• Operation-based tool - Single tool with multiple operations (like tool-issue)
• Bundle configuration - Module loaded via bundle.md
• Real Amplifier sessions - Not mocks, actual AmplifierSession with streaming

Key Files to Study

Common Pitfalls & Solutions

Pitfall 1: Wrong Import Path

from amplifier_foundation.registry import load_bundle

Pitfall 2: Wrong Method Name

composed = bundle.compose()
mount_plan = composed.to_mount_plan()

Pitfall 3: Async Loading Not Awaited

bundle = await load_bundle(path)
prepared = await bundle.prepare()  # Also async!
session = await prepared.create_session(...)  # Also async!

Pitfall 4: Wrong Session Method

response = await session.execute(prompt)
# Returns the full response as a string

Pitfall 5: ToolResult Attribute

# Wrong
return ToolResult(data={"emails": emails})
emails = result.data

# Correct
return ToolResult(output={"emails": emails})
emails = result.output

Pitfall 6: Module File Paths

# bundle.md at root, module in modules/tool-email/, want data/ at root:
tools:
  - module: tool-email
    source: ./modules/tool-email
    config:
      db_path: ../data/emails.db  # ../ goes up from module dir

Pitfall 7: Context Field Format

# Wrong
context:
  - file: context/guidelines.md

# Correct - no context field, use @mentions
---
bundle:
  name: my-bundle
---

# Instructions

@my-bundle:context/guidelines.md

Pitfall 5: Accessing Tools Outside Sessions

# In app startup
bundle = await load_bundle("./bundle.md")
prepared = await bundle.prepare()  # Loads all modules

# Create temp session with prepared bundle
temp_session = await prepared.create_session(
    session_id="temp-for-tool-access"
)

# Get tool from coordinator using .get()
my_tool = temp_session.coordinator.get("tools", "my-tool-name")

# Store globally for API endpoints
app.state.my_tool = my_tool

# Use in API endpoints
@app.post("/api/endpoint")
async def endpoint():
    result = await app.state.my_tool.execute({...})
    return result.data

Module Configuration in Bundle

Local Module (Development)

tools:
  - module: tool-email
    source: file://./modules/tool-email
    config:
      db_path: data/emails.db

Git Module (Distribution)

tools:
  - module: tool-email
    source: git+https://github.com/you/amplifier-module-tool-email@main
    config:
      db_path: data/emails.db

Testing Your Module

Unit Test the Manager

# test_manager.py
from amplifier_module_tool_email.manager import EmailManager

def test_email_manager():
    mgr = EmailManager(":memory:")  # SQLite in-memory
    mgr.connect_imap("imap.gmail.com", "test@test.com", "pass")
    assert mgr.imap is not None

Integration Test the Tool

# test_tool.py
from amplifier_module_tool_email import mount
from amplifier_core import ModuleCoordinator

async def test_email_tool():
    coordinator = ModuleCoordinator()
    await mount(coordinator, {"db_path": ":memory:"})

    email_tool = coordinator.tools.get("email")
    assert email_tool is not None

    result = await email_tool.execute({
        "operation": "list",
        "params": {"limit": 10}
    })
    assert result.success

End-to-End Test with Session

async def test_with_session():
    from amplifier_foundation.registry import load_bundle

    # Load and prepare your bundle
    bundle = await load_bundle("./bundle.md")
    prepared = await bundle.prepare()

    # Create session (modules are loaded and mounted)
    session = await prepared.create_session(session_id="test")

    # AI can now use your tool!
    response = await session.execute("Fetch my emails")
    print(response)

    # Check that your tool was mounted
    my_tool = session.coordinator.get("tools", "email")
    assert my_tool is not None

Distributing Your Module

Option 1: Local Development

# In bundle.md, reference local path
source: file://./modules/tool-email

Option 2: Git Repository

# Push to GitHub
git init
git add .
git commit -m "Initial module"
git push origin main

# In bundle.md
source: git+https://github.com/you/amplifier-module-tool-email@main

Option 3: PyPI Package

# Build and publish
uv build
uv publish

# In bundle.md (no source needed, just module name)
tools:
  - module: tool-email

Real-World Learnings

Lesson 1: Start Simple, Then Modularize

When building the email assistant, we initially created inline tools (EmailReplyTool, EmailArchiveTool as separate classes). This was faster for prototyping. Once the pattern was clear, we refactored into a proper module with an operation-based tool.

Recommendation: Build inline first, extract to module when stable.

Lesson 2: Operation-Based > Multiple Tools

Initially designed as 3 separate tools (email_reply, email_archive, email_label). Refactored to 1 tool with operations.

Why better:

• Single shared state (EmailManager)
• Consistent error handling
• Easier to extend (add new operations)
• Follows ecosystem pattern (tool-issue works this way)

Lesson 3: Test with Real Amplifier, Not Mocks

Lesson 4: Bundle Configuration Matters

The bundle must include your module with correct source path:

# Relative path from bundle.md location
source: file://./modules/tool-email

# Absolute path
source: file:///home/user/project/modules/tool-email

# Git URL
source: git+https://github.com/owner/repo@main

Lesson 5: Coordinator is the Source of Truth

Don't create parallel instances. If your app needs access to a tool's state:

1. Load bundle: bundle = await load_bundle("./bundle.md")
2. Prepare it: prepared = await bundle.prepare()
3. Create session: session = await prepared.create_session(...)
4. Get tool: tool = session.coordinator.get("tools", "tool-name")
5. Store reference and use throughout your app

Critical: prepare() is what actually loads modules from their sources. Without it, you just have config.

Module Patterns in the Ecosystem

Next Steps

• See complete email assistant code
• amplifier-core examples
• amplifier-foundation examples
• Study existing modules in the ecosystem for patterns