@lobehub/chat 0.105.2 → 0.107.0

package/src/database/core/types/db.ts

@@ -0,0 +1,13 @@
+import { z } from 'zod';
+
+export type DBModel<T> = T & {
+  createdAt: number;
+  id: string;
+  updatedAt: number;
+};
+
+export const DBBaseFieldsSchema = z.object({
+  createdAt: z.number(),
+  id: z.string(),
+  updatedAt: z.number(),
+});

package/src/database/models/__DEBUG.ts

@@ -0,0 +1,119 @@
+// This file is for debugging purposes only.
+// DON'T USE IT IN PRODUCTION.
+import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
+import { DBModel } from '@/database/core/types/db';
+import { DB_Message } from '@/database/schemas/message';
+import { DB_Topic } from '@/database/schemas/topic';
+
+import { BaseModel } from '../core';
+import { DB_Session, DB_SessionSchema } from '../schemas/session';
+
+class _DEBUG_MODEL extends BaseModel<'sessions'> {
+  constructor() {
+    super('sessions', DB_SessionSchema);
+  }
+  private getRandomInt(min: number, max: number) {
+    return Math.floor(Math.random() * (max - min + 1)) + min;
+  }
+
+  private randomString(length: number) {
+    const characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+    let result = '';
+    for (let i = 0; i < length; i++) {
+      result += characters.charAt(Math.floor(Math.random() * characters.length));
+    }
+    return result;
+  }
+
+  private randomDate(start: Date, end: Date) {
+    return new Date(start.getTime() + Math.random() * (end.getTime() - start.getTime())).getTime();
+  }
+
+  private randomPick<T>(array: T[]): T {
+    const randomIndex = this.getRandomInt(0, array.length - 1);
+    return array[randomIndex];
+  }
+
+  createRandomData = async ({ sessionCount = 10, topicCount = 2000, messageCount = 10_000 }) => {
+    const numberOfSessions = sessionCount;
+    const numberOfTopics = topicCount;
+    const numberOfMessages = messageCount;
+
+    // Prepare data for batch inserts
+    const sessionsData: DBModel<DB_Session>[] = [];
+    const topicsData: DBModel<DB_Topic>[] = [];
+    const messagesData: DBModel<DB_Message>[] = [];
+
+    // Prepare sessions
+    for (let i = 1; i <= numberOfSessions; i++) {
+      sessionsData.push({
+        config: DEFAULT_AGENT_CONFIG,
+        createdAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+        group: 'default',
+        id: `sess_${i}`,
+        meta: {
+          description: `Session Description ${i}`,
+          title: `Session Title ${i}`,
+        },
+        type: 'agent',
+        updatedAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+      });
+    }
+
+    // Prepare topics
+    for (let i = 1; i <= numberOfTopics; i++) {
+      topicsData.push({
+        createdAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+        favorite: this.getRandomInt(0, 1),
+        id: `topic_${i}`,
+        sessionId: `sess_${this.getRandomInt(1, numberOfSessions)}`,
+        title: `Topic Title ${i}`,
+        updatedAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+      });
+    }
+
+    // Prepare messages
+    for (let i = 1; i <= numberOfMessages; i++) {
+      messagesData.push({
+        content: this.randomString(300),
+        createdAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+        favorite: this.getRandomInt(0, 1),
+        fromModel: 'model',
+        id: `msg_${i}`,
+        parentId: `msg_${this.getRandomInt(1, numberOfMessages)}`,
+        quotaId: `msg_${this.getRandomInt(1, numberOfMessages)}`,
+        role: this.randomPick(['user', 'assistant']),
+        sessionId: `sess_${this.getRandomInt(1, numberOfSessions)}`,
+        topicId: `topic_${this.getRandomInt(1, numberOfTopics)}`,
+        updatedAt: this.randomDate(new Date(2020, 0, 1), new Date()),
+      });
+    }
+
+    // Start a transaction for batch inserts
+    await this.db.transaction(
+      'rw',
+      this.db.sessions,
+      this.db.topics,
+      this.db.messages,
+      async () => {
+        // Batch insert sessions, topics, and messages
+        console.log('开始插入 sessions');
+        console.time('插入sessions');
+        await this.db.sessions.bulkAdd(sessionsData);
+        console.timeEnd('插入sessions');
+
+        console.log('开始插入 topics');
+        console.time('插入topics');
+        await this.db.topics.bulkAdd(topicsData);
+        console.timeEnd('插入topics');
+
+        console.log('开始插入 messages');
+        console.time('插入messages');
+        await this.db.messages.bulkAdd(messagesData);
+        console.timeEnd('插入messages');
+      },
+    );
+  };
+}
+
+export const DEBUG_MODEL = new _DEBUG_MODEL();

package/src/database/models/{file.test.ts → __tests__/file.test.ts}

@@ -1,15 +1,15 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
-import { LocalFile } from '@/types/database/files';
+import { DB_File } from '@/database/schemas/files';
 
-import { LocalDB } from '../core/db';
-import { FileModel } from './file';
+import { LocalDB } from '../../core/db';
+import { FileModel } from '../file';
 
 // Assuming LocalDB is already mocked or using an in-memory database
 // and LocalFileSchema has been imported correctly.
 
 describe('_FileModel', () => {
-  let fileData: LocalFile;
+  let fileData: DB_File;
 
   beforeEach(() => {
     // Set up file data with the correct structure according to LocalFileSchema
@@ -32,7 +32,7 @@ describe('_FileModel', () => {
 
   it('should create a file record', async () => {
     // First, create a file to test the create method
-    const fileData: LocalFile = {
+    const fileData: DB_File = {
       data: new ArrayBuffer(10),
       fileType: 'image/png',
       name: 'test.png',
@@ -68,4 +68,17 @@ describe('_FileModel', () => {
     const fileInDb = await FileModel.findById(createdFile.id);
     expect(fileInDb).toBeUndefined();
   });
+
+  it('should clear all files', async () => {
+    // First, create a file to test the delete method
+    const createdFile = await FileModel.create(fileData);
+    const createdFile2 = await FileModel.create(fileData);
+    await FileModel.clear();
+
+    // Verify that the file has been removed from the database
+    const fileInDb = await FileModel.findById(createdFile.id);
+    expect(fileInDb).toBeUndefined();
+    const fileInDb2 = await FileModel.findById(createdFile2.id);
+    expect(fileInDb2).toBeUndefined();
+  });
 });

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

@@ -0,0 +1,314 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { ChatMessage } from '@/types/chatMessage';
+
+import { CreateMessageParams, MessageModel } from '../message';
+
+describe('MessageModel', () => {
+  let messageData: CreateMessageParams;
+
+  beforeEach(() => {
+    // 设置正确结构的消息数据
+    messageData = {
+      content: 'Test message content',
+      role: 'user',
+      sessionId: 'session1',
+      topicId: 'topic1',
+    };
+  });
+
+  afterEach(async () => {
+    // 每次测试后清理数据库
+    await MessageModel.clearTable();
+  });
+
+  describe('create', () => {
+    it('should create a message record', async () => {
+      const result = await MessageModel.create(messageData);
+
+      expect(result).toHaveProperty('id');
+      // 验证消息是否已添加到数据库
+      const messageInDb = await MessageModel.findById(result.id);
+
+      expect(messageInDb).toEqual(
+        expect.objectContaining({
+          content: messageData.content,
+          role: messageData.role,
+          sessionId: messageData.sessionId,
+          topicId: messageData.topicId,
+        }),
+      );
+    });
+
+    it('should create with tts', async () => {
+      const result = await MessageModel.create({
+        content: 'abc',
+        role: 'assistant',
+        extra: { translate: { content: 'avc', from: 'a', to: 'f' } },
+        sessionId: 'a',
+      });
+
+      // 验证消息是否已添加到数据库
+      const messageInDb = await MessageModel.findById(result.id);
+
+      expect(messageInDb).toEqual(
+        expect.objectContaining({
+          content: 'abc',
+          role: 'assistant',
+          translate: { content: 'avc', from: 'a', to: 'f' },
+          sessionId: 'a',
+        }),
+      );
+    });
+  });
+
+  describe('batchCreate', () => {
+    it('should batch create message records', async () => {
+      const messagesToCreate = [messageData, messageData] as ChatMessage[];
+      const results = await MessageModel.batchCreate(messagesToCreate);
+
+      expect(results.success).toBeTruthy();
+      expect(results.errors).toBeUndefined();
+
+      // 验证消息是否已添加到数据库
+      for (const message of results.ids!) {
+        const messageInDb = await MessageModel.findById(message);
+        expect(messageInDb).toEqual(
+          expect.objectContaining({
+            content: messageData.content,
+            role: messageData.role,
+            sessionId: messageData.sessionId,
+            topicId: messageData.topicId,
+          }),
+        );
+      }
+    });
+  });
+
+  describe('query', () => {
+    it('should query messages with pagination', async () => {
+      // 创建多条消息以测试查询方法
+      await MessageModel.batchCreate([messageData, messageData] as ChatMessage[]);
+
+      const queriedMessages = await MessageModel.query({
+        pageSize: 1,
+        current: 0,
+        sessionId: messageData.sessionId,
+        topicId: messageData.topicId,
+      });
+
+      expect(queriedMessages).toHaveLength(1);
+    });
+
+    it('should query correctly without topic id', async () => {
+      // 创建多条消息以测试查询方法
+      await MessageModel.batchCreate([messageData, messageData] as ChatMessage[]);
+
+      const queriedMessages = await MessageModel.query({ sessionId: messageData.sessionId });
+
+      expect(queriedMessages).toHaveLength(0);
+    });
+
+    it('should query correctly with exactly topic id', async () => {
+      // 创建多条消息以测试查询方法
+      await MessageModel.batchCreate([
+        messageData,
+        { ...messageData, topicId: undefined },
+      ] as ChatMessage[]);
+
+      const queriedMessages = await MessageModel.query({ sessionId: messageData.sessionId });
+
+      expect(queriedMessages).toHaveLength(1);
+    });
+
+    it('should should have correct order', async () => {
+      const data: ChatMessage[] = [
+        {
+          role: 'user',
+          content: '1',
+          createdAt: 1697120044345,
+          id: 'NQ7RscYx',
+          updatedAt: 1697120181827,
+          extra: {},
+          meta: {},
+          sessionId: '1',
+        },
+        {
+          role: 'assistant',
+          content: '2',
+          parentId: 'NQ7RscYx',
+          createdAt: 1697120130973,
+          id: '9tDAumEx',
+          updatedAt: 1697120181827,
+          meta: {},
+          extra: {
+            fromModel: 'gpt-3.5-turbo-16k',
+          },
+          sessionId: '1',
+        },
+        {
+          role: 'assistant',
+          content: '3',
+          parentId: 'tOMH7c5R',
+          meta: {},
+          createdAt: 1697120163272,
+          id: '5Ie5hClg',
+          updatedAt: 1697120181827,
+          extra: {
+            fromModel: 'gpt-3.5-turbo-16k',
+          },
+          sessionId: '1',
+        },
+        {
+          role: 'user',
+          content: '4',
+          meta: {},
+          createdAt: 1697120163272,
+          id: 'tOMH7c5R',
+          updatedAt: 1697120181827,
+          extra: {},
+          sessionId: '1',
+        },
+      ];
+
+      await MessageModel.batchCreate(data);
+
+      const queriedMessages = await MessageModel.query({ sessionId: '1' });
+
+      expect(queriedMessages).toEqual([
+        {
+          role: 'user',
+          content: '1',
+          createdAt: 1697120044345,
+          id: 'NQ7RscYx',
+          updatedAt: 1697120181827,
+          sessionId: '1',
+          extra: {},
+          meta: {},
+        },
+        {
+          role: 'assistant',
+          content: '2',
+          parentId: 'NQ7RscYx',
+          createdAt: 1697120130973,
+          id: '9tDAumEx',
+          sessionId: '1',
+          updatedAt: 1697120181827,
+          meta: {},
+          extra: {
+            fromModel: 'gpt-3.5-turbo-16k',
+          },
+        },
+        {
+          role: 'user',
+          content: '4',
+          sessionId: '1',
+          createdAt: 1697120163272,
+          id: 'tOMH7c5R',
+          updatedAt: 1697120181827,
+          meta: {},
+          extra: {},
+        },
+        {
+          role: 'assistant',
+          content: '3',
+          parentId: 'tOMH7c5R',
+          meta: {},
+          createdAt: 1697120163272,
+          sessionId: '1',
+          id: '5Ie5hClg',
+          updatedAt: 1697120181827,
+          extra: {
+            fromModel: 'gpt-3.5-turbo-16k',
+          },
+        },
+      ]);
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a message by id', async () => {
+      const createdMessage = await MessageModel.create(messageData);
+      const messageInDb = await MessageModel.findById(createdMessage.id);
+
+      expect(messageInDb).toEqual(
+        expect.objectContaining({
+          id: createdMessage.id,
+          content: messageData.content,
+        }),
+      );
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete a message', async () => {
+      const createdMessage = await MessageModel.create(messageData);
+      await MessageModel.delete(createdMessage.id);
+
+      const messageInDb = await MessageModel.findById(createdMessage.id);
+      expect(messageInDb).toBeUndefined();
+    });
+  });
+
+  describe('update', () => {
+    it('should update a message', async () => {
+      const createdMessage = await MessageModel.create(messageData);
+      const updateData = { content: 'Updated content' };
+
+      await MessageModel.update(createdMessage.id, updateData);
+      const updatedMessage = await MessageModel.findById(createdMessage.id);
+
+      expect(updatedMessage).toHaveProperty('content', 'Updated content');
+    });
+  });
+
+  describe('batchUpdate', () => {
+    it('should batch update messages', async () => {
+      const createdMessage1 = await MessageModel.create(messageData);
+      const createdMessage2 = await MessageModel.create(messageData);
+      const updateData = { content: 'Batch updated content' };
+
+      const numUpdated = await MessageModel.batchUpdate(
+        [createdMessage1.id, createdMessage2.id],
+        updateData,
+      );
+
+      expect(numUpdated).toBe(2);
+
+      const updatedMessage1 = await MessageModel.findById(createdMessage1.id);
+      const updatedMessage2 = await MessageModel.findById(createdMessage2.id);
+
+      expect(updatedMessage1).toHaveProperty('content', 'Batch updated content');
+      expect(updatedMessage2).toHaveProperty('content', 'Batch updated content');
+    });
+  });
+
+  describe('batchDelete', () => {
+    it('should batch delete messages by session id', async () => {
+      // 创建多条消息以测试批量删除方法
+      const createdMessage1 = await MessageModel.create(messageData);
+      const createdMessage2 = await MessageModel.create(messageData);
+
+      await MessageModel.batchDelete(messageData.sessionId, undefined);
+
+      // 验证所有具有给定会话 ID 的消息是否已删除
+      const messagesInDb = await MessageModel.query({ sessionId: messageData.sessionId });
+      expect(messagesInDb).toHaveLength(0);
+    });
+
+    it('should batch delete messages by session id and topic id', async () => {
+      // 创建多条消息以测试批量删除方法
+      const createdMessage1 = await MessageModel.create(messageData);
+      const createdMessage2 = await MessageModel.create(messageData);
+
+      await MessageModel.batchDelete(messageData.sessionId, messageData.topicId);
+
+      // 验证所有具有给定会话 ID 和话题 ID 的消息是否已删除
+      const messagesInDb = await MessageModel.query({
+        sessionId: messageData.sessionId,
+        topicId: messageData.topicId,
+      });
+      expect(messagesInDb).toHaveLength(0);
+    });
+  });
+});

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

@@ -0,0 +1,117 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
+import { CreateMessageParams, MessageModel } from '@/database/models/message';
+import { TopicModel } from '@/database/models/topic';
+import { LobeAgentSession, LobeSessionType, SessionGroupKey } from '@/types/session';
+
+import { SessionModel } from '../session';
+
+describe('SessionModel', () => {
+  let sessionData: Partial<LobeAgentSession>;
+
+  beforeEach(() => {
+    // Set up session data with the correct structure
+    sessionData = {
+      type: LobeSessionType.Agent,
+      group: 'testGroup',
+      meta: {},
+      config: DEFAULT_AGENT_CONFIG,
+      // ... other properties based on LobeAgentSession
+    };
+  });
+
+  afterEach(async () => {
+    // Clean up the database after each test
+    await SessionModel.clearTable();
+  });
+
+  it('should create a session record', async () => {
+    const result = await SessionModel.create('agent', sessionData);
+
+    expect(result).toHaveProperty('id');
+    // Verify that the session has been added to the database
+    // Assuming findById is a method that retrieves a session by ID
+    const sessionInDb = await SessionModel.findById(result.id);
+
+    expect(sessionInDb).toEqual(expect.objectContaining(sessionData));
+  });
+
+  it('should batch create session records', async () => {
+    const sessionsToCreate = [sessionData, sessionData];
+    const results = await SessionModel.batchCreate(sessionsToCreate as LobeAgentSession[]);
+
+    expect(results.ids).toHaveLength(sessionsToCreate.length);
+    // Verify that the sessions have been added to the database
+    for (const result of results.ids!) {
+      const sessionInDb = await SessionModel.findById(result);
+      expect(sessionInDb).toEqual(expect.objectContaining(sessionData));
+    }
+  });
+
+  it('should query sessions with pagination', async () => {
+    // Create multiple sessions to test the query method
+    await SessionModel.batchCreate([sessionData, sessionData] as LobeAgentSession[]);
+
+    const queriedSessions = await SessionModel.query({ pageSize: 1, current: 0 });
+
+    expect(queriedSessions).toHaveLength(1);
+  });
+
+  it('should query sessions by group', async () => {
+    // Create multiple sessions to test the queryByGroup method
+    const group: SessionGroupKey = 'testGroup';
+    await SessionModel.batchCreate([sessionData, sessionData] as LobeAgentSession[]);
+
+    const sessionsByGroup = await SessionModel.queryByGroup(group);
+
+    // Assuming all created sessions belong to the same group
+    expect(sessionsByGroup).toHaveLength(2);
+    expect(sessionsByGroup.every((i) => i.group === group)).toBeTruthy();
+  });
+
+  it('should update a session', async () => {
+    const createdSession = await SessionModel.create('agent', sessionData);
+    const updateData = { group: 'newGroup' };
+
+    await SessionModel.update(createdSession.id, updateData);
+    const updatedSession = await SessionModel.findById(createdSession.id);
+
+    expect(updatedSession).toHaveProperty('group', 'newGroup');
+  });
+
+  // 删除一个 session 时，也需要同步删除具有 sessionId 的 topic 和 message
+  it('should delete a session and its associated data', async () => {
+    // create session , topic and message test data
+    const { id: sessionId } = await SessionModel.create('agent', sessionData);
+
+    const topicData = {
+      title: 'Test Topic',
+      sessionId: sessionId,
+      favorite: false,
+    };
+    const createdTopic = await TopicModel.create(topicData);
+
+    const messageData: CreateMessageParams = {
+      content: 'Test Message',
+      sessionId: sessionId,
+      topicId: createdTopic.id,
+      role: 'user',
+    };
+    await MessageModel.create(messageData);
+
+    await SessionModel.delete(sessionId);
+
+    // Verify the session and its related data (topics, messages) are deleted
+    const sessionInDb = await SessionModel.findById(sessionId);
+    expect(sessionInDb).toBeUndefined();
+
+    // You need to verify that topics and messages related to the session are also deleted
+    const topicsInDb = await TopicModel.findBySessionId(sessionId);
+    expect(topicsInDb).toHaveLength(0);
+
+    // Verify all associated messages are deleted
+    const messagesInDb = await MessageModel.query({ sessionId });
+    expect(messagesInDb).toHaveLength(0);
+  });
+});