@lobehub/chat 1.11.9 → 1.12.0

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

@@ -1,9 +1,17 @@
-import { eq } from 'drizzle-orm';
+import { asc, count, eq, ilike, inArray, notExists } from 'drizzle-orm';
 import { and, desc } from 'drizzle-orm/expressions';
 
 import { serverDB } from '@/database/server/core/db';
+import { FilesTabs, QueryFileListParams, SortType } from '@/types/files';
 
-import { FileItem, NewFile, files } from '../schemas/lobechat';
+import {
+  FileItem,
+  NewFile,
+  NewGlobalFile,
+  files,
+  globalFiles,
+  knowledgeBaseFiles,
+} from '../schemas/lobechat';
 
 export class FileModel {
   private readonly userId: string;
@@ -12,27 +20,204 @@ export class FileModel {
     this.userId = userId;
   }
 
-  create = async (params: Omit<NewFile, 'id' | 'userId'>) => {
-    const result = await serverDB
-      .insert(files)
-      .values({ ...params, userId: this.userId })
-      .returning();
+  create = async (params: Omit<NewFile, 'id' | 'userId'> & { knowledgeBaseId?: string }) => {
+    const result = await serverDB.transaction(async (trx) => {
+      const result = await trx
+        .insert(files)
+        .values({ ...params, userId: this.userId })
+        .returning();
+
+      const item = result[0];
+
+      if (params.knowledgeBaseId) {
+        await trx
+          .insert(knowledgeBaseFiles)
+          .values({ fileId: item.id, knowledgeBaseId: params.knowledgeBaseId });
+      }
+
+      return item;
+    });
+
+    return { id: result.id };
+  };
+
+  createGlobalFile = async (file: Omit<NewGlobalFile, 'id' | 'userId'>) => {
+    return serverDB.insert(globalFiles).values(file).returning();
+  };
 
-    return { id: result[0].id };
+  checkHash = async (hash: string) => {
+    const item = await serverDB.query.globalFiles.findFirst({
+      where: eq(globalFiles.hashId, hash),
+    });
+    if (!item) return { isExist: false };
+
+    return {
+      fileType: item.fileType,
+      isExist: true,
+      metadata: item.metadata,
+      size: item.size,
+      url: item.url,
+    };
   };
 
   delete = async (id: string) => {
-    return serverDB.delete(files).where(and(eq(files.id, id), eq(files.userId, this.userId)));
+    const file = await this.findById(id);
+    if (!file) return;
+
+    const fileHash = file.fileHash!;
+
+    return await serverDB.transaction(async (trx) => {
+      await trx.delete(files).where(and(eq(files.id, id), eq(files.userId, this.userId)));
+
+      const result = await trx
+        .select({ count: count() })
+        .from(files)
+        .where(and(eq(files.fileHash, fileHash)));
+
+      const fileCount = result[0].count;
+
+      // delete the file from global file if it is not used by other files
+      if (fileCount === 0) {
+        await trx.delete(globalFiles).where(eq(globalFiles.hashId, fileHash));
+
+        return file;
+      }
+    });
+  };
+
+  deleteGlobalFile = async (hashId: string) => {
+    return serverDB.delete(globalFiles).where(eq(globalFiles.hashId, hashId));
+  };
+
+  deleteMany = async (ids: string[]) => {
+    const fileList = await this.findByIds(ids);
+    const hashList = fileList.map((file) => file.fileHash!);
+
+    return await serverDB.transaction(async (trx) => {
+      // delete the files
+      await trx.delete(files).where(and(inArray(files.id, ids), eq(files.userId, this.userId)));
+
+      // count the files by hash
+      const result = await trx
+        .select({
+          count: count(),
+          hashId: files.fileHash,
+        })
+        .from(files)
+        .where(inArray(files.fileHash, hashList))
+        .groupBy(files.fileHash);
+
+      // Create a Map to store the query result
+      const countMap = new Map(result.map((item) => [item.hashId, item.count]));
+
+      // Ensure that all incoming hashes have a result, even if it is 0
+      const fileHashCounts = hashList.map((hashId) => ({
+        count: countMap.get(hashId) || 0,
+        hashId: hashId,
+      }));
+
+      const needToDeleteList = fileHashCounts.filter((item) => item.count === 0);
+
+      if (needToDeleteList.length === 0) return;
+
+      // delete the file from global file if it is not used by other files
+      await trx.delete(globalFiles).where(
+        inArray(
+          globalFiles.hashId,
+          needToDeleteList.map((item) => item.hashId!),
+        ),
+      );
+
+      return fileList.filter((file) =>
+        needToDeleteList.some((item) => item.hashId === file.fileHash),
+      );
+    });
   };
 
   clear = async () => {
     return serverDB.delete(files).where(eq(files.userId, this.userId));
   };
 
-  query = async () => {
+  query = async ({
+    category,
+    q,
+    sortType,
+    sorter,
+    knowledgeBaseId,
+    showFilesInKnowledgeBase,
+  }: QueryFileListParams = {}) => {
+    // 1. query where
+    let whereClause = and(
+      q ? ilike(files.name, `%${q}%`) : undefined,
+      eq(files.userId, this.userId),
+    );
+    if (category && category !== FilesTabs.All) {
+      const fileTypePrefix = this.getFileTypePrefix(category as FilesTabs);
+      whereClause = and(whereClause, ilike(files.fileType, `${fileTypePrefix}%`));
+    }
+
+    // 2. order part
+
+    let orderByClause = desc(files.createdAt);
+    // create a map for sortable fields
+    const sortableFields = {
+      createdAt: files.createdAt,
+      name: files.name,
+      size: files.size,
+      updatedAt: files.updatedAt,
+    } as const;
+    type SortableField = keyof typeof sortableFields;
+
+    if (sorter && sortType && sorter in sortableFields) {
+      const sortFunction = sortType.toLowerCase() === SortType.Asc ? asc : desc;
+      orderByClause = sortFunction(sortableFields[sorter as SortableField]);
+    }
+
+    // 3. build query
+    let query = serverDB
+      .select({
+        chunkTaskId: files.chunkTaskId,
+        createdAt: files.createdAt,
+        embeddingTaskId: files.embeddingTaskId,
+        fileType: files.fileType,
+        id: files.id,
+        name: files.name,
+        size: files.size,
+        updatedAt: files.updatedAt,
+        url: files.url,
+      })
+      .from(files);
+
+    // 4. add knowledge base query
+    if (knowledgeBaseId) {
+      // if knowledgeBaseId is provided, it means we are querying files in a knowledge-base
+
+      // @ts-ignore
+      query = query.innerJoin(
+        knowledgeBaseFiles,
+        and(
+          eq(files.id, knowledgeBaseFiles.fileId),
+          eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
+        ),
+      );
+    }
+    // 5.if we don't show files in knowledge base, we need exclude files in knowledge base
+    else if (!showFilesInKnowledgeBase) {
+      whereClause = and(
+        whereClause,
+        notExists(
+          serverDB.select().from(knowledgeBaseFiles).where(eq(knowledgeBaseFiles.fileId, files.id)),
+        ),
+      );
+    }
+
+    // or we are just filter in the global files
+    return query.where(whereClause).orderBy(orderByClause);
+  };
+
+  findByIds = async (ids: string[]) => {
     return serverDB.query.files.findMany({
-      orderBy: [desc(files.updatedAt)],
-      where: eq(files.userId, this.userId),
+      where: and(inArray(files.id, ids), eq(files.userId, this.userId)),
     });
   };
 
@@ -42,10 +227,44 @@ export class FileModel {
     });
   };
 
+  countFilesByHash = async (hash: string) => {
+    const result = await serverDB
+      .select({
+        count: count(),
+      })
+      .from(files)
+      .where(and(eq(files.fileHash, hash)));
+
+    return result[0].count;
+  };
+
   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)));
   }
+
+  /**
+   * get the corresponding file type prefix according to FilesTabs
+   */
+  private getFileTypePrefix = (category: FilesTabs): string => {
+    switch (category) {
+      case FilesTabs.Audios: {
+        return 'audio';
+      }
+      case FilesTabs.Documents: {
+        return 'application';
+      }
+      case FilesTabs.Images: {
+        return 'image';
+      }
+      case FilesTabs.Videos: {
+        return 'video';
+      }
+      default: {
+        return '';
+      }
+    }
+  };
 }

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

@@ -0,0 +1,94 @@
+import { eq, inArray } from 'drizzle-orm';
+import { and, desc } from 'drizzle-orm/expressions';
+
+import { serverDB } from '@/database/server';
+import { KnowledgeBaseItem } from '@/types/knowledgeBase';
+
+import { NewKnowledgeBase, knowledgeBaseFiles, knowledgeBases } from '../schemas/lobechat';
+
+export class KnowledgeBaseModel {
+  private userId: string;
+
+  constructor(userId: string) {
+    this.userId = userId;
+  }
+
+  // create
+
+  create = async (params: Omit<NewKnowledgeBase, 'userId'>) => {
+    const [result] = await serverDB
+      .insert(knowledgeBases)
+      .values({ ...params, userId: this.userId })
+      .returning();
+
+    return result;
+  };
+
+  addFilesToKnowledgeBase = async (id: string, fileIds: string[]) => {
+    return serverDB
+      .insert(knowledgeBaseFiles)
+      .values(fileIds.map((fileId) => ({ fileId, knowledgeBaseId: id, userId: this.userId })))
+      .returning();
+  };
+
+  // delete
+  delete = async (id: string) => {
+    return serverDB
+      .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));
+  };
+
+  removeFilesFromKnowledgeBase = async (knowledgeBaseId: string, ids: string[]) => {
+    return serverDB.delete(knowledgeBaseFiles).where(
+      and(
+        eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
+        inArray(knowledgeBaseFiles.fileId, ids),
+        // eq(knowledgeBaseFiles.userId, this.userId),
+      ),
+    );
+  };
+  // query
+  query = async () => {
+    const data = await serverDB
+      .select({
+        avatar: knowledgeBases.avatar,
+        createdAt: knowledgeBases.createdAt,
+        description: knowledgeBases.description,
+        id: knowledgeBases.id,
+        isPublic: knowledgeBases.isPublic,
+        name: knowledgeBases.name,
+        settings: knowledgeBases.settings,
+        type: knowledgeBases.type,
+        updatedAt: knowledgeBases.updatedAt,
+      })
+      .from(knowledgeBases)
+      .where(eq(knowledgeBases.userId, this.userId))
+      .orderBy(desc(knowledgeBases.updatedAt));
+
+    return data as KnowledgeBaseItem[];
+  };
+
+  findById = async (id: string) => {
+    return serverDB.query.knowledgeBases.findFirst({
+      where: and(eq(knowledgeBases.id, id), eq(knowledgeBases.userId, this.userId)),
+    });
+  };
+
+  // update
+  async update(id: string, value: Partial<KnowledgeBaseItem>) {
+    return serverDB
+      .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({
+      where: eq(knowledgeBases.id, id),
+    });
+  }
+}

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

@@ -1,21 +1,33 @@
-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 { count } from 'drizzle-orm';
+import { and, asc, desc, eq, gte, inArray, isNull, like, lt } from 'drizzle-orm/expressions';
 
-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 { getFullFileUrl } from '@/server/utils/files';
+import {
+  ChatFileItem,
+  ChatImageItem,
+  ChatTTS,
+  ChatToolPayload,
+  CreateMessageParams,
+} from '@/types/message';
 import { merge } from '@/utils/merge';
 
 import {
   MessageItem,
   MessagePluginItem,
-  filesToMessages,
+  NewMessageQuery,
+  chunks,
+  embeddings,
+  fileChunks,
+  files,
   messagePlugins,
+  messageQueries,
+  messageQueryChunks,
   messageTTS,
   messageTranslates,
   messages,
+  messagesFiles,
 } from '../schemas/lobechat';
 
 export interface QueryMessageParams {
@@ -41,6 +53,7 @@ export class MessageModel {
   }: QueryMessageParams = {}): Promise<MessageItem[]> {
     const offset = current * pageSize;
 
+    // 1. get basic messages
     const result = await serverDB
       .select({
         /* eslint-disable sort-keys-fix/sort-keys-fix*/
@@ -102,13 +115,55 @@ export class MessageModel {
 
     if (messageIds.length === 0) return result;
 
-    const fileIds = await serverDB
+    // 2. get relative files
+    const rawRelatedFileList = await serverDB
       .select({
-        fileId: filesToMessages.fileId,
-        messageId: filesToMessages.messageId,
+        fileType: files.fileType,
+        id: messagesFiles.fileId,
+        messageId: messagesFiles.messageId,
+        name: files.name,
+        size: files.size,
+        url: files.url,
       })
-      .from(filesToMessages)
-      .where(inArray(filesToMessages.messageId, messageIds));
+      .from(messagesFiles)
+      .leftJoin(files, eq(files.id, messagesFiles.fileId))
+      .where(inArray(messagesFiles.messageId, messageIds));
+
+    const relatedFileList = rawRelatedFileList.map((file) => ({
+      ...file,
+      url: getFullFileUrl(file.url),
+    }));
+
+    const imageList = relatedFileList.filter((i) => (i.fileType || '').startsWith('image'));
+    const fileList = relatedFileList.filter((i) => !(i.fileType || '').startsWith('image'));
+
+    // 3. get relative file chunks
+    const chunksList = await serverDB
+      .select({
+        fileId: files.id,
+        fileType: files.fileType,
+        fileUrl: files.url,
+        filename: files.name,
+        id: chunks.id,
+        messageId: messageQueryChunks.messageId,
+        text: chunks.text,
+      })
+      .from(messageQueryChunks)
+      .leftJoin(chunks, eq(chunks.id, messageQueryChunks.chunkId))
+      .leftJoin(fileChunks, eq(fileChunks.chunkId, chunks.id))
+      .innerJoin(files, eq(fileChunks.fileId, files.id))
+      .where(inArray(messageQueryChunks.messageId, messageIds));
+
+    // 3. get relative message query
+    const messageQueriesList = await serverDB
+      .select({
+        id: messageQueries.id,
+        messageId: messageQueries.messageId,
+        rewriteQuery: messageQueries.rewriteQuery,
+        userQuery: messageQueries.userQuery,
+      })
+      .from(messageQueries)
+      .where(inArray(messageQueries.messageId, messageIds));
 
     return result.map(
       ({
@@ -118,22 +173,45 @@ export class MessageModel {
         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),
-      }),
+      }) => {
+        const messageQuery = messageQueriesList.find((relation) => relation.messageId === item.id);
+        return {
+          ...item,
+          chunksList: chunksList.filter((relation) => relation.messageId === item.id),
+
+          extra: {
+            fromModel: model,
+            fromProvider: provider,
+            translate,
+            tts: ttsId
+              ? {
+                  // contentMd5: ttsContentMd5,
+                  // file: ttsFile,
+                  // voice: ttsVoice,
+                }
+              : undefined,
+          },
+          fileList: fileList
+            .filter((relation) => relation.messageId === item.id)
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            .map<ChatFileItem>(({ id, url, size, fileType, name }) => ({
+              fileType: fileType!,
+              id,
+              name: name!,
+              size: size!,
+              url,
+            })),
+
+          imageList: imageList
+            .filter((relation) => relation.messageId === item.id)
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
+            .map<ChatImageItem>(({ id, url, name }) => ({ alt: name!, id, url })),
+
+          ragQuery: messageQuery?.rewriteQuery,
+          ragQueryId: messageQuery?.id,
+          ragRawQuery: messageQuery?.userQuery,
+        };
+      },
     );
   }
 
@@ -143,6 +221,24 @@ export class MessageModel {
     });
   }
 
+  async findMessageQueriesById(messageId: string) {
+    const result = await serverDB
+      .select({
+        embeddings: embeddings.embeddings,
+        id: messageQueries.id,
+        query: messageQueries.rewriteQuery,
+        rewriteQuery: messageQueries.rewriteQuery,
+        userQuery: messageQueries.userQuery,
+      })
+      .from(messageQueries)
+      .where(and(eq(messageQueries.messageId, messageId)))
+      .leftJoin(embeddings, eq(embeddings.id, messageQueries.embeddingsId));
+
+    if (result.length === 0) return undefined;
+
+    return result[0];
+  }
+
   async queryAll(): Promise<MessageItem[]> {
     return serverDB
       .select()
@@ -195,7 +291,8 @@ export class MessageModel {
       .where(
         and(
           eq(messages.userId, this.userId),
-          sql`${messages.createdAt} >= ${today} AND ${messages.createdAt} < ${tomorrow}`,
+          gte(messages.createdAt, today),
+          lt(messages.createdAt, tomorrow),
         ),
       )
       .execute();
@@ -206,7 +303,16 @@ export class MessageModel {
   // **************** Create *************** //
 
   async create(
-    { fromModel, fromProvider, files, plugin, pluginState, ...message }: CreateMessageParams,
+    {
+      fromModel,
+      fromProvider,
+      files,
+      plugin,
+      pluginState,
+      fileChunks,
+      ragQueryId,
+      ...message
+    }: CreateMessageParams,
     id: string = this.genId(),
   ): Promise<MessageItem> {
     return serverDB.transaction(async (trx) => {
@@ -236,10 +342,21 @@ export class MessageModel {
 
       if (files && files.length > 0) {
         await trx
-          .insert(filesToMessages)
+          .insert(messagesFiles)
           .values(files.map((file) => ({ fileId: file, messageId: id })));
       }
 
+      if (fileChunks && fileChunks.length > 0 && ragQueryId) {
+        await trx.insert(messageQueryChunks).values(
+          fileChunks.map((chunk) => ({
+            chunkId: chunk.id,
+            messageId: id,
+            queryId: ragQueryId,
+            similarity: chunk.similarity.toString(),
+          })),
+        );
+      }
+
       return item;
     });
   }
@@ -252,6 +369,11 @@ export class MessageModel {
     return serverDB.insert(messages).values(messagesToInsert);
   }
 
+  async createMessageQuery(params: NewMessageQuery) {
+    const result = await serverDB.insert(messageQueries).values(params).returning();
+
+    return result[0];
+  }
   // **************** Update *************** //
 
   async update(id: string, message: Partial<MessageItem>) {
@@ -367,6 +489,10 @@ export class MessageModel {
     return serverDB.delete(messageTTS).where(and(eq(messageTTS.id, id)));
   }
 
+  async deleteMessageQuery(id: string) {
+    return serverDB.delete(messageQueries).where(and(eq(messageQueries.id, id)));
+  }
+
   async deleteMessagesBySession(sessionId?: string | null, topicId?: string | null) {
     return serverDB
       .delete(messages)

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

@@ -4,7 +4,7 @@ import { DeepPartial } from 'utility-types';
 
 import { serverDB } from '@/database/server/core/db';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
-import { UserPreference } from '@/types/user';
+import { UserGuide, UserPreference } from '@/types/user';
 import { UserSettings } from '@/types/user/settings';
 import { merge } from '@/utils/merge';
 
@@ -145,4 +145,15 @@ export class UserModel {
       .set({ preference: merge(user.preference, value) })
       .where(eq(users.id, id));
   }
+
+  async updateGuide(id: string, value: Partial<UserGuide>) {
+    const user = await serverDB.query.users.findFirst({ where: eq(users.id, id) });
+    if (!user) return;
+
+    const prevPreference = (user.preference || {}) as UserPreference;
+    return serverDB
+      .update(users)
+      .set({ preference: { ...prevPreference, guide: merge(prevPreference.guide || {}, value) } })
+      .where(eq(users.id, id));
+  }
 }