Skip to content

Commit 0f09a2c

Browse files
pjkrokerpjkroker
committed
feat(graphql-conversation-transformer): add AmplifyAIMetrics and AmplifyAIUsage type definitions
Register explicit GraphQL types for metrics and usage in prepare handler. Update test schema and snapshots.
1 parent 18ef668 commit 0f09a2c

3 files changed

Lines changed: 60 additions & 9 deletions

File tree

‎packages/amplify-graphql-conversation-transformer/src/__tests__/__snapshots__/amplify-graphql-conversation-transformer.test.ts.snap‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,8 @@ export function request(ctx) {
214214
associatedUserMessageId,
215215
accumulatedTurnContent,
216216
errors,
217+
metrics,
218+
usage,
217219
} = ctx.args.input;
218220

219221
const { owner } = ctx.args;
@@ -230,7 +232,7 @@ export function request(ctx) {
230232
const { createdAt, updatedAt } = ctx.stash.defaultValues;
231233

232234
const assistantResponseId = \`\${associatedUserMessageId}#response\`;
233-
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt';
235+
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt, #metrics = :metrics, #usage = :usage';
234236

235237
const expressionValues = util.dynamodb.toMapValues({
236238
':typename': 'ConversationMessagePirateChat',
@@ -241,6 +243,8 @@ export function request(ctx) {
241243
':owner': owner,
242244
':createdAt': createdAt,
243245
':updatedAt': updatedAt,
246+
':metrics': metrics,
247+
':usage': usage,
244248
});
245249

246250
// https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ReservedWords.html
@@ -253,6 +257,8 @@ export function request(ctx) {
253257
'#owner': 'owner',
254258
'#createdAt': 'createdAt',
255259
'#updatedAt': 'updatedAt',
260+
'#metrics': 'metrics',
261+
'#usage': 'usage',
256262
};
257263

258264
return {
@@ -676,7 +682,7 @@ export function request(ctx) {
676682
const { graphqlApiEndpoint } = ctx.stash;
677683
const userAgent = createUserAgent(request);
678684

679-
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p';
685+
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p metrics { latencyMs } usage { inputTokens outputTokens totalTokens }';
680686

681687
const streamingResponseMutation = {
682688
name: 'createAssistantResponseStreamPirateChat',
@@ -1078,6 +1084,8 @@ export function request(ctx) {
10781084
associatedUserMessageId,
10791085
accumulatedTurnContent,
10801086
errors,
1087+
metrics,
1088+
usage,
10811089
} = ctx.args.input;
10821090

10831091
const { owner } = ctx.args;
@@ -1094,7 +1102,7 @@ export function request(ctx) {
10941102
const { createdAt, updatedAt } = ctx.stash.defaultValues;
10951103

10961104
const assistantResponseId = \`\${associatedUserMessageId}#response\`;
1097-
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt';
1105+
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt, #metrics = :metrics, #usage = :usage';
10981106

10991107
const expressionValues = util.dynamodb.toMapValues({
11001108
':typename': 'ConversationMessagePirateChat',
@@ -1105,6 +1113,8 @@ export function request(ctx) {
11051113
':owner': owner,
11061114
':createdAt': createdAt,
11071115
':updatedAt': updatedAt,
1116+
':metrics': metrics,
1117+
':usage': usage,
11081118
});
11091119

11101120
// https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ReservedWords.html
@@ -1117,6 +1127,8 @@ export function request(ctx) {
11171127
'#owner': 'owner',
11181128
'#createdAt': 'createdAt',
11191129
'#updatedAt': 'updatedAt',
1130+
'#metrics': 'metrics',
1131+
'#usage': 'usage',
11201132
};
11211133

11221134
return {
@@ -1540,7 +1552,7 @@ export function request(ctx) {
15401552
const { graphqlApiEndpoint } = ctx.stash;
15411553
const userAgent = createUserAgent(request);
15421554

1543-
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p';
1555+
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p metrics { latencyMs } usage { inputTokens outputTokens totalTokens }';
15441556

15451557
const streamingResponseMutation = {
15461558
name: 'createAssistantResponseStreamPirateChat',
@@ -1942,6 +1954,8 @@ export function request(ctx) {
19421954
associatedUserMessageId,
19431955
accumulatedTurnContent,
19441956
errors,
1957+
metrics,
1958+
usage,
19451959
} = ctx.args.input;
19461960

19471961
const { owner } = ctx.args;
@@ -1958,7 +1972,7 @@ export function request(ctx) {
19581972
const { createdAt, updatedAt } = ctx.stash.defaultValues;
19591973

19601974
const assistantResponseId = \`\${associatedUserMessageId}#response\`;
1961-
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt';
1975+
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt, #metrics = :metrics, #usage = :usage';
19621976

19631977
const expressionValues = util.dynamodb.toMapValues({
19641978
':typename': 'ConversationMessagePirateChat',
@@ -1969,6 +1983,8 @@ export function request(ctx) {
19691983
':owner': owner,
19701984
':createdAt': createdAt,
19711985
':updatedAt': updatedAt,
1986+
':metrics': metrics,
1987+
':usage': usage,
19721988
});
19731989

19741990
// https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ReservedWords.html
@@ -1981,6 +1997,8 @@ export function request(ctx) {
19811997
'#owner': 'owner',
19821998
'#createdAt': 'createdAt',
19831999
'#updatedAt': 'updatedAt',
2000+
'#metrics': 'metrics',
2001+
'#usage': 'usage',
19842002
};
19852003

19862004
return {
@@ -2404,7 +2422,7 @@ export function request(ctx) {
24042422
const { graphqlApiEndpoint } = ctx.stash;
24052423
const userAgent = createUserAgent(request);
24062424

2407-
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p';
2425+
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p metrics { latencyMs } usage { inputTokens outputTokens totalTokens }';
24082426

24092427
const streamingResponseMutation = {
24102428
name: 'createAssistantResponseStreamPirateChat',
@@ -2806,6 +2824,8 @@ export function request(ctx) {
28062824
associatedUserMessageId,
28072825
accumulatedTurnContent,
28082826
errors,
2827+
metrics,
2828+
usage,
28092829
} = ctx.args.input;
28102830

28112831
const { owner } = ctx.args;
@@ -2822,7 +2842,7 @@ export function request(ctx) {
28222842
const { createdAt, updatedAt } = ctx.stash.defaultValues;
28232843

28242844
const assistantResponseId = \`\${associatedUserMessageId}#response\`;
2825-
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt';
2845+
const expression = 'SET #typename = :typename, #conversationId = :conversationId, #associatedUserMessageId = :associatedUserMessageId, #role = :role, #content = :content, #owner = :owner, #createdAt = if_not_exists(#createdAt, :createdAt), #updatedAt = :updatedAt, #metrics = :metrics, #usage = :usage';
28262846

28272847
const expressionValues = util.dynamodb.toMapValues({
28282848
':typename': 'ConversationMessagePirateChat',
@@ -2833,6 +2853,8 @@ export function request(ctx) {
28332853
':owner': owner,
28342854
':createdAt': createdAt,
28352855
':updatedAt': updatedAt,
2856+
':metrics': metrics,
2857+
':usage': usage,
28362858
});
28372859

28382860
// https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ReservedWords.html
@@ -2845,6 +2867,8 @@ export function request(ctx) {
28452867
'#owner': 'owner',
28462868
'#createdAt': 'createdAt',
28472869
'#updatedAt': 'updatedAt',
2870+
'#metrics': 'metrics',
2871+
'#usage': 'usage',
28482872
};
28492873

28502874
return {
@@ -3268,7 +3292,7 @@ export function request(ctx) {
32683292
const { graphqlApiEndpoint } = ctx.stash;
32693293
const userAgent = createUserAgent(request);
32703294

3271-
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p';
3295+
const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message } p metrics { latencyMs } usage { inputTokens outputTokens totalTokens }';
32723296

32733297
const streamingResponseMutation = {
32743298
name: 'createAssistantResponseStreamPirateChat',

‎packages/amplify-graphql-conversation-transformer/src/graphql-types/message-model.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,24 @@ export const createConversationTurnErrorInput = (): InputObjectTypeDefinitionNod
158158
};
159159
};
160160

161+
export const constructMetricsType = (): ObjectTypeDefinitionNode => ({
162+
kind: 'ObjectTypeDefinition',
163+
name: { kind: 'Name', value: 'AmplifyAIMetrics' },
164+
fields: [makeField('latencyMs', [], makeNamedType('Int'))],
165+
directives: [],
166+
});
167+
168+
export const constructUsageType = (): ObjectTypeDefinitionNode => ({
169+
kind: 'ObjectTypeDefinition',
170+
name: { kind: 'Name', value: 'AmplifyAIUsage' },
171+
fields: [
172+
makeField('inputTokens', [], makeNamedType('Int')),
173+
makeField('outputTokens', [], makeNamedType('Int')),
174+
makeField('totalTokens', [], makeNamedType('Int')),
175+
],
176+
directives: [],
177+
});
178+
161179
export const createAssistantStreamingMutationField = (fieldName: string, inputTypeName: string): FieldDefinitionNode => {
162180
const args = [makeInputValueDefinition('input', makeNonNullType(makeNamedType(inputTypeName)))];
163181
const cognitoAuthDirective = makeDirective('aws_cognito_user_pools', []);
@@ -319,6 +337,8 @@ export const constructStreamResponseType = (): ObjectTypeDefinitionNode => {
319337
makeField('contentBlockDoneAtIndex', [], makeNamedType('Int')),
320338

321339
makeField('stopReason', [], makeNamedType('String')),
340+
makeField('metrics', [], makeNamedType('AmplifyAIMetrics')),
341+
makeField('usage', [], makeNamedType('AmplifyAIUsage')),
322342
],
323343
};
324344
};

‎packages/amplify-graphql-conversation-transformer/src/transformer-steps/conversation-prepare-handler.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
22
import { BelongsToTransformer, HasManyTransformer } from '@aws-amplify/graphql-relational-transformer';
33
import { DDB_AMPLIFY_MANAGED_DATASOURCE_STRATEGY, InvalidTransformerError } from '@aws-amplify/graphql-transformer-core';
44
import { ConversationDirectiveConfiguration } from '../conversation-directive-configuration';
5-
import { constructStreamResponseType, createConversationTurnErrorInput } from '../graphql-types/message-model';
5+
import {
6+
constructStreamResponseType,
7+
createConversationTurnErrorInput,
8+
constructMetricsType,
9+
constructUsageType,
10+
} from '../graphql-types/message-model';
611
import { TransformerAuthProvider, TransformerPrepareStepContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
712

813
/**
@@ -50,6 +55,8 @@ export class ConversationPrepareHandler {
5055
// add once per schema
5156
const conversationTurnErrorInput = createConversationTurnErrorInput();
5257
ctx.output.addInput(conversationTurnErrorInput);
58+
ctx.output.addObject(constructMetricsType());
59+
ctx.output.addObject(constructUsageType());
5360

5461
for (const directive of directives) {
5562
this.prepareResourcesForDirective(directive, ctx);

0 commit comments

Comments
 (0)