Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .changeset/fix-escape-double-quotes-ai-prompts.md
Original file line number Diff line number Diff line change
@@ -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.
Original file line number Diff line number Diff line change
@@ -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."');
});
});
});
12 changes: 1 addition & 11 deletions packages/data-schema/src/SchemaProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 ?? {},
);
Expand All @@ -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}`;
Expand Down
19 changes: 8 additions & 11 deletions packages/data-schema/src/ai/ConversationSchemaProcessor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@ import type {
} from './ConversationType';
import type { InferenceConfiguration } from './ModelType';

const escapeGraphQLString = (str: string): string =>
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this follow the same approach?

function escapeGraphQlString(str: string) {
return JSON.stringify(str);
}

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. The @generation path already uses JSON.stringify and I should have been consistent here too. I'll swap the manual regex for JSON.stringify(str).slice(1, -1) to use the same core while keeping the call-site quoting pattern. Pushing the fix now.

str
.replace(/\\/g, '\\\\')
.replace(/"/g, '\\"')
.replace(/\r?\n/g, '\\n');

export const createConversationField = (
typeDef: InternalConversationType,
typeName: string,
Expand All @@ -18,16 +24,7 @@ export const createConversationField = (

const args: Record<string, string> = {
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)
Expand Down Expand Up @@ -126,7 +123,7 @@ const getConversationToolsString = (tools: DataToolDefinition[]) =>
);
}
const toolDefinition = extractToolDefinition(tool);
return `{ name: "${name}", description: "${description}", ${toolDefinition} }`;
return `{ name: "${name}", description: "${escapeGraphQLString(description)}", ${toolDefinition} }`;
})
.join(', ');

Expand Down