Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit a7394b73 by YeYuheng Committed by GitHub
Options

fix: guard embedding token limits (#7269) 

* fix: guard embedding token limits

* perf: split

* feat: test

* fix: avoid header-only table chunks

* chore: hide vector model chunk limit label

* fix: harden embedding token limit handling

---------

Co-authored-by: Archer <545436317@qq.com>
Co-authored-by: Finley Ge <finleyge@fastgpt.io>
parent a76f080e
Showing with 1058 additions and 99 deletions
packages/global/core/dataset/training/utils.ts
...@@ -10,6 +10,7 @@ import type { ChunkSettingsType } from '../type'; ...@@ -10,6 +10,7 @@ import type { ChunkSettingsType } from '../type';
import { cloneDeep } from 'lodash'; import { cloneDeep } from 'lodash';
export const minChunkSize = 64; // min index and chunk size export const minChunkSize = 64; // min index and chunk size
export const maxPreviewChunkCount = 50_000;
// Chunk size // Chunk size
export const chunkAutoChunkSize = 1000; export const chunkAutoChunkSize = 1000;
......
packages/global/openapi/core/dataset/file/api.ts
...@@ -3,12 +3,57 @@ import { ObjectIdSchema } from '../../../../common/type/mongo'; ...@@ -3,12 +3,57 @@ import { ObjectIdSchema } from '../../../../common/type/mongo';
import { DatasetSourceReadTypeEnum } from '../../../../core/dataset/constants'; import { DatasetSourceReadTypeEnum } from '../../../../core/dataset/constants';
import { ChunkSettingsSchema } from '../../../../core/dataset/type'; import { ChunkSettingsSchema } from '../../../../core/dataset/type';
import { CreatePostPresignedUrlResponseSchema } from '../../../../common/file/s3/type'; import { CreatePostPresignedUrlResponseSchema } from '../../../../common/file/s3/type';
import { IntSchema, NumSchema } from '../../../../common/zod';
import { minChunkSize } from '../../../../core/dataset/training/utils';
const PreviewChunkSplitterSchema = z
.string()
.max(200)
.superRefine((value, ctx) => {
if (!value) return;
const separators = value.replace(/\\n/g, '\n').split('|');
if (separators.length > 10) {
ctx.addIssue({
code: 'custom',
message: 'Custom chunk splitter supports at most 10 separators'
});
}
if (separators.some((item) => item.length === 0)) {
ctx.addIssue({
code: 'custom',
message: 'Custom chunk separators cannot be empty'
});
}
})
.meta({
example: '\\n\\n|---',
description: '自定义最高优先分隔符，使用 | 分隔多个非空项，最多 10 项、200 字符'
});
const PreviewChunkSettingsSchema = ChunkSettingsSchema.extend({
chunkTriggerMinSize: IntSchema.optional().meta({ description: '分块触发最小大小' }),
paragraphChunkDeep: IntSchema.max(8).optional().meta({ description: '段落分块深度，最大 8' }),
paragraphChunkMinSize: IntSchema.optional().meta({ description: '段落分块最小大小' }),
chunkSize: IntSchema.min(minChunkSize)
.optional()
.meta({
example: 512,
description: `分块大小，最小 ${minChunkSize}`
}),
chunkSplitter: PreviewChunkSplitterSchema.optional()
});
const PreviewOverlapRatioSchema = NumSchema.min(0).max(0.4).meta({
example: 0.2,
description: '分块重叠比例，范围 0-0.4'
});
/* ============================================================================ /* ============================================================================
* API: 预览文件分块 * API: 预览文件分块
* Route: POST /api/core/dataset/file/getPreviewChunks * Route: POST /api/core/dataset/file/getPreviewChunks
* ============================================================================ */ * ============================================================================ */
export const GetPreviewChunksBodySchema = ChunkSettingsSchema.extend({ export const GetPreviewChunksBodySchema = PreviewChunkSettingsSchema.extend({
datasetId: ObjectIdSchema.meta({ datasetId: ObjectIdSchema.meta({
example: '68ad85a7463006c963799a05', example: '68ad85a7463006c963799a05',
description: '知识库 ID' description: '知识库 ID'
...@@ -24,10 +69,7 @@ export const GetPreviewChunksBodySchema = ChunkSettingsSchema.extend({ ...@@ -24,10 +69,7 @@ export const GetPreviewChunksBodySchema = ChunkSettingsSchema.extend({
customPdfParse: z.boolean().optional().meta({ customPdfParse: z.boolean().optional().meta({
description: '是否启用自定义 PDF 解析' description: '是否启用自定义 PDF 解析'
}), }),
overlapRatio: z.number().meta({ overlapRatio: PreviewOverlapRatioSchema,
example: 0.2,
description: '分块重叠比例'
}),
selector: z.string().optional().meta({ selector: z.string().optional().meta({
example: 'body', example: 'body',
description: '网页抓取的 CSS 选择器' description: '网页抓取的 CSS 选择器'
...@@ -55,6 +97,31 @@ export const GetPreviewChunksResponseSchema = z.object({ ...@@ -55,6 +97,31 @@ export const GetPreviewChunksResponseSchema = z.object({
export type GetPreviewChunksResponse = z.infer<typeof GetPreviewChunksResponseSchema>; export type GetPreviewChunksResponse = z.infer<typeof GetPreviewChunksResponseSchema>;
/* ============================================================================ /* ============================================================================
* API: 预览原始文本分块
* Route: POST /api/core/dataset/file/getRawTextPreviewChunks
* Method: POST
* Description: 对前端已读取到的原始文本执行后端分块预览，用于 fileCustom 导入预览
* Tags: ['Dataset', 'File', 'Read']
* ============================================================================ */
export const GetRawTextPreviewChunksBodySchema = PreviewChunkSettingsSchema.extend({
datasetId: ObjectIdSchema.meta({
example: '68ad85a7463006c963799a05',
description: '知识库 ID'
}),
rawText: z
.string()
.max(10 * 1024 * 1024)
.meta({
example: '# 产品文档\n\n这是待预览分块的原始文本',
description: '前端已读取到的原始文本，最多 10 MiB 字符'
}),
overlapRatio: PreviewOverlapRatioSchema
});
export type GetRawTextPreviewChunksBody = z.infer<typeof GetRawTextPreviewChunksBodySchema>;
export type GetRawTextPreviewChunksResponse = z.infer<typeof GetPreviewChunksResponseSchema>;
/* ============================================================================
* API: 获取知识库文件上传预签名 URL * API: 获取知识库文件上传预签名 URL
* Route: POST /api/core/dataset/file/presignDatasetFilePostUrl * Route: POST /api/core/dataset/file/presignDatasetFilePostUrl
* ============================================================================ */ * ============================================================================ */
......
packages/global/openapi/core/dataset/file/index.ts
...@@ -5,6 +5,7 @@ import { ...@@ -5,6 +5,7 @@ import {
GetSearchTestImagePreviewUrlsResponseSchema, GetSearchTestImagePreviewUrlsResponseSchema,
GetPreviewChunksBodySchema, GetPreviewChunksBodySchema,
GetPreviewChunksResponseSchema, GetPreviewChunksResponseSchema,
GetRawTextPreviewChunksBodySchema,
PresignDatasetFilePostUrlBodySchema, PresignDatasetFilePostUrlBodySchema,
PresignDatasetFilePostUrlResponseSchema, PresignDatasetFilePostUrlResponseSchema,
PresignSearchTestImageBodySchema, PresignSearchTestImageBodySchema,
...@@ -36,6 +37,30 @@ export const DatasetFilePath: OpenAPIPath = { ...@@ -36,6 +37,30 @@ export const DatasetFilePath: OpenAPIPath = {
} }
} }
}, },
'/core/dataset/file/getRawTextPreviewChunks': {
post: {
summary: '预览原始文本分块',
description: '对前端已读取到的原始文本执行后端分块预览，用于自定义文件导入预览',
tags: [DevApiTagsMap.datasetFile],
requestBody: {
content: {
'application/json': {
schema: GetRawTextPreviewChunksBodySchema
}
}
},
responses: {
200: {
description: '成功返回预览分块列表及总数',
content: {
'application/json': {
schema: GetPreviewChunksResponseSchema
}
}
}
}
}
},
'/core/dataset/file/presignDatasetFilePostUrl': { '/core/dataset/file/presignDatasetFilePostUrl': {
post: { post: {
summary: '获取知识库文件上传预签名 URL', summary: '获取知识库文件上传预签名 URL',
......
packages/service/core/ai/embedding/index.ts
import { type EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.schema'; import { type EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.schema';
import { getAIApi } from '../config'; import { getAIApi } from '../config';
import { countPromptTokens } from '../../../common/string/tiktoken/index'; import { countPromptTokens, countPromptTokensBatch } from '../../../common/string/tiktoken/index';
import { EmbeddingTypeEnm } from '@fastgpt/global/core/ai/constants'; import { EmbeddingTypeEnm } from '@fastgpt/global/core/ai/constants';
import { retryFn } from '@fastgpt/global/common/system/utils'; import { retryFn } from '@fastgpt/global/common/system/utils';
import { getLogger, LogCategories } from '../../../common/logger'; import { getLogger, LogCategories } from '../../../common/logger';
import z from 'zod'; import z from 'zod';
import { truncateTextByFormattedTokenLimit } from './tokenLimit';
const logger = getLogger(LogCategories.MODULE.AI.EMBEDDING); const logger = getLogger(LogCategories.MODULE.AI.EMBEDDING);
...@@ -43,13 +44,43 @@ const countInputTokens = async (input: GetVectorInputItem) => { ...@@ -43,13 +44,43 @@ const countInputTokens = async (input: GetVectorInputItem) => {
}; };
export async function getVectors({ model, inputs: rawInputs, type, headers }: GetVectorsProps) { export async function getVectors({ model, inputs: rawInputs, type, headers }: GetVectorsProps) {
const inputs = z const validatedInputs = z
.array(InputItemSchema) .array(InputItemSchema)
.parse(rawInputs) .parse(rawInputs)
.map((item) => ({ .map((item) => ({
...item, ...item,
input: item.input.trim() input: item.input.trim()
})); }));
if (validatedInputs.length === 0 || validatedInputs.some((item) => !item.input)) {
return Promise.reject({
code: 500,
message: 'input is empty'
});
}
const textInputs = validatedInputs
.filter((item) => item.type === 'text')
.map((item) => item.input);
const textTokenCounts = textInputs.length > 0 ? await countPromptTokensBatch(textInputs) : [];
let textIndex = 0;
const inputs = await Promise.all(
validatedInputs.map(async (item) => {
const currentTokens = item.type === 'text' ? textTokenCounts[textIndex++] : undefined;
// getVectors 是所有 embedding 请求的最后入口。这里仅对 text 做单条截断兜底，
// 不做拆分；知识库入库这类需要保留完整内容的场景，应在上游先拆成多条 index。
return {
...item,
input:
item.type === 'text'
? await truncateTextByFormattedTokenLimit({
text: item.input,
maxToken: model.maxToken,
currentTokens
})
: item.input
};
})
);
if (inputs.length === 0 || inputs.some((item) => !item.input)) { if (inputs.length === 0 || inputs.some((item) => !item.input)) {
return Promise.reject({ return Promise.reject({
code: 500, code: 500,
......
packages/service/core/ai/embedding/tokenLimit.ts 0 → 100644
import { countPromptTokens } from '../../../common/string/tiktoken/index';
/**
* 按格式化后的文本 token 上限，从原始文本里二分出最长安全前缀。
*
* 这个函数只做“单条输入截断”，不会把一条文本拆成多条文本。它主要用于
* embedding query 这类不能扩增输入数量的场景；知识库入库索引需要保留内容时，
* 应该在上游按 token 分块生成多条 index。
*
* `formatText` 用于处理“实际送入 embedding 的文本并不等于原文”的场景，
* 例如知识库索引会给正文补充集合标题前缀。这里仍只返回原文前缀，由调用方决定如何组装最终文本。
*/
export const truncateTextByFormattedTokenLimit = async ({
text,
maxToken,
formatText = (text) => text,
currentTokens
}: {
text: string;
maxToken: number;
formatText?: (text: string) => string;
currentTokens?: number;
}) => {
const trimmedText = text.trim();
if (!Number.isFinite(maxToken) || maxToken <= 0) return trimmedText;
const formattedTokens = currentTokens ?? (await countPromptTokens(formatText(trimmedText)));
if (!trimmedText || formattedTokens <= maxToken) {
return trimmedText;
}
const textChars = Array.from(trimmedText);
let left = 1;
let right = textChars.length;
let bestEnd = 0;
while (left <= right) {
const mid = Math.floor((left + right) / 2);
const candidate = textChars.slice(0, mid).join('').trim();
if (!candidate) {
left = mid + 1;
continue;
}
if ((await countPromptTokens(formatText(candidate))) <= maxToken) {
bestEnd = mid;
left = mid + 1;
} else {
right = mid - 1;
}
}
return textChars.slice(0, bestEnd).join('').trim();
};
packages/service/core/dataset/read.ts
...@@ -3,7 +3,7 @@ import { ...@@ -3,7 +3,7 @@ import {
DatasetSourceReadTypeEnum DatasetSourceReadTypeEnum
} from '@fastgpt/global/core/dataset/constants'; } from '@fastgpt/global/core/dataset/constants';
import { urlsFetch } from '../../common/string/cheerio'; import { urlsFetch } from '../../common/string/cheerio';
import { type TextSplitProps } from '@fastgpt/global/common/string/textSplitter'; import { type TextSplitProps } from '../../common/string/textSplitter';
import { axios } from '../../common/api/axios'; import { axios } from '../../common/api/axios';
import { readFileContentByBuffer } from '../../common/file/read/utils'; import { readFileContentByBuffer } from '../../common/file/read/utils';
import { parseFileExtensionFromUrl } from '@fastgpt/global/common/string/tools'; import { parseFileExtensionFromUrl } from '@fastgpt/global/common/string/tools';
......
packages/global/test/common/string/textSplitter.test.ts → packages/service/test/common/string/textSplitter.test.ts
import { it, expect } from 'vitest'; // 必须显式导入 import { it, expect } from 'vitest'; // 必须显式导入
import { splitText2Chunks } from '@fastgpt/global/common/string/textSplitter'; import {
getMaxPrefixByLength,
getMaxSuffixByLength,
splitText2Chunks
} from '@fastgpt/service/common/string/textSplitter';
import { countPromptTokensInWorker } from '@fastgpt/service/worker/countGptMessagesTokens/count';
import fs from 'fs'; import fs from 'fs';
const simpleChunks = (chunks: string[]) => { const simpleChunks = (chunks: string[]) => {
...@@ -1042,3 +1047,222 @@ it(`Test splitText2Chunks 14 - lastText not lost when strategies exhausted`, () ...@@ -1042,3 +1047,222 @@ it(`Test splitText2Chunks 14 - lastText not lost when strategies exhausted`, ()
expect(chunk.length).toBeGreaterThan(0); expect(chunk.length).toBeGreaterThan(0);
}); });
}); });
it(`Test splitText2Chunks 15 - token mode should not append table header only chunk`, () => {
const header = `| id | payload |
| --- | --- |
`;
const text = `${header}| 1 | ${'𠮷'.repeat(20)} |
`;
const chunkSize = countPromptTokensInWorker(header) + 8;
const { chunks } = splitText2Chunks({
text,
chunkSize,
maxSize: chunkSize,
overlapRatio: 0,
lengthUnit: 'token'
});
expect(chunks.length).toBeGreaterThan(0);
expect(chunks).not.toContain('| id | payload |\n| --- | --- |');
expect(chunks.join('\n')).toContain('𠮷');
});
it(`Test splitText2Chunks 15.1 - should not create markdown table header-only chunk`, () => {
const { chunks, chars } = splitText2Chunks({
text: `| id | payload | note |
| --- | --- | --- |`,
chunkSize: 40,
maxSize: 200,
overlapRatio: 0
});
expect(chunks).toEqual([]);
expect(chars).toBe(0);
});
it(`Test splitText2Chunks 15.2 - char mode should not append table header only chunk`, () => {
const header = `| id | payload |
| --- | --- |
`;
const text = `${header}| 1 | ${'a'.repeat(80)} |
| 2 | normal |
`;
const { chunks } = splitText2Chunks({
text,
chunkSize: header.length + 20,
maxSize: 200,
overlapRatio: 0
});
expect(chunks.length).toBeGreaterThan(0);
expect(chunks).not.toContain('| id | payload |\n| --- | --- |');
expect(chunks.join('\n')).toContain('| 1 |');
});
it(`Test splitText2Chunks 16 - token mode table chunks should include header within limit`, () => {
const header = `| id | payload |
| --- | --- |
`;
const text = `${header}| 1 | ${'a'.repeat(40)} |
`;
const chunkSize = countPromptTokensInWorker(header) + 4;
const { chunks } = splitText2Chunks({
text,
chunkSize,
maxSize: chunkSize,
overlapRatio: 0,
lengthUnit: 'token'
});
expect(chunks.length).toBeGreaterThan(1);
expect(chunks.every((chunk) => chunk.includes('| id | payload |'))).toBe(true);
expect(chunks.every((chunk) => countPromptTokensInWorker(chunk) <= chunkSize)).toBe(true);
});
it(`Test splitText2Chunks 17 - token mode table should fail when header has no content budget`, () => {
const text = `| id | payload |
| --- | --- |
| 1 | a |
`;
expect(() =>
splitText2Chunks({
text,
chunkSize: 5,
maxSize: 5,
overlapRatio: 0,
lengthUnit: 'token'
})
).toThrow('Markdown table header exceeds token chunk size');
});
it(`Test getMaxPrefixByLength - returns the longest prefix within custom length limit`, () => {
const countLength = (text: string) => Array.from(text).length;
expect(
getMaxPrefixByLength({
text: 'A𠮷BC',
maxLength: 2,
countLength
})
).toBe('A𠮷');
expect(
getMaxPrefixByLength({
text: 'A𠮷BC',
maxLength: 10,
countLength
})
).toBe('A𠮷BC');
});
it(`Test getMaxPrefixByLength - returns empty when no code point fits`, () => {
const countLength = (text: string) => Array.from(text).length * 2;
expect(
getMaxPrefixByLength({
text: '𠮷',
maxLength: 1,
countLength
})
).toBe('');
});
it(`Test getMaxSuffixByLength - returns the longest suffix within custom length limit`, () => {
const countLength = (text: string) => Array.from(text).length;
expect(
getMaxSuffixByLength({
text: 'AB𠮷C',
maxLength: 2,
countLength
})
).toBe('𠮷C');
expect(
getMaxSuffixByLength({
text: 'AB𠮷C',
maxLength: 10,
countLength
})
).toBe('AB𠮷C');
});
it(`Test getMaxSuffixByLength - returns empty when overlap budget is unavailable`, () => {
const countLength = (text: string) => Array.from(text).length;
expect(
getMaxSuffixByLength({
text: 'AB𠮷C',
maxLength: 0,
countLength
})
).toBe('');
});
it(`Test getMaxPrefixByLength - does not tokenize the complete remainder for a small limit`, () => {
const measuredLengths: number[] = [];
const result = getMaxPrefixByLength({
text: 'a'.repeat(10_000),
maxLength: 10,
countLength: (text) => {
measuredLengths.push(text.length);
return text.length;
}
});
expect(result).toBe('a'.repeat(10));
expect(Math.max(...measuredLengths)).toBeLessThanOrEqual(20);
});
it.each(['|', 'prefix|', '|suffix', 'prefix||suffix'])(
'Test splitText2Chunks - rejects empty custom separators: %s',
(customReg) => {
expect(() =>
splitText2Chunks({
text: 'safe text',
chunkSize: 64,
customReg: [customReg]
})
).toThrow('Custom split separators cannot be empty');
}
);
it('Test splitText2Chunks - rejects an overlap ratio that cannot advance', () => {
expect(() =>
splitText2Chunks({
text: 'a'.repeat(100),
chunkSize: 64,
maxSize: 64,
overlapRatio: 1
})
).toThrow('Overlap ratio must be greater than or equal to 0 and less than 1');
});
it('Test splitText2Chunks - rejects work beyond the configured chunk limit', () => {
expect(() =>
splitText2Chunks({
text: 'a'.repeat(1_000),
chunkSize: 64,
maxSize: 64,
overlapRatio: 0,
maxChunks: 5
})
).toThrow('Text split exceeds the maximum chunk count of 5');
});
it('Test splitText2Chunks - rejects high-frequency custom separators before splitting', () => {
expect(() =>
splitText2Chunks({
text: 'a'.repeat(100_000),
chunkSize: 64,
customReg: ['a'],
maxChunks: 5
})
).toThrow('Text split exceeds the maximum chunk count of 5');
});
packages/service/test/core/ai/embedding/index.test.ts
...@@ -7,11 +7,15 @@ import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; ...@@ -7,11 +7,15 @@ import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
// We control the embeddings.create implementation per-test via `mockCreate`. // We control the embeddings.create implementation per-test via `mockCreate`.
const mockCreate = vi.fn(); const mockCreate = vi.fn();
const mockCountPromptTokens = vi.hoisted(() => vi.fn(async (text: string) => text.length)); const mockCountPromptTokens = vi.hoisted(() => vi.fn(async (text: string) => text.length));
const mockCountPromptTokensBatch = vi.hoisted(() =>
vi.fn(async (texts: string[]) => texts.map((text) => text.length))
);
// getVectors 在缺少 usage 时会回退本地 token 计数；测试里只验证回退路径生效， // getVectors 在缺少 usage 时会回退本地 token 计数；测试里只验证回退路径生效，
// 不启动真实 worker，避免 service 包单测依赖 app/pro 的 worker 构建目录。 // 不启动真实 worker，避免 service 包单测依赖 app/pro 的 worker 构建目录。
vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({ vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({
countPromptTokens: mockCountPromptTokens countPromptTokens: mockCountPromptTokens,
countPromptTokensBatch: mockCountPromptTokensBatch
})); }));
vi.mock('@fastgpt/service/core/ai/config', () => ({ vi.mock('@fastgpt/service/core/ai/config', () => ({
...@@ -378,7 +382,12 @@ describe('getVectors function test', () => { ...@@ -378,7 +382,12 @@ describe('getVectors function test', () => {
beforeEach(() => { beforeEach(() => {
mockCreate.mockReset(); mockCreate.mockReset();
mockCountPromptTokens.mockClear();
mockCountPromptTokensBatch.mockClear();
mockCountPromptTokens.mockImplementation(async (text: string) => text.length); mockCountPromptTokens.mockImplementation(async (text: string) => text.length);
mockCountPromptTokensBatch.mockImplementation(async (texts: string[]) =>
texts.map((text) => text.length)
);
}); });
const buildModel = (overrides: Partial<EmbeddingModelItemType> = {}): EmbeddingModelItemType => const buildModel = (overrides: Partial<EmbeddingModelItemType> = {}): EmbeddingModelItemType =>
...@@ -386,6 +395,7 @@ describe('getVectors function test', () => { ...@@ -386,6 +395,7 @@ describe('getVectors function test', () => {
model: 'text-embedding-3-small', model: 'text-embedding-3-small',
name: 'text-embedding-3-small', name: 'text-embedding-3-small',
batchSize: 10, batchSize: 10,
maxToken: 8192,
normalization: false, normalization: false,
...overrides ...overrides
}) as EmbeddingModelItemType; }) as EmbeddingModelItemType;
...@@ -420,6 +430,51 @@ describe('getVectors function test', () => { ...@@ -420,6 +430,51 @@ describe('getVectors function test', () => {
}); });
expect(mockCreate).not.toHaveBeenCalled(); expect(mockCreate).not.toHaveBeenCalled();
}); });
it('should embed image inputs normally', async () => {
mockCreate.mockResolvedValue(
makeResponse([[0.1, 0.2, 0.3, 0.4]], { usage: { total_tokens: 1 } })
);
const result = await getVectors({
model: buildModel({ maxToken: 1 }),
inputs: [imageInput('data:image/png;base64,aaa')]
});
expect(mockCreate).toHaveBeenCalledTimes(1);
expect(result.vectors).toHaveLength(1);
});
it('should truncate text inputs by model maxToken before requesting embeddings', async () => {
mockCreate.mockResolvedValue(
makeResponse([[0.1, 0.2, 0.3, 0.4]], { usage: { total_tokens: 1 } })
);
await getVectors({
model: buildModel({ maxToken: 12 }),
inputs: [textInput('abcdefghijklmnopqrstuvwxy')]
});
expect(mockCreate).toHaveBeenCalledTimes(1);
expect(mockCreate.mock.calls[0][0].input).toEqual(['abcdefghijkl']);
});
it('should keep astral Unicode characters well formed when truncating', async () => {
mockCreate.mockResolvedValue(
makeResponse([[0.1, 0.2, 0.3, 0.4]], { usage: { total_tokens: 1 } })
);
mockCountPromptTokens.mockImplementation(async (text: string) => text.length);
mockCountPromptTokensBatch.mockResolvedValueOnce([3]);
await getVectors({
model: buildModel({ maxToken: 2 }),
inputs: [textInput('a𠮷')]
});
const providerInput = mockCreate.mock.calls[0][0].input[0] as string;
expect(providerInput).toBe('a');
expect(providerInput).not.toMatch(/[\uD800-\uDFFF]/u);
});
}); });
describe('basic embedding calls', () => { describe('basic embedding calls', () => {
...@@ -543,6 +598,7 @@ describe('getVectors function test', () => { ...@@ -543,6 +598,7 @@ describe('getVectors function test', () => {
]); ]);
expect(result.tokens).toBe(6); expect(result.tokens).toBe(6);
expect(result.vectors).toHaveLength(2); expect(result.vectors).toHaveLength(2);
expect(mockCountPromptTokensBatch).not.toHaveBeenCalled();
}); });
it('should build mixed text and image input parts in order', async () => { it('should build mixed text and image input parts in order', async () => {
......
packages/service/test/core/ai/hooks/useTextCosine.test.ts
import { describe, expect, it, vi, beforeEach } from 'vitest'; import { describe, expect, it, vi, beforeEach } from 'vitest';
const mockCountPromptTokens = vi.hoisted(() => vi.fn(async (text: string) => text.length));
vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({
countPromptTokens: mockCountPromptTokens
}));
import { useTextCosine } from '@fastgpt/service/core/ai/hooks/useTextCosine'; import { useTextCosine } from '@fastgpt/service/core/ai/hooks/useTextCosine';
import { getEmbeddingModel } from '@fastgpt/service/core/ai/model';
import { import {
generateMockEmbedding, generateMockEmbedding,
createMockVectorsResponse, createMockVectorsResponse,
...@@ -11,6 +19,12 @@ import { ...@@ -11,6 +19,12 @@ import {
describe('useTextCosine', () => { describe('useTextCosine', () => {
beforeEach(() => { beforeEach(() => {
vi.clearAllMocks(); vi.clearAllMocks();
mockCountPromptTokens.mockImplementation(async (text: string) => text.length);
vi.mocked(getEmbeddingModel).mockReturnValue({
model: 'text-embedding-ada-002',
name: 'text-embedding-ada-002',
maxToken: 100
} as any);
}); });
describe('lazyGreedyQuerySelection', () => { describe('lazyGreedyQuerySelection', () => {
...@@ -168,6 +182,46 @@ describe('useTextCosine', () => { ...@@ -168,6 +182,46 @@ describe('useTextCosine', () => {
expect(result.selectedData).toEqual(['candidate']); expect(result.selectedData).toEqual(['candidate']);
}); });
it('should pass overlong query and candidates to centralized embedding fallback', async () => {
vi.mocked(getEmbeddingModel).mockReturnValue({
model: 'mock-embedding-model',
name: 'Mock Embedding Model',
maxToken: 12
} as any);
mockGetVectors.mockResolvedValueOnce({
tokens: 10,
vectors: [
generateMockEmbedding('abcdefghijklmnopqrstuvwxy'),
generateMockEmbedding('klmnopqrstuvwxy')
]
});
const { lazyGreedyQuerySelection } = useTextCosine({ embeddingModel: 'custom-model' });
const result = await lazyGreedyQuerySelection({
originalText: 'abcdefghijklmnopqrstuvwxy',
candidates: ['klmnopqrstuvwxy'],
k: 1
});
expect(mockGetVectors).toHaveBeenCalledWith({
model: expect.objectContaining({
model: 'mock-embedding-model'
}),
inputs: [
{
type: 'text',
input: 'abcdefghijklmnopqrstuvwxy'
},
{
type: 'text',
input: 'klmnopqrstuvwxy'
}
],
type: 'query'
});
expect(result.selectedData).toEqual(['klmnopqrstuvwxy']);
});
it('should handle identical candidates correctly', async () => { it('should handle identical candidates correctly', async () => {
const originalVector = generateMockEmbedding('original'); const originalVector = generateMockEmbedding('original');
const identicalVector = generateMockEmbedding('same'); const identicalVector = generateMockEmbedding('same');
......
packages/service/test/core/dataset/search/defaultRecall.test.ts
...@@ -13,6 +13,7 @@ const mockMongoDatasetCollectionFind = vi.hoisted(() => vi.fn()); ...@@ -13,6 +13,7 @@ const mockMongoDatasetCollectionFind = vi.hoisted(() => vi.fn());
const mockMongoDatasetDataFind = vi.hoisted(() => vi.fn()); const mockMongoDatasetDataFind = vi.hoisted(() => vi.fn());
const mockMongoDatasetDataTextAggregate = vi.hoisted(() => vi.fn()); const mockMongoDatasetDataTextAggregate = vi.hoisted(() => vi.fn());
const mockGetImageBase64 = vi.hoisted(() => vi.fn()); const mockGetImageBase64 = vi.hoisted(() => vi.fn());
const mockCountPromptTokens = vi.hoisted(() => vi.fn(async (prompt: string) => prompt.length));
const mockCountPromptTokensBatch = vi.hoisted(() => const mockCountPromptTokensBatch = vi.hoisted(() =>
vi.fn(async (prompts: string[]) => prompts.map((prompt) => prompt.length)) vi.fn(async (prompts: string[]) => prompts.map((prompt) => prompt.length))
); );
...@@ -45,6 +46,7 @@ vi.mock('@fastgpt/service/common/file/image/utils', () => ({ ...@@ -45,6 +46,7 @@ vi.mock('@fastgpt/service/common/file/image/utils', () => ({
// defaultRecall 的结果过滤只关心 token 数的相对大小，测试里用稳定 mock // defaultRecall 的结果过滤只关心 token 数的相对大小，测试里用稳定 mock
// 隔离真实 worker 路径，避免单元测试依赖 app/pro 的 worker 构建产物。 // 隔离真实 worker 路径，避免单元测试依赖 app/pro 的 worker 构建产物。
vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({ vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({
countPromptTokens: mockCountPromptTokens,
countPromptTokensBatch: mockCountPromptTokensBatch countPromptTokensBatch: mockCountPromptTokensBatch
})); }));
...@@ -81,10 +83,12 @@ describe('default recall dataset search', () => { ...@@ -81,10 +83,12 @@ describe('default recall dataset search', () => {
mockCountPromptTokensBatch.mockImplementation(async (prompts: string[]) => mockCountPromptTokensBatch.mockImplementation(async (prompts: string[]) =>
prompts.map((prompt) => prompt.length) prompts.map((prompt) => prompt.length)
); );
mockCountPromptTokens.mockImplementation(async (prompt: string) => prompt.length);
mockGetEmbeddingModel.mockReturnValue({ mockGetEmbeddingModel.mockReturnValue({
model: 'mock-embedding-model', model: 'mock-embedding-model',
name: 'Mock Embedding Model' name: 'Mock Embedding Model',
maxToken: 100
}); });
mockGetDefaultRerankModel.mockReturnValue(undefined); mockGetDefaultRerankModel.mockReturnValue(undefined);
mockGetLLMModel.mockReturnValue({ mockGetLLMModel.mockReturnValue({
...@@ -257,6 +261,45 @@ describe('default recall dataset search', () => { ...@@ -257,6 +261,45 @@ describe('default recall dataset search', () => {
); );
}); });
it('should pass overlong text queries to centralized embedding fallback without creating extra queries', async () => {
mockGetLLMModel.mockReturnValue(undefined);
mockIsImageEmbeddingModel.mockReturnValue(false);
mockGetEmbeddingModel.mockReturnValueOnce({
model: 'mock-embedding-model',
name: 'Mock Embedding Model',
maxToken: 12
});
mockGetVectors.mockImplementationOnce(async ({ inputs }) => ({
tokens: 10,
vectors: inputs.map((_: unknown, index: number) => [index + 1])
}));
await searchDatasetData({
histories: [],
teamId: 'team-1',
model: 'mock-embedding-model',
datasetIds: ['dataset-1'],
reRankQuery: 'abcdefghijklmnopqrstuvwxy',
textQueries: ['abcdefghijklmnopqrstuvwxy'],
imageQueries: [],
limit: 5000,
searchMode: DatasetSearchModeEnum.embedding,
embeddingWeight: 0.5,
usingReRank: false
});
expect(mockGetVectors).toHaveBeenCalledWith(
expect.objectContaining({
inputs: [
{
type: 'text',
input: 'abcdefghijklmnopqrstuvwxy'
}
]
})
);
});
it('should ignore failed image embedding normalization and keep text recall', async () => { it('should ignore failed image embedding normalization and keep text recall', async () => {
mockGetLLMModel.mockReturnValue(undefined); mockGetLLMModel.mockReturnValue(undefined);
mockIsImageEmbeddingModel.mockReturnValue(true); mockIsImageEmbeddingModel.mockReturnValue(true);
......
packages/service/test/core/dataset/textSplitter.test.ts
...@@ -405,3 +405,36 @@ it('should preserve escaped pipe in markdown table cells when splitting', async ...@@ -405,3 +405,36 @@ it('should preserve escaped pipe in markdown table cells when splitting', async
expect(data.map((chunk) => chunk.q).join('\n')).toContain('投资回报率 \\| abcd'); expect(data.map((chunk) => chunk.q).join('\n')).toContain('投资回报率 \\| abcd');
}); });
it('should skip markdown table header-only chunks when building dataset chunks', async () => {
const data = await rawText2Chunks({
rawText: `| id | payload | note |
| --- | --- | --- |`,
chunkTriggerType: ChunkTriggerConfigTypeEnum.forceChunk,
chunkTriggerMinSize: 10,
maxSize: 10000,
chunkSize: 40,
backupParse: false
});
expect(data).toEqual([]);
});
it('should not create header-only chunk for markdown table with a long first row', async () => {
const data = await rawText2Chunks({
rawText: `| id | payload | note |
| --- | --- | --- |
| 1 | ${'𠮷'.repeat(3000)} | old split keeps this single markdown table row as one index chunk |`,
chunkTriggerType: ChunkTriggerConfigTypeEnum.forceChunk,
chunkTriggerMinSize: 10,
maxSize: 10000,
chunkSize: 512,
backupParse: false
});
expect(data.length).toBeGreaterThan(0);
expect(data.map((chunk) => chunk.q)).not.toContain(
'| id | payload | note |\n| --- | --- | --- |'
);
expect(data.map((chunk) => chunk.q).join('\n')).toContain('| 1 |');
});
packages/service/test/worker/function.test.ts
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { WorkerNameEnum } from '@fastgpt/service/worker/utils'; import { WorkerNameEnum } from '@fastgpt/service/worker/utils';
import { countPromptTokensInWorker } from '@fastgpt/service/worker/countGptMessagesTokens/count';
// hoisted: 这些 mock 必须在 vi.mock 工厂里可见 // hoisted: 这些 mock 必须在 vi.mock 工厂里可见
const { mockRun, mockGetWorkerController, mockRunWorker, mockUploadImage2S3Bucket, mockEnv } = const { mockRun, mockGetWorkerController, mockRunWorker, mockUploadImage2S3Bucket, mockEnv } =
...@@ -82,6 +83,79 @@ describe('worker/function', () => { ...@@ -82,6 +83,79 @@ describe('worker/function', () => {
const result = await text2Chunks({ text: '', chunkSize: 100, maxSize: 200 }); const result = await text2Chunks({ text: '', chunkSize: 100, maxSize: 200 });
expect(result.chunks).toEqual([]); expect(result.chunks).toEqual([]);
}); });
it('test 环境下 token 模式按 token 上限切分文本', async () => {
const text = '𠮷'.repeat(8);
const result = await text2Chunks({
text,
chunkSize: 12,
maxSize: 12,
lengthUnit: 'token'
});
expect(result.chunks.length).toBeGreaterThan(1);
expect(result.chunks.every((chunk) => countPromptTokensInWorker(chunk) <= 12)).toBe(true);
expect(result.chunks.join('')).toBe(text);
expect(mockRunWorker).not.toHaveBeenCalled();
expect(mockGetWorkerController).not.toHaveBeenCalled();
});
it('test 环境下 token 模式长文本兜底分割仍不超过 maxSize', async () => {
const text = '𠮷'.repeat(400);
const chunkSize = 96;
const result = await text2Chunks({
text,
chunkSize,
maxSize: chunkSize,
overlapRatio: 0,
lengthUnit: 'token'
});
expect(countPromptTokensInWorker(text)).toBeGreaterThan(chunkSize * 10);
expect(result.chunks.length).toBeGreaterThan(10);
expect(result.chunks.every((chunk) => countPromptTokensInWorker(chunk) <= chunkSize)).toBe(
true
);
expect(result.chunks.join('')).toBe(text);
expect(mockRunWorker).not.toHaveBeenCalled();
expect(mockGetWorkerController).not.toHaveBeenCalled();
});
it('token 模式无法放入单个字符时直接报错', async () => {
await expect(
text2Chunks({
text: '𠮷',
chunkSize: 1,
maxSize: 1,
lengthUnit: 'token'
})
).rejects.toThrow('Text contains a character that exceeds the token length limit');
expect(mockRunWorker).not.toHaveBeenCalled();
expect(mockGetWorkerController).not.toHaveBeenCalled();
});
it('token 模式拆分 markdown 表格时每个最终分块都包含表头且不超过上限', async () => {
const header = `| id | payload |
| --- | --- |
`;
const text = `${header}| 1 | ${'𠮷'.repeat(20)} |
`;
const result = await text2Chunks({
text,
chunkSize: 28,
maxSize: 28,
lengthUnit: 'token'
});
expect(result.chunks.length).toBeGreaterThan(1);
expect(result.chunks.every((chunk) => chunk.startsWith(header))).toBe(true);
expect(result.chunks.every((chunk) => countPromptTokensInWorker(chunk) <= 28)).toBe(true);
expect(result.chunks.join('\n')).toContain('𠮷');
expect(mockRunWorker).not.toHaveBeenCalled();
expect(mockGetWorkerController).not.toHaveBeenCalled();
});
}); });
describe('readRawContentFromBuffer', () => { describe('readRawContentFromBuffer', () => {
......
packages/service/worker/function.ts
import { import type { SplitProps, SplitResponse } from '../common/string/textSplitter';
splitText2Chunks,
type SplitProps,
type SplitResponse
} from '@fastgpt/global/common/string/textSplitter';
import { getWorkerController, WorkerNameEnum } from './utils'; 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';
...@@ -11,9 +7,10 @@ import { uploadImage2S3Bucket } from '../common/s3/utils'; ...@@ -11,9 +7,10 @@ import { uploadImage2S3Bucket } from '../common/s3/utils';
import { normalizeMimeType, resolveMimeType } from '../common/s3/utils/mime'; import { normalizeMimeType, resolveMimeType } from '../common/s3/utils/mime';
import path from 'node:path'; import path from 'node:path';
export const text2Chunks = (props: SplitProps) => { export const text2Chunks = async (props: SplitProps) => {
// Test env, not run worker // Test env, not run worker
if (isTestEnv) { if (isTestEnv) {
const { splitText2Chunks } = await import('../common/string/textSplitter');
return splitText2Chunks(props); return splitText2Chunks(props);
} }
return getWorkerController<SplitProps, SplitResponse>({ return getWorkerController<SplitProps, SplitResponse>({
......
packages/service/worker/readFile/extension/xlsx.ts
import { CUSTOM_SPLIT_SIGN } from '@fastgpt/global/common/string/textSplitter'; import { CUSTOM_SPLIT_SIGN } from '../../../common/string/textSplitter';
import { type ReadRawTextByBuffer, type ReadFileResponse } from '../type'; import { type ReadRawTextByBuffer, type ReadFileResponse } from '../type';
import XLSX from 'xlsx'; import XLSX from 'xlsx';
import { filterEmptyTableData, formatMarkdownTableRow } from './utils'; import { filterEmptyTableData, formatMarkdownTableRow } from './utils';
......
packages/service/worker/text2Chunks/index.ts
import { parentPort } from 'worker_threads'; import { parentPort } from 'worker_threads';
import type { SplitProps } from '@fastgpt/global/common/string/textSplitter'; import type { SplitProps } from '../../common/string/textSplitter';
import { splitText2Chunks } from '@fastgpt/global/common/string/textSplitter'; import { splitText2Chunks } from '../../common/string/textSplitter';
type IncomingMessage = { type IncomingMessage = {
id: string; id: string;
......
projects/app/src/pageComponents/dataset/detail/Import/commonProgress/PreviewData.tsx
...@@ -8,15 +8,13 @@ import FormLabel from '@fastgpt/web/components/common/MyBox/FormLabel'; ...@@ -8,15 +8,13 @@ import FormLabel from '@fastgpt/web/components/common/MyBox/FormLabel';
import EmptyTip from '@fastgpt/web/components/common/EmptyTip'; import EmptyTip from '@fastgpt/web/components/common/EmptyTip';
import { useRequest } from '@fastgpt/web/hooks/useRequest'; import { useRequest } from '@fastgpt/web/hooks/useRequest';
import { ImportDataSourceEnum } from '@fastgpt/global/core/dataset/constants'; import { ImportDataSourceEnum } from '@fastgpt/global/core/dataset/constants';
import { splitText2Chunks } from '@fastgpt/global/common/string/textSplitter'; import { getPreviewChunks, getRawTextPreviewChunks } from '@/web/core/dataset/api/file';
import { getPreviewChunks } from '@/web/core/dataset/api/file';
import { type ImportSourceItemType } from '@/web/core/dataset/type'; import { type ImportSourceItemType } from '@/web/core/dataset/type';
import { getPreviewSourceReadType } from '../utils'; import { getPreviewSourceReadType } from '../utils';
import { DatasetPageContext } from '@/web/core/dataset/context/datasetPageContext'; import { DatasetPageContext } from '@/web/core/dataset/context/datasetPageContext';
import MyBox from '@fastgpt/web/components/common/MyBox'; import MyBox from '@fastgpt/web/components/common/MyBox';
import Markdown from '@/components/Markdown'; import Markdown from '@/components/Markdown';
import { useToast } from '@fastgpt/web/hooks/useToast'; import { useToast } from '@fastgpt/web/hooks/useToast';
import { getLLMMaxChunkSize } from '@fastgpt/global/core/dataset/training/utils';
const PreviewData = () => { const PreviewData = () => {
const { t } = useTranslation(); const { t } = useTranslation();
...@@ -24,7 +22,6 @@ const PreviewData = () => { ...@@ -24,7 +22,6 @@ const PreviewData = () => {
const goToNext = useContextSelector(DatasetImportContext, (v) => v.goToNext); const goToNext = useContextSelector(DatasetImportContext, (v) => v.goToNext);
const datasetId = useContextSelector(DatasetPageContext, (v) => v.datasetId); const datasetId = useContextSelector(DatasetPageContext, (v) => v.datasetId);
const datasetDetail = useContextSelector(DatasetPageContext, (v) => v.datasetDetail);
const sources = useContextSelector(DatasetImportContext, (v) => v.sources); const sources = useContextSelector(DatasetImportContext, (v) => v.sources);
const importSource = useContextSelector(DatasetImportContext, (v) => v.importSource); const importSource = useContextSelector(DatasetImportContext, (v) => v.importSource);
...@@ -39,21 +36,12 @@ const PreviewData = () => { ...@@ -39,21 +36,12 @@ const PreviewData = () => {
const chunkData = processParamsForm.getValues(); const chunkData = processParamsForm.getValues();
if (importSource === ImportDataSourceEnum.fileCustom) { if (importSource === ImportDataSourceEnum.fileCustom) {
const chunkSplitter = processParamsForm.getValues('chunkSplitter'); return getRawTextPreviewChunks({
const { chunks } = splitText2Chunks({ datasetId,
text: previewFile.rawText || '', rawText: previewFile.rawText || '',
chunkSize: chunkData.chunkSize, ...chunkData,
maxSize: getLLMMaxChunkSize(datasetDetail.agentModel), overlapRatio: 0.2
overlapRatio: 0.2,
customReg: chunkSplitter ? [chunkSplitter] : []
}); });
return {
chunks: chunks.map((chunk) => ({
q: chunk,
a: ''
})),
total: chunks.length
};
} }
return getPreviewChunks({ return getPreviewChunks({
...@@ -90,12 +78,12 @@ const PreviewData = () => { ...@@ -90,12 +78,12 @@ const PreviewData = () => {
return ( return (
<Flex flexDirection={'column'} h={'100%'}> <Flex flexDirection={'column'} h={'100%'}>
<Flex flex={'1 0 0'} border={'base'} borderRadius={'md'}> <Flex flex={'1 0 0'} minW={0} overflow={'hidden'} border={'base'} borderRadius={'md'}>
<Flex flexDirection={'column'} flex={'1 0 0'} borderRight={'base'}> <Flex flexDirection={'column'} flex={'0 0 50%'} maxW={'50%'} minW={0} borderRight={'base'}>
<FormLabel fontSize={'md'} py={4} px={5} borderBottom={'base'}> <FormLabel fontSize={'md'} py={4} px={5} borderBottom={'base'}>
{t('dataset:file_list')} {t('dataset:file_list')}
</FormLabel> </FormLabel>
<Box flex={'1 0 0'} overflowY={'auto'} px={5} py={3}> <Box flex={'1 0 0'} minW={0} overflowY={'auto'} px={5} py={3}>
{sources.map((source) => ( {sources.map((source) => (
<HStack <HStack
key={source.id} key={source.id}
...@@ -126,22 +114,22 @@ const PreviewData = () => { ...@@ -126,22 +114,22 @@ const PreviewData = () => {
}} }}
> >
<MyIcon name={source.icon as any} w={'1.25rem'} /> <MyIcon name={source.icon as any} w={'1.25rem'} />
<Box ml={1} flex={'1 0 0'} wordBreak={'break-all'} fontSize={'sm'}> <Box ml={1} flex={'1 1 0'} minW={0} wordBreak={'break-all'} fontSize={'sm'}>
{source.sourceName} {source.sourceName}
</Box> </Box>
</HStack> </HStack>
))} ))}
</Box> </Box>
</Flex> </Flex>
<Flex flexDirection={'column'} flex={'1 0 0'}> <Flex flexDirection={'column'} flex={'0 0 50%'} maxW={'50%'} minW={0}>
<Flex py={4} px={5} borderBottom={'base'} justifyContent={'space-between'}> <Flex py={4} px={5} borderBottom={'base'} justifyContent={'space-between'}>
<FormLabel fontSize={'md'}>{t('dataset:preview_chunk')}</FormLabel> <FormLabel fontSize={'md'}>{t('dataset:preview_chunk')}</FormLabel>
<Box fontSize={'xs'} color={'myGray.500'}> <Box fontSize={'xs'} color={'myGray.500'}>
{t('dataset:preview_chunk_intro', { total: data.total })} {t('dataset:preview_chunk_intro', { total: data.total })}
</Box> </Box>
</Flex> </Flex>
<MyBox isLoading={isLoading} flex={'1 0 0'} h={0}> <MyBox isLoading={isLoading} flex={'1 0 0'} h={0} minW={0}>
<Box h={'100%'} overflowY={'auto'} px={5} py={3}> <Box h={'100%'} minW={0} overflowY={'auto'} overflowX={'auto'} px={5} py={3}>
{previewFile ? ( {previewFile ? (
<> <>
{data.chunks.map((item, index) => ( {data.chunks.map((item, index) => (
......
projects/app/src/pageComponents/dataset/detail/Info/index.tsx
...@@ -28,7 +28,6 @@ import DatasetTypeTag from '@/components/core/dataset/DatasetTypeTag'; ...@@ -28,7 +28,6 @@ import DatasetTypeTag from '@/components/core/dataset/DatasetTypeTag';
import dynamic from 'next/dynamic'; import dynamic from 'next/dynamic';
import type { EditAPIDatasetInfoFormType } from './components/EditApiServiceModal'; import type { EditAPIDatasetInfoFormType } from './components/EditApiServiceModal';
import { type EditResourceInfoFormType } from '@/components/common/Modal/EditResourceModal'; import { type EditResourceInfoFormType } from '@/components/common/Modal/EditResourceModal';
import MyTooltip from '@fastgpt/web/components/common/MyTooltip';
import { ReadRoleVal } from '@fastgpt/global/support/permission/constant'; import { ReadRoleVal } from '@fastgpt/global/support/permission/constant';
const EditResourceModal = dynamic(() => import('@/components/common/Modal/EditResourceModal')); const EditResourceModal = dynamic(() => import('@/components/common/Modal/EditResourceModal'));
...@@ -170,11 +169,6 @@ const Info = ({ datasetId }: { datasetId: string }) => { ...@@ -170,11 +169,6 @@ const Info = ({ datasetId }: { datasetId: string }) => {
<FormLabel fontWeight={'500'} flex={'1 0 0'} fontSize={'mini'}> <FormLabel fontWeight={'500'} flex={'1 0 0'} fontSize={'mini'}>
{t('common:core.ai.model.Vector Model')} {t('common:core.ai.model.Vector Model')}
</FormLabel> </FormLabel>
<MyTooltip label={t('dataset:vector_model_max_tokens_tip')}>
<Box fontSize={'mini'}>
{t('dataset:chunk_max_tokens')}: {vectorModel.maxToken}
</Box>
</MyTooltip>
</Flex> </Flex>
<Box pt={2} minW={0} maxW={'100%'} overflow={'hidden'}> <Box pt={2} minW={0} maxW={'100%'} overflow={'hidden'}>
<AIModelSelector <AIModelSelector
......
projects/app/src/pages/api/core/dataset/file/getPreviewChunks.ts
...@@ -8,7 +8,8 @@ import { authDataset } from '@fastgpt/service/support/permission/dataset/auth'; ...@@ -8,7 +8,8 @@ import { authDataset } from '@fastgpt/service/support/permission/dataset/auth';
import { isAuthorizedDatasetFileS3Key } from '@fastgpt/service/common/s3/sources/dataset/key'; import { isAuthorizedDatasetFileS3Key } from '@fastgpt/service/common/s3/sources/dataset/key';
import { import {
computedCollectionChunkSettings, computedCollectionChunkSettings,
getLLMMaxChunkSize getLLMMaxChunkSize,
maxPreviewChunkCount
} from '@fastgpt/global/core/dataset/training/utils'; } from '@fastgpt/global/core/dataset/training/utils';
import { CommonErrEnum } from '@fastgpt/global/common/error/code/common'; import { CommonErrEnum } from '@fastgpt/global/common/error/code/common';
import { getEmbeddingModel, getLLMModel } from '@fastgpt/service/core/ai/model'; import { getEmbeddingModel, getLLMModel } from '@fastgpt/service/core/ai/model';
...@@ -97,7 +98,8 @@ async function handler( ...@@ -97,7 +98,8 @@ async function handler(
paragraphChunkMinSize: formatChunkSettings.paragraphChunkMinSize, paragraphChunkMinSize: formatChunkSettings.paragraphChunkMinSize,
maxSize: getLLMMaxChunkSize(getLLMModel(dataset.agentModel)), maxSize: getLLMMaxChunkSize(getLLMModel(dataset.agentModel)),
overlapRatio, overlapRatio,
customReg: formatChunkSettings.chunkSplitter ? [formatChunkSettings.chunkSplitter] : [] customReg: formatChunkSettings.chunkSplitter ? [formatChunkSettings.chunkSplitter] : [],
maxChunks: maxPreviewChunkCount
}); });
const chunksWithJWT = chunks.slice(0, 10).map((chunk) => ({ const chunksWithJWT = chunks.slice(0, 10).map((chunk) => ({
......
projects/app/src/pages/api/core/dataset/file/getRawTextPreviewChunks.ts 0 → 100644
import { NextAPI } from '@/service/middleware/entry';
import type { ApiRequestProps } from '@fastgpt/service/type/next';
import { WritePermissionVal } from '@fastgpt/global/support/permission/constant';
import { authDataset } from '@fastgpt/service/support/permission/dataset/auth';
import { rawText2Chunks } from '@fastgpt/service/core/dataset/read';
import {
computedCollectionChunkSettings,
getLLMMaxChunkSize,
maxPreviewChunkCount
} from '@fastgpt/global/core/dataset/training/utils';
import { getEmbeddingModel, getLLMModel } from '@fastgpt/service/core/ai/model';
import { replaceS3KeyToPreviewUrl } from '@fastgpt/service/core/dataset/utils';
import { addDays } from 'date-fns';
import { parseApiInput } from '@fastgpt/service/common/zod/requestParseError';
import {
GetRawTextPreviewChunksBodySchema,
GetPreviewChunksResponseSchema,
type GetRawTextPreviewChunksBody,
type GetRawTextPreviewChunksResponse
} from '@fastgpt/global/openapi/core/dataset/file/api';
async function handler(
req: ApiRequestProps<GetRawTextPreviewChunksBody>
): Promise<GetRawTextPreviewChunksResponse> {
const { datasetId, rawText, overlapRatio, ...chunkSettings } = parseApiInput({
req,
bodySchema: GetRawTextPreviewChunksBodySchema
}).body;
const { dataset } = await authDataset({
req,
authApiKey: true,
authToken: true,
datasetId,
per: WritePermissionVal
});
const formatChunkSettings = computedCollectionChunkSettings({
...chunkSettings,
llmModel: getLLMModel(dataset.agentModel),
vectorModel: getEmbeddingModel(dataset.vectorModel)
});
const chunks = await rawText2Chunks({
rawText,
chunkTriggerType: formatChunkSettings.chunkTriggerType,
chunkTriggerMinSize: formatChunkSettings.chunkTriggerMinSize,
chunkSize: formatChunkSettings.chunkSize,
paragraphChunkDeep: formatChunkSettings.paragraphChunkDeep,
paragraphChunkMinSize: formatChunkSettings.paragraphChunkMinSize,
maxSize: getLLMMaxChunkSize(getLLMModel(dataset.agentModel)),
overlapRatio,
customReg: formatChunkSettings.chunkSplitter ? [formatChunkSettings.chunkSplitter] : [],
maxChunks: maxPreviewChunkCount
});
const chunksWithJWT = chunks.slice(0, 10).map((chunk) => ({
q: replaceS3KeyToPreviewUrl(chunk.q, addDays(new Date(), 1)),
a: replaceS3KeyToPreviewUrl(chunk.a, addDays(new Date(), 1))
}));
return GetPreviewChunksResponseSchema.parse({
chunks: chunksWithJWT,
total: chunks.length
});
}
export const config = {
api: {
bodyParser: {
sizeLimit: '10mb'
}
}
};
export default NextAPI(handler);
projects/app/src/service/core/dataset/data/dataIndex.ts
...@@ -22,6 +22,7 @@ import { isS3ObjectKey } from '@fastgpt/service/common/s3/utils'; ...@@ -22,6 +22,7 @@ import { isS3ObjectKey } from '@fastgpt/service/common/s3/utils';
import { getS3DatasetSource } from '@fastgpt/service/common/s3/sources/dataset'; import { getS3DatasetSource } from '@fastgpt/service/common/s3/sources/dataset';
import { uniqueDatasetDataMarkdownImageUrls } from '@fastgpt/service/core/dataset/data/utils'; import { uniqueDatasetDataMarkdownImageUrls } from '@fastgpt/service/core/dataset/data/utils';
import { isDatasetDataSystemIndexType } from '@fastgpt/global/core/dataset/data/utils'; import { isDatasetDataSystemIndexType } from '@fastgpt/global/core/dataset/data/utils';
import { minChunkSize } from '@fastgpt/global/core/dataset/training/utils';
export type DatasetDataIndexDraft = Omit<DatasetDataIndexItemType, 'dataId'> & { export type DatasetDataIndexDraft = Omit<DatasetDataIndexItemType, 'dataId'> & {
dataId?: string; dataId?: string;
...@@ -64,6 +65,88 @@ const formatIndexTextWithPrefix = (text: string, indexPrefix?: string) => { ...@@ -64,6 +65,88 @@ const formatIndexTextWithPrefix = (text: string, indexPrefix?: string) => {
return text; return text;
}; };
/**
* 按 embedding token 预算拆分索引正文，并给每个最终 chunk 补上集合前缀。
*
* 这个 helper 总是走 text2Chunks，避免绕过索引分块的文本规范化；调用方通过
* `indexSize` 控制期望索引粒度，通过 `maxToken` 控制 embedding provider 硬上限。
* indexSize 下限沿用知识库分块最小值，避免 prefix 挤压后传入过小 token 预算。
*/
const splitIndexTextByTokenLimit = async ({
text,
indexSize,
maxToken,
indexPrefix
}: {
text: string;
indexSize: number;
maxToken: number;
indexPrefix?: string;
}) => {
const trimmedText = text.trim();
if (!trimmedText) return [];
const prefixTokens = indexPrefix ? await countPromptTokens(`${indexPrefix}\n`) : 0;
const maxContentTokens = maxToken - prefixTokens;
if (maxContentTokens <= 0) {
throw new Error('Dataset index prefix is too long for embedding token limit');
}
if (maxContentTokens < minChunkSize) {
throw new Error('Dataset index content token budget is smaller than minimum chunk size');
}
const normalizedIndexSize = Math.max(indexSize, minChunkSize);
const chunkTokenLimit = Math.min(normalizedIndexSize, maxContentTokens);
// 入库索引和 query 不一样：这里允许一条 index 拆成多条，以尽量保留原始内容。
// 每条最终文本都会拼上 indexPrefix，所以正文预算必须先扣掉前缀 token。
const chunks = (
await text2Chunks({
text: trimmedText,
chunkSize: chunkTokenLimit,
maxSize: chunkTokenLimit,
lengthUnit: 'token'
})
).chunks;
return chunks
.map((chunk) => formatIndexTextWithPrefix(chunk, indexPrefix))
.filter((item) => item.trim());
};
/**
* 构建最终可写入 embedding 的索引文本。
*
* 默认保持旧语义：未超过 embedding 上限的既有索引不强行重分块，避免无谓重建向量。
* 超过上限时再按 `min(max(indexSize, 64), maxToken - prefixTokens)` 做 token-safe 二次拆分。
*/
const buildEmbeddingSafeIndexTexts = async ({
text,
indexSize,
maxToken,
indexPrefix
}: {
text: string;
indexSize: number;
maxToken: number;
indexPrefix?: string;
}) => {
const trimmedText = text.trim();
if (!trimmedText) return [];
const formattedText = formatIndexTextWithPrefix(text, indexPrefix);
if ((await countPromptTokens(formattedText)) <= maxToken) {
return [formattedText];
}
return splitIndexTextByTokenLimit({
text: trimmedText,
indexSize,
maxToken,
indexPrefix
});
};
const isImageEmbeddingIndex = (index: DatasetDataIndexDraft) => const isImageEmbeddingIndex = (index: DatasetDataIndexDraft) =>
index.type === DatasetDataIndexTypeEnum.imageEmbedding; index.type === DatasetDataIndexTypeEnum.imageEmbedding;
...@@ -160,30 +243,28 @@ export class DatasetDataIndexOperation { ...@@ -160,30 +243,28 @@ export class DatasetDataIndexOperation {
maxIndexSize?: number; maxIndexSize?: number;
indexPrefix?: string; indexPrefix?: string;
}) { }) {
const qChunks = ( const qIndexTexts = await splitIndexTextByTokenLimit({
await text2Chunks({ text: q,
text: q, indexSize,
chunkSize: indexSize, maxToken: maxIndexSize ?? this.maxToken,
maxSize: maxIndexSize ?? this.maxToken indexPrefix
}) });
).chunks; const aIndexTexts = a
const aChunks = a ? await splitIndexTextByTokenLimit({
? ( text: a,
await text2Chunks({ indexSize,
text: a, maxToken: maxIndexSize ?? this.maxToken,
chunkSize: indexSize, indexPrefix
maxSize: maxIndexSize ?? this.maxToken })
})
).chunks
: []; : [];
return [ return [
...qChunks.map((text) => ({ ...qIndexTexts.map((text) => ({
text: formatIndexTextWithPrefix(text, indexPrefix), text,
type: DatasetDataIndexTypeEnum.default type: DatasetDataIndexTypeEnum.default
})), })),
...aChunks.map((text) => ({ ...aIndexTexts.map((text) => ({
text: formatIndexTextWithPrefix(text, indexPrefix), text,
type: DatasetDataIndexTypeEnum.default type: DatasetDataIndexTypeEnum.default
})), })),
...this.getImageEmbeddingSources({ ...this.getImageEmbeddingSources({
...@@ -284,17 +365,21 @@ export class DatasetDataIndexOperation { ...@@ -284,17 +365,21 @@ export class DatasetDataIndexOperation {
if (item.type === DatasetDataIndexTypeEnum.imageEmbedding) { if (item.type === DatasetDataIndexTypeEnum.imageEmbedding) {
return item; return item;
} }
// 系统文本索引刚由 getSystemIndexes 按最终 prefix 和 token 上限生成，
// 这里直接复用，避免在同一入库请求内再次投递 token worker 计数。
if (isDatasetDataSystemIndexType(item.type)) {
return item;
}
const tokens = await countPromptTokens(item.text); const indexTexts = await buildEmbeddingSafeIndexTexts({
if (tokens > (maxIndexSize ?? this.maxToken)) { text: item.text,
const splitText = ( indexSize,
await text2Chunks({ maxToken: maxIndexSize ?? this.maxToken,
text: item.text, indexPrefix: item.type === DatasetDataIndexTypeEnum.default ? indexPrefix : undefined
chunkSize: indexSize, });
maxSize: maxIndexSize ?? this.maxToken
}) if (indexTexts.length > 1 || indexTexts[0] !== item.text) {
).chunks; return indexTexts.map((text) => ({
return splitText.map((text) => ({
text, text,
type: item.type type: item.type
})); }));
...@@ -307,21 +392,7 @@ export class DatasetDataIndexOperation { ...@@ -307,21 +392,7 @@ export class DatasetDataIndexOperation {
.flat() .flat()
.filter((item) => !!item.text.trim()); .filter((item) => !!item.text.trim());
return indexPrefix return checkedIndexes;
? checkedIndexes.map((index) => {
// 自定义索引与图片向量索引不需要添加前缀
if (
index.type === DatasetDataIndexTypeEnum.custom ||
index.type === DatasetDataIndexTypeEnum.imageEmbedding
) {
return index;
}
return {
...index,
text: formatIndexTextWithPrefix(index.text, indexPrefix)
};
})
: checkedIndexes;
} }
/** /**
......
projects/app/src/web/core/dataset/api/file.ts
...@@ -4,6 +4,8 @@ import type { ...@@ -4,6 +4,8 @@ import type {
GetSearchTestImagePreviewUrlsResponse, GetSearchTestImagePreviewUrlsResponse,
GetPreviewChunksBody, GetPreviewChunksBody,
GetPreviewChunksResponse, GetPreviewChunksResponse,
GetRawTextPreviewChunksBody,
GetRawTextPreviewChunksResponse,
PresignDatasetFilePostUrlBody, PresignDatasetFilePostUrlBody,
PresignSearchTestImageBody, PresignSearchTestImageBody,
PresignSearchTestImageResponse PresignSearchTestImageResponse
...@@ -19,6 +21,12 @@ export const getPreviewChunks = (data: GetPreviewChunksBody) => ...@@ -19,6 +21,12 @@ export const getPreviewChunks = (data: GetPreviewChunksBody) =>
timeout: 600000 timeout: 600000
}); });
export const getRawTextPreviewChunks = (data: GetRawTextPreviewChunksBody) =>
POST<GetRawTextPreviewChunksResponse>('/core/dataset/file/getRawTextPreviewChunks', data, {
maxQuantity: 1,
timeout: 600000
});
export const getUploadSearchTestImagePresignedUrl = (data: PresignSearchTestImageBody) => export const getUploadSearchTestImagePresignedUrl = (data: PresignSearchTestImageBody) =>
POST<PresignSearchTestImageResponse>('/core/dataset/file/presignSearchTestImage', data); POST<PresignSearchTestImageResponse>('/core/dataset/file/presignSearchTestImage', data);
......
projects/app/test/pages/api/core/dataset/file/getPreviewChunks.test.ts
...@@ -42,7 +42,9 @@ vi.mock('@fastgpt/global/core/dataset/training/utils', () => ({ ...@@ -42,7 +42,9 @@ vi.mock('@fastgpt/global/core/dataset/training/utils', () => ({
paragraphChunkMinSize: 100, paragraphChunkMinSize: 100,
chunkSplitter: '' chunkSplitter: ''
})), })),
getLLMMaxChunkSize: vi.fn(() => 1000) getLLMMaxChunkSize: vi.fn(() => 1000),
minChunkSize: 64,
maxPreviewChunkCount: 50_000
})); }));
vi.mock('@fastgpt/service/core/dataset/utils', () => ({ vi.mock('@fastgpt/service/core/dataset/utils', () => ({
......
projects/app/test/pages/api/core/dataset/file/getRawTextPreviewChunks.test.ts 0 → 100644
import type { ApiRequestProps } from '@fastgpt/service/type/next';
import { beforeEach, describe, expect, it, vi } from 'vitest';
const mocks = vi.hoisted(() => ({
authDataset: vi.fn(),
rawText2Chunks: vi.fn(),
replaceS3KeyToPreviewUrl: vi.fn((value: string) => value)
}));
vi.mock('@/service/middleware/entry', () => ({
NextAPI: (handler: unknown) => handler
}));
vi.mock('@fastgpt/service/support/permission/dataset/auth', () => ({
authDataset: mocks.authDataset
}));
vi.mock('@fastgpt/service/core/dataset/read', () => ({
rawText2Chunks: mocks.rawText2Chunks
}));
vi.mock('@fastgpt/service/core/ai/model', () => ({
getEmbeddingModel: vi.fn(() => ({})),
getLLMModel: vi.fn(() => ({}))
}));
vi.mock('@fastgpt/global/core/dataset/training/utils', () => ({
computedCollectionChunkSettings: vi.fn(() => ({
chunkTriggerType: 'minSize',
chunkTriggerMinSize: 100,
chunkSize: 500,
paragraphChunkDeep: 1,
paragraphChunkMinSize: 100,
chunkSplitter: ''
})),
getLLMMaxChunkSize: vi.fn(() => 1000),
minChunkSize: 64,
maxPreviewChunkCount: 50_000
}));
vi.mock('@fastgpt/service/core/dataset/utils', () => ({
replaceS3KeyToPreviewUrl: mocks.replaceS3KeyToPreviewUrl
}));
import handler from '@/pages/api/core/dataset/file/getRawTextPreviewChunks';
const datasetId = '507f1f77bcf86cd799439011';
const previewHandler = handler as unknown as (req: ApiRequestProps) => Promise<unknown>;
const callHandler = (body: Record<string, unknown>) =>
previewHandler({
body
} as ApiRequestProps);
describe('getRawTextPreviewChunks', () => {
beforeEach(() => {
vi.clearAllMocks();
mocks.authDataset.mockResolvedValue({
dataset: {
agentModel: 'gpt',
vectorModel: 'embedding'
}
});
mocks.rawText2Chunks.mockResolvedValue([
{
q: 'hello',
a: ''
}
]);
});
it('previews frontend raw text through backend chunking', async () => {
await expect(
callHandler({
datasetId,
rawText: 'hello world',
overlapRatio: 0.2,
chunkSize: 500,
chunkSplitter: ''
})
).resolves.toEqual({
chunks: [
{
q: 'hello',
a: ''
}
],
total: 1
});
expect(mocks.authDataset).toHaveBeenCalledWith(
expect.objectContaining({
datasetId
})
);
expect(mocks.rawText2Chunks).toHaveBeenCalledWith(
expect.objectContaining({
rawText: 'hello world',
chunkSize: 500,
overlapRatio: 0.2,
maxChunks: 50_000
})
);
});
it('stops before chunking when dataset write permission is denied', async () => {
mocks.authDataset.mockRejectedValueOnce(new Error('forbidden'));
await expect(
callHandler({
datasetId,
rawText: 'hello world',
overlapRatio: 0.2,
chunkSize: 500
})
).rejects.toThrow('forbidden');
expect(mocks.rawText2Chunks).not.toHaveBeenCalled();
});
it.each(['|', 'prefix|', '|suffix', 'prefix||suffix'])(
'rejects empty custom chunk separators: %s',
async (chunkSplitter) => {
await expect(
callHandler({
datasetId,
rawText: 'hello world',
overlapRatio: 0.2,
chunkSize: 500,
chunkSplitter
})
).rejects.toBeDefined();
expect(mocks.authDataset).not.toHaveBeenCalled();
expect(mocks.rawText2Chunks).not.toHaveBeenCalled();
}
);
it.each([
{ overlapRatio: -0.1, chunkSize: 500 },
{ overlapRatio: 0.41, chunkSize: 500 },
{ overlapRatio: 1, chunkSize: 500 },
{ overlapRatio: 0.2, chunkSize: 63 },
{ overlapRatio: 0.2, chunkSize: 64.5 }
])('rejects unsafe numeric chunk settings: %o', async (settings) => {
await expect(
callHandler({
datasetId,
rawText: 'hello world',
...settings
})
).rejects.toBeDefined();
expect(mocks.authDataset).not.toHaveBeenCalled();
expect(mocks.rawText2Chunks).not.toHaveBeenCalled();
});
});
test/mocks/core/ai/embedding.ts
...@@ -124,7 +124,8 @@ vi.mock('@fastgpt/service/core/ai/model', async (importOriginal) => { ...@@ -124,7 +124,8 @@ vi.mock('@fastgpt/service/core/ai/model', async (importOriginal) => {
...actual, ...actual,
getEmbeddingModel: vi.fn().mockReturnValue({ getEmbeddingModel: vi.fn().mockReturnValue({
model: 'text-embedding-ada-002', model: 'text-embedding-ada-002',
name: 'text-embedding-ada-002' name: 'text-embedding-ada-002',
maxToken: 100
}) })
}; };
}); });
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