Pass user agent in conversation handler lambda (#2086)

sobolk · web-flow · commit 3a29d43fee67 · 2024-10-07T15:49:29.000-07:00
* Pass user agent in conversation handler lambda

* Pass user agent in conversation handler lambda
diff --git a/.changeset/soft-gifts-exist.md b/.changeset/soft-gifts-exist.md
@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': minor
+---
+
+Pass user agent in conversation handler lambda
diff --git a/packages/ai-constructs/API.md b/packages/ai-constructs/API.md
@@ -92,9 +92,7 @@ type ConversationTurnEvent = {
         };
     };
     request: {
-        headers: {
-            authorization: string;
-        };
+        headers: Record<string, string>;
     };
     messages?: Array<ConversationMessage>;
     messageHistoryQuery: {
diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts
@@ -840,4 +840,43 @@ void describe('Bedrock converse adapter', () => {
       },
     ]);
   });
+
+  void it('adds user agent middleware', async () => {
+    const event: ConversationTurnEvent = {
+      ...commonEvent,
+    };
+
+    event.request.headers['x-amz-user-agent'] = 'testUserAgent';
+
+    const bedrockClient = new BedrockRuntimeClient();
+    const addMiddlewareMock = mock.method(bedrockClient.middlewareStack, 'add');
+
+    new BedrockConverseAdapter(
+      event,
+      [],
+      bedrockClient,
+      undefined,
+      messageHistoryRetriever
+    );
+
+    assert.strictEqual(addMiddlewareMock.mock.calls.length, 1);
+    const middlewareHandler = addMiddlewareMock.mock.calls[0].arguments[0];
+    const options = addMiddlewareMock.mock.calls[0].arguments[1];
+    assert.strictEqual(options.name, 'amplify-user-agent-injector');
+    const args: {
+      request: {
+        headers: Record<string, string>;
+      };
+    } = {
+      request: {
+        headers: {},
+      },
+    };
+    // @ts-expect-error We mock subset of middleware inputs here.
+    await middlewareHandler(mock.fn(), {})(args);
+    assert.strictEqual(
+      args.request.headers['x-amz-user-agent'],
+      'testUserAgent'
+    );
+  });
 });
diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts
@@ -44,6 +44,21 @@ export class BedrockConverseAdapter {
     ),
     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.executableTools = [
       ...eventToolsProvider.getEventTools(),
       ...additionalTools,
diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.test.ts b/packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.test.ts
@@ -346,7 +346,7 @@ void describe('Conversation message history retriever', () => {
 
   for (const testCase of testCases) {
     void it(testCase.name, async () => {
-      const graphqlRequestExecutor = new GraphqlRequestExecutor('', '');
+      const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
       const executeGraphqlMock = mock.method(
         graphqlRequestExecutor,
         'executeGraphql',
diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.ts b/packages/ai-constructs/src/conversation/runtime/conversation_message_history_retriever.ts
@@ -102,7 +102,8 @@ export class ConversationMessageHistoryRetriever {
     private readonly event: ConversationTurnEvent,
     private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
       event.graphqlApiEndpoint,
-      event.request.headers.authorization
+      event.request.headers.authorization,
+      event.request.headers['x-amz-user-agent']
     )
   ) {}
 
diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.test.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.test.ts
@@ -33,7 +33,7 @@ void describe('Conversation turn response sender', () => {
   };
 
   void it('sends response back to appsync', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '');
+    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
@@ -81,7 +81,7 @@ void describe('Conversation turn response sender', () => {
   });
 
   void it('serializes tool use input to JSON', async () => {
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '');
+    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.ts
@@ -22,7 +22,8 @@ export class ConversationTurnResponseSender {
     private readonly event: ConversationTurnEvent,
     private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
       event.graphqlApiEndpoint,
-      event.request.headers.authorization
+      event.request.headers.authorization,
+      event.request.headers['x-amz-user-agent']
     ),
     private readonly logger = console
   ) {}
diff --git a/packages/ai-constructs/src/conversation/runtime/event-tools-provider/event_tools_provider.ts b/packages/ai-constructs/src/conversation/runtime/event-tools-provider/event_tools_provider.ts
@@ -28,7 +28,8 @@ export class ConversationTurnEventToolsProvider {
         inputSchema,
         graphqlApiEndpoint,
         query,
-        this.event.request.headers.authorization
+        this.event.request.headers.authorization,
+        this.event.request.headers['x-amz-user-agent']
       );
     });
     return tools ?? [];
diff --git a/packages/ai-constructs/src/conversation/runtime/event-tools-provider/graphql_tool.test.ts b/packages/ai-constructs/src/conversation/runtime/event-tools-provider/graphql_tool.test.ts
@@ -22,6 +22,7 @@ void describe('GraphQl tool', () => {
       graphQlEndpoint,
       query,
       accessToken,
+      '',
       graphqlRequestExecutor
     );
   };
@@ -30,7 +31,7 @@ void describe('GraphQl tool', () => {
     const testResponse = {
       test: 'response',
     };
-    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '');
+    const graphqlRequestExecutor = new GraphqlRequestExecutor('', '', '');
     const executeGraphqlMock = mock.method(
       graphqlRequestExecutor,
       'executeGraphql',
diff --git a/packages/ai-constructs/src/conversation/runtime/event-tools-provider/graphql_tool.ts b/packages/ai-constructs/src/conversation/runtime/event-tools-provider/graphql_tool.ts
@@ -17,9 +17,11 @@ export class GraphQlTool implements ExecutableTool<JSONSchema, unknown> {
     readonly graphQlEndpoint: string,
     private readonly query: string,
     readonly accessToken: string,
+    readonly userAgent: string,
     private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
       graphQlEndpoint,
-      accessToken
+      accessToken,
+      userAgent
     )
   ) {}
 
diff --git a/packages/ai-constructs/src/conversation/runtime/graphql_request_executor.test.ts b/packages/ai-constructs/src/conversation/runtime/graphql_request_executor.test.ts
@@ -6,6 +6,7 @@ import { GraphqlRequestExecutor } from './graphql_request_executor';
 void describe('Graphql executor test', () => {
   const graphqlEndpoint = 'http://fake.endpoint/';
   const accessToken = 'testToken';
+  const userAgent = 'testUserAgent';
 
   void it('sends request to appsync', async () => {
     const fetchMock = mock.fn(
@@ -17,6 +18,7 @@ void describe('Graphql executor test', () => {
     const executor = new GraphqlRequestExecutor(
       graphqlEndpoint,
       accessToken,
+      userAgent,
       fetchMock
     );
     const query = 'testQuery';
@@ -37,6 +39,7 @@ void describe('Graphql executor test', () => {
       'application/graphql'
     );
     assert.strictEqual(request.headers.get('Authorization'), accessToken);
+    assert.strictEqual(request.headers.get('x-amz-user-agent'), userAgent);
     assert.ok(request.body);
     assert.deepStrictEqual(JSON.parse(await text(request.body)), {
       query: 'testQuery',
@@ -59,6 +62,7 @@ void describe('Graphql executor test', () => {
     const executor = new GraphqlRequestExecutor(
       graphqlEndpoint,
       accessToken,
+      userAgent,
       fetchMock
     );
     const query = 'testQuery';
@@ -98,6 +102,7 @@ void describe('Graphql executor test', () => {
     const executor = new GraphqlRequestExecutor(
       graphqlEndpoint,
       accessToken,
+      userAgent,
       fetchMock
     );
     const query = 'testQuery';
diff --git a/packages/ai-constructs/src/conversation/runtime/graphql_request_executor.ts b/packages/ai-constructs/src/conversation/runtime/graphql_request_executor.ts
@@ -15,6 +15,7 @@ export class GraphqlRequestExecutor {
   constructor(
     private readonly graphQlEndpoint: string,
     private readonly accessToken: string,
+    private readonly userAgent: string,
     private readonly _fetch = fetch
   ) {}
 
@@ -26,6 +27,7 @@ export class GraphqlRequestExecutor {
       headers: {
         'Content-Type': 'application/graphql',
         Authorization: this.accessToken,
+        'x-amz-user-agent': this.userAgent,
       },
       body: JSON.stringify({
         query: request.query,
diff --git a/packages/ai-constructs/src/conversation/runtime/types.ts b/packages/ai-constructs/src/conversation/runtime/types.ts
@@ -56,9 +56,7 @@ export type ConversationTurnEvent = {
     };
   };
   request: {
-    headers: {
-      authorization: string;
-    };
+    headers: Record<string, string>;
   };
   /**
    * @deprecated This field is going to be removed in upcoming releases.

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'@aws-amplify/ai-constructs': minor
 +---
++
 +Pass user agent in conversation handler lambda
Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,8 @@ export class ConversationMessageHistoryRetriever {`
`102`	`102`	`private readonly event: ConversationTurnEvent,`
`103`	`103`	`private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(`
`104`	`104`	`event.graphqlApiEndpoint,`
`105`		`- event.request.headers.authorization`
	`105`	`+ event.request.headers.authorization,`
	`106`	`+ event.request.headers['x-amz-user-agent']`
`106`	`107`	`)`
`107`	`108`	`) {}`
`108`	`109`