From edc67e9ca0c3121c89e57b387270b46d0fa175b7 Mon Sep 17 00:00:00 2001
From: Santthosh Selvadurai <santthosh@users.noreply.github.com>
Date: Fri, 28 Jun 2024 14:34:41 -0700
Subject: [PATCH] #93 Bug fix for the Open AI BYOK

---
 src/app/api/assistants/route.ts |  4 ++--
 src/app/api/utils/openai.ts     | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/app/api/assistants/route.ts b/src/app/api/assistants/route.ts
index 7bb32fe..68009ec 100644
--- 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
index 125da4e..958aec6 100644
--- 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');
   }