Merge pull request #617 from shivam-pareek/fix/max-completion-tokens

fix: added max_completion_tokens for chat complete
Portkey-AI · Sep 27, 2024 · a65600e · a65600e
2 parents f3ed2b1 + 3b1c139
commit a65600e
Show file tree

Hide file tree

Showing 30 changed files with 178 additions and 0 deletions.
diff --git a/src/providers/ai21/chatComplete.ts b/src/providers/ai21/chatComplete.ts
@@ -49,6 +49,10 @@ export const AI21ChatCompleteConfig: ProviderConfig = {
     param: 'maxTokens',
     default: 16,
   },
+  max_completion_tokens: {
+    param: 'maxTokens',
+    default: 16,
+  },
   minTokens: {
     param: 'minTokens',
     default: 0,

diff --git a/src/providers/anthropic/chatComplete.ts b/src/providers/anthropic/chatComplete.ts
@@ -260,6 +260,10 @@ export const AnthropicChatCompleteConfig: ProviderConfig = {
     param: 'max_tokens',
     required: true,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: true,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/anyscale/chatComplete.ts b/src/providers/anyscale/chatComplete.ts
@@ -32,6 +32,11 @@ export const AnyscaleChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/azure-ai-inference/chatComplete.ts b/src/providers/azure-ai-inference/chatComplete.ts
@@ -20,6 +20,11 @@ export const AzureAIInferenceChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/azure-openai/chatComplete.ts b/src/providers/azure-openai/chatComplete.ts
@@ -27,6 +27,11 @@ export const AzureOpenAIChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_completion_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

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,
@@ -437,6 +452,11 @@ export const BedrockMistralChatCompleteConfig: ProviderConfig = {
     default: 20,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 20,
+    min: 1,
+  },
   temperature: {
     param: 'temperature',
     default: 0.75,
@@ -470,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'];
   }
@@ -508,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),
@@ -542,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/cohere/chatComplete.ts b/src/providers/cohere/chatComplete.ts
@@ -72,6 +72,11 @@ export const CohereChatCompleteConfig: ProviderConfig = {
     default: 20,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 20,
+    min: 1,
+  },
   temperature: {
     param: 'temperature',
     default: 0.75,

diff --git a/src/providers/deepbricks/chatComplete.ts b/src/providers/deepbricks/chatComplete.ts
@@ -29,6 +29,11 @@ export const DeepbricksChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/deepinfra/chatComplete.ts b/src/providers/deepinfra/chatComplete.ts
@@ -34,6 +34,11 @@ export const DeepInfraChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 1,
+  },
   n: {
     param: 'n',
     default: 1,

diff --git a/src/providers/deepseek/chatComplete.ts b/src/providers/deepseek/chatComplete.ts
@@ -25,6 +25,11 @@ export const DeepSeekChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/fireworks-ai/chatComplete.ts b/src/providers/fireworks-ai/chatComplete.ts
@@ -27,6 +27,11 @@ export const FireworksAIChatCompleteConfig: ProviderConfig = {
     default: 200,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 200,
+    min: 1,
+  },
   prompt_truncate_len: {
     param: 'prompt_truncate_len',
     default: 1500,

diff --git a/src/providers/google-vertex-ai/chatComplete.ts b/src/providers/google-vertex-ai/chatComplete.ts
@@ -213,6 +213,10 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
     param: 'generationConfig',
     transform: (params: Params) => transformGenerationConfig(params),
   },
+  max_completion_tokens: {
+    param: 'generationConfig',
+    transform: (params: Params) => transformGenerationConfig(params),
+  },
   stop: {
     param: 'generationConfig',
     transform: (params: Params) => transformGenerationConfig(params),
@@ -509,6 +513,10 @@ export const VertexAnthropicChatCompleteConfig: ProviderConfig = {
     param: 'max_tokens',
     required: true,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: true,
+  },
   temperature: {
     param: 'temperature',
     default: 1,
@@ -663,6 +671,12 @@ export const VertexLlamaChatCompleteConfig: ProviderConfig = {
     min: 1,
     max: 2048,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 512,
+    min: 1,
+    max: 2048,
+  },
   temperature: {
     param: 'temperature',
     default: 0.5,

diff --git a/src/providers/google-vertex-ai/transformGenerationConfig.ts b/src/providers/google-vertex-ai/transformGenerationConfig.ts
@@ -17,6 +17,9 @@ export function 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'];
   }

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'];
   }
@@ -287,6 +290,10 @@ export const GoogleChatCompleteConfig: ProviderConfig = {
     param: 'generationConfig',
     transform: (params: Params) => transformGenerationConfig(params),
   },
+  max_completion_tokens: {
+    param: 'generationConfig',
+    transform: (params: Params) => transformGenerationConfig(params),
+  },
   stop: {
     param: 'generationConfig',
     transform: (params: Params) => transformGenerationConfig(params),

diff --git a/src/providers/groq/chatComplete.ts b/src/providers/groq/chatComplete.ts
@@ -24,6 +24,11 @@ export const GroqChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/huggingface/chatComplete.ts b/src/providers/huggingface/chatComplete.ts
@@ -30,6 +30,11 @@ export const HuggingfaceChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/lingyi/chatComplete.ts b/src/providers/lingyi/chatComplete.ts
@@ -24,6 +24,11 @@ export const LingyiChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/mistral-ai/chatComplete.ts b/src/providers/mistral-ai/chatComplete.ts
@@ -36,6 +36,11 @@ export const MistralAIChatCompleteConfig: ProviderConfig = {
     default: null,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: null,
+    min: 1,
+  },
   stream: {
     param: 'stream',
     default: false,

diff --git a/src/providers/moonshot/chatComplete.ts b/src/providers/moonshot/chatComplete.ts
@@ -25,6 +25,11 @@ export const MoonshotChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/novita-ai/chatComplete.ts b/src/providers/novita-ai/chatComplete.ts
@@ -28,6 +28,12 @@ export const NovitaAIChatCompleteConfig: ProviderConfig = {
     default: 128,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: true,
+    default: 128,
+    min: 1,
+  },
   stop: {
     param: 'stop',
   },

diff --git a/src/providers/ollama/chatComplete.ts b/src/providers/ollama/chatComplete.ts
@@ -56,6 +56,11 @@ export const OllamaChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
 };
 
 export interface OllamaChatCompleteResponse

diff --git a/src/providers/openrouter/chatComplete.ts b/src/providers/openrouter/chatComplete.ts
@@ -24,6 +24,11 @@ export const OpenrouterChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     default: 1,

diff --git a/src/providers/palm/chatComplete.ts b/src/providers/palm/chatComplete.ts
@@ -66,6 +66,11 @@ export const PalmChatCompleteConfig: ProviderConfig = {
     default: 100,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'maxOutputTokens',
+    default: 100,
+    min: 1,
+  },
   stop: {
     param: 'stopSequences',
   },

diff --git a/src/providers/perplexity-ai/chatComplete.ts b/src/providers/perplexity-ai/chatComplete.ts
@@ -27,6 +27,11 @@ export const PerplexityAIChatCompleteConfig: ProviderConfig = {
     required: true,
     min: 1,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: true,
+    min: 1,
+  },
   temperature: {
     param: 'temperature',
     min: 0,

diff --git a/src/providers/predibase/chatComplete.ts b/src/providers/predibase/chatComplete.ts
@@ -35,6 +35,12 @@ export const PredibaseChatCompleteConfig: ProviderConfig = {
     default: 4096,
     min: 0,
   },
+  max_completion_tokens: {
+    param: 'max_tokens',
+    required: false,
+    default: 4096,
+    min: 0,
+  },
   temperature: {
     param: 'temperature',
     required: false,