From 6427b74a27aed2bf038c8abb2d7b0b30f7aedac4 Mon Sep 17 00:00:00 2001
From: Ian Saultz <52051793+atierian@users.noreply.github.com>
Date: Fri, 13 Dec 2024 14:09:33 -0500
Subject: [PATCH 1/5] feat(generation): add handling for models without
 SpecificToolChoice

---
 .../src/grapqhl-generation-transformer.ts     |  2 +-
 .../src/utils/tools.ts                        | 42 +++++++++++++++----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/packages/amplify-graphql-generation-transformer/src/grapqhl-generation-transformer.ts b/packages/amplify-graphql-generation-transformer/src/grapqhl-generation-transformer.ts
index 71d6420263..c6c21be226 100644
--- a/packages/amplify-graphql-generation-transformer/src/grapqhl-generation-transformer.ts
+++ b/packages/amplify-graphql-generation-transformer/src/grapqhl-generation-transformer.ts
@@ -78,7 +78,7 @@ export class GenerationTransformer extends TransformerPluginBase {
 
       const directiveWithToolConfig: GenerationConfigurationWithToolConfig = {
         ...directive,
-        toolConfig: createResponseTypeTool(field, ctx),
+        toolConfig: createResponseTypeTool(directive, ctx),
       };
 
       const stackName = this.bedrockDataSourceName(fieldName) + 'Stack';
diff --git a/packages/amplify-graphql-generation-transformer/src/utils/tools.ts b/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
index 4c990da972..44a019c27e 100644
--- a/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
+++ b/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
@@ -1,7 +1,7 @@
 import { TransformerContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
-import { FieldDefinitionNode } from 'graphql';
 import { generateJSONSchemaFromTypeNode } from './graphql-json-schema-type';
 import { JSONSchema } from '@aws-amplify/graphql-transformer-core';
+import { GenerationDirectiveConfiguration } from '../grapqhl-generation-transformer';
 
 export type Tool = {
   toolSpec: ToolSpec;
@@ -13,13 +13,21 @@ export type Tools = {
 
 export type ToolConfig = {
   tools: Tool[];
-  toolChoice: {
-    tool: {
-      name: string;
-    };
+  toolChoice?: ToolChoice;
+};
+
+type SpecificToolChoice = {
+  tool: {
+    name: string;
   };
 };
 
+type AnyToolChoice = {
+  any: {};
+};
+
+type ToolChoice = SpecificToolChoice | AnyToolChoice | undefined;
+
 type ToolSpec = {
   name: string;
   description: string;
@@ -43,8 +51,8 @@ type ToolSpec = {
  * The returned tool configuration can be used with AI models that support tool-based interactions,
  * ensuring that generated responses match the expected structure of the GraphQL field.
  */
-export const createResponseTypeTool = (field: FieldDefinitionNode, ctx: TransformerContextProvider): ToolConfig => {
-  const { type } = field;
+export const createResponseTypeTool = (config: GenerationDirectiveConfiguration, ctx: TransformerContextProvider): ToolConfig => {
+  const { type } = config.field;
   const schema = generateJSONSchemaFromTypeNode(type, ctx);
 
   // We box the schema to support scalar return types.
@@ -69,8 +77,26 @@ export const createResponseTypeTool = (field: FieldDefinitionNode, ctx: Transfor
       },
     },
   ];
-  const toolChoice = { tool: { name: 'responseType' } };
+
+  const toolChoice = getToolChoice(config);
   const toolConfig = { tools, toolChoice };
 
   return toolConfig;
 };
+
+const getToolChoice = (config: GenerationDirectiveConfiguration): ToolChoice => {
+  switch (config.aiModel) {
+    case 'anthropic.claude-3-opus-20240229-v1:0':
+    case 'anthropic.claude-3-haiku-20240307-v1:0':
+    case 'anthropic.claude-3-sonnet-20240229-v1:0':
+    case 'anthropic.claude-3-5-haiku-20241022-v1:0':
+    case 'anthropic.claude-3-5-sonnet-20240620-v1:0':
+    case 'anthropic.claude-3-5-sonnet-20241022-v2:0':
+      return { tool: { name: 'responseType' } };
+    case 'mistral.mistral-large-2402-v1:0':
+    case 'mistral.mistral-large-2407-v1:0':
+      return { any: {} };
+    default:
+      return undefined;
+  }
+};

From 2192c764c6552c17bcd1281114eab66d4863ae3a Mon Sep 17 00:00:00 2001
From: Ian Saultz <52051793+atierian@users.noreply.github.com>
Date: Fri, 13 Dec 2024 14:11:26 -0500
Subject: [PATCH 2/5] fix(generation): handle single quoted stringifed JSON
 tool input

---
 .../invoke-bedrock-resolver-fn.template.js    | 49 +++++++++++++++++--
 .../src/resolvers/invoke-bedrock.ts           | 12 +----
 2 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
index 4f236ca7d6..e1f84b853e 100644
--- a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
+++ b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
@@ -53,13 +53,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
+  }
+
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is `String` / `String!`, the value is `false` and we don't attempt any fallback parsing.
+  // If the return type isn't `String` / `String!`, the valie is `true` and the toolUse input is a `string`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if ([[NON_STRING_RESPONSE_TYPE]] && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
 
-  [[NON_STRING_RESPONSE_HANDLING]]
   return value;
 }
 
@@ -73,3 +81,38 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // `@aws-appsync/no-try: Try statements are not supported`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains `'` where it should contain `\"`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // `error @aws-appsync/no-regex: Regex literals are not supported`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}
\ No newline at end of file
diff --git a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
index 46e7386af0..c7bbaca9e9 100644
--- a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
+++ b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
@@ -20,15 +20,7 @@ export const createInvokeBedrockResolverFunction = (config: GenerationConfigurat
   const TOOL_CONFIG = JSON.stringify(toolConfig);
   const SYSTEM_PROMPT = JSON.stringify(config.systemPrompt);
   const INFERENCE_CONFIG = getInferenceConfigResolverDefinition(inferenceConfiguration);
-
-  const NON_STRING_RESPONSE_HANDLING = stringTypedScalarTypes.includes(getBaseType(config.field.type))
-    ? ''
-    : `// Added for non-string scalar response types
-  // This catches the occasional stringified JSON response.
-  if (typeof value === 'string') {
-    return JSON.parse(value);
-  }`;
-
+  const NON_STRING_RESPONSE_TYPE = (!stringTypedScalarTypes.includes(getBaseType(config.field.type))).toString()
   const PACKAGE_METADATA = `'${packageName}#${packageVersion}'`;
 
   const resolver = generateResolver('invoke-bedrock-resolver-fn.template.js', {
@@ -36,7 +28,7 @@ export const createInvokeBedrockResolverFunction = (config: GenerationConfigurat
     TOOL_CONFIG,
     SYSTEM_PROMPT,
     INFERENCE_CONFIG,
-    NON_STRING_RESPONSE_HANDLING,
+    NON_STRING_RESPONSE_TYPE,
     PACKAGE_METADATA,
   });
 

From 8f57b789570a82bc497977bab43f5e21ef134f9d Mon Sep 17 00:00:00 2001
From: Ian Saultz <52051793+atierian@users.noreply.github.com>
Date: Fri, 13 Dec 2024 14:11:38 -0500
Subject: [PATCH 3/5] test: update snapshots

---
 ...raphql-generation-transformer.test.ts.snap | 256 ++++++++++++++++--
 .../src/resolvers/invoke-bedrock.ts           |   2 +-
 2 files changed, 228 insertions(+), 30 deletions(-)

diff --git a/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap b/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
index 147b22b833..1121076bd8 100644
--- a/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
+++ b/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
@@ -112,17 +112,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // Added for non-string scalar response types
-  // This catches the occasional stringified JSON response.
-  if (typeof value === 'string') {
-    return JSON.parse(value);
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if (true && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
+
   return value;
 }
 
@@ -136,7 +140,41 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
-",
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // \`@aws-appsync/no-try: Try statements are not supported\`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains \`'\` where it should contain \`\\"\`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // \`error @aws-appsync/no-regex: Regex literals are not supported\`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}",
 }
 `;
 
@@ -252,17 +290,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // Added for non-string scalar response types
-  // This catches the occasional stringified JSON response.
-  if (typeof value === 'string') {
-    return JSON.parse(value);
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if (true && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
+
   return value;
 }
 
@@ -276,7 +318,41 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
-",
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // \`@aws-appsync/no-try: Try statements are not supported\`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains \`'\` where it should contain \`\\"\`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // \`error @aws-appsync/no-regex: Regex literals are not supported\`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}",
 }
 `;
 
@@ -391,17 +467,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // Added for non-string scalar response types
-  // This catches the occasional stringified JSON response.
-  if (typeof value === 'string') {
-    return JSON.parse(value);
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if (true && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
+
   return value;
 }
 
@@ -415,7 +495,41 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
-"
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // \`@aws-appsync/no-try: Try statements are not supported\`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains \`'\` where it should contain \`\\"\`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // \`error @aws-appsync/no-regex: Regex literals are not supported\`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}"
 `;
 
 exports[`generation route scalar type 1`] = `
@@ -547,13 +661,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
+  }
+
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if (false && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
 
-  
   return value;
 }
 
@@ -567,7 +689,41 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
-",
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // \`@aws-appsync/no-try: Try statements are not supported\`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains \`'\` where it should contain \`\\"\`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // \`error @aws-appsync/no-regex: Regex literals are not supported\`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}",
 }
 `;
 
@@ -628,13 +784,21 @@ export function response(ctx) {
   }
 
   const body = JSON.parse(ctx.result.body);
-  const value = body?.output?.message?.content?.[0]?.toolUse?.input?.value;
+  let value = body?.output?.message?.content?.find((content) => !!content.toolUse)?.toolUse?.input?.value;
 
   if (!value) {
-    util.error('Invalid Bedrock response', 'InvalidResponseException');
+    util.error('Invalid foundation model response', 'InvalidResponseException');
+  }
+
+  // The first condition (the boolean literal) in this if statement represents whether the
+  // return type of the generation route is a raw string or not.
+  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
+  if (false && typeof value === 'string') {
+    return parseIncorrectlyStringifiedJSON(value);
   }
 
-  
   return value;
 }
 
@@ -648,6 +812,40 @@ function createUserAgent(request) {
   }
   return userAgent;
 }
-",
+
+function parseIncorrectlyStringifiedJSON(input) {
+  // Try statements are not supported:
+  // \`@aws-appsync/no-try: Try statements are not supported\`
+
+  // This initial attempt covers the case where the tool input is valid stringified JSON
+  let value = JSON.parse(input);
+  // A failed parse attempt doesn't throw an error in resolver functions.
+  // It returns an empty string, so a truthiness check suffices.
+  if (value) return value;
+
+  // Since the tool input wasn't valid stringified JSON, we're assuming that
+  // it contains \`'\` where it should contain \`\\"\`. Some foundation models like to do this.
+  // This is our last fallback attempt and covers the cases observed in the wild.
+
+  // Regular expression is not supported in resolver functions:
+  // \`error @aws-appsync/no-regex: Regex literals are not supported\`
+  // However, raw string inputs are processed by the underlying Java runtime.
+  // So the patterns used are valid Java patterns, and not necessarily valid JavaScript patterns
+
+  // Replaces single quotes with double quotes, handling escaped single quotes.
+  value = input
+    // Replace any escaped single quotes with a marker.
+    .replaceAll("\\\\\\\\'", "___ESCAPED_QUOTE___")
+    // Replace all remaining single quotes with double quotes
+    .replaceAll("'", "\\"")
+    // Restore escaped single quotes
+    .replaceAll("___ESCAPED_QUOTE___", "'");
+
+  value = JSON.parse(value);
+  if (value) return value;
+
+  // Nothing more to do, time to bail.
+  util.error('Unable to parse foundation model response', 'InvalidResponseException')
+}",
 }
 `;
diff --git a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
index c7bbaca9e9..cd7af66659 100644
--- a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
+++ b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
@@ -20,7 +20,7 @@ export const createInvokeBedrockResolverFunction = (config: GenerationConfigurat
   const TOOL_CONFIG = JSON.stringify(toolConfig);
   const SYSTEM_PROMPT = JSON.stringify(config.systemPrompt);
   const INFERENCE_CONFIG = getInferenceConfigResolverDefinition(inferenceConfiguration);
-  const NON_STRING_RESPONSE_TYPE = (!stringTypedScalarTypes.includes(getBaseType(config.field.type))).toString()
+  const NON_STRING_RESPONSE_TYPE = (!stringTypedScalarTypes.includes(getBaseType(config.field.type))).toString();
   const PACKAGE_METADATA = `'${packageName}#${packageVersion}'`;
 
   const resolver = generateResolver('invoke-bedrock-resolver-fn.template.js', {

From ba8ff59da806024b6beb22e7bc02a9189f674251 Mon Sep 17 00:00:00 2001
From: Ian Saultz <52051793+atierian@users.noreply.github.com>
Date: Fri, 17 Jan 2025 17:03:57 -0500
Subject: [PATCH 4/5] generation tool use > allow list for tool choice

---
 ...raphql-generation-transformer.test.ts.snap | 55 ++++++++-------
 ...ify-graphql-generation-transformer.test.ts | 67 ++++++++++++++++++-
 .../invoke-bedrock-resolver-fn.template.js    | 11 +--
 .../src/resolvers/invoke-bedrock.ts           |  4 +-
 .../src/utils/tools.ts                        | 50 ++++++++++----
 5 files changed, 141 insertions(+), 46 deletions(-)

diff --git a/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap b/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
index 1121076bd8..023b2ca6f5 100644
--- a/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
+++ b/packages/amplify-graphql-generation-transformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
@@ -118,12 +118,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
-  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is \`String\`, \`ID\`, or \`AWSJSON\` (including required \`!\` ), the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is \`true\`.
+  const isScalarStringResponseType = false;
+  // When \`isScalarStringResponseType\` is \`false\` and the toolUse input is a \`string\`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if (true && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
@@ -296,12 +297,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
-  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is \`String\`, \`ID\`, or \`AWSJSON\` (including required \`!\` ), the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is \`true\`.
+  const isScalarStringResponseType = false;
+  // When \`isScalarStringResponseType\` is \`false\` and the toolUse input is a \`string\`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if (true && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
@@ -473,12 +475,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
-  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is \`String\`, \`ID\`, or \`AWSJSON\` (including required \`!\` ), the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is \`true\`.
+  const isScalarStringResponseType = false;
+  // When \`isScalarStringResponseType\` is \`false\` and the toolUse input is a \`string\`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if (true && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
@@ -667,12 +670,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
-  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is \`String\`, \`ID\`, or \`AWSJSON\` (including required \`!\` ), the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is \`true\`.
+  const isScalarStringResponseType = true;
+  // When \`isScalarStringResponseType\` is \`false\` and the toolUse input is a \`string\`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if (false && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
@@ -790,12 +794,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is \`String\` / \`String!\`, the value is \`false\` and we don't attempt any fallback parsing.
-  // If the return type isn't \`String\` / \`String!\`, the valie is \`true\` and the toolUse input is a \`string\`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is \`String\`, \`ID\`, or \`AWSJSON\` (including required \`!\` ), the value is \`false\` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is \`true\`.
+  const isScalarStringResponseType = true;
+  // When \`isScalarStringResponseType\` is \`false\` and the toolUse input is a \`string\`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if (false && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
diff --git a/packages/amplify-graphql-generation-transformer/src/__tests__/amplify-graphql-generation-transformer.test.ts b/packages/amplify-graphql-generation-transformer/src/__tests__/amplify-graphql-generation-transformer.test.ts
index 97d6ab58e5..0f8cc245f7 100644
--- a/packages/amplify-graphql-generation-transformer/src/__tests__/amplify-graphql-generation-transformer.test.ts
+++ b/packages/amplify-graphql-generation-transformer/src/__tests__/amplify-graphql-generation-transformer.test.ts
@@ -324,7 +324,72 @@ describe('generation route invalid inference configuration', () => {
     );
   });
 });
-// });
+
+describe('generation route model support', () => {
+  test.each([
+    // standard models
+    'anthropic.claude-3-haiku-20240307-v1:0',
+    'anthropic.claude-3-5-haiku-20241022-v1:0',
+    'anthropic.claude-3-sonnet-20240229-v1:0',
+    'anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'anthropic.claude-3-5-sonnet-20241022-v2:0',
+    'mistral.mistral-large-2402-v1:0',
+    'mistral.mistral-large-2407-v1:0',
+    'amazon.nova-pro-v1:0',
+    'amazon.nova-lite-v1:0',
+    'meta.llama3-1-405b-instruct-v1:0',
+    'ai21.jamba-1-5-mini-v1:0',
+    // cross-region inference model identifiers
+    'us.anthropic.claude-3-haiku-20240307-v1:0',
+    'eu.mistral.mistral-large-2402-v1:0',
+    'ap.amazon.nova-pro-v1:0',
+  ])('supported model: %s', (model) => {
+    const inputSchema = `
+      type Todo {
+        content: String
+        isDone: Boolean
+      }
+
+      type Query {
+          makeTodo(description: String!): Todo
+          @generation(
+            aiModel: "${model}",
+            systemPrompt: "Make a string based on the description.",
+          )
+      }
+    `;
+
+    const out = transform(inputSchema);
+    expect(out).toBeDefined();
+  });
+
+  test.each([
+    'amazon.nova-micro-v1:0',
+    'cohere.command-r-v1:0',
+    'cohere.command-r-plus-v1:0',
+    'meta.llama3-1-70b-instruct-v1:0',
+    'meta.llama3-1-8b-instruct-v1:0',
+    'mistral.mistral-small-2402-v1:0',
+    'some-random-model-v1:0',
+  ])('unsupported model: %s', (model) => {
+    const inputSchema = `
+      type Todo {
+        content: String
+        isDone: Boolean
+      }
+
+      type Query {
+          makeTodo(description: String!): Todo
+          @generation(
+            aiModel: "${model}",
+            systemPrompt: "Make a string based on the description.",
+          )
+      }
+    `;
+
+    expect(() => transform(inputSchema)).toThrow(`Model ${model} is not supported for Generation routes.`);
+  });
+});
 
 const getResolverResource = (queryName: string, resources?: Record<string, any>): Record<string, any> => {
   const resolverName = `Query${queryName}Resolver`;
diff --git a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
index e1f84b853e..4c2ac03817 100644
--- a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
+++ b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock-resolver-fn.template.js
@@ -59,12 +59,13 @@ export function response(ctx) {
     util.error('Invalid foundation model response', 'InvalidResponseException');
   }
 
-  // The first condition (the boolean literal) in this if statement represents whether the
-  // return type of the generation route is a raw string or not.
-  // If the return type is `String` / `String!`, the value is `false` and we don't attempt any fallback parsing.
-  // If the return type isn't `String` / `String!`, the valie is `true` and the toolUse input is a `string`,
+  // Represents whether the return type of the generation route is a raw string or not.
+  // If the return type is `String`, `ID`, or `AWSJSON` (including required `!` ), the value is `false` and we don't attempt any fallback parsing.
+  // If the return type isn't on the the string based scalars above, the value is `true`.
+  const isScalarStringResponseType = [[SCALAR_STRING_RESPONSE_TYPE]];
+  // When `isScalarStringResponseType` is `false` and the toolUse input is a `string`,
   // the foundation model has returned stringified JSON, so we attempt to parse it into a valid object.
-  if ([[NON_STRING_RESPONSE_TYPE]] && typeof value === 'string') {
+  if (!isScalarStringResponseType && typeof value === 'string') {
     return parseIncorrectlyStringifiedJSON(value);
   }
 
diff --git a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
index cd7af66659..cc62fec63d 100644
--- a/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
+++ b/packages/amplify-graphql-generation-transformer/src/resolvers/invoke-bedrock.ts
@@ -20,7 +20,7 @@ export const createInvokeBedrockResolverFunction = (config: GenerationConfigurat
   const TOOL_CONFIG = JSON.stringify(toolConfig);
   const SYSTEM_PROMPT = JSON.stringify(config.systemPrompt);
   const INFERENCE_CONFIG = getInferenceConfigResolverDefinition(inferenceConfiguration);
-  const NON_STRING_RESPONSE_TYPE = (!stringTypedScalarTypes.includes(getBaseType(config.field.type))).toString();
+  const SCALAR_STRING_RESPONSE_TYPE = (stringTypedScalarTypes.includes(getBaseType(config.field.type))).toString();
   const PACKAGE_METADATA = `'${packageName}#${packageVersion}'`;
 
   const resolver = generateResolver('invoke-bedrock-resolver-fn.template.js', {
@@ -28,7 +28,7 @@ export const createInvokeBedrockResolverFunction = (config: GenerationConfigurat
     TOOL_CONFIG,
     SYSTEM_PROMPT,
     INFERENCE_CONFIG,
-    NON_STRING_RESPONSE_TYPE,
+    SCALAR_STRING_RESPONSE_TYPE,
     PACKAGE_METADATA,
   });
 
diff --git a/packages/amplify-graphql-generation-transformer/src/utils/tools.ts b/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
index 44a019c27e..59dcbf7e62 100644
--- a/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
+++ b/packages/amplify-graphql-generation-transformer/src/utils/tools.ts
@@ -85,18 +85,42 @@ export const createResponseTypeTool = (config: GenerationDirectiveConfiguration,
 };
 
 const getToolChoice = (config: GenerationDirectiveConfiguration): ToolChoice => {
-  switch (config.aiModel) {
-    case 'anthropic.claude-3-opus-20240229-v1:0':
-    case 'anthropic.claude-3-haiku-20240307-v1:0':
-    case 'anthropic.claude-3-sonnet-20240229-v1:0':
-    case 'anthropic.claude-3-5-haiku-20241022-v1:0':
-    case 'anthropic.claude-3-5-sonnet-20240620-v1:0':
-    case 'anthropic.claude-3-5-sonnet-20241022-v2:0':
-      return { tool: { name: 'responseType' } };
-    case 'mistral.mistral-large-2402-v1:0':
-    case 'mistral.mistral-large-2407-v1:0':
-      return { any: {} };
-    default:
-      return undefined;
+  // Note: We're checking `includes()` below to avoid throwing a false positive
+  // for cross-region inference model identifiers, e.g. `us.anthropic.claude-3-5-sonnet-20241022-v2:0`.
+  const model = config.aiModel;
+
+  // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_SpecificToolChoice.html
+  // This field is only supported by Anthropic Claude 3 models.
+  const claude3Models = [
+    'anthropic.claude-3-opus-20240229-v1:0',
+    'anthropic.claude-3-haiku-20240307-v1:0',
+    'anthropic.claude-3-sonnet-20240229-v1:0',
+    'anthropic.claude-3-5-haiku-20241022-v1:0',
+    'anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'anthropic.claude-3-5-sonnet-20241022-v2:0',
+  ];
+  if (claude3Models.some((supportedModel) => model.includes(supportedModel))) {
+    return { tool: { name: 'responseType' } };
   }
+
+  // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
+  // ToolChoice is only supported by Anthropic Claude 3 models and by Mistral AI Mistral Large.
+  const mistralModels = ['mistral.mistral-large-2402-v1:0', 'mistral.mistral-large-2407-v1:0'];
+  if (mistralModels.some((supportedModel) => model.includes(supportedModel))) {
+    return { any: {} };
+  }
+
+  // These models do not support toolChoice but are known to work from testing.
+  const sansToolChoiceKnownWorkingModels = [
+    'amazon.nova-pro-v1:0',
+    'amazon.nova-lite-v1:0',
+    'meta.llama3-1-405b-instruct-v1:0',
+    'ai21.jamba-1-5-mini-v1:0',
+  ];
+
+  if (sansToolChoiceKnownWorkingModels.some((supportedModel) => model.includes(supportedModel))) {
+    return undefined;
+  }
+
+  throw new Error(`Model ${model} is not supported for Generation routes.`);
 };

From a9a7de2491f1cbc624f3bf5fd9ed37e7cfa6ad15 Mon Sep 17 00:00:00 2001
From: Ian Saultz <52051793+atierian@users.noreply.github.com>
Date: Fri, 17 Jan 2025 17:04:15 -0500
Subject: [PATCH 5/5] add evaluate code test stub

---
 .../__tests__/generations/generation-evaluate-code.test.ts    | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 packages/amplify-graphql-api-construct-tests/src/__tests__/generations/generation-evaluate-code.test.ts

diff --git a/packages/amplify-graphql-api-construct-tests/src/__tests__/generations/generation-evaluate-code.test.ts b/packages/amplify-graphql-api-construct-tests/src/__tests__/generations/generation-evaluate-code.test.ts
new file mode 100644
index 0000000000..7694dadd9a
--- /dev/null
+++ b/packages/amplify-graphql-api-construct-tests/src/__tests__/generations/generation-evaluate-code.test.ts
@@ -0,0 +1,4 @@
+
+describe('generation-evaluate-code', () => {
+  
+});