fix(conversation): use functionMap for custom handler IFunction reference

dependency_licenses.txt

@@ -14486,6 +14486,32 @@ SOFTWARE.
 
 -----
 
+The following software may be included in this product: esbuild. A copy of the source code may be downloaded from git+https://github.com/evanw/esbuild.git. This software contains the following license and notice below:
+
+MIT License
+
+Copyright (c) 2020 Evan Wallace
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+-----
+
 The following software may be included in this product: escalade. A copy of the source code may be downloaded from https://github.com/lukeed/escalade.git. This software contains the following license and notice below:
 
 MIT License

packages/amplify-graphql-conversation-transformer/API.md

@@ -9,6 +9,7 @@ import { DirectiveNode } from 'graphql';
 import { FieldDefinitionNode } from 'graphql';
 import { HasManyTransformer } from '@aws-amplify/graphql-relational-transformer';
 import { InterfaceTypeDefinitionNode } from 'graphql';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
 import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
 import { ObjectTypeDefinitionNode } from 'graphql';
 import { TransformerAuthProvider } from '@aws-amplify/graphql-transformer-interfaces';
@@ -19,7 +20,7 @@ import { TransformerSchemaVisitStepContextProvider } from '@aws-amplify/graphql-
 
 // @public (undocumented)
 export class ConversationTransformer extends TransformerPluginBase {
-    constructor(modelTransformer: ModelTransformer, hasManyTransformer: HasManyTransformer, belongsToTransformer: BelongsToTransformer, authProvider: TransformerAuthProvider);
+    constructor(modelTransformer: ModelTransformer, hasManyTransformer: HasManyTransformer, belongsToTransformer: BelongsToTransformer, authProvider: TransformerAuthProvider, functionNameMap?: Record<string, lambda.IFunction>);
     // (undocumented)
     field: (parent: ObjectTypeDefinitionNode | InterfaceTypeDefinitionNode, definition: FieldDefinitionNode, directive: DirectiveNode, context: TransformerSchemaVisitStepContextProvider) => void;
     // (undocumented)

packages/amplify-graphql-conversation-transformer/package.json

@@ -37,7 +37,8 @@
     "immer": "^9.0.12"
   },
   "devDependencies": {
-    "@aws-amplify/graphql-transformer-test-utils": "1.0.3"
+    "@aws-amplify/graphql-transformer-test-utils": "1.0.3",
+    "esbuild": "^0.24.0"
   },
   "peerDependencies": {
     "aws-cdk-lib": "^2.158.0",

packages/amplify-graphql-conversation-transformer/src/__tests__/__snapshots__/amplify-graphql-conversation-transformer.test.ts.snap

@@ -22,11 +22,9 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
   ctx.stash.connectionAttributes = {};
   ctx.stash.lambdaFunctionArn = "",
       {
-        "Fn::Sub": [
-          "arn:aws:lambda:\${AWS::Region}:\${AWS::AccountId}:function:\${name}",
-          {
-            "name": "conversation-handler",
-          },
+        "Fn::GetAtt": [
+          "PirateChatConversationDirectiveLambdaStack",
+          "Outputs.transformerrootstackPirateChatConversationDirectiveLambdaStackPirateChatDefaultConversationHandlerconversationHandlerFunction2B526F1AArn",
         ],
       },
       "";
@@ -144,11 +142,9 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
   ctx.stash.connectionAttributes = {};
   ctx.stash.lambdaFunctionArn = "",
       {
-        "Fn::Sub": [
-          "arn:aws:lambda:\${AWS::Region}:\${AWS::AccountId}:function:\${name}",
-          {
-            "name": "conversation-handler",
-          },
+        "Fn::GetAtt": [
+          "PirateChatConversationDirectiveLambdaStack",
+          "Outputs.transformerrootstackPirateChatConversationDirectiveLambdaStackPirateChatDefaultConversationHandlerconversationHandlerFunction2B526F1AArn",
         ],
       },
       "";
@@ -268,11 +264,9 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
   ctx.stash.connectionAttributes = {};
   ctx.stash.lambdaFunctionArn = "",
       {
-        "Fn::Sub": [
-          "arn:aws:lambda:\${AWS::Region}:\${AWS::AccountId}:function:\${name}",
-          {
-            "name": "conversation-handler",
-          },
+        "Fn::GetAtt": [
+          "PirateChatConversationDirectiveLambdaStack",
+          "Outputs.transformerrootstackPirateChatConversationDirectiveLambdaStackPirateChatDefaultConversationHandlerconversationHandlerFunction2B526F1AArn",
         ],
       },
       "";
@@ -392,11 +386,9 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
   ctx.stash.connectionAttributes = {};
   ctx.stash.lambdaFunctionArn = "",
       {
-        "Fn::Sub": [
-          "arn:aws:lambda:\${AWS::Region}:\${AWS::AccountId}:function:\${name}",
-          {
-            "name": "conversation-handler",
-          },
+        "Fn::GetAtt": [
+          "PirateChatConversationDirectiveLambdaStack",
+          "Outputs.transformerrootstackPirateChatConversationDirectiveLambdaStackPirateChatDefaultConversationHandlerconversationHandlerFunction2B526F1AArn",
         ],
       },
       "";

packages/amplify-graphql-conversation-transformer/src/__tests__/amplify-graphql-conversation-transformer.test.ts

@@ -3,12 +3,13 @@ import { IndexTransformer, PrimaryKeyTransformer } from '@aws-amplify/graphql-in
 import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
 import { validateModelSchema } from '@aws-amplify/graphql-transformer-core';
 import { AppSyncAuthConfiguration, ModelDataSourceStrategy } from '@aws-amplify/graphql-transformer-interfaces';
-import { DeploymentResources, testTransform } from '@aws-amplify/graphql-transformer-test-utils';
+import { DeploymentResources, testTransform, TransformManager } from '@aws-amplify/graphql-transformer-test-utils';
 import { parse } from 'graphql';
 import { ConversationTransformer } from '..';
 import { BelongsToTransformer, HasManyTransformer, HasOneTransformer } from '@aws-amplify/graphql-relational-transformer';
 import * as fs from 'fs-extra';
 import * as path from 'path';
+import { Code, Function, IFunction, Runtime } from 'aws-cdk-lib/aws-lambda';
 import { GenerationTransformer } from '@aws-amplify/graphql-generation-transformer';
 import { toUpper } from 'graphql-transformer-common';
 
@@ -44,6 +45,35 @@ describe('ConversationTransformer', () => {
       const schema = parse(out.schema);
       validateModelSchema(schema);
     });
+
+    it('uses functionMap for custom handler', () => {
+      const routeName = 'pirateChat';
+      const inputSchema = getSchema('conversation-route-custom-handler.graphql', { ROUTE_NAME: routeName });
+
+      const transformerManager = new TransformManager();
+      const stack = transformerManager.getTransformScope();
+      const customHandler = new Function(stack, 'conversation-handler', {
+        runtime: Runtime.NODEJS_18_X,
+        code: Code.fromInline('exports.handler = async (event) => { return "Hello World"; }'),
+        handler: 'index.handler',
+      });
+
+      const functionMap = {
+        [`Fn${routeName}`]: customHandler,
+      };
+
+      const out = transform(inputSchema, {}, defaultAuthConfig, functionMap, transformerManager);
+      expect(out).toBeDefined();
+
+      const expectedCustomHandlerArn = out.rawRootStack.resolve(customHandler.functionArn);
+      const conversationLambdaStackName = `${toUpper(routeName)}ConversationDirectiveLambdaStack`;
+      const conversationLambdaDataSourceName = `Fn${routeName}LambdaDataSource`;
+      const conversationLambdaDataSourceFunctionArnRef =
+        out.stacks[conversationLambdaStackName].Resources?.[conversationLambdaDataSourceName].Properties.LambdaConfig.LambdaFunctionArn.Ref;
+      const lambdaDataSourceFunctionArn =
+        out.rootStack.Resources?.[conversationLambdaStackName].Properties?.Parameters?.[conversationLambdaDataSourceFunctionArnRef];
+      expect(lambdaDataSourceFunctionArn).toEqual(expectedCustomHandlerArn);
+    });
   });
 
   describe('invalid schemas', () => {
@@ -114,6 +144,8 @@ function transform(
   inputSchema: string,
   dataSourceStrategies?: Record<string, ModelDataSourceStrategy>,
   authConfig: AppSyncAuthConfiguration = defaultAuthConfig,
+  functionMap?: Record<string, IFunction>,
+  transformerManager?: TransformManager,
 ): DeploymentResources {
   const modelTransformer = new ModelTransformer();
   const authTransformer = new AuthTransformer();
@@ -129,7 +161,7 @@ function transform(
     hasManyTransformer,
     hasOneTransformer,
     belongsToTransformer,
-    new ConversationTransformer(modelTransformer, hasManyTransformer, belongsToTransformer, authTransformer),
+    new ConversationTransformer(modelTransformer, hasManyTransformer, belongsToTransformer, authTransformer, functionMap),
     new GenerationTransformer(),
     authTransformer,
   ];
@@ -139,6 +171,7 @@ function transform(
     authConfig,
     transformers,
     dataSourceStrategies,
+    transformerManager,
   });
 
   return out;

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-custom-handler.graphql

@@ -0,0 +1,13 @@
+type Mutation {
+  ROUTE_NAME(
+    conversationId: ID!,
+    content: [ContentBlockInput],
+    aiContext: AWSJSON,
+    toolConfiguration: ToolConfigurationInput
+  ): ConversationMessage
+  @conversation(
+    aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
+    functionName: "FnROUTE_NAME",
+    systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
+  )
+}

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-custom-query-tool.graphql

@@ -17,7 +17,6 @@ type Mutation {
   ): ConversationMessage
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
     tools: [{ name: "getTemperature", description: "does a thing" }, { name: "plus", description: "does a different thing" }]
   )

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-inference-configuration-template.graphql

@@ -7,7 +7,6 @@ type Mutation {
   ): ConversationMessage
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
     INFERENENCE_CONFIGURATION
   )

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-invalid-missing-ai-model.graphql

@@ -6,7 +6,6 @@ type Mutation {
     toolConfiguration: ToolConfigurationInput
   ): ConversationMessage
   @conversation(
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot."
   )
 }

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-invalid-return-type.graphql

@@ -7,7 +7,6 @@ type Mutation {
   ): String
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot."
   )
 }

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-model-query-tool-with-relationships.graphql

@@ -31,7 +31,6 @@ type Mutation {
   ): ConversationMessage
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
     tools: [{ name: "listCustomers", description: "Provides data about the customer sending a message" }]
   )

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-model-query-tool.graphql

@@ -12,7 +12,6 @@ type Mutation {
   ): ConversationMessage
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
-    functionName: "conversation-handler",
     systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
     tools: [{ name: "listTodos", description: "lists todos" }]
   )

packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/conversation-route-with-inference-configuration.graphql

@@ -8,7 +8,6 @@ type Mutation {
   @conversation(
     aiModel: "anthropic.claude-3-haiku-20240307-v1:0",
     systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
-    functionName: "conversation-handler",
     inferenceConfiguration: {
         temperature: 0.5,
         topP: 0.9,

packages/amplify-graphql-conversation-transformer/src/grapqhl-conversation-transformer.ts

@@ -15,6 +15,7 @@ import { type ToolDefinition, type Tools } from './utils/tools';
 import { ConversationPrepareHandler } from './transformer-steps/conversation-prepare-handler';
 import { ConversationResolverGenerator } from './transformer-steps/conversation-resolver-generator';
 import { ConversationFieldHandler } from './transformer-steps/conversation-field-handler';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
 
 /**
  * Configuration for the Conversation Directive
@@ -58,11 +59,12 @@ export class ConversationTransformer extends TransformerPluginBase {
     hasManyTransformer: HasManyTransformer,
     belongsToTransformer: BelongsToTransformer,
     authProvider: TransformerAuthProvider,
+    functionNameMap?: Record<string, lambda.IFunction>,
   ) {
     super('amplify-conversation-transformer', ConversationDirective.definition);
     this.fieldHandler = new ConversationFieldHandler();
     this.prepareHandler = new ConversationPrepareHandler(modelTransformer, hasManyTransformer, belongsToTransformer, authProvider);
-    this.resolverGenerator = new ConversationResolverGenerator();
+    this.resolverGenerator = new ConversationResolverGenerator(functionNameMap);
   }
 
   /**

packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-resolver-generator.ts

@@ -5,7 +5,7 @@ import { APPSYNC_JS_RUNTIME, TransformerResolver } from '@aws-amplify/graphql-tr
 import { ResolverResourceIDs, FunctionResourceIDs, ResourceConstants, toUpper } from 'graphql-transformer-common';
 import * as cdk from 'aws-cdk-lib';
 import { conversation } from '@aws-amplify/ai-constructs';
-import { IFunction, Function } from 'aws-cdk-lib/aws-lambda';
+import { IFunction } from 'aws-cdk-lib/aws-lambda';
 import { getModelDataSourceNameForTypeName, getTable } from '@aws-amplify/graphql-transformer-core';
 import { initMappingTemplate } from '../resolvers/init-resolver';
 import { authMappingTemplate } from '../resolvers/auth-resolver';
@@ -24,6 +24,8 @@ type KeyAttributeDefinition = {
 
 // TODO: add explanation for the tool model queries
 export class ConversationResolverGenerator {
+  constructor(private readonly functionNameMap?: Record<string, IFunction>) {}
+
   generateResolvers(directives: ConversationDirectiveConfiguration[], ctx: TransformerContextProvider): void {
     for (const directive of directives) {
       this.processToolsForDirective(directive, ctx);
@@ -89,7 +91,7 @@ export class ConversationResolverGenerator {
     capitalizedFieldName: string,
   ): { functionDataSourceId: string; referencedFunction: IFunction } {
     if (directive.functionName) {
-      return this.setupExistingFunctionDataSource(directive.functionName, functionStack);
+      return this.setupExistingFunctionDataSource(directive.functionName);
     } else {
       return this.setupDefaultConversationHandler(functionStack, capitalizedFieldName, directive.aiModel);
     }
@@ -101,28 +103,18 @@ export class ConversationResolverGenerator {
    * @param functionStack - The CDK stack to add the function to
    * @returns An object containing the function data source ID and the referenced function
    */
-  private setupExistingFunctionDataSource(
-    functionName: string,
-    functionStack: cdk.Stack,
-  ): { functionDataSourceId: string; referencedFunction: IFunction } {
+  private setupExistingFunctionDataSource(functionName: string): { functionDataSourceId: string; referencedFunction: IFunction } {
     const functionDataSourceId = FunctionResourceIDs.FunctionDataSourceID(functionName);
-    const referencedFunction = Function.fromFunctionAttributes(functionStack, `${functionDataSourceId}Function`, {
-      functionArn: this.lambdaArnResource(functionName),
-    });
-
+    if (!this.functionNameMap) {
+      throw new Error('Function name map is not provided');
+    }
+    const referencedFunction = this.functionNameMap[functionName];
+    if (!referencedFunction) {
+      throw new Error(`Function ${functionName} not found in function name map`);
+    }
     return { functionDataSourceId, referencedFunction };
   }
 
-  /**
-   * Generates the Lambda ARN resource string
-   * @param name - The name of the Lambda function
-   * @returns The Lambda ARN resource string
-   */
-  private lambdaArnResource(name: string): string {
-    // eslint-disable-next-line no-template-curly-in-string
-    return cdk.Fn.sub('arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${name}', { name });
-  }
-
   /**
    * Sets up a default conversation handler function
    * @param functionStack - The CDK stack to add the function to

packages/amplify-graphql-transformer-test-utils/API.md

@@ -226,6 +226,7 @@ export type TestTransformParameters = RDSLayerMappingProvider & RDSSNSTopicMappi
     transformers: TransformerPluginProvider[];
     transformParameters?: Partial<TransformParameters>;
     userDefinedSlots?: Record<string, UserDefinedSlot[]>;
+    transformerManager?: TransformManager;
 };
 
 // @public (undocumented)

packages/amplify-graphql-transformer-test-utils/src/test-transform.ts

@@ -31,6 +31,7 @@ export type TestTransformParameters = RDSLayerMappingProvider &
     transformers: TransformerPluginProvider[];
     transformParameters?: Partial<TransformParameters>;
     userDefinedSlots?: Record<string, UserDefinedSlot[]>;
+    transformerManager?: TransformManager;
   };
 
 /**
@@ -52,6 +53,7 @@ export const testTransform = (params: TestTransformParameters): DeploymentResour
     transformers,
     transformParameters,
     userDefinedSlots,
+    transformerManager,
   } = params;
 
   const transform = new GraphQLTransform({
@@ -63,7 +65,7 @@ export const testTransform = (params: TestTransformParameters): DeploymentResour
     resolverConfig,
   });
 
-  const transformManager = new TransformManager(overrideConfig);
+  const transformManager = transformerManager ?? new TransformManager(overrideConfig);
 
   const authConfigTypes = [authConfig?.defaultAuthentication, ...(authConfig?.additionalAuthenticationProviders ?? [])].map(
     (authConfigEntry) => authConfigEntry?.authenticationType,

packages/amplify-graphql-transformer/src/graphql-transformer.ts

@@ -78,7 +78,7 @@ export const constructTransformerChain = (options?: TransformerFactoryArgs): Tra
     hasOneTransformer,
     new ManyToManyTransformer(modelTransformer, indexTransformer, hasOneTransformer, authTransformer),
     belongsToTransformer,
-    new ConversationTransformer(modelTransformer, hasManyTransformer, belongsToTransformer, authTransformer),
+    new ConversationTransformer(modelTransformer, hasManyTransformer, belongsToTransformer, authTransformer, options?.functionNameMap),
     new GenerationTransformer(),
     new DefaultValueTransformer(),
     authTransformer,