Restrict known additional model fields to the specific providers

Author: narengogi

src/providers/bedrock/chatComplete.ts

@@ -17,6 +17,9 @@ import {
 import { BedrockErrorResponse } from './embed';
 import {
   transformAdditionalModelRequestFields,
+  transformAI21AdditionalModelRequestFields,
+  transformAnthropicAdditionalModelRequestFields,
+  transformCohereAdditionalModelRequestFields,
   transformInferenceConfig,
 } from './utils';
 
@@ -32,6 +35,29 @@ export interface BedrockChatCompletionsParams extends Params {
   countPenalty?: number;
 }
 
+export interface BedrockConverseAnthropicChatCompletionsParams
+  extends BedrockChatCompletionsParams {
+  anthropic_version?: string;
+  user?: string;
+}
+
+export interface BedrockConverseCohereChatCompletionsParams
+  extends BedrockChatCompletionsParams {
+  frequency_penalty?: number;
+  presence_penalty?: number;
+  logit_bias?: Record<string, number>;
+  n?: number;
+}
+
+export interface BedrockConverseAI21ChatCompletionsParams
+  extends BedrockChatCompletionsParams {
+  frequency_penalty?: number;
+  presence_penalty?: number;
+  frequencyPenalty?: number;
+  presencePenalty?: number;
+  countPenalty?: number;
+}
+
 const getMessageTextContentArray = (message: Message): { text: string }[] => {
   if (message.content && typeof message.content === 'object') {
     return message.content
@@ -232,41 +258,6 @@ export const BedrockConverseChatCompleteConfig: ProviderConfig = {
     transform: (params: BedrockChatCompletionsParams) =>
       transformAdditionalModelRequestFields(params),
   },
-  anthropic_version: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  frequency_penalty: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  presence_penalty: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  logit_bias: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  n: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  stream: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
-  countPenalty: {
-    param: 'additionalModelRequestFields',
-    transform: (params: BedrockChatCompletionsParams) =>
-      transformAdditionalModelRequestFields(params),
-  },
 };
 
 interface BedrockChatCompletionResponse {
@@ -460,6 +451,108 @@ export const BedrockChatCompleteStreamChunkTransform: (
   })}\n\n`;
 };
 
+export const BedrockConverseAnthropicChatCompleteConfig: ProviderConfig = {
+  ...BedrockConverseChatCompleteConfig,
+  additionalModelRequestFields: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAnthropicAdditionalModelRequestFields(params),
+  },
+  top_k: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAnthropicAdditionalModelRequestFields(params),
+  },
+  anthropic_version: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAnthropicAdditionalModelRequestFields(params),
+  },
+  user: {
+    param: 'user',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAnthropicAdditionalModelRequestFields(params),
+  },
+};
+
+export const BedrockConverseCohereChatCompleteConfig: ProviderConfig = {
+  ...BedrockConverseChatCompleteConfig,
+  additionalModelRequestFields: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  top_k: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  frequency_penalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  presence_penalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  logit_bias: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  n: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+  stream: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformCohereAdditionalModelRequestFields(params),
+  },
+};
+
+export const BedrockConverseAI21ChatCompleteConfig: ProviderConfig = {
+  ...BedrockConverseChatCompleteConfig,
+  additionalModelRequestFields: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  top_k: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  frequency_penalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  presence_penalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  frequencyPenalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  presencePenalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+  countPenalty: {
+    param: 'additionalModelRequestFields',
+    transform: (params: BedrockChatCompletionsParams) =>
+      transformAI21AdditionalModelRequestFields(params),
+  },
+};
+
 export const BedrockCohereChatCompleteConfig: ProviderConfig = {
   messages: {
     param: 'prompt',

src/providers/bedrock/index.ts

@@ -12,6 +12,9 @@ import {
   BedrockCohereChatCompleteResponseTransform,
   BedrockAI21ChatCompleteConfig,
   BedrockAI21ChatCompleteResponseTransform,
+  BedrockConverseAnthropicChatCompleteConfig,
+  BedrockConverseCohereChatCompleteConfig,
+  BedrockConverseAI21ChatCompleteConfig,
 } from './chatComplete';
 import {
   BedrockAI21CompleteConfig,
@@ -64,6 +67,7 @@ const BedrockConfig: ProviderConfigs = {
       case ANTHROPIC:
         config = {
           complete: BedrockAnthropicCompleteConfig,
+          chatComplete: BedrockConverseAnthropicChatCompleteConfig,
           api: BedrockAPIConfig,
           responseTransforms: {
             'stream-complete': BedrockAnthropicCompleteStreamChunkTransform,
@@ -74,6 +78,7 @@ const BedrockConfig: ProviderConfigs = {
       case COHERE:
         config = {
           complete: BedrockCohereCompleteConfig,
+          chatComplete: BedrockConverseCohereChatCompleteConfig,
           embed: BedrockCohereEmbedConfig,
           api: BedrockAPIConfig,
           responseTransforms: {
@@ -126,6 +131,7 @@ const BedrockConfig: ProviderConfigs = {
         config = {
           complete: BedrockAI21CompleteConfig,
           api: BedrockAPIConfig,
+          chatComplete: BedrockConverseAI21ChatCompleteConfig,
           responseTransforms: {
             complete: BedrockAI21CompleteResponseTransform,
           },

src/providers/bedrock/utils.ts

@@ -1,6 +1,11 @@
 import { SignatureV4 } from '@smithy/signature-v4';
 import { Sha256 } from '@aws-crypto/sha256-js';
-import { BedrockChatCompletionsParams } from './chatComplete';
+import {
+  BedrockConverseAI21ChatCompletionsParams,
+  BedrockConverseAnthropicChatCompletionsParams,
+  BedrockChatCompletionsParams,
+  BedrockConverseCohereChatCompletionsParams,
+} from './chatComplete';
 
 export const generateAWSHeaders = async (
   body: Record<string, any>,
@@ -66,26 +71,83 @@ export const transformAdditionalModelRequestFields = (
   const additionalModelRequestFields: Record<string, any> =
     params.additionalModelRequestFields || {};
   if (params['top_k']) {
-    additionalModelRequestFields['topK'] = params['top_k'];
+    additionalModelRequestFields['top_k'] = params['top_k'];
+  }
+  return additionalModelRequestFields;
+};
+
+export const transformAnthropicAdditionalModelRequestFields = (
+  params: BedrockConverseAnthropicChatCompletionsParams
+) => {
+  const additionalModelRequestFields: Record<string, any> =
+    params.additionalModelRequestFields || {};
+  if (params['top_k']) {
+    additionalModelRequestFields['top_k'] = params['top_k'];
   }
-  // Backward compatibility
   if (params['anthropic_version']) {
     additionalModelRequestFields['anthropic_version'] =
       params['anthropic_version'];
   }
+  if (params['user']) {
+    additionalModelRequestFields['metadata'] = {
+      user_id: params['user'],
+    };
+  }
+  return additionalModelRequestFields;
+};
+
+export const transformCohereAdditionalModelRequestFields = (
+  params: BedrockConverseCohereChatCompletionsParams
+) => {
+  const additionalModelRequestFields: Record<string, any> =
+    params.additionalModelRequestFields || {};
+  if (params['top_k']) {
+    additionalModelRequestFields['top_k'] = params['top_k'];
+  }
+  if (params['n']) {
+    additionalModelRequestFields['n'] = params['n'];
+  }
   if (params['frequency_penalty']) {
-    additionalModelRequestFields['frequencyPenalty'] =
+    additionalModelRequestFields['frequency_penalty'] =
       params['frequency_penalty'];
   }
   if (params['presence_penalty']) {
-    additionalModelRequestFields['presencePenalty'] =
+    additionalModelRequestFields['presence_penalty'] =
       params['presence_penalty'];
   }
   if (params['logit_bias']) {
     additionalModelRequestFields['logitBias'] = params['logit_bias'];
   }
-  if (params['n']) {
-    additionalModelRequestFields['n'] = params['n'];
+  if (params['stream']) {
+    additionalModelRequestFields['stream'] = params['stream'];
+  }
+  return additionalModelRequestFields;
+};
+
+export const transformAI21AdditionalModelRequestFields = (
+  params: BedrockConverseAI21ChatCompletionsParams
+) => {
+  const additionalModelRequestFields: Record<string, any> =
+    params.additionalModelRequestFields || {};
+  if (params['top_k']) {
+    additionalModelRequestFields['top_k'] = params['top_k'];
+  }
+  if (params['frequency_penalty']) {
+    additionalModelRequestFields['frequencyPenalty'] = {
+      scale: params['frequency_penalty'],
+    };
+  }
+  if (params['presence_penalty']) {
+    additionalModelRequestFields['presencePenalty'] = {
+      scale: params['presence_penalty'],
+    };
+  }
+  if (params['frequencyPenalty']) {
+    additionalModelRequestFields['frequencyPenalty'] =
+      params['frequencyPenalty'];
+  }
+  if (params['presencePenalty']) {
+    additionalModelRequestFields['presencePenalty'] = params['presencePenalty'];
   }
   if (params['countPenalty']) {
     additionalModelRequestFields['countPenalty'] = params['countPenalty'];