Merge pull request #365 from Portkey-AI/feat/support-openai-org-and-p…

…roject Feat: support openai organization and project
Portkey-AI · May 15, 2024 · 602567e · 602567e
2 parents 863fddf + 6a2e031
commit 602567e
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 1 deletion.
diff --git a/src/handlers/handlerUtils.ts b/src/handlers/handlerUtils.ts
@@ -9,6 +9,7 @@ import {
   RESPONSE_HEADER_KEYS,
   RETRY_STATUS_CODES,
   GOOGLE_VERTEX_AI,
+  OPEN_AI,
 } from '../globals';
 import Providers from '../providers';
 import { ProviderAPIConfig, endpointStrings } from '../providers/types';
@@ -989,6 +990,11 @@ export function constructConfigFromRequestHeaders(
     workersAiAccountId: requestHeaders[`x-${POWERED_BY}-workers-ai-account-id`],
   };
 
+  const openAiConfig = {
+    openaiOrganization: requestHeaders[`x-${POWERED_BY}-openai-organization`],
+    openaiProject: requestHeaders[`x-${POWERED_BY}-openai-project`],
+  };
+
   const vertexConfig = {
     vertexProjectId: requestHeaders[`x-${POWERED_BY}-vertex-project-id`],
     vertexRegion: requestHeaders[`x-${POWERED_BY}-vertex-region`],
@@ -1024,6 +1030,13 @@ export function constructConfigFromRequestHeaders(
           ...workersAiConfig,
         };
       }
+
+      if (parsedConfigJson.provider === OPEN_AI) {
+        parsedConfigJson = {
+          ...parsedConfigJson,
+          ...openAiConfig,
+        };
+      }
     }
     return convertKeysToCamelCase(parsedConfigJson, [
       'override_params',
@@ -1042,5 +1055,6 @@ export function constructConfigFromRequestHeaders(
       workersAiConfig),
     ...(requestHeaders[`x-${POWERED_BY}-provider`] === GOOGLE_VERTEX_AI &&
       vertexConfig),
+    ...(requestHeaders[`x-${POWERED_BY}-provider`] === OPEN_AI && openAiConfig),
   };
 }
diff --git a/src/middlewares/requestValidator/schema/config.ts b/src/middlewares/requestValidator/schema/config.ts
@@ -62,6 +62,9 @@ export const configSchema: any = z
     // Google Vertex AI specific
     vertex_project_id: z.string().optional(),
     vertex_region: z.string().optional(),
+    // OpenAI specific
+    openai_project: z.string().optional(),
+    openai_organization: z.string().optional(),
   })
   .refine(
     (value) => {

diff --git a/src/providers/openai/api.ts b/src/providers/openai/api.ts
@@ -3,7 +3,18 @@ import { ProviderAPIConfig } from '../types';
 const OpenAIAPIConfig: ProviderAPIConfig = {
   getBaseURL: () => 'https://api.openai.com/v1',
   headers: ({ providerOptions }) => {
-    return { Authorization: `Bearer ${providerOptions.apiKey}` };
+    const headersObj: Record<string, string> = {
+      Authorization: `Bearer ${providerOptions.apiKey}`,
+    };
+    if (providerOptions.openaiOrganization) {
+      headersObj['OpenAI-Organization'] = providerOptions.openaiOrganization;
+    }
+
+    if (providerOptions.openaiProject) {
+      headersObj['OpenAI-Project'] = providerOptions.openaiProject;
+    }
+
+    return headersObj;
   },
   getEndpoint: ({ fn }) => {
     switch (fn) {

diff --git a/src/types/requestBody.ts b/src/types/requestBody.ts
@@ -63,6 +63,10 @@ export interface Options {
   /** Google Vertex AI specific */
   vertexRegion?: string;
   vertexProjectId?: string;
+
+  /** OpenAI specific */
+  openaiProject?: string;
+  openaiOrganization?: string;
 }
 
 /**