diff --git a/.changeset/bedrock-model-usage.md b/.changeset/bedrock-model-usage.md new file mode 100644 index 000000000..2dd218b7b --- /dev/null +++ b/.changeset/bedrock-model-usage.md @@ -0,0 +1,7 @@ +--- +"@aws-amplify/data-schema": minor +--- + +feat(conversation): add metrics and usage fields to conversation messages and stream events + +Adds support for Bedrock model usage tracking (inputTokens, outputTokens, totalTokens) and metrics (latencyMs) on conversation messages and stream events. diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap index 280ff819e..6d8b72891 100644 --- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap +++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap @@ -43,6 +43,8 @@ interface AmplifyAIConversationMessage { content: [AmplifyAIContentBlock] aiContext: AWSJSON toolConfiguration: AmplifyAIToolConfiguration + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage createdAt: AWSDateTime updatedAt: AWSDateTime owner: String @@ -210,9 +212,21 @@ type AmplifyAIConversationMessageStreamPart @aws_cognito_user_pools { contentBlockDoneAtIndex: Int stopReason: String errors: [AmplifyAIConversationTurnError] + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage p: String } +type AmplifyAIMetrics @aws_cognito_user_pools { + latencyMs: Int +} + +type AmplifyAIUsage @aws_cognito_user_pools { + inputTokens: Int + outputTokens: Int + totalTokens: Int +} + type AmplifyAIConversationTurnError @aws_cognito_user_pools { message: String! errorType: String! @@ -237,6 +251,8 @@ interface AmplifyAIConversationMessage { content: [AmplifyAIContentBlock] aiContext: AWSJSON toolConfiguration: AmplifyAIToolConfiguration + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage createdAt: AWSDateTime updatedAt: AWSDateTime owner: String @@ -404,9 +420,21 @@ type AmplifyAIConversationMessageStreamPart @aws_cognito_user_pools { contentBlockDoneAtIndex: Int stopReason: String errors: [AmplifyAIConversationTurnError] + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage p: String } +type AmplifyAIMetrics @aws_cognito_user_pools { + latencyMs: Int +} + +type AmplifyAIUsage @aws_cognito_user_pools { + inputTokens: Int + outputTokens: Int + totalTokens: Int +} + type AmplifyAIConversationTurnError @aws_cognito_user_pools { message: String! errorType: String! @@ -431,6 +459,8 @@ interface AmplifyAIConversationMessage { content: [AmplifyAIContentBlock] aiContext: AWSJSON toolConfiguration: AmplifyAIToolConfiguration + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage createdAt: AWSDateTime updatedAt: AWSDateTime owner: String @@ -598,9 +628,21 @@ type AmplifyAIConversationMessageStreamPart @aws_cognito_user_pools { contentBlockDoneAtIndex: Int stopReason: String errors: [AmplifyAIConversationTurnError] + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage p: String } +type AmplifyAIMetrics @aws_cognito_user_pools { + latencyMs: Int +} + +type AmplifyAIUsage @aws_cognito_user_pools { + inputTokens: Int + outputTokens: Int + totalTokens: Int +} + type AmplifyAIConversationTurnError @aws_cognito_user_pools { message: String! errorType: String! @@ -638,6 +680,8 @@ interface AmplifyAIConversationMessage { content: [AmplifyAIContentBlock] aiContext: AWSJSON toolConfiguration: AmplifyAIToolConfiguration + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage createdAt: AWSDateTime updatedAt: AWSDateTime owner: String @@ -805,9 +849,21 @@ type AmplifyAIConversationMessageStreamPart @aws_cognito_user_pools { contentBlockDoneAtIndex: Int stopReason: String errors: [AmplifyAIConversationTurnError] + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage p: String } +type AmplifyAIMetrics @aws_cognito_user_pools { + latencyMs: Int +} + +type AmplifyAIUsage @aws_cognito_user_pools { + inputTokens: Int + outputTokens: Int + totalTokens: Int +} + type AmplifyAIConversationTurnError @aws_cognito_user_pools { message: String! errorType: String! diff --git a/packages/data-schema/__tests__/internals/ai/conversationStreamEventDeserializers.test.ts b/packages/data-schema/__tests__/internals/ai/conversationStreamEventDeserializers.test.ts new file mode 100644 index 000000000..7a17769b6 --- /dev/null +++ b/packages/data-schema/__tests__/internals/ai/conversationStreamEventDeserializers.test.ts @@ -0,0 +1,60 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { convertItemToConversationStreamEvent } from '../../../src/runtime/internals/ai/conversationStreamEventDeserializers'; + +describe('convertItemToConversationStreamEvent()', () => { + const mockBaseEvent = { + id: 'message-id', + conversationId: 'conversation-id', + associatedUserMessageId: 'associated-user-message-id', + }; + + it('includes metrics and usage in turn-done event', () => { + const metrics = { latencyMs: 150 }; + const usage = { inputTokens: 50, outputTokens: 100, totalTokens: 150 }; + const { next, error } = convertItemToConversationStreamEvent({ + ...mockBaseEvent, + stopReason: 'end_turn', + metrics, + usage, + }); + + expect(error).toBeUndefined(); + expect(next).toMatchObject({ + ...mockBaseEvent, + stopReason: 'end_turn', + metrics, + usage, + }); + }); + + it('omits metrics and usage when values are undefined', () => { + const { next, error } = convertItemToConversationStreamEvent({ + ...mockBaseEvent, + contentBlockIndex: 0, + contentBlockDeltaIndex: 0, + contentBlockText: 'hello', + }); + + expect(error).toBeUndefined(); + expect(next?.metrics).toBeUndefined(); + expect(next?.usage).toBeUndefined(); + }); + + it('strips null metrics and usage from text event', () => { + const { next, error } = convertItemToConversationStreamEvent({ + ...mockBaseEvent, + contentBlockIndex: 0, + contentBlockDeltaIndex: 0, + contentBlockText: 'hello', + metrics: null, + usage: null, + }); + + expect(error).toBeUndefined(); + expect(next).not.toHaveProperty('metrics'); + expect(next).not.toHaveProperty('usage'); + }); + +}); diff --git a/packages/data-schema/__tests__/internals/ai/convertItemToConversationMessage.test.ts b/packages/data-schema/__tests__/internals/ai/convertItemToConversationMessage.test.ts index 24f3f2987..09cab8ee0 100644 --- a/packages/data-schema/__tests__/internals/ai/convertItemToConversationMessage.test.ts +++ b/packages/data-schema/__tests__/internals/ai/convertItemToConversationMessage.test.ts @@ -36,4 +36,23 @@ describe('convertItemToConversationMessage()', () => { ).toStrictEqual(mockMessageItem); expect(mockDeserializeContent).toHaveBeenCalledWith(mockMessageContent); }); + + it('includes metrics and usage when present', () => { + const metrics = { latencyMs: 123 }; + const usage = { inputTokens: 10, outputTokens: 20, totalTokens: 30 }; + expect( + convertItemToConversationMessage({ + ...mockMessageItem, + metrics, + usage, + }), + ).toStrictEqual({ ...mockMessageItem, metrics, usage }); + }); + + it('omits metrics and usage when not present', () => { + const result = convertItemToConversationMessage(mockMessageItem); + expect(result).toStrictEqual(mockMessageItem); + expect(result).not.toHaveProperty('metrics'); + expect(result).not.toHaveProperty('usage'); + }); }); diff --git a/packages/data-schema/__tests__/internals/ai/createOnStreamEventFunction.test.ts b/packages/data-schema/__tests__/internals/ai/createOnStreamEventFunction.test.ts index 063c42af6..9af3459a5 100644 --- a/packages/data-schema/__tests__/internals/ai/createOnStreamEventFunction.test.ts +++ b/packages/data-schema/__tests__/internals/ai/createOnStreamEventFunction.test.ts @@ -91,6 +91,8 @@ describe('createOnStreamEventFunction()', () => { contentBlockDoneAtIndex: undefined, toolUse: undefined, stopReason: undefined, + metrics: undefined, + usage: undefined, }; onStreamEvent(mockHandler); @@ -131,6 +133,8 @@ describe('createOnStreamEventFunction()', () => { contentBlockDeltaIndex: undefined, text: undefined, role: undefined, + metrics: undefined, + usage: undefined, }; onStreamEvent(mockHandler); expect(mockCustomOpFactory).toHaveBeenCalledWith( diff --git a/packages/data-schema/src/ai/ConversationSchemaGraphQLTypes.ts b/packages/data-schema/src/ai/ConversationSchemaGraphQLTypes.ts index 40159d3af..6752439c8 100644 --- a/packages/data-schema/src/ai/ConversationSchemaGraphQLTypes.ts +++ b/packages/data-schema/src/ai/ConversationSchemaGraphQLTypes.ts @@ -12,6 +12,8 @@ interface AmplifyAIConversationMessage { content: [AmplifyAIContentBlock] aiContext: AWSJSON toolConfiguration: AmplifyAIToolConfiguration + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage createdAt: AWSDateTime updatedAt: AWSDateTime owner: String @@ -179,9 +181,21 @@ type AmplifyAIConversationMessageStreamPart @aws_cognito_user_pools { contentBlockDoneAtIndex: Int stopReason: String errors: [AmplifyAIConversationTurnError] + metrics: AmplifyAIMetrics + usage: AmplifyAIUsage p: String } +type AmplifyAIMetrics @aws_cognito_user_pools { + latencyMs: Int +} + +type AmplifyAIUsage @aws_cognito_user_pools { + inputTokens: Int + outputTokens: Int + totalTokens: Int +} + type AmplifyAIConversationTurnError @aws_cognito_user_pools { message: String! errorType: String! diff --git a/packages/data-schema/src/ai/ConversationType.ts b/packages/data-schema/src/ai/ConversationType.ts index 6031ea613..dc84da1ee 100644 --- a/packages/data-schema/src/ai/ConversationType.ts +++ b/packages/data-schema/src/ai/ConversationType.ts @@ -13,7 +13,7 @@ import { ConversationSendMessageInputContent, } from './types/ConversationMessageContent'; import { ToolConfiguration } from './types/ToolConfiguration'; -import { ConversationStreamErrorEvent, ConversationStreamEvent } from './types/ConversationStreamEvent'; +import { ConversationStreamErrorEvent, ConversationStreamEvent, ConversationStreamMetrics, ConversationStreamUsage } from './types/ConversationStreamEvent'; export const brandName = 'conversationCustomOperation'; @@ -25,6 +25,8 @@ export interface ConversationMessage { id: string; role: 'user' | 'assistant'; associatedUserMessageId?: string; + metrics?: ConversationStreamMetrics; + usage?: ConversationStreamUsage; } // conversation route types diff --git a/packages/data-schema/src/ai/types/ConversationStreamEvent.ts b/packages/data-schema/src/ai/types/ConversationStreamEvent.ts index f2cce47bf..bf5b53a5c 100644 --- a/packages/data-schema/src/ai/types/ConversationStreamEvent.ts +++ b/packages/data-schema/src/ai/types/ConversationStreamEvent.ts @@ -3,6 +3,16 @@ import { ToolUseBlock } from "./contentBlocks"; +export interface ConversationStreamMetrics { + latencyMs?: number; +} + +export interface ConversationStreamUsage { + inputTokens?: number; + outputTokens?: number; + totalTokens?: number; +} + export interface ConversationStreamTextEvent { id: string; conversationId: string; @@ -13,6 +23,8 @@ export interface ConversationStreamTextEvent { text: string; toolUse?: never; stopReason?: never; + metrics?: never; + usage?: never; } export interface ConversationStreamToolUseEvent { @@ -25,6 +37,8 @@ export interface ConversationStreamToolUseEvent { text?: never; toolUse: ToolUseBlock; stopReason?: never; + metrics?: never; + usage?: never; } export interface ConversationStreamDoneAtIndexEvent { @@ -37,6 +51,8 @@ export interface ConversationStreamDoneAtIndexEvent { text?: never; toolUse?: never; stopReason?: never; + metrics?: never; + usage?: never; } export interface ConversationStreamTurnDoneEvent { @@ -49,6 +65,8 @@ export interface ConversationStreamTurnDoneEvent { text?: never; toolUse?: never; stopReason: string; + metrics?: ConversationStreamMetrics; + usage?: ConversationStreamUsage; } export interface ConversationStreamErrorEvent { diff --git a/packages/data-schema/src/runtime/internals/ai/conversationStreamEventDeserializers.ts b/packages/data-schema/src/runtime/internals/ai/conversationStreamEventDeserializers.ts index ab6ee2a23..e424c0e9b 100644 --- a/packages/data-schema/src/runtime/internals/ai/conversationStreamEventDeserializers.ts +++ b/packages/data-schema/src/runtime/internals/ai/conversationStreamEventDeserializers.ts @@ -15,6 +15,8 @@ export const convertItemToConversationStreamEvent = ({ contentBlockText, contentBlockToolUse, stopReason, + metrics, + usage, errors, }: any): { next?: ConversationStreamEvent, error?: ConversationStreamErrorEvent } => { if (errors) { @@ -36,6 +38,8 @@ export const convertItemToConversationStreamEvent = ({ text: contentBlockText, toolUse: deserializeToolUseBlock(contentBlockToolUse), stopReason, + metrics, + usage, }); return { next }; diff --git a/packages/data-schema/src/runtime/internals/ai/convertItemToConversationMessage.ts b/packages/data-schema/src/runtime/internals/ai/convertItemToConversationMessage.ts index 3117bde9a..ad5a847c2 100644 --- a/packages/data-schema/src/runtime/internals/ai/convertItemToConversationMessage.ts +++ b/packages/data-schema/src/runtime/internals/ai/convertItemToConversationMessage.ts @@ -3,16 +3,15 @@ import { deserializeContent } from './conversationMessageDeserializers'; -export const convertItemToConversationMessage = ({ - content, - createdAt, - id, - conversationId, - role, -}: any) => ({ - content: deserializeContent(content ?? []), - conversationId, - createdAt, - id, - role, -}); +export const convertItemToConversationMessage = (item: any) => { + const { content, createdAt, id, conversationId, role, metrics, usage } = item; + return { + content: deserializeContent(content ?? []), + conversationId, + createdAt, + id, + role, + ...(metrics != null && { metrics }), + ...(usage != null && { usage }), + }; +}; diff --git a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/ai-conversation.ts.snap b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/ai-conversation.ts.snap index c6c9cbeff..5f29fe23a 100644 --- a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/ai-conversation.ts.snap +++ b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/ai-conversation.ts.snap @@ -319,7 +319,7 @@ exports[`AI Conversation Routes Messages Subscribe to messages 1`] = ` "endpoint": undefined, "query": " subscription($conversationId: ID!) { - onCreateAssistantResponseChatBot(conversationId: $conversationId) {id owner conversationId associatedUserMessageId contentBlockIndex contentBlockText contentBlockDeltaIndex contentBlockToolUse { toolUseId name input type } contentBlockDoneAtIndex stopReason errors { message errorType } p} + onCreateAssistantResponseChatBot(conversationId: $conversationId) {id owner conversationId associatedUserMessageId contentBlockIndex contentBlockText contentBlockDeltaIndex contentBlockToolUse { toolUseId name input type } contentBlockDoneAtIndex stopReason errors { message errorType } metrics { latencyMs } usage { inputTokens outputTokens totalTokens } p} } ", "variables": { @@ -363,6 +363,14 @@ exports[`AI Conversation Routes Messages Subscribe to messages 2`] = ` message errorType } + metrics { + latencyMs + } + usage { + inputTokens + outputTokens + totalTokens + } p } } diff --git a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/ai-conversation.ts b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/ai-conversation.ts index 40fd229a2..388ac3be2 100644 --- a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/ai-conversation.ts +++ b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/ai-conversation.ts @@ -510,6 +510,55 @@ describe('AI Conversation Routes', () => { // #endregion assertions }); + test('List messages with metrics and usage', async () => { + // #region mocking + const messageWithMetrics = { + ...sampleConversationMessage1, + metrics: { latencyMs: 200 }, + usage: { inputTokens: 10, outputTokens: 25, totalTokens: 35 }, + }; + const { spy, generateClient } = mockedGenerateClient([ + { + data: { + getConversation: sampleConversation, + }, + }, + { + data: { + listMessages: { + items: [messageWithMetrics], + }, + }, + }, + ]); + // simulated amplifyconfiguration.json + const config = await buildAmplifyConfig(schema); + // #endregion mocking + // #region api call + // App.tsx + Amplify.configure(config); + + const client = generateClient(); + // get conversation + const { data: conversation } = await client.conversations.chatBot.get({ + id: sampleConversation.id, + }); + // list conversation messages + const { data: messages, errors: listMessagesErrors } = + (await conversation?.listMessages()) ?? {}; + // #endregion api call + // #region assertions + expect(listMessagesErrors).toBeUndefined(); + expect(messages).toStrictEqual([messageWithMetrics]); + expect(messages?.[0]).toHaveProperty('metrics', { latencyMs: 200 }); + expect(messages?.[0]).toHaveProperty('usage', { + inputTokens: 10, + outputTokens: 25, + totalTokens: 35, + }); + // #endregion assertions + }); + test('Paginate messages', async () => { // #region mocking const sampleNextToken = 'next-token'; @@ -609,6 +658,8 @@ describe('AI Conversation Routes', () => { ...rest, }; expect(mockNextHandler).toHaveBeenCalledWith(expectedConversationStreamEvent); + expect(expectedConversationStreamEvent).not.toHaveProperty('metrics'); + expect(expectedConversationStreamEvent).not.toHaveProperty('usage'); }); test('Tool use event', async () => { @@ -658,6 +709,22 @@ describe('AI Conversation Routes', () => { expect(mockNextHandler).toHaveBeenCalledWith(sampleConversationStreamTurnDoneEvent); }); + test('Turn done event with metrics and usage', async () => { + const turnDoneWithMetrics = { + ...sampleConversationStreamTurnDoneEvent, + metrics: { latencyMs: 150 }, + usage: { inputTokens: 50, outputTokens: 100, totalTokens: 150 }, + }; + subs.onCreateAssistantResponseChatBot.next({ + data: { + onCreateAssistantResponseChatBot: turnDoneWithMetrics, + }, + }); + + await pause(1); + expect(mockNextHandler).toHaveBeenCalledWith(turnDoneWithMetrics); + }); + test('Error event', async () => { subs.onCreateAssistantResponseChatBot.next({ data: {