fix: include custom conversation handlers in lambdaFunctions for downstream processing

.changeset/lazy-guests-tap.md

@@ -0,0 +1,5 @@
+---
+"@aws-amplify/data-schema": patch
+---
+
+include custom conversation handler in lambdaDefinitions

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<typeof schema>;
@@ -1507,8 +1513,12 @@ describe('ai routes', () => {
 
     type test = Expect<Equal<ActualChatBotInterface, Expected>>;
 
-    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', () => {

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 {

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) {

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<string, string> = {
@@ -25,12 +26,11 @@ export const createConversationField = (
     systemPrompt: systemPrompt.replace(/\r?\n/g, '\\n'),
   };
 
+  const 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'] } =>

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