#10 Refactoring to support Global keys

assistants-hub · Feb 5, 2024 · a9c2806 · a9c2806
1 parent ef14228
commit a9c2806
Show file tree

Hide file tree

Showing 17 changed files with 255 additions and 180 deletions.
diff --git a/prisma/migrations/20240204202910_use_accounts_instead_of_credentials/migration.sql b/prisma/migrations/20240204202910_use_accounts_instead_of_credentials/migration.sql
@@ -0,0 +1,39 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `credentialsOwner` on the `Assistant` table. All the data in the column will be lost.
+  - You are about to drop the column `credentialsOwnerType` on the `Assistant` table. All the data in the column will be lost.
+  - You are about to drop the `Credentials` table. If the table is not empty, all the data it contains will be lost.
+  - Added the required column `limitRunsPerDay` to the `Assistant` table without a default value. This is not possible if the table is not empty.
+  - Added the required column `limitThreadsPerDay` to the `Assistant` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- DropForeignKey
+ALTER TABLE "Assistant" DROP CONSTRAINT "Assistant_credentialsOwner_credentialsOwnerType_fkey";
+
+-- AlterTable
+ALTER TABLE "Assistant" DROP COLUMN "credentialsOwner",
+DROP COLUMN "credentialsOwnerType",
+ADD COLUMN     "accountOwner" TEXT,
+ADD COLUMN     "accountOwnerType" TEXT,
+ADD COLUMN     "extendedAssistant" JSONB,
+ADD COLUMN     "limitRunsPerDay" INTEGER NOT NULL,
+ADD COLUMN     "limitThreadsPerDay" INTEGER NOT NULL;
+
+-- DropTable
+DROP TABLE "Credentials";
+
+-- CreateTable
+CREATE TABLE "Account" (
+    "owner" TEXT NOT NULL,
+    "ownerType" TEXT NOT NULL DEFAULT 'personal',
+    "openAIApiKey" TEXT NOT NULL,
+    "assistantsPerAccount" INTEGER NOT NULL,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "Account_pkey" PRIMARY KEY ("owner","ownerType")
+);
+
+-- AddForeignKey
+ALTER TABLE "Assistant" ADD CONSTRAINT "Assistant_accountOwner_accountOwnerType_fkey" FOREIGN KEY ("accountOwner", "accountOwnerType") REFERENCES "Account"("owner", "ownerType") ON DELETE SET NULL ON UPDATE CASCADE;
diff --git a/prisma/migrations/20240204204409_use_consistent_names_for_limits/migration.sql b/prisma/migrations/20240204204409_use_consistent_names_for_limits/migration.sql
@@ -0,0 +1,9 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `assistantsPerAccount` on the `Account` table. All the data in the column will be lost.
+
+*/
+-- AlterTable
+ALTER TABLE "Account" DROP COLUMN "assistantsPerAccount",
+ADD COLUMN     "limitAssistants" INTEGER NOT NULL DEFAULT 2;
diff --git a/prisma/migrations/20240204204845_using_modified_name_for_open_ai_objects/migration.sql b/prisma/migrations/20240204204845_using_modified_name_for_open_ai_objects/migration.sql
@@ -0,0 +1,9 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `extendedAssistant` on the `Assistant` table. All the data in the column will be lost.
+
+*/
+-- AlterTable
+ALTER TABLE "Assistant" DROP COLUMN "extendedAssistant",
+ADD COLUMN     "object" JSONB;
diff --git a/prisma/migrations/20240205001258_removing_the_thread_limits/migration.sql b/prisma/migrations/20240205001258_removing_the_thread_limits/migration.sql
@@ -0,0 +1,10 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `limitRunsPerDay` on the `Assistant` table. All the data in the column will be lost.
+  - You are about to drop the column `limitThreadsPerDay` on the `Assistant` table. All the data in the column will be lost.
+
+*/
+-- AlterTable
+ALTER TABLE "Assistant" DROP COLUMN "limitRunsPerDay",
+DROP COLUMN "limitThreadsPerDay";
diff --git a/prisma/schema.prisma b/prisma/schema.prisma
@@ -10,22 +10,24 @@ datasource db {
   url      = env("POSTGRES_PRISMA_URL")
 }
 
-model Credentials {
-  owner        String
-  ownerType    String      @default("personal")
-  openAIApiKey String
-  created_at   DateTime    @default(now())
-  updated_at   DateTime    @updatedAt
-  Assistant    Assistant[]
+model Account {
+  owner           String
+  ownerType       String      @default("personal")
+  openAIApiKey    String
+  limitAssistants Int         @default(2)
+  created_at      DateTime    @default(now())
+  updated_at      DateTime    @updatedAt
+  Assistant       Assistant[]
 
   @@id([owner, ownerType])
 }
 
 model Assistant {
-  id                   String       @id
-  credentialsOwner     String?
-  credentialsOwnerType String?
-  credentials          Credentials? @relation(fields: [credentialsOwner, credentialsOwnerType], references: [owner, ownerType])
-  created_at           DateTime     @default(now())
-  updated_at           DateTime     @updatedAt
+  id               String   @id
+  object           Json?
+  accountOwner     String?
+  accountOwnerType String?
+  account          Account? @relation(fields: [accountOwner, accountOwnerType], references: [owner, ownerType])
+  created_at       DateTime @default(now())
+  updated_at       DateTime @updatedAt
 }
diff --git a/src/app/api/credentials/route.ts b/src/app/api/credentials/route.ts
@@ -9,16 +9,16 @@ export async function GET(req: NextRequest, res: NextResponse) {
   const token = await getToken({ req });
   if (token) {
     // Signed in
-    let credentials = await prisma.credentials.findFirst({
+    let account = await prisma.account.findFirst({
       where: {
         owner: token.sub,
       },
     });
 
-    if (!credentials) {
+    if (!account) {
       return Response.json([]);
     } else {
-      return Response.json([credentials]);
+      return Response.json([{ hasCredentials: account.openAIApiKey !== null }]);
     }
   } else {
     // Not Signed in
@@ -32,6 +32,10 @@ export async function POST(req: NextRequest, res: NextResponse) {
     const body = await req.json();
 
     if (body.openAiApiKey) {
+      if (body.openAiApiKey.toLowerCase() === 'use-default') {
+        body.openAiApiKey = process.env.OPENAI_API_KEY;
+      }
+
       // Validate the Open API Key
       let models = null;
       try {
@@ -45,7 +49,7 @@ export async function POST(req: NextRequest, res: NextResponse) {
       }
 
       // Signed in
-      let credential = await prisma.credentials.upsert({
+      await prisma.account.upsert({
         where: {
           owner_ownerType: {
             owner: token.sub,

diff --git a/src/app/api/openai/assistants/[id]/route.ts b/src/app/api/openai/assistants/[id]/route.ts
@@ -10,36 +10,32 @@ const getId = (req: Request) => {
   return url.pathname.split('/').splice(-1, 1)[0];
 };
 
+// TODO: Use local DB instead of OpenAI API
 export async function GET(req: NextRequest, res: NextResponse) {
   const token = await getToken({ req });
 
   const id = getId(req);
 
   if (token) {
-    let credential = await prisma.credentials.findFirst({
+    let assistant = await prisma.assistant.findFirst({
       where: {
-        owner: token.sub,
-        ownerType: 'personal',
+        id: id,
+        accountOwner: token.sub,
+        accountOwnerType: 'personal',
+      },
+      select: {
+        id: true,
+        object: true,
       },
     });
 
-    if (credential) {
-      const openai = new OpenAI({
-        apiKey: credential.openAIApiKey,
-      });
-
-      try {
-        const listResponse = await openai.beta.assistants.retrieve(id);
-        return Response.json(listResponse, { status: 200 });
-      } catch (err: any) {
-        return Response.json({ message: err.message }, { status: err.status });
-      }
-    } else {
+    if (!assistant) {
       return Response.json(
-        { message: 'OpenAI API Key does not exist' },
-        { status: 400 }
+        { message: 'Assistant does not exist' },
+        { status: 404 }
       );
     }
+    return Response.json(assistant.object, { status: 200 });
   } else {
     // Not Signed in
     return Response.json({ message: 'Unauthenticated' }, { status: 401 });
@@ -52,16 +48,16 @@ export async function DELETE(req: NextRequest, res: NextResponse) {
   const id = getId(req);
 
   if (token) {
-    let credential = await prisma.credentials.findFirst({
+    let account = await prisma.account.findFirst({
       where: {
         owner: token.sub,
         ownerType: 'personal',
       },
     });
 
-    if (credential) {
+    if (account) {
       const openai = new OpenAI({
-        apiKey: credential.openAIApiKey,
+        apiKey: account.openAIApiKey,
       });
 
       try {

diff --git a/src/app/api/openai/assistants/route.ts b/src/app/api/openai/assistants/route.ts
@@ -9,16 +9,16 @@ export async function POST(req: NextRequest, res: NextResponse) {
   const token = await getToken({ req });
 
   if (token) {
-    let credential = await prisma.credentials.findFirst({
+    let account = await prisma.account.findFirst({
       where: {
         owner: token.sub,
         ownerType: 'personal',
       },
     });
 
-    if (credential) {
+    if (account) {
       const openai = new OpenAI({
-        apiKey: credential.openAIApiKey,
+        apiKey: account.openAIApiKey,
       });
 
       const body = await req.json();
@@ -33,18 +33,21 @@ export async function POST(req: NextRequest, res: NextResponse) {
           },
           update: {
             id: createResponse.id,
-            credentialsOwner: token.sub,
-            credentialsOwnerType: 'personal',
+            accountOwner: token.sub,
+            accountOwnerType: 'personal',
+            object: createResponse as any,
           },
           create: {
             id: createResponse.id,
-            credentialsOwner: token.sub,
-            credentialsOwnerType: 'personal',
+            accountOwner: token.sub,
+            accountOwnerType: 'personal',
+            object: createResponse as any,
           },
         });
 
         return Response.json(createResponse, { status: 201 });
       } catch (err: any) {
+        console.log(err);
         return Response.json({ message: err.message }, { status: err.status });
       }
     } else {
@@ -63,30 +66,20 @@ export async function GET(req: NextRequest, res: NextResponse) {
   const token = await getToken({ req });
 
   if (token) {
-    let credential = await prisma.credentials.findFirst({
+    let assistants = await prisma.assistant.findMany({
       where: {
-        owner: token.sub,
-        ownerType: 'personal',
+        accountOwner: token.sub,
+        accountOwnerType: 'personal',
+      },
+      select: {
+        id: true,
+        object: true,
       },
     });
-
-    if (credential) {
-      const openai = new OpenAI({
-        apiKey: credential.openAIApiKey,
-      });
-
-      try {
-        const listResponse = await openai.beta.assistants.list();
-        return Response.json(listResponse, { status: 200 });
-      } catch (err: any) {
-        return Response.json({ message: err.message }, { status: err.status });
-      }
-    } else {
-      return Response.json(
-        { message: 'OpenAI API Key does not exist' },
-        { status: 400 }
-      );
-    }
+    let assistantsCollection = assistants.map((assistant) => {
+      return assistant.object;
+    });
+    return Response.json(assistantsCollection, { status: 200 });
   } else {
     // Not Signed in
     return Response.json({ message: 'Unauthenticated' }, { status: 401 });

diff --git a/src/app/api/openai/models/route.ts b/src/app/api/openai/models/route.ts
@@ -5,19 +5,20 @@ import { NextRequest, NextResponse } from 'next/server';
 
 const prisma = new PrismaClient();
 
+// Note: We should not cache the models list as it may change frequently for different organizations
 export async function GET(req: NextRequest, res: NextResponse) {
   const token = await getToken({ req });
   if (token) {
-    let credential = await prisma.credentials.findFirst({
+    let account = await prisma.account.findFirst({
       where: {
         owner: token.sub,
         ownerType: 'personal',
       },
     });
 
-    if (credential) {
+    if (account) {
       const openai = new OpenAI({
-        apiKey: credential.openAIApiKey,
+        apiKey: account.openAIApiKey,
       });
       const models = await openai.models.list();
       return Response.json(models);