@lobehub/chat 1.35.9 → 1.35.11

package/src/database/server/models/__tests__/message.test.ts

@@ -2,10 +2,16 @@ import { eq } from 'drizzle-orm';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { uuid } from '@/utils/uuid';
 
 import {
+  chunks,
+  embeddings,
+  fileChunks,
   files,
   messagePlugins,
+  messageQueries,
+  messageQueryChunks,
   messageTTS,
   messageTranslates,
   messages,
@@ -15,17 +21,12 @@ import {
   users,
 } from '../../../schemas';
 import { MessageModel } from '../message';
+import { codeEmbedding } from './fixtures/embedding';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'message-db';
-const messageModel = new MessageModel(userId);
+const messageModel = new MessageModel(serverDB, userId);
 
 beforeEach(async () => {
   // 在每个测试用例之前，清空表
@@ -273,6 +274,93 @@ describe('MessageModel', () => {
       const result3 = await messageModel.query({ current: 2, pageSize: 2 });
       expect(result3).toHaveLength(0);
     });
+
+    // 补充测试复杂查询场景
+    it('should handle complex query with multiple joins and file chunks', async () => {
+      await serverDB.transaction(async (trx) => {
+        const chunk1Id = uuid();
+        const query1Id = uuid();
+        // 创建基础消息
+        await trx.insert(messages).values({
+          id: 'msg1',
+          userId,
+          role: 'user',
+          content: 'test message',
+          createdAt: new Date('2023-01-01'),
+        });
+
+        // 创建文件
+        await trx.insert(files).values([
+          {
+            id: 'file1',
+            userId,
+            name: 'test.txt',
+            url: 'test-url',
+            fileType: 'text/plain',
+            size: 100,
+          },
+        ]);
+
+        // 创建文件块
+        await trx.insert(chunks).values({
+          id: chunk1Id,
+          text: 'chunk content',
+        });
+
+        // 关联消息和文件
+        await trx.insert(messagesFiles).values({
+          messageId: 'msg1',
+          fileId: 'file1',
+        });
+
+        // 创建文件块关联
+        await trx.insert(fileChunks).values({
+          fileId: 'file1',
+          chunkId: chunk1Id,
+        });
+
+        // 创建消息查询
+        await trx.insert(messageQueries).values({
+          id: query1Id,
+          messageId: 'msg1',
+          userQuery: 'original query',
+          rewriteQuery: 'rewritten query',
+        });
+
+        // 创建消息查询块关联
+        await trx.insert(messageQueryChunks).values({
+          messageId: 'msg1',
+          queryId: query1Id,
+          chunkId: chunk1Id,
+          similarity: '0.95',
+        });
+      });
+
+      const result = await messageModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].chunksList).toHaveLength(1);
+      expect(result[0].chunksList[0]).toMatchObject({
+        text: 'chunk content',
+        similarity: 0.95,
+      });
+    });
+
+    it('should return empty arrays for files and chunks if none exist', async () => {
+      await serverDB.insert(messages).values({
+        id: 'msg1',
+        userId,
+        role: 'user',
+        content: 'test message',
+      });
+
+      const result = await messageModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].fileList).toEqual([]);
+      expect(result[0].imageList).toEqual([]);
+      expect(result[0].chunksList).toEqual([]);
+    });
   });
 
   describe('queryAll', () => {
@@ -1021,4 +1109,139 @@ describe('MessageModel', () => {
       expect(result).toBe(2);
     });
   });
+
+  describe('findMessageQueriesById', () => {
+    it('should return undefined for non-existent message query', async () => {
+      const result = await messageModel.findMessageQueriesById('non-existent-id');
+      expect(result).toBeUndefined();
+    });
+
+    it('should return message query with embeddings', async () => {
+      const query1Id = uuid();
+      const embeddings1Id = uuid();
+
+      await serverDB.transaction(async (trx) => {
+        await trx.insert(messages).values({ id: 'msg1', userId, role: 'user', content: 'abc' });
+
+        await trx.insert(embeddings).values({
+          id: embeddings1Id,
+          embeddings: codeEmbedding,
+        });
+
+        await trx.insert(messageQueries).values({
+          id: query1Id,
+          messageId: 'msg1',
+          userQuery: 'test query',
+          rewriteQuery: 'rewritten query',
+          embeddingsId: embeddings1Id,
+        });
+      });
+
+      const result = await messageModel.findMessageQueriesById('msg1');
+
+      expect(result).toBeDefined();
+      expect(result).toMatchObject({
+        id: query1Id,
+        userQuery: 'test query',
+        rewriteQuery: 'rewritten query',
+        embeddings: codeEmbedding,
+      });
+    });
+  });
+
+  describe('deleteMessagesBySession', () => {
+    it('should delete messages by session ID', async () => {
+      await serverDB.insert(sessions).values([
+        { id: 'session1', userId },
+        { id: 'session2', userId },
+      ]);
+
+      await serverDB.insert(messages).values([
+        {
+          id: '1',
+          userId,
+          sessionId: 'session1',
+          role: 'user',
+          content: 'message 1',
+        },
+        {
+          id: '2',
+          userId,
+          sessionId: 'session1',
+          role: 'assistant',
+          content: 'message 2',
+        },
+        {
+          id: '3',
+          userId,
+          sessionId: 'session2',
+          role: 'user',
+          content: 'message 3',
+        },
+      ]);
+
+      await messageModel.deleteMessagesBySession('session1');
+
+      const remainingMessages = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.userId, userId));
+
+      expect(remainingMessages).toHaveLength(1);
+      expect(remainingMessages[0].id).toBe('3');
+    });
+
+    it('should delete messages by session ID and topic ID', async () => {
+      await serverDB.insert(sessions).values([{ id: 'session1', userId }]);
+      await serverDB.insert(topics).values([
+        { id: 'topic1', sessionId: 'session1', userId },
+        { id: 'topic2', sessionId: 'session1', userId },
+      ]);
+
+      await serverDB.insert(messages).values([
+        {
+          id: '1',
+          userId,
+          sessionId: 'session1',
+          topicId: 'topic1',
+          role: 'user',
+          content: 'message 1',
+        },
+        {
+          id: '2',
+          userId,
+          sessionId: 'session1',
+          topicId: 'topic2',
+          role: 'assistant',
+          content: 'message 2',
+        },
+      ]);
+
+      await messageModel.deleteMessagesBySession('session1', 'topic1');
+
+      const remainingMessages = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.userId, userId));
+
+      expect(remainingMessages).toHaveLength(1);
+      expect(remainingMessages[0].id).toBe('2');
+    });
+  });
+
+  describe('genId', () => {
+    it('should generate unique message IDs', () => {
+      const model = new MessageModel(serverDB, userId);
+      // @ts-ignore - accessing private method for testing
+      const id1 = model.genId();
+      // @ts-ignore - accessing private method for testing
+      const id2 = model.genId();
+
+      expect(id1).toHaveLength(18);
+      expect(id2).toHaveLength(18);
+      expect(id1).not.toBe(id2);
+      expect(id1).toMatch(/^msg_/);
+      expect(id2).toMatch(/^msg_/);
+    });
+  });
 });

package/src/database/server/models/__tests__/nextauth.test.ts

@@ -8,7 +8,6 @@ import type {
 import { eq } from 'drizzle-orm';
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import {
   nextauthAccounts,
   nextauthAuthenticators,
@@ -16,17 +15,12 @@ import {
   nextauthVerificationTokens,
   users,
 } from '@/database/schemas';
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import { LobeNextAuthDbAdapter } from '@/libs/next-auth/adapter';
 
 let serverDB = await getTestDBInstance();
 let nextAuthAdapter = LobeNextAuthDbAdapter(serverDB);
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'user-db';
 const user: AdapterUser = {
   id: userId,

package/src/database/server/models/__tests__/plugin.test.ts

@@ -8,14 +8,8 @@ import { PluginModel } from '../plugin';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'plugin-db';
-const pluginModel = new PluginModel(userId);
+const pluginModel = new PluginModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.transaction(async (trx) => {

package/src/database/server/models/__tests__/session.test.ts

@@ -1,7 +1,9 @@
-import { eq, inArray } from 'drizzle-orm';
+import { and, eq, inArray } from 'drizzle-orm';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
+import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { idGenerator } from '@/database/utils/idGenerator';
 
 import {
   NewSession,
@@ -14,19 +16,12 @@ import {
   topics,
   users,
 } from '../../../schemas';
-import { idGenerator } from '@/database/utils/idGenerator';
 import { SessionModel } from '../session';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'session-user';
-const sessionModel = new SessionModel(userId);
+const sessionModel = new SessionModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
@@ -259,7 +254,13 @@ describe('SessionModel', () => {
       ]);
 
       await serverDB.insert(agents).values([
-        { id: 'agent-1', userId, model: 'gpt-3.5-turbo', title: 'Agent 1', description: 'Description with Keyword' },
+        {
+          id: 'agent-1',
+          userId,
+          model: 'gpt-3.5-turbo',
+          title: 'Agent 1',
+          description: 'Description with Keyword',
+        },
         { id: 'agent-2', userId, model: 'gpt-4', title: 'Agent 2' },
       ]);
 
@@ -338,7 +339,7 @@ describe('SessionModel', () => {
     });
   });
 
-  describe.skip('batchCreate', () => {
+  describe('batchCreate', () => {
     it('should batch create sessions', async () => {
       // 调用 batchCreate 方法
       const sessions: NewSession[] = [
@@ -624,4 +625,161 @@ describe('SessionModel', () => {
       expect(await serverDB.select().from(sessions).where(eq(sessions.id, '3'))).toHaveLength(1);
     });
   });
+
+  // 在原有的 describe('SessionModel') 中添加以下测试套件
+
+  describe('createInbox', () => {
+    it('should create inbox session if not exists', async () => {
+      const inbox = await sessionModel.createInbox();
+
+      expect(inbox).toBeDefined();
+      expect(inbox?.slug).toBe('inbox');
+
+      // verify agent config
+      const session = await sessionModel.findByIdOrSlug('inbox');
+      expect(session?.agent).toBeDefined();
+      expect(session?.agent.model).toBe(DEFAULT_AGENT_CONFIG.model);
+    });
+
+    it('should not create duplicate inbox session', async () => {
+      // Create first inbox
+      await sessionModel.createInbox();
+
+      // Try to create another inbox
+      const duplicateInbox = await sessionModel.createInbox();
+
+      // Should return undefined as inbox already exists
+      expect(duplicateInbox).toBeUndefined();
+
+      // Verify only one inbox exists
+      const sessions = await serverDB.query.sessions.findMany();
+
+      const inboxSessions = sessions.filter((s) => s.slug === 'inbox');
+      expect(inboxSessions).toHaveLength(1);
+    });
+  });
+
+  describe('deleteAll', () => {
+    it('should delete all sessions for current user', async () => {
+      // Create test data
+      await serverDB.insert(sessions).values([
+        { id: '1', userId },
+        { id: '2', userId },
+        { id: '3', userId },
+      ]);
+
+      // Create sessions for another user that should not be deleted
+      await serverDB.insert(users).values([{ id: 'other-user' }]);
+      await serverDB.insert(sessions).values([
+        { id: '4', userId: 'other-user' },
+        { id: '5', userId: 'other-user' },
+      ]);
+
+      await sessionModel.deleteAll();
+
+      // Verify all sessions for current user are deleted
+      const remainingSessions = await serverDB
+        .select()
+        .from(sessions)
+        .where(eq(sessions.userId, userId));
+      expect(remainingSessions).toHaveLength(0);
+
+      // Verify other user's sessions are not deleted
+      const otherUserSessions = await serverDB
+        .select()
+        .from(sessions)
+        .where(eq(sessions.userId, 'other-user'));
+      expect(otherUserSessions).toHaveLength(2);
+    });
+
+    it('should delete associated data when deleting all sessions', async () => {
+      // Create test data with associated records
+      await serverDB.transaction(async (trx) => {
+        await trx.insert(sessions).values([
+          { id: '1', userId },
+          { id: '2', userId },
+        ]);
+
+        await trx.insert(topics).values([
+          { id: 't1', sessionId: '1', userId },
+          { id: 't2', sessionId: '2', userId },
+        ]);
+
+        await trx.insert(messages).values([
+          { id: 'm1', sessionId: '1', userId, role: 'user' },
+          { id: 'm2', sessionId: '2', userId, role: 'assistant' },
+        ]);
+      });
+
+      await sessionModel.deleteAll();
+
+      // Verify all associated data is deleted
+      const remainingTopics = await serverDB.select().from(topics).where(eq(topics.userId, userId));
+      expect(remainingTopics).toHaveLength(0);
+
+      const remainingMessages = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.userId, userId));
+      expect(remainingMessages).toHaveLength(0);
+    });
+  });
+
+  describe('updateConfig', () => {
+    it('should update agent config', async () => {
+      // Create test agent
+      const agentId = 'test-agent';
+      await serverDB.insert(agents).values({
+        id: agentId,
+        userId,
+        model: 'gpt-3.5-turbo',
+        title: 'Original Title',
+      });
+
+      // Update config
+      await sessionModel.updateConfig(agentId, {
+        model: 'gpt-4',
+        title: 'Updated Title',
+        description: 'New description',
+      });
+
+      // Verify update
+      const updatedAgent = await serverDB
+        .select()
+        .from(agents)
+        .where(and(eq(agents.id, agentId), eq(agents.userId, userId)));
+
+      expect(updatedAgent[0]).toMatchObject({
+        model: 'gpt-4',
+        title: 'Updated Title',
+        description: 'New description',
+      });
+    });
+
+    it('should not update config for other users agents', async () => {
+      // Create agent for another user
+      const agentId = 'other-agent';
+      await serverDB.insert(users).values([{ id: 'other-user' }]);
+      await serverDB.insert(agents).values({
+        id: agentId,
+        userId: 'other-user',
+        model: 'gpt-3.5-turbo',
+        title: 'Original Title',
+      });
+
+      // Try to update other user's agent
+      await sessionModel.updateConfig(agentId, {
+        model: 'gpt-4',
+        title: 'Updated Title',
+      });
+
+      // Verify no changes were made
+      const agent = await serverDB.select().from(agents).where(eq(agents.id, agentId));
+
+      expect(agent[0]).toMatchObject({
+        model: 'gpt-3.5-turbo',
+        title: 'Original Title',
+      });
+    });
+  });
 });

package/src/database/server/models/__tests__/sessionGroup.test.ts

@@ -10,14 +10,8 @@ import { SessionGroupModel } from '../sessionGroup';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'session-group-model-test-user-id';
-const sessionGroupModel = new SessionGroupModel(userId);
+const sessionGroupModel = new SessionGroupModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
@@ -75,7 +69,7 @@ describe('SessionGroupModel', () => {
       await sessionGroupModel.create({ name: 'Test Group 1' });
       await sessionGroupModel.create({ name: 'Test Group 333' });
 
-      const anotherSessionGroupModel = new SessionGroupModel('user2');
+      const anotherSessionGroupModel = new SessionGroupModel(serverDB, 'user2');
       await anotherSessionGroupModel.create({ name: 'Test Group 2' });
 
       await sessionGroupModel.deleteAll();

package/src/database/server/models/__tests__/topic.test.ts

@@ -8,15 +8,9 @@ import { CreateTopicParams, TopicModel } from '../topic';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'topic-user-test';
 const sessionId = 'topic-session';
-const topicModel = new TopicModel(userId);
+const topicModel = new TopicModel(serverDB, userId);
 
 describe('TopicModel', () => {
   beforeEach(async () => {