@lobehub/chat 0.162.25 → 0.164.0

package/src/database/server/models/__tests__/user.test.ts

@@ -0,0 +1,173 @@
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import { getTestDBInstance } from '@/database/server/core/dbForTest';
+import { KeyVaultsGateKeeper } from '@/server/keyVaultsEncrypt';
+import { UserPreference } from '@/types/user';
+import { UserSettings } from '@/types/user/settings';
+
+import { userSettings, users } from '../../schemas/lobechat';
+import { SessionModel } from '../session';
+import { UserModel } from '../user';
+
+let serverDB = await getTestDBInstance();
+
+vi.mock('@/database/server/core/db', async () => ({
+  get serverDB() {
+    return serverDB;
+  },
+}));
+
+const userId = 'user-db';
+const userModel = new UserModel();
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.delete(userSettings);
+  process.env.KEY_VAULTS_SECRET = 'ofQiJCXLF8mYemwfMWLOHoHimlPu91YmLfU7YZ4lreQ=';
+});
+
+afterEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.delete(userSettings);
+  process.env.KEY_VAULTS_SECRET = undefined;
+});
+
+describe('UserModel', () => {
+  describe('createUser', () => {
+    it('should create a new user and inbox session', async () => {
+      const params = {
+        id: userId,
+        username: 'testuser',
+        email: 'test@example.com',
+      };
+
+      await userModel.createUser(params);
+
+      const user = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
+      expect(user).not.toBeNull();
+      expect(user?.username).toBe('testuser');
+      expect(user?.email).toBe('test@example.com');
+
+      const sessionModel = new SessionModel(userId);
+      const inbox = await sessionModel.findByIdOrSlug(INBOX_SESSION_ID);
+      expect(inbox).not.toBeNull();
+    });
+  });
+
+  describe('deleteUser', () => {
+    it('should delete a user', async () => {
+      await serverDB.insert(users).values({ id: userId });
+
+      await userModel.deleteUser(userId);
+
+      const user = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
+      expect(user).toBeUndefined();
+    });
+  });
+
+  describe('findById', () => {
+    it('should find a user by ID', async () => {
+      await serverDB.insert(users).values({ id: userId, username: 'testuser' });
+
+      const user = await userModel.findById(userId);
+
+      expect(user).not.toBeNull();
+      expect(user?.id).toBe(userId);
+      expect(user?.username).toBe('testuser');
+    });
+  });
+
+  describe('getUserState', () => {
+    it('should get user state with decrypted keyVaults', async () => {
+      const preference = { useCmdEnterToSend: true } as UserPreference;
+      const keyVaults = { apiKey: 'secret' };
+
+      await serverDB.insert(users).values({ id: userId, preference });
+
+      const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+      const encryptedKeyVaults = await gateKeeper.encrypt(JSON.stringify(keyVaults));
+
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        keyVaults: encryptedKeyVaults,
+      });
+
+      const state = await userModel.getUserState(userId);
+
+      expect(state.userId).toBe(userId);
+      expect(state.preference).toEqual(preference);
+      expect(state.settings.keyVaults).toEqual(keyVaults);
+    });
+
+    it('should throw an error if user not found', async () => {
+      await expect(userModel.getUserState('invalid-user-id')).rejects.toThrow('user not found');
+    });
+  });
+
+  describe('updateUser', () => {
+    it('should update user fields', async () => {
+      await serverDB.insert(users).values({ id: userId, username: 'oldname' });
+
+      await userModel.updateUser(userId, { username: 'newname' });
+
+      const updatedUser = await serverDB.query.users.findFirst({
+        where: eq(users.id, userId),
+      });
+      expect(updatedUser?.username).toBe('newname');
+    });
+  });
+
+  describe('deleteSetting', () => {
+    it('should delete user settings', async () => {
+      await serverDB.insert(users).values({ id: userId });
+      await serverDB.insert(userSettings).values({ id: userId });
+
+      await userModel.deleteSetting(userId);
+
+      const settings = await serverDB.query.userSettings.findFirst({
+        where: eq(users.id, userId),
+      });
+
+      expect(settings).toBeUndefined();
+    });
+  });
+
+  describe('updateSetting', () => {
+    it('should update user settings with encrypted keyVaults', async () => {
+      const settings = {
+        general: { language: 'en-US' },
+        keyVaults: { openai: { apiKey: 'secret' } },
+      } as UserSettings;
+      await serverDB.insert(users).values({ id: userId });
+
+      await userModel.updateSetting(userId, settings);
+
+      const updatedSettings = await serverDB.query.userSettings.findFirst({
+        where: eq(users.id, userId),
+      });
+      expect(updatedSettings?.general).toEqual(settings.general);
+      expect(updatedSettings?.keyVaults).not.toBe(JSON.stringify(settings.keyVaults));
+
+      const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+      const { plaintext } = await gateKeeper.decrypt(updatedSettings!.keyVaults!);
+      expect(JSON.parse(plaintext)).toEqual(settings.keyVaults);
+    });
+  });
+
+  describe('updatePreference', () => {
+    it('should update user preference', async () => {
+      const preference = { guide: { topic: false } } as UserPreference;
+      await serverDB.insert(users).values({ id: userId, preference });
+
+      const newPreference: Partial<UserPreference> = {
+        guide: { topic: true, moveSettingsToAvatar: true },
+      };
+      await userModel.updatePreference(userId, newPreference);
+
+      const updatedUser = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
+      expect(updatedUser?.preference).toEqual({ ...preference, ...newPreference });
+    });
+  });
+});

package/src/database/server/models/_template.ts

@@ -0,0 +1,44 @@
+import { eq } from 'drizzle-orm';
+import { and, desc } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+
+import { NewSessionGroup, UserItem, sessionGroups } from '../schemas/lobechat';
+
+export class TemplateModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (params: NewSessionGroup) => {
+    return serverDB.insert(sessionGroups).values({ ...params, userId: this.userId });
+  };
+
+  delete = async (id: string) => {
+    return serverDB
+      .delete(sessionGroups)
+      .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
+  };
+
+  query = async () => {
+    return serverDB.query.sessionGroups.findMany({
+      orderBy: [desc(sessionGroups.updatedAt)],
+      where: eq(sessionGroups.userId, this.userId),
+    });
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.sessionGroups.findFirst({
+      where: and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)),
+    });
+  };
+
+  async update(id: string, value: Partial<UserItem>) {
+    return serverDB
+      .update(sessionGroups)
+      .set({ ...value, updatedAt: new Date() })
+      .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
+  }
+}

package/src/database/server/models/file.ts

@@ -0,0 +1,51 @@
+import { eq } from 'drizzle-orm';
+import { and, desc } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server/core/db';
+
+import { FileItem, NewFile, files } from '../schemas/lobechat';
+
+export class FileModel {
+  private readonly userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (params: Omit<NewFile, 'id' | 'userId'>) => {
+    const result = await serverDB
+      .insert(files)
+      .values({ ...params, userId: this.userId })
+      .returning();
+
+    return { id: result[0].id };
+  };
+
+  delete = async (id: string) => {
+    return serverDB.delete(files).where(and(eq(files.id, id), eq(files.userId, this.userId)));
+  };
+
+  clear = async () => {
+    return serverDB.delete(files).where(eq(files.userId, this.userId));
+  };
+
+  query = async () => {
+    return serverDB.query.files.findMany({
+      orderBy: [desc(files.updatedAt)],
+      where: eq(files.userId, this.userId),
+    });
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.files.findFirst({
+      where: and(eq(files.id, id), eq(files.userId, this.userId)),
+    });
+  };
+
+  async update(id: string, value: Partial<FileItem>) {
+    return serverDB
+      .update(files)
+      .set({ ...value, updatedAt: new Date() })
+      .where(and(eq(files.id, id), eq(files.userId, this.userId)));
+  }
+}

package/src/database/server/models/message.ts

@@ -0,0 +1,378 @@
+import { count, sql } from 'drizzle-orm';
+import { and, asc, desc, eq, isNull, like } from 'drizzle-orm/expressions';
+import { inArray } from 'drizzle-orm/sql/expressions/conditions';
+
+import { CreateMessageParams } from '@/database/client/models/message';
+import { serverDB } from '@/database/server/core/db';
+import { idGenerator } from '@/database/server/utils/idGenerator';
+import { ChatTTS, ChatToolPayload } from '@/types/message';
+import { merge } from '@/utils/merge';
+
+import {
+  MessageItem,
+  filesToMessages,
+  messagePlugins,
+  messageTTS,
+  messageTranslates,
+  messages,
+} from '../schemas/lobechat';
+
+export interface QueryMessageParams {
+  current?: number;
+  pageSize?: number;
+  sessionId?: string | null;
+  topicId?: string | null;
+}
+
+export class MessageModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  // **************** Query *************** //
+  async query({
+    current = 0,
+    pageSize = 1000,
+    sessionId,
+    topicId,
+  }: QueryMessageParams = {}): Promise<MessageItem[]> {
+    const offset = current * pageSize;
+
+    const result = await serverDB
+      .select({
+        /* eslint-disable sort-keys-fix/sort-keys-fix*/
+        id: messages.id,
+        role: messages.role,
+        content: messages.content,
+        error: messages.error,
+
+        model: messages.model,
+        provider: messages.provider,
+
+        createdAt: messages.createdAt,
+        updatedAt: messages.updatedAt,
+
+        parentId: messages.parentId,
+
+        tools: messages.tools,
+        tool_call_id: messagePlugins.toolCallId,
+
+        plugin: {
+          apiName: messagePlugins.apiName,
+          arguments: messagePlugins.arguments,
+          identifier: messagePlugins.identifier,
+          type: messagePlugins.type,
+        },
+        pluginError: messagePlugins.error,
+        pluginState: messagePlugins.state,
+
+        translate: {
+          content: messageTranslates.content,
+          from: messageTranslates.from,
+          to: messageTranslates.to,
+        },
+
+        ttsId: messageTTS.id,
+
+        // TODO: 确认下如何处理 TTS 的读取
+        // ttsContentMd5: messageTTS.contentMd5,
+        // ttsFile: messageTTS.fileId,
+        // ttsVoice: messageTTS.voice,
+        /* eslint-enable */
+      })
+      .from(messages)
+      .where(
+        and(
+          eq(messages.userId, this.userId),
+          this.matchSession(sessionId),
+          this.matchTopic(topicId),
+        ),
+      )
+      .leftJoin(messagePlugins, eq(messagePlugins.id, messages.id))
+      .leftJoin(messageTranslates, eq(messageTranslates.id, messages.id))
+      .leftJoin(messageTTS, eq(messageTTS.id, messages.id))
+      .orderBy(asc(messages.createdAt))
+      .limit(pageSize)
+      .offset(offset);
+
+    const messageIds = result.map((message) => message.id as string);
+
+    if (messageIds.length === 0) return result;
+
+    const fileIds = await serverDB
+      .select({
+        fileId: filesToMessages.fileId,
+        messageId: filesToMessages.messageId,
+      })
+      .from(filesToMessages)
+      .where(inArray(filesToMessages.messageId, messageIds));
+
+    return result.map(
+      ({
+        model,
+        provider,
+        translate,
+        ttsId,
+        // ttsFile, ttsId, ttsContentMd5, ttsVoice,
+        ...item
+      }) => ({
+        ...item,
+        extra: {
+          fromModel: model,
+          fromProvider: provider,
+          translate,
+          tts: ttsId
+            ? {
+                // contentMd5: ttsContentMd5,
+                // file: ttsFile,
+                // voice: ttsVoice,
+              }
+            : undefined,
+        },
+        files: fileIds.filter((relation) => relation.messageId === item.id).map((r) => r.fileId),
+      }),
+    );
+  }
+
+  async findById(id: string) {
+    return serverDB.query.messages.findFirst({
+      where: and(eq(messages.id, id), eq(messages.userId, this.userId)),
+    });
+  }
+
+  async queryAll(): Promise<MessageItem[]> {
+    return serverDB
+      .select()
+      .from(messages)
+      .orderBy(messages.createdAt)
+      .where(eq(messages.userId, this.userId))
+
+      .execute();
+  }
+
+  async queryBySessionId(sessionId?: string | null): Promise<MessageItem[]> {
+    return serverDB.query.messages.findMany({
+      orderBy: [asc(messages.createdAt)],
+      where: and(eq(messages.userId, this.userId), this.matchSession(sessionId)),
+    });
+  }
+
+  async queryByKeyword(keyword: string): Promise<MessageItem[]> {
+    if (!keyword) return [];
+
+    return serverDB.query.messages.findMany({
+      orderBy: [desc(messages.createdAt)],
+      where: and(eq(messages.userId, this.userId), like(messages.content, `%${keyword}%`)),
+    });
+  }
+
+  async count() {
+    const result = await serverDB
+      .select({
+        count: count(),
+      })
+      .from(messages)
+      .where(eq(messages.userId, this.userId))
+      .execute();
+
+    return result[0].count;
+  }
+
+  async countToday() {
+    const today = new Date();
+    today.setHours(0, 0, 0, 0);
+    const tomorrow = new Date(today);
+    tomorrow.setDate(tomorrow.getDate() + 1);
+
+    const result = await serverDB
+      .select({
+        count: count(),
+      })
+      .from(messages)
+      .where(
+        and(
+          eq(messages.userId, this.userId),
+          sql`${messages.createdAt} >= ${today} AND ${messages.createdAt} < ${tomorrow}`,
+        ),
+      )
+      .execute();
+
+    return result[0].count;
+  }
+
+  // **************** Create *************** //
+
+  async create(
+    { fromModel, fromProvider, files, ...message }: CreateMessageParams,
+    id: string = this.genId(),
+  ): Promise<MessageItem> {
+    return serverDB.transaction(async (trx) => {
+      const [item] = (await trx
+        .insert(messages)
+        .values({
+          ...message,
+          id,
+          model: fromModel,
+          provider: fromProvider,
+          userId: this.userId,
+        })
+        .returning()) as MessageItem[];
+
+      // Insert the plugin data if the message is a tool
+      if (message.role === 'tool') {
+        await trx.insert(messagePlugins).values({
+          apiName: message.plugin?.apiName,
+          arguments: message.plugin?.arguments,
+          id,
+          identifier: message.plugin?.identifier,
+          toolCallId: message.tool_call_id,
+          type: message.plugin?.type,
+        });
+      }
+
+      if (files && files.length > 0) {
+        await trx
+          .insert(filesToMessages)
+          .values(files.map((file) => ({ fileId: file, messageId: id })));
+      }
+
+      return item;
+    });
+  }
+
+  async batchCreate(newMessages: MessageItem[]) {
+    const messagesToInsert = newMessages.map((m) => {
+      return { ...m, userId: this.userId };
+    });
+
+    return serverDB.insert(messages).values(messagesToInsert);
+  }
+
+  // **************** Update *************** //
+
+  async update(id: string, message: Partial<MessageItem>) {
+    return serverDB
+      .update(messages)
+      .set(message)
+      .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
+  }
+
+  async updatePluginState(id: string, state: Record<string, any>) {
+    const item = await serverDB.query.messagePlugins.findFirst({
+      where: eq(messagePlugins.id, id),
+    });
+    if (!item) throw new Error('Plugin not found');
+
+    return serverDB
+      .update(messagePlugins)
+      .set({ state: merge(item.state || {}, state) })
+      .where(eq(messagePlugins.id, id));
+  }
+
+  async updateTranslate(id: string, translate: Partial<MessageItem>) {
+    const result = await serverDB.query.messageTranslates.findFirst({
+      where: and(eq(messageTranslates.id, id)),
+    });
+
+    // If the message does not exist in the translate table, insert it
+    if (!result) {
+      return serverDB.insert(messageTranslates).values({ ...translate, id });
+    }
+
+    // or just update the existing one
+    return serverDB.update(messageTranslates).set(translate).where(eq(messageTranslates.id, id));
+  }
+
+  async updateTTS(id: string, tts: Partial<ChatTTS>) {
+    const result = await serverDB.query.messageTTS.findFirst({
+      where: and(eq(messageTTS.id, id)),
+    });
+
+    // If the message does not exist in the translate table, insert it
+    if (!result) {
+      return serverDB
+        .insert(messageTTS)
+        .values({ contentMd5: tts.contentMd5, fileId: tts.file, id, voice: tts.voice });
+    }
+
+    // or just update the existing one
+    return serverDB
+      .update(messageTTS)
+      .set({ contentMd5: tts.contentMd5, fileId: tts.file, voice: tts.voice })
+      .where(eq(messageTTS.id, id));
+  }
+
+  // **************** Delete *************** //
+
+  async deleteMessage(id: string) {
+    return serverDB.transaction(async (tx) => {
+      // 1. 查询要删除的 message 的完整信息
+      const message = await tx
+        .select()
+        .from(messages)
+        .where(and(eq(messages.id, id), eq(messages.userId, this.userId)))
+        .limit(1);
+
+      // 如果找不到要删除的 message,直接返回
+      if (message.length === 0) return;
+
+      // 2. 检查 message 是否包含 tools
+      const toolCallIds = message[0].tools?.map((tool: ChatToolPayload) => tool.id).filter(Boolean);
+
+      let relatedMessageIds: string[] = [];
+
+      if (toolCallIds?.length > 0) {
+        // 3. 如果 message 包含 tools,查询出所有相关联的 message id
+        const res = await tx
+          .select({ id: messagePlugins.id })
+          .from(messagePlugins)
+          .where(inArray(messagePlugins.toolCallId, toolCallIds))
+          .execute();
+
+        relatedMessageIds = res.map((row) => row.id);
+      }
+
+      // 4. 合并要删除的 message id 列表
+      const messageIdsToDelete = [id, ...relatedMessageIds];
+
+      // 5. 删除所有相关的 message
+      await tx.delete(messages).where(inArray(messages.id, messageIdsToDelete));
+    });
+  }
+
+  async deleteMessageTranslate(id: string) {
+    return serverDB.delete(messageTranslates).where(and(eq(messageTranslates.id, id)));
+  }
+
+  async deleteMessageTTS(id: string) {
+    return serverDB.delete(messageTTS).where(and(eq(messageTTS.id, id)));
+  }
+
+  async deleteMessages(sessionId?: string | null, topicId?: string | null) {
+    return serverDB
+      .delete(messages)
+      .where(
+        and(
+          eq(messages.userId, this.userId),
+          this.matchSession(sessionId),
+          this.matchTopic(topicId),
+        ),
+      );
+  }
+
+  async deleteAllMessages() {
+    return serverDB.delete(messages).where(eq(messages.userId, this.userId));
+  }
+
+  // **************** Helper *************** //
+
+  private genId = () => idGenerator('messages', 14);
+
+  private matchSession = (sessionId?: string | null) =>
+    sessionId ? eq(messages.sessionId, sessionId) : isNull(messages.sessionId);
+
+  private matchTopic = (topicId?: string | null) =>
+    topicId ? eq(messages.topicId, topicId) : isNull(messages.topicId);
+}

package/src/database/server/models/plugin.ts

@@ -0,0 +1,63 @@
+import { and, desc, eq } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+
+import { InstalledPluginItem, NewInstalledPlugin, installedPlugins } from '../schemas/lobechat';
+
+export class PluginModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  create = async (
+    params: Pick<NewInstalledPlugin, 'type' | 'identifier' | 'manifest' | 'customParams'>,
+  ) => {
+    const [result] = await serverDB
+      .insert(installedPlugins)
+      .values({ ...params, createdAt: new Date(), updatedAt: new Date(), userId: this.userId })
+      .returning();
+
+    return result;
+  };
+
+  delete = async (id: string) => {
+    return serverDB
+      .delete(installedPlugins)
+      .where(and(eq(installedPlugins.identifier, id), eq(installedPlugins.userId, this.userId)));
+  };
+
+  deleteAll = async () => {
+    return serverDB.delete(installedPlugins).where(eq(installedPlugins.userId, this.userId));
+  };
+
+  query = async () => {
+    return serverDB
+      .select({
+        createdAt: installedPlugins.createdAt,
+        customParams: installedPlugins.customParams,
+        identifier: installedPlugins.identifier,
+        manifest: installedPlugins.manifest,
+        settings: installedPlugins.settings,
+        type: installedPlugins.type,
+        updatedAt: installedPlugins.updatedAt,
+      })
+      .from(installedPlugins)
+      .where(eq(installedPlugins.userId, this.userId))
+      .orderBy(desc(installedPlugins.createdAt));
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.installedPlugins.findFirst({
+      where: and(eq(installedPlugins.identifier, id), eq(installedPlugins.userId, this.userId)),
+    });
+  };
+
+  async update(id: string, value: Partial<InstalledPluginItem>) {
+    return serverDB
+      .update(installedPlugins)
+      .set({ ...value, updatedAt: new Date() })
+      .where(and(eq(installedPlugins.identifier, id), eq(installedPlugins.userId, this.userId)));
+  }
+}