@lobehub/lobehub 2.0.0-next.12 → 2.0.0-next.14

package/src/services/message/client.test.ts

@@ -1,410 +0,0 @@
-import {
-  ChatMessageError,
-  ChatTTS,
-  ChatTranslate,
-  CreateMessageParams,
-  DBMessageItem,
-  UIChatMessage,
-} from '@lobechat/types';
-import { and, eq } from 'drizzle-orm';
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-
-import { clientDB, initializeDB } from '@/database/client/db';
-import {
-  files,
-  messagePlugins,
-  messageTTS,
-  messageTranslates,
-  messages,
-  sessions,
-  topics,
-  users,
-} from '@/database/schemas';
-
-import { ClientService } from './client';
-
-const userId = 'message-db';
-const sessionId = '1';
-const topicId = 'topic-id';
-
-// Mock data
-const mockMessageId = 'mock-message-id';
-const mockMessage = {
-  id: mockMessageId,
-  content: 'Mock message content',
-  sessionId,
-  role: 'user',
-} as UIChatMessage;
-
-const mockMessages = [mockMessage];
-
-beforeEach(async () => {
-  await initializeDB();
-
-  // 在每个测试用例之前，清空表
-  await clientDB.transaction(async (trx) => {
-    await trx.delete(users);
-    await trx.insert(users).values([{ id: userId }, { id: '456' }]);
-
-    await trx.insert(sessions).values([{ id: sessionId, userId }]);
-    await trx.insert(topics).values([{ id: topicId, sessionId, userId }]);
-    await trx.insert(files).values({
-      id: 'f1',
-      userId: userId,
-      url: 'abc',
-      name: 'file-1',
-      fileType: 'image/png',
-      size: 1000,
-    });
-  });
-});
-
-afterEach(async () => {
-  // 在每个测试用例之后，清空表
-  await clientDB.delete(users);
-});
-
-const messageService = new ClientService(userId);
-
-describe('MessageClientService', () => {
-  describe('create', () => {
-    it('should create a message and return its id', async () => {
-      // Setup
-      const createParams: CreateMessageParams = {
-        content: 'New message content',
-        sessionId,
-        role: 'user',
-      };
-
-      // Execute
-      const messageId = await messageService.createMessage(createParams);
-
-      // Assert
-      expect(messageId).toMatch(/^msg_/);
-    });
-  });
-
-  describe('batchCreate', () => {
-    it('should batch create messages', async () => {
-      // Execute
-      await messageService.batchCreateMessages([
-        {
-          content: 'Mock message content',
-          sessionId,
-          role: 'user',
-        },
-        {
-          content: 'Mock message content',
-          sessionId,
-          role: 'user',
-        },
-      ] as DBMessageItem[]);
-      const count = await clientDB.$count(messages);
-
-      // Assert
-      expect(count).toBe(2);
-    });
-  });
-
-  describe('removeMessage', () => {
-    it('should remove a message by id', async () => {
-      // Execute
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      await messageService.removeMessage(mockMessageId);
-
-      // Assert
-      const count = await clientDB.$count(messages);
-
-      expect(count).toBe(0);
-    });
-  });
-  describe('removeMessages', () => {
-    it('should remove a message by id', async () => {
-      // Setup
-      await clientDB.insert(messages).values([
-        { id: mockMessageId, role: 'user', userId },
-        { role: 'assistant', userId },
-      ]);
-
-      // Execute
-      await messageService.removeMessages([mockMessageId]);
-
-      // Assert
-      const count = await clientDB.$count(messages);
-
-      expect(count).toBe(1);
-    });
-  });
-
-  describe('getMessages', () => {
-    it('should retrieve messages by sessionId and topicId', async () => {
-      // Setup
-      await clientDB
-        .insert(messages)
-        .values({ id: mockMessageId, sessionId, topicId, role: 'user', userId });
-
-      // Execute
-      const data = await messageService.getMessages(sessionId, topicId);
-
-      // Assert
-      expect(data[0]).toMatchObject({ id: mockMessageId, role: 'user' });
-    });
-  });
-
-  describe('getAllMessagesInSession', () => {
-    it('should retrieve all messages in a session', async () => {
-      // Setup
-      const sessionId = 'session-id';
-      await clientDB.insert(sessions).values([
-        { id: 'bbb', userId },
-        { id: sessionId, userId },
-      ]);
-      await clientDB.insert(messages).values([
-        { sessionId, topicId, role: 'user', userId },
-        { sessionId, topicId, role: 'assistant', userId },
-        { sessionId: 'bbb', topicId, role: 'assistant', userId },
-      ]);
-
-      // Execute
-      const data = await messageService.getAllMessagesInSession(sessionId);
-
-      // Assert
-      expect(data.length).toBe(2);
-    });
-  });
-
-  describe('removeMessagesByAssistant', () => {
-    it('should batch remove messages by assistantId and topicId', async () => {
-      // Setup
-      const assistantId = 'assistant-id';
-      const sessionId = 'session-id';
-      await clientDB.insert(sessions).values([
-        { id: 'bbb', userId },
-        { id: sessionId, userId },
-      ]);
-      await clientDB.insert(messages).values([
-        { sessionId, topicId, role: 'user', userId },
-        { sessionId, topicId, role: 'assistant', userId },
-        { sessionId: 'bbb', topicId, role: 'assistant', userId },
-      ]);
-
-      // Execute
-      await messageService.removeMessagesByAssistant(sessionId, topicId);
-
-      // Assert
-      const result = await clientDB.query.messages.findMany({
-        where: and(eq(messages.sessionId, sessionId), eq(messages.topicId, topicId)),
-      });
-
-      expect(result.length).toBe(0);
-    });
-  });
-
-  describe('clearAllMessage', () => {
-    it('should clear all messages from the table', async () => {
-      // Setup
-      await clientDB.insert(users).values({ id: 'another' });
-      await clientDB.insert(messages).values([
-        { id: mockMessageId, role: 'user', userId },
-        { role: 'user', userId: 'another' },
-      ]);
-
-      // Execute
-      await messageService.removeAllMessages();
-
-      // Assert
-      const result = await clientDB.query.messages.findMany({
-        where: eq(messages.userId, userId),
-      });
-      expect(result.length).toBe(0);
-    });
-  });
-
-  describe('getAllMessages', () => {
-    it('should retrieve all messages', async () => {
-      await clientDB.insert(messages).values([
-        { sessionId, topicId, content: '1', role: 'user', userId },
-        { sessionId, topicId, content: '2', role: 'assistant', userId },
-      ]);
-
-      // Execute
-      const data = await messageService.getAllMessages();
-
-      // Assert
-      expect(data).toMatchObject([
-        { sessionId, topicId, content: '1', role: 'user', userId },
-        { sessionId, topicId, content: '2', role: 'assistant', userId },
-      ]);
-    });
-  });
-
-  describe('updateMessageError', () => {
-    it('should update the error field of a message', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      const newError = {
-        type: 'InvalidProviderAPIKey',
-        message: 'Error occurred',
-      } as ChatMessageError;
-
-      // Execute
-      await messageService.updateMessageError(mockMessageId, newError);
-
-      // Assert
-      const result = await clientDB.query.messages.findFirst({
-        where: eq(messages.id, mockMessageId),
-      });
-
-      expect(result!.error).toEqual(newError);
-    });
-  });
-
-  // describe('updateMessagePlugin', () => {
-  // it('should update the plugin payload of a message', async () => {
-  //   // Setup
-  //   const newPlugin = {
-  //     type: 'default',
-  //     apiName: 'abc',
-  //     arguments: '',
-  //     identifier: 'plugin1',
-  //   } as ChatPluginPayload;
-  //
-  //   (MessageModel.update as Mock).mockResolvedValue({ ...mockMessage, plugin: newPlugin });
-  //
-  //   // Execute
-  //   const result = await messageService.updateMessagePlugin(mockMessageId, newPlugin);
-  //
-  //   // Assert
-  //   expect(MessageModel.update).toHaveBeenCalledWith(mockMessageId, { plugin: newPlugin });
-  //   expect(result).toEqual({ ...mockMessage, plugin: newPlugin });
-  // });
-  // });
-
-  describe('updateMessagePluginState', () => {
-    it('should update the plugin state of a message', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      await clientDB.insert(messagePlugins).values({ id: mockMessageId, userId });
-      const key = 'stateKey';
-      const value = 'stateValue';
-      const newPluginState = { [key]: value };
-
-      // Execute
-      await messageService.updateMessagePluginState(mockMessageId, { stateKey: value });
-
-      // Assert
-      const result = await clientDB.query.messagePlugins.findFirst({
-        where: eq(messagePlugins.id, mockMessageId),
-      });
-      expect(result!.state).toEqual(newPluginState);
-    });
-  });
-
-  describe('updateMessagePluginArguments', () => {
-    it('should update the plugin arguments object of a message', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      await clientDB.insert(messagePlugins).values({ id: mockMessageId, userId });
-      const value = 'stateValue';
-
-      // Execute
-      await messageService.updateMessagePluginArguments(mockMessageId, { key: value });
-
-      // Assert
-      const result = await clientDB.query.messagePlugins.findFirst({
-        where: eq(messageTTS.id, mockMessageId),
-      });
-      expect(result).toMatchObject({ arguments: '{"key":"stateValue"}' });
-    });
-    it('should update the plugin arguments string of a message', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      await clientDB.insert(messagePlugins).values({ id: mockMessageId, userId });
-      const value = 'stateValue';
-      // Execute
-      await messageService.updateMessagePluginArguments(
-        mockMessageId,
-        JSON.stringify({ abc: value }),
-      );
-
-      // Assert
-      const result = await clientDB.query.messagePlugins.findFirst({
-        where: eq(messageTTS.id, mockMessageId),
-      });
-      expect(result).toMatchObject({ arguments: '{"abc":"stateValue"}' });
-    });
-  });
-
-  describe('countMessages', () => {
-    it('should count the total number of messages', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-
-      // Execute
-      const count = await messageService.countMessages();
-
-      // Assert
-      expect(count).toBe(1);
-    });
-  });
-
-  describe('updateMessageTTS', () => {
-    it('should update the TTS field of a message', async () => {
-      // Setup
-      await clientDB
-        .insert(files)
-        .values({ id: 'file-abc', fileType: 'text', name: 'abc', url: 'abc', size: 100, userId });
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      const newTTS: ChatTTS = { contentMd5: 'abc', file: 'file-abc' };
-
-      // Execute
-      await messageService.updateMessageTTS(mockMessageId, newTTS);
-
-      // Assert
-      const result = await clientDB.query.messageTTS.findFirst({
-        where: eq(messageTTS.id, mockMessageId),
-      });
-
-      expect(result).toMatchObject({ contentMd5: 'abc', fileId: 'file-abc', id: mockMessageId });
-    });
-  });
-
-  describe('updateMessageTranslate', () => {
-    it('should update the translate field of a message', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-      const newTranslate: ChatTranslate = { content: 'Translated text', to: 'es' };
-
-      // Execute
-      await messageService.updateMessageTranslate(mockMessageId, newTranslate);
-
-      // Assert
-      const result = await clientDB.query.messageTranslates.findFirst({
-        where: eq(messageTranslates.id, mockMessageId),
-      });
-
-      expect(result).toMatchObject(newTranslate);
-    });
-  });
-
-  describe('hasMessages', () => {
-    it('should return true if there are messages', async () => {
-      // Setup
-      await clientDB.insert(messages).values({ id: mockMessageId, role: 'user', userId });
-
-      // Execute
-      const result = await messageService.hasMessages();
-
-      // Assert
-      expect(result).toBe(true);
-    });
-
-    it('should return false if there are no messages', async () => {
-      // Execute
-      const result = await messageService.hasMessages();
-
-      // Assert
-      expect(result).toBe(false);
-    });
-  });
-});

package/src/services/message/client.ts

@@ -1,192 +0,0 @@
-import { UIChatMessage } from '@lobechat/types';
-
-import { INBOX_SESSION_ID } from '@/const/session';
-import { clientDB } from '@/database/client/db';
-import { MessageModel } from '@/database/models/message';
-import { BaseClientService } from '@/services/baseClientService';
-import { clientS3Storage } from '@/services/file/ClientS3';
-
-import { IMessageService } from './type';
-
-export class ClientService extends BaseClientService implements IMessageService {
-  private get messageModel(): MessageModel {
-    return new MessageModel(clientDB as any, this.userId);
-  }
-
-  createMessage: IMessageService['createMessage'] = async ({ sessionId, ...params }) => {
-    const { id } = await this.messageModel.create({
-      ...params,
-      sessionId: sessionId ? (this.toDbSessionId(sessionId) as string) : '',
-    });
-
-    return id;
-  };
-
-  createNewMessage: IMessageService['createNewMessage'] = async ({ sessionId, ...params }) => {
-    return await this.messageModel.createNewMessage(
-      {
-        ...params,
-        sessionId: sessionId ? (this.toDbSessionId(sessionId) as string) : '',
-      },
-      {
-        postProcessUrl: this.postProcessUrl,
-      },
-    );
-  };
-
-  batchCreateMessages: IMessageService['batchCreateMessages'] = async (messages) => {
-    return this.messageModel.batchCreate(messages);
-  };
-
-  getMessages: IMessageService['getMessages'] = async (sessionId, topicId) => {
-    const data = await this.messageModel.query(
-      {
-        sessionId: this.toDbSessionId(sessionId),
-        topicId,
-      },
-      {
-        groupAssistantMessages: false,
-        postProcessUrl: this.postProcessUrl,
-      },
-    );
-
-    return data as unknown as UIChatMessage[];
-  };
-
-  getGroupMessages: IMessageService['getGroupMessages'] = async (groupId, topicId) => {
-    // Use full query to hydrate fileList/imageList like single chat
-    const data = await this.messageModel.query(
-      {
-        sessionId: groupId,
-        topicId,
-      },
-      {
-        groupAssistantMessages: false,
-        postProcessUrl: this.postProcessUrl,
-      },
-    );
-
-    return data as unknown as UIChatMessage[];
-  };
-
-  getAllMessages: IMessageService['getAllMessages'] = async () => {
-    const data = await this.messageModel.queryAll();
-
-    return data as unknown as UIChatMessage[];
-  };
-
-  countMessages: IMessageService['countMessages'] = async (params) => {
-    return this.messageModel.count(params);
-  };
-
-  countWords: IMessageService['countWords'] = async (params) => {
-    return this.messageModel.countWords(params);
-  };
-
-  rankModels: IMessageService['rankModels'] = async () => {
-    return this.messageModel.rankModels();
-  };
-
-  getHeatmaps: IMessageService['getHeatmaps'] = async () => {
-    return this.messageModel.getHeatmaps();
-  };
-
-  getAllMessagesInSession: IMessageService['getAllMessagesInSession'] = async (sessionId) => {
-    const data = this.messageModel.queryBySessionId(this.toDbSessionId(sessionId));
-
-    return data as unknown as UIChatMessage[];
-  };
-
-  updateMessageError: IMessageService['updateMessageError'] = async (id, error) => {
-    return this.messageModel.update(id, { error });
-  };
-
-  updateMessage: IMessageService['updateMessage'] = async (id, message, options) => {
-    return this.messageModel.update(id, message, {
-      postProcessUrl: this.postProcessUrl,
-      sessionId: options?.sessionId,
-      topicId: options?.topicId,
-    });
-  };
-
-  updateMessageTTS: IMessageService['updateMessageTTS'] = async (id, tts) => {
-    return this.messageModel.updateTTS(id, tts as any);
-  };
-
-  updateMessageTranslate: IMessageService['updateMessageTranslate'] = async (id, translate) => {
-    return this.messageModel.updateTranslate(id, translate as any);
-  };
-
-  updateMessagePluginState: IMessageService['updateMessagePluginState'] = async (id, value) => {
-    return this.messageModel.updatePluginState(id, value);
-  };
-
-  updateMessagePluginError: IMessageService['updateMessagePluginError'] = async (id, value) => {
-    return this.messageModel.updateMessagePlugin(id, { error: value });
-  };
-
-  updateMessageRAG: IMessageService['updateMessageRAG'] = async (id, value) => {
-    console.log(id, value);
-    throw new Error('not implemented');
-  };
-
-  updateMessagePluginArguments: IMessageService['updateMessagePluginArguments'] = async (
-    id,
-    value,
-  ) => {
-    const args = typeof value === 'string' ? value : JSON.stringify(value);
-
-    return this.messageModel.updateMessagePlugin(id, { arguments: args });
-  };
-
-  removeMessage: IMessageService['removeMessage'] = async (id) => {
-    return this.messageModel.deleteMessage(id);
-  };
-
-  removeMessages: IMessageService['removeMessages'] = async (ids) => {
-    return this.messageModel.deleteMessages(ids);
-  };
-
-  removeMessagesByAssistant: IMessageService['removeMessagesByAssistant'] = async (
-    sessionId,
-    topicId,
-  ) => {
-    return this.messageModel.deleteMessagesBySession(this.toDbSessionId(sessionId), topicId);
-  };
-
-  removeMessagesByGroup: IMessageService['removeMessagesByGroup'] = async (groupId, topicId) => {
-    return this.messageModel.deleteMessagesBySession(groupId, topicId);
-  };
-
-  removeAllMessages: IMessageService['removeAllMessages'] = async () => {
-    return this.messageModel.deleteAllMessages();
-  };
-
-  hasMessages: IMessageService['hasMessages'] = async () => {
-    const number = await this.countMessages();
-    return number > 0;
-  };
-
-  messageCountToCheckTrace: IMessageService['messageCountToCheckTrace'] = async () => {
-    const number = await this.countMessages();
-    return number >= 4;
-  };
-
-  private toDbSessionId = (sessionId: string | undefined) => {
-    return sessionId === INBOX_SESSION_ID ? undefined : sessionId;
-  };
-
-  private postProcessUrl = async (url: string | null, file: any) => {
-    const hash = (url as string).replace('client-s3://', '');
-    const base64 = await this.getBase64ByFileHash(hash);
-
-    return `data:${file.fileType};base64,${base64}`;
-  };
-
-  private getBase64ByFileHash = async (hash: string) => {
-    const fileItem = await clientS3Storage.getObject(hash);
-    if (!fileItem) throw new Error('file not found');
-
-    return Buffer.from(await fileItem.arrayBuffer()).toString('base64');
-  };
-}