图片上传到新建的 dataset_collection_images 表，逻辑跟随更改

Author: ctrlz526

packages/global/core/dataset/imageCollection.d.ts

@@ -0,0 +1,19 @@
+export interface DatasetCollectionImageSchema {
+  _id: string;
+  teamId: string;
+  datasetId: string;
+  collectionId?: string;
+  createTime: Date;
+  expiredTime: Date;
+  size: number;
+  name: string;
+  path: string;
+  contentType: string;
+  metadata?: Record<string, any>;
+}
+
+// API请求参数类型
+export interface UploadDatasetImageProps {
+  datasetId: string;
+  collectionId?: string;
+}

packages/service/core/ai/model.ts

@@ -20,6 +20,15 @@ export const getVlmModel = (model?: string) => {
     ?.find((item) => item.model === model || item.name === model);
 };
 
+export const getVlmModelList = () => {
+  return Array.from(global.llmModelMap.values())?.filter((item) => item.vision) || [];
+};
+
+export const hasAvailableVlmModel = () => {
+  const vlmModels = getVlmModelList();
+  return vlmModels.length > 0;
+};
+
 export const getDefaultEmbeddingModel = () => global?.systemDefaultModel.embedding!;
 export const getEmbeddingModel = (model?: string) => {
   if (!model) return getDefaultEmbeddingModel();

packages/service/core/dataset/collection/controller.ts

@@ -38,6 +38,7 @@ import {
   getLLMMaxChunkSize
 } from '@fastgpt/global/core/dataset/training/utils';
 import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
+import { MongoDatasetCollectionImage } from '../schema';
 
 export const createCollectionAndInsertData = async ({
   dataset,
@@ -403,7 +404,12 @@ export async function delCollection({
           ]
         : []),
       // Delete vector data
-      deleteDatasetDataVector({ teamId, datasetIds, collectionIds })
+      deleteDatasetDataVector({ teamId, datasetIds, collectionIds }),
+      // Delete collection images
+      MongoDatasetCollectionImage.deleteMany({
+        teamId,
+        collectionId: { $in: collectionIds }
+      })
     ]);
 
     // delete collections

packages/service/core/dataset/collection/controller_imageFileId.ts

@@ -13,7 +13,6 @@ import type {
 import { MongoDatasetTraining } from '../training/schema';
 import { MongoDatasetData } from '../data/schema';
 import { delImgByRelatedId } from '../../../common/file/image/controller';
-// import { deleteDatasetDataVector } from '../../../common/vectorStore/controller';
 import { deleteDatasetDataVector } from '../../../common/vectorDB/controller';
 import { delFileByFileIdList } from '../../../common/file/gridfs/controller';
 import { BucketNameEnum } from '@fastgpt/global/common/file/constants';
@@ -38,6 +37,8 @@ import {
   computeChunkSplitter,
   getLLMMaxChunkSize
 } from '@fastgpt/global/core/dataset/training/utils';
+import { DatasetDataIndexTypeEnum } from '@fastgpt/global/core/dataset/data/constants';
+import { MongoDatasetCollectionImage } from '../schema';
 
 export const createCollectionAndInsertData = async ({
   dataset,
@@ -84,8 +85,7 @@ export const createCollectionAndInsertData = async ({
     chunkSize,
     maxSize: getLLMMaxChunkSize(getLLMModel(dataset.agentModel)),
     overlapRatio: trainingType === DatasetCollectionDataProcessModeEnum.chunk ? 0.2 : 0,
-    customReg: chunkSplitter ? [chunkSplitter] : [],
-    isQAImport
+    customReg: chunkSplitter ? [chunkSplitter] : []
   });
 
   // 2. auth limit
@@ -102,8 +102,7 @@ export const createCollectionAndInsertData = async ({
   });
 
   const fn = async (session: ClientSession) => {
-    // 3. create collection,拿着传入的collectionId
-
+    // 3. Use the passed collectionId and do not create a new collection.
     // 4. create training bill
     const traingBillId = await (async () => {
       if (billId) return billId;
@@ -120,19 +119,26 @@ export const createCollectionAndInsertData = async ({
       return newBillId;
     })();
 
-    // 5. insert to training queue
-    console.log(`[Image parse queue] Pushing data to training queue:`, {
-      datasetId: dataset._id,
-      collectionId: collectionId || '',
-      mode: getTrainingModeByCollection({
-        trainingType: trainingType,
-        autoIndexes: createCollectionParams.autoIndexes,
-        imageIndex: createCollectionParams.imageIndex,
-        isImageCollection: createCollectionParams.metadata?.isImageCollection === true
-      }),
-      chunksCount: chunks.length
-    });
+    // 5. Update the collectionId field in the image record
+    if (createCollectionParams.metadata?.isImageCollection && createCollectionParams.fileId) {
+      await MongoDatasetCollectionImage.updateOne(
+        {
+          _id: createCollectionParams.fileId,
+          teamId: teamId
+        },
+        {
+          $set: {
+            collectionId: collectionId
+          },
+          $unset: {
+            expiredTime: 1
+          }
+        },
+        { session }
+      );
+    }
 
+    // 6. insert to training queue
     const insertResults = await pushDataListToTrainingQueue({
       teamId,
       tmbId,
@@ -152,6 +158,10 @@ export const createCollectionAndInsertData = async ({
       billId: traingBillId,
       data: chunks.map((item, index) => ({
         ...item,
+        indexes: item.indexes?.map((text) => ({
+          type: DatasetDataIndexTypeEnum.custom,
+          text
+        })),
         chunkIndex: index
       })),
       session
@@ -177,7 +187,7 @@ export const createCollectionAndInsertData = async ({
     }
 
     return {
-      collectionId,
+      collectionId: collectionId,
       insertResults
     };
   };
@@ -382,7 +392,12 @@ export async function delCollection({
           ]
         : []),
       // Delete vector data
-      deleteDatasetDataVector({ teamId, datasetIds, collectionIds })
+      deleteDatasetDataVector({ teamId, datasetIds, collectionIds }),
+      // Delete collection images
+      MongoDatasetCollectionImage.deleteMany({
+        teamId,
+        collectionId: { $in: collectionIds }
+      })
     ]);
 
     // delete collections
@@ -413,16 +428,20 @@ export async function pushImageFileToTrainingQueue({
   billId?: string;
   model?: string;
 }) {
+  const mongoose = require('mongoose');
+  const ObjectId = mongoose.Types.ObjectId;
+
   await MongoDatasetTraining.create({
-    teamId,
-    tmbId,
-    datasetId,
-    collectionId,
+    teamId: new ObjectId(teamId),
+    tmbId: new ObjectId(tmbId),
+    datasetId: new ObjectId(datasetId),
+    collectionId: new ObjectId(collectionId),
     billId,
     mode: TrainingModeEnum.imageParse,
     model,
     imageFileId,
     retryCount: 5,
+    lockTime: new Date('2000/1/1'),
     indexes: []
   });
 }

packages/service/core/dataset/controller.ts

@@ -1,6 +1,6 @@
 import { type DatasetSchemaType } from '@fastgpt/global/core/dataset/type';
 import { MongoDatasetCollection } from './collection/schema';
-import { MongoDataset } from './schema';
+import { MongoDataset, MongoDatasetCollectionImage } from './schema';
 import { delCollectionRelatedSource } from './collection/controller';
 import { type ClientSession } from '../../common/mongo';
 import { MongoDatasetTraining } from './training/schema';
@@ -9,6 +9,7 @@ import { deleteDatasetDataVector } from '../../common/vectorDB/controller';
 import { MongoDatasetDataText } from './data/dataTextSchema';
 import { DatasetErrEnum } from '@fastgpt/global/common/error/code/dataset';
 import { retryFn } from '@fastgpt/global/common/system/utils';
+import { addMinutes } from 'date-fns';
 
 /* ============= dataset ========== */
 /* find all datasetId by top datasetId */
@@ -105,7 +106,12 @@ export async function delDatasetRelevantData({
       // Delete Image and file
       delCollectionRelatedSource({ collections }),
       // Delete vector data
-      deleteDatasetDataVector({ teamId, datasetIds })
+      deleteDatasetDataVector({ teamId, datasetIds }),
+      // Delete dataset collection images
+      MongoDatasetCollectionImage.deleteMany({
+        teamId,
+        datasetId: { $in: datasetIds }
+      })
     ]);
   });
 
@@ -115,3 +121,57 @@ export async function delDatasetRelevantData({
     datasetId: { $in: datasetIds }
   }).session(session);
 }
+
+/* ============= dataset images ========== */
+
+export async function createDatasetImage({
+  teamId,
+  datasetId,
+  collectionId,
+  name,
+  path,
+  contentType,
+  size,
+  metadata = {}
+}: {
+  teamId: string;
+  datasetId: string;
+  collectionId?: string;
+  name: string;
+  path: string;
+  contentType: string;
+  size: number;
+  metadata?: Record<string, any>;
+}): Promise<string> {
+  // Set TTL to 30min
+  const expiredTime = addMinutes(new Date(), 30);
+
+  const image = await MongoDatasetCollectionImage.create({
+    teamId: String(teamId),
+    datasetId: String(datasetId),
+    collectionId: collectionId ? String(collectionId) : null,
+    name,
+    path,
+    contentType,
+    size,
+    metadata,
+    createTime: new Date(),
+    expiredTime
+  });
+
+  return String(image._id);
+}
+
+export async function getDatasetImage(imageId: string): Promise<any> {
+  try {
+    if (!imageId || imageId.length !== 24) {
+      return null;
+    }
+
+    const result = await MongoDatasetCollectionImage.findById(imageId).lean();
+
+    return result;
+  } catch (error) {
+    return null;
+  }
+}

packages/service/core/dataset/schema.ts

@@ -13,6 +13,8 @@ import {
   TeamMemberCollectionName
 } from '@fastgpt/global/support/user/team/constant';
 import type { DatasetSchemaType } from '@fastgpt/global/core/dataset/type.d';
+import type { DatasetCollectionImageSchema } from '@fastgpt/global/core/dataset/imageCollection';
+import mongoose from 'mongoose';
 
 export const DatasetCollectionName = 'datasets';
 
@@ -146,3 +148,69 @@ try {
 }
 
 export const MongoDataset = getMongoModel<DatasetSchemaType>(DatasetCollectionName, DatasetSchema);
+
+export const DatasetCollectionImageCollectionName = 'dataset_collection_images';
+
+if (!mongoose.modelNames().includes('dataset_collection_images')) {
+  const DatasetCollectionImageSchema = new Schema({
+    teamId: {
+      type: String,
+      required: true
+    },
+    datasetId: {
+      type: String,
+      required: true
+    },
+    collectionId: {
+      type: String,
+      required: false,
+      default: null
+    },
+    name: {
+      type: String,
+      required: true
+    },
+    path: {
+      type: String,
+      required: true
+    },
+    contentType: {
+      type: String,
+      required: true
+    },
+    size: {
+      type: Number,
+      required: true
+    },
+    metadata: {
+      type: Object,
+      default: {}
+    },
+    createTime: {
+      type: Date,
+      default: Date.now
+    },
+    expiredTime: {
+      type: Date,
+      required: true,
+      index: { expireAfterSeconds: 0 }
+    }
+  });
+
+  DatasetCollectionImageSchema.index({ expiredTime: 1 }, { expireAfterSeconds: 0 });
+
+  DatasetCollectionImageSchema.index({ teamId: 1 });
+  DatasetCollectionImageSchema.index({ datasetId: 1 });
+  DatasetCollectionImageSchema.index({ collectionId: 1 });
+
+  mongoose.model(
+    'dataset_collection_images',
+    DatasetCollectionImageSchema,
+    'dataset_collection_images'
+  );
+}
+
+export const MongoDatasetCollectionImage = getMongoModel<DatasetCollectionImageSchema>(
+  DatasetCollectionImageCollectionName,
+  mongoose.model('dataset_collection_images').schema
+);