Merge pull request #467 from keshavkrishna/mistralai-fim-completions

Added Support for the new /fim/completions route on Mistral AI API
Portkey-AI · Sep 19, 2024 · 114895d · 114895d
2 parents 132d6ad + 3b107d0
commit 114895d
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 3 deletions.
diff --git a/src/handlers/handlerUtils.ts b/src/handlers/handlerUtils.ts
@@ -1120,6 +1120,8 @@ export function constructConfigFromRequestHeaders(
       anthropicConfig),
     ...(requestHeaders[`x-${POWERED_BY}-provider`] === HUGGING_FACE &&
       huggingfaceConfig),
+    mistralFimCompletion:
+      requestHeaders[`x-${POWERED_BY}-mistral-fim-completion`],
   };
 }
 

diff --git a/src/providers/mistral-ai/api.ts b/src/providers/mistral-ai/api.ts
@@ -5,12 +5,18 @@ const MistralAIAPIConfig: ProviderAPIConfig = {
   headers: ({ providerOptions }) => {
     return { Authorization: `Bearer ${providerOptions.apiKey}` };
   },
-  getEndpoint: ({ fn }) => {
-    switch (fn) {
+  getEndpoint: ({ fn, providerOptions }) => {
+    let mappedfn = fn;
+    if (providerOptions.mistralFimCompletion === 'true') {
+      mappedfn = 'fimComplete';
+    }
+    switch (mappedfn) {
       case 'chatComplete':
         return '/chat/completions';
       case 'embed':
         return '/embeddings';
+      case 'fimComplete':
+        return '/fim/completions';
       default:
         return '';
     }

diff --git a/src/providers/mistral-ai/chatComplete.ts b/src/providers/mistral-ai/chatComplete.ts
@@ -53,6 +53,16 @@ export const MistralAIChatCompleteConfig: ProviderConfig = {
     param: 'safe_prompt',
     default: false,
   },
+  prompt: {
+    param: 'prompt',
+    required: false,
+    default: '',
+  },
+  suffix: {
+    param: 'suffix',
+    required: false,
+    default: '',
+  },
 };
 
 interface MistralAIChatCompleteResponse extends ChatCompletionResponse {

diff --git a/src/types/requestBody.ts b/src/types/requestBody.ts
@@ -100,7 +100,8 @@ export interface Options {
 
   /** The parameter to determine if extra non-openai compliant fields should be returned in response */
   strictOpenAiCompliance?: boolean;
-
+  /** Parameter to determine if fim/completions endpoint is to be used */
+  mistralFimCompletion?: String;
   /** Anthropic specific headers */
   anthropicBeta?: string;
   anthropicVersion?: string;