@lobehub/chat 1.35.9 → 1.35.11

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

@@ -3,7 +3,7 @@ import { and, desc, like } from 'drizzle-orm/expressions';
 import type { PgTransaction } from 'drizzle-orm/pg-core';
 
 import { serverDBEnv } from '@/config/db';
-import { serverDB } from '@/database/server/core/db';
+import { LobeChatDatabase } from '@/database/type';
 import { FilesTabs, QueryFileListParams, SortType } from '@/types/files';
 
 import {
@@ -20,13 +20,15 @@ import {
 
 export class FileModel {
   private readonly userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   create = async (params: Omit<NewFile, 'id' | 'userId'> & { knowledgeBaseId?: string }) => {
-    const result = await serverDB.transaction(async (trx) => {
+    const result = await this.db.transaction(async (trx) => {
       const result = await trx
         .insert(files)
         .values({ ...params, userId: this.userId })
@@ -47,11 +49,11 @@ export class FileModel {
   };
 
   createGlobalFile = async (file: Omit<NewGlobalFile, 'id' | 'userId'>) => {
-    return serverDB.insert(globalFiles).values(file).returning();
+    return this.db.insert(globalFiles).values(file).returning();
   };
 
   checkHash = async (hash: string) => {
-    const item = await serverDB.query.globalFiles.findFirst({
+    const item = await this.db.query.globalFiles.findFirst({
       where: eq(globalFiles.hashId, hash),
     });
     if (!item) return { isExist: false };
@@ -71,7 +73,7 @@ export class FileModel {
 
     const fileHash = file.fileHash!;
 
-    return await serverDB.transaction(async (trx) => {
+    return await this.db.transaction(async (trx) => {
       // 1. 删除相关的 chunks
       await this.deleteFileChunks(trx as any, [id]);
 
@@ -96,11 +98,11 @@ export class FileModel {
   };
 
   deleteGlobalFile = async (hashId: string) => {
-    return serverDB.delete(globalFiles).where(eq(globalFiles.hashId, hashId));
+    return this.db.delete(globalFiles).where(eq(globalFiles.hashId, hashId));
   };
 
   countUsage = async () => {
-    const result = await serverDB
+    const result = await this.db
       .select({
         totalSize: sum(files.size),
       })
@@ -114,7 +116,7 @@ export class FileModel {
     const fileList = await this.findByIds(ids);
     const hashList = fileList.map((file) => file.fileHash!);
 
-    return await serverDB.transaction(async (trx) => {
+    return await this.db.transaction(async (trx) => {
       // 1. 删除相关的 chunks
       await this.deleteFileChunks(trx as any, ids);
 
@@ -159,7 +161,7 @@ export class FileModel {
   };
 
   clear = async () => {
-    return serverDB.delete(files).where(eq(files.userId, this.userId));
+    return this.db.delete(files).where(eq(files.userId, this.userId));
   };
 
   query = async ({
@@ -198,7 +200,7 @@ export class FileModel {
     }
 
     // 3. build query
-    let query = serverDB
+    let query = this.db
       .select({
         chunkTaskId: files.chunkTaskId,
         createdAt: files.createdAt,
@@ -230,7 +232,7 @@ export class FileModel {
       whereClause = and(
         whereClause,
         notExists(
-          serverDB.select().from(knowledgeBaseFiles).where(eq(knowledgeBaseFiles.fileId, files.id)),
+          this.db.select().from(knowledgeBaseFiles).where(eq(knowledgeBaseFiles.fileId, files.id)),
         ),
       );
     }
@@ -240,19 +242,19 @@ export class FileModel {
   };
 
   findByIds = async (ids: string[]) => {
-    return serverDB.query.files.findMany({
+    return this.db.query.files.findMany({
       where: and(inArray(files.id, ids), eq(files.userId, this.userId)),
     });
   };
 
   findById = async (id: string) => {
-    return serverDB.query.files.findFirst({
+    return this.db.query.files.findFirst({
       where: and(eq(files.id, id), eq(files.userId, this.userId)),
     });
   };
 
   countFilesByHash = async (hash: string) => {
-    const result = await serverDB
+    const result = await this.db
       .select({
         count: count(),
       })
@@ -263,7 +265,7 @@ export class FileModel {
   };
 
   async update(id: string, value: Partial<FileItem>) {
-    return serverDB
+    return this.db
       .update(files)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(files.id, id), eq(files.userId, this.userId)));
@@ -293,7 +295,7 @@ export class FileModel {
   };
 
   async findByNames(fileNames: string[]) {
-    return serverDB.query.files.findMany({
+    return this.db.query.files.findMany({
       where: and(
         or(...fileNames.map((name) => like(files.name, `${name}%`))),
         eq(files.userId, this.userId),

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

@@ -1,22 +1,24 @@
 import { eq, inArray } from 'drizzle-orm';
 import { and, desc } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { KnowledgeBaseItem } from '@/types/knowledgeBase';
 
 import { NewKnowledgeBase, knowledgeBaseFiles, knowledgeBases } from '../../schemas';
 
 export class KnowledgeBaseModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   // create
 
   create = async (params: Omit<NewKnowledgeBase, 'userId'>) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(knowledgeBases)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -25,7 +27,7 @@ export class KnowledgeBaseModel {
   };
 
   addFilesToKnowledgeBase = async (id: string, fileIds: string[]) => {
-    return serverDB
+    return this.db
       .insert(knowledgeBaseFiles)
       .values(fileIds.map((fileId) => ({ fileId, knowledgeBaseId: id, userId: this.userId })))
       .returning();
@@ -33,17 +35,17 @@ export class KnowledgeBaseModel {
 
   // delete
   delete = async (id: string) => {
-    return serverDB
+    return this.db
       .delete(knowledgeBases)
       .where(and(eq(knowledgeBases.id, id), eq(knowledgeBases.userId, this.userId)));
   };
 
   deleteAll = async () => {
-    return serverDB.delete(knowledgeBases).where(eq(knowledgeBases.userId, this.userId));
+    return this.db.delete(knowledgeBases).where(eq(knowledgeBases.userId, this.userId));
   };
 
   removeFilesFromKnowledgeBase = async (knowledgeBaseId: string, ids: string[]) => {
-    return serverDB.delete(knowledgeBaseFiles).where(
+    return this.db.delete(knowledgeBaseFiles).where(
       and(
         eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
         inArray(knowledgeBaseFiles.fileId, ids),
@@ -53,7 +55,7 @@ export class KnowledgeBaseModel {
   };
   // query
   query = async () => {
-    const data = await serverDB
+    const data = await this.db
       .select({
         avatar: knowledgeBases.avatar,
         createdAt: knowledgeBases.createdAt,
@@ -73,21 +75,21 @@ export class KnowledgeBaseModel {
   };
 
   findById = async (id: string) => {
-    return serverDB.query.knowledgeBases.findFirst({
+    return this.db.query.knowledgeBases.findFirst({
       where: and(eq(knowledgeBases.id, id), eq(knowledgeBases.userId, this.userId)),
     });
   };
 
   // update
   async update(id: string, value: Partial<KnowledgeBaseItem>) {
-    return serverDB
+    return this.db
       .update(knowledgeBases)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(knowledgeBases.id, id), eq(knowledgeBases.userId, this.userId)));
   }
 
-  static async findById(id: string) {
-    return serverDB.query.knowledgeBases.findFirst({
+  static async findById(db: LobeChatDatabase, id: string) {
+    return db.query.knowledgeBases.findFirst({
       where: eq(knowledgeBases.id, id),
     });
   }

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

@@ -1,8 +1,8 @@
 import { count } from 'drizzle-orm';
 import { and, asc, desc, eq, gte, inArray, isNull, like, lt } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server/core/db';
 import { idGenerator } from '@/database/utils/idGenerator';
+import { LobeChatDatabase } from '@/database/type';
 import { getFullFileUrl } from '@/server/utils/files';
 import {
   ChatFileItem,
@@ -39,9 +39,11 @@ export interface QueryMessageParams {
 
 export class MessageModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   // **************** Query *************** //
@@ -54,7 +56,7 @@ export class MessageModel {
     const offset = current * pageSize;
 
     // 1. get basic messages
-    const result = await serverDB
+    const result = await this.db
       .select({
         /* eslint-disable sort-keys-fix/sort-keys-fix*/
         id: messages.id,
@@ -115,7 +117,7 @@ export class MessageModel {
     if (messageIds.length === 0) return [];
 
     // 2. get relative files
-    const rawRelatedFileList = await serverDB
+    const rawRelatedFileList = await this.db
       .select({
         fileType: files.fileType,
         id: messagesFiles.fileId,
@@ -139,7 +141,7 @@ export class MessageModel {
     const fileList = relatedFileList.filter((i) => !(i.fileType || '').startsWith('image'));
 
     // 3. get relative file chunks
-    const chunksList = await serverDB
+    const chunksList = await this.db
       .select({
         fileId: files.id,
         fileType: files.fileType,
@@ -157,7 +159,7 @@ export class MessageModel {
       .where(inArray(messageQueryChunks.messageId, messageIds));
 
     // 3. get relative message query
-    const messageQueriesList = await serverDB
+    const messageQueriesList = await this.db
       .select({
         id: messageQueries.id,
         messageId: messageQueries.messageId,
@@ -216,13 +218,13 @@ export class MessageModel {
   }
 
   async findById(id: string) {
-    return serverDB.query.messages.findFirst({
+    return this.db.query.messages.findFirst({
       where: and(eq(messages.id, id), eq(messages.userId, this.userId)),
     });
   }
 
   async findMessageQueriesById(messageId: string) {
-    const result = await serverDB
+    const result = await this.db
       .select({
         embeddings: embeddings.embeddings,
         id: messageQueries.id,
@@ -240,7 +242,7 @@ export class MessageModel {
   }
 
   async queryAll(): Promise<MessageItem[]> {
-    return serverDB
+    return this.db
       .select()
       .from(messages)
       .orderBy(messages.createdAt)
@@ -250,7 +252,7 @@ export class MessageModel {
   }
 
   async queryBySessionId(sessionId?: string | null): Promise<MessageItem[]> {
-    return serverDB.query.messages.findMany({
+    return this.db.query.messages.findMany({
       orderBy: [asc(messages.createdAt)],
       where: and(eq(messages.userId, this.userId), this.matchSession(sessionId)),
     });
@@ -259,14 +261,14 @@ export class MessageModel {
   async queryByKeyword(keyword: string): Promise<MessageItem[]> {
     if (!keyword) return [];
 
-    return serverDB.query.messages.findMany({
+    return this.db.query.messages.findMany({
       orderBy: [desc(messages.createdAt)],
       where: and(eq(messages.userId, this.userId), like(messages.content, `%${keyword}%`)),
     });
   }
 
   async count() {
-    const result = await serverDB
+    const result = await this.db
       .select({
         count: count(),
       })
@@ -283,7 +285,7 @@ export class MessageModel {
     const tomorrow = new Date(today);
     tomorrow.setDate(tomorrow.getDate() + 1);
 
-    const result = await serverDB
+    const result = await this.db
       .select({
         count: count(),
       })
@@ -315,7 +317,7 @@ export class MessageModel {
     }: CreateMessageParams,
     id: string = this.genId(),
   ): Promise<MessageItem> {
-    return serverDB.transaction(async (trx) => {
+    return this.db.transaction(async (trx) => {
       const [item] = (await trx
         .insert(messages)
         .values({
@@ -366,72 +368,72 @@ export class MessageModel {
       return { ...m, userId: this.userId };
     });
 
-    return serverDB.insert(messages).values(messagesToInsert);
+    return this.db.insert(messages).values(messagesToInsert);
   }
 
   async createMessageQuery(params: NewMessageQuery) {
-    const result = await serverDB.insert(messageQueries).values(params).returning();
+    const result = await this.db.insert(messageQueries).values(params).returning();
 
     return result[0];
   }
   // **************** Update *************** //
 
   async update(id: string, message: Partial<MessageItem>) {
-    return serverDB
+    return this.db
       .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({
+    const item = await this.db.query.messagePlugins.findFirst({
       where: eq(messagePlugins.id, id),
     });
     if (!item) throw new Error('Plugin not found');
 
-    return serverDB
+    return this.db
       .update(messagePlugins)
       .set({ state: merge(item.state || {}, state) })
       .where(eq(messagePlugins.id, id));
   }
 
   async updateMessagePlugin(id: string, value: Partial<MessagePluginItem>) {
-    const item = await serverDB.query.messagePlugins.findFirst({
+    const item = await this.db.query.messagePlugins.findFirst({
       where: eq(messagePlugins.id, id),
     });
     if (!item) throw new Error('Plugin not found');
 
-    return serverDB.update(messagePlugins).set(value).where(eq(messagePlugins.id, id));
+    return this.db.update(messagePlugins).set(value).where(eq(messagePlugins.id, id));
   }
 
   async updateTranslate(id: string, translate: Partial<MessageItem>) {
-    const result = await serverDB.query.messageTranslates.findFirst({
+    const result = await this.db.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 });
+      return this.db.insert(messageTranslates).values({ ...translate, id });
     }
 
     // or just update the existing one
-    return serverDB.update(messageTranslates).set(translate).where(eq(messageTranslates.id, id));
+    return this.db.update(messageTranslates).set(translate).where(eq(messageTranslates.id, id));
   }
 
   async updateTTS(id: string, tts: Partial<ChatTTS>) {
-    const result = await serverDB.query.messageTTS.findFirst({
+    const result = await this.db.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
+      return this.db
         .insert(messageTTS)
         .values({ contentMd5: tts.contentMd5, fileId: tts.file, id, voice: tts.voice });
     }
 
     // or just update the existing one
-    return serverDB
+    return this.db
       .update(messageTTS)
       .set({ contentMd5: tts.contentMd5, fileId: tts.file, voice: tts.voice })
       .where(eq(messageTTS.id, id));
@@ -440,7 +442,7 @@ export class MessageModel {
   // **************** Delete *************** //
 
   async deleteMessage(id: string) {
-    return serverDB.transaction(async (tx) => {
+    return this.db.transaction(async (tx) => {
       // 1. 查询要删除的 message 的完整信息
       const message = await tx
         .select()
@@ -476,25 +478,25 @@ export class MessageModel {
   }
 
   async deleteMessages(ids: string[]) {
-    return serverDB
+    return this.db
       .delete(messages)
       .where(and(eq(messages.userId, this.userId), inArray(messages.id, ids)));
   }
 
   async deleteMessageTranslate(id: string) {
-    return serverDB.delete(messageTranslates).where(and(eq(messageTranslates.id, id)));
+    return this.db.delete(messageTranslates).where(and(eq(messageTranslates.id, id)));
   }
 
   async deleteMessageTTS(id: string) {
-    return serverDB.delete(messageTTS).where(and(eq(messageTTS.id, id)));
+    return this.db.delete(messageTTS).where(and(eq(messageTTS.id, id)));
   }
 
   async deleteMessageQuery(id: string) {
-    return serverDB.delete(messageQueries).where(and(eq(messageQueries.id, id)));
+    return this.db.delete(messageQueries).where(and(eq(messageQueries.id, id)));
   }
 
   async deleteMessagesBySession(sessionId?: string | null, topicId?: string | null) {
-    return serverDB
+    return this.db
       .delete(messages)
       .where(
         and(
@@ -506,7 +508,7 @@ export class MessageModel {
   }
 
   async deleteAllMessages() {
-    return serverDB.delete(messages).where(eq(messages.userId, this.userId));
+    return this.db.delete(messages).where(eq(messages.userId, this.userId));
   }
 
   // **************** Helper *************** //

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

@@ -1,20 +1,22 @@
 import { and, desc, eq } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 
 import { InstalledPluginItem, NewInstalledPlugin, installedPlugins } from '../../schemas';
 
 export class PluginModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   create = async (
     params: Pick<NewInstalledPlugin, 'type' | 'identifier' | 'manifest' | 'customParams'>,
   ) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(installedPlugins)
       .values({ ...params, createdAt: new Date(), updatedAt: new Date(), userId: this.userId })
       .returning();
@@ -23,17 +25,17 @@ export class PluginModel {
   };
 
   delete = async (id: string) => {
-    return serverDB
+    return this.db
       .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));
+    return this.db.delete(installedPlugins).where(eq(installedPlugins.userId, this.userId));
   };
 
   query = async () => {
-    return serverDB
+    return this.db
       .select({
         createdAt: installedPlugins.createdAt,
         customParams: installedPlugins.customParams,
@@ -49,13 +51,13 @@ export class PluginModel {
   };
 
   findById = async (id: string) => {
-    return serverDB.query.installedPlugins.findFirst({
+    return this.db.query.installedPlugins.findFirst({
       where: and(eq(installedPlugins.identifier, id), eq(installedPlugins.userId, this.userId)),
     });
   };
 
   async update(id: string, value: Partial<InstalledPluginItem>) {
-    return serverDB
+    return this.db
       .update(installedPlugins)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(installedPlugins.identifier, id), eq(installedPlugins.userId, this.userId)));