@lobehub/chat 1.15.7 → 1.15.9

package/src/locales/default/ragEval.ts

@@ -0,0 +1,93 @@
+export default {
+  addDataset: {
+    confirm: '新建',
+    description: {
+      placeholder: '数据集简介（选填）',
+    },
+    name: {
+      placeholder: '数据集名称',
+      required: '请填写数据集名称',
+    },
+    title: '添加数据集',
+  },
+  dataset: {
+    addNewButton: '创建数据集',
+    emptyGuide: '当前数据集为空，请创建一个数据集。',
+    list: {
+      table: {
+        actions: {
+          importData: '导入数据',
+        },
+        columns: {
+          actions: '操作',
+          ideal: {
+            title: '期望回答',
+          },
+          question: {
+            title: '问题',
+          },
+          referenceFiles: {
+            title: '参考文件',
+          },
+        },
+        notSelected: '请在左侧选择数据集',
+        title: '数据集详情',
+      },
+      title: '数据集',
+    },
+  },
+  evaluation: {
+    addEvaluation: {
+      confirm: '新建',
+      datasetId: {
+        placeholder: '请选择你的评测数据集',
+        required: '请选择评测数据集',
+      },
+      description: {
+        placeholder: '评测任务简介（选填）',
+      },
+      name: {
+        placeholder: '评测任务名称',
+        required: '请填写评测任务名称',
+      },
+      title: '添加评测任务',
+    },
+    addNewButton: '创建评测',
+    emptyGuide: '当前评测任务为空，开始创建评测。',
+
+    table: {
+      columns: {
+        actions: {
+          checkStatus: '检查状态',
+          confirmDelete: '是否删除本条评测',
+          confirmRun:
+            '是否开始运行？开始运行后将在后台异步执行评测任务，关闭页面不影响异步任务的执行',
+          downloadRecords: '下载评测',
+          retry: '重试',
+          run: '运行',
+          title: '操作',
+        },
+        datasetId: {
+          title: '数据集',
+        },
+        name: {
+          title: '评测任务名称',
+        },
+        records: {
+          title: '评测记录数',
+        },
+        referenceFiles: {
+          title: '参考文件',
+        },
+        status: {
+          error: '执行出错',
+          pending: '待运行',
+          processing: '运行中',
+          success: '执行成功',
+          title: '状态',
+        },
+      },
+      title: '评测任务列表',
+    },
+  },
+};

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

@@ -104,4 +104,15 @@ export class S3 {
 
     return getSignedUrl(this.client, command, { expiresIn: 3600 });
   }
+
+  public async uploadContent(path: string, content: string) {
+    const command = new PutObjectCommand({
+      ACL: this.setAcl ? 'public-read' : undefined,
+      Body: content,
+      Bucket: this.bucket,
+      Key: path,
+    });
+
+    return this.client.send(command);
+  }
 }

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

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

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

@@ -0,0 +1,138 @@
+import { TRPCError } from '@trpc/server';
+import OpenAI from 'openai';
+import { z } from 'zod';
+
+import { initAgentRuntimeWithUserPayload } from '@/app/api/chat/agentRuntime';
+import { chainAnswerWithContext } from '@/chains/answerWithContext';
+import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
+import { ChunkModel } from '@/database/server/models/chunk';
+import { EmbeddingModel } from '@/database/server/models/embedding';
+import { FileModel } from '@/database/server/models/file';
+import {
+  EvalDatasetRecordModel,
+  EvalEvaluationModel,
+  EvaluationRecordModel,
+} from '@/database/server/models/ragEval';
+import { ModelProvider } from '@/libs/agent-runtime';
+import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
+import { ChunkService } from '@/server/services/chunk';
+import { AsyncTaskError } from '@/types/asyncTask';
+import { EvalEvaluationStatus } from '@/types/eval';
+
+const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      chunkModel: new ChunkModel(ctx.userId),
+      chunkService: new ChunkService(ctx.userId),
+      datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.userId),
+      evalRecordModel: new EvaluationRecordModel(ctx.userId),
+      evaluationModel: new EvalEvaluationModel(ctx.userId),
+      fileModel: new FileModel(ctx.userId),
+    },
+  });
+});
+
+export const ragEvalRouter = router({
+  runRecordEvaluation: ragEvalProcedure
+    .input(
+      z.object({
+        evalRecordId: z.number(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const evalRecord = await ctx.evalRecordModel.findById(input.evalRecordId);
+
+      if (!evalRecord) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
+      }
+
+      const now = Date.now();
+      try {
+        const agentRuntime = await initAgentRuntimeWithUserPayload(
+          ModelProvider.OpenAI,
+          ctx.jwtPayload,
+        );
+
+        const { question, languageModel, embeddingModel } = evalRecord;
+
+        let questionEmbeddingId = evalRecord.questionEmbeddingId;
+        let context = evalRecord.context;
+
+        // 如果不存在 questionEmbeddingId，那么就需要做一次 embedding
+        if (!questionEmbeddingId) {
+          const embeddings = await agentRuntime.embeddings({
+            dimensions: 1024,
+            input: question,
+            model: !!embeddingModel ? embeddingModel : DEFAULT_EMBEDDING_MODEL,
+          });
+
+          const embeddingId = await ctx.embeddingModel.create({
+            embeddings: embeddings?.[0].embedding,
+            model: embeddingModel,
+          });
+
+          await ctx.evalRecordModel.update(evalRecord.id, {
+            questionEmbeddingId: embeddingId,
+          });
+
+          questionEmbeddingId = embeddingId;
+        }
+
+        // 如果不存在 context，那么就需要做一次检索
+        if (!context || context.length === 0) {
+          const datasetRecord = await ctx.datasetRecordModel.findById(evalRecord.datasetRecordId);
+
+          const embeddingItem = await ctx.embeddingModel.findById(questionEmbeddingId);
+
+          const chunks = await ctx.chunkModel.semanticSearchForChat({
+            embedding: embeddingItem!.embeddings!,
+            fileIds: datasetRecord!.referenceFiles!,
+            query: evalRecord.question,
+          });
+
+          context = chunks.map((item) => item.text).filter(Boolean) as string[];
+          await ctx.evalRecordModel.update(evalRecord.id, { context });
+        }
+
+        // 做一次生成 LLM 答案生成
+        const { messages } = chainAnswerWithContext({ context, knowledge: [], question });
+
+        const response = await agentRuntime.chat({
+          messages: messages!,
+          model: !!languageModel ? languageModel : DEFAULT_MODEL,
+          responseMode: 'json',
+          stream: false,
+          temperature: 1,
+        });
+
+        const data = (await response.json()) as OpenAI.ChatCompletion;
+
+        const answer = data.choices[0].message.content;
+
+        await ctx.evalRecordModel.update(input.evalRecordId, {
+          answer,
+          duration: Date.now() - now,
+          languageModel,
+          status: EvalEvaluationStatus.Success,
+        });
+
+        return { success: true };
+      } catch (e) {
+        await ctx.evalRecordModel.update(input.evalRecordId, {
+          error: new AsyncTaskError((e as Error).name, (e as Error).message),
+          status: EvalEvaluationStatus.Error,
+        });
+
+        await ctx.evaluationModel.update(evalRecord.evaluationId, {
+          status: EvalEvaluationStatus.Error,
+        });
+
+        console.error('[RAGEvaluation] error', e);
+
+        return { success: false };
+      }
+    }),
+});

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

@@ -5,12 +5,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 { ragEvalRouter } from './ragEval';
 import { sessionRouter } from './session';
 import { sessionGroupRouter } from './sessionGroup';
 import { topicRouter } from './topic';
@@ -25,6 +25,7 @@ export const lambdaRouter = router({
   knowledgeBase: knowledgeBaseRouter,
   message: messageRouter,
   plugin: pluginRouter,
+  ragEval: ragEvalRouter,
   session: sessionRouter,
   sessionGroup: sessionGroupRouter,
   topic: topicRouter,

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

@@ -0,0 +1,296 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import { TRPCError } from '@trpc/server';
+import dayjs from 'dayjs';
+import JSONL from 'jsonl-parse-stringify';
+import pMap from 'p-map';
+import { z } from 'zod';
+
+import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
+import { FileModel } from '@/database/server/models/file';
+import {
+  EvalDatasetModel,
+  EvalDatasetRecordModel,
+  EvalEvaluationModel,
+  EvaluationRecordModel,
+} from '@/database/server/models/ragEval';
+import { authedProcedure, router } from '@/libs/trpc';
+import { keyVaults } from '@/libs/trpc/middleware/keyVaults';
+import { S3 } from '@/server/modules/S3';
+import { createAsyncServerClient } from '@/server/routers/async';
+import { getFullFileUrl } from '@/server/utils/files';
+import {
+  EvalDatasetRecord,
+  EvalEvaluationStatus,
+  InsertEvalDatasetRecord,
+  RAGEvalDataSetItem,
+  insertEvalDatasetRecordSchema,
+  insertEvalDatasetsSchema,
+  insertEvalEvaluationSchema,
+} from '@/types/eval';
+
+const ragEvalProcedure = authedProcedure.use(keyVaults).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      datasetModel: new EvalDatasetModel(ctx.userId),
+      fileModel: new FileModel(ctx.userId),
+      datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
+      evaluationModel: new EvalEvaluationModel(ctx.userId),
+      evaluationRecordModel: new EvaluationRecordModel(ctx.userId),
+      s3: new S3(),
+    },
+  });
+});
+
+export const ragEvalRouter = router({
+  createDataset: ragEvalProcedure
+    .input(
+      z.object({
+        description: z.string().optional(),
+        knowledgeBaseId: z.string(),
+        name: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const data = await ctx.datasetModel.create({
+        description: input.description,
+        knowledgeBaseId: input.knowledgeBaseId,
+        name: input.name,
+      });
+
+      return data?.id;
+    }),
+
+  getDatasets: ragEvalProcedure
+    .input(z.object({ knowledgeBaseId: z.string() }))
+
+    .query(async ({ ctx, input }): Promise<RAGEvalDataSetItem[]> => {
+      return ctx.datasetModel.query(input.knowledgeBaseId);
+    }),
+
+  removeDataset: ragEvalProcedure
+    .input(z.object({ id: z.number() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.datasetModel.delete(input.id);
+    }),
+
+  updateDataset: ragEvalProcedure
+    .input(
+      z.object({
+        id: z.number(),
+        value: insertEvalDatasetsSchema.partial(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.datasetModel.update(input.id, input.value);
+    }),
+
+  // Dataset Item operations
+  createDatasetRecords: ragEvalProcedure
+    .input(
+      z.object({
+        datasetId: z.number(),
+        question: z.string(),
+        ideal: z.string().optional(),
+        referenceFiles: z.array(z.string()).optional(),
+        metadata: z.record(z.unknown()).optional(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const data = await ctx.datasetRecordModel.create(input);
+      return data?.id;
+    }),
+
+  getDatasetRecords: ragEvalProcedure
+    .input(z.object({ datasetId: z.number() }))
+    .query(async ({ ctx, input }): Promise<EvalDatasetRecord[]> => {
+      return ctx.datasetRecordModel.query(input.datasetId);
+    }),
+
+  removeDatasetRecords: ragEvalProcedure
+    .input(z.object({ id: z.number() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.datasetRecordModel.delete(input.id);
+    }),
+
+  updateDatasetRecords: ragEvalProcedure
+    .input(
+      z.object({
+        id: z.number(),
+        value: z
+          .object({
+            question: z.string(),
+            ideal: z.string(),
+            referenceFiles: z.array(z.string()),
+            metadata: z.record(z.unknown()),
+          })
+          .partial(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      return ctx.datasetRecordModel.update(input.id, input.value);
+    }),
+
+  importDatasetRecords: ragEvalProcedure
+    .input(
+      z.object({
+        datasetId: z.number(),
+        pathname: z.string(),
+      }),
+    )
+    .mutation(async ({ input, ctx }) => {
+      const dataStr = await ctx.s3.getFileContent(input.pathname);
+      const items = JSONL.parse<InsertEvalDatasetRecord>(dataStr);
+
+      insertEvalDatasetRecordSchema.array().parse(items);
+
+      const data = await Promise.all(
+        items.map(async ({ referenceFiles, question, ideal }) => {
+          const files = typeof referenceFiles === 'string' ? [referenceFiles] : referenceFiles;
+
+          let fileIds: string[] | undefined = undefined;
+
+          if (files) {
+            const items = await ctx.fileModel.findByNames(files);
+
+            fileIds = items.map((item) => item.id);
+          }
+
+          return {
+            question,
+            ideal,
+            referenceFiles: fileIds,
+            datasetId: input.datasetId,
+          };
+        }),
+      );
+
+      return ctx.datasetRecordModel.batchCreate(data);
+    }),
+
+  // Evaluation operations
+  startEvaluationTask: ragEvalProcedure
+    .input(z.object({ id: z.number() }))
+    .mutation(async ({ input, ctx }) => {
+      // Start evaluation task
+      const evaluation = await ctx.evaluationModel.findById(input.id);
+
+      if (!evaluation) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
+      }
+
+      // create evaluation records by dataset records
+      const datasetRecords = await ctx.datasetRecordModel.findByDatasetId(evaluation.datasetId);
+
+      if (datasetRecords.length === 0) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Dataset record is empty' });
+      }
+
+      const evalRecords = await ctx.evaluationRecordModel.batchCreate(
+        datasetRecords.map((record) => ({
+          evaluationId: input.id,
+          datasetRecordId: record.id,
+          question: record.question!,
+          ideal: record.ideal,
+          status: EvalEvaluationStatus.Pending,
+          embeddingModel: DEFAULT_EMBEDDING_MODEL,
+          languageModel: DEFAULT_MODEL,
+        })),
+      );
+
+      const asyncCaller = await createAsyncServerClient(ctx.userId, ctx.jwtPayload);
+
+      await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Processing });
+      try {
+        await pMap(
+          evalRecords,
+          async (record) => {
+            asyncCaller.ragEval.runRecordEvaluation
+              .mutate({ evalRecordId: record.id })
+              .catch(async (e) => {
+                await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Error });
+
+                throw new TRPCError({
+                  code: 'BAD_GATEWAY',
+                  message: `[ASYNC_TASK] Failed to start evaluation task: ${e.message}`,
+                });
+              });
+          },
+          {
+            concurrency: 30,
+          },
+        );
+
+        return { success: true };
+      } catch (e) {
+        console.error('[startEvaluationTask]:', e);
+
+        await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Error });
+
+        return { success: false };
+      }
+    }),
+
+  checkEvaluationStatus: ragEvalProcedure
+    .input(z.object({ id: z.number() }))
+    .query(async ({ input, ctx }) => {
+      const evaluation = await ctx.evaluationModel.findById(input.id);
+
+      if (!evaluation) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Evaluation not found' });
+      }
+
+      const records = await ctx.evaluationRecordModel.findByEvaluationId(input.id);
+
+      const isSuccess = records.every((record) => record.status === EvalEvaluationStatus.Success);
+
+      if (isSuccess) {
+        // 将结果上传到 S3
+
+        const evalRecords = records.map((record) => ({
+          question: record.question,
+          context: record.context,
+          answer: record.answer,
+          ground_truth: record.ideal,
+        }));
+        const date = dayjs().format('YYYY-MM-DD-HH-mm');
+        const filename = `${date}-eval_${evaluation.id}-${evaluation.name}.jsonl`;
+        const path = `rag_eval_records/${filename}`;
+
+        await ctx.s3.uploadContent(path, JSONL.stringify(evalRecords));
+
+        // 保存数据
+        await ctx.evaluationModel.update(input.id, {
+          status: EvalEvaluationStatus.Success,
+          evalRecordsUrl: getFullFileUrl(path),
+        });
+      }
+
+      return { success: isSuccess };
+    }),
+  createEvaluation: ragEvalProcedure
+    .input(insertEvalEvaluationSchema)
+    .mutation(async ({ input, ctx }) => {
+      const data = await ctx.evaluationModel.create({
+        description: input.description,
+        knowledgeBaseId: input.knowledgeBaseId,
+        datasetId: input.datasetId,
+        name: input.name,
+      });
+
+      return data?.id;
+    }),
+
+  removeEvaluation: ragEvalProcedure
+    .input(z.object({ id: z.number() }))
+    .mutation(async ({ input, ctx }) => {
+      return ctx.evaluationModel.delete(input.id);
+    }),
+
+  getEvaluationList: ragEvalProcedure
+    .input(z.object({ knowledgeBaseId: z.string() }))
+    .query(async ({ ctx, input }) => {
+      return ctx.evaluationModel.queryByKnowledgeBaseId(input.knowledgeBaseId);
+    }),
+});

package/src/services/ragEval.ts

@@ -0,0 +1,67 @@
+import { lambdaClient } from '@/libs/trpc/client';
+import { uploadService } from '@/services/upload';
+import {
+  CreateNewEvalDatasets,
+  CreateNewEvalEvaluation,
+  EvalDatasetRecord,
+  RAGEvalDataSetItem,
+  RAGEvalEvaluationItem,
+  insertEvalDatasetsSchema,
+} from '@/types/eval';
+
+class RAGEvalService {
+  // Dataset
+  async createDataset(params: CreateNewEvalDatasets): Promise<number | undefined> {
+    return await lambdaClient.ragEval.createDataset.mutate(params);
+  }
+
+  async getDatasets(knowledgeBaseId: string): Promise<RAGEvalDataSetItem[]> {
+    return lambdaClient.ragEval.getDatasets.query({ knowledgeBaseId });
+  }
+
+  async removeDataset(id: number): Promise<void> {
+    await lambdaClient.ragEval.removeDataset.mutate({ id });
+  }
+
+  async updateDataset(id: number, value: Partial<typeof insertEvalDatasetsSchema>): Promise<void> {
+    await lambdaClient.ragEval.updateDataset.mutate({ id, value });
+  }
+
+  // Dataset Records
+  async getDatasetRecords(datasetId: number): Promise<EvalDatasetRecord[]> {
+    return lambdaClient.ragEval.getDatasetRecords.query({ datasetId });
+  }
+
+  async removeDatasetRecord(id: number): Promise<void> {
+    await lambdaClient.ragEval.removeDatasetRecords.mutate({ id });
+  }
+
+  async importDatasetRecords(datasetId: number, file: File): Promise<void> {
+    const { path } = await uploadService.uploadWithProgress(file, { directory: 'ragEval' });
+
+    await lambdaClient.ragEval.importDatasetRecords.mutate({ datasetId, pathname: path });
+  }
+
+  // Evaluation
+  async createEvaluation(params: CreateNewEvalEvaluation): Promise<number | undefined> {
+    return await lambdaClient.ragEval.createEvaluation.mutate(params);
+  }
+
+  async getEvaluationList(knowledgeBaseId: string): Promise<RAGEvalEvaluationItem[]> {
+    return lambdaClient.ragEval.getEvaluationList.query({ knowledgeBaseId });
+  }
+
+  async startEvaluationTask(id: number) {
+    return lambdaClient.ragEval.startEvaluationTask.mutate({ id });
+  }
+
+  async removeEvaluation(id: number): Promise<void> {
+    await lambdaClient.ragEval.removeEvaluation.mutate({ id });
+  }
+
+  async checkEvaluationStatus(id: number): Promise<{ success: boolean }> {
+    return lambdaClient.ragEval.checkEvaluationStatus.query({ id });
+  }
+}
+
+export const ragEvalService = new RAGEvalService();

package/src/services/upload.ts

@@ -11,11 +11,17 @@ export const UPLOAD_NETWORK_ERROR = 'NetWorkError';
 class UploadService {
   uploadWithProgress = async (
     file: File,
-    onProgress: (status: FileUploadStatus, state: FileUploadState) => void,
+    {
+      onProgress,
+      directory,
+    }: {
+      directory?: string;
+      onProgress?: (status: FileUploadStatus, state: FileUploadState) => void;
+    },
   ): Promise<FileMetadata> => {
     const xhr = new XMLHttpRequest();
 
-    const { preSignUrl, ...result } = await this.getSignedUploadUrl(file);
+    const { preSignUrl, ...result } = await this.getSignedUploadUrl(file, directory);
     let startTime = Date.now();
     xhr.upload.addEventListener('progress', (event) => {
       if (event.lengthComputable) {
@@ -29,7 +35,7 @@ class UploadService {
           // so make it as 99.9 and let users think it's still uploading
           progress: progress === 100 ? 99.9 : progress,
           restTime: (event.total - event.loaded) / speedInByte,
-          speed: speedInByte / 1024,
+          speed: speedInByte,
         });
       }
     });
@@ -41,7 +47,7 @@ class UploadService {
     await new Promise((resolve, reject) => {
       xhr.addEventListener('load', () => {
         if (xhr.status >= 200 && xhr.status < 300) {
-          onProgress('success', {
+          onProgress?.('success', {
             progress: 100,
             restTime: 0,
             speed: file.size / ((Date.now() - startTime) / 1000),
@@ -95,6 +101,7 @@ class UploadService {
 
   private getSignedUploadUrl = async (
     file: File,
+    directory?: string,
   ): Promise<
     FileMetadata & {
       preSignUrl: string;
@@ -104,7 +111,7 @@ class UploadService {
 
     // 精确到以 h 为单位的 path
     const date = (Date.now() / 1000 / 60 / 60).toFixed(0);
-    const dirname = `${fileEnv.NEXT_PUBLIC_S3_FILE_PATH}/${date}`;
+    const dirname = `${directory || fileEnv.NEXT_PUBLIC_S3_FILE_PATH}/${date}`;
     const pathname = `${dirname}/${filename}`;
 
     const preSignUrl = await edgeClient.upload.createS3PreSignedUrl.mutate({ pathname });