feat(conversation): add metrics and usage types and deserializers

.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.

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!

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');
+  });
+
+});

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');
+  });
 });

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(

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!

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

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 {

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 };

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 }),
+  };
+};