Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit 0f55590a by YeYuheng Committed by GitHub
Options

Support image embedding indexes for multimodal dataset search (#6948) 

* support image embedding dataset indexes

* feat(dataset): support image embedding system indexes

* fix: migrate legacy dataset search input

* perf: code

* perf: code

* submodule

* add test

* fix: migrate dataset search input

* doc

* test: update dataset search input expectation

---------

Co-authored-by: archer <545436317@qq.com>
parent 2de10d37
Showing with 1693 additions and 153 deletions
document/content/self-host/upgrading/4-15/41502.mdx
......@@ -21,6 +21,7 @@ description: 'FastGPT V4.15.0-beta2 更新说明'
1. 工作流，单节点调试，存在异常默认值。
2. 模型配置，defaultConfig 覆盖异常。
3. 切换团队时，清除本地 chat 缓存。
## 代码优化
......
document/data/doc-last-modified.json
......@@ -253,8 +253,10 @@
"content/self-host/upgrading/4-14/41419.mdx": "2026-05-09T15:25:23+08:00",
"content/self-host/upgrading/4-14/4142.en.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/4-14/4142.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/4-14/41420.en.mdx": "2026-05-14T17:58:58+08:00",
"content/self-host/upgrading/4-14/41420.en.mdx": "2026-05-21T15:39:08+08:00",
"content/self-host/upgrading/4-14/41420.mdx": "2026-05-19T14:44:21+08:00",
"content/self-host/upgrading/4-14/41421.en.mdx": "2026-05-21T15:39:08+08:00",
"content/self-host/upgrading/4-14/41421.mdx": "2026-05-21T15:39:08+08:00",
"content/self-host/upgrading/4-14/4143.en.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/4-14/4143.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/4-14/4144.en.mdx": "2026-04-26T21:08:47+08:00",
......@@ -415,6 +417,6 @@
"content/self-host/upgrading/outdated/499.mdx": "2026-05-07T15:06:40+08:00",
"content/self-host/upgrading/upgrade-intruction.en.mdx": "2026-04-26T21:08:47+08:00",
"content/self-host/upgrading/upgrade-intruction.mdx": "2026-04-26T21:08:47+08:00",
"content/toc.en.mdx": "2026-05-18T17:26:16+08:00",
"content/toc.mdx": "2026-05-18T17:26:16+08:00"
"content/toc.en.mdx": "2026-05-21T15:39:08+08:00",
"content/toc.mdx": "2026-05-21T15:39:08+08:00"
}
\ No newline at end of file
packages/global/core/dataset/data/constants.ts
import { i18nT } from '../../../common/i18n/utils';
export enum DatasetDataIndexTypeEnum {
default = 'default', // 默认的
default = 'default', // 默认文本索引
imageEmbedding = 'imageEmbedding', // 默认图片向量
summary = 'summary', // 摘要，系统生成
question = 'question', // 补全问题，系统生成
image = 'image', // 图片描述，系统生成
......@@ -34,6 +36,10 @@ export const DatasetDataIndexMap: Record<
[DatasetDataIndexTypeEnum.image]: {
label: i18nT('dataset:data_index_image'),
color: 'purple'
},
[DatasetDataIndexTypeEnum.imageEmbedding]: {
label: i18nT('dataset:data_index_image_embedding'),
color: 'purple'
}
};
export const defaultDatasetIndexData = DatasetDataIndexMap[DatasetDataIndexTypeEnum.custom];
......
packages/global/core/dataset/data/utils.ts 0 → 100644
import { DatasetDataIndexTypeEnum } from './constants';
export const datasetDataSystemIndexTypes = [
DatasetDataIndexTypeEnum.default,
DatasetDataIndexTypeEnum.imageEmbedding
] as const;
const datasetDataSystemIndexTypeSet = new Set<DatasetDataIndexTypeEnum>(
datasetDataSystemIndexTypes
);
/**
* 判断索引类型是否由数据内容自动生成和维护。
*
* 系统索引会随 data 的 q/a/imageId/markdown 图片重新生成，前端和后端都不应把它当作
* 用户可手动编辑的外部索引处理。
*/
export const isDatasetDataSystemIndexType = (type?: DatasetDataIndexTypeEnum) =>
datasetDataSystemIndexTypeSet.has(type || DatasetDataIndexTypeEnum.custom);
packages/global/core/dataset/type.ts
......@@ -94,7 +94,7 @@ export const DatasetSchema = z
deleteTime: z.coerce.date().nullish().meta({ description: '删除时间' }),
autoSync: z.boolean().optional().meta({ description: '自动同步', deprecated: true }),
autoSync: z.boolean().optional().meta({ description: '自动同步' }),
externalReadUrl: z.string().optional().meta({ description: '外部读取 URL', deprecated: true }),
defaultPermission: z.number().optional().meta({ description: '默认权限', deprecated: true }),
apiServer: APIFileServerSchema.optional().meta({
......@@ -171,7 +171,10 @@ export const DatasetDataIndexItemSchema = z.object({
.default(DatasetDataIndexTypeEnum.custom)
.meta({ description: '索引类型' }),
dataId: z.string().meta({ description: 'vectorDB ID' }),
text: z.string().meta({ description: '索引文本' })
text: z.string().meta({
description: `默认就是索引的文本内容，特殊的：
imageEmbedding - 图片的 objectKey/url`
})
});
const DatasetDataIndexOptionalSchema = DatasetDataIndexItemSchema.omit({ dataId: true }).extend({
dataId: z.string().optional().meta({
......
packages/global/core/workflow/template/system/datasetSearch.ts
......@@ -126,11 +126,6 @@ export const DatasetSearchModule: FlowNodeTemplateType = {
},
{
...Input_Template_UserChatInput,
toolDescription: i18nT('workflow:content_to_search'),
deprecated: true
},
{
...Input_Template_UserChatInput,
label: i18nT('workflow:search_query'),
key: NodeInputKeyEnum.datasetSearchInput,
valueType: WorkflowIOValueTypeEnum.arrayString,
......
packages/global/openapi/core/dataset/data/api.ts
......@@ -53,7 +53,16 @@ export type GetDatasetDataDetailResponse = z.infer<typeof GetDatasetDataDetailRe
* ============================================================================ */
export const UpdateDatasetDataBodySchema = UpdateDatasetDataPropsSchema;
export type UpdateDatasetDataBody = z.infer<typeof UpdateDatasetDataBodySchema>;
export const UpdateDatasetDataResponseSchema = z.object({});
export const UpdateDatasetDataResponseSchema = z.object({
q: z.string().optional().meta({
example: '![image.png](/api/system/file/download/xxx?filename=image.png)',
description: '展示态问题/主文本，内部 S3 图片会替换为签名访问地址'
}),
a: z.string().optional().meta({
example: 'FastGPT 是一个 AI Agent 构建平台',
description: '展示态回答/补充文本，内部 S3 图片会替换为签名访问地址'
})
});
export type UpdateDatasetDataResponse = z.infer<typeof UpdateDatasetDataResponseSchema>;
const DatasetDataIndexContentSchema = z.object({
......
packages/global/test/core/workflow/template/datasetSearch.test.ts
......@@ -3,7 +3,7 @@ import { NodeInputKeyEnum, WorkflowIOValueTypeEnum } from '@fastgpt/global/core/
import { DatasetSearchModule } from '@fastgpt/global/core/workflow/template/system/datasetSearch';
describe('DatasetSearchModule', () => {
it('should deprecate legacy user question input and add array search input', () => {
it('should use array search input without legacy user question input', () => {
const legacyUserQuestionInput = DatasetSearchModule.inputs.find(
(input) => input.key === NodeInputKeyEnum.userChatInput
);
......@@ -11,8 +11,7 @@ describe('DatasetSearchModule', () => {
(input) => input.key === NodeInputKeyEnum.datasetSearchInput
);
expect(legacyUserQuestionInput?.valueType).toBe(WorkflowIOValueTypeEnum.string);
expect(legacyUserQuestionInput?.deprecated).toBe(true);
expect(legacyUserQuestionInput).toBeUndefined();
expect(searchInput?.valueType).toBe(WorkflowIOValueTypeEnum.arrayString);
});
});
packages/service/common/logger/categories.ts
......@@ -72,7 +72,8 @@ export const LogCategories = {
IMAGE_PARSE: ['dataset', 'training', 'image-parse'],
IMAGE_INDEX: ['dataset', 'training', 'image-index'],
INDEX_EXTEND: ['dataset', 'training', 'index-extend'],
LLM_PARGRAPH: ['dataset', 'training', 'llm-pargraph']
LLM_PARGRAPH: ['dataset', 'training', 'llm-pargraph'],
WEB_SYNC: ['dataset', 'webSync']
}),
AI: Object.assign(['ai'], {
AGENT: ['ai', 'agent'],
......
packages/service/common/string/utils.ts
import { serviceEnv } from '../../env';
import { WorkerNameEnum, getWorkerController } from '../../worker/utils';
import { type ImageType } from '../../worker/readFile/type';
import { getLogger, LogCategories } from '../logger';
const logger = getLogger(LogCategories.INFRA.WORKER);
const HTML_TO_MARKDOWN_TIMEOUT_MS = 300000;
/**
* 将 HTML 转为 Markdown。实际转换在 worker 中执行，并在派发前记录任务大小，
* 便于定位站点同步中是否卡在 HTML 转换阶段。
*/
export const htmlToMarkdown = async (html?: string | null) => {
const htmlContent = html || '';
const workerController = getWorkerController<
{ html: string },
{
......@@ -12,11 +21,18 @@ export const htmlToMarkdown = async (html?: string | null) => {
>({
name: WorkerNameEnum.htmlStr2Md,
maxReservedThreads: serviceEnv.HTML_TO_MARKDOWN_WORKERS,
taskTimeoutMs: 300000,
taskTimeoutMs: HTML_TO_MARKDOWN_TIMEOUT_MS,
maxTasksPerWorker: 100
});
const md = await workerController.run({ html: html || '' });
logger.info('HTML to markdown worker task started', {
htmlLength: htmlContent.length,
workerName: WorkerNameEnum.htmlStr2Md,
maxReservedThreads: serviceEnv.HTML_TO_MARKDOWN_WORKERS,
timeoutMs: HTML_TO_MARKDOWN_TIMEOUT_MS
});
const md = await workerController.run({ html: htmlContent });
return md.rawText;
};
packages/service/common/system/timerLock/constants.ts
......@@ -13,6 +13,7 @@ export enum TimerIdEnum {
recordTeamQPM = 'recordTeamQPM',
auditLogCleanup = 'auditLogCleanup',
chatHistoryCleanup = 'chatHistoryCleanup',
datasetSyncSchedulerReconcile = 'datasetSyncSchedulerReconcile',
/** 纠正长时间卡在 generating 的会话状态 */
cleanStaleGeneratingChat = 'cleanStaleGeneratingChat'
}
......
packages/service/common/vectorDB/controller.ts
......@@ -4,6 +4,7 @@ import { ObVectorCtrl } from './oceanbase';
import { SeekVectorCtrl } from './seekdb';
import { OpenGaussVectorCtrl } from './opengauss';
import { getVectors } from '../../core/ai/embedding';
import type { GetVectorsProps } from '../../core/ai/embedding';
import type { VectorControllerType, InsertVectorControllerPropsType } from './type';
import { type EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.schema';
import {
......@@ -103,12 +104,21 @@ export const initVectorStore = Vector.init;
export const recallFromVectorStore: VectorControllerType['embRecall'] = (props) =>
retryFn(() => Vector.embRecall(props));
type DatasetVectorInput = string | GetVectorsProps['inputs'][number];
/**
* 统一写入知识库索引向量。
*
* `inputs` 的 text/image 类型只用于告诉 embedding 模型如何生成向量；
* 进入向量库时已经统一成 number[][]，向量库本身不区分文本向量或图片向量。
* 传入 string 时保持旧行为，默认按文本生成 embedding。
*/
export const insertDatasetDataVector = async ({
model,
inputs,
...props
}: Omit<InsertVectorControllerPropsType, 'vectors'> & {
inputs: string[];
inputs: DatasetVectorInput[];
model: EmbeddingModelItemType;
}) => {
if (inputs.length === 0) {
......@@ -118,12 +128,17 @@ export const insertDatasetDataVector = async ({
};
}
const embeddingInputs = inputs.map((input) =>
typeof input === 'string'
? {
type: 'text' as const,
input
}
: input
);
const { vectors, tokens } = await getVectors({
model,
inputs: inputs.map((text) => ({
type: 'text',
input: text
})),
inputs: embeddingInputs,
type: 'db'
});
const { insertIds } = await retryFn(() =>
......
packages/service/core/dataset/collection/controller.ts
......@@ -25,6 +25,7 @@ import { hashStr } from '@fastgpt/global/common/string/tools';
import { MongoDatasetDataText } from '../data/dataTextSchema';
import { retryFn } from '@fastgpt/global/common/system/utils';
import { getTrainingModeByCollection } from './utils';
import { getDatasetImageIndexCapability } from '../utils';
import {
computedCollectionChunkSettings,
getLLMMaxChunkSize
......@@ -77,7 +78,11 @@ export const createCollectionAndInsertData = async ({
const trainingMode = getTrainingModeByCollection({
trainingType: trainingType,
autoIndexes: formatCreateCollectionParams.autoIndexes,
imageIndex: formatCreateCollectionParams.imageIndex
imageIndex: formatCreateCollectionParams.imageIndex,
supportImageIndex: getDatasetImageIndexCapability({
vectorModel: dataset.vectorModel,
vlmModel: dataset.vlmModel
}).supportImageIndex
});
if (
......
packages/service/core/dataset/collection/utils.ts
......@@ -3,7 +3,6 @@ import type { ClientSession } from '../../../common/mongo';
import { MongoDatasetCollectionTags } from '../tag/schema';
import { readFromSecondary } from '../../../common/mongo/utils';
import type { CollectionWithDatasetType } from '@fastgpt/global/core/dataset/type';
import { DatasetCollectionSchemaType } from '@fastgpt/global/core/dataset/type';
import {
DatasetCollectionDataProcessModeEnum,
DatasetCollectionSyncResultEnum,
......@@ -216,11 +215,13 @@ export const syncCollection = async (collection: CollectionWithDatasetType) => {
export const getTrainingModeByCollection = ({
trainingType,
autoIndexes,
imageIndex
imageIndex,
supportImageIndex = false
}: {
trainingType?: DatasetCollectionDataProcessModeEnum;
autoIndexes?: boolean;
imageIndex?: boolean;
supportImageIndex?: boolean;
}) => {
if (
trainingType === DatasetCollectionDataProcessModeEnum.imageParse &&
......@@ -235,6 +236,7 @@ export const getTrainingModeByCollection = ({
if (
trainingType === DatasetCollectionDataProcessModeEnum.chunk &&
imageIndex &&
supportImageIndex &&
global.feConfigs?.isPlus
) {
return TrainingModeEnum.image;
......
packages/service/core/dataset/data/controller.ts
......@@ -4,6 +4,7 @@ import type { DatasetDataSchemaType } from '@fastgpt/global/core/dataset/type';
import { addDays } from 'date-fns';
import { isS3ObjectKey, jwtSignS3DownloadToken } from '../../../common/s3/utils';
import { S3Buckets } from '../../../common/s3/config/constants';
import { matchDatasetDataMarkdownImages } from './utils';
export const formatDatasetDataValue = ({
q,
......@@ -24,15 +25,23 @@ export const formatDatasetDataValue = ({
if (imageDescMap) {
// Helper function to replace image markdown with description
const replaceImageMarkdown = (text: string): string => {
return text.replace(/!\[([^\]]*)\]\(([^)]+)\)/g, (match, altText, url) => {
const description = imageDescMap[url];
const matches = matchDatasetDataMarkdownImages(text);
let content = text;
for (const item of matches.slice().reverse()) {
const description = imageDescMap[item.url];
if (description) {
// Add description to alt text, keeping original if exists
const newAltText = altText ? `${altText} - ${description}` : description;
return `![${newAltText.replace(/\n/g, '')}](${url})`;
const newAltText = item.alt ? `${item.alt} - ${description}` : description;
const replacement = `![${newAltText.replace(/\n/g, '')}](${item.url})`;
content =
content.slice(0, item.index) +
replacement +
content.slice(item.index + item.raw.length);
}
return match; // Return original if no description found
});
}
return content;
};
// Apply replacement to both q and a
......
packages/service/core/dataset/data/schema.ts
......@@ -33,10 +33,7 @@ const DatasetDataSchema = new Schema({
ref: DatasetColCollectionName,
required: true
},
q: {
type: String,
required: true
},
q: String,
a: {
type: String
},
......
packages/service/core/dataset/data/utils.ts 0 → 100644
export type DatasetDataMarkdownImageItem = {
raw: string;
alt: string;
url: string;
index: number;
};
/**
* 从 dataset data 的 markdown 内容中提取图片节点。
*
* 这里只负责识别 `![alt](url)`，用于 VLM 图片描述索引、imageEmbedding 图片向量索引、
* 展示态描述回填等链路共用同一套图片提取语义。图片来源合法性校验、S3/base64 转换、
* 向量生成都在后续链路处理。
*/
export const matchDatasetDataMarkdownImages = (text = ''): DatasetDataMarkdownImageItem[] => {
if (typeof text !== 'string' || !text) return [];
const regex = /!\[([\s\S]*?)\]\((.*?)\)/g;
return Array.from(text.matchAll(regex))
.map((match) => ({
raw: match[0],
alt: match[1] || '',
url: match[2]?.trim() || '',
index: match.index ?? 0
}))
.filter((item) => !!item.url);
};
/**
* 提取 dataset data markdown 图片 URL。
*
* 这是图片描述索引和图片向量索引共同使用的 URL 入口，避免不同训练/重建链路
* 分别维护 markdown 图片提取规则。
*/
export const matchDatasetDataMarkdownImageUrls = (text = '') =>
matchDatasetDataMarkdownImages(text).map((item) => item.url);
/**
* 从多个文本字段中提取并按首次出现顺序去重图片 URL。
*/
export const uniqueDatasetDataMarkdownImageUrls = (texts: Array<string | null | undefined>) =>
Array.from(
new Set(
texts.filter((text): text is string => !!text).flatMap(matchDatasetDataMarkdownImageUrls)
)
);
packages/service/core/dataset/datasetSync/index.ts
import { type Processor } from 'bullmq';
import { type JobSchedulerJson, type Processor } from 'bullmq';
import { getQueue, getWorker, QueueNames } from '../../../common/bullmq';
import { DatasetStatusEnum } from '@fastgpt/global/core/dataset/constants';
import { MongoDataset } from '../schema';
import { getLogger, LogCategories } from '../../../common/logger';
export type DatasetSyncJobData = {
datasetId: string;
};
const logger = getLogger(LogCategories.MODULE.DATASET);
export const datasetSyncQueue = getQueue<DatasetSyncJobData>(QueueNames.datasetSync, {
defaultJobOptions: {
attempts: 3, // retry 3 times
......@@ -99,3 +103,55 @@ export const getDatasetSyncJobScheduler = (datasetId: string) => {
export const removeDatasetSyncJobScheduler = (datasetId: string) => {
return datasetSyncQueue.removeJobScheduler(String(datasetId));
};
export type DatasetSyncSchedulerReconcileResult = {
autoSyncDatasetCount: number;
schedulerCount: number;
createdSchedulerCount: number;
createdDatasetIds: string[];
};
/**
* 以 Mongo `autoSync=true` 作为期望态，补齐缺失的 BullMQ datasetSync scheduler。
*
* 该函数只增加缺失 scheduler，不修改 Mongo `autoSync`，也不移除 Redis 中已有 scheduler/job。
*/
export const reconcileDatasetSyncSchedulers =
async (): Promise<DatasetSyncSchedulerReconcileResult> => {
const autoSyncDatasets = await MongoDataset.find(
{
autoSync: true,
$or: [{ deleteTime: null }, { deleteTime: { $exists: false } }]
},
'_id'
).lean();
const autoSyncDatasetIds = new Set(autoSyncDatasets.map((dataset) => String(dataset._id)));
const schedulers = (await datasetSyncQueue.getJobSchedulers(
0,
-1,
true
)) as JobSchedulerJson<DatasetSyncJobData>[];
const schedulerIds = new Set(
schedulers.map((scheduler) => String(scheduler.key)).filter(Boolean)
);
const createdDatasetIds: string[] = [];
for (const datasetId of autoSyncDatasetIds) {
if (schedulerIds.has(datasetId)) continue;
await upsertDatasetSyncJobScheduler({ datasetId });
createdDatasetIds.push(datasetId);
}
const result = {
autoSyncDatasetCount: autoSyncDatasetIds.size,
schedulerCount: schedulers.length,
createdSchedulerCount: createdDatasetIds.length,
createdDatasetIds
};
logger.info('Dataset sync scheduler reconcile finished', result);
return result;
};
packages/service/core/dataset/schema.ts
......@@ -138,12 +138,16 @@ const DatasetSchema = new Schema({
default: null // null表示未删除，有值表示删除时间
},
// abandoned
autoSync: Boolean,
/** @deprecated */
externalReadUrl: String,
/** @deprecated */
defaultPermission: Number,
/** @deprecated */
apiServer: Object,
/** @deprecated */
feishuServer: Object,
/** @deprecated */
yuqueServer: Object
});
......
packages/service/core/dataset/search/utils.ts
......@@ -5,6 +5,7 @@ import { getLogger, LogCategories } from '../../../common/logger';
import type { OpenaiAccountType } from '@fastgpt/global/support/user/team/type';
import { getImageBase64 } from '../../../common/file/image/utils';
import { serviceEnv } from '../../../env';
import { isS3ObjectKey } from '../../../common/s3/utils';
const logger = getLogger(LogCategories.MODULE.DATASET.DATA);
......@@ -25,6 +26,19 @@ export const computeFilterIntersection = (lists: (string[] | undefined)[]) => {
});
};
export const isValidImageEmbeddingSource = (imageUrl?: string) => {
const url = imageUrl?.trim();
if (!url) return false;
if (url.startsWith('data:image/')) return true;
if (isS3ObjectKey(url, 'dataset')) return true;
if (isS3ObjectKey(url, 'temp')) return true;
if (isS3ObjectKey(url, 'chat')) return true;
if (/^https?:\/\//i.test(url)) return true;
return false;
};
/**
* 按环境开关规范化图片输入。
* data URL 已经是模型可读内容，始终原样返回；普通图片 URL 只有
......
packages/service/core/dataset/training/controller.ts
......@@ -5,7 +5,7 @@ import type {
} from '@fastgpt/global/openapi/core/dataset/data/api';
import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
import { type ClientSession } from '../../../common/mongo';
import { getLLMModel, getEmbeddingModel, getVlmModel } from '../../ai/model';
import { getLLMModel, getEmbeddingModel, getVlmModel, isImageEmbeddingModel } from '../../ai/model';
import { mongoSessionRun } from '../../../common/mongo/sessionRun';
import { i18nT } from '@fastgpt/global/common/i18n/utils';
import { getLLMMaxChunkSize } from '../../../../global/core/dataset/training/utils';
......@@ -79,7 +79,7 @@ export const pushDataListToTrainingQueue = async ({
return Promise.reject(i18nT('common:error_llm_not_config'));
}
const { model, maxToken, weight } = await (async () => {
const { maxToken, weight } = await (async () => {
if (mode === TrainingModeEnum.chunk) {
return {
maxToken: Infinity,
......@@ -97,6 +97,13 @@ export const pushDataListToTrainingQueue = async ({
if (mode === TrainingModeEnum.image || mode === TrainingModeEnum.imageParse) {
const vllmModelData = getVlmModel(vlmModel);
if (!vllmModelData) {
if (mode === TrainingModeEnum.image && isImageEmbeddingModel(vectorModelData)) {
return {
maxToken: Infinity,
model: vectorModelData.model,
weight: vectorModelData.weight
};
}
return Promise.reject(i18nT('common:error_vlm_not_config'));
}
return {
......
packages/service/core/dataset/utils.ts
......@@ -4,6 +4,8 @@ import { S3Sources } from '../../common/s3/contracts/type';
import { jwtSignS3DownloadToken, isS3ObjectKey } from '../../common/s3/utils';
import { getLogger, LogCategories } from '../../common/logger';
import { S3Buckets } from '../../common/s3/config/constants';
import { getVlmModelList, isImageEmbeddingModel } from '../ai/model';
import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
const logger = getLogger(LogCategories.MODULE.DATASET.FILE);
......@@ -82,3 +84,46 @@ export function replaceS3KeyToPreviewUrl(documentQuoteText: string, expiredTime:
return content;
}
const getAvailableDatasetVlmModel = (vlmModel?: string) => {
if (!vlmModel) return;
const vlmModelList = getVlmModelList();
return vlmModelList.find((item) => item.model === vlmModel || item.name === vlmModel);
};
export const getDatasetImageIndexCapability = ({
vectorModel,
vlmModel
}: {
vectorModel?: string;
vlmModel?: string;
}) => {
const availableVlmModel = getAvailableDatasetVlmModel(vlmModel);
const supportVlm = !!availableVlmModel;
const supportImageEmbedding = isImageEmbeddingModel(vectorModel);
return {
availableVlmModel,
supportVlm,
supportImageEmbedding,
supportImageIndex: supportVlm || supportImageEmbedding
};
};
export const getDatasetImageTrainingMode = ({
supportVlm,
supportImageIndex,
imageId,
hasMarkdownImages
}: {
supportVlm: boolean;
supportImageIndex: boolean;
imageId?: string;
hasMarkdownImages: boolean;
}) => {
if (supportVlm && imageId) return TrainingModeEnum.imageParse;
if (supportImageIndex && hasMarkdownImages) return TrainingModeEnum.image;
return TrainingModeEnum.chunk;
};
packages/service/test/common/vectorDB/controller.test.ts
......@@ -289,6 +289,51 @@ describe('VectorDB Controller', () => {
});
});
it('should pass explicit image inputs to embedding generation', async () => {
const mockVectors = [[0.1, 0.2]];
mockGetVectors.mockResolvedValue({
tokens: 1,
vectors: mockVectors
});
mockVectorInsert.mockResolvedValue({
insertIds: ['image_id']
});
const result = await insertDatasetDataVector({
teamId: 'team_123',
datasetId: 'dataset_456',
collectionId: 'col_789',
inputs: [
{
type: 'image',
input: 'data:image/png;base64,image'
}
],
model: mockModel as any
});
expect(mockGetVectors).toHaveBeenCalledWith({
model: mockModel,
inputs: [
{
type: 'image',
input: 'data:image/png;base64,image'
}
],
type: 'db'
});
expect(mockVectorInsert).toHaveBeenCalledWith({
teamId: 'team_123',
datasetId: 'dataset_456',
collectionId: 'col_789',
vectors: mockVectors
});
expect(result).toEqual({
tokens: 1,
insertIds: ['image_id']
});
});
it('should invalidate team vector cache after insert', async () => {
mockGetVectors.mockResolvedValue({
tokens: 50,
......
packages/service/test/core/dataset/data/controller.test.ts 0 → 100644
import { describe, expect, it } from 'vitest';
import { formatDatasetDataValue } from '@fastgpt/service/core/dataset/data/controller';
describe('formatDatasetDataValue', () => {
it('should append image descriptions to markdown image alt text in question and answer', () => {
const result = formatDatasetDataValue({
q: 'Question ![cat]( https://example.com/cat.png ) and ![bird](https://example.com/bird.png)',
a: 'Answer ![](https://example.com/dog.png)',
imageDescMap: {
'https://example.com/cat.png': 'cat desc\nline',
'https://example.com/dog.png': 'dog desc'
}
});
expect(result).toEqual({
q: 'Question ![cat - cat descline](https://example.com/cat.png) and ![bird](https://example.com/bird.png)',
a: 'Answer ![dog desc](https://example.com/dog.png)'
});
});
});
packages/service/test/core/dataset/search/utils.test.ts
......@@ -15,6 +15,7 @@ vi.mock('@fastgpt/service/common/file/image/utils', () => ({
import {
computeFilterIntersection,
datasetSearchQueryExtension,
isValidImageEmbeddingSource,
normalizeImageToBase64
} from '../../../../core/dataset/search/utils';
......@@ -86,6 +87,21 @@ describe('normalizeImageToBase64', () => {
});
});
describe('isValidImageEmbeddingSource', () => {
it('should accept model-readable image sources', () => {
expect(isValidImageEmbeddingSource('data:image/png;base64,input')).toBe(true);
expect(isValidImageEmbeddingSource('dataset/team/file.png')).toBe(true);
expect(isValidImageEmbeddingSource('temp/team/file.png')).toBe(true);
expect(isValidImageEmbeddingSource('chat/app/user/file.png')).toBe(true);
expect(isValidImageEmbeddingSource('https://example.com/file.png')).toBe(true);
});
it('should reject empty or local non-url image sources', () => {
expect(isValidImageEmbeddingSource('')).toBe(false);
expect(isValidImageEmbeddingSource('/local/file.png')).toBe(false);
});
});
describe('computeFilterIntersection', () => {
it('should return undefined when every filter is absent', () => {
expect(computeFilterIntersection([])).toBeUndefined();
......
packages/service/test/core/dataset/utils.test.ts
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { replaceS3KeyToPreviewUrl } from '@fastgpt/service/core/dataset/utils';
import {
getDatasetImageIndexCapability,
getDatasetImageTrainingMode,
replaceS3KeyToPreviewUrl
} from '@fastgpt/service/core/dataset/utils';
import {
matchDatasetDataMarkdownImages,
matchDatasetDataMarkdownImageUrls,
uniqueDatasetDataMarkdownImageUrls
} from '@fastgpt/service/core/dataset/data/utils';
import { getTrainingModeByCollection } from '@fastgpt/service/core/dataset/collection/utils';
import {
DatasetCollectionDataProcessModeEnum,
TrainingModeEnum
} from '@fastgpt/global/core/dataset/constants';
vi.mock('@fastgpt/service/common/s3/utils', () => ({
jwtSignS3DownloadToken: vi.fn(
......@@ -439,3 +453,144 @@ describe('replaceS3KeyToPreviewUrl', () => {
});
});
});
describe('matchDatasetDataMarkdownImageUrls', () => {
it('应提取统一的 markdown 图片节点结构', () => {
const result = matchDatasetDataMarkdownImages(
'文字 ![猫]( dataset/team/cat.png ) 和 ![dog](https://example.com/dog.png)'
);
expect(result).toEqual([
{
raw: '![猫]( dataset/team/cat.png )',
alt: '猫',
url: 'dataset/team/cat.png',
index: expect.any(Number)
},
{
raw: '![dog](https://example.com/dog.png)',
alt: 'dog',
url: 'https://example.com/dog.png',
index: expect.any(Number)
}
]);
});
it('应提取 markdown 图片 URL 并忽略普通链接', () => {
const result = matchDatasetDataMarkdownImageUrls(
'![a](dataset/team/a.png) [普通链接](https://example.com) ![b](https://img.test/b.jpg)'
);
expect(result).toEqual(['dataset/team/a.png', 'https://img.test/b.jpg']);
});
it('应从多个文本字段按首次出现顺序去重图片 URL', () => {
const result = uniqueDatasetDataMarkdownImageUrls([
'new ![a](dataset/team/a.png) ![a again](dataset/team/a.png)',
undefined,
'old ![b](https://example.com/b.jpg)'
]);
expect(result).toEqual(['dataset/team/a.png', 'https://example.com/b.jpg']);
});
});
describe('getDatasetImageTrainingMode', () => {
it('有 VLM 且是图片数据时应走 imageParse', () => {
expect(
getDatasetImageTrainingMode({
supportVlm: true,
supportImageIndex: true,
imageId: 'dataset/team/image.png',
hasMarkdownImages: false
})
).toBe(TrainingModeEnum.imageParse);
});
it('有图片索引能力且正文有 markdown 图片时应走 image', () => {
expect(
getDatasetImageTrainingMode({
supportVlm: false,
supportImageIndex: true,
hasMarkdownImages: true
})
).toBe(TrainingModeEnum.image);
});
it('没有图片索引能力时应回退 chunk', () => {
expect(
getDatasetImageTrainingMode({
supportVlm: false,
supportImageIndex: false,
hasMarkdownImages: true
})
).toBe(TrainingModeEnum.chunk);
});
});
describe('getTrainingModeByCollection', () => {
beforeEach(() => {
global.feConfigs = {
...global.feConfigs,
isPlus: true
};
});
it('图片自动索引有 VLM 或原生 embedding 图片索引能力时进入 image 队列', () => {
expect(
getTrainingModeByCollection({
trainingType: DatasetCollectionDataProcessModeEnum.chunk,
imageIndex: true,
supportImageIndex: true
})
).toBe(TrainingModeEnum.image);
expect(
getTrainingModeByCollection({
trainingType: DatasetCollectionDataProcessModeEnum.chunk,
imageIndex: true,
supportImageIndex: false
})
).toBe(TrainingModeEnum.chunk);
});
});
describe('getDatasetImageIndexCapability', () => {
beforeEach(() => {
global.embeddingModelMap.set('vision-embedding-model', {
...global.systemDefaultModel.embedding,
model: 'vision-embedding-model',
name: 'vision-embedding-model',
vision: true
});
global.llmModelMap.set('dataset-vlm-model', {
...global.systemDefaultModel.llm,
model: 'dataset-vlm-model',
name: 'dataset-vlm-model',
vision: true
});
});
it('未配置 VLM 时不应自动回退到默认 VLM', () => {
const result = getDatasetImageIndexCapability({
vectorModel: 'vision-embedding-model'
});
expect(result.supportVlm).toBe(false);
expect(result.supportImageEmbedding).toBe(true);
expect(result.supportImageIndex).toBe(true);
expect(result.availableVlmModel).toBeUndefined();
});
it('配置 VLM 时应同时返回 VLM 和多模态索引能力', () => {
const result = getDatasetImageIndexCapability({
vectorModel: 'vision-embedding-model',
vlmModel: 'dataset-vlm-model'
});
expect(result.supportVlm).toBe(true);
expect(result.supportImageEmbedding).toBe(true);
expect(result.supportImageIndex).toBe(true);
expect(result.availableVlmModel?.model).toBe('dataset-vlm-model');
});
});
packages/web/i18n/en/dataset.json
......@@ -50,6 +50,8 @@
"data_amount": "{{dataAmount}} chunks, {{indexAmount}} indexes",
"data_error_amount": "{{errorAmount}} Group training exception",
"data_index_image": "Image index",
"data_index_image_embedding": "Multimodal image index",
"image_embedding_index_default_desc": "An image vector has been generated by the multimodal model and can be used for image search",
"data_parsing": "Data analysis",
"data_uploading": "Data is being uploaded: {{num}}%",
"dataset.Chunk_Number": "Block number",
......
packages/web/i18n/zh-CN/dataset.json
......@@ -50,6 +50,8 @@
"data_amount": "{{dataAmount}} 组数据, {{indexAmount}} 组索引",
"data_error_amount": "{{errorAmount}} 组训练异常",
"data_index_image": "图片索引",
"data_index_image_embedding": "多模态图片索引",
"image_embedding_index_default_desc": "已通过多模态模型生成图片向量，支持以图搜图",
"data_parsing": "数据解析中",
"data_uploading": "数据上传中: {{num}}%",
"dataset.Chunk_Number": "分块号",
......
packages/web/i18n/zh-Hant/dataset.json
......@@ -50,6 +50,8 @@
"data_amount": "{{dataAmount}} 組資料，{{indexAmount}} 組索引",
"data_error_amount": "{{errorAmount}} 組訓練異常",
"data_index_image": "圖片索引",
"data_index_image_embedding": "多模態圖片索引",
"image_embedding_index_default_desc": "已透過多模態模型生成圖片向量，支援以圖搜圖",
"data_parsing": "數據解析中",
"data_uploading": "數據上傳中: {{num}}%",
"dataset.Chunk_Number": "分塊號",
......
pro @ 3cd897c7
Subproject commit 821e0d5d30280d6ed4ea5c42f8682d07a6f87bf6
Subproject commit 3cd897c756ca2c0d7c7fee9f68c0d88ad89b7009
projects/app/src/pageComponents/dataset/detail/components/InputDataModal/IndexInputPanel.tsx
......@@ -11,6 +11,7 @@ import {
DatasetDataIndexTypeEnum,
getDatasetIndexMapData
} from '@fastgpt/global/core/dataset/data/constants';
import { isDatasetDataSystemIndexType } from '@fastgpt/global/core/dataset/data/utils';
import type { InputDataType } from './useInputDataModal';
type IndexField = UseFieldArrayReturn<InputDataType, 'indexes'>['fields'][number];
......@@ -91,10 +92,14 @@ const IndexInputPanel = ({
const canFoldIndex = indexes.length > 1;
const hasIndexDataId = !!index.dataId;
const isDeletingCurrentIndex = deletingIndexClientId === index.clientId;
const isDefaultIndex = index.type === DatasetDataIndexTypeEnum.default;
const isSystem = isDatasetDataSystemIndexType(index.type);
const canDeleteIndex =
canWrite && !isDefaultIndex && hasIndexDataId && !isDeletingCurrentIndex;
canWrite && !isSystem && hasIndexDataId && !isDeletingCurrentIndex;
const canToggleFold = canFoldIndex && !isDeletingCurrentIndex;
const isImageEmbeddingIndex = index.type === DatasetDataIndexTypeEnum.imageEmbedding;
const indexText = isImageEmbeddingIndex
? t('dataset:image_embedding_index_default_desc')
: index.text;
return (
<MyBox
......@@ -152,11 +157,11 @@ const IndexInputPanel = ({
)}
</Flex>
<DataIndexTextArea
disabled={!canWrite || isDefaultIndex}
disabled={!canWrite || isSystem}
canClickMark={hasIndexDataId}
autoFocus={focusIndexClientId === index.clientId}
index={i}
value={index.text}
value={indexText}
isFolder={index.fold && canFoldIndex}
maxToken={maxToken}
register={register}
......
projects/app/src/pageComponents/dataset/detail/components/InputDataModal/useInputDataModal.ts
......@@ -16,6 +16,7 @@ import { useRequest } from '@fastgpt/web/hooks/useRequest';
import { useToast } from '@fastgpt/web/hooks/useToast';
import type { DatasetDataIndexItemType } from '@fastgpt/global/core/dataset/type';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
import { isDatasetDataSystemIndexType } from '@fastgpt/global/core/dataset/data/utils';
import { DatasetCollectionTypeEnum } from '@fastgpt/global/core/dataset/constants';
export type InputDataType = {
......@@ -45,8 +46,8 @@ const clearEditingIndexDelay = 600;
const sortIndexesForDisplay = (indexes: InputDataIndexType[] = []) => {
const getOrder = (index: InputDataIndexType) => {
// Keep the editable custom indexes before the generated default index.
if (index.type === DatasetDataIndexTypeEnum.default) return 1;
// Keep editable indexes before the generated system indexes.
if (isDatasetDataSystemIndexType(index.type)) return 1;
return 0;
};
......@@ -82,7 +83,7 @@ const formatIndexesForForm = (
const formatIndexesForRequest = (indexes: InputDataType['indexes'] = []) =>
indexes
.filter((item) => !!item.text?.trim())
.filter((item) => !isDatasetDataSystemIndexType(item.type) && !!item.text?.trim())
.map((item) => ({
// Strip UI-only fields before submitting to the import API.
type: item.type,
......@@ -346,12 +347,17 @@ export const useInputDataModal = ({
async (e: InputDataType) => {
if (!dataId) return Promise.reject(t('common:error.unKnow'));
await putDatasetDataById({
const updateResult = await putDatasetDataById({
dataId,
q: e.q,
a: currentTab === TabEnum.qa ? e.a : ''
});
return refreshDataForm(dataId);
const refreshedData = await refreshDataForm(dataId);
return {
...refreshedData,
q: updateResult.q ?? refreshedData.q,
a: updateResult.a ?? refreshedData.a
};
},
{
refreshDeps: [currentTab, refreshDataForm],
......@@ -377,7 +383,7 @@ export const useInputDataModal = ({
const text = targetIndex?.text?.trim() || '';
const type = targetIndex?.type || DatasetDataIndexTypeEnum.custom;
if (type === DatasetDataIndexTypeEnum.default) {
if (isDatasetDataSystemIndexType(type)) {
return;
}
......@@ -456,6 +462,12 @@ export const useInputDataModal = ({
if (successData) {
onSuccess(successData);
}
if (!shouldSaveLatest) {
toast({
title: t('common:save_success'),
status: 'success'
});
}
} catch (error) {
saveError = error;
} finally {
......@@ -476,7 +488,16 @@ export const useInputDataModal = ({
}
},
{
refreshDeps: [dataId, findIndexByClientId, getSuccessData, getValues, removeIndexes, setValue]
refreshDeps: [
dataId,
findIndexByClientId,
getSuccessData,
getValues,
removeIndexes,
setValue,
t,
toast
]
}
);
......
projects/app/src/pages/api/core/dataset/collection/create/images.ts
......@@ -22,6 +22,7 @@ import { getTeamPlanStatus } from '@fastgpt/service/support/wallet/sub/utils';
import { datasetImageCollectionFileType } from '@fastgpt/global/common/file/constants';
import { parseAllowedExtensions } from '@fastgpt/service/common/s3/utils/uploadConstraints';
import { checkDatasetIndexLimit } from '@fastgpt/service/support/permission/teamLimit';
import { getDatasetImageIndexCapability } from '@fastgpt/service/core/dataset/utils';
async function handler(req: ApiRequestProps): Promise<CreateCollectionWithResultResponseType> {
const filepaths: string[] = [];
......@@ -59,7 +60,12 @@ async function handler(req: ApiRequestProps): Promise<CreateCollectionWithResult
num: result.fileMetadata.length
});
if (!dataset.vlmModel) {
const { supportVlm, supportImageEmbedding } = getDatasetImageIndexCapability({
vectorModel: dataset.vectorModel,
vlmModel: dataset.vlmModel
});
if (!supportVlm && !supportImageEmbedding) {
return Promise.reject(i18nT('file:Image_dataset_requires_VLM_model_to_be_configured'));
}
......@@ -87,7 +93,9 @@ async function handler(req: ApiRequestProps): Promise<CreateCollectionWithResult
datasetId,
type: DatasetCollectionTypeEnum.images,
name: collectionName,
trainingType: DatasetCollectionDataProcessModeEnum.imageParse
trainingType: supportVlm
? DatasetCollectionDataProcessModeEnum.imageParse
: DatasetCollectionDataProcessModeEnum.chunk
}
});
} catch (error) {
......
projects/app/src/pages/api/core/dataset/data/insertData.ts
......@@ -47,6 +47,8 @@ async function handler(req: ApiRequestProps): Promise<InsertDataResponse> {
{
dataset: { _id: datasetId, vectorModel },
indexPrefixTitle,
imageIndex,
indexSize,
name
}
] = await Promise.all([getCollectionWithDataset(collectionId)]);
......@@ -76,8 +78,10 @@ async function handler(req: ApiRequestProps): Promise<InsertDataResponse> {
q: formatQ,
a: formatA,
chunkIndex: 0,
indexSize,
indexPrefix: indexPrefixTitle ? `# ${name}` : undefined,
embeddingModel: vectorModelData.model,
imageIndex: !!imageIndex,
indexes: formatIndexes
});
......
projects/app/src/pages/api/core/dataset/data/insertImages.ts
......@@ -7,7 +7,7 @@ import { WritePermissionVal } from '@fastgpt/global/support/permission/constant'
import { mongoSessionRun } from '@fastgpt/service/common/mongo/sessionRun';
import { createTrainingUsage } from '@fastgpt/service/support/wallet/usage/controller';
import { UsageSourceEnum } from '@fastgpt/global/support/wallet/usage/constants';
import { getEmbeddingModel, getLLMModel, getVlmModel } from '@fastgpt/service/core/ai/model';
import { getEmbeddingModel, getLLMModel } from '@fastgpt/service/core/ai/model';
import { pushDataListToTrainingQueue } from '@fastgpt/service/core/dataset/training/controller';
import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
import path from 'node:path';
......@@ -21,6 +21,8 @@ import {
} from '@fastgpt/global/openapi/core/dataset/data/api';
import { datasetImageCollectionFileType } from '@fastgpt/global/common/file/constants';
import { parseAllowedExtensions } from '@fastgpt/service/common/s3/utils/uploadConstraints';
import { i18nT } from '@fastgpt/global/common/i18n/utils';
import { getDatasetImageIndexCapability } from '@fastgpt/service/core/dataset/utils';
async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> {
const filepaths: string[] = [];
......@@ -42,6 +44,16 @@ async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> {
authApiKey: true
});
const dataset = collection.dataset;
const { availableVlmModel, supportVlm, supportImageEmbedding } = getDatasetImageIndexCapability(
{
vectorModel: dataset.vectorModel,
vlmModel: dataset.vlmModel
}
);
if (!supportVlm && !supportImageEmbedding) {
return Promise.reject(i18nT('file:Image_dataset_requires_VLM_model_to_be_configured'));
}
const planStatus = await getTeamPlanStatus({ teamId });
await authFrequencyLimit({
......@@ -75,7 +87,7 @@ async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> {
billSource: UsageSourceEnum.training,
vectorModel: getEmbeddingModel(dataset.vectorModel)?.name,
agentModel: getLLMModel(dataset.agentModel)?.name,
vllmModel: getVlmModel(dataset.vlmModel)?.name,
vllmModel: availableVlmModel?.name,
session
});
return usageId;
......@@ -89,9 +101,9 @@ async function handler(req: ApiRequestProps): Promise<InsertImagesResponse> {
agentModel: dataset.agentModel,
vectorModel: dataset.vectorModel,
vlmModel: dataset.vlmModel,
mode: TrainingModeEnum.imageParse,
mode: supportVlm ? TrainingModeEnum.imageParse : TrainingModeEnum.chunk,
billId: traingBillId,
data: imageIds.map((item, index) => ({
data: imageIds.map((item) => ({
imageId: item
})),
session
......
projects/app/src/pages/api/core/dataset/data/pushData.ts
......@@ -6,6 +6,7 @@ import { pushDataListToTrainingQueue } from '@fastgpt/service/core/dataset/train
import { NextAPI } from '@/service/middleware/entry';
import { WritePermissionVal } from '@fastgpt/global/support/permission/constant';
import { getTrainingModeByCollection } from '@fastgpt/service/core/dataset/collection/utils';
import { getDatasetImageIndexCapability } from '@fastgpt/service/core/dataset/utils';
import type { ApiRequestProps } from '@fastgpt/service/type/next';
import {
PushDataBodySchema,
......@@ -17,7 +18,6 @@ import { getLLMModel } from '@fastgpt/service/core/ai/model';
import { getVlmModel } from '@fastgpt/service/core/ai/model';
import { createTrainingUsage } from '@fastgpt/service/support/wallet/usage/controller';
import { mongoSessionRun } from '@fastgpt/service/common/mongo/sessionRun';
import { DatasetCollectionDataProcessModeEnum } from '@fastgpt/global/core/dataset/constants';
import { parseApiInput } from '@fastgpt/service/common/zod/requestParseError';
async function handler(req: ApiRequestProps): Promise<PushDataResponseType> {
......@@ -36,7 +36,13 @@ async function handler(req: ApiRequestProps): Promise<PushDataResponseType> {
per: WritePermissionVal
});
const mode = getTrainingModeByCollection(collection);
const mode = getTrainingModeByCollection({
...collection,
supportImageIndex: getDatasetImageIndexCapability({
vectorModel: collection.dataset.vectorModel,
vlmModel: collection.dataset.vlmModel
}).supportImageIndex
});
// auth dataset limit
await checkDatasetIndexLimit({
......
projects/app/src/pages/api/core/dataset/data/update.ts
import {
updateDatasetDataByIndexes,
updateDatasetDataDefaultIndexes
updateDatasetDataSystemIndexes,
updateDatasetDataByIndexes
} from '@/service/core/dataset/data/data';
import { pushGenerateVectorUsage } from '@/service/support/wallet/usage/push';
import { NextAPI } from '@/service/middleware/entry';
......@@ -16,6 +16,9 @@ import {
UpdateDatasetDataResponseSchema,
type UpdateDatasetDataResponse
} from '@fastgpt/global/openapi/core/dataset/data/api';
import { replaceS3KeyToPreviewUrl } from '@fastgpt/service/core/dataset/utils';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
import { addHours } from 'date-fns';
async function handler(req: ApiRequestProps): Promise<UpdateDatasetDataResponse> {
const { dataId, q, a, indexes } = parseApiInput({
......@@ -26,11 +29,7 @@ async function handler(req: ApiRequestProps): Promise<UpdateDatasetDataResponse>
// auth data permission
const {
collection: {
dataset: { vectorModel },
name,
indexPrefixTitle
},
collection: { name, indexPrefixTitle },
teamId,
tmbId,
collection,
......@@ -43,13 +42,38 @@ async function handler(req: ApiRequestProps): Promise<UpdateDatasetDataResponse>
per: WritePermissionVal
});
const dataset = collection.dataset;
const vectorModel = dataset.vectorModel;
const nextQ = q ?? datasetData.q ?? '';
const nextA = a ?? datasetData.a ?? '';
const pushUpdateDataAuditLog = () => {
addAuditLog({
tmbId,
teamId,
event: AuditEventEnum.UPDATE_DATA,
params: {
collectionName: collection.name,
datasetName: collection.dataset?.name || '',
datasetType: getI18nDatasetType(collection.dataset?.type || '')
}
});
};
if (hasIndexes) {
// 兼容旧 API：调用方显式传 indexes 时仍按完整索引更新。
// imageEmbedding 是系统索引，不能被外部 indexes 覆盖，统一由 dataIndex 根据数据内容重建。
const manualIndexes = indexes?.filter(
(index) => index.type !== DatasetDataIndexTypeEnum.imageEmbedding
);
const { tokens } = await updateDatasetDataByIndexes({
dataId,
q,
a,
indexes: indexes ?? [],
q: nextQ,
a: nextA,
imageId: datasetData.imageId,
imageIndex: !!collection.imageIndex,
indexes: manualIndexes || [],
model: vectorModel,
indexSize: collection.indexSize,
indexPrefix: indexPrefixTitle ? `# ${name}` : undefined
});
......@@ -61,14 +85,13 @@ async function handler(req: ApiRequestProps): Promise<UpdateDatasetDataResponse>
model: vectorModel
});
}
} else {
const nextQ = q || datasetData.q || '';
const nextA = a ?? datasetData.a ?? '';
const { tokens } = await updateDatasetDataDefaultIndexes({
} else if (!!nextQ || !!datasetData.imageId) {
const { tokens } = await updateDatasetDataSystemIndexes({
dataId,
q: nextQ,
a: nextA,
imageId: datasetData.imageId,
imageIndex: !!collection.imageIndex,
model: vectorModel,
indexSize: collection.indexSize,
indexPrefix: indexPrefixTitle ? `# ${name}` : undefined
......@@ -84,20 +107,12 @@ async function handler(req: ApiRequestProps): Promise<UpdateDatasetDataResponse>
}
}
(() => {
addAuditLog({
tmbId,
teamId,
event: AuditEventEnum.UPDATE_DATA,
params: {
collectionName: collection.name,
datasetName: collection.dataset?.name || '',
datasetType: getI18nDatasetType(collection.dataset?.type || '')
}
});
})();
pushUpdateDataAuditLog();
return UpdateDatasetDataResponseSchema.parse({});
return UpdateDatasetDataResponseSchema.parse({
q: replaceS3KeyToPreviewUrl(nextQ, addHours(new Date(), 1)),
a: nextA ? replaceS3KeyToPreviewUrl(nextA, addHours(new Date(), 1)) : undefined
});
}
export default NextAPI(handler);
projects/app/src/pages/api/core/dataset/training/rebuildEmbedding.ts
......@@ -3,10 +3,16 @@ import { authDataset } from '@fastgpt/service/support/permission/dataset/auth';
import { mongoSessionRun } from '@fastgpt/service/common/mongo/sessionRun';
import { MongoDataset } from '@fastgpt/service/core/dataset/schema';
import { MongoDatasetData } from '@fastgpt/service/core/dataset/data/schema';
import { MongoDatasetCollection } from '@fastgpt/service/core/dataset/collection/schema';
import { MongoDatasetTraining } from '@fastgpt/service/core/dataset/training/schema';
import { createTrainingUsage } from '@fastgpt/service/support/wallet/usage/controller';
import { UsageSourceEnum } from '@fastgpt/global/support/wallet/usage/constants';
import { getLLMModel, getEmbeddingModel, getVlmModel } from '@fastgpt/service/core/ai/model';
import { getLLMModel, getEmbeddingModel } from '@fastgpt/service/core/ai/model';
import {
getDatasetImageIndexCapability,
getDatasetImageTrainingMode
} from '@fastgpt/service/core/dataset/utils';
import { uniqueDatasetDataMarkdownImageUrls } from '@fastgpt/service/core/dataset/data/utils';
import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
import { type ApiRequestProps } from '@fastgpt/service/type/next';
import { OwnerPermissionVal } from '@fastgpt/global/support/permission/constant';
......@@ -46,14 +52,19 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
return Promise.reject('数据集正在训练或者重建中，请稍后再试');
}
const { availableVlmModel, supportVlm, supportImageIndex } = getDatasetImageIndexCapability({
vectorModel,
vlmModel: dataset.vlmModel
});
const { usageId } = await createTrainingUsage({
teamId,
tmbId,
appName: '切换索引模型',
billSource: UsageSourceEnum.training,
vectorModel: getEmbeddingModel(dataset.vectorModel)?.name,
vectorModel: getEmbeddingModel(vectorModel)?.name || vectorModel,
agentModel: getLLMModel(dataset.agentModel)?.name,
vllmModel: getVlmModel(dataset.vlmModel)?.name
vllmModel: availableVlmModel?.name
});
// update vector model and dataset.data rebuild field
......@@ -61,10 +72,27 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
await MongoDataset.findByIdAndUpdate(
datasetId,
{
vectorModel
$set: {
vectorModel,
...(!supportImageIndex && { 'chunkSettings.imageIndex': false })
}
},
{ session }
);
if (!supportImageIndex) {
await MongoDatasetCollection.updateMany(
{
teamId,
datasetId
},
{
$set: {
imageIndex: false
}
},
{ session }
);
}
await MongoDatasetData.updateMany(
{
teamId,
......@@ -85,7 +113,7 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
const max = global.systemEnv?.vectorMaxProcess || 10;
const arr = new Array(max * 2).fill(0);
for await (const _ of arr) {
for (let i = 0; i < arr.length; i++) {
try {
const hasNext = await mongoSessionRun(async (session) => {
// get next dataset.data
......@@ -106,10 +134,25 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
}
).select({
_id: 1,
collectionId: 1
collectionId: 1,
imageId: 1,
q: 1,
indexes: 1
});
if (data) {
const collection = await MongoDatasetCollection.findById(data.collectionId)
.select('imageIndex')
.session(session);
const hasMarkdownImages =
!!collection?.imageIndex && uniqueDatasetDataMarkdownImageUrls([data.q]).length > 0;
const mode = getDatasetImageTrainingMode({
supportVlm,
supportImageIndex,
imageId: data.imageId,
hasMarkdownImages
});
await MongoDatasetTraining.create(
[
{
......@@ -118,9 +161,19 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
datasetId,
collectionId: data.collectionId,
billId: usageId,
mode: TrainingModeEnum.chunk,
model: vectorModel,
mode,
model:
(mode === TrainingModeEnum.imageParse || mode === TrainingModeEnum.image) &&
supportVlm &&
availableVlmModel
? availableVlmModel.model
: vectorModel,
dataId: data._id,
...(data.imageId && { imageId: data.imageId }),
...(mode === TrainingModeEnum.image && {
q: data.q,
indexes: data.indexes
}),
retryCount: 50
}
],
......@@ -137,7 +190,7 @@ async function handler(req: ApiRequestProps): Promise<RebuildEmbeddingResponse>
if (!hasNext) {
break;
}
} catch (error) {}
} catch {}
}
return RebuildEmbeddingResponseSchema.parse({});
......
projects/app/src/service/core/dataset/queues/datasetParse.ts
......@@ -23,6 +23,7 @@ import { getLLMMaxChunkSize } from '@fastgpt/global/core/dataset/training/utils'
import { checkDatasetIndexLimit } from '@fastgpt/service/support/permission/teamLimit';
import { predictDataLimitLength } from '@fastgpt/global/core/dataset/utils';
import { getTrainingModeByCollection } from '@fastgpt/service/core/dataset/collection/utils';
import { getDatasetImageIndexCapability } from '@fastgpt/service/core/dataset/utils';
import { pushDataListToTrainingQueue } from '@fastgpt/service/core/dataset/training/controller';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
import { mongoSessionRun } from '@fastgpt/service/common/mongo/sessionRun';
......@@ -150,7 +151,7 @@ export const datasetParseQueue = async (): Promise<any> => {
return {
data
};
} catch (error) {
} catch {
return {
error: true
};
......@@ -197,7 +198,11 @@ export const datasetParseQueue = async (): Promise<any> => {
const trainingMode = getTrainingModeByCollection({
trainingType: collection.trainingType ?? DatasetCollectionDataProcessModeEnum.chunk,
autoIndexes: collection.autoIndexes,
imageIndex: collection.imageIndex
imageIndex: collection.imageIndex,
supportImageIndex: getDatasetImageIndexCapability({
vectorModel: dataset.vectorModel,
vlmModel: dataset.vlmModel
}).supportImageIndex
});
// 1. Parse rawtext
......
projects/app/src/web/core/workflow/utils.ts
......@@ -41,6 +41,32 @@ import { useSystemStore } from '@/web/common/system/useSystemStore';
import type { LLMModelItemType } from '@fastgpt/global/core/ai/model.schema';
/* ====== node ======= */
/**
* 适配从数据库读取出的节点输入。
* 旧知识库搜索节点使用 userChatInput；当前节点改为 datasetSearchInput 数组。
* 这里仅处理旧字段到新字段的 key 和 valueType 迁移。
*/
export const adaptStoreNodeInputs = (storeNode: StoreNodeItemType): FlowNodeInputItemType[] => {
if (storeNode.flowNodeType !== FlowNodeTypeEnum.datasetSearchNode) {
return storeNode.inputs;
}
return storeNode.inputs.map((input) => {
if (input.key !== NodeInputKeyEnum.userChatInput) return input;
const isReferenceValue = isValidReferenceValueFormat(input.value);
return {
...input,
key: NodeInputKeyEnum.datasetSearchInput,
label: 'workflow:search_query',
value: isReferenceValue ? [input.value] : input.value,
valueType: WorkflowIOValueTypeEnum.arrayString,
selectedTypeIndex: isReferenceValue ? 0 : 1
};
});
};
export const nodeTemplate2FlowNode = ({
template,
position,
......@@ -100,6 +126,7 @@ export const storeNode2FlowNode = ({
const dynamicInput = template.inputs.find(
(input) => input.renderTypeList[0] === FlowNodeInputTypeEnum.addInputParam
);
const adaptedStoreInputs = adaptStoreNodeInputs(storeNode);
// replace item data
const nodeItem: FlowNodeItemType = {
......@@ -113,7 +140,7 @@ export const storeNode2FlowNode = ({
inputs: templateInputs
.map<FlowNodeInputItemType>((templateInput) => {
const storeInput =
storeNode.inputs.find((item) => item.key === templateInput.key) || templateInput;
adaptedStoreInputs.find((item) => item.key === templateInput.key) || templateInput;
return {
...storeInput,
......@@ -126,7 +153,7 @@ export const storeNode2FlowNode = ({
})
.concat(
// 合并 store 中有，template 中没有的输入
storeNode.inputs
adaptedStoreInputs
.filter((item) => !templateInputs.find((input) => input.key === item.key))
.map((item) => {
const templateInput = template.inputs.find((input) => input.key === item.key);
......
projects/app/test/api/core/dataset/collection/create/images.test.ts 0 → 100644
import { beforeEach, describe, expect, it, vi } from 'vitest';
import {
DatasetCollectionDataProcessModeEnum,
DatasetCollectionTypeEnum
} from '@fastgpt/global/core/dataset/constants';
const {
mockResolveMultipleFormData,
mockClearDiskTempFiles,
mockAuthDataset,
mockCheckDatasetIndexLimit,
mockAuthFrequencyLimit,
mockGetTeamPlanStatus,
mockReadFile,
mockGetFileS3Key,
mockUploadImage2S3Bucket,
mockCreateCollectionAndInsertData,
mockGetDatasetImageIndexCapability
} = vi.hoisted(() => ({
mockResolveMultipleFormData: vi.fn(),
mockClearDiskTempFiles: vi.fn(),
mockAuthDataset: vi.fn(),
mockCheckDatasetIndexLimit: vi.fn(),
mockAuthFrequencyLimit: vi.fn(),
mockGetTeamPlanStatus: vi.fn(),
mockReadFile: vi.fn(),
mockGetFileS3Key: {
dataset: vi.fn()
},
mockUploadImage2S3Bucket: vi.fn(),
mockCreateCollectionAndInsertData: vi.fn(),
mockGetDatasetImageIndexCapability: vi.fn()
}));
vi.mock('@/service/middleware/entry', () => ({
NextAPI: (handler: any) => handler
}));
vi.mock('@fastgpt/service/common/file/multer', () => ({
multer: {
resolveMultipleFormData: mockResolveMultipleFormData,
clearDiskTempFiles: mockClearDiskTempFiles
}
}));
vi.mock('@fastgpt/service/support/permission/dataset/auth', () => ({
authDataset: mockAuthDataset
}));
vi.mock('@fastgpt/service/support/permission/teamLimit', () => ({
checkDatasetIndexLimit: mockCheckDatasetIndexLimit
}));
vi.mock('@fastgpt/service/common/system/frequencyLimit/utils', () => ({
authFrequencyLimit: mockAuthFrequencyLimit
}));
vi.mock('@fastgpt/service/support/wallet/sub/utils', () => ({
getTeamPlanStatus: mockGetTeamPlanStatus
}));
vi.mock('node:fs', () => ({
default: {
promises: {
readFile: mockReadFile
}
},
promises: {
readFile: mockReadFile
}
}));
vi.mock('@fastgpt/service/common/s3/utils', () => ({
getFileS3Key: mockGetFileS3Key,
uploadImage2S3Bucket: mockUploadImage2S3Bucket
}));
vi.mock('@fastgpt/service/core/dataset/collection/controller', () => ({
createCollectionAndInsertData: mockCreateCollectionAndInsertData
}));
vi.mock('@fastgpt/service/core/dataset/utils', async (importOriginal) => {
const actual = (await importOriginal()) as any;
return {
...actual,
getDatasetImageIndexCapability: mockGetDatasetImageIndexCapability
};
});
import handler from '@/pages/api/core/dataset/collection/create/images';
const datasetId = '68ad85a7463006c963799a07';
const parentId = '68ad85a7463006c963799a08';
describe('POST /api/core/dataset/collection/create/images', () => {
beforeEach(() => {
vi.clearAllMocks();
mockResolveMultipleFormData.mockResolvedValue({
data: {
parentId,
datasetId,
collectionName: 'Native image embedding collection'
},
fileMetadata: [
{
path: '/tmp/cat.png',
filename: 'cat.png',
mimetype: 'image/png'
}
]
});
mockAuthDataset.mockResolvedValue({
teamId: 'team-id',
tmbId: 'tmb-id',
dataset: {
_id: datasetId,
vectorModel: 'vision-embedding',
agentModel: 'gpt-5'
}
});
mockGetDatasetImageIndexCapability.mockReturnValue({
availableVlmModel: undefined,
supportVlm: false,
supportImageEmbedding: true,
supportImageIndex: true
});
mockGetTeamPlanStatus.mockResolvedValue({ standard: { maxUploadFileCount: 10 } });
mockReadFile.mockResolvedValue(Buffer.from('image-bytes'));
mockGetFileS3Key.dataset.mockReturnValue({ fileKey: 'dataset/team/cat.png' });
mockUploadImage2S3Bucket.mockResolvedValue('dataset/team/cat.png');
mockCreateCollectionAndInsertData.mockResolvedValue({
collectionId: 'collection-id',
results: {
insertLen: 1
}
});
});
it('should create an image collection with chunk training when only native image embedding is available', async () => {
const result = await handler({} as any);
expect(result).toEqual({
collectionId: 'collection-id',
results: {
insertLen: 1
}
});
expect(mockCreateCollectionAndInsertData).toHaveBeenCalledWith({
dataset: {
_id: datasetId,
vectorModel: 'vision-embedding',
agentModel: 'gpt-5'
},
imageIds: ['dataset/team/cat.png'],
createCollectionParams: {
parentId,
teamId: 'team-id',
tmbId: 'tmb-id',
datasetId,
type: DatasetCollectionTypeEnum.images,
name: 'Native image embedding collection',
trainingType: DatasetCollectionDataProcessModeEnum.chunk
}
});
expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']);
});
});
projects/app/test/api/core/dataset/data/insertImages.test.ts 0 → 100644
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { TrainingModeEnum } from '@fastgpt/global/core/dataset/constants';
const {
mockResolveMultipleFormData,
mockClearDiskTempFiles,
mockAuthDatasetCollection,
mockAuthFrequencyLimit,
mockGetTeamPlanStatus,
mockReadFile,
mockGetFileS3Key,
mockUploadImage2S3Bucket,
mockMongoSessionRun,
mockCreateTrainingUsage,
mockPushDataListToTrainingQueue,
mockGetDatasetImageIndexCapability
} = vi.hoisted(() => ({
mockResolveMultipleFormData: vi.fn(),
mockClearDiskTempFiles: vi.fn(),
mockAuthDatasetCollection: vi.fn(),
mockAuthFrequencyLimit: vi.fn(),
mockGetTeamPlanStatus: vi.fn(),
mockReadFile: vi.fn(),
mockGetFileS3Key: {
dataset: vi.fn()
},
mockUploadImage2S3Bucket: vi.fn(),
mockMongoSessionRun: vi.fn(),
mockCreateTrainingUsage: vi.fn(),
mockPushDataListToTrainingQueue: vi.fn(),
mockGetDatasetImageIndexCapability: vi.fn()
}));
vi.mock('@/service/middleware/entry', () => ({
NextAPI: (handler: any) => handler
}));
vi.mock('@fastgpt/service/common/file/multer', () => ({
multer: {
resolveMultipleFormData: mockResolveMultipleFormData,
clearDiskTempFiles: mockClearDiskTempFiles
}
}));
vi.mock('@fastgpt/service/support/permission/dataset/auth', () => ({
authDatasetCollection: mockAuthDatasetCollection
}));
vi.mock('@fastgpt/service/common/system/frequencyLimit/utils', () => ({
authFrequencyLimit: mockAuthFrequencyLimit
}));
vi.mock('@fastgpt/service/support/wallet/sub/utils', () => ({
getTeamPlanStatus: mockGetTeamPlanStatus
}));
vi.mock('node:fs', () => ({
default: {
promises: {
readFile: mockReadFile
}
},
promises: {
readFile: mockReadFile
}
}));
vi.mock('@fastgpt/service/common/s3/utils', () => ({
getFileS3Key: mockGetFileS3Key,
uploadImage2S3Bucket: mockUploadImage2S3Bucket
}));
vi.mock('@fastgpt/service/common/mongo/sessionRun', () => ({
mongoSessionRun: mockMongoSessionRun
}));
vi.mock('@fastgpt/service/support/wallet/usage/controller', () => ({
createTrainingUsage: mockCreateTrainingUsage
}));
vi.mock('@fastgpt/service/core/dataset/training/controller', () => ({
pushDataListToTrainingQueue: mockPushDataListToTrainingQueue
}));
vi.mock('@fastgpt/service/core/ai/model', async (importOriginal) => {
const actual = (await importOriginal()) as any;
return {
...actual,
getEmbeddingModel: vi.fn((model: string) => ({ name: model, model })),
getLLMModel: vi.fn((model: string) => ({ name: model, model }))
};
});
vi.mock('@fastgpt/service/core/dataset/utils', async (importOriginal) => {
const actual = (await importOriginal()) as any;
return {
...actual,
getDatasetImageIndexCapability: mockGetDatasetImageIndexCapability
};
});
import handler from '@/pages/api/core/dataset/data/insertImages';
const collectionId = '68ad85a7463006c963799a06';
const datasetId = '68ad85a7463006c963799a07';
describe('POST /api/core/dataset/data/insertImages', () => {
beforeEach(() => {
vi.clearAllMocks();
mockResolveMultipleFormData.mockResolvedValue({
data: { collectionId },
fileMetadata: [
{
path: '/tmp/cat.png',
filename: 'cat.png',
mimetype: 'image/png'
}
]
});
mockAuthDatasetCollection.mockResolvedValue({
teamId: 'team-id',
tmbId: 'tmb-id',
collection: {
_id: collectionId,
name: 'Images',
datasetId,
dataset: {
_id: datasetId,
vectorModel: 'vision-embedding',
agentModel: 'gpt-5'
}
}
});
mockGetDatasetImageIndexCapability.mockReturnValue({
availableVlmModel: undefined,
supportVlm: false,
supportImageEmbedding: true,
supportImageIndex: true
});
mockGetTeamPlanStatus.mockResolvedValue({ standard: { maxUploadFileCount: 10 } });
mockReadFile.mockResolvedValue(Buffer.from('image-bytes'));
mockGetFileS3Key.dataset.mockReturnValue({ fileKey: 'dataset/team/cat.png' });
mockUploadImage2S3Bucket.mockResolvedValue('dataset/team/cat.png');
mockMongoSessionRun.mockImplementation((fn: any) => fn('session'));
mockCreateTrainingUsage.mockResolvedValue({ usageId: 'usage-id' });
mockPushDataListToTrainingQueue.mockResolvedValue({ insertLen: 1 });
});
it('should upload images with chunk mode when only native image embedding is available', async () => {
const result = await handler({} as any);
expect(result).toEqual({});
expect(mockCreateTrainingUsage).toHaveBeenCalledWith(
expect.objectContaining({
vectorModel: 'vision-embedding',
agentModel: 'gpt-5',
vllmModel: undefined,
session: 'session'
})
);
expect(mockPushDataListToTrainingQueue).toHaveBeenCalledWith({
teamId: 'team-id',
tmbId: 'tmb-id',
datasetId,
collectionId,
agentModel: 'gpt-5',
vectorModel: 'vision-embedding',
vlmModel: undefined,
mode: TrainingModeEnum.chunk,
billId: 'usage-id',
data: [{ imageId: 'dataset/team/cat.png' }],
session: 'session'
});
expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']);
});
it('should reject image upload when neither VLM nor native image embedding is available', async () => {
mockGetDatasetImageIndexCapability.mockReturnValueOnce({
availableVlmModel: undefined,
supportVlm: false,
supportImageEmbedding: false,
supportImageIndex: false
});
await expect(handler({} as any)).rejects.toBeTruthy();
expect(mockUploadImage2S3Bucket).not.toHaveBeenCalled();
expect(mockPushDataListToTrainingQueue).not.toHaveBeenCalled();
expect(mockClearDiskTempFiles).toHaveBeenCalledWith(['/tmp/cat.png']);
});
});
projects/app/test/api/core/dataset/data/update.test.ts 0 → 100644
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { DatasetCollectionTypeEnum } from '@fastgpt/global/core/dataset/constants';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
const {
mockAuthDatasetData,
mockUpdateDatasetDataByIndexes,
mockUpdateDatasetDataSystemIndexes,
mockPushGenerateVectorUsage,
mockAddAuditLog,
mockReplaceS3KeyToPreviewUrl
} = vi.hoisted(() => ({
mockAuthDatasetData: vi.fn(),
mockUpdateDatasetDataByIndexes: vi.fn(),
mockUpdateDatasetDataSystemIndexes: vi.fn(),
mockPushGenerateVectorUsage: vi.fn(),
mockAddAuditLog: vi.fn(),
mockReplaceS3KeyToPreviewUrl: vi.fn()
}));
vi.mock('@/service/middleware/entry', () => ({
NextAPI: (handler: any) => handler
}));
vi.mock('@fastgpt/service/support/permission/dataset/auth', () => ({
authDatasetData: mockAuthDatasetData
}));
vi.mock('@/service/core/dataset/data/data', () => ({
updateDatasetDataByIndexes: mockUpdateDatasetDataByIndexes,
updateDatasetDataSystemIndexes: mockUpdateDatasetDataSystemIndexes
}));
vi.mock('@/service/support/wallet/usage/push', () => ({
pushGenerateVectorUsage: mockPushGenerateVectorUsage
}));
vi.mock('@fastgpt/service/support/user/audit/util', () => ({
addAuditLog: mockAddAuditLog,
getI18nDatasetType: vi.fn((type: string) => type)
}));
vi.mock('@fastgpt/service/core/dataset/utils', () => ({
replaceS3KeyToPreviewUrl: mockReplaceS3KeyToPreviewUrl
}));
import handler from '@/pages/api/core/dataset/data/update';
const dataId = '68ad85a7463006c963799a05';
const buildAuthResult = () => ({
teamId: 'team-id',
tmbId: 'tmb-id',
collection: {
name: 'Collection',
indexPrefixTitle: true,
indexSize: 256,
imageIndex: true,
type: DatasetCollectionTypeEnum.images,
dataset: {
name: 'Dataset',
type: 'dataset',
vectorModel: 'vision-embedding',
vlmModel: 'vlm-model'
}
},
datasetData: {
q: 'old question',
a: 'old answer',
imageId: 'dataset/team/main.png',
indexes: [
{
type: DatasetDataIndexTypeEnum.custom,
text: 'old custom',
dataId: 'custom_old'
},
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/main.png',
dataId: 'image_embedding_old'
}
]
}
});
describe('PUT /api/core/dataset/data/update', () => {
beforeEach(() => {
vi.clearAllMocks();
mockAuthDatasetData.mockResolvedValue(buildAuthResult());
mockReplaceS3KeyToPreviewUrl.mockImplementation((text: string) => text);
mockUpdateDatasetDataByIndexes.mockResolvedValue({ tokens: 12 });
mockUpdateDatasetDataSystemIndexes.mockResolvedValue({ tokens: 0 });
});
it('should keep legacy indexes update API compatible while rebuilding system image embedding indexes', async () => {
const result = await handler({
body: {
dataId,
q: 'new question',
a: 'new answer',
indexes: [
{
type: DatasetDataIndexTypeEnum.custom,
text: 'new custom'
},
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/client-should-not-replace.png'
}
]
}
} as any);
expect(mockUpdateDatasetDataByIndexes).toHaveBeenCalledWith({
dataId,
q: 'new question',
a: 'new answer',
imageId: 'dataset/team/main.png',
imageIndex: true,
indexes: [
{
type: DatasetDataIndexTypeEnum.custom,
text: 'new custom'
}
],
model: 'vision-embedding',
indexSize: 256,
indexPrefix: '# Collection'
});
expect(mockUpdateDatasetDataSystemIndexes).not.toHaveBeenCalled();
expect(mockPushGenerateVectorUsage).toHaveBeenCalledWith({
teamId: 'team-id',
tmbId: 'tmb-id',
inputTokens: 12,
model: 'vision-embedding'
});
expect(result).toEqual({
q: 'new question',
a: 'new answer'
});
});
it('should pass an explicit empty question to the index update path', async () => {
await handler({
body: {
dataId,
q: '',
a: '',
indexes: []
}
} as any);
expect(mockUpdateDatasetDataByIndexes).toHaveBeenCalledWith(
expect.objectContaining({
dataId,
q: '',
a: '',
imageId: 'dataset/team/main.png',
imageIndex: true,
indexes: []
})
);
});
it('should pass image context when rebuilding generated indexes', async () => {
await handler({
body: {
dataId,
q: 'new question ![new](dataset/team/new.png)'
}
} as any);
expect(mockUpdateDatasetDataSystemIndexes).toHaveBeenCalledWith({
dataId,
q: 'new question ![new](dataset/team/new.png)',
a: 'old answer',
imageId: 'dataset/team/main.png',
imageIndex: true,
model: 'vision-embedding',
indexSize: 256,
indexPrefix: '# Collection'
});
expect(mockUpdateDatasetDataByIndexes).not.toHaveBeenCalled();
});
});
projects/app/test/api/core/dataset/training/rebuildEmbedding.test.ts 0 → 100644
import { beforeEach, describe, expect, it } from 'vitest';
import handler from '@/pages/api/core/dataset/training/rebuildEmbedding';
import { MongoDatasetCollection } from '@fastgpt/service/core/dataset/collection/schema';
import { MongoDatasetData } from '@fastgpt/service/core/dataset/data/schema';
import { MongoDataset } from '@fastgpt/service/core/dataset/schema';
import { MongoDatasetTraining } from '@fastgpt/service/core/dataset/training/schema';
import {
DatasetCollectionTypeEnum,
TrainingModeEnum
} from '@fastgpt/global/core/dataset/constants';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
import { getRootUser } from '@test/datas/users';
import { Call } from '@test/utils/request';
const registerEmbeddingModel = ({ model, vision = false }: { model: string; vision?: boolean }) => {
global.embeddingModelMap.set(model, {
...global.systemDefaultModel.embedding,
model,
name: model,
vision
});
};
const registerVlmModel = (model: string) => {
global.llmModelMap.set(model, {
...global.systemDefaultModel.llm,
model,
name: model,
vision: true
});
};
const createDatasetContext = async ({ vlmModel }: { vlmModel?: string } = {}) => {
const root = await getRootUser();
const dataset = await MongoDataset.create({
name: 'test dataset',
teamId: root.teamId,
tmbId: root.tmbId,
vectorModel: 'old-embedding',
agentModel: 'gpt-5',
vlmModel
});
const collection = await MongoDatasetCollection.create({
name: 'test collection',
type: DatasetCollectionTypeEnum.file,
teamId: root.teamId,
tmbId: root.tmbId,
datasetId: dataset._id,
imageIndex: true
});
return { root, dataset, collection };
};
describe('POST /api/core/dataset/training/rebuildEmbedding', () => {
beforeEach(() => {
global.systemEnv = {
...global.systemEnv,
vectorMaxProcess: 1
};
registerEmbeddingModel({ model: 'old-embedding' });
registerEmbeddingModel({ model: 'vision-embedding', vision: true });
registerEmbeddingModel({ model: 'text-only-embedding' });
registerVlmModel('dataset-vlm-model');
});
it('should keep image index and enqueue image mode when the new embedding model supports images', async () => {
const { root, dataset, collection } = await createDatasetContext();
const data = await MongoDatasetData.create({
teamId: root.teamId,
tmbId: root.tmbId,
datasetId: dataset._id,
collectionId: collection._id,
q: 'question with ![cat](dataset/team/cat.png)',
indexes: [
{
type: DatasetDataIndexTypeEnum.custom,
text: 'manual index',
dataId: 'manual_id'
}
]
});
const res = await Call(handler, {
auth: root,
body: {
datasetId: String(dataset._id),
vectorModel: 'vision-embedding'
}
});
const updatedDataset = await MongoDataset.findById(dataset._id).lean();
const updatedCollection = await MongoDatasetCollection.findById(collection._id).lean();
const training = await MongoDatasetTraining.findOne({ dataId: data._id }).lean();
expect(res.code).toBe(200);
expect(updatedDataset?.vectorModel).toBe('vision-embedding');
expect(updatedCollection?.imageIndex).toBe(true);
expect(training).toEqual(
expect.objectContaining({
mode: TrainingModeEnum.image,
q: 'question with ![cat](dataset/team/cat.png)',
retryCount: 50
})
);
expect(training?.indexes).toEqual([
expect.objectContaining({
type: DatasetDataIndexTypeEnum.custom,
text: 'manual index'
})
]);
});
it('should disable image index and enqueue chunk mode when the new embedding model has no image capability', async () => {
const { root, dataset, collection } = await createDatasetContext();
const data = await MongoDatasetData.create({
teamId: root.teamId,
tmbId: root.tmbId,
datasetId: dataset._id,
collectionId: collection._id,
q: 'question with ![cat](dataset/team/cat.png)'
});
const res = await Call(handler, {
auth: root,
body: {
datasetId: String(dataset._id),
vectorModel: 'text-only-embedding'
}
});
const updatedDataset = await MongoDataset.findById(dataset._id).lean();
const updatedCollection = await MongoDatasetCollection.findById(collection._id).lean();
const training = await MongoDatasetTraining.findOne({ dataId: data._id }).lean();
expect(res.code).toBe(200);
expect(updatedDataset?.vectorModel).toBe('text-only-embedding');
expect(updatedDataset?.chunkSettings?.imageIndex).toBe(false);
expect(updatedCollection?.imageIndex).toBe(false);
expect(training).toEqual(
expect.objectContaining({
mode: TrainingModeEnum.chunk,
retryCount: 50
})
);
expect(training?.q).toBe('');
});
it('should enqueue imageParse mode with VLM model for image data when VLM is configured', async () => {
const { root, dataset, collection } = await createDatasetContext({
vlmModel: 'dataset-vlm-model'
});
const data = await MongoDatasetData.create({
teamId: root.teamId,
tmbId: root.tmbId,
datasetId: dataset._id,
collectionId: collection._id,
q: '',
imageId: 'dataset/team/main.png'
});
const res = await Call(handler, {
auth: root,
body: {
datasetId: String(dataset._id),
vectorModel: 'text-only-embedding'
}
});
const training = await MongoDatasetTraining.findOne({ dataId: data._id }).lean();
expect(res.code).toBe(200);
expect(training).toEqual(
expect.objectContaining({
mode: TrainingModeEnum.imageParse,
imageId: 'dataset/team/main.png',
retryCount: 50
})
);
});
});
projects/app/test/service/core/dataset/queues/generateVector.test.ts 0 → 100644
import { beforeEach, describe, expect, it } from 'vitest';
import { getRebuildBaseIndexes } from '@/service/core/dataset/queues/generateVector';
import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
const visionEmbeddingModel = {
model: 'vision-embedding',
name: 'vision-embedding',
maxToken: 100,
vision: true
} as any;
describe('generateVector image embedding helpers', () => {
beforeEach(() => {
global.embeddingModelMap.set(visionEmbeddingModel.model, visionEmbeddingModel);
global.llmModelMap.set('vlm-model', {
...global.systemDefaultModel.llm,
model: 'vlm-model',
name: 'vlm-model',
vision: true
});
});
it('should drop system indexes and keep supported external image description indexes when rebuilding', () => {
const result = getRebuildBaseIndexes({
indexes: [
{ type: DatasetDataIndexTypeEnum.default, text: 'old default', dataId: 'default_id' },
{ type: DatasetDataIndexTypeEnum.custom, text: 'manual', dataId: 'manual_id' },
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/main.png',
dataId: 'main_vector_id'
},
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/stale.png',
dataId: 'stale_vector_id'
},
{
type: DatasetDataIndexTypeEnum.image,
text: 'image description',
dataId: 'image_desc_id'
}
],
q: 'content ![markdown](dataset/team/markdown.png)',
dataset: {
vectorModel: visionEmbeddingModel.model,
vlmModel: 'vlm-model'
},
collection: {
imageIndex: true
},
data: {
imageId: 'dataset/team/main.png',
indexes: []
}
} as any);
expect(result).toEqual([
{ type: DatasetDataIndexTypeEnum.custom, text: 'manual', dataId: 'manual_id' },
{
type: DatasetDataIndexTypeEnum.image,
text: 'image description',
dataId: 'image_desc_id'
}
]);
});
it('should drop VLM image description indexes when collection image index is disabled', () => {
const result = getRebuildBaseIndexes({
indexes: [
{ type: DatasetDataIndexTypeEnum.custom, text: 'manual', dataId: 'manual_id' },
{
type: DatasetDataIndexTypeEnum.image,
text: 'image description',
dataId: 'image_desc_id'
},
{
type: DatasetDataIndexTypeEnum.imageEmbedding,
text: 'dataset/team/main.png',
dataId: 'main_vector_id'
}
],
dataset: {
vectorModel: visionEmbeddingModel.model,
vlmModel: 'vlm-model'
},
collection: {
imageIndex: false
},
data: {
imageId: 'dataset/team/main.png',
indexes: []
}
} as any);
expect(result).toEqual([
{ type: DatasetDataIndexTypeEnum.custom, text: 'manual', dataId: 'manual_id' }
]);
});
});
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