yx-dataset-server/app/bll/impl/internal/b_dataset.go

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
//}