refactor(conversation): consolidate resolver generation logic

.prettierignore

@@ -14,7 +14,7 @@ packages/amplify-graphql-api-construct-tests/amplify-e2e-reports
 packages/amplify-graphql-api-construct/README.md
 packages/amplify-graphql-api-construct/tsconfig.json
 packages/amplify-graphql-conversation-transformer/src/__tests__/schemas/*.graphql
-packages/amplify-graphql-conversation-transformer/src/resolvers/*.template.js
+packages/amplify-graphql-conversation-transformer/src/resolvers/templates/*.template.js
 packages/amplify-graphql-generation-transformer/src/resolvers/*.template.js
 packages/amplify-data-construct/README.md
 packages/amplify-data-construct/tsconfig.json

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

@@ -27,8 +27,6 @@ export class ConversationTransformer extends TransformerPluginBase {
     generateResolvers: (ctx: TransformerContextProvider) => void;
     // (undocumented)
     prepare: (ctx: TransformerPrepareStepContextProvider) => void;
-    // (undocumented)
-    validate: () => void;
 }
 
 // (No @packageDocumentation comment for this package)

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

@@ -19,7 +19,7 @@
     "build": "tsc && yarn copy-js-resolver-templates",
     "watch": "tsc -w",
     "clean": "rimraf ./lib",
-    "copy-js-resolver-templates": "cp ./src/resolvers/*.template.js ./lib/resolvers",
+    "copy-js-resolver-templates": "cd lib/resolvers && mkdir -p templates && cp ../../src/resolvers/templates/*.template.js templates",
     "test": "jest",
     "extract-api": "ts-node ../../scripts/extract-api.ts"
   },

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

@@ -152,7 +152,6 @@ export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
 
-
   const selectionSet = 'id conversationId content { image { format source { bytes }} text toolUse { toolUseId name input } toolResult { status toolUseId content { json text image { format source { bytes }} document { format name source { bytes }} }}} role owner createdAt updatedAt';
 
   const responseMutation = {
@@ -161,18 +160,18 @@ export function request(ctx) {
     selectionSet,
   };
   const currentMessageId = ctx.stash.defaultValues.id;
+
   const modelConfiguration = {
-  modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
-  systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
-  inferenceConfiguration: {"temperature":0.5,"topP":0.9,"maxTokens":100},
-};
+    modelId: "anthropic.claude-3-haiku-20240307-v1:0",
+    systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
+    inferenceConfiguration: {"temperature":0.5,"topP":0.9,"maxTokens":100},
+  };
 
   const clientTools = args.toolConfiguration?.tools?.map((tool) => {
     return { ...tool.toolSpec };
   });
-  const toolsConfiguration = {
-  clientTools
-};
+  const dataTools = undefined;
+  const toolsConfiguration = { dataTools, clientTools };
 
   const messageHistoryQuery = {
     getQueryName: 'getConversationMessagePirateChat',
@@ -373,7 +372,6 @@ export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
 
-  const toolDefinitions = {"tools":[{"name":"listTodos","description":"lists todos","inputSchema":{"json":{"type":"object","properties":{},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"items { content isDone id createdAt updatedAt owner } nextToken","propertyTypes":{},"queryName":"listTodos"}}]};
   const selectionSet = 'id conversationId content { image { format source { bytes }} text toolUse { toolUseId name input } toolResult { status toolUseId content { json text image { format source { bytes }} document { format name source { bytes }} }}} role owner createdAt updatedAt';
 
   const responseMutation = {
@@ -382,20 +380,18 @@ export function request(ctx) {
     selectionSet,
   };
   const currentMessageId = ctx.stash.defaultValues.id;
+
   const modelConfiguration = {
-  modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
-  systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
-
-};
+    modelId: "anthropic.claude-3-haiku-20240307-v1:0",
+    systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
+    inferenceConfiguration: undefined,
+  };
 
   const clientTools = args.toolConfiguration?.tools?.map((tool) => {
     return { ...tool.toolSpec };
   });
-  const dataTools = toolDefinitions.tools;
- const toolsConfiguration = {
-  dataTools,
-  clientTools,
-};
+  const dataTools = [{"name":"listTodos","description":"lists todos","inputSchema":{"json":{"type":"object","properties":{},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"items { content isDone id createdAt updatedAt owner } nextToken","propertyTypes":{},"queryName":"listTodos"}}];
+  const toolsConfiguration = { dataTools, clientTools };
 
   const messageHistoryQuery = {
     getQueryName: 'getConversationMessagePirateChat',
@@ -596,7 +592,6 @@ export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
 
-  const toolDefinitions = {"tools":[{"name":"listCustomers","description":"Provides data about the customer sending a message","inputSchema":{"json":{"type":"object","properties":{},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"items { name email activeCart { products { name price } customerId id createdAt updatedAt owner } orderHistory { items { products { name price } customerId id createdAt updatedAt owner } nextToken } id createdAt updatedAt owner } nextToken","propertyTypes":{},"queryName":"listCustomers"}}]};
   const selectionSet = 'id conversationId content { image { format source { bytes }} text toolUse { toolUseId name input } toolResult { status toolUseId content { json text image { format source { bytes }} document { format name source { bytes }} }}} role owner createdAt updatedAt';
 
   const responseMutation = {
@@ -605,20 +600,18 @@ export function request(ctx) {
     selectionSet,
   };
   const currentMessageId = ctx.stash.defaultValues.id;
+
   const modelConfiguration = {
-  modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
-  systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
-
-};
+    modelId: "anthropic.claude-3-haiku-20240307-v1:0",
+    systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
+    inferenceConfiguration: undefined,
+  };
 
   const clientTools = args.toolConfiguration?.tools?.map((tool) => {
     return { ...tool.toolSpec };
   });
-  const dataTools = toolDefinitions.tools;
- const toolsConfiguration = {
-  dataTools,
-  clientTools,
-};
+  const dataTools = [{"name":"listCustomers","description":"Provides data about the customer sending a message","inputSchema":{"json":{"type":"object","properties":{},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"items { name email activeCart { products { name price } customerId id createdAt updatedAt owner } orderHistory { items { products { name price } customerId id createdAt updatedAt owner } nextToken } id createdAt updatedAt owner } nextToken","propertyTypes":{},"queryName":"listCustomers"}}];
+  const toolsConfiguration = { dataTools, clientTools };
 
   const messageHistoryQuery = {
     getQueryName: 'getConversationMessagePirateChat',
@@ -819,7 +812,6 @@ export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
 
-  const toolDefinitions = {"tools":[{"name":"getTemperature","description":"does a thing","inputSchema":{"json":{"type":"object","properties":{"city":{"type":"string","description":"A UTF-8 character sequence."}},"required":["city"]}},"graphqlRequestInputDescriptor":{"selectionSet":"value unit","propertyTypes":{"city":"String!"},"queryName":"getTemperature"}},{"name":"plus","description":"does a different thing","inputSchema":{"json":{"type":"object","properties":{"a":{"type":"number","description":"A signed 32-bit integer value."},"b":{"type":"number","description":"A signed 32-bit integer value."}},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"","propertyTypes":{"a":"Int","b":"Int"},"queryName":"plus"}}]};
   const selectionSet = 'id conversationId content { image { format source { bytes }} text toolUse { toolUseId name input } toolResult { status toolUseId content { json text image { format source { bytes }} document { format name source { bytes }} }}} role owner createdAt updatedAt';
 
   const responseMutation = {
@@ -828,20 +820,18 @@ export function request(ctx) {
     selectionSet,
   };
   const currentMessageId = ctx.stash.defaultValues.id;
+
   const modelConfiguration = {
-  modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
-  systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
-
-};
+    modelId: "anthropic.claude-3-haiku-20240307-v1:0",
+    systemPrompt: "You are a helpful chatbot. Answer questions to the best of your ability.",
+    inferenceConfiguration: undefined,
+  };
 
   const clientTools = args.toolConfiguration?.tools?.map((tool) => {
     return { ...tool.toolSpec };
   });
-  const dataTools = toolDefinitions.tools;
- const toolsConfiguration = {
-  dataTools,
-  clientTools,
-};
+  const dataTools = [{"name":"getTemperature","description":"does a thing","inputSchema":{"json":{"type":"object","properties":{"city":{"type":"string","description":"A UTF-8 character sequence."}},"required":["city"]}},"graphqlRequestInputDescriptor":{"selectionSet":"value unit","propertyTypes":{"city":"String!"},"queryName":"getTemperature"}},{"name":"plus","description":"does a different thing","inputSchema":{"json":{"type":"object","properties":{"a":{"type":"number","description":"A signed 32-bit integer value."},"b":{"type":"number","description":"A signed 32-bit integer value."}},"required":[]}},"graphqlRequestInputDescriptor":{"selectionSet":"","propertyTypes":{"a":"Int","b":"Int"},"queryName":"plus"}}];
+  const toolsConfiguration = { dataTools, clientTools };
 
   const messageHistoryQuery = {
     getQueryName: 'getConversationMessagePirateChat',

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

@@ -1,17 +1,17 @@
 import { AuthTransformer } from '@aws-amplify/graphql-auth-transformer';
+import { GenerationTransformer } from '@aws-amplify/graphql-generation-transformer';
 import { IndexTransformer, PrimaryKeyTransformer } from '@aws-amplify/graphql-index-transformer';
 import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
+import { BelongsToTransformer, HasManyTransformer, HasOneTransformer } from '@aws-amplify/graphql-relational-transformer';
 import { validateModelSchema } from '@aws-amplify/graphql-transformer-core';
 import { AppSyncAuthConfiguration, ModelDataSourceStrategy } from '@aws-amplify/graphql-transformer-interfaces';
 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 * as fs from 'fs-extra';
+import { parse } from 'graphql';
 import { toUpper } from 'graphql-transformer-common';
+import * as path from 'path';
+import { ConversationTransformer } from '..';
 
 const conversationSchemaTypes = fs.readFileSync(path.join(__dirname, 'schemas/conversation-schema-types.graphql'), 'utf8');
 
@@ -157,16 +157,6 @@ const getResolverResource = (mutationName: string, resources?: Record<string, an
   return resources?.[resolverName];
 };
 
-const getResolverFnResource = (mutationName: string, resources: DeploymentResources): string => {
-  const resolverFnCode =
-    resources.rootStack.Resources &&
-    Object.entries(resources.rootStack.Resources).find(([key, _]) => key.startsWith(`Mutation${toUpper(mutationName)}DataResolverFn`))?.[1][
-      'Properties'
-    ]['Code'];
-
-  return resolverFnCode;
-};
-
 const defaultAuthConfig: AppSyncAuthConfiguration = {
   defaultAuthentication: {
     authenticationType: 'AMAZON_COGNITO_USER_POOLS',

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

@@ -1,7 +1,7 @@
-import { ConversationFieldHandler } from '../transformer-steps/conversation-field-handler';
-import { DirectiveNode, FieldDefinitionNode, ObjectTypeDefinitionNode } from 'graphql';
 import { InvalidDirectiveError } from '@aws-amplify/graphql-transformer-core';
 import { TransformerSchemaVisitStepContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
+import { DirectiveNode, FieldDefinitionNode, ObjectTypeDefinitionNode } from 'graphql';
+import { ConversationFieldHandler } from '../transformer-steps/conversation-field-handler';
 
 describe('ConversationFieldHandler', () => {
   let handler: ConversationFieldHandler;
@@ -55,8 +55,8 @@ describe('ConversationFieldHandler', () => {
       const config = handler.getDirectiveConfig(parent, field, directive, mockContext);
 
       expect(config).toBeDefined();
-      expect(config.messageModel).toBeDefined();
-      expect(config.conversationModel).toBeDefined();
+      expect(config.message).toBeDefined();
+      expect(config.conversation).toBeDefined();
     });
 
     it('should throw an error if field type is not ConversationMessage', () => {
@@ -74,12 +74,10 @@ describe('ConversationFieldHandler', () => {
       const field = createMockField('testField', 'ConversationMessage');
       const directive = createMockDirective();
 
-      const config = handler.getDirectiveConfig(parent, field, directive, mockContext);
-      const { messageModel } = config.messageModel;
-      const { conversationModel } = config.conversationModel;
+      const { message, conversation } = handler.getDirectiveConfig(parent, field, directive, mockContext);
 
-      expect(messageModel.name.value).toBe('ConversationMessageTestField');
-      expect(conversationModel.name.value).toBe('ConversationTestField');
+      expect(message.model.name.value).toBe('ConversationMessageTestField');
+      expect(conversation.model.name.value).toBe('ConversationTestField');
     });
   });
 });

packages/amplify-graphql-conversation-transformer/src/conversation-directive-configuration.ts

@@ -0,0 +1,70 @@
+import { DataSourceProvider } from '@aws-amplify/graphql-transformer-interfaces';
+import { DirectiveNode, FieldDefinitionNode, InputObjectTypeDefinitionNode, ObjectTypeDefinitionNode } from 'graphql';
+import { ConversationModel } from './graphql-types/conversation-model';
+import { MessageModel } from './graphql-types/message-model';
+import { Tool } from './tools/process-tools';
+
+/**
+ * Configuration for the Conversation Directive
+ */
+export type ConversationDirectiveConfiguration = {
+  // From the GraphQL Schema
+  parent: ObjectTypeDefinitionNode;
+  field: FieldDefinitionNode;
+  directive: DirectiveNode;
+
+  // Raw Directive Input
+  aiModel: string;
+  systemPrompt: string;
+  inferenceConfiguration: ConversationInferenceConfiguration | undefined;
+  tools: ToolDefinition[];
+  /**
+   * Custom handler function name.
+   *
+   * @deprecated Replaced by 'handler'
+   */
+  functionName: string | undefined;
+  handler: ConversationHandlerFunctionConfiguration | undefined;
+
+  // Generated within the Conversation Transformer
+  toolSpec: Tool[] | undefined;
+  conversation: ConversationModel;
+  message: MessageModel;
+  assistantResponseMutation: { field: FieldDefinitionNode; input: InputObjectTypeDefinitionNode };
+  assistantResponseSubscriptionField: FieldDefinitionNode;
+  dataSources: ConversationDirectiveDataSources;
+};
+
+/**
+ * Conversation Directive Handler Function Configuration
+ */
+export type ConversationHandlerFunctionConfiguration = {
+  functionName: string;
+  eventVersion: string;
+};
+
+/**
+ * Conversation Directive Inference Configuration
+ */
+export type ConversationInferenceConfiguration = {
+  maxTokens?: number;
+  temperature?: number;
+  topP?: number;
+};
+
+/**
+ * Conversation Directive Tool Definition
+ */
+export type ToolDefinition = {
+  name: string;
+  description: string;
+};
+
+/**
+ * Conversation Directive Data Sources
+ */
+export type ConversationDirectiveDataSources = {
+  lambdaFunctionDataSource: DataSourceProvider;
+  messageTableDataSource: DataSourceProvider;
+  conversationTableDataSource: DataSourceProvider;
+};