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 }