Overview

Middleware in Magma allows you to intercept and modify agent behavior at key points in the execution flow. Using middleware, you can add custom logic that runs before or after completions are generated and tools are executed. You can manipulate the inputs of middleware directly, or throw an error to modify the flow of execution.

Middleware Types

Magma supports four types of middleware:

  • preCompletion: Runs before an LLM completion is generated
  • onCompletion: Runs after an LLM completion is generated
  • preToolExecution: Runs before a tool is executed
  • onToolExecution: Runs after a tool is executed
@middleware("preCompletion")
async preCompletionMiddleware(message: MagmaUserMessage) {
    // Modify the message directly
    message.content = message.content.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

@middleware("onCompletion")
async onCompletionMiddleware(message: MagmaAssistantMessage) {
    // Modify the message directly
    message.content = message.content.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

@middleware("preToolExecution")
async preToolExecutionMiddleware(message: MagmaToolCall) {
    // Modify the call directly
    message.fn_name = message.fn_name.toUpperCase();

    // or throw an error to modify the flow of execution
    throw new Error("Tool call contains bad words");
}

@middleware("onToolExecution")
async onToolExecutionMiddleware(message: MagmaToolResult) {
    // Modify the result directly
    message.result = message.result.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Tool response contains bad words");
}

Middleware Error handling

If a middleware function throws an error, the agent flow will be modified to handle it accordingly.

  • preCompletion: The agent will not generate a completion, and the error message will be returned as if it were the agent’s response.
  • onCompletion: The agent will be notified of the error, and a completion will be re-generated.
  • preToolExecution: The agent will not execute the tool, and the error message will be returned as if it were the tool’s response.
  • onToolExecution: The agent will replace the tool’s response with the error message.

Implementation

To add middleware to your agent, use the @middleware decorator:

import { MagmaAgent } from "@pompeii-labs/magma";
import { middleware } from "@pompeii-labs/magma/decorators";

class MyAgent extends MagmaAgent {
    @middleware("preCompletion")
    async validateMessage(message: string) {
        if (message.length > 1000) {
            throw new Error(
                "Message is too long. Please keep responses under 1000 characters."
            );
        }
        return message;
    }
}

Chaining Middleware

You can add as many middleware functions as you want to your agent. Middleware functions of the same type are executed in order of definition.

class MyAgent extends MagmaAgent {

    @middleware("preCompletion")
    async filterContent(message: string) {
        return message.replace(/badword/gi, '');
    }
    
    @middleware("onCompletion")
    async logResponse(message: string) {
        console.log(Agent response: ${message});
        return message;
    }
}

Common Use Cases

Middleware is particularly useful for:

  • Input validation and sanitization
  • Response filtering and modification
  • Logging and monitoring
  • Rate limiting
  • Access control
  • Response formatting

Best Practices

  • Keep middleware functions focused and single-purpose
  • Handle errors appropriately within middleware
  • Use TypeScript types for better type safety
  • Add clear documentation about what each middleware does
  • Consider the order of middleware execution
  • Return meaningful error messages when validation fails