chore: added bedrock and vertex providers with google fix

Portkey-AI · Sep 17, 2024 · 95fedc7 · 95fedc7
1 parent 9bf1d21
commit 95fedc7
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/src/providers/bedrock/chatComplete.ts b/src/providers/bedrock/chatComplete.ts
@@ -255,6 +255,10 @@ export const BedrockAnthropicChatCompleteConfig: ProviderConfig = {
     param: 'max_tokens',
     required: true,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: true,
+  },
   temperature: {
     param: 'temperature',
     default: 1,
@@ -318,6 +322,11 @@ export const BedrockCohereChatCompleteConfig: ProviderConfig = {
     default: 20,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 20,
+    min: 1,
+  },
   temperature: {
     param: 'temperature',
     default: 0.75,
@@ -394,6 +403,12 @@ export const BedrockLLamaChatCompleteConfig: ProviderConfig = {
     min: 1,
     max: 2048,
   },
+  max_completion_tokens: {
+    param: 'max_gen_len',
+    default: 512,
+    min: 1,
+    max: 2048,
+  },
   temperature: {
     param: 'temperature',
     default: 0.5,
@@ -475,6 +490,9 @@ const transformTitanGenerationConfig = (params: Params) => {
   if (params['max_tokens']) {
     generationConfig['maxTokenCount'] = params['max_tokens'];
   }
+  if (params['max_completion_tokens']) {
+    generationConfig['maxTokenCount'] = params['max_completion_tokens'];
+  }
   if (params['stop']) {
     generationConfig['stopSequences'] = params['stop'];
   }
@@ -513,6 +531,10 @@ export const BedrockTitanChatompleteConfig: ProviderConfig = {
     param: 'textGenerationConfig',
     transform: (params: Params) => transformTitanGenerationConfig(params),
   },
+  max_completion_tokens: {
+    param: 'textGenerationConfig',
+    transform: (params: Params) => transformTitanGenerationConfig(params),
+  },
   top_p: {
     param: 'textGenerationConfig',
     transform: (params: Params) => transformTitanGenerationConfig(params),
@@ -547,6 +569,10 @@ export const BedrockAI21ChatCompleteConfig: ProviderConfig = {
     param: 'maxTokens',
     default: 200,
   },
+  max_completion_tokens: {
+    param: 'maxTokens',
+    default: 200,
+  },
   temperature: {
     param: 'temperature',
     default: 0.7,

diff --git a/src/providers/google/chatComplete.ts b/src/providers/google/chatComplete.ts
@@ -32,6 +32,9 @@ const transformGenerationConfig = (params: Params) => {
   if (params['max_tokens']) {
     generationConfig['maxOutputTokens'] = params['max_tokens'];
   }
+  if (params['max_completion_tokens']) {
+    generationConfig['maxOutputTokens'] = params['max_completion_tokens'];
+  }
   if (params['stop']) {
     generationConfig['stopSequences'] = params['stop'];
   }