@lobehub/chat 1.36.46 → 1.37.0

package/src/services/import/pglite.ts

@@ -0,0 +1,34 @@
+import { clientDB } from '@/database/client/db';
+import { DataImporterRepos } from '@/database/repositories/dataImporter';
+import { BaseClientService } from '@/services/baseClientService';
+import { useUserStore } from '@/store/user';
+import { ImportStage, ImporterEntryData, OnImportCallbacks } from '@/types/importer';
+import { UserSettings } from '@/types/user/settings';
+
+export class ClientService extends BaseClientService {
+  private get dataImporter(): DataImporterRepos {
+    return new DataImporterRepos(clientDB as any, this.userId);
+  }
+
+  importSettings = async (settings: UserSettings) => {
+    await useUserStore.getState().importAppSettings(settings);
+  };
+
+  importData = async (data: ImporterEntryData, callbacks?: OnImportCallbacks) => {
+    callbacks?.onStageChange?.(ImportStage.Importing);
+    const time = Date.now();
+    try {
+      const result = await this.dataImporter.importData(data);
+      const duration = Date.now() - time;
+
+      callbacks?.onStageChange?.(ImportStage.Success);
+      callbacks?.onSuccess?.(result, duration);
+    } catch (e) {
+      console.error(e);
+      callbacks?.onStageChange?.(ImportStage.Error);
+      const error = e as Error;
+
+      callbacks?.onError?.({ code: 'ImportError', httpStatus: 0, message: error.message });
+    }
+  };
+}

package/src/services/message/client.ts

@@ -21,6 +21,7 @@ export class ClientService implements IMessageService {
     return id;
   }
 
+  // @ts-ignore
   async batchCreateMessages(messages: ChatMessage[]) {
     return MessageModel.batchCreate(messages);
   }
@@ -70,6 +71,7 @@ export class ClientService implements IMessageService {
     return MessageModel.update(id, { error });
   }
 
+  // @ts-ignore
   async updateMessage(id: string, message: Partial<DB_Message>) {
     return MessageModel.update(id, message);
   }

package/src/services/message/index.ts

@@ -1,5 +1,9 @@
-import { ClientService } from './client';
+import { ClientService as DeprecatedService } from './client';
+import { ClientService } from './pglite';
 import { ServerService } from './server';
 
+const clientService =
+  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
+
 export const messageService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : new ClientService();
+  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : clientService;

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

@@ -0,0 +1,430 @@
+import dayjs from 'dayjs';
+import { and, eq } from 'drizzle-orm';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/_deprecated/models/message';
+import { clientDB, initializeDB } from '@/database/client/db';
+import {
+  MessageItem,
+  files,
+  messagePlugins,
+  messageTTS,
+  messageTranslates,
+  messages,
+  sessions,
+  topics,
+  users,
+} from '@/database/schemas';
+import {
+  ChatMessage,
+  ChatMessageError,
+  ChatTTS,
+  ChatTranslate,
+  CreateMessageParams,
+} from '@/types/message';
+
+import { ClientService } from './pglite';
+
+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 ChatMessage;
+
+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 MessageItem[]);
+      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 });
+      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 });
+      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 });
+      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('countTodayMessages', () => {
+    it('should count the number of messages created today', async () => {
+      // Setup
+      const mockMessages = [
+        { ...mockMessage, id: undefined, createdAt: new Date(), userId },
+        { ...mockMessage, id: undefined, createdAt: new Date(), userId },
+        { ...mockMessage, id: undefined, createdAt: new Date('2023-01-01'), userId },
+      ];
+      await clientDB.insert(messages).values(mockMessages);
+
+      // Execute
+      const count = await messageService.countTodayMessages();
+
+      // Assert
+      expect(count).toBe(2);
+    });
+  });
+
+  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/pglite.ts

@@ -0,0 +1,118 @@
+import dayjs from 'dayjs';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import { clientDB } from '@/database/client/db';
+import { MessageItem } from '@/database/schemas';
+import { MessageModel } from '@/database/server/models/message';
+import { BaseClientService } from '@/services/baseClientService';
+import {
+  ChatMessage,
+  ChatMessageError,
+  ChatTTS,
+  ChatTranslate,
+  CreateMessageParams,
+} from '@/types/message';
+
+import { IMessageService } from './type';
+
+export class ClientService extends BaseClientService implements IMessageService {
+  private get messageModel(): MessageModel {
+    return new MessageModel(clientDB as any, this.userId);
+  }
+
+  async createMessage({ sessionId, ...params }: CreateMessageParams) {
+    const { id } = await this.messageModel.create({
+      ...params,
+      sessionId: this.toDbSessionId(sessionId) as string,
+    });
+
+    return id;
+  }
+
+  async batchCreateMessages(messages: MessageItem[]) {
+    return this.messageModel.batchCreate(messages);
+  }
+
+  async getMessages(sessionId: string, topicId?: string) {
+    const data = await this.messageModel.query({
+      sessionId: this.toDbSessionId(sessionId),
+      topicId,
+    });
+
+    return data as unknown as ChatMessage[];
+  }
+
+  async getAllMessages() {
+    const data = await this.messageModel.queryAll();
+
+    return data as unknown as ChatMessage[];
+  }
+
+  async countMessages() {
+    return this.messageModel.count();
+  }
+
+  async countTodayMessages() {
+    const topics = await this.messageModel.queryAll();
+    return topics.filter(
+      (item) => dayjs(item.createdAt).format('YYYY-MM-DD') === dayjs().format('YYYY-MM-DD'),
+    ).length;
+  }
+
+  async getAllMessagesInSession(sessionId: string) {
+    const data = this.messageModel.queryBySessionId(this.toDbSessionId(sessionId));
+
+    return data as unknown as ChatMessage[];
+  }
+
+  async updateMessageError(id: string, error: ChatMessageError) {
+    return this.messageModel.update(id, { error });
+  }
+
+  async updateMessage(id: string, message: Partial<MessageItem>) {
+    return this.messageModel.update(id, message);
+  }
+
+  async updateMessageTTS(id: string, tts: Partial<ChatTTS> | false) {
+    return this.messageModel.updateTTS(id, tts as any);
+  }
+
+  async updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false) {
+    return this.messageModel.updateTranslate(id, translate as any);
+  }
+
+  async updateMessagePluginState(id: string, value: Record<string, any>) {
+    return this.messageModel.updatePluginState(id, value);
+  }
+
+  async updateMessagePluginArguments(id: string, value: string | Record<string, any>) {
+    const args = typeof value === 'string' ? value : JSON.stringify(value);
+
+    return this.messageModel.updateMessagePlugin(id, { arguments: args });
+  }
+
+  async removeMessage(id: string) {
+    return this.messageModel.deleteMessage(id);
+  }
+
+  async removeMessages(ids: string[]) {
+    return this.messageModel.deleteMessages(ids);
+  }
+
+  async removeMessagesByAssistant(sessionId: string, topicId?: string) {
+    return this.messageModel.deleteMessagesBySession(this.toDbSessionId(sessionId), topicId);
+  }
+
+  async removeAllMessages() {
+    return this.messageModel.deleteAllMessages();
+  }
+
+  async hasMessages() {
+    const number = await this.countMessages();
+    return number > 0;
+  }
+
+  private toDbSessionId(sessionId: string | undefined) {
+    return sessionId === INBOX_SESSION_ID ? undefined : sessionId;
+  }
+}

package/src/services/message/server.ts

@@ -1,5 +1,6 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 import { INBOX_SESSION_ID } from '@/const/session';
+import { MessageItem } from '@/database/schemas';
 import { lambdaClient } from '@/libs/trpc/client';
 import {
   ChatMessage,
@@ -19,20 +20,23 @@ export class ServerService implements IMessageService {
     });
   }
 
-  batchCreateMessages(messages: ChatMessage[]): Promise<any> {
+  batchCreateMessages(messages: MessageItem[]): Promise<any> {
     return lambdaClient.message.batchCreateMessages.mutate(messages);
   }
 
-  getMessages(sessionId?: string, topicId?: string | undefined): Promise<ChatMessage[]> {
-    return lambdaClient.message.getMessages.query({
+  getMessages = async (sessionId?: string, topicId?: string | undefined) => {
+    const data = await lambdaClient.message.getMessages.query({
       sessionId: this.toDbSessionId(sessionId),
       topicId,
     });
-  }
+
+    return data as unknown as ChatMessage[];
+  };
 
   getAllMessages(): Promise<ChatMessage[]> {
     return lambdaClient.message.getAllMessages.query();
   }
+
   getAllMessagesInSession(sessionId: string): Promise<ChatMessage[]> {
     return lambdaClient.message.getAllMessagesInSession.query({
       sessionId: this.toDbSessionId(sessionId),
@@ -63,7 +67,7 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updateMessagePlugin.mutate({ id, value: { arguments: args } });
   }
 
-  updateMessage(id: string, message: Partial<ChatMessage>): Promise<any> {
+  updateMessage(id: string, message: Partial<MessageItem>): Promise<any> {
     return lambdaClient.message.update.mutate({ id, value: message });
   }
 
@@ -79,10 +83,6 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updatePluginState.mutate({ id, value });
   }
 
-  bindMessagesToTopic(_topicId: string, _messageIds: string[]): Promise<any> {
-    throw new Error('Method not implemented.');
-  }
-
   removeMessage(id: string): Promise<any> {
     return lambdaClient.message.removeMessage.mutate({ id });
   }

package/src/services/message/type.ts

@@ -1,4 +1,4 @@
-import { DB_Message } from '@/database/_deprecated/schemas/message';
+import { MessageItem } from '@/database/schemas';
 import {
   ChatMessage,
   ChatMessageError,
@@ -11,7 +11,7 @@ import {
 
 export interface IMessageService {
   createMessage(data: CreateMessageParams): Promise<string>;
-  batchCreateMessages(messages: ChatMessage[]): Promise<any>;
+  batchCreateMessages(messages: MessageItem[]): Promise<any>;
 
   getMessages(sessionId: string, topicId?: string): Promise<ChatMessage[]>;
   getAllMessages(): Promise<ChatMessage[]>;
@@ -20,11 +20,10 @@ export interface IMessageService {
   countTodayMessages(): Promise<number>;
 
   updateMessageError(id: string, error: ChatMessageError): Promise<any>;
-  updateMessage(id: string, message: Partial<DB_Message>): Promise<any>;
+  updateMessage(id: string, message: Partial<MessageItem>): Promise<any>;
   updateMessageTTS(id: string, tts: Partial<ChatTTS> | false): Promise<any>;
   updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false): Promise<any>;
   updateMessagePluginState(id: string, value: Record<string, any>): Promise<any>;
-  bindMessagesToTopic(topicId: string, messageIds: string[]): Promise<any>;
 
   removeMessage(id: string): Promise<any>;
   removeMessages(ids: string[]): Promise<any>;