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,
	mRole model.IRole,
	mRelation model.IDatasetRelation,
	mOrg model.IOrganization,
) *Dataset {
	return &Dataset{
		DatasetModel:  mDataset,
		fileModel:     mFile,
		transModel:    mTrans,
		userModel:     mUser,
		roleModel:     mRole,
		relationModel: mRelation,
		orgModel:      mOrg,
	}
}

// Dataset 创建Dataset对象
type Dataset struct {
	DatasetModel  model.IDataset
	fileModel     model.IDatasetFile
	transModel    model.ITrans
	userModel     model.IUser
	roleModel     model.IRole
	relationModel model.IDatasetRelation
	orgModel      model.IOrganization
}

// Query 查询数据
func (a *Dataset) Query(ctx context.Context, params schema.DatasetQueryParam, opts ...schema.DatasetQueryOptions) (*schema.DatasetQueryResult, error) {

	if isSys, _ := IsSystemAdmin(ctx, a.userModel, a.roleModel); !isSys {
		// 非系统管理员：仅能看到自己有授权的知识库（user 维度 + 所属组织维度）
		userID := GetUserID(ctx)
		bizIds := []string{userID}
		if u, err := a.userModel.Get(ctx, userID); err == nil && u != nil && u.OrgId != "" {
			bizIds = append(bizIds, u.OrgId)
		}
		rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{BizIds: bizIds})
		if err != nil {
			return nil, err
		}
		if len(rel.Data) == 0 {
			return &schema.DatasetQueryResult{
				Data:       make(schema.Datasets, 0),
				PageResult: &schema.PaginationResult{Total: 0},
			}, nil
		}
		params.RecordIds = rel.Data.ToDatasetIds()
	}

	result, err := a.DatasetModel.Query(ctx, params, opts...)
	if err != nil {
		return nil, err
	}
	users, err := a.userModel.Query(ctx, schema.UserQueryParam{RoleCode: []string{RoleCodeSystemAdmin, RoleCodeEnterpriseAdmin}})
	if err != nil {
		return nil, err
	}
	result.Data.FillCreator(users.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 err != nil {
		return nil, err
	}
	if len(file.Data) > 0 {
		item.Files = file.Data
	}
	return item, nil
}

// Create 创建数据
// 依据角色决定 type：
//   - 系统管理员 / root → 公共/共享知识库 (type=1)
//   - 企业管理员           → 企业知识库 (type=2)，biz_id=org_id
//   - 员工                  → 个人知识库 (type=3)，biz_id=user_id
func (a *Dataset) Create(ctx context.Context, item schema.Dataset) error {
	item.RecordID = guid.S()
	item.CreatorId = GetUserID(ctx)

	// 非 root：取创建者所属组织
	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

	return ExecTrans(ctx, a.transModel, func(ctx context.Context) error {
		if err := a.DatasetModel.Create(ctx, item); err != nil {
			return err
		}
		// 根据类型建立初始关系映射
		switch item.Type {
		case schema.DatasetTypePublic:
			// 共享知识库：由系统管理员建立，初始不绑定到任何业务实体；
			// 通过新建/编辑组织或分配给员工时再写入映射记录。
			return nil
		case schema.DatasetTypeOrg:
			// 企业知识库：自动给所属组织建立映射
			if err := a.relationModel.Create(ctx, schema.DatasetRelation{
				RecordID:  guid.S(),
				DatasetId: item.RecordID,
				BizId:     item.OrgId,
				Type:      schema.DatasetTypeOrg,
				CreatorId: item.CreatorId,
			}); err != nil {
				return err
			}
			// 组织内的企业管理员自动获得该知识库访问权
			admins, err := a.userModel.Query(ctx, schema.UserQueryParam{OrgId: item.OrgId, RoleCode: []string{RoleCodeEnterpriseAdmin}})
			if err != nil {
				return err
			}
			for _, v := range admins.Data {
				if err := a.relationModel.Create(ctx, schema.DatasetRelation{
					RecordID:  guid.S(),
					DatasetId: item.RecordID,
					BizId:     v.RecordID,
					Type:      schema.DatasetTypeOrg,
					CreatorId: item.CreatorId,
				}); err != nil {
					return err
				}
			}
			return nil
		case schema.DatasetTypePersonal:
			// 个人知识库：仅给创建者本人建立映射
			return a.relationModel.Create(ctx, schema.DatasetRelation{
				RecordID:  guid.S(),
				DatasetId: item.RecordID,
				BizId:     item.CreatorId,
				Type:      schema.DatasetTypePersonal,
				CreatorId: item.CreatorId,
			})
		}
		return nil
	})
}

// Update 更新数据
func (a *Dataset) Update(ctx context.Context, recordID string, item schema.Dataset) error {
	isSys, err := IsSystemAdmin(ctx, a.userModel, a.roleModel)
	if err != nil {
		return err
	}
	oldItem, err := a.DatasetModel.Get(ctx, recordID)
	if err != nil {
		return err
	} else if oldItem == nil {
		return errors.ErrNotFound
	}
	if !isSys {
		u, err := a.userModel.Get(ctx, GetUserID(ctx))
		if err != nil {
			return err
		}
		if u == nil {
			return errors.ErrUserNotFound
		}
		// 企业自建：限本组织；个人：限本人
		switch oldItem.Type {
		case schema.DatasetTypeOrg:
			if oldItem.OrgId != u.OrgId {
				return errors.New400Response("仅能修改本组织自建知识库")
			}
		case schema.DatasetTypePersonal:
			if oldItem.CreatorId != u.RecordID {
				return errors.New400Response("仅能修改本人个人知识库")
			}
		default:
			return errors.New400Response("仅系统管理员可修改共享知识库")
		}
	}
	// 保留 type/creator_id，不允许变更
	item.Type = oldItem.Type
	item.CreatorId = oldItem.CreatorId
	item.OrgId = oldItem.OrgId

	if item.Name != oldItem.Name {
		if _, err := ragflow.GetHttpClient().UpdateDataset(ctx, oldItem.RagDataId, &ragflow.UpdateDatasetReq{
			Name:        item.Name,
			ChunkMethod: "naive",
		}); err != nil {
			return err
		}
	}
	return a.DatasetModel.Update(ctx, recordID, item)
}

// Delete 删除数据
func (a *Dataset) Delete(ctx context.Context, recordID string) error {
	isSys, err := IsSystemAdmin(ctx, a.userModel, a.roleModel)
	if err != nil {
		return err
	}
	oldItem, err := a.DatasetModel.Get(ctx, recordID)
	if err != nil {
		return err
	} else if oldItem == nil {
		return errors.ErrNotFound
	}
	if !isSys {
		u, err := a.userModel.Get(ctx, GetUserID(ctx))
		if err != nil {
			return err
		}
		if u == nil {
			return errors.ErrUserNotFound
		}
		switch oldItem.Type {
		case schema.DatasetTypeOrg:
			if oldItem.OrgId != u.OrgId {
				return errors.New400Response("仅能删除本组织自建知识库")
			}
		case schema.DatasetTypePersonal:
			if oldItem.CreatorId != u.RecordID {
				return errors.New400Response("仅能删除本人个人知识库")
			}
		default:
			return errors.New400Response("仅系统管理员可删除共享知识库")
		}
	}

	return ExecTrans(ctx, a.transModel, func(ctx context.Context) error {
		if err := a.relationModel.DeleteByDatasetId(ctx, recordID); err != nil {
			return err
		}
		if err := a.fileModel.DeleteByDatasetIds(ctx, []string{recordID}); err != nil {
			return err
		}
		if err := a.DatasetModel.Delete(ctx, recordID); err != nil {
			return err
		}
		_, err = ragflow.GetHttpClient().DeleteDataset(ctx, []string{oldItem.RagDataId})
		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)
}

// GetPermissionDatasets 获取有权限的知识库（按组织分组）
// - orgId 非空：返回该组织维度映射到的知识库（用于前端“组织可选知识库”列表）
// - orgId 空：返回当前登录用户有权限的知识库
func (a *Dataset) GetPermissionDatasets(ctx context.Context, orgId string) (schema.Organizations, error) {
	var datasetQueryParm schema.DatasetQueryParam
	if orgId != "" {
		rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{BizId: orgId})
		if err != nil {
			return nil, err
		}
		datasetQueryParm.RecordIds = rel.Data.ToDatasetIds()
	} else if !CheckIsRootUser(ctx) {
		userID := GetUserID(ctx)
		bizIds := []string{userID}
		if u, err := a.userModel.Get(ctx, userID); err == nil && u != nil && u.OrgId != "" {
			bizIds = append(bizIds, u.OrgId)
		}
		rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{BizIds: bizIds})
		if err != nil {
			return nil, err
		}
		datasetQueryParm.RecordIds = rel.Data.ToDatasetIds()
	}

	datasets, err := a.DatasetModel.Query(ctx, datasetQueryParm)
	if err != nil {
		return nil, err
	}
	orgs, err := a.orgModel.Query(ctx, schema.OrganizationQueryParam{RecordIds: datasets.Data.ToOrgIds()})
	if err != nil {
		return nil, err
	}
	orgs.Data.FillDataset(datasets.Data)

	return orgs.Data, nil
}

// GetAvailableDatasets 按类型查询可用知识库
// - orgId 为空：返回全部公共/共享知识库（PublicDatasets）
// - orgId 非空：返回该企业可访问的全部知识库（被分配的共享 PublicDatasets + 企业自建 OrgDatasets）
func (a *Dataset) GetAvailableDatasets(ctx context.Context, orgId string) (*schema.AvailableDatasets, error) {
	result := &schema.AvailableDatasets{
		PublicDatasets: make(schema.Datasets, 0),
		OrgDatasets:    make(schema.Datasets, 0),
	}

	if orgId == "" {
		// 新建企业场景：只返回全部公共/共享知识库
		shared, err := a.DatasetModel.Query(ctx, schema.DatasetQueryParam{Type: schema.DatasetTypePublic})
		if err != nil {
			return nil, err
		}
		result.PublicDatasets = shared.Data
		return result, nil
	}

	// 该企业被分配的共享知识库（通过组织级关系映射获得）
	rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{
		BizId: orgId,
		Type:  schema.DatasetTypePublic,
	})
	if err != nil {
		return nil, err
	}
	if ids := rel.Data.ToDatasetIds(); len(ids) > 0 {
		shared, err := a.DatasetModel.Query(ctx, schema.DatasetQueryParam{RecordIds: ids})
		if err != nil {
			return nil, err
		}
		result.PublicDatasets = shared.Data
	}

	// 该企业自建的企业知识库
	orgDs, err := a.DatasetModel.Query(ctx, schema.DatasetQueryParam{
		OrgId: orgId,
		Type:  schema.DatasetTypeOrg,
	})
	if err != nil {
		return nil, err
	}
	result.OrgDatasets = orgDs.Data
	return result, nil
}