@lobehub/lobehub 2.0.0-next.40 → 2.0.0-next.41

package/packages/database/src/models/__tests__/messages/message.create.test.ts

@@ -0,0 +1,492 @@
+import { DBMessageItem } from '@lobechat/types';
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { uuid } from '@/utils/uuid';
+
+import {
+  chunks,
+  embeddings,
+  files,
+  messagePlugins,
+  messageQueries,
+  messageQueryChunks,
+  messages,
+  messagesFiles,
+  sessions,
+  users,
+} from '../../../schemas';
+import { LobeChatDatabase } from '../../../type';
+import { MessageModel } from '../../message';
+import { getTestDB } from '../_util';
+import { codeEmbedding } from '../fixtures/embedding';
+
+const serverDB: LobeChatDatabase = await getTestDB();
+
+const userId = 'message-create-test';
+const otherUserId = 'message-create-test-other';
+const messageModel = new MessageModel(serverDB, userId);
+const embeddingsId = uuid();
+
+beforeEach(async () => {
+  // Clear tables before each test case
+  await serverDB.transaction(async (trx) => {
+    await trx.delete(users).where(eq(users.id, userId));
+    await trx.delete(users).where(eq(users.id, otherUserId));
+    await trx.insert(users).values([{ id: userId }, { id: otherUserId }]);
+
+    await trx.insert(sessions).values([
+      // { id: 'session1', userId },
+      // { id: 'session2', userId },
+      { id: '1', userId },
+    ]);
+    await trx.insert(files).values({
+      id: 'f1',
+      userId: userId,
+      url: 'abc',
+      name: 'file-1',
+      fileType: 'image/png',
+      size: 1000,
+    });
+
+    await trx.insert(embeddings).values({
+      id: embeddingsId,
+      embeddings: codeEmbedding,
+      userId,
+    });
+  });
+});
+
+afterEach(async () => {
+  // Clear tables after each test case
+  await serverDB.delete(users).where(eq(users.id, userId));
+  await serverDB.delete(users).where(eq(users.id, otherUserId));
+});
+
+describe('MessageModel Create Tests', () => {
+  describe('createMessage', () => {
+    it('should create a new message', async () => {
+      // Call createMessage method
+      await messageModel.create({ role: 'user', content: 'new message', sessionId: '1' });
+
+      // Assert result
+      const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
+      expect(result).toHaveLength(1);
+      expect(result[0].content).toBe('new message');
+    });
+
+    it('should create a message', async () => {
+      const sessionId = 'session1';
+      await serverDB.insert(sessions).values([{ id: sessionId, userId }]);
+
+      const result = await messageModel.create({
+        content: 'message 1',
+        role: 'user',
+        sessionId: 'session1',
+      });
+
+      expect(result.id).toBeDefined();
+      expect(result.content).toBe('message 1');
+      expect(result.role).toBe('user');
+      expect(result.sessionId).toBe('session1');
+      expect(result.userId).toBe(userId);
+    });
+
+    it('should generate message ID automatically', async () => {
+      // Call createMessage method
+      await messageModel.create({
+        role: 'user',
+        content: 'new message',
+        sessionId: '1',
+      });
+
+      // Assert result
+      const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
+      expect(result[0].id).toBeDefined();
+      expect(result[0].id).toHaveLength(18);
+    });
+
+    it('should create a tool message and insert into messagePlugins table', async () => {
+      // Call create method
+      const result = await messageModel.create({
+        content: 'message 1',
+        role: 'tool',
+        sessionId: '1',
+        tool_call_id: 'tool1',
+        plugin: {
+          apiName: 'api1',
+          arguments: 'arg1',
+          identifier: 'plugin1',
+          type: 'default',
+        },
+      });
+
+      // Assert result
+      expect(result.id).toBeDefined();
+      expect(result.content).toBe('message 1');
+      expect(result.role).toBe('tool');
+      expect(result.sessionId).toBe('1');
+
+      const pluginResult = await serverDB
+        .select()
+        .from(messagePlugins)
+        .where(eq(messagePlugins.id, result.id));
+      expect(pluginResult).toHaveLength(1);
+      expect(pluginResult[0].identifier).toBe('plugin1');
+    });
+
+    it('should create tool message ', async () => {
+      // Call create method
+      const state = {
+        query: 'Composio',
+        answers: [],
+        results: [
+          {
+            url: 'https://www.composio.dev/',
+            score: 16,
+            title: 'Composio - Connect 90+ tools to your AI agents',
+            engine: 'bing',
+            content:
+              'Faster DevelopmentHigher ReliabilityBetter Integrations. Get Started Now. Our platform lets you ditch the specs and seamlessly integrate any tool you need in less than 5 mins.',
+            engines: ['bing', 'qwant', 'brave', 'duckduckgo'],
+            category: 'general',
+            template: 'default.html',
+            positions: [1, 1, 1, 1],
+            thumbnail: '',
+            parsed_url: ['https', 'www.composio.dev', '/', '', '', ''],
+            publishedDate: null,
+          },
+          {
+            url: 'https://www.composio.co/',
+            score: 10.75,
+            title: 'Composio',
+            engine: 'bing',
+            content:
+              'Composio was created to help streamline the entire book creation process! Writing. Take time out to write / Make a schedule to write consistently. We have writing software that optimizes your books for printing or ebook format. Figure out what you want to write. Collaborate and write with others. Professional editing is a necessity.',
+            engines: ['qwant', 'duckduckgo', 'google', 'bing', 'brave'],
+            category: 'general',
+            template: 'default.html',
+            positions: [5, 2, 1, 5, 4],
+            thumbnail: null,
+            parsed_url: ['https', 'www.composio.co', '/', '', '', ''],
+            publishedDate: null,
+          },
+        ],
+        unresponsive_engines: [],
+      };
+      const result = await messageModel.create({
+        content: '[{}]',
+        plugin: {
+          apiName: 'searchWithSearXNG',
+          arguments: '{\n  "query": "Composio"\n}',
+          identifier: 'lobe-web-browsing',
+          type: 'builtin',
+        },
+        pluginState: state,
+        role: 'tool',
+        tool_call_id: 'tool_call_ymxXC2J0',
+        sessionId: '1',
+      });
+
+      // Assert result
+      expect(result.id).toBeDefined();
+      expect(result.content).toBe('[{}]');
+      expect(result.role).toBe('tool');
+      expect(result.sessionId).toBe('1');
+
+      const pluginResult = await serverDB
+        .select()
+        .from(messagePlugins)
+        .where(eq(messagePlugins.id, result.id));
+      expect(pluginResult).toHaveLength(1);
+      expect(pluginResult[0].identifier).toBe('lobe-web-browsing');
+      expect(pluginResult[0].state!).toMatchObject(state);
+    });
+
+    describe('create with advanced parameters', () => {
+      it('should create a message with custom ID', async () => {
+        const customId = 'custom-msg-id';
+
+        const result = await messageModel.create(
+          {
+            role: 'user',
+            content: 'message with custom ID',
+            sessionId: '1',
+          },
+          customId,
+        );
+
+        expect(result.id).toBe(customId);
+
+        // Verify database records
+        const dbResult = await serverDB.select().from(messages).where(eq(messages.id, customId));
+        expect(dbResult).toHaveLength(1);
+        expect(dbResult[0].id).toBe(customId);
+      });
+
+      it.skip('should create a message with file chunks and RAG query ID', async () => {
+        // Create test data
+        const chunkId1 = uuid();
+        const chunkId2 = uuid();
+        const ragQueryId = uuid();
+
+        await serverDB.insert(chunks).values([
+          { id: chunkId1, text: 'chunk text 1' },
+          { id: chunkId2, text: 'chunk text 2' },
+        ]);
+
+        // Call create method
+        const result = await messageModel.create({
+          role: 'assistant',
+          content: 'message with file chunks',
+          fileChunks: [
+            { id: chunkId1, similarity: 0.95 },
+            { id: chunkId2, similarity: 0.85 },
+          ],
+          ragQueryId,
+          sessionId: '1',
+        });
+
+        // Verify message created successfully
+        expect(result.id).toBeDefined();
+
+        // Verify message query chunk associations created successfully
+        const queryChunks = await serverDB
+          .select()
+          .from(messageQueryChunks)
+          .where(eq(messageQueryChunks.messageId, result.id));
+
+        expect(queryChunks).toHaveLength(2);
+        expect(queryChunks[0].chunkId).toBe(chunkId1);
+        expect(queryChunks[0].queryId).toBe(ragQueryId);
+        expect(queryChunks[0].similarity).toBe('0.95');
+        expect(queryChunks[1].chunkId).toBe(chunkId2);
+        expect(queryChunks[1].similarity).toBe('0.85');
+      });
+
+      it('should create a message with files', async () => {
+        // Create test data
+        await serverDB.insert(files).values([
+          {
+            id: 'file1',
+            name: 'file1.txt',
+            fileType: 'text/plain',
+            size: 100,
+            url: 'url1',
+            userId,
+          },
+          {
+            id: 'file2',
+            name: 'file2.jpg',
+            fileType: 'image/jpeg',
+            size: 200,
+            url: 'url2',
+            userId,
+          },
+        ]);
+
+        // Call create method
+        const result = await messageModel.create({
+          role: 'user',
+          content: 'message with files',
+          files: ['file1', 'file2'],
+          sessionId: '1',
+        });
+
+        // Verify message created successfully
+        expect(result.id).toBeDefined();
+
+        // Verify message file associations created successfully
+        const messageFiles = await serverDB
+          .select()
+          .from(messagesFiles)
+          .where(eq(messagesFiles.messageId, result.id));
+
+        expect(messageFiles).toHaveLength(2);
+        expect(messageFiles[0].fileId).toBe('file1');
+        expect(messageFiles[1].fileId).toBe('file2');
+      });
+
+      it('should create a message with custom timestamps', async () => {
+        const customCreatedAt = '2022-05-15T10:30:00Z';
+        const customUpdatedAt = '2022-05-16T11:45:00Z';
+
+        const result = await messageModel.create({
+          role: 'user',
+          content: 'message with custom timestamps',
+          createdAt: customCreatedAt as any,
+          updatedAt: customUpdatedAt as any,
+          sessionId: '1',
+        });
+
+        // Verify database records
+        const dbResult = await serverDB.select().from(messages).where(eq(messages.id, result.id));
+
+        // Date comparison needs to consider timezone and formatting, so use toISOString for comparison
+        expect(new Date(dbResult[0].createdAt!).toISOString()).toBe(
+          new Date(customCreatedAt).toISOString(),
+        );
+        expect(new Date(dbResult[0].updatedAt!).toISOString()).toBe(
+          new Date(customUpdatedAt).toISOString(),
+        );
+      });
+    });
+  });
+
+  describe('batchCreateMessages', () => {
+    it('should batch create messages', async () => {
+      // Prepare test data
+      const newMessages = [
+        { id: '1', role: 'user', content: 'message 1' },
+        { id: '2', role: 'assistant', content: 'message 2' },
+      ] as DBMessageItem[];
+
+      // Call batchCreateMessages method
+      await messageModel.batchCreate(newMessages);
+
+      // Assert result
+      const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
+      expect(result).toHaveLength(2);
+      expect(result[0].content).toBe('message 1');
+      expect(result[1].content).toBe('message 2');
+    });
+  });
+
+  describe('createMessageQuery', () => {
+    it('should create a new message query', async () => {
+      // Create test data
+      await serverDB.insert(messages).values({
+        id: 'msg1',
+        userId,
+        role: 'user',
+        content: 'test message',
+      });
+
+      // 调用 createMessageQuery 方法
+      const result = await messageModel.createMessageQuery({
+        messageId: 'msg1',
+        userQuery: 'original query',
+        rewriteQuery: 'rewritten query',
+        embeddingsId,
+      });
+
+      // Assert result
+      expect(result).toBeDefined();
+      expect(result.id).toBeDefined();
+      expect(result.messageId).toBe('msg1');
+      expect(result.userQuery).toBe('original query');
+      expect(result.rewriteQuery).toBe('rewritten query');
+      expect(result.userId).toBe(userId);
+
+      // 验证数据库中的记录
+      const dbResult = await serverDB
+        .select()
+        .from(messageQueries)
+        .where(eq(messageQueries.id, result.id));
+
+      expect(dbResult).toHaveLength(1);
+      expect(dbResult[0].messageId).toBe('msg1');
+      expect(dbResult[0].userQuery).toBe('original query');
+      expect(dbResult[0].rewriteQuery).toBe('rewritten query');
+    });
+
+    it('should create a message query with embeddings ID', async () => {
+      // Create test data
+      await serverDB.insert(messages).values({
+        id: 'msg2',
+        userId,
+        role: 'user',
+        content: 'test message',
+      });
+
+      // 调用 createMessageQuery 方法
+      const result = await messageModel.createMessageQuery({
+        messageId: 'msg2',
+        userQuery: 'test query',
+        rewriteQuery: 'test rewritten query',
+        embeddingsId,
+      });
+
+      // Assert result
+      expect(result).toBeDefined();
+      expect(result.embeddingsId).toBe(embeddingsId);
+
+      // 验证数据库中的记录
+      const dbResult = await serverDB
+        .select()
+        .from(messageQueries)
+        .where(eq(messageQueries.id, result.id));
+
+      expect(dbResult[0].embeddingsId).toBe(embeddingsId);
+    });
+
+    it('should generate a unique ID for each message query', async () => {
+      // Create test data
+      await serverDB.insert(messages).values({
+        id: 'msg3',
+        userId,
+        role: 'user',
+        content: 'test message',
+      });
+
+      // 连续创建两个消息查询
+      const result1 = await messageModel.createMessageQuery({
+        messageId: 'msg3',
+        userQuery: 'query 1',
+        rewriteQuery: 'rewritten query 1',
+        embeddingsId,
+      });
+
+      const result2 = await messageModel.createMessageQuery({
+        messageId: 'msg3',
+        userQuery: 'query 2',
+        rewriteQuery: 'rewritten query 2',
+        embeddingsId,
+      });
+
+      // Assert result
+      expect(result1.id).not.toBe(result2.id);
+    });
+  });
+
+  describe('updateMessageRAG', () => {
+    it('should insert message query chunks for RAG', async () => {
+      // prepare message and query
+      const messageId = 'rag-msg-1';
+      const queryId = uuid();
+      const chunk1 = uuid();
+      const chunk2 = uuid();
+
+      await serverDB.transaction(async (trx) => {
+        await trx.insert(messages).values({ id: messageId, role: 'user', userId, content: 'c' });
+        await trx.insert(chunks).values([
+          { id: chunk1, text: 'a' },
+          { id: chunk2, text: 'b' },
+        ]);
+        await trx
+          .insert(messageQueries)
+          .values({ id: queryId, messageId, userId, userQuery: 'q', rewriteQuery: 'rq' });
+      });
+
+      await messageModel.updateMessageRAG(messageId, {
+        ragQueryId: queryId,
+        fileChunks: [
+          { id: chunk1, similarity: 0.9 },
+          { id: chunk2, similarity: 0.8 },
+        ],
+      });
+
+      const rows = await serverDB
+        .select()
+        .from(messageQueryChunks)
+        .where(eq(messageQueryChunks.messageId, messageId));
+
+      expect(rows).toHaveLength(2);
+      const s1 = rows.find((r) => r.chunkId === chunk1)!;
+      const s2 = rows.find((r) => r.chunkId === chunk2)!;
+      expect(s1.queryId).toBe(queryId);
+      expect(s1.similarity).toBe('0.90000');
+      expect(s2.similarity).toBe('0.80000');
+    });
+  });
+});