@lobehub/chat 0.105.2 → 0.107.0

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

@@ -0,0 +1,250 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { CreateMessageParams, MessageModel } from '@/database/models/message';
+
+import { CreateTopicParams, QueryTopicParams, TopicModel } from '../topic';
+
+describe('TopicModel', () => {
+  let topicData: CreateTopicParams;
+
+  beforeEach(() => {
+    // Set up topic data with the correct structure
+    topicData = {
+      sessionId: 'session1',
+      title: 'Test Topic',
+      favorite: false,
+    };
+  });
+
+  afterEach(async () => {
+    // Clean up the database after each test
+    await TopicModel.clearTable();
+  });
+  describe('create', () => {
+    it('should create a topic record', async () => {
+      const result = await TopicModel.create(topicData);
+
+      expect(result).toHaveProperty('id');
+      // Verify that the topic has been added to the database
+      const topicInDb = await TopicModel.findById(result.id);
+
+      expect(topicInDb).toEqual(
+        expect.objectContaining({
+          title: topicData.title,
+          favorite: topicData.favorite ? 1 : 0,
+          sessionId: topicData.sessionId,
+        }),
+      );
+    });
+
+    it('should create a topic with favorite set to true', async () => {
+      const favoriteTopicData: CreateTopicParams = {
+        ...topicData,
+        favorite: true,
+      };
+      const result = await TopicModel.create(favoriteTopicData);
+
+      expect(result).toHaveProperty('id');
+      const topicInDb = await TopicModel.findById(result.id);
+      expect(topicInDb).toEqual(
+        expect.objectContaining({
+          title: favoriteTopicData.title,
+          favorite: 1,
+          sessionId: favoriteTopicData.sessionId,
+        }),
+      );
+    });
+  });
+  describe('batch create', () => {
+    it('should batch create topic records', async () => {
+      const topicsToCreate = [topicData, topicData];
+      const results = await TopicModel.batchCreate(topicsToCreate);
+
+      expect(results.ids).toHaveLength(topicsToCreate.length);
+      // Verify that the topics have been added to the database
+      for (const result of results.ids!) {
+        const topicInDb = await TopicModel.findById(result);
+        expect(topicInDb).toEqual(
+          expect.objectContaining({
+            title: topicData.title,
+            favorite: topicData.favorite ? 1 : 0,
+            sessionId: topicData.sessionId,
+          }),
+        );
+      }
+    });
+
+    it('should batch create topics with mixed favorite values', async () => {
+      const mixedTopicsData: CreateTopicParams[] = [
+        { ...topicData, favorite: true },
+        { ...topicData, favorite: false },
+      ];
+
+      const results = await TopicModel.batchCreate(mixedTopicsData);
+
+      expect(results.ids).toHaveLength(mixedTopicsData.length);
+      for (const id of results.ids!) {
+        const topicInDb = await TopicModel.findById(id);
+        expect(topicInDb).toBeDefined();
+        expect(topicInDb.favorite).toBeGreaterThanOrEqual(0);
+        expect(topicInDb.favorite).toBeLessThanOrEqual(1);
+      }
+    });
+  });
+
+  it('should query topics with pagination', async () => {
+    // Create multiple topics to test the query method
+    await TopicModel.batchCreate([topicData, topicData]);
+
+    const queryParams: QueryTopicParams = { pageSize: 1, current: 0, sessionId: 'session1' };
+    const queriedTopics = await TopicModel.query(queryParams);
+
+    expect(queriedTopics).toHaveLength(1);
+  });
+
+  it('should find topics by session id', async () => {
+    // Create multiple topics to test the findBySessionId method
+    await TopicModel.batchCreate([topicData, topicData]);
+
+    const topicsBySessionId = await TopicModel.findBySessionId(topicData.sessionId);
+
+    expect(topicsBySessionId).toHaveLength(2);
+    expect(topicsBySessionId.every((i) => i.sessionId === topicData.sessionId)).toBeTruthy();
+  });
+
+  it('should delete a topic and its associated messages', async () => {
+    const createdTopic = await TopicModel.create(topicData);
+
+    await TopicModel.delete(createdTopic.id);
+
+    // Verify the topic and its related messages are deleted
+    const topicInDb = await TopicModel.findById(createdTopic.id);
+    expect(topicInDb).toBeUndefined();
+
+    // You need to verify that messages related to the topic are also deleted
+    // This would require additional setup to create messages associated with the topic
+    // and then assertions to check that they're deleted after the topic itself is deleted
+  });
+
+  it('should batch delete topics by session id', async () => {
+    // Create multiple topics to test the batchDeleteBySessionId method
+    await TopicModel.batchCreate([topicData, topicData]);
+
+    await TopicModel.batchDeleteBySessionId(topicData.sessionId);
+
+    // Verify that all topics with the given session id are deleted
+    const topicsInDb = await TopicModel.findBySessionId(topicData.sessionId);
+    expect(topicsInDb).toHaveLength(0);
+  });
+
+  it('should update a topic', async () => {
+    const createdTopic = await TopicModel.create(topicData);
+    const updateData = { title: 'New Title' };
+
+    await TopicModel.update(createdTopic.id, updateData);
+    const updatedTopic = await TopicModel.findById(createdTopic.id);
+
+    expect(updatedTopic).toHaveProperty('title', 'New Title');
+  });
+
+  describe('toggleFavorite', () => {
+    it('should toggle favorite status of a topic', async () => {
+      const createdTopic = await TopicModel.create(topicData);
+
+      const newState = await TopicModel.toggleFavorite(createdTopic.id);
+
+      expect(newState).toBe(true);
+      const topicInDb = await TopicModel.findById(createdTopic.id);
+      expect(topicInDb).toHaveProperty('favorite', 1);
+    });
+
+    it('should handle toggleFavorite when topic does not exist', async () => {
+      const nonExistentTopicId = 'non-existent-id';
+      await expect(TopicModel.toggleFavorite(nonExistentTopicId)).rejects.toThrow(
+        `Topic with id ${nonExistentTopicId} not found`,
+      );
+    });
+
+    it('should set favorite to specific state using toggleFavorite', async () => {
+      const createdTopic = await TopicModel.create(topicData);
+
+      // Set favorite to true regardless of current state
+      await TopicModel.toggleFavorite(createdTopic.id, true);
+      let topicInDb = await TopicModel.findById(createdTopic.id);
+      expect(topicInDb.favorite).toBe(1);
+
+      // Set favorite to false regardless of current state
+      await TopicModel.toggleFavorite(createdTopic.id, false);
+      topicInDb = await TopicModel.findById(createdTopic.id);
+      expect(topicInDb.favorite).toBe(0);
+    });
+  });
+
+  it('should delete a topic and its associated messages', async () => {
+    // 创建话题和相关联的消息
+    const createdTopic = await TopicModel.create(topicData);
+    const messageData: CreateMessageParams = {
+      content: 'Test Message',
+      topicId: createdTopic.id,
+      sessionId: topicData.sessionId,
+      role: 'user',
+    };
+    await MessageModel.create(messageData);
+
+    // 删除话题
+    await TopicModel.delete(createdTopic.id);
+
+    // 验证话题是否被删除
+    const topicInDb = await TopicModel.findById(createdTopic.id);
+    expect(topicInDb).toBeUndefined();
+
+    // 验证与话题关联的消息是否也被删除
+    const messagesInDb = await MessageModel.query({
+      sessionId: topicData.sessionId,
+      topicId: createdTopic.id,
+    });
+    expect(messagesInDb).toHaveLength(0);
+  });
+
+  it('should batch delete topics and their associated messages', async () => {
+    // 创建多个话题和相关联的消息
+    const createdTopic1 = await TopicModel.create(topicData);
+    const createdTopic2 = await TopicModel.create(topicData);
+
+    const messageData1: CreateMessageParams = {
+      content: 'Test Message 1',
+      topicId: createdTopic1.id,
+      sessionId: topicData.sessionId,
+      role: 'user',
+    };
+    const messageData2: CreateMessageParams = {
+      content: 'Test Message 2',
+      topicId: createdTopic2.id,
+      sessionId: topicData.sessionId,
+      role: 'user',
+    };
+    await MessageModel.create(messageData1);
+    await MessageModel.create(messageData2);
+
+    // 执行批量删除
+    await TopicModel.batchDelete([createdTopic1.id, createdTopic2.id]);
+
+    // 验证话题是否被删除
+    const topicInDb1 = await TopicModel.findById(createdTopic1.id);
+    const topicInDb2 = await TopicModel.findById(createdTopic2.id);
+    expect(topicInDb1).toBeUndefined();
+    expect(topicInDb2).toBeUndefined();
+
+    // 验证与话题关联的消息是否也被删除
+    const messagesInDb1 = await MessageModel.query({
+      sessionId: topicData.sessionId,
+      topicId: createdTopic1.id,
+    });
+    const messagesInDb2 = await MessageModel.query({
+      sessionId: topicData.sessionId,
+      topicId: createdTopic2.id,
+    });
+    expect(messagesInDb1).toHaveLength(0);
+    expect(messagesInDb2).toHaveLength(0);
+  });
+});

package/src/database/models/file.ts

@@ -1,17 +1,17 @@
-import { LocalFile, LocalFileSchema } from '@/types/database/files';
+import { DB_File, DB_FileSchema } from '@/database/schemas/files';
 import { nanoid } from '@/utils/uuid';
 
 import { BaseModel } from '../core';
 
-class _FileModel extends BaseModel {
+class _FileModel extends BaseModel<'files'> {
   constructor() {
-    super('files', LocalFileSchema);
+    super('files', DB_FileSchema);
   }
 
-  async create(file: LocalFile) {
+  async create(file: DB_File) {
     const id = nanoid();
 
-    return this.add(file, `file-${id}`);
+    return this._add(file, `file-${id}`);
   }
 
   async findById(id: string) {
@@ -21,6 +21,10 @@ class _FileModel extends BaseModel {
   async delete(id: string) {
     return this.table.delete(id);
   }
+
+  async clear() {
+    return this.table.clear();
+  }
 }
 
 export const FileModel = new _FileModel();

package/src/database/models/message.ts

@@ -0,0 +1,186 @@
+import { DeepPartial } from 'utility-types';
+
+import { BaseModel } from '@/database/core';
+import { DBModel } from '@/database/core/types/db';
+import { DB_Message, DB_MessageSchema } from '@/database/schemas/message';
+import { ChatMessage } from '@/types/chatMessage';
+import { nanoid } from '@/utils/uuid';
+
+export interface CreateMessageParams
+  extends Partial<Omit<ChatMessage, 'content' | 'role'>>,
+    Pick<ChatMessage, 'content' | 'role'> {
+  sessionId: string;
+}
+
+export interface QueryMessageParams {
+  current?: number;
+  pageSize?: number;
+  sessionId: string;
+  topicId?: string;
+}
+
+class _MessageModel extends BaseModel {
+  constructor() {
+    super('messages', DB_MessageSchema);
+  }
+  async create(data: CreateMessageParams) {
+    const id = nanoid();
+
+    const messageData: DB_Message = this.mapChatMessageToDBMessage(data as ChatMessage);
+
+    return this._add(messageData, id);
+  }
+
+  async batchCreate(messages: ChatMessage[]) {
+    const data: DB_Message[] = messages.map((m) => this.mapChatMessageToDBMessage(m));
+
+    return this._batchAdd(data);
+  }
+
+  async query({
+    sessionId,
+    topicId,
+    pageSize = 9999,
+    current = 0,
+  }: QueryMessageParams): Promise<ChatMessage[]> {
+    const offset = current * pageSize;
+
+    const query =
+      topicId !== undefined
+        ? // TODO: The query {"sessionId":"xxx","topicId":"xxx"} on messages would benefit of a compound index [sessionId+topicId]
+          this.table.where({ sessionId, topicId }) // Use a compound index
+        : this.table
+            .where('sessionId')
+            .equals(sessionId)
+            .and((message) => !message.topicId);
+
+    const dbMessages: DBModel<DB_Message>[] = await query
+      .sortBy('createdAt')
+      // handle page size
+      .then((sortedArray) => sortedArray.slice(offset, offset + pageSize));
+
+    const messages = dbMessages.map((msg) => this.mapToChatMessage(msg));
+
+    const finalList: ChatMessage[] = [];
+
+    const addItem = (item: ChatMessage) => {
+      const isExist = finalList.findIndex((i) => item.id === i.id) > -1;
+      if (!isExist) {
+        finalList.push(item);
+      }
+    };
+    const messageMap = new Map<string, ChatMessage>();
+    for (const item of messages) messageMap.set(item.id, item);
+
+    for (const item of messages) {
+      if (!item.parentId || !messageMap.has(item.parentId)) {
+        // 如果消息没有父消息或者父消息不在列表中，直接添加
+        addItem(item);
+      } else {
+        // 如果消息有父消息，确保先添加父消息
+        addItem(messageMap.get(item.parentId)!);
+        addItem(item);
+      }
+    }
+    return finalList;
+  }
+
+  async findById(id: string): Promise<DBModel<DB_Message>> {
+    return this.table.get(id);
+  }
+
+  async delete(id: string) {
+    return this.table.delete(id);
+  }
+
+  async clearTable() {
+    return this.table.clear();
+  }
+
+  async update(id: string, data: DeepPartial<DB_Message>) {
+    return super._update(id, data);
+  }
+
+  /**
+   * Batch updates multiple fields of the specified messages.
+   *
+   * @param {string[]} messageIds - The identifiers of the messages to be updated.
+   * @param {Partial<DB_Message>} updateFields - An object containing the fields to update and their new values.
+   * @returns {Promise<number>} - The number of updated messages.
+   */
+  async batchUpdate(messageIds: string[], updateFields: Partial<DB_Message>): Promise<number> {
+    // Retrieve the messages by their IDs
+    const messagesToUpdate = await this.table.where(':id').anyOf(messageIds).toArray();
+
+    // Update the specified fields of each message
+    const updatedMessages = messagesToUpdate.map((message) => ({
+      ...message,
+      ...updateFields,
+    }));
+
+    // Use the bulkPut method to update the messages in bulk
+    await this.table.bulkPut(updatedMessages);
+
+    return updatedMessages.length;
+  }
+
+  /**
+   * Deletes multiple messages based on the assistantId and optionally the topicId.
+   * If topicId is not provided, it deletes messages where topicId is undefined.
+   * If topicId is provided, it deletes messages with that specific topicId.
+   *
+   * @param {string} sessionId - The identifier of the assistant associated with the messages.
+   * @param {string | undefined} topicId - The identifier of the topic associated with the messages (optional).
+   * @returns {Promise<void>}
+   */
+  async batchDelete(sessionId: string, topicId: string | undefined): Promise<void> {
+    // If topicId is specified, use both assistantId and topicId as the filter criteria in the query.
+    // Otherwise, filter by assistantId and require that topicId is undefined.
+    const query =
+      topicId !== undefined
+        ? this.table.where({ sessionId, topicId }) // Use a compound index
+        : this.table
+            .where('sessionId')
+            .equals(sessionId)
+            .and((message) => message.topicId === undefined);
+
+    // Retrieve a collection of message IDs that satisfy the criteria
+    const messageIds = await query.primaryKeys();
+
+    // Use the bulkDelete method to delete all selected messages in bulk
+    return this.table.bulkDelete(messageIds);
+  }
+
+  async queryAll() {
+    const data: DBModel<DB_Message>[] = await this.table.orderBy('updatedAt').toArray();
+
+    return data.map((element) => this.mapToChatMessage(element));
+  }
+
+  async isEmpty() {
+    const count = await this.table.count();
+
+    return count === 0;
+  }
+
+  private mapChatMessageToDBMessage(message: ChatMessage): DB_Message {
+    const { extra, ...messageData } = message;
+
+    return { ...messageData, ...extra } as DB_Message;
+  }
+
+  private mapToChatMessage = ({
+    fromModel,
+    translate,
+    tts,
+    ...item
+  }: DBModel<DB_Message>): ChatMessage => {
+    return {
+      ...item,
+      extra: { fromModel: fromModel, translate: translate, tts: tts },
+      meta: {},
+    };
+  };
+}
+
+export const MessageModel = new _MessageModel();

package/src/database/models/session.ts

@@ -0,0 +1,160 @@
+import { DeepPartial } from 'utility-types';
+
+import { DEFAULT_AGENT_LOBE_SESSION } from '@/const/session';
+import { BaseModel } from '@/database/core';
+import { DB_Session, DB_SessionSchema } from '@/database/schemas/session';
+import { LobeAgentConfig } from '@/types/agent';
+import { LobeAgentSession, LobeSessions, SessionGroupKey } from '@/types/session';
+import { merge } from '@/utils/merge';
+import { uuid } from '@/utils/uuid';
+
+class _SessionModel extends BaseModel {
+  constructor() {
+    super('sessions', DB_SessionSchema);
+  }
+
+  async create(type: 'agent' | 'group', defaultValue: Partial<LobeAgentSession>, id = uuid()) {
+    const data = merge(DEFAULT_AGENT_LOBE_SESSION, { type, ...defaultValue });
+    return this._add(data, id);
+  }
+
+  async batchCreate(sessions: LobeAgentSession[]) {
+    return this._batchAdd(sessions, { idGenerator: uuid });
+  }
+
+  async query({
+    pageSize = 9999,
+    current = 0,
+  }: { current?: number; pageSize?: number } = {}): Promise<LobeSessions> {
+    const offset = current * pageSize;
+
+    return this.table.orderBy('updatedAt').reverse().offset(offset).limit(pageSize).toArray();
+  }
+
+  /**
+   * get sessions by group
+   * @param group
+   */
+  async queryByGroup(group: SessionGroupKey) {
+    return this.table.where('group').equals(group).toArray();
+  }
+
+  async update(id: string, data: Partial<DB_Session>) {
+    return super._update(id, data);
+  }
+
+  async updateConfig(id: string, data: DeepPartial<LobeAgentConfig>) {
+    const session = await this.findById(id);
+    if (!session) return;
+
+    const config = merge(session.config, data);
+
+    return this.update(id, { config });
+  }
+  /**
+   * Delete a session , also delete all messages and topic associated with it.
+   */
+  async delete(id: string) {
+    return this.db.transaction('rw', [this.table, this.db.topics, this.db.messages], async () => {
+      // Delete all topics associated with the session
+      const topics = await this.db.topics.where('sessionId').equals(id).toArray();
+      const topicIds = topics.map((topic) => topic.id);
+      if (topicIds.length > 0) {
+        await this.db.topics.bulkDelete(topicIds);
+      }
+
+      // Delete all messages associated with the session
+      const messages = await this.db.messages.where('sessionId').equals(id).toArray();
+      const messageIds = messages.map((message) => message.id);
+      if (messageIds.length > 0) {
+        await this.db.messages.bulkDelete(messageIds);
+      }
+
+      // Finally, delete the session itself
+      await this.table.delete(id);
+    });
+  }
+
+  async clearTable() {
+    return this.table.clear();
+  }
+
+  async findById(id: string) {
+    return this.table.get(id);
+  }
+
+  async isEmpty() {
+    return (await this.table.count()) === 0;
+  }
+
+  /**
+   * Query sessions by keyword in title, description, content, or translated content
+   * @param keyword The keyword to search for
+   */
+  async queryByKeyword(keyword: string): Promise<LobeSessions> {
+    if (!keyword) return [];
+
+    console.time('queryByKeyword');
+    const keywordLowerCase = keyword.toLowerCase();
+
+    // First, filter sessions by title and description
+    const matchingSessionsPromise = this.table
+      .filter((session) => {
+        return (
+          session.meta.title.toLowerCase().includes(keywordLowerCase) ||
+          session.meta.description.toLowerCase().includes(keywordLowerCase)
+        );
+      })
+      .toArray();
+
+    // Next, find message IDs that contain the keyword in content or translated content
+    const matchingMessagesPromise = this.db.messages
+      .filter((message) => {
+        // check content
+        if (message.content.toLowerCase().includes(keywordLowerCase)) return true;
+
+        // check translate content
+        if (message.translate && message.translate.content) {
+          return message.translate.content.toLowerCase().includes(keywordLowerCase);
+        }
+
+        return false;
+      })
+      .toArray();
+
+    //  match topics
+    const matchingTopicsPromise = this.db.topics
+      .filter((topic) => {
+        return topic.title.toLowerCase().includes(keywordLowerCase);
+      })
+      .toArray();
+
+    // Resolve both promises
+    const [matchingSessions, matchingMessages, matchingTopics] = await Promise.all([
+      matchingSessionsPromise,
+      matchingMessagesPromise,
+      matchingTopicsPromise,
+    ]);
+
+    const sessionIdsFromMessages = matchingMessages.map((message) => message.sessionId);
+    const sessionIdsFromTopics = matchingTopics.map((topic) => topic.sessionId);
+
+    // Combine session IDs from both sources
+    const combinedSessionIds = new Set([
+      ...sessionIdsFromMessages,
+      ...sessionIdsFromTopics,
+      ...matchingSessions.map((session) => session.id),
+    ]);
+
+    // Retrieve unique sessions by IDs
+    const data = await this.table
+      .where('id')
+      .anyOf([...combinedSessionIds])
+      .toArray();
+
+    console.timeEnd('queryByKeyword');
+    return data;
+  }
+}
+
+export const SessionModel = new _SessionModel();