@lobehub/chat 1.11.8 → 1.12.0

package/src/locales/default/knowledgeBase.ts

@@ -0,0 +1,31 @@
+export default {
+  addToKnowledgeBase: {
+    addSuccess: '文件添加成功，<1>立即查看</1>',
+    confirm: '添加',
+    id: {
+      placeholder: '请选择要添加的知识库',
+      required: '请选择知识库',
+      title: '目标知识库',
+    },
+    title: '添加到知识库',
+    totalFiles: '已选择 {{count}} 个文件',
+  },
+  createNew: {
+    confirm: '新建',
+    description: {
+      placeholder: '知识库简介（选填）',
+    },
+    formTitle: '基本信息',
+    name: {
+      placeholder: '知识库名称',
+      required: '请填写知识库名称',
+    },
+    title: '新建知识库',
+  },
+  tab: {
+    files: '文档',
+    settings: '设置',
+    testing: '召回测试',
+  },
+  title: '知识库',
+};

package/src/locales/default/portal.ts

@@ -6,5 +6,6 @@ export default {
     summaryTooltip: '总结当前内容',
   },
   emptyArtifactList: '当前 Artifacts 列表为空，请在会话中按需使用插件后再查看',
+  files: '文件',
   title: '工作区',
 };

package/src/middleware.ts

@@ -14,6 +14,8 @@ export const config = {
     '/',
     '/chat(.*)',
     '/settings(.*)',
+    '/files(.*)',
+    '/repos(.*)',
     // ↓ cloud ↓
   ],
 };
@@ -46,6 +48,7 @@ const nextAuthMiddleware = NextAuthEdge.auth((req) => {
 
 const isProtectedRoute = createRouteMatcher([
   '/settings(.*)',
+  '/files(.*)',
   // ↓ cloud ↓
 ]);
 

package/src/server/asyncContext.ts

@@ -0,0 +1,40 @@
+import { NextRequest } from 'next/server';
+
+import { JWTPayload, LOBE_CHAT_AUTH_HEADER } from '@/const/auth';
+import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
+
+export interface AsyncAuthContext {
+  jwtPayload: JWTPayload;
+  secret: string;
+  userId?: string | null;
+}
+
+/**
+ * Inner function for `createContext` where we create the context.
+ * This is useful for testing when we don't want to mock Next.js' request/response
+ */
+export const createAsyncContextInner = async (params?: {
+  jwtPayload?: JWTPayload;
+  secret?: string;
+  userId?: string | null;
+}): Promise<AsyncAuthContext> => ({
+  jwtPayload: params?.jwtPayload || {},
+  secret: params?.secret || '',
+  userId: params?.userId,
+});
+
+export type AsyncContext = Awaited<ReturnType<typeof createAsyncContextInner>>;
+
+export const createAsyncRouteContext = async (request: NextRequest): Promise<AsyncContext> => {
+  // for API-response caching see https://trpc.io/docs/v11/caching
+
+  const authorization = request.headers.get('Authorization');
+  const lobeChatAuthorization = request.headers.get(LOBE_CHAT_AUTH_HEADER);
+
+  const secret = authorization?.split(' ')[1];
+  const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+  const { plaintext } = await gateKeeper.decrypt(lobeChatAuthorization || '');
+
+  const { userId, payload } = JSON.parse(plaintext);
+  return createAsyncContextInner({ jwtPayload: payload, secret, userId });
+};

package/src/server/modules/ContentChunk/index.ts

@@ -0,0 +1,135 @@
+import { ChunkingLoader } from 'src/libs/langchain';
+import { Strategy } from 'unstructured-client/sdk/models/shared';
+
+import { NewChunkItem, NewUnstructuredChunkItem } from '@/database/server/schemas/lobechat';
+import { ChunkingStrategy, Unstructured } from '@/libs/unstructured';
+
+export interface ChunkContentParams {
+  content: Uint8Array;
+  fileType: string;
+  filename: string;
+  mode?: 'fast' | 'hi-res';
+}
+
+interface ChunkResult {
+  chunks: NewChunkItem[];
+  unstructuredChunks?: NewUnstructuredChunkItem[];
+}
+
+export class ContentChunk {
+  private unstructuredClient: Unstructured;
+  private langchainClient: ChunkingLoader;
+
+  constructor() {
+    this.unstructuredClient = new Unstructured();
+    this.langchainClient = new ChunkingLoader();
+  }
+
+  isUsingUnstructured(params: ChunkContentParams) {
+    return params.fileType === 'application/pdf' && params.mode === 'hi-res';
+  }
+
+  async chunkContent(params: ChunkContentParams): Promise<ChunkResult> {
+    if (this.isUsingUnstructured(params))
+      return await this.chunkByUnstructured(params.filename, params.content);
+
+    return await this.chunkByLangChain(params.filename, params.content);
+  }
+
+  private chunkByUnstructured = async (
+    filename: string,
+    content: Uint8Array,
+  ): Promise<ChunkResult> => {
+    const result = await this.unstructuredClient.partition({
+      chunkingStrategy: ChunkingStrategy.ByPage,
+      fileContent: content,
+      filename,
+      strategy: Strategy.Auto,
+    });
+
+    // after finish partition, we need to filter out some elements
+    const documents = result.compositeElements
+      .filter((e) => !new Set(['PageNumber', 'Footer']).has(e.type))
+      .map((item, index): NewChunkItem => {
+        const {
+          text_as_html,
+          page_number,
+          page_name,
+          image_mime_type,
+          image_base64,
+          parent_id,
+          languages,
+          coordinates,
+        } = item.metadata;
+
+        return {
+          id: item.element_id,
+          index,
+          metadata: {
+            coordinates,
+            image_base64,
+            image_mime_type,
+            languages,
+            page_name,
+            page_number,
+            parent_id,
+            text_as_html,
+          },
+          text: item.text,
+          type: item.type,
+        };
+      });
+
+    const chunks = result.originElements
+      .filter((e) => !new Set(['PageNumber', 'Footer']).has(e.type))
+      .map((item, index): NewUnstructuredChunkItem => {
+        const {
+          text_as_html,
+          page_number,
+          page_name,
+          image_mime_type,
+          image_base64,
+          parent_id,
+          languages,
+          coordinates,
+        } = item.metadata;
+
+        return {
+          compositeId: item.compositeId,
+          id: item.element_id,
+          index,
+          metadata: {
+            coordinates,
+            image_base64,
+            image_mime_type,
+            languages,
+            page_name,
+            page_number,
+            text_as_html,
+          },
+          parentId: parent_id,
+          text: item.text,
+          type: item.type,
+        };
+      });
+
+    return { chunks: documents, unstructuredChunks: chunks };
+  };
+
+  private chunkByLangChain = async (
+    filename: string,
+    content: Uint8Array,
+  ): Promise<ChunkResult> => {
+    const res = await this.langchainClient.partitionContent(filename, content);
+
+    const documents = res.map((item, index) => ({
+      id: item.id,
+      index,
+      metadata: item.metadata,
+      text: item.pageContent,
+      type: 'LangChainElement',
+    }));
+
+    return { chunks: documents };
+  };
+}

package/src/server/modules/S3/index.ts

@@ -1,6 +1,7 @@
 import {
+  DeleteObjectCommand,
+  DeleteObjectsCommand,
   GetObjectCommand,
-  ListObjectsCommand,
   PutObjectCommand,
   S3Client,
 } from '@aws-sdk/client-s3';
@@ -45,13 +46,22 @@ export class S3 {
     });
   }
 
-  public async getImages(): Promise<FileType[]> {
-    const command = new ListObjectsCommand({
+  public async deleteFile(key: string) {
+    const command = new DeleteObjectCommand({
       Bucket: this.bucket,
+      Key: key,
+    });
+
+    return this.client.send(command);
+  }
+
+  public async deleteFiles(keys: string[]) {
+    const command = new DeleteObjectsCommand({
+      Bucket: this.bucket,
+      Delete: { Objects: keys.map((key) => ({ Key: key })) },
     });
 
-    const res = await this.client.send(command);
-    return listFileSchema.parse(res.Contents);
+    return this.client.send(command);
   }
 
   public async getFileContent(key: string): Promise<string> {
@@ -69,6 +79,21 @@ export class S3 {
     return response.Body.transformToString();
   }
 
+  public async getFileByteArray(key: string): Promise<Uint8Array> {
+    const command = new GetObjectCommand({
+      Bucket: this.bucket,
+      Key: key,
+    });
+
+    const response = await this.client.send(command);
+
+    if (!response.Body) {
+      throw new Error(`No body in response with ${key}`);
+    }
+
+    return response.Body.transformToByteArray();
+  }
+
   public async createPreSignedUrl(key: string): Promise<string> {
     const command = new PutObjectCommand({
       ACL: this.setAcl ? 'public-read' : undefined,

package/src/server/routers/async/caller.ts

@@ -0,0 +1,27 @@
+import { createTRPCClient, httpBatchLink } from '@trpc/client';
+import superjson from 'superjson';
+import urlJoin from 'url-join';
+
+import { appEnv } from '@/config/app';
+import { serverDBEnv } from '@/config/db';
+import { JWTPayload, LOBE_CHAT_AUTH_HEADER } from '@/const/auth';
+import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
+
+import type { AsyncRouter } from './index';
+
+export const createAsyncServerClient = async (userId: string, payload: JWTPayload) => {
+  const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+
+  return createTRPCClient<AsyncRouter>({
+    links: [
+      httpBatchLink({
+        headers: {
+          Authorization: `Bearer ${serverDBEnv.KEY_VAULTS_SECRET}`,
+          [LOBE_CHAT_AUTH_HEADER]: await gateKeeper.encrypt(JSON.stringify({ payload, userId })),
+        },
+        transformer: superjson,
+        url: urlJoin(appEnv.APP_URL!, '/trpc/async'),
+      }),
+    ],
+  });
+};

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

@@ -0,0 +1,247 @@
+import { TRPCError } from '@trpc/server';
+import { chunk } from 'lodash-es';
+import pMap from 'p-map';
+import { z } from 'zod';
+
+import { initAgentRuntimeWithUserPayload } from '@/app/api/chat/agentRuntime';
+import { fileEnv } from '@/config/file';
+import { DEFAULT_EMBEDDING_MODEL } from '@/const/settings';
+import { ASYNC_TASK_TIMEOUT, 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 { NewChunkItem, NewEmbeddingsItem } from '@/database/server/schemas/lobechat';
+import { ModelProvider } from '@/libs/agent-runtime';
+import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
+import { S3 } from '@/server/modules/S3';
+import { ChunkService } from '@/server/services/chunk';
+import { AsyncTaskError, AsyncTaskErrorType, AsyncTaskStatus } from '@/types/asyncTask';
+import { safeParseJSON } from '@/utils/safeParseJSON';
+
+const fileProcedure = asyncAuthedProcedure.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),
+    },
+  });
+});
+
+export const fileRouter = router({
+  embeddingChunks: fileProcedure
+    .input(
+      z.object({
+        fileId: z.string(),
+        model: z.string().default(DEFAULT_EMBEDDING_MODEL),
+        taskId: z.string(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const file = await ctx.fileModel.findById(input.fileId);
+
+      if (!file) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
+      }
+
+      const asyncTask = await ctx.asyncTaskModel.findById(input.taskId);
+
+      if (!asyncTask) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Async Task not found' });
+
+      try {
+        const timeoutPromise = new Promise((_, reject) => {
+          setTimeout(() => {
+            reject({
+              body: { detail: 'embedding task is timeout, please try again' },
+              name: AsyncTaskErrorType.Timeout,
+            } as AsyncTaskError);
+          }, ASYNC_TASK_TIMEOUT);
+        });
+
+        const embeddingPromise = async () => {
+          // update the task status to success
+          await ctx.asyncTaskModel.update(input.taskId, {
+            status: AsyncTaskStatus.Processing,
+          });
+
+          const startAt = Date.now();
+
+          const CHUNK_SIZE = 50;
+          const CONCURRENCY = 10;
+
+          const chunks = await ctx.chunkModel.getChunksTextByFileId(input.fileId);
+          const requestArray = chunk(chunks, CHUNK_SIZE);
+
+          await pMap(
+            requestArray,
+            async (chunks, index) => {
+              const agentRuntime = await initAgentRuntimeWithUserPayload(
+                ModelProvider.OpenAI,
+                ctx.jwtPayload,
+              );
+
+              const number = index + 1;
+              console.log(`执行第 ${number} 个任务`);
+
+              console.time(`任务[${number}]: embeddings`);
+
+              const embeddings = await agentRuntime.embeddings({
+                dimensions: 1024,
+                input: chunks.map((c) => c.text),
+                model: input.model,
+              });
+              console.timeEnd(`任务[${number}]: embeddings`);
+
+              const items: NewEmbeddingsItem[] =
+                embeddings?.map((e) => ({
+                  chunkId: chunks[e.index].id,
+                  embeddings: e.embedding,
+                  fileId: input.fileId,
+                  model: input.model,
+                })) || [];
+
+              console.time(`任务[${number}]: insert db`);
+              await ctx.embeddingModel.bulkCreate(items);
+              console.timeEnd(`任务[${number}]: insert db`);
+            },
+            { concurrency: CONCURRENCY },
+          );
+
+          const duration = Date.now() - startAt;
+          // update the task status to success
+          await ctx.asyncTaskModel.update(input.taskId, {
+            duration,
+            status: AsyncTaskStatus.Success,
+          });
+
+          return { success: true };
+        };
+
+        // Race between the chunking process and the timeout
+        return await Promise.race([embeddingPromise(), timeoutPromise]);
+      } catch (e) {
+        console.error('embeddingChunks error', e);
+        await ctx.asyncTaskModel.update(input.taskId, {
+          error: e,
+          status: AsyncTaskStatus.Error,
+        });
+
+        return {
+          message: `File ${file.name}(${input.taskId}) failed to embedding: ${(e as Error).message}`,
+          success: false,
+        };
+      }
+    }),
+
+  parseFileToChunks: fileProcedure
+    .input(
+      z.object({
+        fileId: z.string(),
+        taskId: z.string(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const file = await ctx.fileModel.findById(input.fileId);
+      if (!file) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
+      }
+
+      const s3 = new S3();
+
+      let content: Uint8Array | undefined;
+      try {
+        content = await s3.getFileByteArray(file.url);
+      } catch (e) {
+        console.error(e);
+        // if file not found, delete it from db
+        if ((e as any).Code === 'NoSuchKey') {
+          await ctx.fileModel.delete(input.fileId);
+          throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
+        }
+      }
+
+      if (!content) return;
+
+      const asyncTask = await ctx.asyncTaskModel.findById(input.taskId);
+
+      if (!asyncTask) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Async Task not found' });
+
+      try {
+        const startAt = Date.now();
+
+        const timeoutPromise = new Promise((_, reject) => {
+          setTimeout(() => {
+            reject({
+              body: { detail: 'chunking task is timeout, please try again' },
+              name: AsyncTaskErrorType.Timeout,
+            } as AsyncTaskError);
+          }, ASYNC_TASK_TIMEOUT);
+        });
+
+        const chunkingPromise = async () => {
+          const chunkService = ctx.chunkService;
+          // update the task status to processing
+          await ctx.asyncTaskModel.update(input.taskId, { status: AsyncTaskStatus.Processing });
+
+          // partition file to chunks
+          const chunkResult = await chunkService.chunkContent({
+            content,
+            fileType: file.fileType,
+            filename: file.name,
+          });
+
+          // after finish partition, we need to filter out some elements
+          const chunks = chunkResult.chunks.map(
+            (item): NewChunkItem => ({ ...item, userId: ctx.userId }),
+          );
+
+          const duration = Date.now() - startAt;
+
+          await ctx.chunkModel.bulkCreate(chunks, input.fileId);
+
+          if (chunkResult.unstructuredChunks) {
+            const unstructuredChunks = chunkResult.unstructuredChunks.map(
+              (item): NewChunkItem => ({ ...item, fileId: input.fileId, userId: ctx.userId }),
+            );
+            await ctx.chunkModel.bulkCreateUnstructuredChunks(unstructuredChunks);
+          }
+
+          // update the task status to success
+          await ctx.asyncTaskModel.update(input.taskId, {
+            duration,
+            status: AsyncTaskStatus.Success,
+          });
+
+          // if enable auto embedding, trigger the embedding task
+          if (fileEnv.CHUNKS_AUTO_EMBEDDING) {
+            await chunkService.asyncEmbeddingFileChunks(input.fileId, ctx.jwtPayload);
+          }
+
+          return { success: true };
+        };
+        // Race between the chunking process and the timeout
+        return await Promise.race([chunkingPromise(), timeoutPromise]);
+      } catch (e) {
+        const error = e as any;
+
+        const asyncTaskError: AsyncTaskError = error.body
+          ? { body: safeParseJSON(error.body) ?? error.body, name: error.name }
+          : { body: { detail: error.message }, name: (error as Error).name };
+
+        console.error('[Chunking Error]', asyncTaskError);
+        await ctx.asyncTaskModel.update(input.taskId, {
+          error: asyncTaskError,
+          status: AsyncTaskStatus.Error,
+        });
+
+        return {
+          message: `File ${file.name}(${input.taskId}) failed to chunking: ${(e as Error).message}`,
+          success: false,
+        };
+      }
+    }),
+});

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

@@ -0,0 +1,12 @@
+import { publicProcedure, asyncRouter as router } from '@/libs/trpc/async';
+
+import { fileRouter } from './file';
+
+export const asyncRouter = router({
+  file: fileRouter,
+  healthcheck: publicProcedure.query(() => "i'm live!"),
+});
+
+export type AsyncRouter = typeof asyncRouter;
+
+export * from './caller';

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

@@ -0,0 +1,77 @@
+import { z } from 'zod';
+
+import { SessionGroupModel } from '@/database/server/models/sessionGroup';
+import { insertSessionGroupSchema } from '@/database/server/schemas/lobechat';
+import { authedProcedure, router } from '@/libs/trpc';
+import { SessionGroupItem } from '@/types/session';
+
+const sessionProcedure = authedProcedure.use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      sessionGroupModel: new SessionGroupModel(ctx.userId),
+    },
+  });
+});
+
+export const sessionGroupRouter = router({
+  createSessionGroup: sessionProcedure
+    .input(
+      z.object({
+        name: z.string(),
+        sort: z.number().optional(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const data = await ctx.sessionGroupModel.create({
+        name: input.name,
+        sort: input.sort,
+      });
+
+      return data?.id;
+    }),
+
+  getSessionGroup: sessionProcedure.query(async ({ ctx }): Promise<SessionGroupItem[]> => {
+    return ctx.sessionGroupModel.query() as any;
+  }),
+
+  removeAllSessionGroups: sessionProcedure.mutation(async ({ ctx }) => {
+    return ctx.sessionGroupModel.deleteAll();
+  }),
+
+  removeSessionGroup: sessionProcedure
+    .input(z.object({ id: z.string(), removeChildren: z.boolean().optional() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.sessionGroupModel.delete(input.id);
+    }),
+
+  updateSessionGroup: sessionProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        value: insertSessionGroupSchema.partial(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.sessionGroupModel.update(input.id, input.value);
+    }),
+  updateSessionGroupOrder: sessionProcedure
+    .input(
+      z.object({
+        sortMap: z.array(
+          z.object({
+            id: z.string(),
+            sort: z.number(),
+          }),
+        ),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      console.log('sortMap:', input.sortMap);
+
+      return ctx.sessionGroupModel.updateOrder(input.sortMap);
+    }),
+});
+
+export type SessionGroupRouter = typeof sessionGroupRouter;