Add metadata to user agent in conversation handler runtime (#2181)

* Add metadata to user agent in conversation handler runtime

* tests

* remove validation.

Author: sobolk

.changeset/thin-candles-perform.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': patch
+---
+
+Add metadata to user agent in conversation handler runtime.

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts

@@ -25,6 +25,7 @@ import {
 import { ConversationTurnEventToolsProvider } from './event-tools-provider';
 import { randomBytes, randomUUID } from 'node:crypto';
 import { ConversationMessageHistoryRetriever } from './conversation_message_history_retriever';
+import { UserAgentProvider } from './user_agent_provider';
 
 void describe('Bedrock converse adapter', () => {
   const commonEvent: Readonly<ConversationTurnEvent> = {
@@ -897,17 +898,20 @@ void describe('Bedrock converse adapter', () => {
       ...commonEvent,
     };
 
-    event.request.headers['x-amz-user-agent'] = 'testUserAgent';
-
     const bedrockClient = new BedrockRuntimeClient();
     const addMiddlewareMock = mock.method(bedrockClient.middlewareStack, 'add');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    mock.method(userAgentProvider, 'getUserAgent', () => 'testUserAgent');
 
     new BedrockConverseAdapter(
       event,
       [],
       bedrockClient,
       undefined,
-      messageHistoryRetriever
+      messageHistoryRetriever,
+      userAgentProvider
     );
 
     assert.strictEqual(addMiddlewareMock.mock.calls.length, 1);

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts

@@ -22,6 +22,7 @@ import { ConversationTurnEventToolsProvider } from './event-tools-provider';
 import { ConversationMessageHistoryRetriever } from './conversation_message_history_retriever';
 import * as bedrock from '@aws-sdk/client-bedrock-runtime';
 import { ValidationError } from './errors';
+import { UserAgentProvider } from './user_agent_provider';
 
 /**
  * This class is responsible for interacting with Bedrock Converse API
@@ -48,23 +49,22 @@ export class BedrockConverseAdapter {
     private readonly messageHistoryRetriever = new ConversationMessageHistoryRetriever(
       event
     ),
+    userAgentProvider = new UserAgentProvider(event),
     private readonly logger = console
   ) {
-    if (event.request.headers['x-amz-user-agent']) {
-      this.bedrockClient.middlewareStack.add(
-        (next) => (args) => {
-          // @ts-expect-error Request is typed as unknown.
-          // But this is recommended way to alter headers per https://github.com/aws/aws-sdk-js-v3/blob/main/README.md.
-          args.request.headers['x-amz-user-agent'] =
-            event.request.headers['x-amz-user-agent'];
-          return next(args);
-        },
-        {
-          step: 'build',
-          name: 'amplify-user-agent-injector',
-        }
-      );
-    }
+    this.bedrockClient.middlewareStack.add(
+      (next) => (args) => {
+        // @ts-expect-error Request is typed as unknown.
+        // But this is recommended way to alter headers per https://github.com/aws/aws-sdk-js-v3/blob/main/README.md.
+        args.request.headers['x-amz-user-agent'] =
+          userAgentProvider.getUserAgent();
+        return next(args);
+      },
+      {
+        step: 'build',
+        name: 'amplify-user-agent-injector',
+      }
+    );
     this.executableTools = [
       ...eventToolsProvider.getEventTools(),
       ...additionalTools,

packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.test.ts

@@ -12,6 +12,7 @@ import {
   GetQueryOutput,
   ListQueryOutput,
 } from './conversation_message_history_retriever';
+import { UserAgentProvider } from './user_agent_provider';
 
 type TestCase = {
   name: string;
@@ -704,7 +705,15 @@ void describe('Conversation message history retriever', () => {
 
   for (const testCase of testCases) {
     void it(testCase.name, async () => {
-      const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+      const userAgentProvider = new UserAgentProvider(
+        {} as unknown as ConversationTurnEvent
+      );
+      mock.method(userAgentProvider, 'getUserAgent', () => '');
+      const graphqlRequestExecutor = new GraphqlRequestExecutor(
+        '',
+        '',
+        userAgentProvider
+      );
       const executeGraphqlMock = mock.method(
         graphqlRequestExecutor,
         'executeGraphql',

packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.ts

@@ -4,6 +4,7 @@ import {
   ConversationTurnEvent,
 } from './types';
 import { GraphqlRequestExecutor } from './graphql_request_executor';
+import { UserAgentProvider } from './user_agent_provider';
 
 export type ConversationHistoryMessageItem = ConversationMessage & {
   id: string;
@@ -107,7 +108,7 @@ export class ConversationMessageHistoryRetriever {
     private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
       event.graphqlApiEndpoint,
       event.request.headers.authorization,
-      event.request.headers['x-amz-user-agent']
+      new UserAgentProvider(event)
     )
   ) {}
 

packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.test.ts

@@ -15,6 +15,7 @@ import {
   GraphqlRequest,
   GraphqlRequestExecutor,
 } from './graphql_request_executor';
+import { UserAgentProvider } from './user_agent_provider';
 
 void describe('Conversation turn response sender', () => {
   const event: ConversationTurnEvent = {
@@ -37,7 +38,19 @@ void describe('Conversation turn response sender', () => {
   };
 
   void it('sends response back to appsync', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    const userAgentProviderMock = mock.method(
+      userAgentProvider,
+      'getUserAgent',
+      () => 'testUserAgent'
+    );
+    const graphqlRequestExecutor = new GraphqlRequestExecutor(
+      '',
+      '',
+      userAgentProvider
+    );
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -47,6 +60,7 @@ void describe('Conversation turn response sender', () => {
     );
     const sender = new ConversationTurnResponseSender(
       event,
+      userAgentProvider,
       graphqlRequestExecutor
     );
     const response: Array<ContentBlock> = [
@@ -57,7 +71,14 @@ void describe('Conversation turn response sender', () => {
     ];
     await sender.sendResponse(response);
 
+    assert.strictEqual(userAgentProviderMock.mock.calls.length, 1);
+    assert.deepStrictEqual(userAgentProviderMock.mock.calls[0].arguments[0], {
+      'turn-response-type': 'single',
+    });
     assert.strictEqual(executeGraphqlMock.mock.calls.length, 1);
+    assert.deepStrictEqual(executeGraphqlMock.mock.calls[0].arguments[1], {
+      userAgent: 'testUserAgent',
+    });
     const request = executeGraphqlMock.mock.calls[0]
       .arguments[0] as GraphqlRequest<MutationResponseInput>;
     assert.deepStrictEqual(request, {
@@ -85,7 +106,15 @@ void describe('Conversation turn response sender', () => {
   });
 
   void it('serializes tool use input to JSON', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    mock.method(userAgentProvider, 'getUserAgent', () => '');
+    const graphqlRequestExecutor = new GraphqlRequestExecutor(
+      '',
+      '',
+      userAgentProvider
+    );
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -95,6 +124,7 @@ void describe('Conversation turn response sender', () => {
     );
     const sender = new ConversationTurnResponseSender(
       event,
+      userAgentProvider,
       graphqlRequestExecutor
     );
     const toolUseBlock: ContentBlock.ToolUseMember = {
@@ -140,7 +170,19 @@ void describe('Conversation turn response sender', () => {
   });
 
   void it('sends streaming response chunk back to appsync', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    const userAgentProviderMock = mock.method(
+      userAgentProvider,
+      'getUserAgent',
+      () => 'testUserAgent'
+    );
+    const graphqlRequestExecutor = new GraphqlRequestExecutor(
+      '',
+      '',
+      userAgentProvider
+    );
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -150,6 +192,7 @@ void describe('Conversation turn response sender', () => {
     );
     const sender = new ConversationTurnResponseSender(
       event,
+      userAgentProvider,
       graphqlRequestExecutor
     );
     const chunk: StreamingResponseChunk = {
@@ -162,7 +205,14 @@ void describe('Conversation turn response sender', () => {
     };
     await sender.sendResponseChunk(chunk);
 
+    assert.strictEqual(userAgentProviderMock.mock.calls.length, 1);
+    assert.deepStrictEqual(userAgentProviderMock.mock.calls[0].arguments[0], {
+      'turn-response-type': 'streaming',
+    });
     assert.strictEqual(executeGraphqlMock.mock.calls.length, 1);
+    assert.deepStrictEqual(executeGraphqlMock.mock.calls[0].arguments[1], {
+      userAgent: 'testUserAgent',
+    });
     const request = executeGraphqlMock.mock.calls[0]
       .arguments[0] as GraphqlRequest<MutationStreamingResponseInput>;
     assert.deepStrictEqual(request, {
@@ -181,7 +231,15 @@ void describe('Conversation turn response sender', () => {
   });
 
   void it('serializes tool use input to JSON when streaming', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    mock.method(userAgentProvider, 'getUserAgent', () => '');
+    const graphqlRequestExecutor = new GraphqlRequestExecutor(
+      '',
+      '',
+      userAgentProvider
+    );
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -191,6 +249,7 @@ void describe('Conversation turn response sender', () => {
     );
     const sender = new ConversationTurnResponseSender(
       event,
+      userAgentProvider,
       graphqlRequestExecutor
     );
     const toolUseBlock: ContentBlock.ToolUseMember = {
@@ -242,7 +301,19 @@ void describe('Conversation turn response sender', () => {
   });
 
   void it('sends errors response back to appsync', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
+    const userAgentProvider = new UserAgentProvider(
+      {} as unknown as ConversationTurnEvent
+    );
+    const userAgentProviderMock = mock.method(
+      userAgentProvider,
+      'getUserAgent',
+      () => 'testUserAgent'
+    );
+    const graphqlRequestExecutor = new GraphqlRequestExecutor(
+      '',
+      '',
+      userAgentProvider
+    );
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -252,6 +323,7 @@ void describe('Conversation turn response sender', () => {
     );
     const sender = new ConversationTurnResponseSender(
       event,
+      userAgentProvider,
       graphqlRequestExecutor
     );
     const errors: Array<ConversationTurnError> = [
@@ -266,6 +338,14 @@ void describe('Conversation turn response sender', () => {
     ];
     await sender.sendErrors(errors);
 
+    assert.strictEqual(userAgentProviderMock.mock.calls.length, 1);
+    assert.deepStrictEqual(userAgentProviderMock.mock.calls[0].arguments[0], {
+      'turn-response-type': 'error',
+    });
+    assert.strictEqual(executeGraphqlMock.mock.calls.length, 1);
+    assert.deepStrictEqual(executeGraphqlMock.mock.calls[0].arguments[1], {
+      userAgent: 'testUserAgent',
+    });
     assert.strictEqual(executeGraphqlMock.mock.calls.length, 1);
     const request = executeGraphqlMock.mock.calls[0]
       .arguments[0] as GraphqlRequest<MutationResponseInput>;

packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.ts

@@ -5,6 +5,7 @@ import {
 } from './types.js';
 import type { ContentBlock } from '@aws-sdk/client-bedrock-runtime';
 import { GraphqlRequestExecutor } from './graphql_request_executor';
+import { UserAgentProvider } from './user_agent_provider';
 
 export type MutationResponseInput = {
   input: {
@@ -36,10 +37,11 @@ export class ConversationTurnResponseSender {
    */
   constructor(
     private readonly event: ConversationTurnEvent,
+    private readonly userAgentProvider = new UserAgentProvider(event),
     private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
       event.graphqlApiEndpoint,
       event.request.headers.authorization,
-      event.request.headers['x-amz-user-agent']
+      userAgentProvider
     ),
     private readonly logger = console
   ) {}
@@ -50,7 +52,11 @@ export class ConversationTurnResponseSender {
     await this.graphqlRequestExecutor.executeGraphql<
       MutationResponseInput,
       void
-    >(responseMutationRequest);
+    >(responseMutationRequest, {
+      userAgent: this.userAgentProvider.getUserAgent({
+        'turn-response-type': 'single',
+      }),
+    });
   };
 
   sendResponseChunk = async (chunk: StreamingResponseChunk) => {
@@ -59,7 +65,11 @@ export class ConversationTurnResponseSender {
     await this.graphqlRequestExecutor.executeGraphql<
       MutationStreamingResponseInput,
       void
-    >(responseMutationRequest);
+    >(responseMutationRequest, {
+      userAgent: this.userAgentProvider.getUserAgent({
+        'turn-response-type': 'streaming',
+      }),
+    });
   };
 
   sendErrors = async (errors: ConversationTurnError[]) => {
@@ -71,7 +81,11 @@ export class ConversationTurnResponseSender {
     await this.graphqlRequestExecutor.executeGraphql<
       MutationErrorsResponseInput,
       void
-    >(responseMutationRequest);
+    >(responseMutationRequest, {
+      userAgent: this.userAgentProvider.getUserAgent({
+        'turn-response-type': 'error',
+      }),
+    });
   };
 
   private createMutationErrorsRequest = (errors: ConversationTurnError[]) => {

packages/ai-constructs/src/conversation/runtime/event-tools-provider/event_tools_provider.ts

@@ -1,6 +1,7 @@
 import { ConversationTurnEvent, ExecutableTool } from '../types';
 import { GraphQlTool } from './graphql_tool';
 import { GraphQlQueryFactory } from './graphql_query_factory';
+import { UserAgentProvider } from '../user_agent_provider';
 
 /**
  * Creates executable tools from definitions in conversation turn event.
@@ -29,7 +30,7 @@ export class ConversationTurnEventToolsProvider {
         graphqlApiEndpoint,
         query,
         this.event.request.headers.authorization,
-        this.event.request.headers['x-amz-user-agent']
+        new UserAgentProvider(this.event)
       );
     });
     return tools ?? [];