Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit 7bcee82f by Archer Committed by GitHub
Options

perf: memory leak (#5370) 

* perf: memory leak

* perf: workflow share buffer;Circle checker;Get file from stream

* doc

* remove report.md
parent baf18b14
Showing with 518 additions and 342 deletions
SECURITY.md
...@@ -5,7 +5,7 @@ ...@@ -5,7 +5,7 @@
如果您发现了 FastGPT 的安全漏洞，请按照以下步骤进行报告： 如果您发现了 FastGPT 的安全漏洞，请按照以下步骤进行报告：
1. **报告方式** 1. **报告方式**
发送邮件至：dennis@sealos.io 发送邮件至：yujinlong@sealos.io
请备注版本以及您的 GitHub 账号 请备注版本以及您的 GitHub 账号
3. **响应时间** 3. **响应时间**
......
document/content/docs/protocol/privacy.en.mdx
...@@ -58,5 +58,5 @@ Due to servers potentially being located in different countries/regions, you agr ...@@ -58,5 +58,5 @@ Due to servers potentially being located in different countries/regions, you agr
**Contact Us** **Contact Us**
1. For any questions, suggestions, or complaints about this policy, contact us at: dennis@sealos.io. 1. For any questions, suggestions, or complaints about this policy, contact us at: yujinlong@sealos.io.
2. We will respond promptly and address your concerns. 2. We will respond promptly and address your concerns.
document/content/docs/protocol/privacy.mdx
...@@ -58,5 +58,5 @@ description: ' FastGPT 隐私政策' ...@@ -58,5 +58,5 @@ description: ' FastGPT 隐私政策'
**联系我们** **联系我们**
1. 如您对本隐私政策有任何疑问、建议或投诉，请通过以下方式与我们联系：dennis@sealos.io。 1. 如您对本隐私政策有任何疑问、建议或投诉，请通过以下方式与我们联系：yujinlong@sealos.io。
2. 我们将尽快回复并解决您提出的问题。 2. 我们将尽快回复并解决您提出的问题。
document/content/docs/protocol/terms.en.mdx
...@@ -72,4 +72,4 @@ This FastGPT Service Agreement constitutes the terms and conditions agreed betwe ...@@ -72,4 +72,4 @@ This FastGPT Service Agreement constitutes the terms and conditions agreed betwe
**Article 7 Additional Provisions** **Article 7 Additional Provisions**
1. If any clause is deemed unlawful or invalid, the remaining provisions shall remain enforceable. 1. If any clause is deemed unlawful or invalid, the remaining provisions shall remain enforceable.
2. Sealos retains final authority in interpreting this Agreement and privacy policies. For any inquiries, please contact us at dennis@sealos.io. 2. Sealos retains final authority in interpreting this Agreement and privacy policies. For any inquiries, please contact us at yujinlong@sealos.io.
document/content/docs/protocol/terms.mdx
...@@ -67,4 +67,4 @@ FastGPT 服务协议是您与珠海环界云计算有限公司（以下简称“ ...@@ -67,4 +67,4 @@ FastGPT 服务协议是您与珠海环界云计算有限公司（以下简称“
**第7条 其他条款** **第7条 其他条款**
1. 如本协议中部分条款因违反法律法规而被视为无效，不影响其他条款的效力。 1. 如本协议中部分条款因违反法律法规而被视为无效，不影响其他条款的效力。
2. 本公司保留对本协议及隐私政策的最终解释权。如您对本协议或隐私政策有任何疑问，请联系我们：dennis@sealos.io。 2. 本公司保留对本协议及隐私政策的最终解释权。如您对本协议或隐私政策有任何疑问，请联系我们：yujinlong@sealos.io。
document/content/docs/upgrading/4-11/4112.mdx 0 → 100644
---
title: 'V4.11.2'
description: 'FastGPT V4.11.2 更新说明'
---
## 🚀 新增内容
## ⚙️ 优化
1. 优化 3 处存在潜在内存泄露的代码。
2. 优化工作流部分递归检查，避免无限递归。
3. 优化文档阅读 Worker，采用 ShareBuffer 避免数据拷贝。
## 🐛 修复
1. Doc2x API 更新，导致解析失败。
## 🔨 工具更新
document/content/docs/upgrading/4-11/meta.json
{ {
"title": "4.11.x", "title": "4.11.x",
"description": "", "description": "",
"pages": ["4111", "4110"] "pages": ["4112", "4111", "4110"]
} }
document/data/doc-last-modified.json
{ {
"document/content/docs/api/api1.mdx": "2025-07-23T21:35:03+08:00",
"document/content/docs/api/api2.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/api/index.mdx": "2025-07-30T15:38:30+08:00",
"document/content/docs/api/test/api3.mdx": "2025-07-23T21:35:03+08:00",
"document/content/docs/faq/app.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/faq/app.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/faq/chat.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/faq/chat.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/faq/dataset.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/faq/dataset.mdx": "2025-08-02T19:38:37+08:00",
...@@ -96,14 +92,15 @@ ...@@ -96,14 +92,15 @@
"document/content/docs/protocol/index.mdx": "2025-07-30T15:38:30+08:00", "document/content/docs/protocol/index.mdx": "2025-07-30T15:38:30+08:00",
"document/content/docs/protocol/open-source.en.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/open-source.en.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/protocol/open-source.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/open-source.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/protocol/privacy.en.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/privacy.en.mdx": "2025-08-03T14:03:08+08:00",
"document/content/docs/protocol/privacy.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/privacy.mdx": "2025-08-03T14:03:08+08:00",
"document/content/docs/protocol/terms.en.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/terms.en.mdx": "2025-08-03T14:03:08+08:00",
"document/content/docs/protocol/terms.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/protocol/terms.mdx": "2025-08-03T14:03:08+08:00",
"document/content/docs/upgrading/4-10/4100.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-10/4100.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-10/4101.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-10/4101.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-11/4110.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-11/4110.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-11/4111.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-11/4111.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-11/4112.mdx": "2025-08-03T14:03:08+08:00",
"document/content/docs/upgrading/4-8/40.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-8/40.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-8/41.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-8/41.mdx": "2025-08-02T19:38:37+08:00",
"document/content/docs/upgrading/4-8/42.mdx": "2025-08-02T19:38:37+08:00", "document/content/docs/upgrading/4-8/42.mdx": "2025-08-02T19:38:37+08:00",
......
packages/global/core/workflow/runtime/utils.ts
...@@ -23,13 +23,20 @@ import type { RuntimeEdgeItemType, RuntimeNodeItemType } from './type'; ...@@ -23,13 +23,20 @@ import type { RuntimeEdgeItemType, RuntimeNodeItemType } from './type';
export const extractDeepestInteractive = ( export const extractDeepestInteractive = (
interactive: WorkflowInteractiveResponseType interactive: WorkflowInteractiveResponseType
): WorkflowInteractiveResponseType => { ): WorkflowInteractiveResponseType => {
if ( const MAX_DEPTH = 100;
(interactive?.type === 'childrenInteractive' || interactive?.type === 'loopInteractive') && let current = interactive;
interactive.params?.childrenResponse let depth = 0;
while (
depth < MAX_DEPTH &&
(current?.type === 'childrenInteractive' || current?.type === 'loopInteractive') &&
current.params?.childrenResponse
) { ) {
return extractDeepestInteractive(interactive.params.childrenResponse); current = current.params.childrenResponse;
depth++;
} }
return interactive;
return current;
}; };
export const getMaxHistoryLimitFromNodes = (nodes: StoreNodeItemType[]): number => { export const getMaxHistoryLimitFromNodes = (nodes: StoreNodeItemType[]): number => {
let limit = 10; let limit = 10;
...@@ -294,22 +301,42 @@ export const checkNodeRunStatus = ({ ...@@ -294,22 +301,42 @@ export const checkNodeRunStatus = ({
const commonEdges: RuntimeEdgeItemType[] = []; const commonEdges: RuntimeEdgeItemType[] = [];
const recursiveEdges: RuntimeEdgeItemType[] = []; const recursiveEdges: RuntimeEdgeItemType[] = [];
const checkIsCircular = (edge: RuntimeEdgeItemType, visited: Set<string>): boolean => { const checkIsCircular = (startEdge: RuntimeEdgeItemType, initialVisited: string[]): boolean => {
if (edge.source === currentNode.nodeId) { const stack: Array<{ edge: RuntimeEdgeItemType; visited: Set<string> }> = [
return true; // 检测到环,并且环中包含当前节点 { edge: startEdge, visited: new Set(initialVisited) }
} ];
if (visited.has(edge.source)) {
return false; // 检测到环,但不包含当前节点(子节点成环) const MAX_DEPTH = 3000;
let iterations = 0;
while (stack.length > 0 && iterations < MAX_DEPTH) {
iterations++;
const { edge, visited } = stack.pop()!;
if (edge.source === currentNode.nodeId) {
return true; // 检测到环,并且环中包含当前节点
}
if (visited.has(edge.source)) {
continue; // 已访问过此节点，跳过（避免子环干扰）
}
const newVisited = new Set(visited);
newVisited.add(edge.source);
// 查找目标节点的 source edges 并加入栈中
const nextEdges = allEdges.filter((item) => item.target === edge.source);
for (const nextEdge of nextEdges) {
stack.push({ edge: nextEdge, visited: newVisited });
}
} }
visited.add(edge.source);
// 递归检测后面的 edge，如果有其中一个成环，则返回 true return false;
const nextEdges = allEdges.filter((item) => item.target === edge.source);
return nextEdges.some((nextEdge) => checkIsCircular(nextEdge, new Set(visited)));
}; };
sourceEdges.forEach((edge) => { sourceEdges.forEach((edge) => {
if (checkIsCircular(edge, new Set([currentNode.nodeId]))) { if (checkIsCircular(edge, [currentNode.nodeId])) {
recursiveEdges.push(edge); recursiveEdges.push(edge);
} else { } else {
commonEdges.push(edge); commonEdges.push(edge);
......
packages/service/common/file/gridfs/controller.ts
...@@ -79,6 +79,8 @@ export async function uploadFile({ ...@@ -79,6 +79,8 @@ export async function uploadFile({
.pipe(stream as any) .pipe(stream as any)
.on('finish', resolve) .on('finish', resolve)
.on('error', reject); .on('error', reject);
}).finally(() => {
readStream.destroy();
}); });
return String(stream.id); return String(stream.id);
......
packages/service/common/file/utils.ts
...@@ -2,6 +2,11 @@ import { isProduction } from '@fastgpt/global/common/system/constants'; ...@@ -2,6 +2,11 @@ import { isProduction } from '@fastgpt/global/common/system/constants';
import fs from 'fs'; import fs from 'fs';
import path from 'path'; import path from 'path';
export const getFileMaxSize = () => {
const mb = global.feConfigs?.uploadFileMaxSize || 1000;
return mb * 1024 * 1024;
};
export const removeFilesByPaths = (paths: string[]) => { export const removeFilesByPaths = (paths: string[]) => {
paths.forEach((path) => { paths.forEach((path) => {
fs.unlink(path, (err) => { fs.unlink(path, (err) => {
......
packages/service/common/mongo/init.ts
...@@ -13,17 +13,22 @@ export async function connectMongo(db: Mongoose, url: string): Promise<Mongoose> ...@@ -13,17 +13,22 @@ export async function connectMongo(db: Mongoose, url: string): Promise<Mongoose>
return db; return db;
} }
const RemoveListeners = () => {
db.connection.removeAllListeners('error');
db.connection.removeAllListeners('disconnected');
};
console.log('MongoDB start connect'); console.log('MongoDB start connect');
try { try {
// Remove existing listeners to prevent duplicates // Remove existing listeners to prevent duplicates
db.connection.removeAllListeners('error'); RemoveListeners();
db.connection.removeAllListeners('disconnected');
db.set('strictQuery', 'throw'); db.set('strictQuery', 'throw');
db.connection.on('error', async (error) => { db.connection.on('error', async (error) => {
console.log('mongo error', error); console.log('mongo error', error);
try { try {
if (db.connection.readyState !== 0) { if (db.connection.readyState !== 0) {
RemoveListeners();
await db.disconnect(); await db.disconnect();
await delay(1000); await delay(1000);
await connectMongo(db, url); await connectMongo(db, url);
...@@ -34,6 +39,7 @@ export async function connectMongo(db: Mongoose, url: string): Promise<Mongoose> ...@@ -34,6 +39,7 @@ export async function connectMongo(db: Mongoose, url: string): Promise<Mongoose>
console.log('mongo disconnected'); console.log('mongo disconnected');
try { try {
if (db.connection.readyState !== 0) { if (db.connection.readyState !== 0) {
RemoveListeners();
await db.disconnect(); await db.disconnect();
await delay(1000); await delay(1000);
await connectMongo(db, url); await connectMongo(db, url);
......
packages/service/core/dataset/read.ts
...@@ -13,6 +13,9 @@ import { getApiDatasetRequest } from './apiDataset'; ...@@ -13,6 +13,9 @@ import { getApiDatasetRequest } from './apiDataset';
import Papa from 'papaparse'; import Papa from 'papaparse';
import type { ApiDatasetServerType } from '@fastgpt/global/core/dataset/apiDataset/type'; import type { ApiDatasetServerType } from '@fastgpt/global/core/dataset/apiDataset/type';
import { text2Chunks } from '../../worker/function'; import { text2Chunks } from '../../worker/function';
import { addLog } from '../../common/system/log';
import { retryFn } from '@fastgpt/global/common/system/utils';
import { getFileMaxSize } from '../../common/file/utils';
export const readFileRawTextByUrl = async ({ export const readFileRawTextByUrl = async ({
teamId, teamId,
...@@ -20,7 +23,8 @@ export const readFileRawTextByUrl = async ({ ...@@ -20,7 +23,8 @@ export const readFileRawTextByUrl = async ({
url, url,
customPdfParse, customPdfParse,
getFormatText, getFormatText,
relatedId relatedId,
maxFileSize = getFileMaxSize()
}: { }: {
teamId: string; teamId: string;
tmbId: string; tmbId: string;
...@@ -28,30 +32,113 @@ export const readFileRawTextByUrl = async ({ ...@@ -28,30 +32,113 @@ export const readFileRawTextByUrl = async ({
customPdfParse?: boolean; customPdfParse?: boolean;
getFormatText?: boolean; getFormatText?: boolean;
relatedId: string; // externalFileId / apiFileId relatedId: string; // externalFileId / apiFileId
maxFileSize?: number;
}) => { }) => {
const extension = parseFileExtensionFromUrl(url);
// Check file size
try {
const headResponse = await axios.head(url, { timeout: 10000 });
const contentLength = parseInt(headResponse.headers['content-length'] || '0');
if (contentLength > 0 && contentLength > maxFileSize) {
return Promise.reject(
`File too large. Size: ${Math.round(contentLength / 1024 / 1024)}MB, Maximum allowed: ${Math.round(maxFileSize / 1024 / 1024)}MB`
);
}
} catch (error) {
addLog.warn('Check file HEAD request failed');
}
// Use stream response type, avoid double memory usage
const response = await axios({ const response = await axios({
method: 'get', method: 'get',
url: url, url: url,
responseType: 'arraybuffer' responseType: 'stream',
maxContentLength: maxFileSize,
timeout: 30000
}); });
const extension = parseFileExtensionFromUrl(url);
const buffer = Buffer.from(response.data, 'binary'); // 优化：直接从 stream 转换为 buffer，避免 arraybuffer 中间步骤
const chunks: Buffer[] = [];
let totalLength = 0;
const { rawText } = await readRawContentByFileBuffer({ return new Promise<string>((resolve, reject) => {
customPdfParse, let isAborted = false;
getFormatText,
extension, const cleanup = () => {
teamId, if (!isAborted) {
tmbId, isAborted = true;
buffer, chunks.length = 0; // 清理内存
encoding: 'utf-8', response.data.destroy();
metadata: { }
relatedId };
}
}); // Stream timeout
const timeoutId = setTimeout(() => {
cleanup();
reject('File download timeout after 30 seconds');
}, 600000);
return rawText; response.data.on('data', (chunk: Buffer) => {
if (isAborted) return;
totalLength += chunk.length;
if (totalLength > maxFileSize) {
clearTimeout(timeoutId);
cleanup();
return reject(
`File too large. Maximum size allowed is ${Math.round(maxFileSize / 1024 / 1024)}MB.`
);
}
chunks.push(chunk);
});
response.data.on('end', async () => {
if (isAborted) return;
clearTimeout(timeoutId);
try {
// 合并所有 chunks 为单个 buffer
const buffer = Buffer.concat(chunks);
// 立即清理 chunks 数组释放内存
chunks.length = 0;
const { rawText } = await retryFn(() =>
readRawContentByFileBuffer({
customPdfParse,
getFormatText,
extension,
teamId,
tmbId,
buffer,
encoding: 'utf-8',
metadata: {
relatedId
}
})
);
resolve(rawText);
} catch (error) {
cleanup();
reject(error);
}
});
response.data.on('error', (error: Error) => {
clearTimeout(timeoutId);
cleanup();
reject(error);
});
response.data.on('close', () => {
clearTimeout(timeoutId);
cleanup();
});
});
}; };
/* /*
......
packages/service/core/workflow/dispatch/index.ts
...@@ -180,6 +180,7 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons ...@@ -180,6 +180,7 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons
} }
let workflowRunTimes = 0; let workflowRunTimes = 0;
let streamCheckTimer: NodeJS.Timeout | null = null;
// Init // Init
if (isRootRuntime) { if (isRootRuntime) {
...@@ -198,18 +199,14 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons ...@@ -198,18 +199,14 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons
res.setHeader('Cache-Control', 'no-cache, no-transform'); res.setHeader('Cache-Control', 'no-cache, no-transform');
// 10s sends a message to prevent the browser from thinking that the connection is disconnected // 10s sends a message to prevent the browser from thinking that the connection is disconnected
const sendStreamTimerSign = () => { streamCheckTimer = setInterval(() => {
setTimeout(() => { props?.workflowStreamResponse?.({
props?.workflowStreamResponse?.({ event: SseResponseEventEnum.answer,
event: SseResponseEventEnum.answer, data: textAdaptGptResponse({
data: textAdaptGptResponse({ text: ''
text: '' })
}) });
}); }, 10000);
sendStreamTimerSign();
}, 10000);
};
sendStreamTimerSign();
} }
// Get default variables // Get default variables
...@@ -841,6 +838,10 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons ...@@ -841,6 +838,10 @@ export async function dispatchWorkFlow(data: Props): Promise<DispatchFlowRespons
}; };
} catch (error) { } catch (error) {
return Promise.reject(error); return Promise.reject(error);
} finally {
if (streamCheckTimer) {
clearInterval(streamCheckTimer);
}
} }
} }
......
packages/service/thirdProvider/doc2x/index.ts
...@@ -37,15 +37,15 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -37,15 +37,15 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
if (typeof err === 'string') { if (typeof err === 'string') {
return Promise.reject({ message: `[Doc2x] ${err}` }); return Promise.reject({ message: `[Doc2x] ${err}` });
} }
if (typeof err.message === 'string') {
return Promise.reject({ message: `[Doc2x] ${err.message}` });
}
if (typeof err.data === 'string') { if (typeof err.data === 'string') {
return Promise.reject({ message: `[Doc2x] ${err.data}` }); return Promise.reject({ message: `[Doc2x] ${err.data}` });
} }
if (err?.response?.data) { if (err?.response?.data) {
return Promise.reject({ message: `[Doc2x] ${getErrText(err?.response?.data)}` }); return Promise.reject({ message: `[Doc2x] ${getErrText(err?.response?.data)}` });
} }
if (typeof err.message === 'string') {
return Promise.reject({ message: `[Doc2x] ${err.message}` });
}
addLog.error('[Doc2x] Unknown error', err); addLog.error('[Doc2x] Unknown error', err);
return Promise.reject({ message: `[Doc2x] ${getErrText(err)}` }); return Promise.reject({ message: `[Doc2x] ${getErrText(err)}` });
...@@ -78,7 +78,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -78,7 +78,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
code, code,
msg, msg,
data: preupload_data data: preupload_data
} = await request<{ uid: string; url: string }>('/v2/parse/preupload', null, 'POST'); } = await request<{ uid: string; url: string }>('/v2/parse/preupload', {}, 'POST');
if (!['ok', 'success'].includes(code)) { if (!['ok', 'success'].includes(code)) {
return Promise.reject(`[Doc2x] Failed to get pre-upload URL: ${msg}`); return Promise.reject(`[Doc2x] Failed to get pre-upload URL: ${msg}`);
} }
...@@ -96,6 +96,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -96,6 +96,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
.catch((error) => { .catch((error) => {
return Promise.reject(`[Doc2x] Failed to upload file: ${getErrText(error)}`); return Promise.reject(`[Doc2x] Failed to upload file: ${getErrText(error)}`);
}); });
if (response.status !== 200) { if (response.status !== 200) {
return Promise.reject( return Promise.reject(
`[Doc2x] Upload failed with status ${response.status}: ${response.statusText}` `[Doc2x] Upload failed with status ${response.status}: ${response.statusText}`
......
packages/service/worker/function.ts
...@@ -20,5 +20,17 @@ export const readRawContentFromBuffer = (props: { ...@@ -20,5 +20,17 @@ export const readRawContentFromBuffer = (props: {
encoding: string; encoding: string;
buffer: Buffer; buffer: Buffer;
}) => { }) => {
return runWorker<ReadFileResponse>(WorkerNameEnum.readFile, props); const bufferSize = props.buffer.length;
// 使用 SharedArrayBuffer，避免数据复制
const sharedBuffer = new SharedArrayBuffer(bufferSize);
const sharedArray = new Uint8Array(sharedBuffer);
sharedArray.set(props.buffer);
return runWorker<ReadFileResponse>(WorkerNameEnum.readFile, {
extension: props.extension,
encoding: props.encoding,
sharedBuffer: sharedBuffer,
bufferSize: bufferSize
});
}; };
packages/service/worker/readFile/extension/pdf.ts
...@@ -56,16 +56,16 @@ export const readPdfFile = async ({ buffer }: ReadRawTextByBuffer): Promise<Read ...@@ -56,16 +56,16 @@ export const readPdfFile = async ({ buffer }: ReadRawTextByBuffer): Promise<Read
} }
}; };
// @ts-ignore // Create a completely new ArrayBuffer to avoid SharedArrayBuffer transferList issues
const loadingTask = pdfjs.getDocument(buffer.buffer); const uint8Array = new Uint8Array(buffer.byteLength);
uint8Array.set(new Uint8Array(buffer.buffer, buffer.byteOffset, buffer.byteLength));
const loadingTask = pdfjs.getDocument({ data: uint8Array });
const doc = await loadingTask.promise; const doc = await loadingTask.promise;
// Avoid OOM.
let result = '';
const pageArr = Array.from({ length: doc.numPages }, (_, i) => i + 1); const pageArr = Array.from({ length: doc.numPages }, (_, i) => i + 1);
for (let i = 0; i < pageArr.length; i++) { const result = (
result += await readPDFPage(doc, i + 1); await Promise.all(pageArr.map(async (page) => await readPDFPage(doc, page)))
} ).join('');
loadingTask.destroy(); loadingTask.destroy();
......
packages/service/worker/readFile/index.ts
...@@ -9,49 +9,60 @@ import { readXlsxRawText } from './extension/xlsx'; ...@@ -9,49 +9,60 @@ import { readXlsxRawText } from './extension/xlsx';
import { readCsvRawText } from './extension/csv'; import { readCsvRawText } from './extension/csv';
import { workerResponse } from '../controller'; import { workerResponse } from '../controller';
parentPort?.on('message', async (props: ReadRawTextProps<Uint8Array>) => { parentPort?.on(
const read = async (params: ReadRawTextByBuffer) => { 'message',
switch (params.extension) { async (
case 'txt': props: Omit<ReadRawTextProps<any>, 'buffer'> & {
case 'md': sharedBuffer: SharedArrayBuffer;
return readFileRawText(params); bufferSize: number;
case 'html':
return readHtmlRawText(params);
case 'pdf':
return readPdfFile(params);
case 'docx':
return readDocsFile(params);
case 'pptx':
return readPptxRawText(params);
case 'xlsx':
return readXlsxRawText(params);
case 'csv':
return readCsvRawText(params);
default:
return Promise.reject(
`Only support .txt, .md, .html, .pdf, .docx, pptx, .csv, .xlsx. "${params.extension}" is not supported.`
);
} }
}; ) => {
const read = async (params: ReadRawTextByBuffer) => {
switch (params.extension) {
case 'txt':
case 'md':
return readFileRawText(params);
case 'html':
return readHtmlRawText(params);
case 'pdf':
return readPdfFile(params);
case 'docx':
return readDocsFile(params);
case 'pptx':
return readPptxRawText(params);
case 'xlsx':
return readXlsxRawText(params);
case 'csv':
return readCsvRawText(params);
default:
return Promise.reject(
`Only support .txt, .md, .html, .pdf, .docx, pptx, .csv, .xlsx. "${params.extension}" is not supported.`
);
}
};
// params.buffer: Uint8Array -> buffer // 使用 SharedArrayBuffer，零拷贝共享内存
const buffer = Buffer.from(props.buffer); const sharedArray = new Uint8Array(props.sharedBuffer);
const newProps: ReadRawTextByBuffer = { const buffer = Buffer.from(sharedArray.buffer, 0, props.bufferSize);
...props,
buffer
};
try { const newProps: ReadRawTextByBuffer = {
workerResponse({ extension: props.extension,
parentPort, encoding: props.encoding,
status: 'success', buffer
data: await read(newProps) };
});
} catch (error) { try {
workerResponse({ workerResponse({
parentPort, parentPort,
status: 'error', status: 'success',
data: error data: await read(newProps)
}); });
} catch (error) {
workerResponse({
parentPort,
status: 'error',
data: error
});
}
} }
}); );
packages/service/worker/utils.ts
...@@ -198,6 +198,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -198,6 +198,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
if (item) { if (item) {
item.reject?.('error'); item.reject?.('error');
clearTimeout(item.timeoutId); clearTimeout(item.timeoutId);
item.worker.removeAllListeners();
item.worker.terminate(); item.worker.terminate();
} }
......
projects/app/src/service/core/dataset/queues/generateQA.ts
...@@ -44,157 +44,161 @@ export async function generateQA(): Promise<any> { ...@@ -44,157 +44,161 @@ export async function generateQA(): Promise<any> {
if (global.qaQueueLen >= max) return; if (global.qaQueueLen >= max) return;
global.qaQueueLen++; global.qaQueueLen++;
while (true) { try {
const startTime = Date.now(); while (true) {
// get training data const startTime = Date.now();
const { // get training data
data, const {
text, data,
done = false, text,
error = false done = false,
} = await (async () => { error = false
try { } = await (async () => {
const data = await MongoDatasetTraining.findOneAndUpdate( try {
{ const data = await MongoDatasetTraining.findOneAndUpdate(
mode: TrainingModeEnum.qa,
retryCount: { $gt: 0 },
lockTime: { $lte: addMinutes(new Date(), -10) }
},
{
lockTime: new Date(),
$inc: { retryCount: -1 }
}
)
.populate<PopulateType>([
{ {
path: 'dataset', mode: TrainingModeEnum.qa,
select: 'agentModel vectorModel vlmModel' retryCount: { $gt: 0 },
lockTime: { $lte: addMinutes(new Date(), -10) }
}, },
{ {
path: 'collection', lockTime: new Date(),
select: 'qaPrompt' $inc: { retryCount: -1 }
} }
]) )
.lean(); .populate<PopulateType>([
{
// task preemption path: 'dataset',
if (!data) { select: 'agentModel vectorModel vlmModel'
},
{
path: 'collection',
select: 'qaPrompt'
}
])
.lean();
// task preemption
if (!data) {
return {
done: true
};
}
return {
data,
text: data.q
};
} catch (error) {
return { return {
done: true error: true
}; };
} }
return { })();
data,
text: data.q
};
} catch (error) {
return {
error: true
};
}
})();
if (done || !data) { if (done || !data) {
break; break;
} }
if (error) { if (error) {
addLog.error(`[QA Queue] Error`, error); addLog.error(`[QA Queue] Error`, error);
await delay(500); await delay(500);
continue; continue;
} }
if (!data.dataset || !data.collection) { if (!data.dataset || !data.collection) {
addLog.info(`[QA Queue] Dataset or collection not found`, data); addLog.info(`[QA Queue] Dataset or collection not found`, data);
// Delete data // Delete data
await MongoDatasetTraining.deleteOne({ _id: data._id }); await MongoDatasetTraining.deleteOne({ _id: data._id });
continue; continue;
} }
// auth balance // auth balance
if (!(await checkTeamAiPointsAndLock(data.teamId))) { if (!(await checkTeamAiPointsAndLock(data.teamId))) {
continue; continue;
} }
addLog.info(`[QA Queue] Start`); addLog.info(`[QA Queue] Start`);
try { try {
const modelData = getLLMModel(data.dataset.agentModel); const modelData = getLLMModel(data.dataset.agentModel);
const prompt = `${data.collection.qaPrompt || Prompt_AgentQA.description} const prompt = `${data.collection.qaPrompt || Prompt_AgentQA.description}
${replaceVariable(Prompt_AgentQA.fixedText, { text })}`; ${replaceVariable(Prompt_AgentQA.fixedText, { text })}`;
// request LLM to get QA // request LLM to get QA
const messages: ChatCompletionMessageParam[] = [ const messages: ChatCompletionMessageParam[] = [
{ {
role: 'user', role: 'user',
content: prompt content: prompt
} }
]; ];
const { response: chatResponse } = await createChatCompletion({ const { response: chatResponse } = await createChatCompletion({
body: llmCompletionsBodyFormat( body: llmCompletionsBodyFormat(
{
model: modelData.model,
temperature: 0.3,
messages: await loadRequestMessages({ messages, useVision: false }),
stream: true
},
modelData
)
});
const { text: answer, usage } = await formatLLMResponse(chatResponse);
const inputTokens = usage?.prompt_tokens || (await countGptMessagesTokens(messages));
const outputTokens = usage?.completion_tokens || (await countPromptTokens(answer));
const qaArr = await formatSplitText({ answer, rawText: text, llmModel: modelData }); // 格式化后的QA对
// get vector and insert
await pushDataListToTrainingQueue({
teamId: data.teamId,
tmbId: data.tmbId,
datasetId: data.datasetId,
collectionId: data.collectionId,
mode: TrainingModeEnum.chunk,
data: qaArr.map((item) => ({
...item,
chunkIndex: data.chunkIndex
})),
billId: data.billId,
vectorModel: data.dataset.vectorModel,
agentModel: data.dataset.agentModel,
vlmModel: data.dataset.vlmModel
});
// delete data from training
await MongoDatasetTraining.findByIdAndDelete(data._id);
// Push usage
pushLLMTrainingUsage({
teamId: data.teamId,
tmbId: data.tmbId,
inputTokens,
outputTokens,
billId: data.billId,
model: modelData.model,
mode: 'qa'
});
addLog.info(`[QA Queue] Finish`, {
time: Date.now() - startTime,
splitLength: qaArr.length,
usage
});
} catch (err: any) {
addLog.error(`[QA Queue] Error`, err);
await MongoDatasetTraining.updateOne(
{ {
model: modelData.model, _id: data._id
temperature: 0.3,
messages: await loadRequestMessages({ messages, useVision: false }),
stream: true
}, },
modelData {
) errorMsg: getErrText(err, 'unknown error')
}); }
const { text: answer, usage } = await formatLLMResponse(chatResponse); );
const inputTokens = usage?.prompt_tokens || (await countGptMessagesTokens(messages));
const outputTokens = usage?.completion_tokens || (await countPromptTokens(answer));
const qaArr = await formatSplitText({ answer, rawText: text, llmModel: modelData }); // 格式化后的QA对
// get vector and insert
await pushDataListToTrainingQueue({
teamId: data.teamId,
tmbId: data.tmbId,
datasetId: data.datasetId,
collectionId: data.collectionId,
mode: TrainingModeEnum.chunk,
data: qaArr.map((item) => ({
...item,
chunkIndex: data.chunkIndex
})),
billId: data.billId,
vectorModel: data.dataset.vectorModel,
agentModel: data.dataset.agentModel,
vlmModel: data.dataset.vlmModel
});
// delete data from training
await MongoDatasetTraining.findByIdAndDelete(data._id);
// Push usage
pushLLMTrainingUsage({
teamId: data.teamId,
tmbId: data.tmbId,
inputTokens,
outputTokens,
billId: data.billId,
model: modelData.model,
mode: 'qa'
});
addLog.info(`[QA Queue] Finish`, {
time: Date.now() - startTime,
splitLength: qaArr.length,
usage
});
} catch (err: any) {
addLog.error(`[QA Queue] Error`, err);
await MongoDatasetTraining.updateOne(
{
_id: data._id
},
{
errorMsg: getErrText(err, 'unknown error')
}
);
await delay(100); await delay(100);
}
} }
} catch (error) {
addLog.error(`[QA Queue] Error`, error);
} }
if (reduceQueue()) { if (reduceQueue()) {
......
projects/app/src/service/core/dataset/queues/generateVector.ts
...@@ -42,116 +42,120 @@ export async function generateVector(): Promise<any> { ...@@ -42,116 +42,120 @@ export async function generateVector(): Promise<any> {
if (global.vectorQueueLen >= max) return; if (global.vectorQueueLen >= max) return;
global.vectorQueueLen++; global.vectorQueueLen++;
while (true) { try {
const start = Date.now(); while (true) {
const start = Date.now();
// get training data // get training data
const { const {
data, data,
done = false, done = false,
error = false error = false
} = await (async () => { } = await (async () => {
try { try {
const data = await MongoDatasetTraining.findOneAndUpdate( const data = await MongoDatasetTraining.findOneAndUpdate(
{
mode: TrainingModeEnum.chunk,
retryCount: { $gt: 0 },
lockTime: { $lte: addMinutes(new Date(), -3) }
},
{
lockTime: new Date(),
$inc: { retryCount: -1 }
}
)
.populate<PopulateType>([
{
path: 'dataset',
select: 'vectorModel'
},
{ {
path: 'collection', mode: TrainingModeEnum.chunk,
select: 'name indexPrefixTitle' retryCount: { $gt: 0 },
lockTime: { $lte: addMinutes(new Date(), -3) }
}, },
{ {
path: 'data', lockTime: new Date(),
select: '_id indexes' $inc: { retryCount: -1 }
} }
]) )
.lean(); .populate<PopulateType>([
{
path: 'dataset',
select: 'vectorModel'
},
{
path: 'collection',
select: 'name indexPrefixTitle'
},
{
path: 'data',
select: '_id indexes'
}
])
.lean();
// task preemption // task preemption
if (!data) { if (!data) {
return {
done: true
};
}
return {
data
};
} catch (error) {
return { return {
done: true error: true
}; };
} }
return { })();
data
};
} catch (error) {
return {
error: true
};
}
})();
// Break loop // Break loop
if (done || !data) { if (done || !data) {
break; break;
} }
if (error) { if (error) {
addLog.error(`[Vector Queue] Error`, error); addLog.error(`[Vector Queue] Error`, error);
await delay(500); await delay(500);
continue; continue;
} }
if (!data.dataset || !data.collection) { if (!data.dataset || !data.collection) {
addLog.info(`[Vector Queue] Dataset or collection not found`, data); addLog.info(`[Vector Queue] Dataset or collection not found`, data);
// Delete data // Delete data
await MongoDatasetTraining.deleteOne({ _id: data._id }); await MongoDatasetTraining.deleteOne({ _id: data._id });
continue; continue;
} }
// auth balance // auth balance
if (!(await checkTeamAiPointsAndLock(data.teamId))) { if (!(await checkTeamAiPointsAndLock(data.teamId))) {
continue; continue;
} }
addLog.info(`[Vector Queue] Start`); addLog.info(`[Vector Queue] Start`);
try { try {
const { tokens } = await (async () => { const { tokens } = await (async () => {
if (data.dataId) { if (data.dataId) {
return rebuildData({ trainingData: data }); return rebuildData({ trainingData: data });
} else { } else {
return insertData({ trainingData: data }); return insertData({ trainingData: data });
} }
})(); })();
// push usage // push usage
pushGenerateVectorUsage({ pushGenerateVectorUsage({
teamId: data.teamId, teamId: data.teamId,
tmbId: data.tmbId, tmbId: data.tmbId,
inputTokens: tokens, inputTokens: tokens,
model: data.dataset.vectorModel, model: data.dataset.vectorModel,
billId: data.billId billId: data.billId
}); });
addLog.info(`[Vector Queue] Finish`, { addLog.info(`[Vector Queue] Finish`, {
time: Date.now() - start time: Date.now() - start
}); });
} catch (err: any) { } catch (err: any) {
addLog.error(`[Vector Queue] Error`, err); addLog.error(`[Vector Queue] Error`, err);
await MongoDatasetTraining.updateOne( await MongoDatasetTraining.updateOne(
{ {
_id: data._id _id: data._id
}, },
{ {
errorMsg: getErrText(err, 'unknown error') errorMsg: getErrText(err, 'unknown error')
} }
); );
await delay(100); await delay(100);
}
} }
} catch (error) {
addLog.error(`[Vector Queue] Error`, error);
} }
if (reduceQueue()) { if (reduceQueue()) {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment