#93 Bug fix for the Open AI BYOK

assistants-hub · Jun 28, 2024 · edc67e9 · edc67e9
1 parent 8a7cf81
commit edc67e9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/src/app/api/assistants/route.ts b/src/app/api/assistants/route.ts
@@ -2,7 +2,7 @@ import { NextRequest, NextResponse } from 'next/server';
 import { ulid } from 'ulidx';
 import prisma from '@/app/api/utils/prisma';
 import { getSession } from '@auth0/nextjs-auth0';
-import { getOpenAIForOrganization } from '@/app/api/utils/openai';
+import { getOpenAIWithKey } from '@/app/api/utils/openai';
 
 export async function GET(req: NextRequest, res: NextResponse) {
   const session = await getSession();
@@ -72,7 +72,7 @@ export async function POST(req: NextRequest, res: NextResponse) {
       try {
         let createResponse = null;
         if (modelProviderId === 'openai') {
-          const openai = getOpenAIForOrganization(organization);
+          const openai = await getOpenAIWithKey(modelProviderKeyId);
 
           body.model = modelId;
 

diff --git a/src/app/api/utils/openai.ts b/src/app/api/utils/openai.ts
@@ -1,4 +1,5 @@
 import OpenAI from 'openai';
+import prisma from '@/app/api/utils/prisma';
 
 export function getOpenAI(assistant: any): OpenAI {
   let openAIAPIKey = process.env.OPENAI_API_KEY
@@ -18,11 +19,29 @@ export function getOpenAI(assistant: any): OpenAI {
   });
 }
 
-export function getOpenAIForOrganization(organization: any): OpenAI {
+export async function getOpenAIWithKey(modelProviderKeyId:string) {
   let openAIAPIKey = process.env.OPENAI_API_KEY
     ? process.env.OPENAI_API_KEY
     : null;
 
+  if (modelProviderKeyId) {
+    let modelProviderKey = await prisma.modelProviderKey.findFirst({
+      where: {
+        id: modelProviderKeyId,
+      },
+      select: {
+        id: true,
+        name: true,
+        key: true,
+      },
+    });
+
+    if(modelProviderKey && modelProviderKey.key) {
+      // @ts-ignore
+      openAIAPIKey = modelProviderKey.key['apiKey']
+    }
+  }
+
   if (!openAIAPIKey) {
     throw new Error('OpenAI API key missing');
   }