Merge pull request #988 from Ajay-Satish-01/cortex

Author: VisargD

src/globals.ts

@@ -77,6 +77,7 @@ export const UPSTAGE: string = 'upstage';
 export const LAMBDA: string = 'lambda';
 export const DASHSCOPE: string = 'dashscope';
 export const X_AI: string = 'x-ai';
+export const CORTEX: string = 'cortex';
 export const SAGEMAKER: string = 'sagemaker';
 export const NEBIUS: string = 'nebius';
 export const RECRAFTAI: string = 'recraft-ai';
@@ -131,6 +132,7 @@ export const VALID_PROVIDERS = [
   LAMBDA,
   DASHSCOPE,
   X_AI,
+  CORTEX,
   SAGEMAKER,
   NEBIUS,
   RECRAFTAI,

src/handlers/handlerUtils.ts

@@ -16,6 +16,7 @@ import {
   STABILITY_AI,
   SAGEMAKER,
   FIREWORKS_AI,
+  CORTEX,
 } from '../globals';
 import Providers from '../providers';
 import { ProviderAPIConfig, endpointStrings } from '../providers/types';
@@ -1072,6 +1073,10 @@ export function constructConfigFromRequestHeaders(
     }
   }
 
+  const cortexConfig = {
+    snowflakeAccount: requestHeaders[`x-${POWERED_BY}-snowflake-account`],
+  };
+
   const defaultsConfig = {
     input_guardrails: requestHeaders[`x-portkey-default-input-guardrails`]
       ? JSON.parse(requestHeaders[`x-portkey-default-input-guardrails`])
@@ -1171,6 +1176,13 @@ export function constructConfigFromRequestHeaders(
           ...stabilityAiConfig,
         };
       }
+
+      if (parsedConfigJson.provider === CORTEX) {
+        parsedConfigJson = {
+          ...parsedConfigJson,
+          ...cortexConfig,
+        };
+      }
     }
     return convertKeysToCamelCase(parsedConfigJson, [
       'override_params',
@@ -1215,6 +1227,7 @@ export function constructConfigFromRequestHeaders(
       stabilityAiConfig),
     ...(requestHeaders[`x-${POWERED_BY}-provider`] === FIREWORKS_AI &&
       fireworksConfig),
+    ...(requestHeaders[`x-${POWERED_BY}-provider`] === CORTEX && cortexConfig),
   };
 }
 

src/providers/cortex/api.ts

@@ -0,0 +1,22 @@
+import { ProviderAPIConfig } from '../types';
+
+const CortexAPIConfig: ProviderAPIConfig = {
+  getBaseURL: ({ providerOptions }) =>
+    `https://${(providerOptions as any).snowflakeAccount}.snowflakecomputing.com/api/v2`,
+  headers: ({ providerOptions }) => ({
+    'X-Snowflake-Authorization-Token-Type': 'KEYPAIR_JWT',
+    Authorization: `Bearer ${providerOptions.apiKey}`,
+    'Content-Type': 'application/json',
+    Accept: 'application/json, text/event-stream',
+  }),
+  getEndpoint: ({ fn }) => {
+    switch (fn) {
+      case 'chatComplete':
+        return '/cortex/inference:complete';
+      default:
+        return '';
+    }
+  },
+};
+
+export default CortexAPIConfig;

src/providers/cortex/index.ts

@@ -0,0 +1,21 @@
+import { ProviderConfigs } from '../types';
+import CortexAPIConfig from './api';
+import { CORTEX } from '../../globals';
+import {
+  chatCompleteParams,
+  completeParams,
+  embedParams,
+  responseTransformers,
+} from '../open-ai-base';
+
+const CortexConfig: ProviderConfigs = {
+  chatComplete: chatCompleteParams([], { model: 'mistral-large' }),
+  complete: completeParams([], { model: 'mistral-large' }),
+  embed: embedParams([], { model: 'mistral-large' }),
+  api: CortexAPIConfig,
+  responseTransforms: responseTransformers(CORTEX, {
+    chatComplete: true,
+  }),
+};
+
+export default CortexConfig;

src/types/requestBody.ts

@@ -150,6 +150,9 @@ export interface Options {
 
   /** Fireworks finetune required fields */
   fireworksAccountId?: string;
+
+  /** Cortex specific fields */
+  snowflakeAccount?: string;
 }
 
 /**