Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit 4c4384a4 by Archer Committed by GitHub
Options

feat: upload parsed docx images from read file worker (#7072) 

* feat: upload parsed docx images from read file worker

* fix: pass docx image upload as array buffer

* fix: reuse read files parser in agent file read

* fix: upload dataset images from buffers

* docgit add .

* refactor: upload parsed images in workers

* fix: guard parsed image uploads
parent 1d7b8768
Showing with 2055 additions and 427 deletions
document/content/self-host/upgrading/4-15/41504.mdx
...@@ -53,8 +53,10 @@ fastgpt-plugin: ...@@ -53,8 +53,10 @@ fastgpt-plugin:
1. 模型获取多模态文件链接异常。 1. 模型获取多模态文件链接异常。
2. 修复 training 接口存在的潜在越权风险。 2. 修复 training 接口存在的潜在越权风险。
3. HTTP tool parse 的 SSRF 风险。
## 🛠️ 代码优化 ## 🛠️ 代码优化
1. 插件服务从旧 `runtime` 结构调整为 pnpm workspace monorepo，拆分为 HTTP 服务入口、领域模型、用例、API adapter、基础设施、SDK 和 CLI。 1. 插件服务从旧 `runtime` 结构调整为 pnpm workspace monorepo，拆分为 HTTP 服务入口、领域模型、用例、API adapter、基础设施、SDK 和 CLI。
2. 将 app API 接口全部用 zod schema 编写并生成文档。 2. 将 app API 接口全部用 zod schema 编写并生成文档。
3. 及时处理 worker 内图片，不再存留 base64，降低内存消耗。
document/data/doc-last-modified.json
...@@ -274,8 +274,8 @@ ...@@ -274,8 +274,8 @@
"content/self-host/upgrading/4-15/41502.mdx": "2026-05-25T11:21:30+08:00", "content/self-host/upgrading/4-15/41502.mdx": "2026-05-25T11:21:30+08:00",
"content/self-host/upgrading/4-15/41503.en.mdx": "2026-05-28T16:21:09+08:00", "content/self-host/upgrading/4-15/41503.en.mdx": "2026-05-28T16:21:09+08:00",
"content/self-host/upgrading/4-15/41503.mdx": "2026-05-28T16:21:09+08:00", "content/self-host/upgrading/4-15/41503.mdx": "2026-05-28T16:21:09+08:00",
"content/self-host/upgrading/4-15/41504.en.mdx": "2026-06-01T17:19:55+08:00", "content/self-host/upgrading/4-15/41504.en.mdx": "2026-06-07T17:54:48+08:00",
"content/self-host/upgrading/4-15/41504.mdx": "2026-06-05T18:12:32+08:00", "content/self-host/upgrading/4-15/41504.mdx": "2026-06-07T17:54:48+08:00",
"content/self-host/upgrading/outdated/40.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/40.en.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/outdated/40.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/40.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/outdated/41.en.mdx": "2026-04-26T21:08:47+08:00", "content/self-host/upgrading/outdated/41.en.mdx": "2026-04-26T21:08:47+08:00",
......
packages/global/common/string/markdown.ts
import { batchRun } from '../system/utils'; import { batchRun } from '../system/utils';
import { getNanoid, simpleText } from './tools'; import { simpleText } from './tools';
import type { ImageType } from '../../../service/worker/readFile/type';
/* Delete redundant text in markdown */ /* Delete redundant text in markdown */
export const simpleMarkdownText = (rawText: string) => { export const simpleMarkdownText = (rawText: string) => {
...@@ -27,7 +26,7 @@ export const simpleMarkdownText = (rawText: string) => { ...@@ -27,7 +26,7 @@ export const simpleMarkdownText = (rawText: string) => {
rawText = rawText.replace(/\\\\n/g, '\\n'); rawText = rawText.replace(/\\\\n/g, '\\n');
// Remove headings and code blocks front spaces // Remove headings and code blocks front spaces
['####', '###', '##', '#', '```', '~~~'].forEach((item, i) => { ['####', '###', '##', '#', '```', '~~~'].forEach((item) => {
const reg = new RegExp(`\\n\\s*${item}`, 'g'); const reg = new RegExp(`\\n\\s*${item}`, 'g');
if (reg.test(rawText)) { if (reg.test(rawText)) {
rawText = rawText.replace(new RegExp(`(\\n)( *)(${item})`, 'g'), '$1$3'); rawText = rawText.replace(new RegExp(`(\\n)( *)(${item})`, 'g'), '$1$3');
...@@ -46,7 +45,7 @@ export const htmlTable2Md = (content: string): string => { ...@@ -46,7 +45,7 @@ export const htmlTable2Md = (content: string): string => {
if (!rows) return htmlTable; if (!rows) return htmlTable;
// Parse table data // Parse table data
let tableData: string[][] = []; const tableData: string[][] = [];
let maxColumns = 0; let maxColumns = 0;
// Try to convert to markdown table // Try to convert to markdown table
...@@ -109,92 +108,224 @@ export const htmlTable2Md = (content: string): string => { ...@@ -109,92 +108,224 @@ export const htmlTable2Md = (content: string): string => {
}); });
return chunks.join('\n'); return chunks.join('\n');
} catch (error) { } catch {
return htmlTable; return htmlTable;
} }
}); });
}; };
/** export type MatchedImageUploadResult = {
* format markdown key: string;
* 1. upload base64 previewUrl?: string;
* 2. replace \ };
*/
export const uploadMarkdownBase64 = async ({ type MarkdownImageBase = {
rawText, altText: string;
uploadImgController url: string;
}: { fullMatch: string;
rawText: string; index: number;
uploadImgController?: (base64: string) => Promise<string>; };
}) => {
if (uploadImgController) { export type MarkdownImage = MarkdownImageBase &
// match base64, upload and replace it (
const base64Regex = /data:image\/.*;base64,([^\)]+)/g; | {
const base64Arr = rawText.match(base64Regex) || []; type: 'base64';
dataUrl: string;
// upload base64 and replace it mime: string;
await batchRun( base64: string;
base64Arr, }
async (base64Img) => { | {
try { type: 'http';
const str = await uploadImgController(base64Img);
rawText = rawText.replace(base64Img, str);
} catch (error) {
rawText = rawText.replace(base64Img, '');
rawText = rawText.replace(/!\[.*\]\(\)/g, '');
} }
},
20
); );
type MarkdownImageUploadController = (image: MarkdownImage) => Promise<MatchedImageUploadResult>;
export type MarkdownImageParseOptions = {
parseBase64?: boolean;
parseHttp?: boolean;
controller?: MarkdownImageUploadController;
controler?: MarkdownImageUploadController;
};
const mdBase64ImageSrcRegex = /^data:image\/([^;]+);base64,([A-Za-z0-9+/=]+)$/;
const mdHttpImageSrcRegex = /^https?:\/\/.+/;
const markdownImageUploadConcurrency = 5;
const unescapeMarkdownUrl = (url: string) => url.replace(/\\([\\()])/g, '$1');
const findClosingBracket = (text: string, startIndex: number) => {
for (let i = startIndex; i < text.length; i++) {
if (text[i] === '\\') {
i++;
continue;
} }
// Remove white space on both sides of the picture if (text[i] === ']') return i;
// const trimReg = /(!\[.*\]\(.*\))\s*/g; }
// if (trimReg.test(rawText)) {
// rawText = rawText.replace(trimReg, '$1');
// }
return rawText; return -1;
}; };
export const markdownProcess = async ({ const findMarkdownImageUrlEnd = (text: string, startIndex: number) => {
rawText, let depth = 0;
uploadImgController
}: { for (let i = startIndex; i < text.length; i++) {
rawText: string; const char = text[i];
uploadImgController?: (base64: string) => Promise<string>;
}) => { if (char === '\\') {
const imageProcess = await uploadMarkdownBase64({ i++;
rawText, continue;
uploadImgController }
});
if (char === '(') {
depth++;
continue;
}
if (char === ')') {
if (depth === 0) return i;
depth--;
}
}
return simpleMarkdownText(imageProcess); return -1;
}; };
export const matchMdImg = (text: string) => { /**
// 优化后的正则: * 扫描 markdown 图片节点，支持 URL 中包含未转义括号或转义右括号的场景。
// 1. 使用 [^\]]* 匹配 alt 文本(更精确) *
// 2. 使用 [A-Za-z0-9+/=]+ 匹配 base64 数据(避免回溯) * 普通正则 `!\[...\]\(([^)]+)\)` 会在 `https://a.com/img(1).png` 的第一个 `)` 截断，
// 3. 明确匹配 data:image/ 前缀 * 导致 http 图片转存失败；这里用轻量扫描保留完整节点范围。
const base64Regex = /!\[([^\]]*)\]\((data:image\/([^;]+);base64,([A-Za-z0-9+/=]+))\)/g; */
const imageList: ImageType[] = []; const matchMarkdownImages = (text: string) => {
const matches: MarkdownImageBase[] = [];
text = text.replace(base64Regex, (_match, altText, _fullDataUrl, mime, base64Data) => { let start = 0;
const uuid = `IMAGE_${getNanoid(12)}_IMAGE`;
while (start < text.length) {
imageList.push({ const imageStart = text.indexOf('![', start);
uuid, if (imageStart === -1) break;
base64: base64Data,
mime: `image/${mime}` const altStart = imageStart + 2;
const altEnd = findClosingBracket(text, altStart);
if (altEnd === -1 || text[altEnd + 1] !== '(') {
start = imageStart + 2;
continue;
}
const urlStart = altEnd + 2;
const urlEnd = findMarkdownImageUrlEnd(text, urlStart);
if (urlEnd === -1) {
start = imageStart + 2;
continue;
}
const fullMatch = text.slice(imageStart, urlEnd + 1);
matches.push({
altText: text.slice(altStart, altEnd),
url: text.slice(urlStart, urlEnd),
fullMatch,
index: imageStart
}); });
// 保持原有的 alt 文本，只替换 base64 部分 start = urlEnd + 1;
return `![${altText}](${uuid})`; }
return matches;
};
/**
* 处理 markdown 图片语法中的图片，并统一执行 markdown 文本清理。
*
* base64 图片默认会被解析：传入上传回调时替换成对象存储 key，不传回调或上传失败时删除，
* 避免大体积 base64 继续在解析链路中流转。http 图片默认不处理，开启后可复用同一个
* 上传回调转存；没有回调或转存失败时保留原 URL。
*/
export const parseMarkdownBase64Images = async (
text: string,
imageOptions: MarkdownImageParseOptions = {}
) => {
const {
parseBase64 = true,
parseHttp = false,
controller = imageOptions.controler
} = imageOptions;
const images = matchMarkdownImages(text).flatMap<MarkdownImage>((match) => {
const { fullMatch, altText, url: rawUrl, index } = match;
const url = unescapeMarkdownUrl(rawUrl);
const base64Match = url.match(mdBase64ImageSrcRegex);
if (parseBase64 && base64Match) {
const [, mime, base64] = base64Match;
return [
{
type: 'base64',
altText,
url,
dataUrl: url,
mime: `image/${mime}`,
base64,
fullMatch,
index
}
];
}
if (parseHttp && mdHttpImageSrcRegex.test(url)) {
return [
{
type: 'http',
altText,
url,
fullMatch,
index
}
];
}
return [];
}); });
return { if (images.length === 0) return simpleMarkdownText(text);
text,
imageList const preservedMarkdownImages = new Map<string, string>();
const preserveMarkdownImage = (image: MarkdownImage, index: number) => {
const token = `__FASTGPT_MARKDOWN_IMAGE_${index}_PLACEHOLDER__`;
preservedMarkdownImages.set(token, image.fullMatch);
return token;
}; };
const uploadResults = controller
? await batchRun(
images,
async (image, index) => {
try {
// 上传回调返回的是对象存储 key，markdown 中先保留 key，后续业务层再决定是否签名成 URL。
const { key } = await controller(image);
return key ? `![${image.altText}](${key})` : '';
} catch {
return image.type === 'http' ? preserveMarkdownImage(image, index) : '';
}
},
markdownImageUploadConcurrency
)
: images.map((image, index) =>
image.type === 'http' ? preserveMarkdownImage(image, index) : ''
);
let result = '';
let lastIndex = 0;
for (const [index, image] of images.entries()) {
result += text.slice(lastIndex, image.index);
result += uploadResults[index];
lastIndex = image.index + image.fullMatch.length;
}
const cleanedText = simpleMarkdownText(result + text.slice(lastIndex));
return Array.from(preservedMarkdownImages.entries()).reduce(
(text, [token, rawMarkdown]) => text.replaceAll(token, rawMarkdown),
cleanedText
);
}; };
packages/service/common/file/image/utils.ts
import { axios } from '../../api/axios'; import { axios } from '../../api/axios';
import { serverRequestBaseUrl } from '../../api/serverRequest'; import { serverRequestBaseUrl } from '../../api/serverRequest';
import { retryFn } from '@fastgpt/global/common/system/utils'; import { getAxiosContentType, getAxiosHeaderValue } from '@fastgpt/global/common/axios/utils';
import { getAxiosContentType } from '@fastgpt/global/common/axios/utils';
import { getLogger, LogCategories } from '../../logger'; import { getLogger, LogCategories } from '../../logger';
import { serviceEnv } from '../../../env'; import { serviceEnv } from '../../../env';
...@@ -62,6 +61,25 @@ const BASE64_PREFIX_MAP: Record<string, string> = { ...@@ -62,6 +61,25 @@ const BASE64_PREFIX_MAP: Record<string, string> = {
}; };
const DEFAULT_IMAGE_TYPE = 'image/jpeg'; const DEFAULT_IMAGE_TYPE = 'image/jpeg';
const DEFAULT_IMAGE_DOWNLOAD_TIMEOUT_MS = 180 * 1000;
const DEFAULT_IMAGE_DOWNLOAD_MAX_SIZE = 10 * 1024 * 1024;
const DEFAULT_IMAGE_BASE64_MAX_BUFFER_SIZE = DEFAULT_IMAGE_DOWNLOAD_MAX_SIZE;
export class ImageDownloadTooLargeError extends Error {
constructor(size: number, maxSize: number) {
super(`Image download too large. Size: ${size} bytes, maximum allowed: ${maxSize} bytes`);
this.name = 'ImageDownloadTooLargeError';
}
}
export class ImageBase64TooLargeError extends Error {
constructor(size: number, maxSize: number) {
super(
`Image buffer too large to convert to base64. Size: ${size} bytes, maximum allowed: ${maxSize} bytes`
);
this.name = 'ImageBase64TooLargeError';
}
}
export const isValidImageContentType = (contentType: string): boolean => { export const isValidImageContentType = (contentType: string): boolean => {
if (!contentType) return false; if (!contentType) return false;
...@@ -97,19 +115,60 @@ export const guessBase64ImageType = (str: string): string => { ...@@ -97,19 +115,60 @@ export const guessBase64ImageType = (str: string): string => {
return BASE64_PREFIX_MAP[str.charAt(0)] || DEFAULT_IMAGE_TYPE; return BASE64_PREFIX_MAP[str.charAt(0)] || DEFAULT_IMAGE_TYPE;
}; };
export const getImageBase64 = async (url: string) => { /**
logger.debug('Load image to base64', { url }); * 下载远程图片并返回 Buffer。
*
* 该函数用于文档解析链路中转存 markdown http 图片，因此默认限制为 180 秒和 10MB：
* 先用 Content-Length 快速拒绝明显超限的资源，下载过程中再按累计字节数中断，避免
* 第三方图片拖住解析任务或把大响应一次性读进内存。
*/
export const getImageBuffer = async (
url: string,
options: {
timeoutMs?: number;
maxSize?: number;
} = {}
) => {
logger.debug('Load image to buffer', { url });
try { try {
const response = await retryFn(() => const timeoutMs = options.timeoutMs ?? DEFAULT_IMAGE_DOWNLOAD_TIMEOUT_MS;
axios.get(url, { const maxSize = options.maxSize ?? DEFAULT_IMAGE_DOWNLOAD_MAX_SIZE;
const response = await axios.get(url, {
baseURL: serverRequestBaseUrl, baseURL: serverRequestBaseUrl,
responseType: 'arraybuffer' responseType: 'stream',
}) timeout: timeoutMs,
); maxContentLength: maxSize
});
const contentLength = Number(getAxiosHeaderValue(response?.headers?.['content-length']) || 0);
if (contentLength > maxSize) {
response.data?.destroy?.();
throw new ImageDownloadTooLargeError(contentLength, maxSize);
}
const chunks: Buffer[] = [];
let totalLength = 0;
const buffer = await new Promise<Buffer>((resolve, reject) => {
response.data.on('data', (chunk: Buffer) => {
totalLength += chunk.length;
if (totalLength > maxSize) {
response.data.destroy();
return reject(new ImageDownloadTooLargeError(totalLength, maxSize));
}
chunks.push(chunk);
});
response.data.on('end', () => {
resolve(Buffer.concat(chunks as unknown as Uint8Array[]));
});
response.data.on('error', reject);
});
const buffer = Buffer.from(response.data);
const base64 = buffer.toString('base64');
const headerContentType = getAxiosContentType(response?.headers?.['content-type']); const headerContentType = getAxiosContentType(response?.headers?.['content-type']);
// 检测图片类型的优先级策略 // 检测图片类型的优先级策略
...@@ -126,15 +185,49 @@ export const getImageBase64 = async (url: string) => { ...@@ -126,15 +185,49 @@ export const getImageBase64 = async (url: string) => {
} }
// 3. 回退到 base64 推断 // 3. 回退到 base64 推断
const base64 = buffer.toString('base64');
return guessBase64ImageType(base64); return guessBase64ImageType(base64);
})(); })();
return { return {
completeBase64: `data:${imageType};base64,${base64}`, buffer,
base64,
mime: imageType mime: imageType
}; };
} catch (error) { } catch (error) {
logger.warn('Load image to buffer failed', { url, error });
return Promise.reject(error);
}
};
export const getImageBase64 = async (
url: string,
options: {
timeoutMs?: number;
maxSize?: number;
maxBase64BufferSize?: number;
} = {}
) => {
logger.debug('Load image to base64', { url });
try {
const { buffer, mime } = await getImageBuffer(url, {
timeoutMs: options.timeoutMs,
maxSize: options.maxSize
});
const maxBase64BufferSize = options.maxBase64BufferSize ?? DEFAULT_IMAGE_BASE64_MAX_BUFFER_SIZE;
if (buffer.length > maxBase64BufferSize) {
throw new ImageBase64TooLargeError(buffer.length, maxBase64BufferSize);
}
const base64 = buffer.toString('base64');
return {
completeBase64: `data:${mime};base64,${base64}`,
base64,
mime
};
} catch (error) {
logger.warn('Load image to base64 failed', { url, error }); logger.warn('Load image to base64 failed', { url, error });
return Promise.reject(error); return Promise.reject(error);
} }
......
packages/service/common/file/read/image.ts 0 → 100644
import { uploadImage2S3Bucket } from '../../s3/utils';
import { normalizeMimeType, resolveMimeExtension, resolveMimeType } from '../../s3/utils/mime';
export type ParsedPdfImageUploadParams =
| {
type: 'base64';
mime: string;
dataUrl: string;
}
| {
type: 'http';
mime: string;
buffer: Buffer;
};
export type ParsedPdfImageKeyOptions = {
prefix: string;
expiredTime?: Date;
};
/**
* 将 PDF 解析过程中得到的图片统一写入文件解析图片目录。
*
* base64 图片直接使用 dataUrl 上传；http 图片由上游先下载成 buffer 后上传。返回值保持
* worker/provider 图片处理回调的统一契约，便于 markdown 中直接替换成对象存储 key。
*/
export const uploadParsedPdfImage = async (
image: ParsedPdfImageUploadParams,
imageKeyOptions?: ParsedPdfImageKeyOptions
) => {
if (!imageKeyOptions?.prefix) return { key: '' };
const { prefix, expiredTime } = imageKeyOptions;
const mimetype = normalizeMimeType(image.mime);
const ext = resolveMimeExtension(mimetype);
const filename = `${crypto.randomUUID()}${ext}`;
const commonParams = {
uploadKey: `${prefix}/${filename}`,
mimetype: resolveMimeType([filename], mimetype),
filename,
expiredTime
};
const key = await uploadImage2S3Bucket(
'private',
image.type === 'base64'
? {
...commonParams,
base64Img: image.dataUrl
}
: {
...commonParams,
buffer: image.buffer
}
);
return { key };
};
packages/service/common/file/read/utils.ts
...@@ -2,15 +2,14 @@ import FormData from 'form-data'; ...@@ -2,15 +2,14 @@ import FormData from 'form-data';
import fs from 'fs'; import fs from 'fs';
import type { ReadFileResponse } from '../../../worker/readFile/type'; import type { ReadFileResponse } from '../../../worker/readFile/type';
import { axios } from '../../api/axios'; import { axios } from '../../api/axios';
import { batchRun } from '@fastgpt/global/common/system/utils'; import { parseMarkdownBase64Images } from '@fastgpt/global/common/string/markdown';
import { matchMdImg } from '@fastgpt/global/common/string/markdown';
import { createPdfParseUsage } from '../../../support/wallet/usage/controller'; import { createPdfParseUsage } from '../../../support/wallet/usage/controller';
import { useDoc2xServer } from '../../../thirdProvider/doc2x'; import { useDoc2xServer } from '../../../thirdProvider/doc2x';
import { useTextinServer } from '../../../thirdProvider/textin'; import { useTextinServer } from '../../../thirdProvider/textin';
import { readRawContentFromBuffer } from '../../../worker/function'; import { readRawContentFromBuffer } from '../../../worker/function';
import { uploadImage2S3Bucket } from '../../s3/utils';
import { normalizeMimeType, resolveMimeExtension, resolveMimeType } from '../../s3/utils/mime';
import { getLogger, LogCategories } from '../../logger'; import { getLogger, LogCategories } from '../../logger';
import { getImageBuffer } from '../image/utils';
import { uploadParsedPdfImage } from './image';
const logger = getLogger(LogCategories.MODULE.DATASET.FILE); const logger = getLogger(LogCategories.MODULE.DATASET.FILE);
...@@ -75,11 +74,42 @@ export const readFileContentByBuffer = async ({ ...@@ -75,11 +74,42 @@ export const readFileContentByBuffer = async ({
}): Promise<{ }): Promise<{
rawText: string; rawText: string;
}> => { }> => {
const parseMarkdownImages = (rawText: string) =>
parseMarkdownBase64Images(rawText, {
parseBase64: true,
parseHttp: true,
controller: imageKeyOptions?.prefix
? async (image) => {
if (image.type === 'base64') {
return uploadParsedPdfImage(
{
type: 'base64',
mime: image.mime,
dataUrl: image.dataUrl
},
imageKeyOptions
);
}
const { buffer, mime } = await getImageBuffer(image.url);
return uploadParsedPdfImage(
{
type: 'http',
mime,
buffer
},
imageKeyOptions
);
}
: undefined
});
const systemParse = () => const systemParse = () =>
readRawContentFromBuffer({ readRawContentFromBuffer({
extension, extension,
encoding, encoding,
buffer buffer,
imageKeyOptions
}); });
const parsePdfFromCustomService = async (): Promise<ReadFileResponse> => { const parsePdfFromCustomService = async (): Promise<ReadFileResponse> => {
const url = global.systemEnv.customPdfParse?.url; const url = global.systemEnv.customPdfParse?.url;
...@@ -114,8 +144,7 @@ export const readFileContentByBuffer = async ({ ...@@ -114,8 +144,7 @@ export const readFileContentByBuffer = async ({
durationMs: Date.now() - start durationMs: Date.now() - start
}); });
const rawText = response.markdown; const text = await parseMarkdownImages(response.markdown);
const { text, imageList } = matchMdImg(rawText);
createPdfParseUsage({ createPdfParseUsage({
teamId, teamId,
...@@ -126,8 +155,7 @@ export const readFileContentByBuffer = async ({ ...@@ -126,8 +155,7 @@ export const readFileContentByBuffer = async ({
return { return {
rawText: text, rawText: text,
formatText: text, formatText: text
imageList
}; };
}; };
// Textin api // Textin api
...@@ -136,10 +164,28 @@ export const readFileContentByBuffer = async ({ ...@@ -136,10 +164,28 @@ export const readFileContentByBuffer = async ({
const secretCode = global.systemEnv.customPdfParse?.textinSecretCode; const secretCode = global.systemEnv.customPdfParse?.textinSecretCode;
if (!appId || !secretCode) return systemParse(); if (!appId || !secretCode) return systemParse();
const { pages, text, imageList } = await useTextinServer({ const { pages, text } = await useTextinServer({
appId, appId,
secretCode secretCode
}).parsePDF(buffer); }).parsePDF(buffer, {
uploadImage: imageKeyOptions?.prefix
? async (image) =>
uploadParsedPdfImage(
image.type === 'base64'
? {
type: 'base64',
mime: image.mime,
dataUrl: image.dataUrl
}
: {
type: 'http',
mime: image.mime,
buffer: image.buffer
},
imageKeyOptions
)
: undefined
});
createPdfParseUsage({ createPdfParseUsage({
teamId, teamId,
...@@ -150,8 +196,7 @@ export const readFileContentByBuffer = async ({ ...@@ -150,8 +196,7 @@ export const readFileContentByBuffer = async ({
return { return {
rawText: text, rawText: text,
formatText: text, formatText: text
imageList
}; };
}; };
// Doc2x api // Doc2x api
...@@ -159,7 +204,25 @@ export const readFileContentByBuffer = async ({ ...@@ -159,7 +204,25 @@ export const readFileContentByBuffer = async ({
const doc2xKey = global.systemEnv.customPdfParse?.doc2xKey; const doc2xKey = global.systemEnv.customPdfParse?.doc2xKey;
if (!doc2xKey) return systemParse(); if (!doc2xKey) return systemParse();
const { pages, text, imageList } = await useDoc2xServer({ apiKey: doc2xKey }).parsePDF(buffer); const { pages, text } = await useDoc2xServer({ apiKey: doc2xKey }).parsePDF(buffer, {
uploadImage: imageKeyOptions?.prefix
? async (image) =>
uploadParsedPdfImage(
image.type === 'base64'
? {
type: 'base64',
mime: image.mime,
dataUrl: image.dataUrl
}
: {
type: 'http',
mime: image.mime,
buffer: image.buffer
},
imageKeyOptions
)
: undefined
});
createPdfParseUsage({ createPdfParseUsage({
teamId, teamId,
...@@ -170,8 +233,7 @@ export const readFileContentByBuffer = async ({ ...@@ -170,8 +233,7 @@ export const readFileContentByBuffer = async ({
return { return {
rawText: text, rawText: text,
formatText: text, formatText: text
imageList
}; };
}; };
// Custom read file service // Custom read file service
...@@ -187,56 +249,15 @@ export const readFileContentByBuffer = async ({ ...@@ -187,56 +249,15 @@ export const readFileContentByBuffer = async ({
const start = Date.now(); const start = Date.now();
logger.debug('Start parsing file', { extension }); logger.debug('Start parsing file', { extension });
const parsedFile = await (async () => { const { rawText, formatText } = await (async () => {
if (extension === 'pdf') { if (extension === 'pdf') {
return await pdfParseFn(); return await pdfParseFn();
} }
return await systemParse(); return await systemParse();
})(); })();
const { imageList } = parsedFile;
let { rawText, formatText } = parsedFile;
logger.debug('File parsing completed', { extension, durationMs: Date.now() - start }); logger.debug('File parsing completed', { extension, durationMs: Date.now() - start });
// markdown data format
if (imageList && imageList.length > 0) {
logger.debug('Processing parsed document images', {
extension,
imageCount: imageList.length
});
await batchRun(imageList, async (item) => {
const src = await (async () => {
if (!imageKeyOptions) return '';
try {
const { prefix, expiredTime } = imageKeyOptions;
const mimetype = normalizeMimeType(item.mime);
const ext = resolveMimeExtension(mimetype);
const filename = `${item.uuid}${ext}`;
return await uploadImage2S3Bucket('private', {
base64Img: `data:${mimetype};base64,${item.base64}`,
uploadKey: `${prefix}/${filename}`,
mimetype: resolveMimeType([filename], mimetype),
filename,
expiredTime
});
} catch (error) {
logger.warn('Failed to upload parsed image to S3', {
extension,
imageUuid: item.uuid,
error
});
return `[Image Upload Failed: ${item.uuid}]`;
}
})();
rawText = rawText.replace(item.uuid, src);
if (formatText) {
formatText = formatText.replace(item.uuid, src);
}
});
}
return { return {
rawText: getFormatText ? formatText || rawText : rawText rawText: getFormatText ? formatText || rawText : rawText
}; };
......
packages/service/common/s3/contracts/type.ts
...@@ -70,13 +70,18 @@ export const CreateGetPresignedUrlParamsSchema = z.object({ ...@@ -70,13 +70,18 @@ export const CreateGetPresignedUrlParamsSchema = z.object({
}); });
export type createPreviewUrlParams = z.infer<typeof CreateGetPresignedUrlParamsSchema>; export type createPreviewUrlParams = z.infer<typeof CreateGetPresignedUrlParamsSchema>;
export const UploadImage2S3BucketParamsSchema = z.object({ export const UploadImage2S3BucketParamsSchema = z
base64Img: z.string().nonempty(), .object({
base64Img: z.string().nonempty().optional(),
buffer: z.instanceof(Buffer).optional(),
uploadKey: z.string().nonempty(), uploadKey: z.string().nonempty(),
mimetype: z.string().nonempty(), mimetype: z.string().nonempty(),
filename: z.string().nonempty(), filename: z.string().nonempty(),
expiredTime: z.coerce.date().optional() expiredTime: z.coerce.date().optional()
}); })
.refine((value) => value.base64Img || value.buffer, {
message: 'base64Img or buffer is required'
});
export type UploadImage2S3BucketParams = z.infer<typeof UploadImage2S3BucketParamsSchema>; export type UploadImage2S3BucketParams = z.infer<typeof UploadImage2S3BucketParamsSchema>;
export const UploadFileByBodySchema = z.object({ export const UploadFileByBodySchema = z.object({
......
packages/service/common/s3/utils.ts
...@@ -87,12 +87,18 @@ export async function uploadImage2S3Bucket( ...@@ -87,12 +87,18 @@ export async function uploadImage2S3Bucket(
bucketName: keyof typeof S3Buckets, bucketName: keyof typeof S3Buckets,
params: UploadImage2S3BucketParams params: UploadImage2S3BucketParams
) { ) {
const { base64Img, filename, mimetype, uploadKey, expiredTime } = params; const { base64Img, buffer: inputBuffer, filename, mimetype, uploadKey, expiredTime } = params;
const bucket = bucketName === 'private' ? new S3PrivateBucket() : new S3PublicBucket(); const bucket = bucketName === 'private' ? new S3PrivateBucket() : new S3PublicBucket();
const base64Data = base64Img.split(',')[1] || base64Img; const buffer = (() => {
const buffer = Buffer.from(base64Data, 'base64'); if (inputBuffer) return inputBuffer;
const base64Data = base64Img?.split(',')[1] || base64Img;
if (!base64Data) {
throw new Error('base64Img or buffer is required');
}
return Buffer.from(base64Data, 'base64');
})();
await bucket.client.uploadObject({ await bucket.client.uploadObject({
key: uploadKey, key: uploadKey,
......
packages/service/common/string/utils.ts
import { serviceEnv } from '../../env'; import { serviceEnv } from '../../env';
import { WorkerNameEnum, getWorkerController } from '../../worker/utils'; import { WorkerNameEnum, getWorkerController } from '../../worker/utils';
import { type ImageType } from '../../worker/readFile/type';
import { getLogger, LogCategories } from '../logger'; import { getLogger, LogCategories } from '../logger';
const logger = getLogger(LogCategories.INFRA.WORKER); const logger = getLogger(LogCategories.INFRA.WORKER);
...@@ -16,7 +15,6 @@ export const htmlToMarkdown = async (html?: string | null) => { ...@@ -16,7 +15,6 @@ export const htmlToMarkdown = async (html?: string | null) => {
{ html: string }, { html: string },
{ {
rawText: string; rawText: string;
imageList: ImageType[];
} }
>({ >({
name: WorkerNameEnum.htmlStr2Md, name: WorkerNameEnum.htmlStr2Md,
......
packages/service/core/workflow/dispatch/ai/agent/sub/file/index.ts
import { isInternalAddress, PRIVATE_URL_TEXT } from '../../../../../../../common/system/utils';
import { pickOutboundAxios } from '../../../../../../../common/api/axios';
import { parseFileExtensionFromUrl } from '@fastgpt/global/common/string/tools';
import {
detectFileEncoding,
parseContentDispositionFilename
} from '@fastgpt/global/common/file/tools';
import { getErrText } from '@fastgpt/global/common/error/utils'; import { getErrText } from '@fastgpt/global/common/error/utils';
import { getS3RawTextSource } from '../../../../../../../common/s3/sources/rawText/index';
import { readFileContentByBuffer } from '../../../../../../../common/file/read/utils';
import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant'; import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant';
import { i18nT } from '@fastgpt/global/common/i18n/utils'; import { i18nT } from '@fastgpt/global/common/i18n/utils';
import { getAxiosHeaderValue } from '@fastgpt/global/common/axios/utils';
import type { DispatchSubAppResponse } from '../../type'; import type { DispatchSubAppResponse } from '../../type';
import { getFileContentByUrl } from '../../../../../utils/file';
type FileReadParams = { type FileReadParams = {
files: { id: string; url: string }[]; files: { id: string; url: string }[];
...@@ -30,78 +21,18 @@ export const dispatchFileRead = async ({ ...@@ -30,78 +21,18 @@ export const dispatchFileRead = async ({
try { try {
const readFilesResult = await Promise.all( const readFilesResult = await Promise.all(
files.map(async ({ id, url }) => { files.map(async ({ id, url }) => {
// Get from buffer
const fileBuffer = await getS3RawTextSource().getRawTextBuffer({
sourceId: url,
customPdfParse
});
if (fileBuffer) {
return {
id,
name: fileBuffer.filename,
content: fileBuffer.text
};
}
try { try {
if (await isInternalAddress(url)) { const { name, content } = await getFileContentByUrl({
return { url,
id,
name: '',
content: PRIVATE_URL_TEXT
};
}
const response = await pickOutboundAxios(url).get(url, {
responseType: 'arraybuffer'
});
const buffer = Buffer.from(response.data, 'binary');
// Get file name
const filename = (() => {
const contentDisposition = getAxiosHeaderValue(response.headers['content-disposition']);
return parseContentDispositionFilename(contentDisposition) || url;
})();
// Extension
const extension = parseFileExtensionFromUrl(filename);
// Get encoding
const encoding = (() => {
const contentType = getAxiosHeaderValue(response.headers['content-type']);
if (contentType) {
const charsetRegex = /charset=([^;]*)/;
const matches = charsetRegex.exec(contentType);
if (matches != null && matches[1]) {
return matches[1];
}
}
return detectFileEncoding(buffer);
})();
// Read file
const { rawText } = await readFileContentByBuffer({
extension,
teamId, teamId,
tmbId, tmbId,
buffer,
encoding,
customPdfParse,
getFormatText: true
});
// Add to buffer
getS3RawTextSource().addRawTextBuffer({
sourceId: url,
sourceName: filename,
text: rawText,
customPdfParse customPdfParse
}); });
return { return {
id, id,
name: filename, name,
content: rawText content
}; };
} catch (error) { } catch (error) {
return { return {
......
packages/service/test/common/file/image/utils.test.ts
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { Readable } from 'node:stream';
import { import {
isValidImageContentType, isValidImageContentType,
detectImageTypeFromBuffer, detectImageTypeFromBuffer,
guessBase64ImageType, guessBase64ImageType,
getImageBuffer,
getImageBase64 getImageBase64
} from '@fastgpt/service/common/file/image/utils'; } from '@fastgpt/service/common/file/image/utils';
...@@ -23,6 +25,11 @@ const loadUtilsModule = async () => { ...@@ -23,6 +25,11 @@ const loadUtilsModule = async () => {
return import('@fastgpt/service/common/file/image/utils'); return import('@fastgpt/service/common/file/image/utils');
}; };
const mockImageResponse = (buffer: Buffer, headers: Record<string, string> = {}) => ({
data: Readable.from([buffer]),
headers
});
describe('isValidImageContentType', () => { describe('isValidImageContentType', () => {
it('should return true for valid image MIME types', () => { it('should return true for valid image MIME types', () => {
expect(isValidImageContentType('image/jpeg')).toBe(true); expect(isValidImageContentType('image/jpeg')).toBe(true);
...@@ -211,18 +218,29 @@ describe('getImageBase64', () => { ...@@ -211,18 +218,29 @@ describe('getImageBase64', () => {
mockAxiosGet.mockReset(); mockAxiosGet.mockReset();
}); });
it('should return base64 with correct mime when header has valid image content-type', async () => { it('should return image buffer without base64 encoding', async () => {
const pngBytes = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); const pngBytes = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(mockImageResponse(pngBytes, { 'content-type': 'image/png' }));
data: pngBytes,
headers: { 'content-type': 'image/png' } const result = await getImageBuffer('/api/system/img/test.png');
expect(result).toEqual({
buffer: pngBytes,
mime: 'image/png'
}); });
});
it('should return base64 with correct mime when header has valid image content-type', async () => {
const pngBytes = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
mockAxiosGet.mockResolvedValue(mockImageResponse(pngBytes, { 'content-type': 'image/png' }));
const result = await getImageBase64('/api/system/img/test.png'); const result = await getImageBase64('/api/system/img/test.png');
expect(mockAxiosGet).toHaveBeenCalledWith('/api/system/img/test.png', { expect(mockAxiosGet).toHaveBeenCalledWith('/api/system/img/test.png', {
baseURL: 'http://localhost:3000', baseURL: 'http://localhost:3000',
responseType: 'arraybuffer' responseType: 'stream',
timeout: 180000,
maxContentLength: 10 * 1024 * 1024
}); });
expect(result.mime).toBe('image/png'); expect(result.mime).toBe('image/png');
expect(result.base64).toBe(pngBytes.toString('base64')); expect(result.base64).toBe(pngBytes.toString('base64'));
...@@ -231,10 +249,9 @@ describe('getImageBase64', () => { ...@@ -231,10 +249,9 @@ describe('getImageBase64', () => {
it('should detect type from buffer when header content-type is not a valid image type', async () => { it('should detect type from buffer when header content-type is not a valid image type', async () => {
const jpegBytes = Buffer.from([0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10]); const jpegBytes = Buffer.from([0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(
data: jpegBytes, mockImageResponse(jpegBytes, { 'content-type': 'application/octet-stream' })
headers: { 'content-type': 'application/octet-stream' } );
});
const result = await getImageBase64('/img/photo.jpg'); const result = await getImageBase64('/img/photo.jpg');
...@@ -244,10 +261,7 @@ describe('getImageBase64', () => { ...@@ -244,10 +261,7 @@ describe('getImageBase64', () => {
it('should detect type from buffer when header content-type is missing', async () => { it('should detect type from buffer when header content-type is missing', async () => {
const gifBytes = Buffer.from([0x47, 0x49, 0x46, 0x38, 0x39, 0x61]); const gifBytes = Buffer.from([0x47, 0x49, 0x46, 0x38, 0x39, 0x61]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(mockImageResponse(gifBytes));
data: gifBytes,
headers: {}
});
const result = await getImageBase64('/img/anim.gif'); const result = await getImageBase64('/img/anim.gif');
...@@ -257,10 +271,9 @@ describe('getImageBase64', () => { ...@@ -257,10 +271,9 @@ describe('getImageBase64', () => {
it('should fallback to guessBase64ImageType when buffer detection fails', async () => { it('should fallback to guessBase64ImageType when buffer detection fails', async () => {
// Unknown magic bytes that do not match any signature // Unknown magic bytes that do not match any signature
const unknownBytes = Buffer.from([0xaa, 0xbb, 0xcc, 0xdd]); const unknownBytes = Buffer.from([0xaa, 0xbb, 0xcc, 0xdd]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(
data: unknownBytes, mockImageResponse(unknownBytes, { 'content-type': 'text/html' })
headers: { 'content-type': 'text/html' } );
});
const result = await getImageBase64('/img/unknown.dat'); const result = await getImageBase64('/img/unknown.dat');
...@@ -273,10 +286,9 @@ describe('getImageBase64', () => { ...@@ -273,10 +286,9 @@ describe('getImageBase64', () => {
it('should handle content-type header with charset parameter', async () => { it('should handle content-type header with charset parameter', async () => {
const svgBytes = Buffer.from([0x3c, 0x73, 0x76, 0x67, 0x20, 0x78, 0x6d, 0x6c]); const svgBytes = Buffer.from([0x3c, 0x73, 0x76, 0x67, 0x20, 0x78, 0x6d, 0x6c]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(
data: svgBytes, mockImageResponse(svgBytes, { 'content-type': 'image/svg+xml; charset=utf-8' })
headers: { 'content-type': 'image/svg+xml; charset=utf-8' } );
});
const result = await getImageBase64('/img/icon.svg'); const result = await getImageBase64('/img/icon.svg');
...@@ -300,10 +312,7 @@ describe('getImageBase64', () => { ...@@ -300,10 +312,7 @@ describe('getImageBase64', () => {
it('should handle empty arraybuffer response', async () => { it('should handle empty arraybuffer response', async () => {
const emptyBuffer = Buffer.from([]); const emptyBuffer = Buffer.from([]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(mockImageResponse(emptyBuffer, { 'content-type': 'image/png' }));
data: emptyBuffer,
headers: { 'content-type': 'image/png' }
});
const result = await getImageBase64('/img/empty.png'); const result = await getImageBase64('/img/empty.png');
...@@ -315,10 +324,7 @@ describe('getImageBase64', () => { ...@@ -315,10 +324,7 @@ describe('getImageBase64', () => {
it('should prefer header content-type over buffer detection when header is valid', async () => { it('should prefer header content-type over buffer detection when header is valid', async () => {
// JPEG magic bytes but header says image/webp // JPEG magic bytes but header says image/webp
const jpegBytes = Buffer.from([0xff, 0xd8, 0xff, 0xe0]); const jpegBytes = Buffer.from([0xff, 0xd8, 0xff, 0xe0]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(mockImageResponse(jpegBytes, { 'content-type': 'image/webp' }));
data: jpegBytes,
headers: { 'content-type': 'image/webp' }
});
const result = await getImageBase64('/img/test.webp'); const result = await getImageBase64('/img/test.webp');
...@@ -328,15 +334,53 @@ describe('getImageBase64', () => { ...@@ -328,15 +334,53 @@ describe('getImageBase64', () => {
it('should construct completeBase64 in correct data URI format', async () => { it('should construct completeBase64 in correct data URI format', async () => {
const bmpBytes = Buffer.from([0x42, 0x4d, 0x00, 0x00, 0x00, 0x00]); const bmpBytes = Buffer.from([0x42, 0x4d, 0x00, 0x00, 0x00, 0x00]);
mockAxiosGet.mockResolvedValue({ mockAxiosGet.mockResolvedValue(mockImageResponse(bmpBytes, { 'content-type': 'image/bmp' }));
data: bmpBytes,
headers: { 'content-type': 'image/bmp' }
});
const result = await getImageBase64('/img/test.bmp'); const result = await getImageBase64('/img/test.bmp');
expect(result.completeBase64).toMatch(/^data:image\/bmp;base64,[A-Za-z0-9+/=]+$/); expect(result.completeBase64).toMatch(/^data:image\/bmp;base64,[A-Za-z0-9+/=]+$/);
}); });
it('should reject before reading when content-length exceeds maxSize', async () => {
const pngBytes = Buffer.from([0x89, 0x50, 0x4e, 0x47]);
mockAxiosGet.mockResolvedValue(
mockImageResponse(pngBytes, {
'content-type': 'image/png',
'content-length': '11'
})
);
await expect(getImageBuffer('/img/too-large.png', { maxSize: 10 })).rejects.toThrow(
'Image download too large'
);
});
it('should reject while streaming when image exceeds maxSize', async () => {
mockAxiosGet.mockResolvedValue(
mockImageResponse(Buffer.from('12345678901'), {
'content-type': 'image/png'
})
);
await expect(getImageBuffer('/img/stream-too-large.png', { maxSize: 10 })).rejects.toThrow(
'Image download too large'
);
});
it('should reject before converting to base64 when buffer exceeds base64 limit', async () => {
mockAxiosGet.mockResolvedValue(
mockImageResponse(Buffer.from('12345678901'), {
'content-type': 'image/png'
})
);
await expect(
getImageBase64('/img/base64-too-large.png', {
maxSize: 20,
maxBase64BufferSize: 10
})
).rejects.toThrow('Image buffer too large to convert to base64');
});
}); });
describe('addEndpointToImageUrl', () => { describe('addEndpointToImageUrl', () => {
......
packages/service/test/core/workflow/dispatch/ai/agent/sub/file.test.ts
import { beforeEach, describe, expect, it, vi } from 'vitest'; import { beforeEach, describe, expect, it, vi } from 'vitest';
import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant'; import { FlowNodeTypeEnum } from '@fastgpt/global/core/workflow/node/constant';
const { getRawTextBufferMock } = vi.hoisted(() => ({ const { getFileContentByUrlMock } = vi.hoisted(() => ({
getRawTextBufferMock: vi.fn() getFileContentByUrlMock: vi.fn()
})); }));
vi.mock('@fastgpt/service/common/s3/sources/rawText/index', () => ({ vi.mock('@fastgpt/service/core/workflow/utils/file', async (importOriginal) => {
getS3RawTextSource: () => ({ const mod = await importOriginal<typeof import('@fastgpt/service/core/workflow/utils/file')>();
getRawTextBuffer: getRawTextBufferMock, return {
addRawTextBuffer: vi.fn() ...mod,
}) getFileContentByUrl: getFileContentByUrlMock
})); };
});
import { dispatchFileRead } from '@fastgpt/service/core/workflow/dispatch/ai/agent/sub/file'; import { dispatchFileRead } from '@fastgpt/service/core/workflow/dispatch/ai/agent/sub/file';
describe('dispatchFileRead', () => { describe('dispatchFileRead', () => {
beforeEach(() => { beforeEach(() => {
vi.clearAllMocks(); vi.clearAllMocks();
getFileContentByUrlMock.mockResolvedValue({
name: 'doc.txt',
content: 'large file content'
}); });
it('returns raw file content and leaves compression to tool response compression', async () => {
getRawTextBufferMock.mockResolvedValue({
filename: 'doc.txt',
text: 'large file content'
}); });
it('复用 readFiles 的文件读取逻辑，并保留 agent 返回格式', async () => {
const result = await dispatchFileRead({ const result = await dispatchFileRead({
files: [ files: [
{ {
id: 'file_0', id: 'file_0',
url: 'file_raw_text_id' url: 'https://example.com/doc.txt'
} }
], ],
teamId: 'team_1', teamId: 'team_1',
tmbId: 'tmb_1' tmbId: 'tmb_1',
customPdfParse: true
}); });
expect(getFileContentByUrlMock).toHaveBeenCalledWith({
url: 'https://example.com/doc.txt',
teamId: 'team_1',
tmbId: 'tmb_1',
customPdfParse: true
});
expect(result.response).toBe( expect(result.response).toBe(
JSON.stringify([ JSON.stringify([
{ {
...@@ -51,4 +58,29 @@ describe('dispatchFileRead', () => { ...@@ -51,4 +58,29 @@ describe('dispatchFileRead', () => {
moduleName: 'chat:read_file' moduleName: 'chat:read_file'
}); });
}); });
it('读取失败时返回对应文件的错误内容', async () => {
getFileContentByUrlMock.mockRejectedValue(new Error('download failed'));
const result = await dispatchFileRead({
files: [
{
id: 'file_1',
url: 'https://example.com/error.docx'
}
],
teamId: 'team_1',
tmbId: 'tmb_1'
});
expect(result.response).toBe(
JSON.stringify([
{
id: 'file_1',
name: '',
content: 'download failed'
}
])
);
});
}); });
packages/service/test/thirdProvider/doc2x.test.ts 0 → 100644
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
const { requestMock, axiosPutMock, getImageBufferMock } = vi.hoisted(() => ({
requestMock: vi.fn(),
axiosPutMock: vi.fn(),
getImageBufferMock: vi.fn()
}));
vi.mock('@fastgpt/service/common/api/axios', () => ({
axios: {
put: axiosPutMock
},
createProxyAxios: vi.fn(() => ({
request: requestMock
}))
}));
vi.mock('@fastgpt/service/common/file/image/utils', () => ({
getImageBuffer: getImageBufferMock
}));
const { useDoc2xServer } = await import('@fastgpt/service/thirdProvider/doc2x');
const mockDoc2xSuccess = (md: string) => {
requestMock
.mockResolvedValueOnce({
data: {
code: 'ok',
data: {
uid: 'uid-1',
url: 'https://upload.example.com/file'
}
}
})
.mockResolvedValueOnce({
data: {
code: 'ok',
data: {
status: 'success',
result: {
pages: [
{
md
}
]
}
}
}
});
};
describe('useDoc2xServer', () => {
beforeEach(() => {
vi.clearAllMocks();
vi.useFakeTimers();
axiosPutMock.mockResolvedValue({
status: 200,
statusText: 'OK'
});
mockDoc2xSuccess('hello ![](https://img.example.com/a.png)');
getImageBufferMock.mockResolvedValue({
buffer: Buffer.from('image-bytes'),
mime: 'image/png'
});
});
afterEach(() => {
vi.useRealTimers();
});
it('转存 Doc2x 图片 URL 到 S3 key，不再返回 imageList', async () => {
const uploadImage = vi.fn().mockResolvedValue({ key: 'dataset/ds1/file-parsed/image.png' });
const resultPromise = useDoc2xServer({ apiKey: 'api-key' }).parsePDF(Buffer.from('pdf'), {
uploadImage
});
await vi.runAllTimersAsync();
const result = await resultPromise;
expect(getImageBufferMock).toHaveBeenCalledWith('https://img.example.com/a.png');
expect(uploadImage).toHaveBeenCalledWith({
type: 'http',
url: 'https://img.example.com/a.png',
mime: 'image/png',
buffer: Buffer.from('image-bytes')
});
expect(result).toEqual({
pages: 1,
text: 'hello ![](dataset/ds1/file-parsed/image.png)'
});
});
it('按匹配顺序逐张转存 Doc2x 图片，不预先收集 imageList', async () => {
requestMock.mockReset();
mockDoc2xSuccess('a ![](https://img.example.com/a.png) b ![](https://img.example.com/b.png)');
const uploadImage = vi
.fn()
.mockResolvedValueOnce({ key: 'dataset/ds1/file-parsed/a.png' })
.mockResolvedValueOnce({ key: 'dataset/ds1/file-parsed/b.png' });
const resultPromise = useDoc2xServer({ apiKey: 'api-key' }).parsePDF(Buffer.from('pdf'), {
uploadImage
});
await vi.runAllTimersAsync();
const result = await resultPromise;
expect(getImageBufferMock).toHaveBeenNthCalledWith(1, 'https://img.example.com/a.png');
expect(getImageBufferMock).toHaveBeenNthCalledWith(2, 'https://img.example.com/b.png');
expect(uploadImage).toHaveBeenCalledTimes(2);
expect(result.text).toBe(
'a ![](dataset/ds1/file-parsed/a.png) b ![](dataset/ds1/file-parsed/b.png)'
);
});
it('兜底处理 Doc2x markdown base64 图片并替换成上传返回 key', async () => {
requestMock.mockReset();
mockDoc2xSuccess('hello ![img](data:image/png;base64,iVBORw0KGgo=)');
const uploadImage = vi.fn().mockResolvedValue({ key: 'dataset/ds1/file-parsed/base64.png' });
const resultPromise = useDoc2xServer({ apiKey: 'api-key' }).parsePDF(Buffer.from('pdf'), {
uploadImage
});
await vi.runAllTimersAsync();
const result = await resultPromise;
expect(getImageBufferMock).not.toHaveBeenCalled();
expect(uploadImage).toHaveBeenCalledWith({
type: 'base64',
mime: 'image/png',
base64: 'iVBORw0KGgo=',
dataUrl: 'data:image/png;base64,iVBORw0KGgo='
});
expect(result.text).toBe('hello ![img](dataset/ds1/file-parsed/base64.png)');
});
it('未传 uploadImage 时删除 Doc2x markdown base64 图片', async () => {
requestMock.mockReset();
mockDoc2xSuccess('hello ![img](data:image/png;base64,iVBORw0KGgo=)');
const resultPromise = useDoc2xServer({ apiKey: 'api-key' }).parsePDF(Buffer.from('pdf'));
await vi.runAllTimersAsync();
const result = await resultPromise;
expect(getImageBufferMock).not.toHaveBeenCalled();
expect(result.text).toBe('hello');
});
it('未传 uploadImage 时保留 Doc2x 图片 URL 且不下载图片', async () => {
const resultPromise = useDoc2xServer({ apiKey: 'api-key' }).parsePDF(Buffer.from('pdf'));
await vi.runAllTimersAsync();
const result = await resultPromise;
expect(getImageBufferMock).not.toHaveBeenCalled();
expect(result.text).toBe('hello ![](https://img.example.com/a.png)');
});
});
packages/service/test/thirdProvider/textin.test.ts 0 → 100644
import { beforeEach, describe, expect, it, vi } from 'vitest';
const { postMock, getImageBufferMock } = vi.hoisted(() => ({
postMock: vi.fn(),
getImageBufferMock: vi.fn()
}));
vi.mock('@fastgpt/service/common/api/axios', () => ({
createProxyAxios: vi.fn(() => ({
post: postMock
}))
}));
vi.mock('@fastgpt/service/common/file/image/utils', () => ({
getImageBuffer: getImageBufferMock
}));
const { useTextinServer } = await import('@fastgpt/service/thirdProvider/textin');
const mockTextinSuccess = (markdown: string) => {
postMock.mockResolvedValueOnce({
data: {
code: 200,
result: {
markdown,
total_page_number: 1
}
}
});
};
describe('useTextinServer', () => {
beforeEach(() => {
vi.clearAllMocks();
getImageBufferMock.mockResolvedValue({
buffer: Buffer.from('image-bytes'),
mime: 'image/png'
});
});
it('解析 Textin markdown base64 图片并替换成上传返回 key', async () => {
const markdown = 'hello ![img](data:image/png;base64,iVBORw0KGgo=)';
mockTextinSuccess(markdown);
const uploadImage = vi.fn().mockResolvedValue({ key: 'dataset/ds1/file-parsed/image.png' });
const result = await useTextinServer({
appId: 'app-id',
secretCode: 'secret-code'
}).parsePDF(Buffer.from('pdf'), {
uploadImage
});
expect(postMock).toHaveBeenCalledWith(
'/pdf_to_markdown',
expect.any(Buffer),
expect.objectContaining({
params: expect.objectContaining({
get_image: 'objects',
image_output_type: 'default'
})
})
);
expect(uploadImage).toHaveBeenCalledWith({
type: 'base64',
mime: 'image/png',
base64: 'iVBORw0KGgo=',
dataUrl: 'data:image/png;base64,iVBORw0KGgo='
});
expect(result).toEqual({
pages: 1,
text: 'hello ![img](dataset/ds1/file-parsed/image.png)'
});
});
it('解析 Textin markdown http 图片并替换成上传返回 key', async () => {
mockTextinSuccess('hello ![img](https://textin.example.com/image.png)');
const uploadImage = vi.fn().mockResolvedValue({ key: 'dataset/ds1/file-parsed/http.png' });
const result = await useTextinServer({
appId: 'app-id',
secretCode: 'secret-code'
}).parsePDF(Buffer.from('pdf'), {
uploadImage
});
expect(getImageBufferMock).toHaveBeenCalledWith('https://textin.example.com/image.png');
expect(uploadImage).toHaveBeenCalledWith({
type: 'http',
url: 'https://textin.example.com/image.png',
mime: 'image/png',
buffer: Buffer.from('image-bytes')
});
expect(result).toEqual({
pages: 1,
text: 'hello ![img](dataset/ds1/file-parsed/http.png)'
});
});
it('未传 uploadImage 时删除 Textin markdown base64 图片', async () => {
mockTextinSuccess('hello ![img](data:image/png;base64,iVBORw0KGgo=)');
const result = await useTextinServer({
appId: 'app-id',
secretCode: 'secret-code'
}).parsePDF(Buffer.from('pdf'));
expect(result).toEqual({
pages: 1,
text: 'hello'
});
});
it('未传 uploadImage 时保留 Textin markdown http 图片且不下载', async () => {
mockTextinSuccess('hello ![img](https://textin.example.com/image.png)');
const result = await useTextinServer({
appId: 'app-id',
secretCode: 'secret-code'
}).parsePDF(Buffer.from('pdf'));
expect(getImageBufferMock).not.toHaveBeenCalled();
expect(result).toEqual({
pages: 1,
text: 'hello ![img](https://textin.example.com/image.png)'
});
});
});
packages/service/test/worker/function.test.ts
...@@ -2,12 +2,14 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; ...@@ -2,12 +2,14 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { WorkerNameEnum } from '@fastgpt/service/worker/utils'; import { WorkerNameEnum } from '@fastgpt/service/worker/utils';
// hoisted: 这些 mock 必须在 vi.mock 工厂里可见 // hoisted: 这些 mock 必须在 vi.mock 工厂里可见
const { mockRun, mockGetWorkerController, mockRunWorker, mockEnv } = vi.hoisted(() => { const { mockRun, mockGetWorkerController, mockRunWorker, mockUploadImage2S3Bucket, mockEnv } =
vi.hoisted(() => {
const mockRun = vi.fn(); const mockRun = vi.fn();
return { return {
mockRun, mockRun,
mockGetWorkerController: vi.fn(() => ({ run: mockRun })), mockGetWorkerController: vi.fn(() => ({ run: mockRun })),
mockRunWorker: vi.fn(), mockRunWorker: vi.fn(),
mockUploadImage2S3Bucket: vi.fn(),
mockEnv: { mockEnv: {
PARSE_FILE_WORKERS: 10, PARSE_FILE_WORKERS: 10,
HTML_TO_MARKDOWN_WORKERS: 10, HTML_TO_MARKDOWN_WORKERS: 10,
...@@ -20,7 +22,7 @@ const { mockRun, mockGetWorkerController, mockRunWorker, mockEnv } = vi.hoisted( ...@@ -20,7 +22,7 @@ const { mockRun, mockGetWorkerController, mockRunWorker, mockEnv } = vi.hoisted(
PARSE_FILE_TIMEOUT_SECONDS: number; PARSE_FILE_TIMEOUT_SECONDS: number;
} }
}; };
}); });
// 拦截 getWorkerController / runWorker，保留 WorkerNameEnum 等枚举 // 拦截 getWorkerController / runWorker，保留 WorkerNameEnum 等枚举
vi.mock('@fastgpt/service/worker/utils', async (importOriginal) => { vi.mock('@fastgpt/service/worker/utils', async (importOriginal) => {
...@@ -37,6 +39,14 @@ vi.mock('@fastgpt/service/env', () => ({ ...@@ -37,6 +39,14 @@ vi.mock('@fastgpt/service/env', () => ({
serviceEnv: mockEnv serviceEnv: mockEnv
})); }));
vi.mock('@fastgpt/service/common/s3/utils', async (importOriginal) => {
const mod = await importOriginal<typeof import('@fastgpt/service/common/s3/utils')>();
return {
...mod,
uploadImage2S3Bucket: mockUploadImage2S3Bucket
};
});
// 必须在 vi.mock 之后再 import 被测模块 // 必须在 vi.mock 之后再 import 被测模块
const { text2Chunks, readRawContentFromBuffer } = await import('@fastgpt/service/worker/function'); const { text2Chunks, readRawContentFromBuffer } = await import('@fastgpt/service/worker/function');
const { htmlToMarkdown } = await import('@fastgpt/service/common/string/utils'); const { htmlToMarkdown } = await import('@fastgpt/service/common/string/utils');
...@@ -47,6 +57,7 @@ describe('worker/function', () => { ...@@ -47,6 +57,7 @@ describe('worker/function', () => {
mockGetWorkerController.mockReset(); mockGetWorkerController.mockReset();
mockGetWorkerController.mockImplementation(() => ({ run: mockRun })); mockGetWorkerController.mockImplementation(() => ({ run: mockRun }));
mockRunWorker.mockReset(); mockRunWorker.mockReset();
mockUploadImage2S3Bucket.mockReset();
}); });
describe('text2Chunks', () => { describe('text2Chunks', () => {
...@@ -206,6 +217,59 @@ describe('worker/function', () => { ...@@ -206,6 +217,59 @@ describe('worker/function', () => {
).rejects.toThrow('parse failed'); ).rejects.toThrow('parse failed');
}); });
it('传入 imageKeyOptions 时为 readFile worker 注册通用 uploadFile handler', async () => {
const expected = { rawText: 'parsed docx' };
const expiredTime = new Date('2030-01-01T00:00:00.000Z');
mockRun.mockResolvedValueOnce(expected);
mockUploadImage2S3Bucket.mockResolvedValueOnce('dataset/ds1/file-parsed/image.png');
const result = await readRawContentFromBuffer({
extension: 'docx',
encoding: 'utf-8',
buffer: Buffer.from('docx'),
imageKeyOptions: {
prefix: 'dataset/ds1/file-parsed',
expiredTime
}
});
expect(result).toEqual(expected);
const runArg = mockRun.mock.calls[0][0];
expect(runArg.imageKeyOptions).toEqual({
prefix: 'dataset/ds1/file-parsed',
expiredTime
});
const handlers = mockRun.mock.calls[0][2];
expect(handlers?.uploadFile).toBeInstanceOf(Function);
const uploadResult = await handlers.uploadFile({
name: '../image.png',
mime: 'image/png',
buffer: new Uint8Array([1, 2, 3]).buffer
});
expect(uploadResult).toEqual({
key: 'dataset/ds1/file-parsed/image.png'
});
expect(mockUploadImage2S3Bucket).toHaveBeenCalledWith('private', {
buffer: Buffer.from([1, 2, 3]),
uploadKey: 'dataset/ds1/file-parsed/image.png',
mimetype: 'image/png',
filename: 'image.png',
expiredTime
});
await expect(
handlers.uploadFile({
name: 'file.txt',
mime: 'text/plain',
buffer: new Uint8Array([1, 2, 3]).buffer
})
).rejects.toThrow('Unsupported worker uploadFile mime type: text/plain');
});
it('并发文件解析直接交给 readFile worker pool，并发数由 PARSE_FILE_WORKERS 决定', async () => { it('并发文件解析直接交给 readFile worker pool，并发数由 PARSE_FILE_WORKERS 决定', async () => {
let activeCount = 0; let activeCount = 0;
let maxActiveCount = 0; let maxActiveCount = 0;
...@@ -308,7 +372,7 @@ describe('worker/function', () => { ...@@ -308,7 +372,7 @@ describe('worker/function', () => {
}); });
it('通过 htmlStr2Md worker pool 派发并返回 rawText', async () => { it('通过 htmlStr2Md worker pool 派发并返回 rawText', async () => {
mockRun.mockResolvedValueOnce({ rawText: '# Title', imageList: [] }); mockRun.mockResolvedValueOnce({ rawText: '# Title' });
const result = await htmlToMarkdown('<h1>Title</h1>'); const result = await htmlToMarkdown('<h1>Title</h1>');
...@@ -326,7 +390,7 @@ describe('worker/function', () => { ...@@ -326,7 +390,7 @@ describe('worker/function', () => {
}); });
it('空 html 统一传空字符串', async () => { it('空 html 统一传空字符串', async () => {
mockRun.mockResolvedValueOnce({ rawText: '', imageList: [] }); mockRun.mockResolvedValueOnce({ rawText: '' });
const result = await htmlToMarkdown(null); const result = await htmlToMarkdown(null);
...@@ -336,7 +400,7 @@ describe('worker/function', () => { ...@@ -336,7 +400,7 @@ describe('worker/function', () => {
it('HTML_TO_MARKDOWN_WORKERS 自定义值生效', async () => { it('HTML_TO_MARKDOWN_WORKERS 自定义值生效', async () => {
mockEnv.HTML_TO_MARKDOWN_WORKERS = 6; mockEnv.HTML_TO_MARKDOWN_WORKERS = 6;
mockRun.mockResolvedValueOnce({ rawText: 'ok', imageList: [] }); mockRun.mockResolvedValueOnce({ rawText: 'ok' });
await htmlToMarkdown('<p>ok</p>'); await htmlToMarkdown('<p>ok</p>');
......
packages/service/test/worker/readFile/extension/docx.test.ts 0 → 100644
import { describe, it, expect, vi, beforeEach } from 'vitest';
const { mockConvertToHtml } = vi.hoisted(() => ({
mockConvertToHtml: vi.fn()
}));
vi.mock('mammoth', () => ({
default: {
convertToHtml: mockConvertToHtml
},
images: {
imgElement: (handler: any) => handler
}
}));
const { readDocsFile } = await import('@fastgpt/service/worker/readFile/extension/docx');
const createImage = () => ({
contentType: 'image/png',
read: vi.fn(async () => Buffer.from([1, 2, 3]))
});
describe('readDocsFile', () => {
beforeEach(() => {
mockConvertToHtml.mockReset();
});
it('docx 图片通过 uploadFile 回调上传，并把返回 key 写入 markdown', async () => {
const image = createImage();
mockConvertToHtml.mockImplementation(async (_input, options) => {
const { src } = await options.convertImage(image);
return {
value: `<p>hello</p><img src="${src}" />`
};
});
const uploadFile = vi.fn(async () => ({
key: 'dataset/file-parsed/image.png'
}));
const result = await readDocsFile(
{
buffer: Buffer.from('docx'),
encoding: 'utf-8',
extension: 'docx'
},
{ uploadFile }
);
expect(uploadFile).toHaveBeenCalledWith({
name: expect.stringMatching(/\.png$/),
mime: 'image/png',
buffer: new Uint8Array([1, 2, 3]).buffer
});
expect(result.rawText).toContain('dataset/file-parsed/image.png');
expect(result).not.toHaveProperty('imageList');
});
it('docx 包含图片但没有 uploadFile 时在 worker 内报错', async () => {
const image = createImage();
mockConvertToHtml.mockImplementation(async (_input, options) => {
await options.convertImage(image);
return {
value: '<p>never</p>'
};
});
await expect(
readDocsFile({
buffer: Buffer.from('docx'),
encoding: 'utf-8',
extension: 'docx'
})
).rejects.toBe('Can not read doc file, please convert to PDF');
});
});
packages/service/test/worker/readFile/extension/html.test.ts 0 → 100644
import { describe, expect, it, vi } from 'vitest';
import { readHtmlRawText } from '@fastgpt/service/worker/readFile/extension/html';
describe('readHtmlRawText', () => {
it('uploads base64 html images in worker and replaces image src with key', async () => {
const base64Data =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
const uploadFile = vi.fn(async () => ({
key: 'dataset/file-parsed/html-image.png'
}));
const result = await readHtmlRawText(
{
extension: 'html',
buffer: Buffer.from(`<p>hello</p><img src="data:image/png;base64,${base64Data}">`),
encoding: 'utf-8'
},
{ uploadFile }
);
expect(uploadFile).toHaveBeenCalledWith({
name: expect.stringMatching(/\.png$/),
mime: 'image/png',
buffer: expect.any(ArrayBuffer)
});
expect(result.rawText).toContain('dataset/file-parsed/html-image.png');
expect(result.rawText).not.toContain('data:image/png;base64');
expect(result).not.toHaveProperty('imageList');
});
});
packages/service/test/worker/readFile/extension/rawText.test.ts
import { describe, expect, it } from 'vitest'; import { describe, expect, it, vi } from 'vitest';
import { readFileRawText } from '@fastgpt/service/worker/readFile/extension/rawText'; import { readFileRawText } from '@fastgpt/service/worker/readFile/extension/rawText';
import { detectFileEncoding } from '@fastgpt/global/common/file/tools'; import { detectFileEncoding } from '@fastgpt/global/common/file/tools';
...@@ -104,12 +104,78 @@ describe('readFileRawText performance', () => { ...@@ -104,12 +104,78 @@ describe('readFileRawText performance', () => {
}); });
const duration = performance.now() - start; const duration = performance.now() - start;
expect(result.rawText.length).toBe(text.length); expect(result.rawText.length).toBe(text.trim().length);
expect(result.imageList ?? []).toHaveLength(0); expect(result).not.toHaveProperty('imageList');
expect(duration).toBeLessThan(PERFORMANCE_THRESHOLDS.largeUtf8Text); expect(duration).toBeLessThan(PERFORMANCE_THRESHOLDS.largeUtf8Text);
}); });
it('should extract 200 base64 images without pathological regex cost', async () => { it('should upload base64 images in worker and replace markdown image src with key', async () => {
const base64Data =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
const content = `段落\n\n![alt](data:image/png;base64,${base64Data})\n`;
const buffer = Buffer.from(content, 'utf8');
const uploadFile = vi.fn(async () => ({
key: 'dataset/file-parsed/image.png'
}));
const result = await readFileRawText(
{
extension: 'md',
buffer,
encoding: 'utf-8'
},
{ uploadFile }
);
expect(uploadFile).toHaveBeenCalledWith({
name: expect.stringMatching(/\.png$/),
mime: 'image/png',
buffer: expect.any(ArrayBuffer)
});
expect(result.rawText).toContain('![alt](dataset/file-parsed/image.png)');
expect(result).not.toHaveProperty('imageList');
});
it('should remove base64 images when uploadFile handler is missing', async () => {
const base64Data =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
const content = `段落\n\n![alt](data:image/png;base64,${base64Data})\n`;
const buffer = Buffer.from(content, 'utf8');
const result = await readFileRawText({
extension: 'md',
buffer,
encoding: 'utf-8'
});
expect(result.rawText).toContain('段落');
expect(result.rawText).not.toContain('data:image/png;base64');
expect(result.rawText).not.toContain('![alt]');
expect(result).not.toHaveProperty('imageList');
});
it('should reject oversized base64 image before upload', async () => {
const oversizedBase64 = 'A'.repeat(Math.ceil((40 * 1024 * 1024 + 1) / 3) * 4);
const content = `段落\n\n![alt](data:image/png;base64,${oversizedBase64})\n`;
const uploadFile = vi.fn(async () => ({
key: 'dataset/file-parsed/image.png'
}));
const result = await readFileRawText(
{
extension: 'md',
buffer: Buffer.from(content, 'utf8'),
encoding: 'utf-8'
},
{ uploadFile }
);
expect(uploadFile).not.toHaveBeenCalled();
expect(result.rawText).toContain('段落');
expect(result.rawText).not.toContain('data:image/png;base64');
});
it('should process 200 base64 images without carrying imageList', async () => {
const base64Data = const base64Data =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=='; 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
const imageCount = 200; const imageCount = 200;
...@@ -118,19 +184,24 @@ describe('readFileRawText performance', () => { ...@@ -118,19 +184,24 @@ describe('readFileRawText performance', () => {
(_, i) => `段落 ${i}\n\n![alt-${i}](data:image/png;base64,${base64Data})\n` (_, i) => `段落 ${i}\n\n![alt-${i}](data:image/png;base64,${base64Data})\n`
).join('\n'); ).join('\n');
const buffer = Buffer.from(content, 'utf8'); const buffer = Buffer.from(content, 'utf8');
const uploadFile = vi.fn(async ({ name }: { name: string }) => ({
key: `dataset/file-parsed/${name}`
}));
const start = performance.now(); const start = performance.now();
const result = await readFileRawText({ const result = await readFileRawText(
{
extension: 'md', extension: 'md',
buffer, buffer,
encoding: 'utf-8' encoding: 'utf-8'
}); },
{ uploadFile }
);
const duration = performance.now() - start; const duration = performance.now() - start;
const imageList = result.imageList ?? []; expect(uploadFile).toHaveBeenCalledTimes(imageCount);
expect(imageList).toHaveLength(imageCount); expect(result).not.toHaveProperty('imageList');
expect(imageList[0].base64).toBe(base64Data); expect(result.rawText).not.toContain('data:image/png;base64');
expect(imageList[0].mime).toBe('image/png');
expect(duration).toBeLessThan(PERFORMANCE_THRESHOLDS.manyBase64Images); expect(duration).toBeLessThan(PERFORMANCE_THRESHOLDS.manyBase64Images);
}); });
}); });
packages/service/test/worker/readFile/integration.test.ts
import { describe, it, expect, beforeAll, afterEach, afterAll, vi } from 'vitest'; import { describe, it, expect, beforeAll, afterEach, afterAll, vi } from 'vitest';
import path from 'path'; import path from 'path';
import { existsSync, readFileSync } from 'fs'; import { existsSync, readFileSync } from 'fs';
import { JSZip } from '@fastgpt/service/core/ai/skill/package';
const { mockUploadImage2S3Bucket } = vi.hoisted(() => ({
mockUploadImage2S3Bucket: vi.fn()
}));
vi.mock('@fastgpt/service/common/s3/utils', async (importOriginal) => {
const mod = await importOriginal<typeof import('@fastgpt/service/common/s3/utils')>();
return {
...mod,
uploadImage2S3Bucket: mockUploadImage2S3Bucket
};
});
/* /*
* 真实 spawn 测试：使用 projects/app/worker/readFile.js 构建产物， * 真实 spawn 测试：使用 projects/app/worker/readFile.js 构建产物，
...@@ -54,6 +67,73 @@ const getPositiveIntegerEnv = (name: string, defaultValue: number) => { ...@@ -54,6 +67,73 @@ const getPositiveIntegerEnv = (name: string, defaultValue: number) => {
return Number.isInteger(value) && value > 0 ? value : defaultValue; return Number.isInteger(value) && value > 0 ? value : defaultValue;
}; };
const createDocxWithImage = async () => {
const zip = new JSZip();
const png = Buffer.from(
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=',
'base64'
);
zip.file(
'[Content_Types].xml',
`<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
<Default Extension="xml" ContentType="application/xml"/>
<Default Extension="png" ContentType="image/png"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>`
);
zip.file(
'_rels/.rels',
`<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>`
);
zip.file(
'word/_rels/document.xml.rels',
`<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rIdImage1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/image" Target="media/image1.png"/>
</Relationships>`
);
zip.file(
'word/document.xml',
`<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"
xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"
xmlns:wp="http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing"
xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"
xmlns:pic="http://schemas.openxmlformats.org/drawingml/2006/picture">
<w:body>
<w:p><w:r><w:t>hello docx image</w:t></w:r></w:p>
<w:p>
<w:r>
<w:drawing>
<wp:inline>
<wp:docPr id="1" name="Picture 1"/>
<a:graphic>
<a:graphicData uri="http://schemas.openxmlformats.org/drawingml/2006/picture">
<pic:pic>
<pic:blipFill>
<a:blip r:embed="rIdImage1"/>
</pic:blipFill>
</pic:pic>
</a:graphicData>
</a:graphic>
</wp:inline>
</w:drawing>
</w:r>
</w:p>
</w:body>
</w:document>`
);
zip.file('word/media/image1.png', png);
return zip.generateAsync({ type: 'nodebuffer' });
};
const destroyReadFilePool = async () => { const destroyReadFilePool = async () => {
const workerPoll = (global as any).workerPoll; const workerPoll = (global as any).workerPoll;
const pool = workerPoll?.[WorkerNameEnum.readFile]; const pool = workerPoll?.[WorkerNameEnum.readFile];
...@@ -110,6 +190,34 @@ describeIfEnabled('readFile worker (real spawn integration)', () => { ...@@ -110,6 +190,34 @@ describeIfEnabled('readFile worker (real spawn integration)', () => {
expect(result.rawText).toContain('item 1'); expect(result.rawText).toContain('item 1');
}); });
it('解析带 base64 图片的 md 时通过主线程 uploadFile handler 上传图片', async () => {
mockUploadImage2S3Bucket.mockResolvedValueOnce('dataset/test/md-parsed/image.png');
const base64Data =
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
const result = await readRawContentFromBuffer({
extension: 'md',
encoding: 'utf-8',
buffer: Buffer.from(`hello\n\n![alt](data:image/png;base64,${base64Data})`, 'utf-8'),
imageKeyOptions: {
prefix: 'dataset/test/md-parsed'
}
});
expect(result.rawText).toContain('hello');
expect(result.rawText).toContain('![alt](dataset/test/md-parsed/image.png)');
expect(result).not.toHaveProperty('imageList');
expect(mockUploadImage2S3Bucket).toHaveBeenCalledWith(
'private',
expect.objectContaining({
buffer: expect.any(Buffer),
uploadKey: expect.stringMatching(/^dataset\/test\/md-parsed\/.+\.png$/),
mimetype: 'image/png',
filename: expect.stringMatching(/\.png$/)
})
);
});
it('解析 csv', async () => { it('解析 csv', async () => {
const csv = 'name,age,city\nAlice,30,Beijing\nBob,25,Shanghai'; const csv = 'name,age,city\nAlice,30,Beijing\nBob,25,Shanghai';
const result = await readRawContentFromBuffer({ const result = await readRawContentFromBuffer({
...@@ -123,6 +231,31 @@ describeIfEnabled('readFile worker (real spawn integration)', () => { ...@@ -123,6 +231,31 @@ describeIfEnabled('readFile worker (real spawn integration)', () => {
expect(result.rawText).toContain('Shanghai'); expect(result.rawText).toContain('Shanghai');
}); });
it('解析带图片 docx 时通过主线程 uploadFile handler 上传图片', async () => {
mockUploadImage2S3Bucket.mockResolvedValueOnce('dataset/test/docx-parsed/image.png');
const result = await readRawContentFromBuffer({
extension: 'docx',
encoding: 'utf-8',
buffer: await createDocxWithImage(),
imageKeyOptions: {
prefix: 'dataset/test/docx-parsed'
}
});
expect(result.rawText).toContain('hello docx image');
expect(result.rawText).toContain('dataset/test/docx-parsed/image.png');
expect(mockUploadImage2S3Bucket).toHaveBeenCalledWith(
'private',
expect.objectContaining({
buffer: expect.any(Buffer),
uploadKey: expect.stringMatching(/^dataset\/test\/docx-parsed\/.+\.png$/),
mimetype: 'image/png',
filename: expect.stringMatching(/\.png$/)
})
);
});
itIfPdfFixture( itIfPdfFixture(
'解析 pdf（真实 worker + LiteParse）', '解析 pdf（真实 worker + LiteParse）',
async () => { async () => {
......
packages/service/test/worker/utils.test.ts 0 → 100644
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import fs from 'fs';
import os from 'os';
import path from 'path';
const { WorkerPool, WorkerNameEnum } = await import('@fastgpt/service/worker/utils');
const workerScript = `
const { parentPort } = require('worker_threads');
parentPort.on('message', (message) => {
const { id } = message;
parentPort.once('message', (response) => {
if (response.type === 'uploadFileResult') {
parentPort.postMessage({
id,
type: 'success',
data: response.data
});
return;
}
parentPort.postMessage({
id,
type: 'error',
data: response.data
});
});
parentPort.postMessage({
id,
type: 'uploadFile',
requestId: 'upload-1',
data: {
name: 'image.png',
mime: 'image/png',
buffer: new Uint8Array([1, 2, 3]).buffer
}
});
});
`;
describe('worker/utils WorkerPool', () => {
let tmpDir: string;
let cwdSpy: ReturnType<typeof vi.spyOn>;
beforeEach(() => {
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'fastgpt-worker-test-'));
fs.mkdirSync(path.join(tmpDir, 'worker'), { recursive: true });
fs.writeFileSync(path.join(tmpDir, 'worker', 'readFile.js'), workerScript);
cwdSpy = vi.spyOn(process, 'cwd').mockReturnValue(tmpDir);
});
afterEach(() => {
cwdSpy.mockRestore();
fs.rmSync(tmpDir, { recursive: true, force: true });
});
it('处理 worker 通用 uploadFile 中间事件，不提前结束任务', async () => {
const pool = new WorkerPool<{ payload: string }, { key: string }>({
name: WorkerNameEnum.readFile,
maxReservedThreads: 1
});
const uploadFile = vi.fn().mockResolvedValue({
key: 'parsed/image.png'
});
const result = await pool.run({ payload: 'run' }, undefined, { uploadFile });
expect(uploadFile).toHaveBeenCalledWith({
name: 'image.png',
mime: 'image/png',
buffer: expect.any(ArrayBuffer)
});
expect(result).toEqual({
key: 'parsed/image.png'
});
expect(pool.workerQueue[0].status).toBe('idle');
});
it('uploadFile handler 失败时把错误回传给 worker', async () => {
const pool = new WorkerPool<{ payload: string }, { key: string; src: string }>({
name: WorkerNameEnum.readFile,
maxReservedThreads: 1
});
const uploadError = new Error('upload failed');
const uploadFile = vi.fn().mockRejectedValue(uploadError);
await expect(pool.run({ payload: 'run' }, undefined, { uploadFile })).rejects.toEqual(
uploadError
);
});
});
packages/service/thirdProvider/doc2x/index.ts
import { batchRun, delay } from '@fastgpt/global/common/system/utils'; import { delay } from '@fastgpt/global/common/system/utils';
import { htmlTable2Md } from '@fastgpt/global/common/string/markdown'; import { htmlTable2Md, parseMarkdownBase64Images } from '@fastgpt/global/common/string/markdown';
import { type Method } from 'axios'; import { type Method } from 'axios';
import { getNanoid } from '@fastgpt/global/common/string/tools';
import { getErrText } from '@fastgpt/global/common/error/utils'; import { getErrText } from '@fastgpt/global/common/error/utils';
import { type ImageType } from '../../worker/readFile/type'; import { getImageBuffer } from '../../common/file/image/utils';
import { getImageBase64 } from '../../common/file/image/utils';
import { createProxyAxios, axios } from '../../common/api/axios'; import { createProxyAxios, axios } from '../../common/api/axios';
import { getLogger, LogCategories } from '../../common/logger'; import { getLogger, LogCategories } from '../../common/logger';
import { type UploadedFileResult } from '../../worker/readFile/type';
type ApiResponseDataType<T = any> = { type ApiResponseDataType<T = any> = {
code: string; code: string;
...@@ -14,6 +13,22 @@ type ApiResponseDataType<T = any> = { ...@@ -14,6 +13,22 @@ type ApiResponseDataType<T = any> = {
data: T; data: T;
}; };
type Doc2xImageUploadHandler = (
params:
| {
type: 'http';
url: string;
mime: string;
buffer: Buffer;
}
| {
type: 'base64';
mime: string;
base64: string;
dataUrl: string;
}
) => Promise<UploadedFileResult>;
export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
const logger = getLogger(LogCategories.MODULE.DATASET.FILE); const logger = getLogger(LogCategories.MODULE.DATASET.FILE);
// Init request // Init request
...@@ -72,7 +87,12 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -72,7 +87,12 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
.catch((err) => responseError(err)); .catch((err) => responseError(err));
}; };
const parsePDF = async (fileBuffer: Buffer) => { const parsePDF = async (
fileBuffer: Buffer,
options: {
uploadImage?: Doc2xImageUploadHandler;
} = {}
) => {
logger.debug('Doc2x PDF parse started'); logger.debug('Doc2x PDF parse started');
const startTime = Date.now(); const startTime = Date.now();
...@@ -185,45 +205,35 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -185,45 +205,35 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
const { text, pages } = await checkResult(); const { text, pages } = await checkResult();
// ![](url) => ![](base64) const formatText = await parseMarkdownBase64Images(htmlTable2Md(text), {
const parseTextImage = async (text: string) => { parseBase64: true,
// Extract image links and convert to base64 parseHttp: true,
const imageList: { id: string; url: string }[] = []; controller: options.uploadImage
let processedText = text.replace(/!\[.*?\]\((http[^)]+)\)/g, (match, url) => { ? async (image) => {
const id = `IMAGE_${getNanoid()}_IMAGE`; if (image.type === 'base64') {
imageList.push({ return options.uploadImage!({
id, type: 'base64',
url mime: image.mime,
}); base64: image.base64,
return `![](${id})`; dataUrl: image.dataUrl
}); });
}
// Get base64 from image url
const resultImageList: ImageType[] = [];
await batchRun(
imageList,
async (item) => {
try { try {
const { base64, mime } = await getImageBase64(item.url); const { buffer, mime } = await getImageBuffer(image.url);
resultImageList.push({ return options.uploadImage!({
uuid: item.id, type: 'http',
url: image.url,
mime, mime,
base64 buffer
}); });
} catch (error) { } catch (error) {
processedText = processedText.replace(item.id, item.url); logger.warn('Doc2x image transfer failed', { url: image.url, error });
logger.warn('Doc2x image fetch failed', { url: item.url, error }); throw error;
} }
}, }
5 : undefined
); });
return {
text: processedText,
imageList: resultImageList
};
};
const { text: formatText, imageList } = await parseTextImage(htmlTable2Md(text));
logger.debug('Doc2x PDF parse finished', { logger.debug('Doc2x PDF parse finished', {
durationMs: Date.now() - startTime, durationMs: Date.now() - startTime,
...@@ -232,8 +242,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => { ...@@ -232,8 +242,7 @@ export const useDoc2xServer = ({ apiKey }: { apiKey: string }) => {
return { return {
pages, pages,
text: formatText, text: formatText
imageList
}; };
}; };
......
packages/service/thirdProvider/textin/index.ts
import { matchMdImg } from '@fastgpt/global/common/string/markdown';
import { createProxyAxios } from '../../common/api/axios'; import { createProxyAxios } from '../../common/api/axios';
import { getErrText } from '@fastgpt/global/common/error/utils'; import { getErrText } from '@fastgpt/global/common/error/utils';
import { getLogger, LogCategories } from '../../common/logger'; import { getLogger, LogCategories } from '../../common/logger';
import { parseMarkdownBase64Images } from '@fastgpt/global/common/string/markdown';
import { type UploadedFileResult } from '../../worker/readFile/type';
import { getImageBuffer } from '../../common/file/image/utils';
type TextinImageUploadHandler = (
params:
| {
type: 'http';
url: string;
mime: string;
buffer: Buffer;
}
| {
type: 'base64';
mime: string;
base64: string;
dataUrl: string;
}
) => Promise<UploadedFileResult>;
export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCode: string }) => { export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCode: string }) => {
const logger = getLogger(LogCategories.MODULE.DATASET.FILE); const logger = getLogger(LogCategories.MODULE.DATASET.FILE);
...@@ -36,7 +54,12 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo ...@@ -36,7 +54,12 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo
return Promise.reject({ message: `[Textin] ${getErrText(err)}` }); return Promise.reject({ message: `[Textin] ${getErrText(err)}` });
}; };
const parsePDF = async (fileBuffer: Buffer) => { const parsePDF = async (
fileBuffer: Buffer,
options: {
uploadImage?: TextinImageUploadHandler;
} = {}
) => {
logger.debug('Textin PDF parse started'); logger.debug('Textin PDF parse started');
const startTime = Date.now(); const startTime = Date.now();
...@@ -44,7 +67,7 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo ...@@ -44,7 +67,7 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo
// Build request parameters (https://docs.textin.com/xparse/parse-quickstart#url%E5%8F%82%E6%95%B0%E8%AF%B4%E6%98%8E) // Build request parameters (https://docs.textin.com/xparse/parse-quickstart#url%E5%8F%82%E6%95%B0%E8%AF%B4%E6%98%8E)
const params = { const params = {
get_image: 'objects', // 返回页面内的子图像 get_image: 'objects', // 返回页面内的子图像
image_output_type: 'base64str', // 图片对象以base64字符串返回 image_output_type: 'default', // 图片对象返回临时链接，避免响应体携带大体积 base64
parse_mode: 'auto', // 自动模式：直接提取pdf中的文字 parse_mode: 'auto', // 自动模式：直接提取pdf中的文字
dpi: 144, // 坐标基准144 dpi dpi: 144, // 坐标基准144 dpi
markdown_details: 1, // 返回detail字段（markdown元素详细信息） markdown_details: 1, // 返回detail字段（markdown元素详细信息）
...@@ -73,8 +96,30 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo ...@@ -73,8 +96,30 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo
return Promise.reject('[Textin] No markdown content in response'); return Promise.reject('[Textin] No markdown content in response');
} }
logger.debug('Textin markdown content received', { length: rawMarkdown.length }); logger.debug('Textin markdown content received', { length: rawMarkdown.length });
// Process tables and images (reuse existing utility functions) const text = await parseMarkdownBase64Images(rawMarkdown, {
const { text, imageList } = matchMdImg(rawMarkdown); parseBase64: true,
parseHttp: true,
controller: options.uploadImage
? async (image) => {
if (image.type === 'base64') {
return options.uploadImage!({
type: 'base64',
mime: image.mime,
base64: image.base64,
dataUrl: image.dataUrl
});
}
const { buffer, mime } = await getImageBuffer(image.url);
return options.uploadImage!({
type: 'http',
url: image.url,
mime,
buffer
});
}
: undefined
});
// Get page count // Get page count
const pages = data.result?.pages?.length || data.result?.total_page_number || 1; const pages = data.result?.pages?.length || data.result?.total_page_number || 1;
...@@ -86,8 +131,7 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo ...@@ -86,8 +131,7 @@ export const useTextinServer = ({ appId, secretCode }: { appId: string; secretCo
return { return {
pages, pages,
text, text
imageList
}; };
} catch (error) { } catch (error) {
return responseError(error); return responseError(error);
......
packages/service/types/joplin-turndown-plugin-gfm.d.ts 0 → 100644
declare module 'joplin-turndown-plugin-gfm' {
import type TurndownService from 'turndown';
export function gfm(turndownService: TurndownService): void;
}
packages/service/worker/function.ts
...@@ -7,6 +7,9 @@ import { getWorkerController, WorkerNameEnum } from './utils'; ...@@ -7,6 +7,9 @@ import { getWorkerController, WorkerNameEnum } from './utils';
import type { ReadFileResponse } from './readFile/type'; import type { ReadFileResponse } from './readFile/type';
import { isTestEnv } from '@fastgpt/global/common/system/constants'; import { isTestEnv } from '@fastgpt/global/common/system/constants';
import { serviceEnv } from '../env'; import { serviceEnv } from '../env';
import { uploadImage2S3Bucket } from '../common/s3/utils';
import { normalizeMimeType, resolveMimeType } from '../common/s3/utils/mime';
import path from 'node:path';
export const text2Chunks = (props: SplitProps) => { export const text2Chunks = (props: SplitProps) => {
// Test env, not run worker // Test env, not run worker
...@@ -27,6 +30,10 @@ type ReadFileWorkerProps = { ...@@ -27,6 +30,10 @@ type ReadFileWorkerProps = {
buffer?: ArrayBuffer; buffer?: ArrayBuffer;
sharedBuffer?: SharedArrayBuffer; sharedBuffer?: SharedArrayBuffer;
bufferSize: number; bufferSize: number;
imageKeyOptions?: {
prefix: string;
expiredTime?: Date;
};
}; };
const getReadFileWorker = () => const getReadFileWorker = () =>
...@@ -43,6 +50,10 @@ export const readRawContentFromBuffer = (props: { ...@@ -43,6 +50,10 @@ export const readRawContentFromBuffer = (props: {
extension: string; extension: string;
encoding: string; encoding: string;
buffer: Buffer; buffer: Buffer;
imageKeyOptions?: {
prefix: string;
expiredTime?: Date;
};
}) => { }) => {
const bufferSize = props.buffer.length; const bufferSize = props.buffer.length;
const sourceArrayBuffer = props.buffer.buffer; const sourceArrayBuffer = props.buffer.buffer;
...@@ -51,6 +62,28 @@ export const readRawContentFromBuffer = (props: { ...@@ -51,6 +62,28 @@ export const readRawContentFromBuffer = (props: {
props.buffer.byteLength === sourceArrayBuffer.byteLength && props.buffer.byteLength === sourceArrayBuffer.byteLength &&
sourceArrayBuffer instanceof ArrayBuffer; sourceArrayBuffer instanceof ArrayBuffer;
const uploadFile = props.imageKeyOptions?.prefix
? async ({ name, mime, buffer }: { name: string; mime: string; buffer: ArrayBuffer }) => {
const mimetype = normalizeMimeType(mime);
if (!mimetype.startsWith('image/')) {
throw new Error(`Unsupported worker uploadFile mime type: ${mimetype}`);
}
// uploadFile 是 worker 通用能力，主线程只接受文件名，避免 worker 传入路径片段越过 prefix。
const filename = path.basename(name);
const key = await uploadImage2S3Bucket('private', {
buffer: Buffer.from(buffer),
uploadKey: `${props.imageKeyOptions!.prefix}/${filename}`,
mimetype: resolveMimeType([filename], mimetype),
filename,
expiredTime: props.imageKeyOptions?.expiredTime
});
return {
key
};
}
: undefined;
if (canTransferBuffer) { if (canTransferBuffer) {
/** /**
* 大文件解析时优先 transfer 独占 ArrayBuffer，避免再复制一份 SharedArrayBuffer。 * 大文件解析时优先 transfer 独占 ArrayBuffer，避免再复制一份 SharedArrayBuffer。
...@@ -61,9 +94,11 @@ export const readRawContentFromBuffer = (props: { ...@@ -61,9 +94,11 @@ export const readRawContentFromBuffer = (props: {
extension: props.extension, extension: props.extension,
encoding: props.encoding, encoding: props.encoding,
buffer: sourceArrayBuffer, buffer: sourceArrayBuffer,
bufferSize bufferSize,
imageKeyOptions: props.imageKeyOptions
}, },
[sourceArrayBuffer] [sourceArrayBuffer],
{ uploadFile }
); );
} }
...@@ -71,10 +106,15 @@ export const readRawContentFromBuffer = (props: { ...@@ -71,10 +106,15 @@ export const readRawContentFromBuffer = (props: {
const sharedArray = new Uint8Array(sharedBuffer); const sharedArray = new Uint8Array(sharedBuffer);
sharedArray.set(props.buffer); sharedArray.set(props.buffer);
return getReadFileWorker().run({ return getReadFileWorker().run(
{
extension: props.extension, extension: props.extension,
encoding: props.encoding, encoding: props.encoding,
sharedBuffer, sharedBuffer,
bufferSize bufferSize,
}); imageKeyOptions: props.imageKeyOptions
},
undefined,
{ uploadFile }
);
}; };
packages/service/worker/htmlStr2Md/index.ts
import { parentPort } from 'worker_threads'; import { parentPort } from 'worker_threads';
import { html2md } from './utils'; import { html2md } from './utils';
import {
createWorkerUploadFileHandler,
handleWorkerUploadFileResponse,
isWorkerUploadFileResponse
} from '../utils/uploadFile';
type IncomingMessage = { type IncomingMessage = {
id: string; id: string;
html: string; html: string;
uploadImages?: boolean;
type?: 'uploadFileResult' | 'uploadFileError';
requestId?: string;
data?: any;
}; };
parentPort?.on('message', (params: IncomingMessage) => { parentPort?.on('message', async (params: IncomingMessage) => {
const { id, html } = params; const { id, html, requestId, data, type } = params;
if (isWorkerUploadFileResponse(type)) {
handleWorkerUploadFileResponse({
taskId: id,
type,
requestId,
data
});
return;
}
const uploadFileHandler = createWorkerUploadFileHandler({
taskId: id,
parentPort
});
try { try {
const md = html2md(html || ''); const md = await html2md(html || '', {
uploadFile: params.uploadImages ? uploadFileHandler.uploadFile : undefined
});
parentPort?.postMessage({ id, type: 'success', data: md }); parentPort?.postMessage({ id, type: 'success', data: md });
} catch (error) { } catch (error) {
parentPort?.postMessage({ id, type: 'error', data: error }); parentPort?.postMessage({ id, type: 'error', data: error });
} finally {
uploadFileHandler.cleanup();
} }
}); });
packages/service/worker/htmlStr2Md/utils.ts
import TurndownService from 'turndown'; import TurndownService from 'turndown';
import { type ImageType } from '../readFile/type';
import { getNanoid } from '@fastgpt/global/common/string/tools';
import { simpleMarkdownText } from '@fastgpt/global/common/string/markdown'; import { simpleMarkdownText } from '@fastgpt/global/common/string/markdown';
import { getLogger, LogCategories } from '../../common/logger'; import { getLogger, LogCategories } from '../../common/logger';
import { workerEnv } from '../env'; import { workerEnv } from '../env';
// @ts-ignore import { gfm } from 'joplin-turndown-plugin-gfm';
const turndownPluginGfm = require('joplin-turndown-plugin-gfm'); import { uploadBase64Image } from '../utils/base64ImageUpload';
import { type UploadFileHandler } from '../readFile/type';
import { batchRun } from '@fastgpt/global/common/system/utils';
const MAX_HTML_SIZE = workerEnv.MAX_HTML_TRANSFORM_CHARS; const MAX_HTML_SIZE = workerEnv.MAX_HTML_TRANSFORM_CHARS;
const logger = getLogger(LogCategories.INFRA.WORKER); const logger = getLogger(LogCategories.INFRA.WORKER);
const htmlBase64UploadConcurrency = 5;
const processBase64Images = (htmlContent: string) => { const htmlBase64SrcRegex = /\bsrc\s*=\s*(["'])data:([^;]+);base64,([A-Za-z0-9+/=]+)\1/gi;
// 优化后的正则:
// 1. 使用精确的 base64 字符集 [A-Za-z0-9+/=]+ 避免回溯 /**
// 2. 明确捕获 mime 类型和 base64 数据 * HTML 转 markdown 前实时处理 base64 图片。
// 3. 减少不必要的捕获组 *
const base64Regex = /src="data:([^;]+);base64,([A-Za-z0-9+/=]+)"/g; * 有 uploadFile 时上传为对象存储 key；没有 uploadFile 时删除 src，避免大体积 base64
const images: ImageType[] = []; * 进入 turndown 或被 worker 结果回传。
*/
const processedHtml = htmlContent.replace(base64Regex, (_match, mime, base64Data) => { const processBase64Images = async (
const uuid = `IMAGE_${getNanoid(12)}_IMAGE`; htmlContent: string,
images.push({ options: {
uuid, uploadFile?: UploadFileHandler;
} = {}
) => {
const matches = Array.from(htmlContent.matchAll(htmlBase64SrcRegex));
if (matches.length === 0) return htmlContent;
const replacements = await batchRun(
matches,
async (match) => {
const [, quote, mime, base64Data] = match;
if (!options.uploadFile) {
return `src=${quote}${quote}`;
}
try {
const { key } = await uploadBase64Image({
mime,
base64: base64Data, base64: base64Data,
mime uploadFile: options.uploadFile
});
return `src="${uuid}"`;
}); });
return `src=${quote}${key}${quote}`;
} catch (error) {
logger.warn('Failed to upload parsed HTML base64 image', { mime, error });
return `src=${quote}${quote}`;
}
},
htmlBase64UploadConcurrency
);
let result = '';
let lastIndex = 0;
for (const [matchIndex, match] of matches.entries()) {
const [fullMatch] = match;
const index = match.index ?? 0;
result += htmlContent.slice(lastIndex, index);
result += replacements[matchIndex];
lastIndex = index + fullMatch.length;
}
return { processedHtml, images }; return result + htmlContent.slice(lastIndex);
}; };
export const html2md = ( export const html2md = async (
html: string html: string,
): { options: {
uploadFile?: UploadFileHandler;
} = {}
): Promise<{
rawText: string; rawText: string;
imageList: ImageType[]; }> => {
} => {
const turndownService = new TurndownService({ const turndownService = new TurndownService({
headingStyle: 'atx', headingStyle: 'atx',
bulletListMarker: '-', bulletListMarker: '-',
...@@ -50,7 +88,7 @@ export const html2md = ( ...@@ -50,7 +88,7 @@ export const html2md = (
try { try {
turndownService.remove(['i', 'script', 'iframe', 'style']); turndownService.remove(['i', 'script', 'iframe', 'style']);
turndownService.use(turndownPluginGfm.gfm); turndownService.use(gfm);
// add custom handling for media tag // add custom handling for media tag
turndownService.addRule('media', { turndownService.addRule('media', {
...@@ -71,25 +109,28 @@ export const html2md = ( ...@@ -71,25 +109,28 @@ export const html2md = (
}); });
// Base64 img to id, otherwise it will occupy memory when going to md // Base64 img to id, otherwise it will occupy memory when going to md
const { processedHtml, images } = processBase64Images(html); const processedHtml = await processBase64Images(html, {
uploadFile: options.uploadFile
});
// if html is too large, return the original html (but preserve image list) // if html is too large, return the original html
if (processedHtml.length > MAX_HTML_SIZE) { if (processedHtml.length > MAX_HTML_SIZE) {
return { rawText: processedHtml, imageList: images }; return { rawText: processedHtml };
} }
const md = turndownService.turndown(processedHtml); const md = turndownService.turndown(processedHtml);
// const { text, imageList } = matchMdImg(md);
return { return {
rawText: simpleMarkdownText(md), rawText: simpleMarkdownText(md)
imageList: images
}; };
} catch (error) { } catch (error) {
if (options.uploadFile) {
throw error;
}
logger.error('HTML to markdown conversion failed', { error }); logger.error('HTML to markdown conversion failed', { error });
return { return {
rawText: '', rawText: ''
imageList: []
}; };
} }
}; };
packages/service/worker/readFile/extension/docx.ts
import mammoth, { images } from 'mammoth'; import mammoth, { images } from 'mammoth';
import { type ReadRawTextByBuffer, type ReadFileResponse, type ImageType } from '../type'; import { type ReadRawTextByBuffer, type ReadFileResponse, type UploadFileHandler } from '../type';
import { html2md } from '../../htmlStr2Md/utils'; import { html2md } from '../../htmlStr2Md/utils';
import { getLogger, LogCategories } from '../../../common/logger'; import { getLogger, LogCategories } from '../../../common/logger';
import { resolveMimeExtension } from '../../../common/s3/utils/mime';
/** /**
* read docx to markdown * read docx to markdown
*/ */
export const readDocsFile = async ({ buffer }: ReadRawTextByBuffer): Promise<ReadFileResponse> => { export const readDocsFile = async (
const imageList: ImageType[] = []; { buffer }: ReadRawTextByBuffer,
options: {
uploadFile?: UploadFileHandler;
} = {}
): Promise<ReadFileResponse> => {
const logger = getLogger(LogCategories.INFRA.WORKER); const logger = getLogger(LogCategories.INFRA.WORKER);
try { try {
const { value: html } = await mammoth.convertToHtml( const { value: html } = await mammoth.convertToHtml(
...@@ -17,26 +22,39 @@ export const readDocsFile = async ({ buffer }: ReadRawTextByBuffer): Promise<Rea ...@@ -17,26 +22,39 @@ export const readDocsFile = async ({ buffer }: ReadRawTextByBuffer): Promise<Rea
{ {
ignoreEmptyParagraphs: false, ignoreEmptyParagraphs: false,
convertImage: images.imgElement(async (image) => { convertImage: images.imgElement(async (image) => {
const imageBase64 = await image.readAsBase64String();
const uuid = crypto.randomUUID();
const mime = image.contentType; const mime = image.contentType;
imageList.push({ const name = `${crypto.randomUUID()}${resolveMimeExtension(mime)}`;
uuid,
base64: imageBase64, if (!options.uploadFile) {
mime logger.warn('Missing image upload handler when parsing docx image', { name, mime });
throw new Error('Missing imageKeyOptions.prefix for parsed document image upload');
}
const imageBuffer = await image.read();
const imageArrayBuffer = new Uint8Array(imageBuffer.byteLength);
imageArrayBuffer.set(imageBuffer);
const { key } = await options
.uploadFile({
name,
mime,
buffer: imageArrayBuffer.buffer
})
.catch((error) => {
logger.warn('Failed to upload docx image from worker', { name, mime, error });
throw error;
}); });
return { return {
src: uuid src: key
}; };
}) })
} }
); );
const { rawText } = html2md(html); const { rawText } = await html2md(html);
return { return {
rawText, rawText
imageList
}; };
} catch (error) { } catch (error) {
logger.error('Failed to parse docx file', { error }); logger.error('Failed to parse docx file', { error });
......
packages/service/worker/readFile/extension/html.ts
import { type ReadRawTextByBuffer, type ReadFileResponse } from '../type'; import { type ReadRawTextByBuffer, type ReadFileResponse, type UploadFileHandler } from '../type';
import { readFileRawText } from './rawText'; import { readFileRawText } from './rawText';
import { html2md } from '../../htmlStr2Md/utils'; import { html2md } from '../../htmlStr2Md/utils';
export const readHtmlRawText = async (params: ReadRawTextByBuffer): Promise<ReadFileResponse> => { export const readHtmlRawText = async (
const { rawText: html } = await readFileRawText(params); params: ReadRawTextByBuffer,
options: {
const { rawText, imageList } = html2md(html); uploadFile?: UploadFileHandler;
} = {}
): Promise<ReadFileResponse> => {
const { rawText: html } = await readFileRawText(params, {
uploadFile: options.uploadFile
});
const { rawText } = await html2md(html, {
uploadFile: options.uploadFile
});
return { return {
rawText, rawText
imageList
}; };
}; };
packages/service/worker/readFile/extension/rawText.ts
import iconv from 'iconv-lite'; import iconv from 'iconv-lite';
import { type ReadRawTextByBuffer, type ReadFileResponse } from '../type'; import { type ReadRawTextByBuffer, type ReadFileResponse, type UploadFileHandler } from '../type';
import { matchMdImg } from '@fastgpt/global/common/string/markdown'; import { parseMarkdownBase64Images } from '@fastgpt/global/common/string/markdown';
import { uploadBase64Image } from '../../utils/base64ImageUpload';
const hasNonAsciiByte = (buffer: Buffer) => { const hasNonAsciiByte = (buffer: Buffer) => {
for (let i = 0; i < buffer.length; i++) { for (let i = 0; i < buffer.length; i++) {
...@@ -25,10 +26,12 @@ const rawEncodingList = [ ...@@ -25,10 +26,12 @@ const rawEncodingList = [
]; ];
// 加载源文件内容 // 加载源文件内容
export const readFileRawText = async ({ export const readFileRawText = async (
buffer, { buffer, encoding }: ReadRawTextByBuffer,
encoding options: {
}: ReadRawTextByBuffer): Promise<ReadFileResponse> => { uploadFile?: UploadFileHandler;
} = {}
): Promise<ReadFileResponse> => {
const content = (() => { const content = (() => {
try { try {
const normalizedEncoding = encoding?.toLowerCase?.() || ''; const normalizedEncoding = encoding?.toLowerCase?.() || '';
...@@ -47,15 +50,24 @@ export const readFileRawText = async ({ ...@@ -47,15 +50,24 @@ export const readFileRawText = async ({
} }
return buffer.toString('utf-8'); return buffer.toString('utf-8');
} catch (error) { } catch {
return buffer.toString('utf-8'); return buffer.toString('utf-8');
} }
})(); })();
const { text, imageList } = matchMdImg(content); const rawText = await parseMarkdownBase64Images(content, {
controller: (image) => {
if (image.type !== 'base64') return Promise.resolve({ key: '' });
return uploadBase64Image({
mime: image.mime,
base64: image.base64,
uploadFile: options.uploadFile
});
}
});
return { return {
rawText: text, rawText
imageList
}; };
}; };
packages/service/worker/readFile/index.ts
...@@ -7,26 +7,45 @@ import { readDocsFile } from './extension/docx'; ...@@ -7,26 +7,45 @@ import { readDocsFile } from './extension/docx';
import { readPptxRawText } from './extension/pptx'; import { readPptxRawText } from './extension/pptx';
import { readXlsxRawText } from './extension/xlsx'; import { readXlsxRawText } from './extension/xlsx';
import { readCsvRawText } from './extension/csv'; import { readCsvRawText } from './extension/csv';
import { type UploadFileHandler } from './type';
import {
createWorkerUploadFileHandlerWithListener,
isWorkerUploadFileResponse
} from '../utils/uploadFile';
type IncomingMessage = { type IncomingMessage = {
id: string; id: string;
type?: string;
} & Omit<ReadRawTextProps<any>, 'buffer'> & { } & Omit<ReadRawTextProps<any>, 'buffer'> & {
buffer?: ArrayBuffer; buffer?: ArrayBuffer;
sharedBuffer?: SharedArrayBuffer; sharedBuffer?: SharedArrayBuffer;
bufferSize: number; bufferSize: number;
imageKeyOptions?: {
prefix: string;
expiredTime?: Date;
};
}; };
const read = async (params: ReadRawTextByBuffer) => { const read = async (
params: ReadRawTextByBuffer,
options: { uploadFile?: UploadFileHandler } = {}
) => {
switch (params.extension) { switch (params.extension) {
case 'txt': case 'txt':
case 'md': case 'md':
return readFileRawText(params); return readFileRawText(params, {
uploadFile: options.uploadFile
});
case 'html': case 'html':
return readHtmlRawText(params); return readHtmlRawText(params, {
uploadFile: options.uploadFile
});
case 'pdf': case 'pdf':
return readPdfFile(params); return readPdfFile(params);
case 'docx': case 'docx':
return readDocsFile(params); return readDocsFile(params, {
uploadFile: options.uploadFile
});
case 'pptx': case 'pptx':
return readPptxRawText(params); return readPptxRawText(params);
case 'xlsx': case 'xlsx':
...@@ -41,7 +60,19 @@ const read = async (params: ReadRawTextByBuffer) => { ...@@ -41,7 +60,19 @@ const read = async (params: ReadRawTextByBuffer) => {
}; };
parentPort?.on('message', async (props: IncomingMessage) => { parentPort?.on('message', async (props: IncomingMessage) => {
const { id, buffer: transferredBuffer, sharedBuffer, bufferSize, extension, encoding } = props; if (isWorkerUploadFileResponse(props.type)) {
return;
}
const {
id,
buffer: transferredBuffer,
sharedBuffer,
bufferSize,
extension,
encoding,
imageKeyOptions
} = props;
try { try {
const rawBuffer = transferredBuffer ?? sharedBuffer; const rawBuffer = transferredBuffer ?? sharedBuffer;
...@@ -52,9 +83,22 @@ parentPort?.on('message', async (props: IncomingMessage) => { ...@@ -52,9 +83,22 @@ parentPort?.on('message', async (props: IncomingMessage) => {
// 优先使用 transfer 进来的 ArrayBuffer；兼容旧的 SharedArrayBuffer 零拷贝路径。 // 优先使用 transfer 进来的 ArrayBuffer；兼容旧的 SharedArrayBuffer 零拷贝路径。
const buffer = Buffer.from(rawBuffer, 0, bufferSize); const buffer = Buffer.from(rawBuffer, 0, bufferSize);
const data = await read({ extension, encoding, buffer }); const uploadFileHandler = createWorkerUploadFileHandlerWithListener({
taskId: id,
parentPort,
enabled: Boolean(imageKeyOptions?.prefix)
});
try {
const data = await read(
{ extension, encoding, buffer },
{ uploadFile: uploadFileHandler.uploadFile }
);
parentPort?.postMessage({ id, type: 'success', data }); parentPort?.postMessage({ id, type: 'success', data });
} finally {
uploadFileHandler.cleanup();
}
} catch (error) { } catch (error) {
parentPort?.postMessage({ id, type: 'error', data: error }); parentPort?.postMessage({ id, type: 'error', data: error });
} }
......
packages/service/worker/readFile/type.ts
...@@ -6,12 +6,17 @@ export type ReadRawTextProps<T> = { ...@@ -6,12 +6,17 @@ export type ReadRawTextProps<T> = {
export type ReadRawTextByBuffer = ReadRawTextProps<Buffer>; export type ReadRawTextByBuffer = ReadRawTextProps<Buffer>;
export type ImageType = { export type UploadedFileResult = {
uuid: string; key: string;
base64: string; previewUrl?: string;
mime: string;
}; };
export type UploadFileHandler = (data: {
name: string;
mime: string;
buffer: ArrayBuffer;
}) => Promise<UploadedFileResult>;
export type TextItem = { export type TextItem = {
text: string; text: string;
x: number; x: number;
...@@ -34,5 +39,4 @@ export type ParsedPage = { ...@@ -34,5 +39,4 @@ export type ParsedPage = {
export type ReadFileResponse = { export type ReadFileResponse = {
rawText: string; rawText: string;
formatText?: string; formatText?: string;
imageList?: ImageType[];
}; };
packages/service/worker/utils.ts
...@@ -68,9 +68,27 @@ export const runWorker = <T = any>(name: WorkerNameEnum, params?: Record<string, ...@@ -68,9 +68,27 @@ export const runWorker = <T = any>(name: WorkerNameEnum, params?: Record<string,
type WorkerRunTaskType<T> = { type WorkerRunTaskType<T> = {
data: T; data: T;
transferList?: TransferListItem[]; transferList?: TransferListItem[];
handlers?: WorkerRunHandlers;
resolve: (e: any) => void; resolve: (e: any) => void;
reject: (e: any) => void; reject: (e: any) => void;
}; };
export type WorkerUploadFileRequest = {
name: string;
mime: string;
buffer: ArrayBuffer;
};
export type WorkerUploadFileResult = {
key: string;
};
/**
* Worker 任务运行期间可发起的通用主线程能力。
*
* 这些 handler 只服务当前 run 调用，worker 发送中间事件时不会释放任务槽位；
* 只有最终 success/error 消息才会完成任务。
*/
export type WorkerRunHandlers = {
uploadFile?: (data: WorkerUploadFileRequest) => Promise<WorkerUploadFileResult>;
};
type WorkerQueueItem = { type WorkerQueueItem = {
id: string; id: string;
worker: NodeWorker; worker: NodeWorker;
...@@ -78,12 +96,14 @@ type WorkerQueueItem = { ...@@ -78,12 +96,14 @@ type WorkerQueueItem = {
taskTime: number; taskTime: number;
tasksCompleted: number; tasksCompleted: number;
timeoutId?: NodeJS.Timeout; timeoutId?: NodeJS.Timeout;
handlers?: WorkerRunHandlers;
resolve: (e: any) => void; resolve: (e: any) => void;
reject: (e: any) => void; reject: (e: any) => void;
}; };
type WorkerResponse<T = any> = { type WorkerResponse<T = any> = {
id: string; id: string;
type: 'success' | 'error'; type: 'success' | 'error' | 'uploadFile';
requestId?: string;
data: T; data: T;
}; };
...@@ -121,7 +141,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -121,7 +141,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
this.maxTasksPerWorker = maxTasksPerWorker; this.maxTasksPerWorker = maxTasksPerWorker;
} }
private runTask({ data, transferList, resolve, reject }: WorkerRunTaskType<Props>) { private runTask({ data, transferList, handlers, resolve, reject }: WorkerRunTaskType<Props>) {
// Get idle worker or create a new worker // Get idle worker or create a new worker
const runningWorker = (() => { const runningWorker = (() => {
// @ts-ignore // @ts-ignore
...@@ -144,6 +164,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -144,6 +164,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
runningWorker.taskTime = Date.now(); runningWorker.taskTime = Date.now();
runningWorker.resolve = resolve; runningWorker.resolve = resolve;
runningWorker.reject = reject; runningWorker.reject = reject;
runningWorker.handlers = handlers;
runningWorker.timeoutId = setTimeout(() => { runningWorker.timeoutId = setTimeout(() => {
reject('Worker timeout'); reject('Worker timeout');
// 超时即销毁，避免占着 idle 槽位永远不释放 // 超时即销毁，避免占着 idle 槽位永远不释放
...@@ -159,11 +180,11 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -159,11 +180,11 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
); );
} else { } else {
// Not enough worker, push to wait queue // Not enough worker, push to wait queue
this.waitQueue.push({ data, transferList, resolve, reject }); this.waitQueue.push({ data, transferList, handlers, resolve, reject });
} }
} }
run(data: Props, transferList?: TransferListItem[]) { run(data: Props, transferList?: TransferListItem[], handlers?: WorkerRunHandlers) {
return new Promise<Response>((resolve, reject) => { return new Promise<Response>((resolve, reject) => {
/* /*
Whether the task is executed immediately or delayed, the promise callback will dispatch after task complete. Whether the task is executed immediately or delayed, the promise callback will dispatch after task complete.
...@@ -171,6 +192,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -171,6 +192,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
this.runTask({ this.runTask({
data, data,
transferList, transferList,
handlers,
resolve, resolve,
reject reject
}); });
...@@ -195,13 +217,21 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -195,13 +217,21 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
status: 'running', status: 'running',
taskTime: Date.now(), taskTime: Date.now(),
tasksCompleted: 0, tasksCompleted: 0,
handlers: undefined,
resolve: () => {}, resolve: () => {},
reject: () => {} reject: () => {}
}; };
this.workerQueue.push(item); this.workerQueue.push(item);
// watch response // watch response
worker.on('message', ({ id, type, data }: WorkerResponse<Response>) => { worker.on('message', ({ id, type, requestId, data }: WorkerResponse<Response>) => {
if (id !== item.id) return;
if (type === 'uploadFile') {
this.handleUploadFileMessage({ item, requestId, data });
return;
}
if (type === 'success') { if (type === 'success') {
item.resolve(data); item.resolve(data);
} else if (type === 'error') { } else if (type === 'error') {
...@@ -217,6 +247,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -217,6 +247,7 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
this.deleteWorker(item.id); this.deleteWorker(item.id);
} else { } else {
item.status = 'idle'; item.status = 'idle';
item.handlers = undefined;
} }
}); });
...@@ -233,11 +264,55 @@ export class WorkerPool<Props = Record<string, any>, Response = any> { ...@@ -233,11 +264,55 @@ export class WorkerPool<Props = Record<string, any>, Response = any> {
return item; return item;
} }
private handleUploadFileMessage({
item,
requestId,
data
}: {
item: WorkerQueueItem;
requestId?: string;
data: any;
}) {
const reply = (type: 'uploadFileResult' | 'uploadFileError', payload: any) => {
if (!this.workerQueue.includes(item) || item.status !== 'running') return;
try {
item.worker.postMessage({
id: item.id,
type,
requestId,
data: payload
});
} catch (error) {
getLogger(LogCategories.INFRA.WORKER).warn('Failed to reply worker uploadFile request', {
workerId: item.id,
name: this.name,
error
});
}
};
if (!requestId) {
reply('uploadFileError', 'Missing uploadFile requestId');
return;
}
const handler = item.handlers?.uploadFile;
if (!handler) {
reply('uploadFileError', 'Missing uploadFile handler');
return;
}
handler(data)
.then((result) => reply('uploadFileResult', result))
.catch((error) => reply('uploadFileError', error));
}
private deleteWorker(workerId: string) { private deleteWorker(workerId: string) {
const item = this.workerQueue.find((item) => item.id === workerId); const item = this.workerQueue.find((item) => item.id === workerId);
if (item) { if (item) {
item.reject?.('error'); item.reject?.('error');
clearTimeout(item.timeoutId); clearTimeout(item.timeoutId);
item.handlers = undefined;
item.worker.removeAllListeners(); item.worker.removeAllListeners();
item.worker.terminate(); item.worker.terminate();
} }
......
packages/service/worker/utils/base64ImageUpload.ts 0 → 100644
import { type UploadFileHandler } from '../readFile/type';
import { resolveMimeExtension } from '../../common/s3/utils/mime';
import { getLogger, LogCategories } from '../../common/logger';
const logger = getLogger(LogCategories.INFRA.WORKER);
const MAX_PARSED_IMAGE_BUFFER_SIZE = 40 * 1024 * 1024;
export class ParsedImageTooLargeError extends Error {
constructor(size: number, maxSize: number) {
super(`Parsed image too large. Size: ${size} bytes, maximum allowed: ${maxSize} bytes`);
this.name = 'ParsedImageTooLargeError';
}
}
const getBase64DecodedSize = (base64: string) => {
const normalizedBase64 = base64.replace(/\s/g, '');
const padding = normalizedBase64.endsWith('==') ? 2 : normalizedBase64.endsWith('=') ? 1 : 0;
return Math.max(0, Math.floor((normalizedBase64.length * 3) / 4) - padding);
};
const toTransferableArrayBuffer = (buffer: Buffer): ArrayBuffer => {
if (
buffer.buffer instanceof ArrayBuffer &&
buffer.byteOffset === 0 &&
buffer.byteLength === buffer.buffer.byteLength
) {
return buffer.buffer;
}
const imageArrayBuffer = new Uint8Array(buffer.byteLength);
imageArrayBuffer.set(buffer);
return imageArrayBuffer.buffer;
};
export const uploadBase64Image = async ({
mime,
base64,
uploadFile
}: {
mime: string;
base64: string;
uploadFile?: UploadFileHandler;
}) => {
if (!uploadFile) {
logger.warn('Missing image upload handler when parsing document image', { mime });
throw new Error('Missing imageKeyOptions.prefix for parsed document image upload');
}
const decodedSize = getBase64DecodedSize(base64);
if (decodedSize > MAX_PARSED_IMAGE_BUFFER_SIZE) {
throw new ParsedImageTooLargeError(decodedSize, MAX_PARSED_IMAGE_BUFFER_SIZE);
}
const imageBuffer = Buffer.from(base64, 'base64');
const filename = `${crypto.randomUUID()}${resolveMimeExtension(mime)}`;
return uploadFile({
name: filename,
mime,
buffer: toTransferableArrayBuffer(imageBuffer)
}).catch((error) => {
logger.warn('Failed to upload parsed document image from worker', {
filename,
mime,
error
});
throw error;
});
};
packages/service/worker/utils/uploadFile.ts 0 → 100644
import type { MessagePort } from 'worker_threads';
import { type UploadFileHandler, type UploadedFileResult } from '../readFile/type';
type WorkerUploadFileResponse = {
id: string;
type?: string;
requestId?: string;
data?: UploadedFileResult | any;
};
type PendingUploadFileRequest = {
taskId: string;
resolve: (value: UploadedFileResult) => void;
reject: (error: any) => void;
};
const pendingUploadFileRequests = new Map<string, PendingUploadFileRequest>();
export const isWorkerUploadFileResponse = (type?: string) =>
type === 'uploadFileResult' || type === 'uploadFileError';
/**
* 处理 worker uploadFile 的主线程回包。
*
* 多个 worker 入口共用同一套 requestId -> Promise 映射，按 taskId 防串扰。
*/
export const handleWorkerUploadFileResponse = ({
taskId,
type,
requestId,
data
}: {
taskId: string;
type?: string;
requestId?: string;
data?: any;
}) => {
if (!isWorkerUploadFileResponse(type) || !requestId) return false;
const pending = pendingUploadFileRequests.get(requestId);
if (!pending || pending.taskId !== taskId) return true;
pendingUploadFileRequests.delete(requestId);
if (type === 'uploadFileError') {
pending.reject(data);
} else {
pending.resolve(data as UploadedFileResult);
}
return true;
};
export const cleanupWorkerUploadFileRequests = (taskId: string, reason: Error) => {
for (const [requestId, pending] of pendingUploadFileRequests.entries()) {
if (pending.taskId !== taskId) continue;
pending.reject(reason);
pendingUploadFileRequests.delete(requestId);
}
};
/**
* 为 worker 内单个任务创建 uploadFile handler。
*
* 返回的 handler 会向主线程发送 `uploadFile` 中间事件，并等待同 requestId 的结果回包；
* cleanup 必须在任务结束时调用，避免 dangling promise。
*/
export const createWorkerUploadFileHandler = ({
taskId,
parentPort
}: {
taskId: string;
parentPort?: MessagePort | null;
}): {
uploadFile: UploadFileHandler;
cleanup: () => void;
} => {
const uploadFile: UploadFileHandler = (data) =>
new Promise((resolve, reject) => {
const requestId = crypto.randomUUID();
pendingUploadFileRequests.set(requestId, { taskId, resolve, reject });
try {
parentPort?.postMessage(
{
id: taskId,
type: 'uploadFile',
requestId,
data
},
[data.buffer]
);
} catch (error) {
pendingUploadFileRequests.delete(requestId);
reject(error);
}
});
return {
uploadFile,
cleanup: () =>
cleanupWorkerUploadFileRequests(
taskId,
new Error('Worker upload request cancelled before completion')
)
};
};
/**
* 兼容需要在 worker 入口内部监听回包的任务，例如 readFile 会忽略主 message handler 中的回包。
*/
export const createWorkerUploadFileHandlerWithListener = ({
taskId,
parentPort,
enabled
}: {
taskId: string;
parentPort?: MessagePort | null;
enabled: boolean;
}): {
uploadFile?: UploadFileHandler;
cleanup: () => void;
} => {
if (!enabled) return { cleanup: () => {} };
const onMessage = ({ id, type, requestId, data }: WorkerUploadFileResponse) => {
if (id !== taskId) return;
handleWorkerUploadFileResponse({
taskId,
type,
requestId,
data
});
};
parentPort?.on('message', onMessage);
const bridge = createWorkerUploadFileHandler({ taskId, parentPort });
return {
uploadFile: bridge.uploadFile,
cleanup: () => {
parentPort?.off('message', onMessage);
bridge.cleanup();
}
};
};
projects/app/src/pages/api/core/dataset/collection/create/images.ts
...@@ -74,7 +74,7 @@ async function handler(req: ApiRequestProps): Promise<CreateCollectionWithResult ...@@ -74,7 +74,7 @@ async function handler(req: ApiRequestProps): Promise<CreateCollectionWithResult
const filename = path.basename(file.filename); const filename = path.basename(file.filename);
const { fileKey } = getFileS3Key.dataset({ datasetId, filename }); const { fileKey } = getFileS3Key.dataset({ datasetId, filename });
return uploadImage2S3Bucket('private', { return uploadImage2S3Bucket('private', {
base64Img: (await fs.promises.readFile(file.path)).toString('base64'), buffer: await fs.promises.readFile(file.path),
uploadKey: fileKey, uploadKey: fileKey,
mimetype: file.mimetype, mimetype: file.mimetype,
filename, filename,
......
projects/app/src/pages/api/core/dataset/data/insertImages.ts
...@@ -67,7 +67,7 @@ async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> { ...@@ -67,7 +67,7 @@ async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> {
const imageIds = await Promise.all( const imageIds = await Promise.all(
result.fileMetadata.map(async (file) => result.fileMetadata.map(async (file) =>
uploadImage2S3Bucket('private', { uploadImage2S3Bucket('private', {
base64Img: (await fs.promises.readFile(file.path)).toString('base64'), buffer: await fs.promises.readFile(file.path),
uploadKey: getFileS3Key.dataset({ uploadKey: getFileS3Key.dataset({
datasetId: dataset._id, datasetId: dataset._id,
filename: path.basename(file.filename) filename: path.basename(file.filename)
......
projects/app/test/api/core/dataset/collection/create/images.test.ts
...@@ -162,6 +162,14 @@ describe('POST /api/core/dataset/collection/create/images', () => { ...@@ -162,6 +162,14 @@ describe('POST /api/core/dataset/collection/create/images', () => {
trainingType: DatasetCollectionDataProcessModeEnum.chunk trainingType: DatasetCollectionDataProcessModeEnum.chunk
} }
}); });
expect(mockUploadImage2S3Bucket).toHaveBeenCalledWith('private', {
buffer: Buffer.from('image-bytes'),
uploadKey: 'dataset/team/cat.png',
mimetype: 'image/png',
filename: 'cat.png',
expiredTime: expect.any(Date)
});
expect(mockUploadImage2S3Bucket.mock.calls[0][1]).not.toHaveProperty('base64Img');
expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']); expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']);
}); });
}); });
projects/app/test/api/core/dataset/data/insertImages.test.ts
...@@ -171,6 +171,14 @@ describe('POST /api/core/dataset/data/insertImages', () => { ...@@ -171,6 +171,14 @@ describe('POST /api/core/dataset/data/insertImages', () => {
data: [{ imageId: 'dataset/team/cat.png' }], data: [{ imageId: 'dataset/team/cat.png' }],
session: 'session' session: 'session'
}); });
expect(mockUploadImage2S3Bucket).toHaveBeenCalledWith('private', {
buffer: Buffer.from('image-bytes'),
uploadKey: 'dataset/team/cat.png',
mimetype: 'image/png',
filename: 'cat.png',
expiredTime: expect.any(Date)
});
expect(mockUploadImage2S3Bucket.mock.calls[0][1]).not.toHaveProperty('base64Img');
expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']); expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']);
}); });
......
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