update assistant response resolver pipeline to reflect 1:1 message:it…

…em mapping

packages/amplify-graphql-conversation-transformer/src/resolvers/assistant-mutation-resolver-fn.js

@@ -10,6 +10,7 @@ export function request(ctx) {
   const { conversationId, content, associatedUserMessageId } = ctx.args.input;
   const { owner } = ctx.args;
   const defaultValues = ctx.stash.defaultValues ?? {};
+  const id = defaultValues.id;
 
   const message = {
     __typename: '[[CONVERSATION_MESSAGE_TYPE_NAME]]',
@@ -23,26 +24,6 @@ export function request(ctx) {
   };
 
   return ddb.put({ key: { id }, item: message });
-
-  // const expression = 'SET #assistantContent = :assistantContent, #updatedAt = :updatedAt';
-  // const expressionNames = { '#assistantContent': 'assistantContent', '#updatedAt': 'updatedAt' };
-  // const expressionValues = { ':assistantContent': content, ':updatedAt': updatedAt };
-  // const condition = JSON.parse(
-  //   util.transform.toDynamoDBConditionExpression({
-  //     owner: { eq: owner },
-  //     conversationId: { eq: conversationId },
-  //   }),
-  // );
-  // return {
-  //   operation: 'UpdateItem',
-  //   key: util.dynamodb.toMapValues({ id: associatedUserMessageId }),
-  //   condition,
-  //   update: {
-  //     expression,
-  //     expressionNames,
-  //     expressionValues: util.dynamodb.toMapValues(expressionValues),
-  //   },
-  // };
 }
 
 /**
@@ -56,16 +37,5 @@ export function response(ctx) {
     util.error(ctx.error.message, ctx.error.type);
   }
 
-  const { conversationId, content, associatedUserMessageId } = ctx.args.input;
-  const { createdAt, updatedAt } = ctx.result;
-
-  return {
-    id: associatedUserMessageId,
-    content,
-    conversationId,
-    role: 'assistant',
-    owner: ctx.stash.owner,
-    createdAt,
-    updatedAt,
-  };
+  return ctx.result;
 }

packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver-fn.js

@@ -1,10 +1,11 @@
 export function request(ctx) {
   const { authFilter } = ctx.stash;
+  const { conversationId } = [[CONVERSATION_ID_PARENT]];
 
   const query = {
     expression: 'id = :id',
     expressionValues: util.dynamodb.toMapValues({
-      ':id': ctx.args.conversationId,
+      ':id': conversationId,
     }),
   };
 

packages/amplify-graphql-conversation-transformer/src/resolvers/verify-session-owner-resolver.ts

@@ -9,8 +9,33 @@ import { ConversationDirectiveConfiguration } from '../grapqhl-conversation-tran
  *
  * @returns {MappingTemplateProvider} An object containing request and response MappingTemplateProviders.
  */
-export const verifySessionOwnerMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => {
-  const resolver = fs.readFileSync(path.join(__dirname, 'verify-session-owner-resolver-fn.js'), 'utf8');
+export const verifySessionOwnerSendMessageMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => {
+  const substitutions = {
+    CONVERSATION_ID_PARENT: 'ctx.args',
+  };
   const templateName = `Mutation.${config.field.name.value}.verify-session-owner.js`;
+  return verifySessionOwnerMappingTemplate(templateName, substitutions);
+};
+
+/**
+ * Creates a mapping template for verifying the session owner in a conversation.
+ *
+ * @returns {MappingTemplateProvider} An object containing request and response MappingTemplateProviders.
+ */
+export const verifySessionOwnerAssistantResponseMappingTemplate = (config: ConversationDirectiveConfiguration): MappingTemplateProvider => {
+  const substitutions = {
+    CONVERSATION_ID_PARENT: 'ctx.args.input',
+  };
+  const templateName = `Mutation.${config.field.name.value}AssistantResponse.verify-session-owner.js`;
+  return verifySessionOwnerMappingTemplate(templateName, substitutions);
+};
+
+const verifySessionOwnerMappingTemplate = (name: string, substitute: Record<string, string>) => {
+  let resolver = fs.readFileSync(path.join(__dirname, 'verify-session-owner-resolver-fn.js'), 'utf8');
+  Object.entries(substitute).forEach(([key, value]) => {
+    const replaced = resolver.replace(new RegExp(`\\[\\[${key}\\]\\]`, 'g'), value);
+    resolver = replaced;
+  });
+  const templateName = `Mutation.${name}.verify-session-owner.js`;
   return MappingTemplate.s3MappingFunctionCodeFromString(resolver, templateName);
 };

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

@@ -9,7 +9,10 @@ import { IFunction, Function } 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';
-import { verifySessionOwnerMappingTemplate } from '../resolvers/verify-session-owner-resolver';
+import {
+  verifySessionOwnerSendMessageMappingTemplate,
+  verifySessionOwnerAssistantResponseMappingTemplate,
+} from '../resolvers/verify-session-owner-resolver';
 import { writeMessageToTableMappingTemplate } from '../resolvers/write-message-to-table-resolver';
 import { invokeLambdaMappingTemplate } from '../resolvers/invoke-lambda-resolver';
 import { assistantMutationResolver } from '../resolvers/assistant-mutation-resolver';
@@ -54,8 +57,8 @@ export class ConversationResolverGenerator {
     this.setupMessageTableIndex(ctx, directive);
     const initResolverFunction = initMappingTemplate(directive);
     const authResolverFunction = authMappingTemplate(directive);
-    const verifySessionOwnerResolverFunction = verifySessionOwnerMappingTemplate(directive);
-
+    const verifySessionOwnerSendMessageResolverFunction = verifySessionOwnerSendMessageMappingTemplate(directive);
+    const verifySessionOwnerAssistantResponseResolverFunction = verifySessionOwnerAssistantResponseMappingTemplate(directive);
 
     this.createConversationPipelineResolver(
       ctx,
@@ -66,10 +69,17 @@ export class ConversationResolverGenerator {
       invokeLambdaFunction,
       initResolverFunction,
       authResolverFunction,
-      verifySessionOwnerResolverFunction,
+      verifySessionOwnerSendMessageResolverFunction,
     );
 
-    this.createAssistantResponseResolver(ctx, directive, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction);
+    this.createAssistantResponseResolver(
+      ctx,
+      directive,
+      capitalizedFieldName,
+      initResolverFunction,
+      authResolverFunction,
+      verifySessionOwnerAssistantResponseResolverFunction,
+    );
     this.createAssistantResponseSubscriptionResolver(ctx, directive, capitalizedFieldName);
   }
 
@@ -193,7 +203,14 @@ export class ConversationResolverGenerator {
       runtime,
     );
 
-    this.addPipelineResolverFunctions(ctx, conversationPipelineResolver, capitalizedFieldName, initResolverFunction, authResolverFunction, verifySessionOwnerResolverFunction);
+    this.addPipelineResolverFunctions(
+      ctx,
+      conversationPipelineResolver,
+      capitalizedFieldName,
+      initResolverFunction,
+      authResolverFunction,
+      verifySessionOwnerResolverFunction,
+    );
 
     ctx.resolvers.addResolver(parentName, fieldName, conversationPipelineResolver);
   }
@@ -256,13 +273,12 @@ export class ConversationResolverGenerator {
       directive.responseMutationName,
       assistantResponseResolverResourceId,
       { codeMappingTemplate: assistantResponseResolverFunction },
-      ['init', 'auth', 'verifySessionOwner',],
+      ['init', 'auth', 'verifySessionOwner'],
       [],
       conversationMessageDataSource as any,
       APPSYNC_JS_RUNTIME,
     );
 
-
     // Add init function
     resolver.addJsFunctionToSlot('init', initResolverFunction);