package internal import ( "context" "github.com/gogf/gf/util/guid" "yx-dataset-server/app/errors" "yx-dataset-server/app/model" "yx-dataset-server/app/schema" "yx-dataset-server/library/ragflow" ) // NewDataset 创建Dataset func NewDataset( mDataset model.IDataset, mFile model.IDatasetFile, mTrans model.ITrans, mUser model.IUser, ) *Dataset { return &Dataset{ DatasetModel: mDataset, fileModel: mFile, transModel: mTrans, userModel: mUser, } } // Dataset 创建Dataset对象 type Dataset struct { DatasetModel model.IDataset fileModel model.IDatasetFile transModel model.ITrans userModel model.IUser } // Query 查询数据 func (a *Dataset) Query(ctx context.Context, params schema.DatasetQueryParam, opts ...schema.DatasetQueryOptions) (*schema.DatasetQueryResult, error) { result, err := a.DatasetModel.Query(ctx, params, opts...) if err != nil { return nil, err } user, err := a.userModel.Query(ctx, schema.UserQueryParam{RoleCode: []string{"11", "12"}}) if err != nil { return nil, err } result.Data.FillCreator(user.Data) return result, nil } // Get 查询指定数据 func (a *Dataset) Get(ctx context.Context, recordID string, opts ...schema.DatasetQueryOptions) (*schema.Dataset, error) { item, err := a.DatasetModel.Get(ctx, recordID, opts...) if err != nil { return nil, err } else if item == nil { return nil, errors.ErrNotFound } file, err := a.fileModel.Query(ctx, schema.DatasetFileQueryParam{DatasetId: recordID}) if len(file.Data) > 0 { item.Files = file.Data } return item, nil } func (a *Dataset) getUpdate(ctx context.Context, recordID string) (*schema.Dataset, error) { return a.Get(ctx, recordID) } // Create 创建数据 func (a *Dataset) Create(ctx context.Context, item schema.Dataset) error { item.RecordID = guid.S() //item.Sequence = 0 //sequence, err := a.DatasetModel.GetMaxSequence(ctx) //if err != nil { // return err //} //if sequence > 0 { // item.Sequence = sequence + 1 //} item.CreatorId = GetUserID(ctx) if !CheckIsRootUser(ctx) { user, err := a.userModel.Get(ctx, GetUserID(ctx)) if err != nil { return err } item.OrgId = user.OrgId } data, err := ragflow.GetHttpClient().CreateDataset(ctx, &ragflow.CreateDatasetReq{ Name: item.Name, ChunkMethod: "naive", EmbeddingModel: "text-embedding-v4@Tongyi-Qianwen", }) if err != nil { return err } item.RagDataId = data.Data.ID _, err = ragflow.GetHttpClient().UpdateDataset(ctx, item.RagDataId, &ragflow.UpdateDatasetReq{ Language: "Chinese", }) if err != nil { return err } return a.DatasetModel.Create(ctx, item) } // Update 更新数据 func (a *Dataset) Update(ctx context.Context, recordID string, item schema.Dataset) error { oldItem, err := a.DatasetModel.Get(ctx, recordID) if err != nil { return err } else if oldItem == nil { return errors.ErrNotFound } if item.Name != oldItem.Name { _, err = ragflow.GetHttpClient().UpdateDataset(ctx, item.RagDataId, &ragflow.UpdateDatasetReq{ Name: item.Name, ChunkMethod: "naive", }) if err != nil { return err } } err = a.DatasetModel.Update(ctx, recordID, item) return err } // Delete 删除数据 func (a *Dataset) Delete(ctx context.Context, recordID string) error { oldItem, err := a.DatasetModel.Get(ctx, recordID) if err != nil { return err } else if oldItem == nil { return errors.ErrNotFound } err = ExecTrans(ctx, a.transModel, func(ctx context.Context) error { err = a.fileModel.DeleteByDatasetIds(ctx, []string{recordID}) if err != nil { return err } err = a.DatasetModel.Delete(ctx, recordID) if err != nil { return err } _, err = ragflow.GetHttpClient().DeleteDataset(ctx, []string{oldItem.RagDataId}) return err }) return err } // UpdateStatus 更新状态 func (a *Dataset) UpdateStatus(ctx context.Context, recordID string, status int) error { oldItem, err := a.DatasetModel.Get(ctx, recordID) if err != nil { return err } else if oldItem == nil { return errors.ErrNotFound } return a.DatasetModel.UpdateStatus(ctx, recordID, status) } // UpdateSequence 更新排序 //func (a *Dataset) UpdateSequence(ctx context.Context, recordID string, sequence int) error { // item, err := a.DatasetModel.Get(ctx, recordID) // if err != nil { // return err // } else if item == nil { // return errors.ErrNotFound // } // err = ExecTrans(ctx, a.transModel, func(ctx context.Context) error { // err = a.DatasetModel.UpdateBatchSequenceMinus(ctx, item.Sequence, sequence) // if err != nil { // return err // } // return a.DatasetModel.UpdateSequence(ctx, recordID, sequence) // }) // return err //}