@lobehub/chat 1.11.8 → 1.12.0

package/src/server/routers/lambda/agent.ts

@@ -0,0 +1,159 @@
+import { z } from 'zod';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import { AgentModel } from '@/database/server/models/agent';
+import { FileModel } from '@/database/server/models/file';
+import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
+import { SessionModel } from '@/database/server/models/session';
+import { pino } from '@/libs/logger';
+import { authedProcedure, router } from '@/libs/trpc';
+import { KnowledgeItem, KnowledgeType } from '@/types/knowledgeBase';
+
+const agentProcedure = authedProcedure.use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      agentModel: new AgentModel(ctx.userId),
+      fileModel: new FileModel(ctx.userId),
+      knowledgeBaseModel: new KnowledgeBaseModel(ctx.userId),
+      sessionModel: new SessionModel(ctx.userId),
+    },
+  });
+});
+
+export const agentRouter = router({
+  createAgentFiles: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        enabled: z.boolean().optional(),
+        fileIds: z.array(z.string()),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.createAgentFiles(input.agentId, input.fileIds, input.enabled);
+    }),
+
+  createAgentKnowledgeBase: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        enabled: z.boolean().optional(),
+        knowledgeBaseId: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.createAgentKnowledgeBase(
+        input.agentId,
+        input.knowledgeBaseId,
+        input.enabled,
+      );
+    }),
+
+  deleteAgentFile: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        fileId: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.deleteAgentFile(input.agentId, input.fileId);
+    }),
+
+  deleteAgentKnowledgeBase: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        knowledgeBaseId: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.deleteAgentKnowledgeBase(input.agentId, input.knowledgeBaseId);
+    }),
+
+  getAgentConfig: agentProcedure
+    .input(
+      z.object({
+        sessionId: z.string(),
+      }),
+    )
+    .query(async ({ input, ctx }) => {
+      if (input.sessionId === INBOX_SESSION_ID) {
+        const item = await ctx.sessionModel.findByIdOrSlug(INBOX_SESSION_ID);
+        // if there is no session for user, create one
+        if (!item) {
+          const res = await ctx.sessionModel.createInbox();
+          pino.info('create inbox session', res);
+        }
+      }
+
+      const session = await ctx.sessionModel.findByIdOrSlug(input.sessionId);
+
+      if (!session) throw new Error('Session not found');
+      const sessionId = session.id;
+
+      return await ctx.agentModel.findBySessionId(sessionId);
+    }),
+
+  getKnowledgeBasesAndFiles: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+      }),
+    )
+    .query(async ({ ctx, input }): Promise<KnowledgeItem[]> => {
+      const knowledgeBases = await ctx.knowledgeBaseModel.query();
+
+      const files = await ctx.fileModel.query({ showFilesInKnowledgeBase: false });
+
+      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,
+        })),
+        ...knowledgeBases.map((knowledgeBase) => ({
+          avatar: knowledgeBase.avatar,
+          description: knowledgeBase.description,
+          enabled: knowledge.knowledgeBases.some((item) => item.id === knowledgeBase.id),
+          id: knowledgeBase.id,
+          name: knowledgeBase.name,
+          type: KnowledgeType.KnowledgeBase,
+        })),
+      ];
+    }),
+
+  toggleFile: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        enabled: z.boolean().optional(),
+        fileId: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.toggleFile(input.agentId, input.fileId, input.enabled);
+    }),
+
+  toggleKnowledgeBase: agentProcedure
+    .input(
+      z.object({
+        agentId: z.string(),
+        enabled: z.boolean().optional(),
+        knowledgeBaseId: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.agentModel.toggleKnowledgeBase(
+        input.agentId,
+        input.knowledgeBaseId,
+        input.enabled,
+      );
+    }),
+});

package/src/server/routers/lambda/chunk.ts

@@ -0,0 +1,189 @@
+import { inArray } from 'drizzle-orm';
+import { z } from 'zod';
+
+import { initAgentRuntimeWithUserPayload } from '@/app/api/chat/agentRuntime';
+import { DEFAULT_EMBEDDING_MODEL } from '@/const/settings';
+import { serverDB } from '@/database/server';
+import { AsyncTaskModel } from '@/database/server/models/asyncTask';
+import { ChunkModel } from '@/database/server/models/chunk';
+import { EmbeddingModel } from '@/database/server/models/embedding';
+import { FileModel } from '@/database/server/models/file';
+import { MessageModel } from '@/database/server/models/message';
+import { knowledgeBaseFiles } from '@/database/server/schemas/lobechat';
+import { ModelProvider } from '@/libs/agent-runtime';
+import { authedProcedure, router } from '@/libs/trpc';
+import { keyVaults } from '@/libs/trpc/middleware/keyVaults';
+import { ChunkService } from '@/server/services/chunk';
+import { SemanticSearchSchema } from '@/types/rag';
+
+const chunkProcedure = authedProcedure.use(keyVaults).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      asyncTaskModel: new AsyncTaskModel(ctx.userId),
+      chunkModel: new ChunkModel(ctx.userId),
+      chunkService: new ChunkService(ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.userId),
+      fileModel: new FileModel(ctx.userId),
+      messageModel: new MessageModel(ctx.userId),
+    },
+  });
+});
+
+export const chunkRouter = router({
+  createEmbeddingChunksTask: chunkProcedure
+    .input(
+      z.object({
+        id: z.string(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const asyncTaskId = await ctx.chunkService.asyncEmbeddingFileChunks(input.id, ctx.jwtPayload);
+
+      return { id: asyncTaskId, success: true };
+    }),
+
+  createParseFileTask: chunkProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        skipExist: z.boolean().optional(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(
+        input.id,
+        ctx.jwtPayload,
+        input.skipExist,
+      );
+
+      return { id: asyncTaskId, success: true };
+    }),
+
+  getChunksByFileId: chunkProcedure
+    .input(
+      z.object({
+        cursor: z.number().nullish(),
+        id: z.string(),
+      }),
+    )
+    .query(async ({ ctx, input }) => {
+      return {
+        items: await ctx.chunkModel.findByFileId(input.id, input.cursor || 0),
+        nextCursor: input.cursor ? input.cursor + 1 : 1,
+      };
+    }),
+
+  retryParseFileTask: chunkProcedure
+    .input(
+      z.object({
+        id: z.string(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const result = await ctx.fileModel.findById(input.id);
+
+      if (!result) return;
+
+      // 1. delete the previous task if exist
+      if (result.chunkTaskId) {
+        await ctx.asyncTaskModel.delete(result.chunkTaskId);
+      }
+
+      // 2. create a new asyncTask for chunking
+      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(input.id, ctx.jwtPayload);
+
+      return { id: asyncTaskId, success: true };
+    }),
+
+  semanticSearch: chunkProcedure
+    .input(
+      z.object({
+        fileIds: z.array(z.string()).optional(),
+        model: z.string().default(DEFAULT_EMBEDDING_MODEL),
+        query: z.string(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      console.time('embedding');
+      const agentRuntime = await initAgentRuntimeWithUserPayload(
+        ModelProvider.OpenAI,
+        ctx.jwtPayload,
+      );
+
+      const embeddings = await agentRuntime.embeddings({
+        dimensions: 1024,
+        input: input.query,
+        model: input.model,
+      });
+      console.timeEnd('embedding');
+
+      return ctx.chunkModel.semanticSearch({
+        embedding: embeddings![0].embedding,
+        fileIds: input.fileIds,
+        query: input.query,
+      });
+    }),
+
+  semanticSearchForChat: chunkProcedure
+    .input(SemanticSearchSchema)
+    .mutation(async ({ ctx, input }) => {
+      const item = await ctx.messageModel.findMessageQueriesById(input.messageId);
+      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(
+          ModelProvider.OpenAI,
+          ctx.jwtPayload,
+        );
+
+        const embeddings = await agentRuntime.embeddings({
+          dimensions: 1024,
+          input: input.rewriteQuery,
+          model: input.model || DEFAULT_EMBEDDING_MODEL,
+        });
+
+        embedding = embeddings![0].embedding;
+        const embeddingsId = await ctx.embeddingModel.create({
+          embeddings: embedding,
+          model: input.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;
+      }
+
+      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),
+        });
+
+        finalFileIds = knowledgeFiles.map((f) => f.fileId).concat(finalFileIds);
+      }
+
+      const chunks = await ctx.chunkModel.semanticSearchForChat({
+        embedding,
+        fileIds: finalFileIds,
+        query: input.rewriteQuery,
+      });
+      console.timeEnd('semanticSearch');
+
+      return { chunks, queryId: ragQueryId };
+    }),
+});

package/src/server/routers/lambda/file.ts

@@ -1,32 +1,66 @@
+import { TRPCError } from '@trpc/server';
+import urlJoin from 'url-join';
 import { z } from 'zod';
 
+import { fileEnv } from '@/config/file';
+import { AsyncTaskModel } from '@/database/server/models/asyncTask';
+import { ChunkModel } from '@/database/server/models/chunk';
 import { FileModel } from '@/database/server/models/file';
 import { authedProcedure, router } from '@/libs/trpc';
-import { UploadFileSchema } from '@/types/files';
+import { S3 } from '@/server/modules/S3';
+import { getFullFileUrl } from '@/server/utils/files';
+import { AsyncTaskStatus, AsyncTaskType } from '@/types/asyncTask';
+import { FileListItem, QueryFileListSchema, UploadFileSchema } from '@/types/files';
 
 const fileProcedure = authedProcedure.use(async (opts) => {
   const { ctx } = opts;
 
   return opts.next({
-    ctx: { fileModel: new FileModel(ctx.userId) },
+    ctx: {
+      asyncTaskModel: new AsyncTaskModel(ctx.userId),
+      chunkModel: new ChunkModel(ctx.userId),
+      fileModel: new FileModel(ctx.userId),
+    },
   });
 });
 
 export const fileRouter = router({
+  checkFileHash: fileProcedure
+    .input(z.object({ hash: z.string() }))
+    .mutation(async ({ ctx, input }) => {
+      return ctx.fileModel.checkHash(input.hash);
+    }),
+
   createFile: fileProcedure
     .input(
       UploadFileSchema.omit({ data: true, saveMode: true, url: true }).extend({ url: z.string() }),
     )
     .mutation(async ({ ctx, input }) => {
-      return ctx.fileModel.create({
+      const { isExist } = await ctx.fileModel.checkHash(input.hash!);
+
+      // if the file is not exist in global file, create a new one
+      if (!isExist) {
+        await ctx.fileModel.createGlobalFile({
+          fileType: input.fileType,
+          hashId: input.hash!,
+          metadata: input.metadata,
+          size: input.size,
+          url: input.url,
+        });
+      }
+
+      const { id } = await ctx.fileModel.create({
+        fileHash: input.hash,
         fileType: input.fileType,
+        knowledgeBaseId: input.knowledgeBaseId,
         metadata: input.metadata,
         name: input.name,
         size: input.size,
         url: input.url,
       });
-    }),
 
+      return { id, url: getFullFileUrl(input.url) };
+    }),
   findById: fileProcedure
     .input(
       z.object({
@@ -37,13 +71,103 @@ export const fileRouter = router({
       return ctx.fileModel.findById(input.id);
     }),
 
+  getFileItemById: fileProcedure
+    .input(
+      z.object({
+        id: z.string(),
+      }),
+    )
+    .query(async ({ ctx, input }): Promise<FileListItem | undefined> => {
+      const item = await ctx.fileModel.findById(input.id);
+
+      if (!item) throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
+
+      let embeddingTask = null;
+      if (item.embeddingTaskId) {
+        embeddingTask = await ctx.asyncTaskModel.findById(item.embeddingTaskId);
+      }
+      let chunkingTask = null;
+      if (item.chunkTaskId) {
+        chunkingTask = await ctx.asyncTaskModel.findById(item.chunkTaskId);
+      }
+
+      const chunkCount = await ctx.chunkModel.countByFileId(input.id);
+
+      return {
+        ...item,
+        chunkCount,
+        chunkingError: chunkingTask?.error,
+        chunkingStatus: chunkingTask?.status as AsyncTaskStatus,
+        embeddingError: embeddingTask?.error,
+        embeddingStatus: embeddingTask?.status as AsyncTaskStatus,
+        finishEmbedding: embeddingTask?.status === AsyncTaskStatus.Success,
+        url: getFullFileUrl(item.url!),
+      };
+    }),
+
+  getFiles: fileProcedure.input(QueryFileListSchema).query(async ({ ctx, input }) => {
+    const fileList = await ctx.fileModel.query(input);
+
+    const fileIds = fileList.map((item) => item.id);
+    const chunks = await ctx.chunkModel.countByFileIds(fileIds);
+
+    const chunkTaskIds = fileList.map((result) => result.chunkTaskId).filter(Boolean) as string[];
+
+    const chunkTasks = await ctx.asyncTaskModel.findByIds(chunkTaskIds, AsyncTaskType.Chunking);
+
+    const embeddingTaskIds = fileList
+      .map((result) => result.embeddingTaskId)
+      .filter(Boolean) as string[];
+    const embeddingTasks = await ctx.asyncTaskModel.findByIds(
+      embeddingTaskIds,
+      AsyncTaskType.Embedding,
+    );
+
+    return fileList.map(({ chunkTaskId, embeddingTaskId, ...item }): FileListItem => {
+      const chunkTask = chunkTaskId ? chunkTasks.find((task) => task.id === chunkTaskId) : null;
+      const embeddingTask = embeddingTaskId
+        ? embeddingTasks.find((task) => task.id === embeddingTaskId)
+        : null;
+
+      return {
+        ...item,
+        chunkCount: chunks.find((chunk) => chunk.id === item.id)?.count ?? null,
+        chunkingError: chunkTask?.error ?? null,
+        chunkingStatus: chunkTask?.status as AsyncTaskStatus,
+        embeddingError: embeddingTask?.error ?? null,
+        embeddingStatus: embeddingTask?.status as AsyncTaskStatus,
+        finishEmbedding: embeddingTask?.status === AsyncTaskStatus.Success,
+        url: urlJoin(fileEnv.NEXT_PUBLIC_S3_DOMAIN!, item.url!),
+      };
+    });
+  }),
+
   removeAllFiles: fileProcedure.mutation(async ({ ctx }) => {
     return ctx.fileModel.clear();
   }),
 
   removeFile: fileProcedure.input(z.object({ id: z.string() })).mutation(async ({ input, ctx }) => {
-    return ctx.fileModel.delete(input.id);
+    const file = await ctx.fileModel.delete(input.id);
+
+    if (!file) return;
+
+    // delele the file from remove from S3 if it is not used by other files
+    const s3Client = new S3();
+    await s3Client.deleteFile(file.url!);
   }),
+
+  removeFiles: fileProcedure
+    .input(z.object({ ids: z.array(z.string()) }))
+    .mutation(async ({ input, ctx }) => {
+      const needToRemoveFileList = await ctx.fileModel.deleteMany(input.ids);
+
+      if (!needToRemoveFileList || needToRemoveFileList.length === 0) return;
+
+      // remove from S3
+      const s3Client = new S3();
+
+      await s3Client.deleteFiles(needToRemoveFileList.map((file) => file.url!));
+    }),
 });
 
 export type FileRouter = typeof fileRouter;

package/src/server/routers/lambda/index.ts

@@ -3,9 +3,12 @@
  */
 import { publicProcedure, router } from '@/libs/trpc';
 
+import { agentRouter } from './agent';
+import { chunkRouter } from './chunk';
 // router that connect to db
 import { fileRouter } from './file';
 import { importerRouter } from './importer';
+import { knowledgeBaseRouter } from './knowledgeBase';
 import { messageRouter } from './message';
 import { pluginRouter } from './plugin';
 import { sessionRouter } from './session';
@@ -14,9 +17,12 @@ import { topicRouter } from './topic';
 import { userRouter } from './user';
 
 export const lambdaRouter = router({
+  agent: agentRouter,
+  chunk: chunkRouter,
   file: fileRouter,
   healthcheck: publicProcedure.query(() => "i'm live!"),
   importer: importerRouter,
+  knowledgeBase: knowledgeBaseRouter,
   message: messageRouter,
   plugin: pluginRouter,
   session: sessionRouter,

package/src/server/routers/lambda/knowledgeBase.ts

@@ -0,0 +1,79 @@
+import { z } from 'zod';
+
+import { KnowledgeBaseModel } from '@/database/server/models/knowledgeBase';
+import { insertKnowledgeBasesSchema } from '@/database/server/schemas/lobechat';
+import { authedProcedure, router } from '@/libs/trpc';
+import { KnowledgeBaseItem } from '@/types/knowledgeBase';
+
+const knowledgeBaseProcedure = authedProcedure.use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      knowledgeBaseModel: new KnowledgeBaseModel(ctx.userId),
+    },
+  });
+});
+
+export const knowledgeBaseRouter = router({
+  addFilesToKnowledgeBase: knowledgeBaseProcedure
+    .input(z.object({ ids: z.array(z.string()), knowledgeBaseId: z.string() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.knowledgeBaseModel.addFilesToKnowledgeBase(input.knowledgeBaseId, input.ids);
+    }),
+
+  createKnowledgeBase: knowledgeBaseProcedure
+    .input(
+      z.object({
+        avatar: z.string().optional(),
+        description: z.string().optional(),
+        name: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const data = await ctx.knowledgeBaseModel.create({
+        avatar: input.avatar,
+        description: input.description,
+        name: input.name,
+      });
+
+      return data?.id;
+    }),
+
+  getKnowledgeBaseById: knowledgeBaseProcedure
+    .input(z.object({ id: z.string() }))
+    .query(async ({ ctx, input }): Promise<KnowledgeBaseItem | undefined> => {
+      return ctx.knowledgeBaseModel.findById(input.id);
+    }),
+
+  getKnowledgeBases: knowledgeBaseProcedure.query(async ({ ctx }): Promise<KnowledgeBaseItem[]> => {
+    return ctx.knowledgeBaseModel.query();
+  }),
+
+  removeAllKnowledgeBases: knowledgeBaseProcedure.mutation(async ({ ctx }) => {
+    return ctx.knowledgeBaseModel.deleteAll();
+  }),
+
+  removeFilesFromKnowledgeBase: knowledgeBaseProcedure
+    .input(z.object({ ids: z.array(z.string()), knowledgeBaseId: z.string() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.knowledgeBaseModel.removeFilesFromKnowledgeBase(input.knowledgeBaseId, input.ids);
+    }),
+
+  removeKnowledgeBase: knowledgeBaseProcedure
+    .input(z.object({ id: z.string(), removeFiles: z.boolean().optional() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.knowledgeBaseModel.delete(input.id);
+    }),
+
+  updateKnowledgeBase: knowledgeBaseProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        value: insertKnowledgeBasesSchema.partial(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.knowledgeBaseModel.update(input.id, input.value);
+    }),
+});

package/src/server/routers/lambda/message.ts

@@ -81,6 +81,12 @@ export const messageRouter = router({
       return ctx.messageModel.deleteMessage(input.id);
     }),
 
+  removeMessageQuery: messageProcedure
+    .input(z.object({ id: z.string() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.messageModel.deleteMessageQuery(input.id);
+    }),
+
   removeMessages: messageProcedure
     .input(z.object({ ids: z.array(z.string()) }))
     .mutation(async ({ input, ctx }) => {

package/src/server/routers/lambda/session.ts

@@ -1,10 +1,8 @@
 import { z } from 'zod';
 
-import { INBOX_SESSION_ID } from '@/const/session';
 import { SessionModel } from '@/database/server/models/session';
 import { SessionGroupModel } from '@/database/server/models/sessionGroup';
 import { insertAgentSchema, insertSessionSchema } from '@/database/server/schemas/lobechat';
-import { pino } from '@/libs/logger';
 import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
 import { AgentChatConfigSchema } from '@/types/agent';
 import { LobeMetaDataSchema } from '@/types/meta';
@@ -91,29 +89,6 @@ export const sessionRouter = router({
     return sessionModel.queryWithGroups();
   }),
 
-  getSessionConfig: sessionProcedure
-    .input(
-      z.object({
-        id: z.string(),
-      }),
-    )
-    .query(async ({ input, ctx }) => {
-      if (input.id === INBOX_SESSION_ID) {
-        const item = await ctx.sessionModel.findByIdOrSlug(INBOX_SESSION_ID);
-        // if there is no session for user, create one
-        if (!item) {
-          const res = await ctx.sessionModel.createInbox();
-          pino.info('create inbox session', res);
-        }
-      }
-
-      const session = await ctx.sessionModel.findByIdOrSlug(input.id);
-
-      if (!session) throw new Error('Session not found');
-
-      return session.agent;
-    }),
-
   getSessions: sessionProcedure
     .input(
       z.object({