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 2640 additions and 648 deletions
.agents/design/core/dataset/embedding-token-limit-guard-功能开发文档.md 0 → 100644
# Embedding 输入超长治理功能说明
# Embedding 输入超长治理功能说明
## 1. 背景
知识库入库、检索和 query extension 等链路都会调用 embedding 模型。此前部分链路只按字符长度或已有分块规则控制文本大小，但 embedding provider 实际按 token 限制输入长度，因此会出现以下问题：
- 知识库入库时，某个 index 文本超过 embedding 模型 `maxToken`，导致建索引失败。
- 知识库检索时，用户 query 或扩展 query 过长，导致 query embedding 失败。
- query extension / 文本相似度计算时，原始文本或候选文本过长，导致相似度 embedding 失败。
本次优化目标是：让知识库文本索引按 token 预算分块，并保证最终进入 embedding 的文本不超过模型 `maxToken`。原始数据内容保持不变，索引数量和粒度可能随 token 计数结果变化。
## 2. 处理原则
本次优化遵循以下原则：
- 不改变 `DatasetData.q` / `DatasetData.a` 的原文内容。
- 不改变现有用户交互和向量库结构；新增内部 rawText 预览 API，把原先浏览器内的分块计算迁到服务端。
- 不做存量数据迁移。
- 不专项优化 Markdown 表格、代码块、PDF 等内容的语义分块质量；只处理 token 上限和 header-only 等安全性问题。
- 入库阶段允许把一条 index 拆成多条 index。
- 检索 query 和相似度 query 不扩增数量；统一由 embedding 调用入口截断到安全长度。
- `getVectors` 是最后一道硬兜底：所有 text input 在请求 provider 前按模型 `maxToken` 截断；image input 不参与文本截断。
## 3. 本次改动范围
### 3.1 分块函数的长度计数能力
文件：
```text
packages/service/common/string/textSplitter.ts
packages/service/worker/text2Chunks/index.ts
packages/service/worker/function.ts
```
职责：
- `splitText2Chunks` 放在 `packages/service`，前端导入预览通过后端接口调用，不再直接依赖分块实现。
- `splitText2Chunks` 新增可选 `lengthUnit`。默认仍按字符长度执行，`token` 模式使用统一的 `o200k_base` tokenizer。
- token fallback 通过单调 code-point 游标，只在当前 chunk 附近做指数探测和二分，避免重复 tokenize 完整剩余文本。
- `maxChunks` 为预览等不可信输入提供工作量上限；空自定义分隔符、非法 overlap 和非正 chunk size 会直接报错。
- `text2Chunks` worker 承载 tokenizer 常驻内存，主 API 进程只提交可序列化的 `lengthUnit` 参数。
边界：
- 不把 tokenizer 放到 `packages/global`，避免 global 包承担后端运行时依赖。
- 前端 `fileCustom` 分块预览改为调用后端 rawText 预览接口，因此 `textSplitter.ts` 可以整体迁到 service。
- token 模式只用于知识库 index 分块这类后端场景。
### 3.2 知识库入库建索引
文件：
```text
projects/app/src/service/core/dataset/data/dataIndex.ts
```
新增核心流程：
```text
原始 q/a/index text
-> 系统 q/a index 直接调用 text2Chunks 的 token 模式
-> 按 min(indexSize, embedding maxToken - prefixTokens) 生成 index
-> 外部 index 未超限时保持原文，超限时按同一 token 模式拆分
-> 再写入向量库
```
对应 helper：
```ts
splitIndexTextByTokenLimit
buildEmbeddingSafeIndexTexts
```
影响范围：
- `getSystemIndexes`：系统默认索引直接按 token 模式和 `indexSize` 生成。
- `formatIndexes`：系统索引复用上一步的安全结果；外部文本索引只在超过 embedding 上限时拆分。
- `imageEmbedding` 类型索引不参与文本 token 拆分。
注意：
- 拆分只作用于 `indexes[].text`，不会改 `q/a` 原文。
- 拆出来的新 index 都保持原索引类型。
- 如果存在 `indexPrefix`，先拼出最终索引文本，再走 token 模式分块，保证最终送 embedding 的文本不超过模型上限。
- 如果一条超长 index 被拆成多条，新 index 会重新生成向量，不复用旧 `dataId`。
### 3.3 知识库检索阶段
文件：
```text
packages/service/core/dataset/search/defaultRecall/embeddingRecall.ts
```
处理对象：
- `textQueries`
- `imageCaptionQueries`
处理流程：
```text
query
-> trim/filter
-> 调用 getVectors
-> getVectors 内部按 embedding maxToken 截断 text input
```
注意：
- 不把一个 query 拆成多个 query。
- 图片 URL query 不走文本 token 截断。
- 检索结果合并、rerank、RRF、limit、similarity 等逻辑不变。
### 3.4 文本相似度 / query extension 阶段
文件：
```text
packages/service/core/ai/hooks/useTextCosine.ts
```
处理对象：
- `originalText`
- `candidates[]`
处理流程：
```text
originalText / candidates
-> 过滤空文本
-> 调用 getVectors 计算向量
-> getVectors 内部按 embedding maxToken 截断 text input
-> 继续执行原 lazy greedy selection 逻辑
```
注意：
- 候选文本数量不会因为超长而扩增。
- 返回的 `selectedData` 仍是原始候选文本的 trim 结果，不因为 embedding 截断改变展示内容。
- 如果原始文本为空、无有效候选或 `k <= 0`，直接返回空结果和 `embeddingTokens = 0`。
### 3.5 embedding provider 调用
文件：
```text
packages/service/core/ai/embedding/index.ts
```
处理流程：
```text
getVectors inputs
-> 校验输入结构和空输入
-> 批量统计 text input token 数
-> 只对超限 text input 按 model.maxToken 截断
-> 按 text / image 组装 provider 请求
-> 调用 embedding provider
-> provider 返回错误时按原错误链路抛出
```
作用：
- 集中处理检索 query、相似度 query、其他直接调用 embedding 的超长文本输入。
- 避免每个调用点都重复写 token 截断逻辑。
- 入库 index 仍优先在分块阶段拆成多条索引；`getVectors` 只作为最后硬兜底。
注意：
- `getVectors` 只截断 text input，不会把单条输入拆成多条。
- Unicode 截断按 code point 二分，不会生成孤立 surrogate。
- image input 保持原有 `image_url` 结构，不走文本 token 统计。
- provider tokenizer 和本地 tokenizer 可能存在差异，如果 provider 仍返回超限错误，继续沿用原错误链路抛出。
## 4. 测试关注点
### 4.1 入库建索引测试
测试文件：
```text
projects/app/test/service/core/dataset/data/dataIndex.test.ts
```
建议验证：
- 超长 `q` 会生成多条 token-safe 默认索引。
- 超长 `a` 会生成多条 token-safe 默认索引。
- 超长自定义文本 index 会拆成多条同类型 index。
- `imageEmbedding` index 不参与文本拆分。
- 短文本不会被额外拆分。
- 带 `indexPrefix` 时，按“前缀 + 正文”的最终文本计算 token。
验收标准：
```text
每条最终写入向量库的文本 index token <= 当前 embedding model.maxToken
```
### 4.2 检索阶段测试
测试文件：
```text
packages/service/test/core/dataset/search/defaultRecall.test.ts
```
建议验证：
- 超长 `textQueries` 不会在检索阶段扩增为多条 query。
- 超长 `imageCaptionQueries` 不会在检索阶段扩增为多条 query。
- query 不会被拆成多条。
- 图片 query 不受文本截断逻辑影响。
验收标准：
```text
检索阶段把有效文本 query 原样交给 getVectors，最终截断由 getVectors 统一兜底
```
### 4.3 文本相似度测试
测试文件：
```text
packages/service/test/core/ai/hooks/useTextCosine.test.ts
```
建议验证：
- `originalText` 超长时不会在相似度阶段扩增为多条输入。
- `candidates[]` 中超长候选不会在相似度阶段扩增为多条输入。
- 返回的 `selectedData` 仍使用原候选文本。
- `k <= 0` 或无有效文本时，不调用 embedding。
验收标准：
```text
useTextCosine 只做 trim/filter 和选择逻辑，最终截断由 getVectors 统一兜底
```
### 4.4 provider 调用入口测试
测试文件：
```text
packages/service/test/core/ai/embedding/index.test.ts
```
建议验证：
- 空 inputs 或空文本 input 会被本地拒绝。
- 超长 text input 会在请求 provider 前按 `model.maxToken` 截断。
- image input 会按 image_url 结构进入 provider。
- provider 返回错误时沿用原错误链路抛出。
- 包含 emoji、生僻汉字等 astral Unicode 字符时，截断结果保持 well-formed。
验收标准：
```text
getVectors 是 embedding 超长文本的最后硬兜底；provider 错误仍能正常透出
```
### 4.5 rawText 预览 API 安全测试
测试文件：
```text
projects/app/test/pages/api/core/dataset/file/getRawTextPreviewChunks.test.ts
packages/service/test/common/string/textSplitter.test.ts
```
建议验证：
- 无知识库写权限时不执行分块。
- 自定义分隔符不允许空项、首尾 `|` 或连续 `||`。
- `overlapRatio`、`chunkSize` 必须落在安全业务范围。
- 超过 `maxChunks` 的工作量会在继续分配大数组前终止。
## 5. 手工测试建议
### 5.1 测试前准备
因为本次改动在服务端逻辑，手工测试前需要：
```text
1. 重启 app 服务。
2. 确认运行的是包含本次改动的分支。
3. 对已有数据测试时，需要重新触发更新索引或重新上传文件。
```
如果服务未重启，或仍查看旧数据索引，页面可能仍显示旧逻辑生成的索引数量。
### 5.2 入库阶段手工验证
操作：
```text
1. 创建或选择知识库。
2. 上传包含超长文本的文件。
3. 等待训练任务完成。
4. 打开数据详情，查看数据索引。
```
预期：
```text
1. 训练任务不应因为 embedding max-token 超限失败。
2. 数据可以正常完成索引生成。
3. 超长 index 可能被拆成多条默认索引。
4. 拆分结果不保证保持 Markdown 表格头、代码块边界或语义完整性。
```
### 5.3 检索阶段手工验证
操作：
```text
1. 在知识库检索测试或应用对话中输入超长 query。
2. 触发语义检索。
```
预期：
```text
1. 不应因为 query embedding 超长导致请求失败。
2. 检索会基于截断后的 query 执行。
3. 不会因为一个超长 query 扩增成多个 query。
```
## 6. 非目标说明
本次不解决以下问题：
- Markdown 表格如何按行、列、单元格拆得更适合检索。
- PDF、Excel、代码块等内容的语义化分块质量。
- 超长 query 如何摘要后再检索。
- 存量超长索引的自动迁移。
- 前端展示索引时的折叠、摘要或虚拟滚动体验。
如果后续要优化索引质量，应单独设计“知识库分块质量优化”方案，不应混在 embedding max-token 兜底里。
## 7. 风险与注意事项
- 入库阶段超长 index 拆成多条后，向量数量会增加，可能带来更多 embedding token 消耗。
- query 截断会丢弃尾部信息，极端情况下可能影响召回准确性。
- 当前 token 统计使用 FastGPT 统一 token worker，与具体 provider 的 tokenizer 可能存在轻微差异。
- rawText 预览最多生成 50,000 个 chunk，超过上限会返回分块错误，避免单请求耗尽 worker 内存。
## 8. 修改文件清单
```text
packages/service/core/ai/embedding/tokenLimit.ts
packages/service/core/ai/embedding/index.ts
packages/service/core/ai/hooks/useTextCosine.ts
packages/service/core/dataset/search/defaultRecall/embeddingRecall.ts
packages/service/common/string/textSplitter.ts
packages/service/worker/function.ts
packages/service/worker/text2Chunks/index.ts
projects/app/src/service/core/dataset/data/dataIndex.ts
packages/global/openapi/core/dataset/file/api.ts
packages/global/openapi/core/dataset/file/index.ts
packages/global/core/dataset/training/utils.ts
projects/app/src/pages/api/core/dataset/file/getPreviewChunks.ts
projects/app/src/pages/api/core/dataset/file/getRawTextPreviewChunks.ts
projects/app/src/web/core/dataset/api/file.ts
projects/app/src/pageComponents/dataset/detail/Import/commonProgress/PreviewData.tsx
test/mocks/core/ai/embedding.ts
packages/service/test/core/ai/embedding/index.test.ts
packages/service/test/core/ai/hooks/useTextCosine.test.ts
packages/service/test/core/dataset/search/defaultRecall.test.ts
packages/service/test/common/string/textSplitter.test.ts
packages/service/test/worker/function.test.ts
projects/app/test/pages/api/core/dataset/file/getPreviewChunks.test.ts
projects/app/test/pages/api/core/dataset/file/getRawTextPreviewChunks.test.ts
projects/app/test/service/core/dataset/data/dataIndex.test.ts
```
## 9. 验收结论
本次功能验收只看一个核心结果：
```text
FastGPT 已知上游入口会尽量把 text input 控制在对应 embedding model.maxToken 内；最终是否超限以 provider 返回为准。
```
入库阶段允许拆分为多条 index；检索和相似度阶段只截断，不扩增 query 数量；`getVectors` 批量预判 token 数并只截断超限文本。
packages/global/common/string/textSplitter.ts deleted 100644 → 0
import { defaultMaxChunkSize } from '../../core/dataset/training/utils';
import { getErrText } from '../error/utils';
import { simpleText } from './tools';
import { getTextValidLength } from './utils';
export const CUSTOM_SPLIT_SIGN = '-----CUSTOM_SPLIT_SIGN-----';
export type SplitProps = {
text: string;
chunkSize: number;
paragraphChunkDeep?: number; // Paragraph deep
paragraphChunkMinSize?: number; // Paragraph min size, if too small, it will merge
maxSize?: number;
overlapRatio?: number;
customReg?: string[];
};
export type TextSplitProps = Omit<SplitProps, 'text' | 'chunkSize'> & {
chunkSize?: number;
};
export type SplitResponse = {
chunks: string[];
chars: number;
};
// 判断字符串是否为markdown的表格形式
const strIsMdTable = (str: string) => {
// 检查是否包含表格分隔符 |
if (!str.includes('|')) {
return false;
}
const lines = str.split('\n');
// 检查表格是否至少有两行
if (lines.length < 2) {
return false;
}
// 检查表头行是否包含 |
const headerLine = lines[0].trim();
if (!headerLine.startsWith('|') || !headerLine.endsWith('|')) {
return false;
}
// 检查分隔行是否由 | 和 - 组成
const separatorLine = lines[1].trim();
const separatorRegex = /^(\|[\s:]*-+[\s:]*)+\|$/;
if (!separatorRegex.test(separatorLine)) {
return false;
}
// 检查数据行是否包含 |
for (let i = 2; i < lines.length; i++) {
const dataLine = lines[i].trim();
if (dataLine && (!dataLine.startsWith('|') || !dataLine.endsWith('|'))) {
return false;
}
}
return true;
};
const markdownTableSplit = (props: SplitProps): SplitResponse => {
const { text = '', chunkSize, maxSize = defaultMaxChunkSize } = props;
// split by rows
const splitText2Lines = text.split('\n').filter((line) => line.trim());
// If there are not enough rows to form a table, return directly
if (splitText2Lines.length < 2) {
return { chunks: [text], chars: text.length };
}
const header = splitText2Lines[0];
const mdSplitString = splitText2Lines[1];
const chunks: string[] = [];
const defaultChunk = `${header}
${mdSplitString}
`;
let chunk = defaultChunk;
for (let i = 2; i < splitText2Lines.length; i++) {
const chunkLength = getTextValidLength(chunk);
const nextLineLength = getTextValidLength(splitText2Lines[i]);
// Over size
if (chunkLength + nextLineLength > chunkSize) {
// 单行非常的长，直接分割
if (chunkLength > maxSize) {
const newChunks = commonSplit({
...props,
text: chunk.replace(defaultChunk, '').trim()
}).chunks;
chunks.push(...newChunks);
} else {
chunks.push(chunk);
}
chunk = defaultChunk;
}
chunk += `${splitText2Lines[i]}\n`;
}
if (chunk) {
chunks.push(chunk);
}
return {
chunks,
chars: chunks.reduce((sum, chunk) => sum + chunk.length, 0)
};
};
/*
1. 自定义分隔符：不需要重叠，不需要小块合并
2. Markdown 标题：不需要重叠；标题嵌套共享，需要小块合并
3. 特殊 markdown 语法：不需要重叠，需要小块合并
4. 段落：尽可能保证它是一个完整的段落。
5. 标点分割：重叠
*/
const commonSplit = (props: SplitProps): SplitResponse => {
const {
text: rawText = '',
chunkSize,
paragraphChunkDeep = 5,
paragraphChunkMinSize = 100,
maxSize = defaultMaxChunkSize,
overlapRatio = 0.15,
customReg = []
} = props;
let text = rawText;
const splitMarker = 'SPLIT_HERE_SPLIT_HERE';
const codeBlockMarker = 'CODE_BLOCK_LINE_MARKER';
const overlapLen = Math.round(chunkSize * overlapRatio);
// 代码块需要尽量保留完整性，但不能直接使用模型 maxSize，否则大段正文包在 ```json/markdown``` 中会绕过 chunkSize 形成超大分块。
const maxCodeBlockChunks = 4;
const codeBlockMaxLen = Math.min(maxSize, chunkSize * maxCodeBlockChunks);
const strIsCodeBlock = (str: string) => /^(```[\s\S]*```|~~~[\s\S]*~~~)$/.test(str.trim());
// 特殊模块处理
// 1. 代码块处理 - 去除空字符
// replace code block all \n to codeBlockMarker
text = text.replace(/(```[\s\S]*?```|~~~[\s\S]*?~~~)/g, function (match) {
return match.replace(/\n/g, codeBlockMarker);
});
// replace invalid \n
text = text.replace(/(\r?\n|\r){3,}/g, '\n\n\n');
// The larger maxLen is, the next sentence is less likely to trigger splitting
const customRegLen = customReg.length;
const markdownIndex = paragraphChunkDeep - 1;
const forbidOverlapIndex = customRegLen + markdownIndex + 4;
const markdownHeaderRules = ((deep?: number): { reg: RegExp; maxLen: number }[] => {
if (!deep || deep === 0) return [];
const maxDeep = Math.min(deep, 8); // Maximum 8 levels
const rules: { reg: RegExp; maxLen: number }[] = [];
for (let i = 1; i <= maxDeep; i++) {
const hashSymbols = '#'.repeat(i);
rules.push({
reg: new RegExp(`^(${hashSymbols}\\s[^\\n]+\\n)`, 'gm'),
maxLen: chunkSize
});
}
return rules;
})(paragraphChunkDeep);
const stepReges: { reg: RegExp | string; maxLen: number; splitAround?: boolean }[] = [
...customReg.map((text) => ({
reg: text.replace(/\\n/g, '\n'),
maxLen: maxSize
})),
...markdownHeaderRules,
// 代码块需要独立成段，避免吞掉前面大段正文；短代码块仍尽量保持完整。
{ reg: /(^|\n)(```[\s\S]*?```|~~~[\s\S]*?~~~)/g, maxLen: codeBlockMaxLen, splitAround: true },
// HTML Table tag 尽可能保障完整
{
reg: /(\n\|(?:[^\n|]*\|)+\n\|(?:[:\-\s]*\|)+\n(?:\|(?:[^\n|]*\|)*\n)*)/g,
maxLen: chunkSize
}, // Markdown Table 尽可能保证完整性
{ reg: /(\n{2,})/g, maxLen: chunkSize },
{ reg: /([\n])/g, maxLen: chunkSize },
// ------ There's no overlap on the top
{ reg: /([。]|([a-zA-Z])\.\s)/g, maxLen: chunkSize },
{ reg: /([！]|!\s)/g, maxLen: chunkSize },
{ reg: /([？]|\?\s)/g, maxLen: chunkSize },
{ reg: /([；]|;\s)/g, maxLen: chunkSize },
{ reg: /([，]|,\s)/g, maxLen: chunkSize }
];
const checkIsCustomStep = (step: number) => step < customRegLen;
const checkIsMarkdownSplit = (step: number) =>
step >= customRegLen && step <= markdownIndex + customRegLen;
const checkForbidOverlap = (step: number) => step <= forbidOverlapIndex;
// if use markdown title split, Separate record title
const getSplitTexts = ({ text, step }: { text: string; step: number }) => {
if (step >= stepReges.length) {
return [
{
text,
title: '',
chunkMaxSize: chunkSize
}
];
}
const isCustomStep = checkIsCustomStep(step);
const isMarkdownSplit = checkIsMarkdownSplit(step);
const { reg, maxLen, splitAround } = stepReges[step];
const replaceText = (() => {
if (typeof reg === 'string') {
let tmpText = text;
reg.split('|').forEach((itemReg) => {
tmpText = tmpText.replaceAll(
itemReg,
(() => {
if (isCustomStep) return splitMarker;
if (isMarkdownSplit) return `${splitMarker}$1`;
return `$1${splitMarker}`;
})()
);
});
return tmpText;
}
return text.replace(
reg,
(() => {
if (isCustomStep) return splitMarker;
if (isMarkdownSplit) return `${splitMarker}$1`;
if (splitAround) return `${splitMarker}$&${splitMarker}`;
return `$1${splitMarker}`;
})()
);
})();
const splitTexts = replaceText.split(splitMarker).filter((part) => part.trim());
return splitTexts
.map((text) => {
const matchTitle = isMarkdownSplit ? text.match(reg)?.[0] || '' : '';
// 如果一个分块没有匹配到，则使用默认块大小，否则使用最大块大小
const chunkMaxSize = (() => {
if (isCustomStep) return maxLen;
return text.match(reg) === null ? chunkSize : maxLen;
})();
return {
text: isMarkdownSplit ? text.replace(matchTitle, '') : text,
title: matchTitle,
chunkMaxSize
};
})
.filter((item) => !!item.title || !!item.text?.trim());
};
/* Gets the overlap at the end of a text as the beginning of the next block */
const getOneTextOverlapText = ({ text, step }: { text: string; step: number }): string => {
const forbidOverlap = checkForbidOverlap(step);
const maxOverlapLen = chunkSize * 0.4;
// step >= stepReges.length: Do not overlap incomplete sentences
if (forbidOverlap || overlapLen === 0 || step >= stepReges.length) return '';
const splitTexts = getSplitTexts({ text, step });
let overlayText = '';
for (let i = splitTexts.length - 1; i >= 0; i--) {
const currentText = splitTexts[i].text;
const newText = currentText + overlayText;
const newTextLen = getTextValidLength(newText);
if (newTextLen > overlapLen) {
if (newTextLen > maxOverlapLen) {
const text = getOneTextOverlapText({ text: newText, step: step + 1 });
return text || overlayText;
}
return newText;
}
overlayText = newText;
}
return overlayText;
};
const splitTextRecursively = ({
text = '',
step,
lastText,
parentTitle = ''
}: {
text: string;
step: number;
lastText: string; // 上一个分块末尾数据会通过这个参数传入。
parentTitle: string;
}): string[] => {
const isMarkdownStep = checkIsMarkdownSplit(step);
const isCustomStep = checkIsCustomStep(step);
const forbidConcat = isCustomStep; // forbid=true时候，lastText肯定为空
// Over step
if (step >= stepReges.length) {
// Merge lastText with current text to prevent data loss
const combinedText = lastText + text;
const combinedLength = getTextValidLength(combinedText);
if (combinedLength < maxSize) {
return [combinedText];
}
// use slice-chunkSize to split text
// Note: Use combinedText.length for slicing, not combinedLength
const chunks: string[] = [];
for (let i = 0; i < combinedText.length; i += chunkSize - overlapLen) {
chunks.push(combinedText.slice(i, i + chunkSize));
}
return chunks;
}
// split text by special char
const splitTexts = getSplitTexts({ text, step });
const chunks: string[] = [];
for (let i = 0; i < splitTexts.length; i++) {
const item = splitTexts[i];
const maxLen = item.chunkMaxSize; // 当前块最大长度
const lastTextLen = getTextValidLength(lastText);
const currentText = item.text;
const newText = lastText + currentText;
const newTextLen = getTextValidLength(newText);
// 代码块独立处理，避免“前面正文 + 代码块”被 maxSize 合成超大分块。
if (strIsCodeBlock(currentText)) {
if (lastTextLen > 0) {
chunks.push(lastText);
lastText = '';
}
if (getTextValidLength(currentText) > maxLen) {
const restoredCodeBlock = currentText.replaceAll(codeBlockMarker, '\n');
for (let i = 0; i < restoredCodeBlock.length; i += chunkSize) {
chunks.push(restoredCodeBlock.slice(i, i + chunkSize));
}
} else {
chunks.push(currentText);
}
continue;
}
// split the current table if it will exceed after adding
if (strIsMdTable(currentText) && newTextLen > maxLen) {
if (lastTextLen > 0) {
chunks.push(lastText);
lastText = '';
}
const { chunks: tableChunks } = markdownTableSplit({
text: currentText,
chunkSize: chunkSize * 1.2
});
chunks.push(...tableChunks);
continue;
}
// Markdown 模式下，会强制向下拆分最小块，并再最后一个标题深度，给小块都补充上所有标题（包含父级标题）
if (isMarkdownStep) {
// split new Text, split chunks must will greater 1 (small lastText)
const innerChunks = splitTextRecursively({
text: newText,
step: step + 1,
lastText: '',
parentTitle: parentTitle + item.title
});
// 只有标题，没有内容。
if (innerChunks.length === 0) {
chunks.push(`${parentTitle}${item.title}`);
continue;
}
// 在合并最深级标题时，需要补充标题
chunks.push(
...innerChunks.map(
(chunk) =>
step === markdownIndex + customRegLen ? `${parentTitle}${item.title}${chunk}` : chunk // 合并进 Markdown 分块时，需要补标题
)
);
continue;
}
// newText is too large(now, The lastText must be smaller than chunkSize)
if (newTextLen > maxLen) {
const minChunkLen = maxLen * 0.8; // 当前块最小长度
const maxChunkLen = maxLen * 1.2; // 当前块最大长度
// 新文本没有非常大，直接认为它是一个新的块
if (newTextLen < maxChunkLen) {
chunks.push(newText);
lastText = getOneTextOverlapText({ text: newText, step }); // next chunk will start with overlayText
continue;
}
// 上一个文本块已经挺大的，单独做一个块
if (lastTextLen > minChunkLen) {
chunks.push(lastText);
lastText = getOneTextOverlapText({ text: lastText, step }); // next chunk will start with overlayText
i--;
continue;
}
// 说明是当前文本比较大，需要进一步拆分
// 把新的文本块进行一个拆分，并追加到 latestText 中
const innerChunks = splitTextRecursively({
text: currentText,
step: step + 1,
lastText,
parentTitle: parentTitle + item.title
});
const lastChunk = innerChunks[innerChunks.length - 1];
if (!lastChunk) continue;
// last chunk is too small, concat it to lastText(next chunk start)
if (getTextValidLength(lastChunk) < minChunkLen) {
chunks.push(...innerChunks.slice(0, -1));
lastText = lastChunk;
continue;
}
// Last chunk is large enough
chunks.push(...innerChunks);
// compute new overlapText
lastText = getOneTextOverlapText({
text: lastChunk,
step
});
continue;
}
// New text is small
// Not overlap
if (forbidConcat) {
chunks.push(currentText);
continue;
}
lastText = newText;
}
/* If the last chunk is independent, it needs to be push chunks. */
if (lastText && chunks[chunks.length - 1] && !chunks[chunks.length - 1].endsWith(lastText)) {
if (
getTextValidLength(lastText) < chunkSize * 0.4 &&
!strIsCodeBlock(chunks[chunks.length - 1])
) {
chunks[chunks.length - 1] = chunks[chunks.length - 1] + lastText;
} else {
chunks.push(lastText);
}
} else if (lastText && chunks.length === 0) {
// 只分出一个很小的块，则直接追加到末尾（如果大于 1 个块，说明这个小块内容已经被上一个块拿到了）
chunks.push(lastText);
}
return chunks;
};
try {
const chunks = splitTextRecursively({
text,
step: 0,
lastText: '',
parentTitle: ''
}).map((chunk) => chunk?.replaceAll(codeBlockMarker, '\n')?.trim() || ''); // restore code block
const chars = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
return {
chunks,
chars
};
} catch (err) {
throw new Error(getErrText(err));
}
};
/**
* text split into chunks
* chunkSize - one chunk len. max: 3500
* overlapLen - The size of the before and after Text
* chunkSize > overlapLen
* markdown
*/
export const splitText2Chunks = (props: SplitProps): SplitResponse => {
const { text = '' } = props;
const splitWithCustomSign = text.split(CUSTOM_SPLIT_SIGN);
const splitResult = splitWithCustomSign.map((item) => {
if (strIsMdTable(item)) {
return markdownTableSplit({ ...props, text: item });
}
return commonSplit({ ...props, text: item });
});
return {
chunks: splitResult
.map((item) => item.chunks)
.flat()
.map((chunk) => simpleText(chunk)),
chars: splitResult.reduce((sum, item) => sum + item.chars, 0)
};
};
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/common/string/textSplitter.ts 0 → 100644
import { defaultMaxChunkSize } from '@fastgpt/global/core/dataset/training/utils';
import { getErrText } from '@fastgpt/global/common/error/utils';
import { simpleText } from '@fastgpt/global/common/string/tools';
import { getTextValidLength } from '@fastgpt/global/common/string/utils';
import { countPromptTokensInWorker } from '../../worker/countGptMessagesTokens/count';
export const CUSTOM_SPLIT_SIGN = '-----CUSTOM_SPLIT_SIGN-----';
export type SplitProps = {
text: string;
chunkSize: number;
paragraphChunkDeep?: number; // Paragraph deep
paragraphChunkMinSize?: number; // Paragraph min size, if too small, it will merge
maxSize?: number;
overlapRatio?: number;
customReg?: string[];
lengthUnit?: 'char' | 'token';
maxChunks?: number;
};
export type TextSplitProps = Omit<SplitProps, 'text' | 'chunkSize'> & {
chunkSize?: number;
};
export type SplitResponse = {
chunks: string[];
chars: number;
};
type TextLengthCounter = (text: string) => number;
type SplitTextByLengthLimit = (props: {
text: string;
maxLength: number;
stepLength: number;
countLength: TextLengthCounter;
maxChunks?: number;
}) => string[];
const assertChunkLimit = (count: number, maxChunks?: number) => {
if (maxChunks !== undefined && count > maxChunks) {
throw new Error(`Text split exceeds the maximum chunk count of ${maxChunks}`);
}
};
const pushChunks = (target: string[], items: string[], maxChunks?: number) => {
assertChunkLimit(target.length + items.length, maxChunks);
target.push(...items);
};
const countOccurrencesUpTo = (text: string, search: string, limit: number) => {
let count = 0;
let start = 0;
while (start <= text.length) {
const index = text.indexOf(search, start);
if (index === -1) return count;
count++;
if (count >= limit) return count;
start = index + search.length;
}
return count;
};
const countRegexMatchesUpTo = (text: string, expression: RegExp, limit: number) => {
const regex = new RegExp(expression.source, expression.flags);
let count = 0;
if (!regex.global) {
return regex.test(text) ? 1 : 0;
}
while (count < limit) {
const match = regex.exec(text);
if (!match) return count;
count++;
if (match[0].length === 0) regex.lastIndex++;
}
return count;
};
const splitTextByCharLengthLimit: SplitTextByLengthLimit = ({
text,
maxLength,
stepLength,
maxChunks
}) => {
const chunks: string[] = [];
const chunkLength = Math.max(1, Math.floor(maxLength));
const chunkStep = Math.max(1, Math.floor(stepLength));
assertChunkLimit(Math.ceil(text.length / chunkStep), maxChunks);
for (let i = 0; i < text.length; i += chunkStep) {
chunks.push(text.slice(i, i + chunkLength));
}
return chunks;
};
const getMaxPrefixEndByLength = ({
textChars,
start,
maxLength,
countLength
}: {
textChars: string[];
start: number;
maxLength: number;
countLength: TextLengthCounter;
}) => {
if (start >= textChars.length || maxLength <= 0) return start;
let bestEnd = start;
let probeEnd = Math.min(textChars.length, start + Math.max(1, Math.floor(maxLength)));
let exceeded = false;
while (probeEnd <= textChars.length) {
const candidate = textChars.slice(start, probeEnd).join('');
if (countLength(candidate) > maxLength) {
exceeded = true;
break;
}
bestEnd = probeEnd;
if (probeEnd === textChars.length) return probeEnd;
const currentSpan = probeEnd - start;
probeEnd = Math.min(textChars.length, start + currentSpan * 2);
}
if (!exceeded) return bestEnd;
let left = bestEnd + 1;
let right = probeEnd - 1;
while (left <= right) {
const mid = Math.floor((left + right) / 2);
const candidate = textChars.slice(start, mid).join('');
if (countLength(candidate) <= maxLength) {
bestEnd = mid;
left = mid + 1;
} else {
right = mid - 1;
}
}
return bestEnd;
};
/**
* 从文本开头取不超过指定长度的最长前缀。
*
* token 模式下字符数和 token 数没有固定比例，不能直接用字符下标推算边界；
* 这里用二分查找减少 tokenizer 调用次数，并通过 Array.from 按 Unicode code point
* 切分，避免把代理对字符截断成非法字符串。
*
* 如果单个 code point 都超过 maxLength，会返回空字符串，由调用方决定是报错还是降级。
*/
export const getMaxPrefixByLength = ({
text,
maxLength,
countLength
}: {
text: string;
maxLength: number;
countLength: TextLengthCounter;
}) => {
const textChars = Array.from(text);
const bestEnd = getMaxPrefixEndByLength({
textChars,
start: 0,
maxLength,
countLength
});
return textChars.slice(0, bestEnd).join('');
};
/**
* 从文本结尾取不超过指定长度的最长后缀。
*
* 该函数专门服务于 overlap：下一块应复用上一块断点附近的尾部上下文，而不是上一块
* 开头内容。和前缀查找一样，这里按 Unicode code point 二分，避免 token 模式下用
* 字符长度误判边界或截断代理对字符。
*
* maxLength 小于等于 0 时没有可用 overlap 预算，直接返回空字符串。
*/
export const getMaxSuffixByLength = ({
text,
maxLength,
countLength
}: {
text: string;
maxLength: number;
countLength: TextLengthCounter;
}) => {
if (!text || maxLength <= 0) return '';
if (countLength(text) <= maxLength) return text;
const textChars = Array.from(text);
let left = 0;
let right = textChars.length - 1;
let bestStart = textChars.length;
while (left <= right) {
const mid = Math.floor((left + right) / 2);
const candidate = textChars.slice(mid).join('');
if (countLength(candidate) <= maxLength) {
bestStart = mid;
right = mid - 1;
} else {
left = mid + 1;
}
}
return textChars.slice(bestStart).join('');
};
const splitTextByCounterLengthLimit: SplitTextByLengthLimit = ({
text,
maxLength,
stepLength,
countLength,
maxChunks
}) => {
if (!text) return [];
if (!Number.isFinite(maxLength) || maxLength <= 0) return [text];
const chunks: string[] = [];
const textChars = Array.from(text);
let start = 0;
// token 模式下 overlap 也必须按同一个计数器计算，不能退回字符长度。
const overlapLength = Math.max(0, maxLength - Math.max(1, stepLength));
while (start < textChars.length) {
// 指针只在当前 chunk 附近做指数探测和二分，不再重复 tokenize 完整剩余文本。
const end = getMaxPrefixEndByLength({
textChars,
start,
maxLength,
countLength
});
const safeText = textChars.slice(start, end).join('');
if (!safeText) {
throw new Error('Text contains a character that exceeds the token length limit');
}
pushChunks(chunks, [safeText], maxChunks);
if (end >= textChars.length) break;
let nextStart = end;
if (overlapLength > 0) {
const overlapText = getMaxSuffixByLength({
text: safeText,
maxLength: overlapLength,
countLength
});
nextStart = end - Array.from(overlapText).length;
}
// 极端情况下 overlap 可能导致没有推进，直接丢弃 overlap 保证循环收敛。
start = nextStart <= start ? end : nextStart;
}
return chunks;
};
const getTextLengthCounter = (props: SplitProps): TextLengthCounter =>
props.lengthUnit === 'token' ? countPromptTokensInWorker : getTextValidLength;
const getSplitTextByLengthLimit = (props: SplitProps): SplitTextByLengthLimit =>
props.lengthUnit === 'token' ? splitTextByCounterLengthLimit : splitTextByCharLengthLimit;
// 判断字符串是否为markdown的表格形式
const strIsMdTable = (str: string) => {
// 检查是否包含表格分隔符 |
if (!str.includes('|')) {
return false;
}
const lines = str.split('\n');
// 检查表格是否至少有两行
if (lines.length < 2) {
return false;
}
// 检查表头行是否包含 |
const headerLine = lines[0].trim();
if (!headerLine.startsWith('|') || !headerLine.endsWith('|')) {
return false;
}
// 检查分隔行是否由 | 和 - 组成
const separatorLine = lines[1].trim();
const separatorRegex = /^(\|[\s:]*-+[\s:]*)+\|$/;
if (!separatorRegex.test(separatorLine)) {
return false;
}
// 检查数据行是否包含 |
for (let i = 2; i < lines.length; i++) {
const dataLine = lines[i].trim();
if (dataLine && (!dataLine.startsWith('|') || !dataLine.endsWith('|'))) {
return false;
}
}
return true;
};
const markdownTableSplit = (props: SplitProps): SplitResponse => {
const { text = '', chunkSize, maxSize = defaultMaxChunkSize, maxChunks } = props;
const countLength = getTextLengthCounter(props);
// split by rows
const splitText2Lines = text.split('\n').filter((line) => line.trim());
// If there are not enough rows to form a table, return directly
if (splitText2Lines.length < 2) {
return { chunks: [text], chars: text.length };
}
const header = splitText2Lines[0];
const mdSplitString = splitText2Lines[1];
const chunks: string[] = [];
const defaultChunk = `${header}
${mdSplitString}
`;
let chunk = defaultChunk;
// 只有表头和分隔行，没有数据行的 markdown table 不生成分块。
// 这种 chunk 没有可检索内容，继续入库只会生成空语义索引。
if (splitText2Lines.length === 2) {
return { chunks: [], chars: 0 };
}
/**
* token 模式下表格行拆分后还要补回表头；这里按“表头 + 内容”的最终文本
* 做二分兜底，避免只按行内容拆分后，拼回表头又超过 embedding 上限。
*/
const splitTextWithHeaderLimit = (text: string) => {
const result: string[] = [];
let restText = text;
while (restText) {
const restChars = Array.from(restText);
let left = 1;
let right = restChars.length;
let bestEnd = 0;
while (left <= right) {
const mid = Math.floor((left + right) / 2);
const candidate = restChars.slice(0, mid).join('');
if (countLength(`${defaultChunk}${candidate}`) <= chunkSize) {
bestEnd = mid;
left = mid + 1;
} else {
right = mid - 1;
}
}
if (bestEnd === 0) {
throw new Error('Markdown table header leaves no token budget for row content');
}
const safeText = restChars.slice(0, bestEnd).join('');
pushChunks(result, [`${defaultChunk}${safeText}`], maxChunks);
restText = restText.slice(safeText.length);
}
return result;
};
const splitTableLineWithHeader = (line: string) => {
const contentChunkSize = chunkSize - countLength(defaultChunk);
if (contentChunkSize <= 0) {
throw new Error('Markdown table header exceeds token chunk size');
}
return commonSplit({
...props,
text: line,
chunkSize: contentChunkSize,
maxSize: Math.min(maxSize, contentChunkSize)
}).chunks.flatMap(splitTextWithHeaderLimit);
};
for (let i = 2; i < splitText2Lines.length; i++) {
const chunkLength = countLength(chunk);
const nextLineLength = countLength(splitText2Lines[i]);
const defaultChunkLength = countLength(defaultChunk);
if (props.lengthUnit === 'token' && defaultChunkLength + nextLineLength > chunkSize) {
if (chunk !== defaultChunk) {
pushChunks(chunks, [chunk], maxChunks);
}
pushChunks(chunks, splitTableLineWithHeader(splitText2Lines[i]), maxChunks);
chunk = defaultChunk;
continue;
}
// Over size
if (chunkLength + nextLineLength > chunkSize) {
// 单行非常的长，直接分割
if (chunkLength > maxSize) {
const newChunks = commonSplit({
...props,
text: chunk.replace(defaultChunk, '').trim()
}).chunks;
pushChunks(chunks, newChunks, maxChunks);
} else if (chunk !== defaultChunk) {
// 第一条表格数据行就超过 chunkSize 时，chunk 仍然只有表头。
// 这时不能先推出 header-only chunk，否则会生成没有可检索内容的分块。
pushChunks(chunks, [chunk], maxChunks);
}
chunk = defaultChunk;
}
chunk += `${splitText2Lines[i]}\n`;
}
if (chunk && (chunk !== defaultChunk || chunks.length === 0)) {
pushChunks(chunks, [chunk], maxChunks);
}
return {
chunks,
chars: chunks.reduce((sum, chunk) => sum + chunk.length, 0)
};
};
/*
1. 自定义分隔符：不需要重叠，不需要小块合并
2. Markdown 标题：不需要重叠；标题嵌套共享，需要小块合并
3. 特殊 markdown 语法：不需要重叠，需要小块合并
4. 段落：尽可能保证它是一个完整的段落。
5. 标点分割：重叠
*/
const commonSplit = (props: SplitProps): SplitResponse => {
const {
text: rawText = '',
chunkSize,
paragraphChunkDeep = 5,
paragraphChunkMinSize = 100,
maxSize = defaultMaxChunkSize,
overlapRatio = 0.15,
customReg = [],
maxChunks
} = props;
if (!Number.isFinite(chunkSize) || chunkSize <= 0) {
throw new Error('Chunk size must be a positive finite number');
}
if (!Number.isFinite(overlapRatio) || overlapRatio < 0 || overlapRatio >= 1) {
throw new Error('Overlap ratio must be greater than or equal to 0 and less than 1');
}
if (maxChunks !== undefined && (!Number.isInteger(maxChunks) || maxChunks <= 0)) {
throw new Error('Maximum chunk count must be a positive integer');
}
customReg.forEach((reg) => {
if (
reg
.replace(/\\n/g, '\n')
.split('|')
.some((item) => item.length === 0)
) {
throw new Error('Custom split separators cannot be empty');
}
});
const countLength = getTextLengthCounter(props);
const splitTextByLengthLimit = getSplitTextByLengthLimit(props);
let text = rawText;
const splitMarker = 'SPLIT_HERE_SPLIT_HERE';
const codeBlockMarker = 'CODE_BLOCK_LINE_MARKER';
const overlapLen = Math.round(chunkSize * overlapRatio);
// 代码块需要尽量保留完整性，但不能直接使用模型 maxSize，否则大段正文包在 ```json/markdown``` 中会绕过 chunkSize 形成超大分块。
const maxCodeBlockChunks = 4;
const codeBlockMaxLen = Math.min(maxSize, chunkSize * maxCodeBlockChunks);
const strIsCodeBlock = (str: string) => /^(```[\s\S]*```|~~~[\s\S]*~~~)$/.test(str.trim());
// 特殊模块处理
// 1. 代码块处理 - 去除空字符
// replace code block all \n to codeBlockMarker
text = text.replace(/(```[\s\S]*?```|~~~[\s\S]*?~~~)/g, function (match) {
return match.replace(/\n/g, codeBlockMarker);
});
// replace invalid \n
text = text.replace(/(\r?\n|\r){3,}/g, '\n\n\n');
// The larger maxLen is, the next sentence is less likely to trigger splitting
const customRegLen = customReg.length;
const markdownIndex = paragraphChunkDeep - 1;
const forbidOverlapIndex = customRegLen + markdownIndex + 4;
const markdownHeaderRules = ((deep?: number): { reg: RegExp; maxLen: number }[] => {
if (!deep || deep === 0) return [];
const maxDeep = Math.min(deep, 8); // Maximum 8 levels
const rules: { reg: RegExp; maxLen: number }[] = [];
for (let i = 1; i <= maxDeep; i++) {
const hashSymbols = '#'.repeat(i);
rules.push({
reg: new RegExp(`^(${hashSymbols}\\s[^\\n]+\\n)`, 'gm'),
maxLen: chunkSize
});
}
return rules;
})(paragraphChunkDeep);
const stepReges: { reg: RegExp | string; maxLen: number; splitAround?: boolean }[] = [
...customReg.map((text) => ({
reg: text.replace(/\\n/g, '\n'),
maxLen: maxSize
})),
...markdownHeaderRules,
// 代码块需要独立成段，避免吞掉前面大段正文；短代码块仍尽量保持完整。
{ reg: /(^|\n)(```[\s\S]*?```|~~~[\s\S]*?~~~)/g, maxLen: codeBlockMaxLen, splitAround: true },
// HTML Table tag 尽可能保障完整
{
reg: /(\n\|(?:[^\n|]*\|)+\n\|(?:[:\-\s]*\|)+\n(?:\|(?:[^\n|]*\|)*\n)*)/g,
maxLen: chunkSize
}, // Markdown Table 尽可能保证完整性
{ reg: /(\n{2,})/g, maxLen: chunkSize },
{ reg: /([\n])/g, maxLen: chunkSize },
// ------ There's no overlap on the top
{ reg: /([。]|([a-zA-Z])\.\s)/g, maxLen: chunkSize },
{ reg: /([！]|!\s)/g, maxLen: chunkSize },
{ reg: /([？]|\?\s)/g, maxLen: chunkSize },
{ reg: /([；]|;\s)/g, maxLen: chunkSize },
{ reg: /([，]|,\s)/g, maxLen: chunkSize }
];
const checkIsCustomStep = (step: number) => step < customRegLen;
const checkIsMarkdownSplit = (step: number) =>
step >= customRegLen && step <= markdownIndex + customRegLen;
const checkForbidOverlap = (step: number) => step <= forbidOverlapIndex;
// if use markdown title split, Separate record title
const getSplitTexts = ({ text, step }: { text: string; step: number }) => {
if (step >= stepReges.length) {
return [
{
text,
title: '',
chunkMaxSize: chunkSize
}
];
}
const isCustomStep = checkIsCustomStep(step);
const isMarkdownSplit = checkIsMarkdownSplit(step);
const { reg, maxLen, splitAround } = stepReges[step];
const replaceText = (() => {
if (typeof reg === 'string') {
let tmpText = text;
reg.split('|').forEach((itemReg) => {
if (maxChunks !== undefined) {
const occurrenceCount = countOccurrencesUpTo(tmpText, itemReg, maxChunks);
assertChunkLimit(occurrenceCount + 1, maxChunks);
}
tmpText = tmpText.replaceAll(
itemReg,
(() => {
if (isCustomStep) return splitMarker;
if (isMarkdownSplit) return `${splitMarker}$1`;
return `$1${splitMarker}`;
})()
);
});
return tmpText;
}
if (maxChunks !== undefined) {
const markersPerMatch = splitAround ? 2 : 1;
const matchCount = countRegexMatchesUpTo(text, reg, Math.ceil(maxChunks / markersPerMatch));
assertChunkLimit(matchCount * markersPerMatch + 1, maxChunks);
}
return text.replace(
reg,
(() => {
if (isCustomStep) return splitMarker;
if (isMarkdownSplit) return `${splitMarker}$1`;
if (splitAround) return `${splitMarker}$&${splitMarker}`;
return `$1${splitMarker}`;
})()
);
})();
if (maxChunks !== undefined) {
const markerCount = countOccurrencesUpTo(replaceText, splitMarker, maxChunks);
assertChunkLimit(markerCount + 1, maxChunks);
}
const splitTexts = replaceText.split(splitMarker).filter((part) => part.trim());
return splitTexts
.map((text) => {
const matchTitle = isMarkdownSplit ? text.match(reg)?.[0] || '' : '';
// 如果一个分块没有匹配到，则使用默认块大小，否则使用最大块大小
const chunkMaxSize = (() => {
if (isCustomStep) return maxLen;
return text.match(reg) === null ? chunkSize : maxLen;
})();
return {
text: isMarkdownSplit ? text.replace(matchTitle, '') : text,
title: matchTitle,
chunkMaxSize
};
})
.filter((item) => !!item.title || !!item.text?.trim());
};
/* Gets the overlap at the end of a text as the beginning of the next block */
const getOneTextOverlapText = ({ text, step }: { text: string; step: number }): string => {
const forbidOverlap = checkForbidOverlap(step);
const maxOverlapLen = chunkSize * 0.4;
// step >= stepReges.length: Do not overlap incomplete sentences
if (forbidOverlap || overlapLen === 0 || step >= stepReges.length) return '';
const splitTexts = getSplitTexts({ text, step });
let overlayText = '';
for (let i = splitTexts.length - 1; i >= 0; i--) {
const currentText = splitTexts[i].text;
const newText = currentText + overlayText;
const newTextLen = countLength(newText);
if (newTextLen > overlapLen) {
if (newTextLen > maxOverlapLen) {
const text = getOneTextOverlapText({ text: newText, step: step + 1 });
return text || overlayText;
}
return newText;
}
overlayText = newText;
}
return overlayText;
};
const splitTextRecursively = ({
text = '',
step,
lastText,
parentTitle = ''
}: {
text: string;
step: number;
lastText: string; // 上一个分块末尾数据会通过这个参数传入。
parentTitle: string;
}): string[] => {
const isMarkdownStep = checkIsMarkdownSplit(step);
const isCustomStep = checkIsCustomStep(step);
const forbidConcat = isCustomStep; // forbid=true时候，lastText肯定为空
let lastTextIsOverlap = false;
// Over step
if (step >= stepReges.length) {
// Merge lastText with current text to prevent data loss
const combinedText = lastText + text;
const combinedLength = countLength(combinedText);
if (combinedLength < maxSize) {
return [combinedText];
}
return splitTextByLengthLimit({
text: combinedText,
maxLength: chunkSize,
stepLength: chunkSize - overlapLen,
countLength,
maxChunks
});
}
// split text by special char
const splitTexts = getSplitTexts({ text, step });
const chunks: string[] = [];
for (let i = 0; i < splitTexts.length; i++) {
const item = splitTexts[i];
const maxLen = item.chunkMaxSize; // 当前块最大长度
const lastTextLen = countLength(lastText);
const currentText = item.text;
const newText = lastText + currentText;
const newTextLen = countLength(newText);
// 代码块独立处理，避免“前面正文 + 代码块”被 maxSize 合成超大分块。
if (strIsCodeBlock(currentText)) {
if (lastTextLen > 0) {
pushChunks(chunks, [lastText], maxChunks);
lastText = '';
lastTextIsOverlap = false;
}
if (countLength(currentText) > maxLen) {
const restoredCodeBlock = currentText.replaceAll(codeBlockMarker, '\n');
pushChunks(
chunks,
splitTextByLengthLimit({
text: restoredCodeBlock,
maxLength: chunkSize,
stepLength: chunkSize,
countLength,
maxChunks
}),
maxChunks
);
} else {
pushChunks(chunks, [currentText], maxChunks);
}
continue;
}
// split the current table if it will exceed after adding
if (strIsMdTable(currentText) && newTextLen > maxLen) {
if (lastTextLen > 0) {
pushChunks(chunks, [lastText], maxChunks);
lastText = '';
lastTextIsOverlap = false;
}
const { chunks: tableChunks } = markdownTableSplit({
text: currentText,
chunkSize: props.lengthUnit === 'token' ? chunkSize : chunkSize * 1.2,
maxSize,
lengthUnit: props.lengthUnit,
maxChunks
});
pushChunks(chunks, tableChunks, maxChunks);
continue;
}
// Markdown 模式下，会强制向下拆分最小块，并再最后一个标题深度，给小块都补充上所有标题（包含父级标题）
if (isMarkdownStep) {
// split new Text, split chunks must will greater 1 (small lastText)
const innerChunks = splitTextRecursively({
text: newText,
step: step + 1,
lastText: '',
parentTitle: parentTitle + item.title
});
// 只有标题，没有内容。
if (innerChunks.length === 0) {
pushChunks(chunks, [`${parentTitle}${item.title}`], maxChunks);
continue;
}
// 在合并最深级标题时，需要补充标题
pushChunks(
chunks,
innerChunks.map(
(chunk) =>
step === markdownIndex + customRegLen ? `${parentTitle}${item.title}${chunk}` : chunk // 合并进 Markdown 分块时，需要补标题
),
maxChunks
);
continue;
}
// newText is too large(now, The lastText must be smaller than chunkSize)
if (newTextLen > maxLen) {
const minChunkLen = maxLen * 0.8; // 当前块最小长度
const maxChunkLen = maxLen * 1.2; // 当前块最大长度
// 新文本没有非常大，直接认为它是一个新的块
if (newTextLen < maxChunkLen && (props.lengthUnit !== 'token' || newTextLen <= maxSize)) {
pushChunks(chunks, [newText], maxChunks);
lastText = getOneTextOverlapText({ text: newText, step }); // next chunk will start with overlayText
lastTextIsOverlap = true;
continue;
}
// 上一个文本块已经挺大的，单独做一个块
if (lastTextLen > minChunkLen) {
pushChunks(chunks, [lastText], maxChunks);
lastText = getOneTextOverlapText({ text: lastText, step }); // next chunk will start with overlayText
lastTextIsOverlap = true;
i--;
continue;
}
// 说明是当前文本比较大，需要进一步拆分
// 把新的文本块进行一个拆分，并追加到 latestText 中
const innerChunks = splitTextRecursively({
text: currentText,
step: step + 1,
lastText,
parentTitle: parentTitle + item.title
});
const lastChunk = innerChunks[innerChunks.length - 1];
if (!lastChunk) continue;
// last chunk is too small, concat it to lastText(next chunk start)
if (countLength(lastChunk) < minChunkLen) {
pushChunks(chunks, innerChunks.slice(0, -1), maxChunks);
lastText = lastChunk;
lastTextIsOverlap = false;
continue;
}
// Last chunk is large enough
pushChunks(chunks, innerChunks, maxChunks);
// compute new overlapText
lastText = getOneTextOverlapText({
text: lastChunk,
step
});
lastTextIsOverlap = true;
continue;
}
// New text is small
// Not overlap
if (forbidConcat) {
pushChunks(chunks, [currentText], maxChunks);
continue;
}
lastText = newText;
lastTextIsOverlap = false;
}
/* If the last chunk is independent, it needs to be push chunks. */
const lastChunk = chunks[chunks.length - 1];
const shouldPushLastText =
props.lengthUnit === 'token'
? !lastTextIsOverlap || !lastChunk?.endsWith(lastText)
: !lastChunk?.endsWith(lastText);
if (lastText && lastChunk && shouldPushLastText) {
if (
countLength(lastText) < chunkSize * 0.4 &&
!strIsCodeBlock(lastChunk) &&
(props.lengthUnit !== 'token' || countLength(lastChunk + lastText) <= maxSize)
) {
chunks[chunks.length - 1] = lastChunk + lastText;
} else {
pushChunks(chunks, [lastText], maxChunks);
}
} else if (lastText && chunks.length === 0) {
// 只分出一个很小的块，则直接追加到末尾（如果大于 1 个块，说明这个小块内容已经被上一个块拿到了）
pushChunks(chunks, [lastText], maxChunks);
}
return chunks;
};
try {
const chunks = splitTextRecursively({
text,
step: 0,
lastText: '',
parentTitle: ''
}).map((chunk) => chunk?.replaceAll(codeBlockMarker, '\n')?.trim() || ''); // restore code block
const chars = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
return {
chunks,
chars
};
} catch (err) {
throw new Error(getErrText(err));
}
};
/**
* text split into chunks
* chunkSize - one chunk len. max: 3500
* overlapLen - The size of the before and after Text
* chunkSize > overlapLen
* markdown
*/
export const splitText2Chunks = (props: SplitProps): SplitResponse => {
const { text = '', maxChunks } = props;
if (maxChunks !== undefined) {
const customSignCount = countOccurrencesUpTo(text, CUSTOM_SPLIT_SIGN, maxChunks);
assertChunkLimit(customSignCount + 1, maxChunks);
}
const splitWithCustomSign = text.split(CUSTOM_SPLIT_SIGN);
const splitResult = splitWithCustomSign.map((item) => {
if (strIsMdTable(item)) {
return markdownTableSplit({ ...props, text: item });
}
return commonSplit({ ...props, text: item });
});
const chunks = splitResult
.map((item) => item.chunks)
.flat()
.map((chunk) => simpleText(chunk));
assertChunkLimit(chunks.length, maxChunks);
return {
chunks,
chars: splitResult.reduce((sum, item) => sum + item.chars, 0)
};
};
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,
chunkSize: indexSize, indexSize,
maxSize: maxIndexSize ?? this.maxToken maxToken: maxIndexSize ?? this.maxToken,
}) indexPrefix
).chunks; });
const aChunks = a const aIndexTexts = a
? ( ? await splitIndexTextByTokenLimit({
await text2Chunks({
text: a, text: a,
chunkSize: indexSize, indexSize,
maxSize: maxIndexSize ?? this.maxToken maxToken: maxIndexSize ?? this.maxToken,
indexPrefix
}) })
).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)) {
const splitText = (
await text2Chunks({
text: item.text, text: item.text,
chunkSize: indexSize, indexSize,
maxSize: maxIndexSize ?? this.maxToken maxToken: maxIndexSize ?? this.maxToken,
}) indexPrefix: item.type === DatasetDataIndexTypeEnum.default ? indexPrefix : undefined
).chunks; });
return splitText.map((text) => ({
if (indexTexts.length > 1 || indexTexts[0] !== item.text) {
return indexTexts.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();
});
});
projects/app/test/service/core/dataset/data/dataIndex.test.ts
...@@ -14,6 +14,8 @@ import { getRootUser } from '@test/datas/users'; ...@@ -14,6 +14,8 @@ import { getRootUser } from '@test/datas/users';
import { mockGetVectors, createMockVectorsResponse } from '@test/mocks/core/ai/embedding'; import { mockGetVectors, createMockVectorsResponse } from '@test/mocks/core/ai/embedding';
import { mockVectorDelete, mockVectorInsert, resetVectorMocks } from '@test/mocks/common/vector'; import { mockVectorDelete, mockVectorInsert, resetVectorMocks } from '@test/mocks/common/vector';
import { serviceEnv } from '@fastgpt/service/env'; import { serviceEnv } from '@fastgpt/service/env';
import { countPromptTokensInWorker } from '@fastgpt/service/worker/countGptMessagesTokens/count';
import { minChunkSize } from '@fastgpt/global/core/dataset/training/utils';
import { import {
createDatasetDataIndex, createDatasetDataIndex,
DatasetDataIndexOperation, DatasetDataIndexOperation,
...@@ -22,7 +24,7 @@ import { ...@@ -22,7 +24,7 @@ import {
} from '@/service/core/dataset/data/dataIndex'; } from '@/service/core/dataset/data/dataIndex';
const { mockCountPromptTokens } = vi.hoisted(() => ({ const { mockCountPromptTokens } = vi.hoisted(() => ({
mockCountPromptTokens: vi.fn(async (text: string) => text.length) mockCountPromptTokens: vi.fn()
})); }));
const { mockGetDatasetBase64Image } = vi.hoisted(() => ({ const { mockGetDatasetBase64Image } = vi.hoisted(() => ({
...@@ -39,12 +41,39 @@ vi.mock('@fastgpt/service/common/string/tiktoken', () => ({ ...@@ -39,12 +41,39 @@ vi.mock('@fastgpt/service/common/string/tiktoken', () => ({
countPromptTokens: mockCountPromptTokens countPromptTokens: mockCountPromptTokens
})); }));
vi.mock('@fastgpt/service/common/string/tiktoken/index', () => ({
countPromptTokens: mockCountPromptTokens
}));
const embeddingModel = { const embeddingModel = {
model: 'text-embedding-3-small', model: 'text-embedding-3-small',
name: 'text-embedding-3-small', name: 'text-embedding-3-small',
maxToken: 12 maxToken: 128
} as any; } as any;
const originalMultipleDataToBase64 = serviceEnv.MULTIPLE_DATA_TO_BASE64; const originalMultipleDataToBase64 = serviceEnv.MULTIPLE_DATA_TO_BASE64;
const tokenHeavyText = Array.from({ length: 30 }, (_, index) => `𠮷${index}`).join('');
const largeTokenHeavyText = Array.from(
{ length: 320 },
(_, index) => `第${index}段𠮷内容${index}`
).join('');
const expectIndexesWithinTokenLimit = (
indexes: Pick<DatasetDataIndexItemType, 'text'>[],
maxToken: number
) => {
expect(indexes.every((index) => countPromptTokensInWorker(index.text) <= maxToken)).toBe(true);
};
const mergeChunksByOverlap = (chunks: string[]) =>
chunks.reduce((mergedText, chunk) => {
const maxOverlapLength = Math.min(mergedText.length, chunk.length);
let overlapLength = maxOverlapLength;
while (overlapLength > 0 && !mergedText.endsWith(chunk.slice(0, overlapLength))) {
overlapLength--;
}
return `${mergedText}${chunk.slice(overlapLength)}`;
}, '');
const createDatasetContext = async () => { const createDatasetContext = async () => {
const root = await getRootUser(); const root = await getRootUser();
...@@ -117,7 +146,10 @@ describe('DatasetDataIndexOperation', () => { ...@@ -117,7 +146,10 @@ describe('DatasetDataIndexOperation', () => {
resetVectorMocks(); resetVectorMocks();
mockGetVectors.mockClear(); mockGetVectors.mockClear();
mockGetDatasetBase64Image.mockClear(); mockGetDatasetBase64Image.mockClear();
mockCountPromptTokens.mockClear(); mockCountPromptTokens.mockReset();
mockCountPromptTokens.mockImplementation(async (text: string) =>
countPromptTokensInWorker(text)
);
vi.mocked(getEmbeddingModel).mockReturnValue(embeddingModel); vi.mocked(getEmbeddingModel).mockReturnValue(embeddingModel);
mockGetVectors.mockImplementation(async ({ inputs }) => mockGetVectors.mockImplementation(async ({ inputs }) =>
createMockVectorsResponse(inputs.map((input) => input.input)) createMockVectorsResponse(inputs.map((input) => input.input))
...@@ -173,6 +205,7 @@ describe('DatasetDataIndexOperation', () => { ...@@ -173,6 +205,7 @@ describe('DatasetDataIndexOperation', () => {
q: 'question text', q: 'question text',
a: 'answer text', a: 'answer text',
indexSize: 50, indexSize: 50,
maxIndexSize: 200,
indexPrefix: 'collection title' indexPrefix: 'collection title'
}); });
...@@ -188,6 +221,171 @@ describe('DatasetDataIndexOperation', () => { ...@@ -188,6 +221,171 @@ describe('DatasetDataIndexOperation', () => {
]); ]);
}); });
it('should clamp default index size to minimum chunk size when text is below embedding limit', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const result = await operation.getSystemIndexes({
q: tokenHeavyText,
indexSize: 12,
maxIndexSize: 200
});
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.default)).toBe(true);
expect(result.some((index) => countPromptTokensInWorker(index.text) > 12)).toBe(true);
expect(result.every((index) => countPromptTokensInWorker(index.text) <= minChunkSize)).toBe(
true
);
expectIndexesWithinTokenLimit(result, 200);
expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(tokenHeavyText);
});
it('should clamp prefixed default index content to minimum chunk size when embedding limit is larger', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const indexPrefix = '# LongTitle';
const result = await operation.getSystemIndexes({
q: tokenHeavyText,
indexSize: 12,
maxIndexSize: 200,
indexPrefix
});
const contentChunks = result.map((index) => index.text.replace(`${indexPrefix}\n`, ''));
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.text.startsWith(`${indexPrefix}\n`))).toBe(true);
expect(contentChunks.some((chunk) => countPromptTokensInWorker(chunk) > 12)).toBe(true);
expect(contentChunks.every((chunk) => countPromptTokensInWorker(chunk) <= minChunkSize)).toBe(
true
);
expectIndexesWithinTokenLimit(result, 200);
expect(mergeChunksByOverlap(contentChunks)).toBe(tokenHeavyText);
});
it('should split default text indexes when text2Chunks result still exceeds embedding limit', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const result = await operation.getSystemIndexes({
q: tokenHeavyText,
indexSize: 100,
maxIndexSize: 70
});
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.default)).toBe(true);
expectIndexesWithinTokenLimit(result, 70);
expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(tokenHeavyText);
});
it('should split an overlong answer into token-safe default indexes', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const result = await operation.getSystemIndexes({
q: '',
a: tokenHeavyText,
indexSize: 100,
maxIndexSize: 70
});
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.default)).toBe(true);
expectIndexesWithinTokenLimit(result, 70);
expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(tokenHeavyText);
});
it('should split large default text indexes by indexSize when model limit is larger', async () => {
const indexSize = 96;
const operation = new DatasetDataIndexOperation({
...embeddingModel,
maxToken: 512
});
const result = await operation.getSystemIndexes({
q: largeTokenHeavyText,
indexSize,
maxIndexSize: 512
});
expect(countPromptTokensInWorker(largeTokenHeavyText)).toBeGreaterThan(indexSize * 10);
expect(result.length).toBeGreaterThan(10);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.default)).toBe(true);
expect(result.some((index) => countPromptTokensInWorker(index.text) > minChunkSize)).toBe(
true
);
expect(result.every((index) => countPromptTokensInWorker(index.text) <= indexSize)).toBe(
true
);
expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(largeTokenHeavyText);
});
it('should split prefixed default indexes by final embedding token size', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const indexPrefix = '# LongTitle';
const result = await operation.getSystemIndexes({
q: tokenHeavyText,
indexSize: 100,
maxIndexSize: 70,
indexPrefix
});
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.text.startsWith(`${indexPrefix}\n`))).toBe(true);
expectIndexesWithinTokenLimit(result, 70);
expect(
mergeChunksByOverlap(result.map((index) => index.text.replace(`${indexPrefix}\n`, '')))
).toBe(tokenHeavyText);
});
it('should fail fast when prefix leaves no room for index content', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
await expect(
operation.getSystemIndexes({
q: 'abc',
indexSize: 50,
maxIndexSize: 12,
indexPrefix: '𠮷'.repeat(4)
})
).rejects.toThrow('Dataset index prefix is too long for embedding token limit');
});
it('should fail fast when prefix leaves less than minimum index content budget', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
await expect(
operation.getSystemIndexes({
q: 'abc',
indexSize: 50,
maxIndexSize: 20,
indexPrefix: '# LongTitle'
})
).rejects.toThrow('Dataset index content token budget is smaller than minimum chunk size');
});
it('should not check prefix token budget when text is empty and only image index is generated', async () => {
const operation = new DatasetDataIndexOperation({
...embeddingModel,
vision: true
});
const result = await operation.getSystemIndexes({
q: '',
imageId: 'dataset/team/main.png',
indexSize: 50,
maxIndexSize: 12,
indexPrefix: '𠮷'.repeat(4)
});
expect(mockCountPromptTokens).not.toHaveBeenCalled();
expect(result).toEqual([
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/main.png'
}
]);
});
it('should create image embedding indexes from image id and markdown images', async () => { it('should create image embedding indexes from image id and markdown images', async () => {
const operation = new DatasetDataIndexOperation({ const operation = new DatasetDataIndexOperation({
...embeddingModel, ...embeddingModel,
...@@ -350,50 +548,94 @@ describe('DatasetDataIndexOperation', () => { ...@@ -350,50 +548,94 @@ describe('DatasetDataIndexOperation', () => {
it('should split a custom index when token count exceeds max token', async () => { it('should split a custom index when token count exceeds max token', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel); const operation = new DatasetDataIndexOperation(embeddingModel);
mockCountPromptTokens.mockResolvedValueOnce(30);
const result = await operation.formatIndexes({ const result = await operation.formatIndexes({
q: '', q: '',
a: '', a: '',
indexSize: 8, indexSize: 8,
maxIndexSize: 12, maxIndexSize: 70,
indexes: [ indexes: [
{ {
type: DatasetDataIndexTypeEnum.custom, type: DatasetDataIndexTypeEnum.custom,
text: 'first sentence. second sentence. third sentence.' text: tokenHeavyText
} }
] ]
}); });
expect(result.length).toBeGreaterThan(1); expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.custom)).toBe(true); expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.custom)).toBe(true);
const mergedText = result.map((index) => index.text).join(' '); expectIndexesWithinTokenLimit(result, 70);
expect(mergedText).toContain('first'); expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(tokenHeavyText);
expect(mergedText).toContain('third');
}); });
it('should check default index token size after prefix is applied', async () => { it('should keep short custom indexes unchanged before token split', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel); const operation = new DatasetDataIndexOperation(embeddingModel);
mockCountPromptTokens.mockResolvedValueOnce(21); const customText = ' first second\n\n\nthird ';
const result = await operation.formatIndexes({ const result = await operation.formatIndexes({
q: 'short content', q: '',
a: '', a: '',
indexSize: 10, indexSize: 50,
maxIndexSize: 20, maxIndexSize: 100,
indexPrefix: '# LongTitle', indexes: [
indexes: [] {
type: DatasetDataIndexTypeEnum.custom,
text: customText
}
]
}); });
expect(mockCountPromptTokens).toHaveBeenCalledWith('# LongTitle\nshort content');
expect(result).toEqual([ expect(result).toEqual([
{ {
type: DatasetDataIndexTypeEnum.default, type: DatasetDataIndexTypeEnum.custom,
text: '# LongTitle\nshort content' text: customText
} }
]); ]);
}); });
it('should force split custom indexes to embedding-safe token size after text2Chunks', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const result = await operation.formatIndexes({
q: '',
a: '',
indexSize: 100,
maxIndexSize: 70,
indexes: [
{
type: DatasetDataIndexTypeEnum.custom,
text: tokenHeavyText
}
]
});
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.type === DatasetDataIndexTypeEnum.custom)).toBe(true);
expectIndexesWithinTokenLimit(result, 70);
expect(mergeChunksByOverlap(result.map((index) => index.text))).toBe(tokenHeavyText);
});
it('should check default index token size after prefix is applied', async () => {
const operation = new DatasetDataIndexOperation(embeddingModel);
const result = await operation.formatIndexes({
q: tokenHeavyText,
a: '',
indexSize: 100,
maxIndexSize: 70,
indexPrefix: '# LongTitle',
indexes: []
});
expect(mockCountPromptTokens).toHaveBeenCalledWith('# LongTitle\n');
expect(result.length).toBeGreaterThan(1);
expect(result.every((index) => index.text.startsWith('# LongTitle\n'))).toBe(true);
expectIndexesWithinTokenLimit(result, 70);
expect(
mergeChunksByOverlap(result.map((index) => index.text.replace('# LongTitle\n', '')))
).toBe(tokenHeavyText);
});
it('should keep image embedding indexes unsplit even when text looks too long', async () => { it('should keep image embedding indexes unsplit even when text looks too long', async () => {
const operation = new DatasetDataIndexOperation({ const operation = new DatasetDataIndexOperation({
...embeddingModel, ...embeddingModel,
...@@ -859,7 +1101,7 @@ describe('DatasetDataIndexOperation', () => { ...@@ -859,7 +1101,7 @@ describe('DatasetDataIndexOperation', () => {
it('should reject custom index text longer than model maxToken', async () => { it('should reject custom index text longer than model maxToken', async () => {
const { dataItem } = await createData(); const { dataItem } = await createData();
mockCountPromptTokens.mockResolvedValueOnce(13); mockCountPromptTokens.mockResolvedValueOnce(129);
await expect( await expect(
createDatasetDataIndex({ createDatasetDataIndex({
...@@ -973,7 +1215,7 @@ describe('DatasetDataIndexOperation', () => { ...@@ -973,7 +1215,7 @@ describe('DatasetDataIndexOperation', () => {
it('should use the resolved embedding model when only a model name is provided', () => { it('should use the resolved embedding model when only a model name is provided', () => {
const operation = new DatasetDataIndexOperation('unknown-model'); const operation = new DatasetDataIndexOperation('unknown-model');
expect(operation.maxToken).toBe(12); expect(operation.maxToken).toBe(128);
}); });
it('keeps object id generation available for data fixtures', () => { it('keeps object id generation available for data fixtures', () => {
......
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