feat: Add Snowflake Cortex LLM API

src/globals.ts

@@ -76,6 +76,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';
@@ -129,6 +130,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 {
@@ -1030,6 +1031,11 @@ export function constructConfigFromRequestHeaders(
     }
   }
 
+  const cortexConfig = {
+    snowflakeAccount: requestHeaders[`x-${POWERED_BY}-snowflake-account`],
+    snowflakeApiKey: requestHeaders[`x-${POWERED_BY}-snowflake-api-key`],
+  };
+
   if (requestHeaders[`x-${POWERED_BY}-config`]) {
     let parsedConfigJson = JSON.parse(requestHeaders[`x-${POWERED_BY}-config`]);
 
@@ -1117,6 +1123,13 @@ export function constructConfigFromRequestHeaders(
           ...stabilityAiConfig,
         };
       }
+
+      if (parsedConfigJson.provider === CORTEX) {
+        parsedConfigJson = {
+          ...parsedConfigJson,
+          ...cortexConfig,
+        };
+      }
     }
     return convertKeysToCamelCase(parsedConfigJson, [
       'override_params',
@@ -1157,6 +1170,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 as any).snowflakeApiKey || 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,23 @@
+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,
+    complete: true,
+    embed: true,
+  }),
+};
+
+export default CortexConfig;

src/types/requestBody.ts

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