diff --git a/.changeset/fix-escape-double-quotes-ai-prompts.md b/.changeset/fix-escape-double-quotes-ai-prompts.md new file mode 100644 index 000000000..2073372b5 --- /dev/null +++ b/.changeset/fix-escape-double-quotes-ai-prompts.md @@ -0,0 +1,7 @@ +--- +"@aws-amplify/data-schema": patch +--- + +fix: escape double quotes and backslashes in AI directive string arguments + +`@conversation` and `@generation` directives interpolate user-supplied strings (`systemPrompt`, tool `description`) directly into GraphQL SDL without escaping special characters. Any prompt containing a double quote or backslash produces invalid SDL, breaking schema compilation. This fix escapes all GraphQL special characters before interpolation. diff --git a/packages/data-schema/__tests__/ai/ConversationSchemaProcessor.test.ts b/packages/data-schema/__tests__/ai/ConversationSchemaProcessor.test.ts new file mode 100644 index 000000000..18f8d783d --- /dev/null +++ b/packages/data-schema/__tests__/ai/ConversationSchemaProcessor.test.ts @@ -0,0 +1,95 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { a } from '../../src/index'; +import { defineFunctionStub } from '../utils'; + +describe('GraphQL string escaping in AI schema directives', () => { + describe('@generation', () => { + test('escapes double quotes in systemPrompt', () => { + const schema = a.schema({ + Result: a.customType({ value: a.string() }), + makeResult: a + .generation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: 'Always say "yes" or "no".', + }) + .returns(a.ref('Result')), + }); + + const { schema: graphql } = schema.transform(); + + expect(graphql).toContain('systemPrompt: "Always say \\"yes\\" or \\"no\\"."'); + }); + }); + + describe('@conversation', () => { + test('escapes double quotes in systemPrompt', () => { + const schema = a.schema({ + ChatBot: a.conversation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: 'Say "hello" and "goodbye" to users.', + }).authorization((allow) => allow.owner()), + }); + + const { schema: graphql } = schema.transform(); + + expect(graphql).toContain('systemPrompt: "Say \\"hello\\" and \\"goodbye\\" to users."'); + }); + + test('escapes double quotes in tool description', () => { + const handler = defineFunctionStub({}); + const schema = a.schema({ + Profile: a.customType({ value: a.integer() }), + infoQuery: a + .query() + .returns(a.ref('Profile')) + .authorization((allow) => allow.publicApiKey()) + .handler(a.handler.function(handler)), + + ChatBot: a.conversation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: 'You are helpful.', + tools: [ + a.ai.dataTool({ + query: a.ref('infoQuery'), + name: 'infoQuery', + description: 'Fetches "live" profile data.', + }), + ], + }).authorization((allow) => allow.owner()), + }); + + const { schema: graphql } = schema.transform(); + + expect(graphql).toContain('description: "Fetches \\"live\\" profile data."'); + }); + + test('escapes backslashes in systemPrompt', () => { + const schema = a.schema({ + ChatBot: a.conversation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: 'Use path C:\\\\docs for all outputs.', + }).authorization((allow) => allow.owner()), + }); + + const { schema: graphql } = schema.transform(); + + expect(graphql).toContain('systemPrompt: "Use path C:\\\\\\\\docs for all outputs."'); + }); + + test('preserves newline escaping in multiline systemPrompt', () => { + const schema = a.schema({ + ChatBot: a.conversation({ + aiModel: a.ai.model('Claude 3 Haiku'), + systemPrompt: `You are helpful. +Respond in haiku.`, + }).authorization((allow) => allow.owner()), + }); + + const { schema: graphql } = schema.transform(); + + expect(graphql).toContain('systemPrompt: "You are helpful.\\nRespond in haiku."'); + }); + }); +}); diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts index bb2a31dc7..56636cf2d 100644 --- a/packages/data-schema/src/SchemaProcessor.ts +++ b/packages/data-schema/src/SchemaProcessor.ts @@ -539,16 +539,6 @@ function customOperationToGql( const { aiModel, systemPrompt, inferenceConfiguration } = typeDef.data.input; - // This is done to escape newlines in potentially multi-line system prompts - // e.g. - // generateStuff: a.generation({ - // aiModel: a.ai.model('Claude 3 Haiku'), - // systemPrompt: `Generate a haiku - // make it multiline`, - // }), - // - // It doesn't affect non multi-line string inputs for system prompts - const escapedSystemPrompt = systemPrompt.replace(/\r?\n/g, '\\n'); const inferenceConfigurationEntries = Object.entries( inferenceConfiguration ?? {}, ); @@ -558,7 +548,7 @@ function customOperationToGql( .map(([key, value]) => `${key}: ${value}`) .join(', ')} }` : ''; - gqlHandlerContent += `@generation(aiModel: "${aiModel.resourcePath}", systemPrompt: "${escapedSystemPrompt}"${inferenceConfigurationGql}) `; + gqlHandlerContent += `@generation(aiModel: "${aiModel.resourcePath}", systemPrompt: ${escapeGraphQlString(systemPrompt)}${inferenceConfigurationGql}) `; } const gqlField = `${callSignature}: ${returnTypeName} ${gqlHandlerContent}${authString}`; diff --git a/packages/data-schema/src/ai/ConversationSchemaProcessor.ts b/packages/data-schema/src/ai/ConversationSchemaProcessor.ts index 8d4e52a55..7b9f5fe36 100644 --- a/packages/data-schema/src/ai/ConversationSchemaProcessor.ts +++ b/packages/data-schema/src/ai/ConversationSchemaProcessor.ts @@ -8,6 +8,9 @@ import type { } from './ConversationType'; import type { InferenceConfiguration } from './ModelType'; +const escapeGraphQLString = (str: string): string => + JSON.stringify(str).slice(1, -1); + export const createConversationField = ( typeDef: InternalConversationType, typeName: string, @@ -18,16 +21,7 @@ export const createConversationField = ( const args: Record = { aiModel: aiModel.resourcePath, - // This is done to escape newlines in potentially multi-line system prompts - // e.g. - // realtorChat: a.conversation({ - // aiModel: a.ai.model('Claude 3 Haiku'), - // systemPrompt: `You are a helpful real estate assistant - // Respond in the poetic form of haiku.`, - // }), - // - // It doesn't affect non multi-line string inputs for system prompts - systemPrompt: systemPrompt.replace(/\r?\n/g, '\\n'), + systemPrompt: escapeGraphQLString(systemPrompt), }; // Add each arg with quotes (aiModel and systemPrompt) @@ -126,7 +120,7 @@ const getConversationToolsString = (tools: DataToolDefinition[]) => ); } const toolDefinition = extractToolDefinition(tool); - return `{ name: "${name}", description: "${description}", ${toolDefinition} }`; + return `{ name: "${name}", description: "${escapeGraphQLString(description)}", ${toolDefinition} }`; }) .join(', ');