From 53dacd1e661ce71d2db3ed50231df453de43d6dd Mon Sep 17 00:00:00 2001 From: Ian Saultz <52051793+atierian@users.noreply.github.com> Date: Thu, 3 Oct 2024 22:49:10 -0400 Subject: [PATCH 1/5] fix: include custom handler in lambdaFunctions for functionMap --- packages/data-schema/src/SchemaProcessor.ts | 4 +++- .../data-schema/src/ai/ConversationSchemaTypes.ts | 15 ++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts index 23baae8fa..bb3081554 100644 --- a/packages/data-schema/src/SchemaProcessor.ts +++ b/packages/data-schema/src/SchemaProcessor.ts @@ -1442,7 +1442,9 @@ const schemaPreprocessor = ( } } else if (isConversationRoute(typeDef)) { // TODO: add inferenceConfiguration values to directive. - customMutations.push(createConversationField(typeDef, typeName)); + const { field, functionHandler } = createConversationField(typeDef, typeName); + customMutations.push(field); + Object.assign(lambdaFunctions, functionHandler); shouldAddConversationTypes = true; } } else if (staticSchema) { diff --git a/packages/data-schema/src/ai/ConversationSchemaTypes.ts b/packages/data-schema/src/ai/ConversationSchemaTypes.ts index f497994f4..542e88244 100644 --- a/packages/data-schema/src/ai/ConversationSchemaTypes.ts +++ b/packages/data-schema/src/ai/ConversationSchemaTypes.ts @@ -1,3 +1,4 @@ +import { LambdaFunctionDefinition } from '@aws-amplify/data-schema-types'; import { InternalRef } from '../RefType'; import { capitalize } from '../runtime/utils'; import type { @@ -8,7 +9,7 @@ import type { export const createConversationField = ( typeDef: InternalConversationType, typeName: string, -): string => { +): { field: string, functionHandler: LambdaFunctionDefinition } => { const { aiModel, systemPrompt, handler, tools } = typeDef; const args: Record = { @@ -25,12 +26,11 @@ export const createConversationField = ( systemPrompt: systemPrompt.replace(/\r?\n/g, '\\n'), }; + let functionHandler: LambdaFunctionDefinition = {}; if (handler) { - if (typeof handler === 'string') { - args['functionName'] = handler; - } else if (typeof handler.getInstance === 'function') { - args['functionName'] = `Fn${capitalize(typeName)}`; - } + const functionName = `Fn${capitalize(typeName)}`; + args['functionName'] = functionName; + functionHandler[functionName] = handler; } const argsString = Object.entries(args) @@ -43,7 +43,8 @@ export const createConversationField = ( const conversationDirective = `@conversation(${argsString}${toolsString})`; - return `${typeName}(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage ${conversationDirective} @aws_cognito_user_pools`; + const field = `${typeName}(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage ${conversationDirective} @aws_cognito_user_pools`; + return { field, functionHandler }; }; const isRef = (query: unknown): query is { data: InternalRef['data'] } => From 735754d015f149bb98e4e329a99beec5da15ff80 Mon Sep 17 00:00:00 2001 From: Ian Saultz <52051793+atierian@users.noreply.github.com> Date: Thu, 3 Oct 2024 22:49:46 -0400 Subject: [PATCH 2/5] remove string handler argument for conversation routes --- packages/data-schema/src/ai/ConversationType.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/data-schema/src/ai/ConversationType.ts b/packages/data-schema/src/ai/ConversationType.ts index 4b25da9ab..e4c72b717 100644 --- a/packages/data-schema/src/ai/ConversationType.ts +++ b/packages/data-schema/src/ai/ConversationType.ts @@ -125,7 +125,7 @@ export interface ConversationInput { systemPrompt: string; inferenceConfiguration?: InferenceConfiguration; tools?: ToolDefinition[]; - handler?: DefineFunction | string; + handler?: DefineFunction; } export interface InternalConversationType From 8a67be61912a564dbe5c1cb494967d0d30a9d071 Mon Sep 17 00:00:00 2001 From: Ian Saultz <52051793+atierian@users.noreply.github.com> Date: Thu, 3 Oct 2024 22:52:42 -0400 Subject: [PATCH 3/5] add changeset --- .changeset/lazy-guests-tap.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/lazy-guests-tap.md diff --git a/.changeset/lazy-guests-tap.md b/.changeset/lazy-guests-tap.md new file mode 100644 index 000000000..75459b3e2 --- /dev/null +++ b/.changeset/lazy-guests-tap.md @@ -0,0 +1,5 @@ +--- +"@aws-amplify/data-schema": patch +--- + +include custom conversation handler in lambdaDefinitions From 3e3c1dbcd2177bd9ffb853f5d0a6d52dc9f1e56f Mon Sep 17 00:00:00 2001 From: Ian Saultz <52051793+atierian@users.noreply.github.com> Date: Thu, 3 Oct 2024 23:23:53 -0400 Subject: [PATCH 4/5] lint --- packages/data-schema/src/ai/ConversationSchemaTypes.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/data-schema/src/ai/ConversationSchemaTypes.ts b/packages/data-schema/src/ai/ConversationSchemaTypes.ts index 542e88244..ff7d1d8d7 100644 --- a/packages/data-schema/src/ai/ConversationSchemaTypes.ts +++ b/packages/data-schema/src/ai/ConversationSchemaTypes.ts @@ -26,7 +26,7 @@ export const createConversationField = ( systemPrompt: systemPrompt.replace(/\r?\n/g, '\\n'), }; - let functionHandler: LambdaFunctionDefinition = {}; + const functionHandler: LambdaFunctionDefinition = {}; if (handler) { const functionName = `Fn${capitalize(typeName)}`; args['functionName'] = functionName; From 7be441b77019fa888f0b9236e10493afc6ba7409 Mon Sep 17 00:00:00 2001 From: Ian Saultz <52051793+atierian@users.noreply.github.com> Date: Fri, 4 Oct 2024 14:03:45 -0400 Subject: [PATCH 5/5] assert custom conversation handlers are including in lambda functions map --- packages/data-schema/__tests__/ClientSchema.test.ts | 12 +++++++++++- .../__snapshots__/ClientSchema.test.ts.snap | 1 + 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/packages/data-schema/__tests__/ClientSchema.test.ts b/packages/data-schema/__tests__/ClientSchema.test.ts index 19cc0c10c..db6cce643 100644 --- a/packages/data-schema/__tests__/ClientSchema.test.ts +++ b/packages/data-schema/__tests__/ClientSchema.test.ts @@ -1451,6 +1451,7 @@ describe('SQL Schema with sql statement references', () => { describe('ai routes', () => { test('conversations', () => { const handler = defineFunctionStub({}); + const customConversationHandler = defineFunctionStub({}); const schema = a.schema({ Profile: a.customType({ value: a.integer(), @@ -1494,6 +1495,11 @@ describe('ai routes', () => { ], }), + CustomHandlerChatBot: a.conversation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: 'Hello, world!', + handler: customConversationHandler, + }), }); type Schema = ClientSchema; @@ -1507,8 +1513,12 @@ describe('ai routes', () => { type test = Expect>; - const graphql = schema.transform().schema; + const derivedApiDefinition = schema.transform(); + const graphql = derivedApiDefinition.schema; expect(graphql).toMatchSnapshot(); + + const lambdaFunctions = derivedApiDefinition.lambdaFunctions; + expect(lambdaFunctions['FnCustomHandlerChatBot']).toBeDefined(); }); test('generations', () => { diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap index accd1cc94..e9520a02e 100644 --- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap +++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap @@ -35,6 +35,7 @@ exports[`ai routes conversations 1`] = ` type Mutation { ChatBot(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage @conversation(aiModel: "anthropic.claude-3-haiku-20240307-v1:0", systemPrompt: "Hello, world!", tools: [{ name: "myToolQuery", description: "does a thing" }, { name: "anotherToolQuery", description: "does a different thing" }]) @aws_cognito_user_pools MultilinePromptChatBot(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage @conversation(aiModel: "anthropic.claude-3-haiku-20240307-v1:0", systemPrompt: "You are a helpful assistant.\\n Respond in the poetic form of haiku.", tools: [{ name: "myToolQuery", description: "does a thing" }, { name: "anotherToolQuery", description: "does a different thing" }]) @aws_cognito_user_pools + CustomHandlerChatBot(conversationId: ID!, content: [ContentBlockInput], aiContext: AWSJSON, toolConfiguration: ToolConfigurationInput): ConversationMessage @conversation(aiModel: "anthropic.claude-3-haiku-20240307-v1:0", systemPrompt: "Hello, world!", functionName: "FnCustomHandlerChatBot") @aws_cognito_user_pools } type Query {