fix(ai): forward user-agent with package metadata in generation and c…

…onversation requests (#3029)
aws-amplify · Nov 18, 2024 · 9a0ac8b · 9a0ac8b
1 parent 1dd63da
commit 9a0ac8b
Show file tree

Hide file tree

Showing 6 changed files with 196 additions and 15 deletions.
diff --git a/...sformer/src/__tests__/__snapshots__/amplify-graphql-conversation-transformer.test.ts.snap b/...sformer/src/__tests__/__snapshots__/amplify-graphql-conversation-transformer.test.ts.snap
@@ -667,6 +667,7 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
 export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
+  const userAgent = createUserAgent(request);
 
   const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message }';
 
@@ -705,7 +706,12 @@ export function request(ctx) {
     responseMutation: streamingResponseMutation,
     graphqlApiEndpoint,
     modelConfiguration,
-    request: { headers: { authorization: authHeader } },
+    request: {
+      headers: {
+        authorization: authHeader,
+        'x-amz-user-agent': userAgent,
+      }
+    },
     messageHistoryQuery,
     toolsConfiguration,
     streamResponse: true,
@@ -735,7 +741,17 @@ export function response(ctx) {
   };
   return response;
 }
-"
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-conversation-transformer#0.7.1';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}"
 `;
 
 exports[`ConversationTransformer valid schemas should transform conversation route with inference configuration: SendMessageMutation resolver code 1`] = `
@@ -1508,6 +1524,7 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
 export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
+  const userAgent = createUserAgent(request);
 
   const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message }';
 
@@ -1546,7 +1563,12 @@ export function request(ctx) {
     responseMutation: streamingResponseMutation,
     graphqlApiEndpoint,
     modelConfiguration,
-    request: { headers: { authorization: authHeader } },
+    request: {
+      headers: {
+        authorization: authHeader,
+        'x-amz-user-agent': userAgent,
+      }
+    },
     messageHistoryQuery,
     toolsConfiguration,
     streamResponse: true,
@@ -1576,7 +1598,17 @@ export function response(ctx) {
   };
   return response;
 }
-"
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-conversation-transformer#0.7.1';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}"
 `;
 
 exports[`ConversationTransformer valid schemas should transform conversation route with model query tool including relationships: SendMessageMutation resolver code 1`] = `
@@ -2349,6 +2381,7 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
 export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
+  const userAgent = createUserAgent(request);
 
   const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message }';
 
@@ -2387,7 +2420,12 @@ export function request(ctx) {
     responseMutation: streamingResponseMutation,
     graphqlApiEndpoint,
     modelConfiguration,
-    request: { headers: { authorization: authHeader } },
+    request: {
+      headers: {
+        authorization: authHeader,
+        'x-amz-user-agent': userAgent,
+      }
+    },
     messageHistoryQuery,
     toolsConfiguration,
     streamResponse: true,
@@ -2417,7 +2455,17 @@ export function response(ctx) {
   };
   return response;
 }
-"
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-conversation-transformer#0.7.1';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}"
 `;
 
 exports[`ConversationTransformer valid schemas should transform conversation route with model query tool: SendMessageMutation resolver code 1`] = `
@@ -3190,6 +3238,7 @@ exports[`ConversationTransformer valid schemas should transform conversation rou
 export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
+  const userAgent = createUserAgent(request);
 
   const selectionSet = 'associatedUserMessageId contentBlockDeltaIndex contentBlockDoneAtIndex contentBlockIndex contentBlockText contentBlockToolUse { toolUseId name input } conversationId id stopReason owner errors { errorType message }';
 
@@ -3228,7 +3277,12 @@ export function request(ctx) {
     responseMutation: streamingResponseMutation,
     graphqlApiEndpoint,
     modelConfiguration,
-    request: { headers: { authorization: authHeader } },
+    request: {
+      headers: {
+        authorization: authHeader,
+        'x-amz-user-agent': userAgent,
+      }
+    },
     messageHistoryQuery,
     toolsConfiguration,
     streamResponse: true,
@@ -3258,7 +3312,17 @@ export function response(ctx) {
   };
   return response;
 }
-"
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-conversation-transformer#0.7.1';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}"
 `;
 
 exports[`ConversationTransformer valid schemas should transform conversation route with query tools: SendMessageMutation resolver code 1`] = `

diff --git a/...mplify-graphql-conversation-transformer/src/resolvers/send-message-pipeline-definition.ts b/...mplify-graphql-conversation-transformer/src/resolvers/send-message-pipeline-definition.ts
@@ -93,6 +93,9 @@ function invokeLambda(): ResolverFunctionDefinition {
   });
 }
 
+const packageName = 'amplify-graphql-conversation-transformer';
+const packageVersion = require('../../package.json').version;
+
 /**
  * The substitutions for the invoke lambda resolver function.
  */
@@ -113,6 +116,7 @@ function invokeLambdaResolverSubstitutions(config: ConversationDirectiveConfigur
     LIST_QUERY_LIMIT: 'undefined',
     STREAMING_RESPONSE_MUTATION_NAME: config.assistantResponseStreamingMutation.field.name.value,
     STREAMING_RESPONSE_MUTATION_INPUT_TYPE_NAME: config.assistantResponseStreamingMutation.input.name.value,
+    PACKAGE_METADATA: `'${packageName}#${packageVersion}'`,
   };
 }
 

diff --git a/...ql-conversation-transformer/src/resolvers/templates/invoke-lambda-resolver-fn.template.js b/...ql-conversation-transformer/src/resolvers/templates/invoke-lambda-resolver-fn.template.js
@@ -3,6 +3,7 @@ import { util } from '@aws-appsync/utils';
 export function request(ctx) {
   const { args, request } = ctx;
   const { graphqlApiEndpoint } = ctx.stash;
+  const userAgent = createUserAgent(request);
 
   const selectionSet = '[[SELECTION_SET]]';
 
@@ -41,7 +42,12 @@ export function request(ctx) {
     responseMutation: streamingResponseMutation,
     graphqlApiEndpoint,
     modelConfiguration,
-    request: { headers: { authorization: authHeader } },
+    request: {
+      headers: {
+        authorization: authHeader,
+        'x-amz-user-agent': userAgent,
+      }
+    },
     messageHistoryQuery,
     toolsConfiguration,
     streamResponse: true,
@@ -71,3 +77,14 @@ export function response(ctx) {
   };
   return response;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = [[PACKAGE_METADATA]];
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = `${userAgent} md/${packageMetadata}`;
+  } else {
+    userAgent = `lib/${packageMetadata}`;
+  }
+  return userAgent;
+}
diff --git a/...ansformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap b/...ansformer/src/__tests__/__snapshots__/amplify-graphql-generation-transformer.test.ts.snap
@@ -64,12 +64,16 @@ export function request(ctx) {
   const prompt = "";
   const args = JSON.stringify(ctx.args);
   const inferenceConfig = undefined;
+  const userAgent = createUserAgent(ctx.request);
 
   return {
     resourcePath: '/model/anthropic.claude-3-haiku-20240307-v1:0/converse',
     method: 'POST',
     params: {
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        'x-amz-user-agent': userAgent,
+      },
       body: {
         messages: [
           {
@@ -121,6 +125,17 @@ export function response(ctx) {
   }
   return value;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-generation-transformer#0.2.7';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}
 ",
 }
 `;
@@ -189,12 +204,16 @@ export function request(ctx) {
   const prompt = "You are a helpful assistant that generates recipes.";
   const args = JSON.stringify(ctx.args);
   const inferenceConfig = undefined;
+  const userAgent = createUserAgent(ctx.request);
 
   return {
     resourcePath: '/model/anthropic.claude-3-haiku-20240307-v1:0/converse',
     method: 'POST',
     params: {
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        'x-amz-user-agent': userAgent,
+      },
       body: {
         messages: [
           {
@@ -246,6 +265,17 @@ export function response(ctx) {
   }
   return value;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-generation-transformer#0.2.7';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}
 ",
 }
 `;
@@ -313,12 +343,16 @@ export function request(ctx) {
   const prompt = "Make a string based on the description.";
   const args = JSON.stringify(ctx.args);
   const inferenceConfig = undefined;
+  const userAgent = createUserAgent(ctx.request);
 
   return {
     resourcePath: '/model/anthropic.claude-3-haiku-20240307-v1:0/converse',
     method: 'POST',
     params: {
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        'x-amz-user-agent': userAgent,
+      },
       body: {
         messages: [
           {
@@ -370,6 +404,17 @@ export function response(ctx) {
   }
   return value;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-generation-transformer#0.2.7';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}
 "
 `;
 
@@ -454,12 +499,16 @@ export function request(ctx) {
   const prompt = "Make a string based on the description.";
   const args = JSON.stringify(ctx.args);
   const inferenceConfig = undefined;
+  const userAgent = createUserAgent(ctx.request);
 
   return {
     resourcePath: '/model/anthropic.claude-3-haiku-20240307-v1:0/converse',
     method: 'POST',
     params: {
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        'x-amz-user-agent': userAgent,
+      },
       body: {
         messages: [
           {
@@ -507,6 +556,17 @@ export function response(ctx) {
   
   return value;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-generation-transformer#0.2.7';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}
 ",
 }
 `;
@@ -520,12 +580,16 @@ export function request(ctx) {
   const prompt = "Generate a string based on the description.";
   const args = JSON.stringify(ctx.args);
   const inferenceConfig = { inferenceConfig: {"maxTokens":100,"temperature":0.7,"topP":0.9} };
+  const userAgent = createUserAgent(ctx.request);
 
   return {
     resourcePath: '/model/anthropic.claude-3-haiku-20240307-v1:0/converse',
     method: 'POST',
     params: {
-      headers: { 'Content-Type': 'application/json' },
+      headers: {
+        'Content-Type': 'application/json',
+        'x-amz-user-agent': userAgent,
+      },
       body: {
         messages: [
           {
@@ -573,6 +637,17 @@ export function response(ctx) {
   
   return value;
 }
+
+function createUserAgent(request) {
+  const packageMetadata = 'amplify-graphql-generation-transformer#0.2.7';
+  let userAgent = request.headers['x-amz-user-agent'];
+  if (userAgent) {
+    userAgent = \`\${userAgent} md/\${packageMetadata}\`;
+  } else {
+    userAgent = \`lib/\${packageMetadata}\`;
+  }
+  return userAgent;
+}
 ",
 }
 `;