fix KnowledgeBase search issue

Author: arvinxx

src/database/server/models/chunk.ts

@@ -207,7 +207,8 @@ export class ChunkModel {
       .leftJoin(files, eq(files.id, fileChunks.fileId))
       .where(inArray(fileChunks.fileId, fileIds))
       .orderBy((t) => desc(t.similarity))
-      .limit(5);
+      // 先放宽到 15
+      .limit(15);
 
     return result.map((item) => {
       return {

src/server/routers/lambda/agent.ts

@@ -122,13 +122,16 @@ export const agentRouter = router({
       const knowledge = await ctx.agentModel.getAgentAssignedKnowledge(input.agentId);
 
       return [
-        ...files.map((file) => ({
-          enabled: knowledge.files.some((item) => item.id === file.id),
-          fileType: file.fileType,
-          id: file.id,
-          name: file.name,
-          type: KnowledgeType.File,
-        })),
+        ...files
+          // 过滤掉所有图片
+          .filter((file) => !file.fileType.startsWith('image'))
+          .map((file) => ({
+            enabled: knowledge.files.some((item) => item.id === file.id),
+            fileType: file.fileType,
+            id: file.id,
+            name: file.name,
+            type: KnowledgeType.File,
+          })),
         ...knowledgeBases.map((knowledgeBase) => ({
           avatar: knowledgeBase.avatar,
           description: knowledgeBase.description,

src/server/routers/lambda/chunk.ts

@@ -1,3 +1,4 @@
+import { TRPCError } from '@trpc/server';
 import { inArray } from 'drizzle-orm/expressions';
 import { z } from 'zod';
 
@@ -126,60 +127,75 @@ export const chunkRouter = router({
   semanticSearchForChat: chunkProcedure
     .input(SemanticSearchSchema)
     .mutation(async ({ ctx, input }) => {
-      const item = await ctx.messageModel.findMessageQueriesById(input.messageId);
-      const { model, provider } =
-        getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
-      let embedding: number[];
-      let ragQueryId: string;
-      // if there is no message rag or it's embeddings, then we need to create one
-      if (!item || !item.embeddings) {
-        // TODO: need to support customize
-        const agentRuntime = await initAgentRuntimeWithUserPayload(provider, ctx.jwtPayload);
-
-        const embeddings = await agentRuntime.embeddings({
-          dimensions: 1024,
-          input: input.rewriteQuery,
-          model,
-        });
-
-        embedding = embeddings![0];
-        const embeddingsId = await ctx.embeddingModel.create({
-          embeddings: embedding,
-          model,
+      try {
+        const item = await ctx.messageModel.findMessageQueriesById(input.messageId);
+        const { model, provider } =
+          getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
+        let embedding: number[];
+        let ragQueryId: string;
+
+        // if there is no message rag or it's embeddings, then we need to create one
+        if (!item || !item.embeddings) {
+          // TODO: need to support customize
+          const agentRuntime = await initAgentRuntimeWithUserPayload(provider, ctx.jwtPayload);
+
+          // slice content to make sure in the context window limit
+          const query =
+            input.rewriteQuery.length > 8000
+              ? input.rewriteQuery.slice(0, 8000)
+              : input.rewriteQuery;
+
+          const embeddings = await agentRuntime.embeddings({
+            dimensions: 1024,
+            input: query,
+            model,
+          });
+
+          embedding = embeddings![0];
+          const embeddingsId = await ctx.embeddingModel.create({
+            embeddings: embedding,
+            model,
+          });
+
+          const result = await ctx.messageModel.createMessageQuery({
+            embeddingsId,
+            messageId: input.messageId,
+            rewriteQuery: input.rewriteQuery,
+            userQuery: input.userQuery,
+          });
+
+          ragQueryId = result.id;
+        } else {
+          embedding = item.embeddings;
+          ragQueryId = item.id;
+        }
+
+        let finalFileIds = input.fileIds ?? [];
+
+        if (input.knowledgeIds && input.knowledgeIds.length > 0) {
+          const knowledgeFiles = await serverDB.query.knowledgeBaseFiles.findMany({
+            where: inArray(knowledgeBaseFiles.knowledgeBaseId, input.knowledgeIds),
+          });
+
+          finalFileIds = knowledgeFiles.map((f) => f.fileId).concat(finalFileIds);
+        }
+
+        const chunks = await ctx.chunkModel.semanticSearchForChat({
+          embedding,
+          fileIds: finalFileIds,
+          query: input.rewriteQuery,
         });
 
-        const result = await ctx.messageModel.createMessageQuery({
-          embeddingsId,
-          messageId: input.messageId,
-          rewriteQuery: input.rewriteQuery,
-          userQuery: input.userQuery,
-        });
+        // TODO: need to rerank the chunks
 
-        ragQueryId = result.id;
-      } else {
-        embedding = item.embeddings;
-        ragQueryId = item.id;
-      }
+        return { chunks, queryId: ragQueryId };
+      } catch (e) {
+        console.error(e);
 
-      console.time('semanticSearch');
-      let finalFileIds = input.fileIds ?? [];
-
-      if (input.knowledgeIds && input.knowledgeIds.length > 0) {
-        const knowledgeFiles = await serverDB.query.knowledgeBaseFiles.findMany({
-          where: inArray(knowledgeBaseFiles.knowledgeBaseId, input.knowledgeIds),
+        throw new TRPCError({
+          code: 'INTERNAL_SERVER_ERROR',
+          message: (e as any).errorType || JSON.stringify(e),
         });
-
-        finalFileIds = knowledgeFiles.map((f) => f.fileId).concat(finalFileIds);
       }
-
-      const chunks = await ctx.chunkModel.semanticSearchForChat({
-        embedding,
-        fileIds: finalFileIds,
-        query: input.rewriteQuery,
-      });
-      // TODO: need to rerank the chunks
-      console.timeEnd('semanticSearch');
-
-      return { chunks, queryId: ragQueryId };
     }),
 });

src/store/chat/slices/aiChat/actions/rag.ts

@@ -21,7 +21,7 @@ export interface ChatRAGAction {
     id: string,
     userQuery: string,
     messages: string[],
-  ) => Promise<{ chunks: ChatSemanticSearchChunk[]; queryId: string; rewriteQuery?: string }>;
+  ) => Promise<{ chunks: ChatSemanticSearchChunk[]; queryId?: string; rewriteQuery?: string }>;
   /**
    * Rewrite user content to better RAG query
    */
@@ -74,17 +74,23 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
 
     // 2. retrieve chunks from semantic search
     const files = chatSelectors.currentUserFiles(get()).map((f) => f.id);
-    const { chunks, queryId } = await ragService.semanticSearchForChat({
-      fileIds: knowledgeIds().fileIds.concat(files),
-      knowledgeIds: knowledgeIds().knowledgeBaseIds,
-      messageId: id,
-      rewriteQuery: rewriteQuery || userQuery,
-      userQuery,
-    });
-
-    get().internal_toggleMessageRAGLoading(false, id);
-
-    return { chunks, queryId, rewriteQuery };
+    try {
+      const { chunks, queryId } = await ragService.semanticSearchForChat({
+        fileIds: knowledgeIds().fileIds.concat(files),
+        knowledgeIds: knowledgeIds().knowledgeBaseIds,
+        messageId: id,
+        rewriteQuery: rewriteQuery || userQuery,
+        userQuery,
+      });
+
+      get().internal_toggleMessageRAGLoading(false, id);
+
+      return { chunks, queryId, rewriteQuery };
+    } catch {
+      get().internal_toggleMessageRAGLoading(false, id);
+
+      return { chunks: [] };
+    }
   },
   internal_rewriteQuery: async (id, content, messages) => {
     let rewriteQuery = content;