feat: add chat session limit

toeverything · Apr 10, 2024 · b04f5e1 · b04f5e1
1 parent ebbdcea
commit b04f5e1
Showing 1 changed file with 54 additions and 30 deletions.
diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts
@@ -111,36 +111,56 @@ export class ChatSessionService {
     private readonly prompt: PromptService
   ) {}
 
-  private async setSession(state: ChatSessionState): Promise<void> {
-    await this.db.aiSession.upsert({
-      where: {
-        id: state.sessionId,
-        userId: state.userId,
-      },
-      update: {
-        messages: {
-          // delete old messages
-          deleteMany: {},
-          create: state.messages.map(m => ({
-            ...m,
-            params: m.params || undefined,
-          })),
+  private async setSession(state: ChatSessionState): Promise<string> {
+    return await this.db.$transaction(async tx => {
+      let sessionId = state.sessionId;
+
+      // find existing session if session is chat session
+      if (!state.prompt.action) {
+        const { id } =
+          (await tx.aiSession.findFirst({
+            where: {
+              userId: state.userId,
+              workspaceId: state.workspaceId,
+              docId: state.docId,
+              prompt: { action: { equals: null } },
+            },
+            select: { id: true },
+          })) || {};
+        if (id) sessionId = id;
+      }
+
+      await tx.aiSession.upsert({
+        where: {
+          id: sessionId,
+          userId: state.userId,
+        },
+        update: {
+          messages: {
+            // delete old messages
+            deleteMany: {},
+            create: state.messages.map(m => ({
+              ...m,
+              params: m.params || undefined,
+            })),
+          },
         },
-      },
-      create: {
-        id: state.sessionId,
-        workspaceId: state.workspaceId,
-        docId: state.docId,
-        messages: {
-          create: state.messages.map(m => ({
-            ...m,
-            params: m.params || undefined,
-          })),
+        create: {
+          id: sessionId,
+          workspaceId: state.workspaceId,
+          docId: state.docId,
+          messages: {
+            create: state.messages.map(m => ({
+              ...m,
+              params: m.params || undefined,
+            })),
+          },
+          // connect
+          user: { connect: { id: state.userId } },
+          prompt: { connect: { name: state.prompt.name } },
         },
-        // connect
-        user: { connect: { id: state.userId } },
-        prompt: { connect: { name: state.prompt.name } },
-      },
+      });
+      return sessionId;
     });
   }
 
@@ -293,8 +313,12 @@ export class ChatSessionService {
       this.logger.error(`Prompt not found: ${options.promptName}`);
       throw new Error('Prompt not found');
     }
-    await this.setSession({ ...options, sessionId, prompt, messages: [] });
-    return sessionId;
+    return await this.setSession({
+      ...options,
+      sessionId,
+      prompt,
+      messages: [],
+    });
   }
 
   /**