External dataset (#1497)

* perf: read rawText and chunk code

* perf: read raw text

* perf: read rawtext

* perf: token count

* log

.vscode/nextapi.code-snippets

@@ -20,9 +20,9 @@
 			"export type ${TM_FILENAME_BASE}Response = {};",
 			"",
 			"async function handler(",
-			"  req: ApiRequestProps<getDatasetTrainingQueueBody, getDatasetTrainingQueueQuery>,",
+			"  req: ApiRequestProps<${TM_FILENAME_BASE}Body, ${TM_FILENAME_BASE}Query>,",
 			"  res: ApiResponseType<any>",
-			"): Promise<getDatasetTrainingQueueResponse> {",
+			"): Promise<${TM_FILENAME_BASE}Response> {",
 			"    $1",
 			"    return {}",
 			"}",

packages/global/common/string/textSplitter.ts

@@ -9,6 +9,9 @@ type SplitProps = {
   overlapRatio?: number;
   customReg?: string[];
 };
+export type TextSplitProps = Omit<SplitProps, 'text' | 'chunkLen'> & {
+  chunkLen?: number;
+};
 
 type SplitResponse = {
   chunks: string[];
@@ -49,6 +52,7 @@ const strIsMdTable = (str: string) => {
       return false;
     }
   }
+
   return true;
 };
 const markdownTableSplit = (props: SplitProps): SplitResponse => {
@@ -77,6 +81,10 @@ ${mdSplitString}
     chunk += `${splitText2Lines[i]}\n`;
   }
 
+  if (chunk) {
+    chunks.push(chunk);
+  }
+
   return {
     chunks,
     chars: chunks.reduce((sum, chunk) => sum + chunk.length, 0)

packages/global/common/system/types/index.d.ts

@@ -66,6 +66,8 @@ export type SystemEnvType = {
   vectorMaxProcess: number;
   qaMaxProcess: number;
   pgHNSWEfSearch: number;
+  tokenWorkers: number; // token count max worker
+
   oneapiUrl?: string;
   chatApiKey?: string;
 };

packages/global/core/dataset/constants.ts

@@ -170,3 +170,10 @@ export const SearchScoreTypeMap = {
 
 export const CustomCollectionIcon = 'common/linkBlue';
 export const LinkCollectionIcon = 'common/linkBlue';
+
+/* source prefix */
+export enum DatasetSourceReadTypeEnum {
+  fileLocal = 'fileLocal',
+  link = 'link',
+  externalFile = 'externalFile'
+}

packages/global/core/dataset/read.ts

@@ -0,0 +1,16 @@
+import { DatasetSourceReadTypeEnum, ImportDataSourceEnum } from './constants';
+
+export const rawTextBackupPrefix = 'index,content';
+
+export const importType2ReadType = (type: ImportDataSourceEnum) => {
+  if (type === ImportDataSourceEnum.csvTable || type === ImportDataSourceEnum.fileLocal) {
+    return DatasetSourceReadTypeEnum.fileLocal;
+  }
+  if (type === ImportDataSourceEnum.fileLink) {
+    return DatasetSourceReadTypeEnum.link;
+  }
+  if (type === ImportDataSourceEnum.externalFile) {
+    return DatasetSourceReadTypeEnum.externalFile;
+  }
+  return DatasetSourceReadTypeEnum.link;
+};

packages/service/common/file/gridfs/controller.ts

@@ -151,12 +151,12 @@ export const readFileContentFromMongo = async ({
   teamId,
   bucketName,
   fileId,
-  csvFormat = false
+  isQAImport = false
 }: {
   teamId: string;
   bucketName: `${BucketNameEnum}`;
   fileId: string;
-  csvFormat?: boolean;
+  isQAImport?: boolean;
 }): Promise<{
   rawText: string;
   filename: string;
@@ -198,7 +198,7 @@ export const readFileContentFromMongo = async ({
 
   const { rawText } = await readFileRawContent({
     extension,
-    csvFormat,
+    isQAImport,
     teamId,
     buffer: fileBuffers,
     encoding,

packages/service/common/file/read/utils.ts

@@ -5,6 +5,7 @@ import { addHours } from 'date-fns';
 
 import { WorkerNameEnum, runWorker } from '../../../worker/utils';
 import { ReadFileResponse } from '../../../worker/file/type';
+import { rawTextBackupPrefix } from '@fastgpt/global/core/dataset/read';
 
 export const initMarkdownText = ({
   teamId,
@@ -29,36 +30,44 @@ export const initMarkdownText = ({
 
 export const readFileRawContent = async ({
   extension,
-  csvFormat,
+  isQAImport,
   teamId,
   buffer,
   encoding,
   metadata
 }: {
-  csvFormat?: boolean;
+  isQAImport?: boolean;
   extension: string;
   teamId: string;
   buffer: Buffer;
   encoding: string;
   metadata?: Record<string, any>;
 }) => {
-  const result = await runWorker<ReadFileResponse>(WorkerNameEnum.readFile, {
+  let { rawText, formatText } = await runWorker<ReadFileResponse>(WorkerNameEnum.readFile, {
     extension,
-    csvFormat,
     encoding,
     buffer
   });
 
   // markdown data format
   if (['md', 'html', 'docx'].includes(extension)) {
-    result.rawText = await initMarkdownText({
+    rawText = await initMarkdownText({
       teamId: teamId,
-      md: result.rawText,
+      md: rawText,
       metadata: metadata
     });
   }
 
-  return result;
+  if (['csv', 'xlsx'].includes(extension)) {
+    // qa data
+    if (isQAImport) {
+      rawText = rawText || '';
+    } else {
+      rawText = formatText || '';
+    }
+  }
+
+  return { rawText };
 };
 
 export const htmlToMarkdown = async (html?: string | null) => {

packages/service/common/string/cheerio.ts

@@ -77,9 +77,8 @@ export const urlsFetch = async ({
           $,
           selector
         });
-        console.log('html====', html);
+
         const md = await htmlToMarkdown(html);
-        console.log('html====', md);
 
         return {
           url,

packages/service/common/string/tiktoken/index.ts

@@ -12,27 +12,34 @@ import { getNanoid } from '@fastgpt/global/common/string/tools';
 import { addLog } from '../../system/log';
 
 export const getTiktokenWorker = () => {
-  if (global.tiktokenWorker) {
-    return global.tiktokenWorker;
+  const maxWorkers = global.systemEnv?.tokenWorkers || 20;
+
+  if (!global.tiktokenWorkers) {
+    global.tiktokenWorkers = [];
+  }
+
+  if (global.tiktokenWorkers.length >= maxWorkers) {
+    return global.tiktokenWorkers[Math.floor(Math.random() * global.tiktokenWorkers.length)];
   }
 
   const worker = getWorker(WorkerNameEnum.countGptMessagesTokens);
 
+  const i = global.tiktokenWorkers.push({
+    index: global.tiktokenWorkers.length,
+    worker,
+    callbackMap: {}
+  });
+
   worker.on('message', ({ id, data }: { id: string; data: number }) => {
-    const callback = global.tiktokenWorker?.callbackMap?.[id];
+    const callback = global.tiktokenWorkers[i - 1]?.callbackMap?.[id];
 
     if (callback) {
       callback?.(data);
-      delete global.tiktokenWorker.callbackMap[id];
+      delete global.tiktokenWorkers[i - 1].callbackMap[id];
     }
   });
 
-  global.tiktokenWorker = {
-    worker,
-    callbackMap: {}
-  };
-
-  return global.tiktokenWorker;
+  return global.tiktokenWorkers[i - 1];
 };
 
 export const countGptMessagesTokens = (
@@ -44,20 +51,29 @@ export const countGptMessagesTokens = (
     const start = Date.now();
 
     const { worker, callbackMap } = getTiktokenWorker();
+
     const id = getNanoid();
 
     const timer = setTimeout(() => {
-      resolve(0);
+      console.log('Count token Time out');
+      resolve(
+        messages.reduce((sum, item) => {
+          if (item.content) {
+            return sum + item.content.length * 0.5;
+          }
+          return sum;
+        }, 0)
+      );
       delete callbackMap[id];
-    }, 300);
+    }, 60000);
 
     callbackMap[id] = (data) => {
-      resolve(data);
-      clearTimeout(timer);
-
       // 检测是否有内存泄漏
-      // addLog.info(`Count token time: ${Date.now() - start}, token: ${data}`);
+      addLog.info(`Count token time: ${Date.now() - start}, token: ${data}`);
       // console.log(process.memoryUsage());
+
+      resolve(data);
+      clearTimeout(timer);
     };
 
     worker.postMessage({

packages/service/core/dataset/read.ts

@@ -0,0 +1,99 @@
+import { BucketNameEnum } from '@fastgpt/global/common/file/constants';
+import { DatasetSourceReadTypeEnum } from '@fastgpt/global/core/dataset/constants';
+import { readFileContentFromMongo } from '../../common/file/gridfs/controller';
+import { urlsFetch } from '../../common/string/cheerio';
+import { rawTextBackupPrefix } from '@fastgpt/global/core/dataset/read';
+import { parseCsvTable2Chunks } from './training/utils';
+import { TextSplitProps, splitText2Chunks } from '@fastgpt/global/common/string/textSplitter';
+import axios from 'axios';
+import { readFileRawContent } from '../../common/file/read/utils';
+
+export const readFileRawTextByUrl = async ({ teamId, url }: { teamId: string; url: string }) => {
+  const response = await axios({
+    method: 'get',
+    url: url,
+    responseType: 'arraybuffer'
+  });
+  const extension = url.split('.')?.pop()?.toLowerCase() || '';
+
+  const buffer = Buffer.from(response.data, 'binary');
+
+  const { rawText } = await readFileRawContent({
+    extension,
+    teamId,
+    buffer,
+    encoding: 'utf-8'
+  });
+
+  return rawText;
+};
+
+/* 
+    fileId - local file, read from mongo
+    link - request
+    externalFile = request read
+*/
+export const readDatasetSourceRawText = async ({
+  teamId,
+  type,
+  sourceId,
+  isQAImport,
+  selector
+}: {
+  teamId: string;
+  type: DatasetSourceReadTypeEnum;
+  sourceId: string;
+  isQAImport?: boolean;
+  selector?: string;
+}): Promise<string> => {
+  if (type === DatasetSourceReadTypeEnum.fileLocal) {
+    const { rawText } = await readFileContentFromMongo({
+      teamId,
+      bucketName: BucketNameEnum.dataset,
+      fileId: sourceId,
+      isQAImport
+    });
+    return rawText;
+  } else if (type === DatasetSourceReadTypeEnum.link) {
+    const result = await urlsFetch({
+      urlList: [sourceId],
+      selector
+    });
+
+    return result[0]?.content || '';
+  } else if (type === DatasetSourceReadTypeEnum.externalFile) {
+    const rawText = await readFileRawTextByUrl({
+      teamId,
+      url: sourceId
+    });
+    return rawText;
+  }
+
+  return '';
+};
+
+export const rawText2Chunks = ({
+  rawText,
+  isQAImport,
+  chunkLen = 512,
+  ...splitProps
+}: {
+  rawText: string;
+  isQAImport?: boolean;
+} & TextSplitProps) => {
+  if (isQAImport) {
+    const { chunks } = parseCsvTable2Chunks(rawText);
+    return chunks;
+  }
+
+  const { chunks } = splitText2Chunks({
+    text: rawText,
+    chunkLen,
+    ...splitProps
+  });
+
+  return chunks.map((item) => ({
+    q: item,
+    a: ''
+  }));
+};

packages/service/core/workflow/dispatch/tools/http468.ts

@@ -71,7 +71,7 @@ export const dispatchHttp468Request = async (props: HttpRequestProps): Promise<H
     chatId,
     responseChatItemId,
     ...variables,
-    histories: histories.slice(-10),
+    histories: histories?.slice(-10) || [],
     ...body,
     ...dynamicInput
   };

packages/service/core/workflow/dispatch/utils.ts

@@ -62,7 +62,10 @@ export const valueTypeFormat = (value: any, type?: WorkflowIOValueTypeEnum) => {
     return JSON.stringify(value);
   }
   if (type === 'number') return Number(value);
-  if (type === 'boolean') return value === 'true' ? true : false;
+  if (type === 'boolean') {
+    if (typeof value === 'string') return value === 'true';
+    return Boolean(value);
+  }
   try {
     if (type === WorkflowIOValueTypeEnum.datasetQuote && !Array.isArray(value)) {
       return JSON.parse(value);

packages/service/package.json

@@ -13,10 +13,10 @@
     "decompress": "^4.2.1",
     "domino-ext": "^2.1.4",
     "encoding": "^0.1.13",
+    "fastgpt-js-tiktoken": "^1.0.12",
     "file-type": "^19.0.0",
     "iconv-lite": "^0.6.3",
     "joplin-turndown-plugin-gfm": "^1.0.12",
-    "js-tiktoken": "^1.0.7",
     "json5": "^2.2.3",
     "jsonwebtoken": "^9.0.2",
     "mammoth": "^1.6.0",

packages/service/type.d.ts

@@ -20,8 +20,9 @@ declare global {
   var whisperModel: WhisperModelType;
   var reRankModels: ReRankModelItemType[];
 
-  var tiktokenWorker: {
+  var tiktokenWorkers: {
+    index: number;
     worker: Worker;
     callbackMap: Record<string, (e: number) => void>;
-  };
+  }[];
 }