Add memory setting to conversation handler (#2143)

* Add memory setting to conversation handler

* pr feedback

Author: Kamil Sobol

.changeset/red-candles-fly.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/backend-function': patch
+---
+
+Fix jsdocs that incorrectly state default memory settings

.changeset/tiny-ears-fly.md

@@ -0,0 +1,6 @@
+---
+'@aws-amplify/ai-constructs': patch
+'@aws-amplify/backend-ai': patch
+---
+
+Add memory setting to conversation handler

packages/ai-constructs/API.md

@@ -54,6 +54,7 @@ type ConversationHandlerFunctionProps = {
         modelId: string;
         region?: string;
     }>;
+    memoryMB?: number;
     outputStorageStrategy?: BackendOutputStorageStrategy<AIConversationOutput>;
 };
 

packages/ai-constructs/src/conversation/conversation_handler_construct.test.ts

@@ -222,4 +222,66 @@ void describe('Conversation Handler Function construct', () => {
       Handler: 'index.handler',
     });
   });
+
+  void describe('memory property', () => {
+    void it('sets valid memory', () => {
+      const app = new App();
+      const stack = new Stack(app);
+      new ConversationHandlerFunction(stack, 'conversationHandler', {
+        models: [],
+        memoryMB: 234,
+      });
+      const template = Template.fromStack(stack);
+
+      template.hasResourceProperties('AWS::Lambda::Function', {
+        MemorySize: 234,
+      });
+    });
+
+    void it('sets default memory', () => {
+      const app = new App();
+      const stack = new Stack(app);
+      new ConversationHandlerFunction(stack, 'conversationHandler', {
+        models: [],
+      });
+      const template = Template.fromStack(stack);
+
+      template.hasResourceProperties('AWS::Lambda::Function', {
+        MemorySize: 512,
+      });
+    });
+
+    void it('throws on memory below 128 MB', () => {
+      assert.throws(() => {
+        const app = new App();
+        const stack = new Stack(app);
+        new ConversationHandlerFunction(stack, 'conversationHandler', {
+          models: [],
+          memoryMB: 127,
+        });
+      }, new Error('memoryMB must be a whole number between 128 and 10240 inclusive'));
+    });
+
+    void it('throws on memory above 10240 MB', () => {
+      assert.throws(() => {
+        const app = new App();
+        const stack = new Stack(app);
+        new ConversationHandlerFunction(stack, 'conversationHandler', {
+          models: [],
+          memoryMB: 10241,
+        });
+      }, new Error('memoryMB must be a whole number between 128 and 10240 inclusive'));
+    });
+
+    void it('throws on fractional memory', () => {
+      assert.throws(() => {
+        const app = new App();
+        const stack = new Stack(app);
+        new ConversationHandlerFunction(stack, 'conversationHandler', {
+          models: [],
+          memoryMB: 256.2,
+        });
+      }, new Error('memoryMB must be a whole number between 128 and 10240 inclusive'));
+    });
+  });
 });

packages/ai-constructs/src/conversation/conversation_handler_construct.ts

@@ -34,6 +34,12 @@ export type ConversationHandlerFunctionProps = {
     modelId: string;
     region?: string;
   }>;
+  /**
+   * An amount of memory (RAM) to allocate to the function between 128 and 10240 MB.
+   * Must be a whole number.
+   * Default is 512MB.
+   */
+  memoryMB?: number;
   /**
    * @internal
    */
@@ -86,6 +92,7 @@ export class ConversationHandlerFunction
         timeout: Duration.seconds(60),
         entry: this.props.entry ?? defaultHandlerFilePath,
         handler: 'handler',
+        memorySize: this.resolveMemory(),
         bundling: {
           // Do not bundle SDK if conversation handler is using our default implementation which is
           // compatible with Lambda provided SDK.
@@ -153,4 +160,27 @@ export class ConversationHandlerFunction
       },
     });
   };
+
+  private resolveMemory = () => {
+    const memoryMin = 128;
+    const memoryMax = 10240;
+    const memoryDefault = 512;
+    if (this.props.memoryMB === undefined) {
+      return memoryDefault;
+    }
+    if (
+      !isWholeNumberBetweenInclusive(this.props.memoryMB, memoryMin, memoryMax)
+    ) {
+      throw new Error(
+        `memoryMB must be a whole number between ${memoryMin} and ${memoryMax} inclusive`
+      );
+    }
+    return this.props.memoryMB;
+  };
 }
+
+const isWholeNumberBetweenInclusive = (
+  test: number,
+  min: number,
+  max: number
+) => min <= test && test <= max && test % 1 === 0;

packages/backend-ai/API.md

@@ -53,6 +53,7 @@ type DefineConversationHandlerFunctionProps = {
         modelId: string | AiModel;
         region?: string;
     }>;
+    memoryMB?: number;
 };
 
 // @public (undocumented)

packages/backend-ai/src/conversation/factory.test.ts

@@ -188,4 +188,19 @@ void describe('ConversationHandlerFactory', () => {
       });
     });
   });
+
+  void it('passes memory setting to construct', () => {
+    const factory = defineConversationHandlerFunction({
+      entry: './test-assets/with-default-entry/handler.ts',
+      name: 'testHandlerName',
+      models: [],
+      memoryMB: 271,
+    });
+    const lambda = factory.getInstance(getInstanceProps);
+    const template = Template.fromStack(Stack.of(lambda.resources.lambda));
+    template.resourceCountIs('AWS::Lambda::Function', 1);
+    template.hasResourceProperties('AWS::Lambda::Function', {
+      MemorySize: 271,
+    });
+  });
 });

packages/backend-ai/src/conversation/factory.ts

@@ -43,6 +43,7 @@ class ConversationHandlerFunctionGenerator
         };
       }),
       outputStorageStrategy: this.outputStorageStrategy,
+      memoryMB: this.props.memoryMB,
     };
     const conversationHandlerFunction = new ConversationHandlerFunction(
       scope,
@@ -119,6 +120,12 @@ export type DefineConversationHandlerFunctionProps = {
     modelId: string | AiModel;
     region?: string;
   }>;
+  /**
+   * An amount of memory (RAM) to allocate to the function between 128 and 10240 MB.
+   * Must be a whole number.
+   * Default is 512MB.
+   */
+  memoryMB?: number;
 };
 
 /**

packages/backend-function/src/factory.ts

@@ -104,7 +104,7 @@ export type FunctionProps = {
   /**
    * An amount of memory (RAM) to allocate to the function between 128 and 10240 MB.
    * Must be a whole number.
-   * Default is 128MB.
+   * Default is 512MB.
    */
   memoryMB?: number;