Add receive decorator

Author: coachjackferg

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pompeii-labs/magma",
-  "version": "1.8.3",
+  "version": "1.9.0",
   "description": "The Typescript framework to build AI agents quickly and easily",
   "keywords": [
     "Agents",

src/agent.ts

@@ -23,12 +23,14 @@ import {
     MagmaToolCallBlock,
     MagmaToolCall,
     MagmaSendFunction,
+    MagmaReceive,
 } from './types';
 import { Provider } from './providers';
 import {
     loadHooks,
     loadJobs,
     loadMiddleware,
+    loadReceivers,
     loadTools,
     parseErrorToError,
     parseErrorToString,
@@ -106,11 +108,6 @@ export class MagmaAgent {
      */
     public async receive?(message: any, send: MagmaSendFunction): Promise<void> {}
 
-    /**
-     * Sends data to the connected client depending on the medium (ws, SSE, etc)
-     * @param message any data object to be sent to the client
-     */
-
     public async onWsClose(code: number, reason?: string): Promise<void> {}
 
     public async cleanup(): Promise<void> {
@@ -148,6 +145,15 @@ export class MagmaAgent {
      * @throws Will rethrow the error if no `onError` handler is defined
      */
     public async main(args: {
+        config?: MagmaProviderConfig;
+        send: MagmaSendFunction;
+        userMessage: MagmaUserMessage;
+        onTrace?: (trace: TraceEvent[]) => void;
+    }) {
+        return await this._main(args);
+    }
+
+    private async _main(args: {
         config?: MagmaProviderConfig;
         send: MagmaSendFunction;
         userMessage: MagmaUserMessage;
@@ -290,7 +296,7 @@ export class MagmaAgent {
 
                             // Trigger another completion with the error message as context
                             return resolve(
-                                await this.main({
+                                await this._main({
                                     config,
                                     send,
                                     messages: localMessages,
@@ -335,7 +341,7 @@ export class MagmaAgent {
                                 );
 
                                 return resolve(
-                                    await this.main({
+                                    await this._main({
                                         config,
                                         send,
                                         messages: localMessages,
@@ -376,7 +382,7 @@ export class MagmaAgent {
 
                             // Trigger another completion with the tool result because last message was a tool call
                             return resolve(
-                                await this.main({
+                                await this._main({
                                     config,
                                     send,
                                     messages: localMessages,
@@ -410,7 +416,7 @@ export class MagmaAgent {
                             );
 
                             return resolve(
-                                await this.main({
+                                await this._main({
                                     config,
                                     send,
                                     messages: localMessages,
@@ -1283,6 +1289,10 @@ export class MagmaAgent {
         return [];
     }
 
+    public getReceivers(): MagmaReceive[] {
+        return [];
+    }
+
     public get tools(): MagmaTool[] {
         const agentTools = loadTools(this);
         const loadedTools = this.getTools();
@@ -1314,6 +1324,13 @@ export class MagmaAgent {
         return agentJobs.concat(loadedJobs).concat(utilityJobs);
     }
 
+    public get receivers(): MagmaReceive[] {
+        const agentReceivers = loadReceivers(this);
+        const loadedReceivers = this.getReceivers();
+        const utilityReceivers = this.utilities.flatMap((u) => u.receivers.filter(Boolean));
+        return agentReceivers.concat(loadedReceivers).concat(utilityReceivers);
+    }
+
     /* EVENT HANDLERS */
 
     getSystemPrompts(): MagmaSystemMessageType[] {

src/decorators.ts

@@ -205,3 +205,21 @@ export function job(cron: string, options: { timezone?: string } = {}) {
         descriptor.value._options = options;
     };
 }
+
+export function receive(messageType: string) {
+    return function <R extends void>(
+        target: object,
+        propertyKey: string,
+        descriptor: TypedPropertyDescriptor<
+            ((data: unknown, send: MagmaSendFunction, agent: MagmaAgent) => R) & {
+                _messageType?: string;
+            }
+        >
+    ) {
+        if (!descriptor.value) {
+            throw new Error('Receive decorator must be used on a function');
+        }
+
+        descriptor.value._messageType = messageType;
+    };
+}

src/helpers/utilities.ts

@@ -2,6 +2,7 @@ import {
     MagmaHook,
     MagmaJob,
     MagmaMiddleware,
+    MagmaReceive,
     MagmaTool,
     MagmaToolParam,
     MagmaUtilities,
@@ -44,8 +45,9 @@ export function loadUtilities(target: any): MagmaUtilities {
     const hooks = loadHooks(target);
     const jobs = loadJobs(target);
     const middleware = loadMiddleware(target);
+    const receivers = loadReceivers(target);
 
-    return { tools, hooks, jobs, middleware };
+    return { tools, hooks, jobs, middleware, receivers };
 }
 
 /**
@@ -156,3 +158,27 @@ export function loadMiddleware(target: any): MagmaMiddleware[] {
 
     return middleware;
 }
+
+/**
+ * Helper function to load middleware from a class or instance of a class
+ * If the target is an instance, it will load both the static and instance middleware
+ * @param target class or instance of a class to load middleware from
+ * @returns array of MagmaMiddleware objects
+ */
+export function loadReceivers(target: any): MagmaReceive[] {
+    const receivers: MagmaReceive[] = [];
+
+    const { staticMethods, instanceMethods } = getMethodsFromClassOrInstance(target);
+    const methods = [...staticMethods, ...instanceMethods];
+
+    for (const method of methods) {
+        if (typeof method === 'function' && '_messageType' in method) {
+            receivers.push({
+                handler: method.bind(target),
+                messageType: (method as any)['_messageType'],
+            });
+        }
+    }
+
+    return receivers;
+}

src/types/utilities/index.ts

@@ -1,16 +1,19 @@
 import { MagmaHook } from './hooks';
 import { MagmaJob } from './jobs';
 import { MagmaMiddleware } from './middleware';
+import { MagmaReceive } from './receive';
 import { MagmaTool } from './tools';
 
 export * from './hooks';
 export * from './jobs';
 export * from './middleware';
 export * from './tools';
+export * from './receive';
 
 export type MagmaUtilities = {
     tools: MagmaTool[];
     middleware: MagmaMiddleware[];
     hooks: MagmaHook[];
     jobs: MagmaJob[];
+    receivers: MagmaReceive[];
 };

src/types/utilities/receive.ts

@@ -0,0 +1,16 @@
+import { MagmaAgent } from '../../agent';
+import { MagmaSendFunction } from '../agent';
+
+export type MagmaWSMessage = {
+    type: string;
+    data: unknown;
+};
+
+export type MagmaReceive = {
+    handler: (
+        wsMessage: MagmaWSMessage,
+        send: MagmaSendFunction,
+        agent: MagmaAgent
+    ) => Promise<void> | void;
+    messageType: string;
+};