diff --git a/internal/handler/datasets.go b/internal/handler/datasets.go index 19ccc194b..ed75b42f8 100644 --- a/internal/handler/datasets.go +++ b/internal/handler/datasets.go @@ -34,8 +34,8 @@ import ( // DatasetsHandler handles the RESTful dataset endpoints. type DatasetsHandler struct { - datasetsService *service.DatasetService - metadataService *service.MetadataService + datasetsService *service.DatasetService + metadataService *service.MetadataService } type listDatasetsExt struct { @@ -47,8 +47,8 @@ type listDatasetsExt struct { // NewDatasetsHandler creates a new datasets handler. func NewDatasetsHandler(datasetsService *service.DatasetService, metadataService *service.MetadataService) *DatasetsHandler { return &DatasetsHandler{ - datasetsService: datasetsService, - metadataService: metadataService, + datasetsService: datasetsService, + metadataService: metadataService, } } @@ -171,6 +171,52 @@ func (h *DatasetsHandler) GetDataset(c *gin.Context) { }) } +func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + datasetID := c.Param("dataset_id") + result, code, err := h.datasetsService.GetMetadataConfig(datasetID, user.ID) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": result, + }) +} + +func (h *DatasetsHandler) UpdateMetadataConfig(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + datasetID := c.Param("dataset_id") + var req service.MetadataConfigRequest + if err := c.ShouldBindJSON(&req); err != nil { + jsonError(c, common.CodeDataError, err.Error()) + return + } + + result, code, err := h.datasetsService.UpdateMetadataConfig(datasetID, user.ID, &req) + if err != nil { + jsonError(c, code, err.Error()) + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": result, + }) +} + // GetIngestionSummary handles GET /api/v1/datasets/:dataset_id/ingestions/summary. func (h *DatasetsHandler) GetIngestionSummary(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) diff --git a/internal/router/router.go b/internal/router/router.go index b0289ae6f..498555f0f 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -214,6 +214,10 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.GET("/:dataset_id/ingestions", r.datasetsHandler.ListIngestionLogs) datasets.GET("/:dataset_id/ingestions/:log_id", r.datasetsHandler.GetIngestionLog) + // Metadata Config + datasets.GET("/:dataset_id/metadata/config", r.datasetsHandler.GetMetadataConfig) + datasets.PUT("/:dataset_id/metadata/config", r.datasetsHandler.UpdateMetadataConfig) + // Dataset documents datasets.GET("/:dataset_id/documents", r.documentHandler.ListDocuments) diff --git a/internal/service/dataset.go b/internal/service/dataset.go index 99fb5b3e6..51b0f4bfe 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -55,6 +55,12 @@ var ( "create_time": {}, "update_time": {}, } + datasetAllowedMetadataTypes = map[string]struct{}{ + "string": {}, + "list": {}, + "time": {}, + "number": {}, + } datasetChunkMethodErrorMessage = "Input should be 'naive', 'book', 'email', 'laws', 'manual', 'one', 'paper', 'picture', 'presentation', 'qa', 'resume', 'table' or 'tag'" ) @@ -95,6 +101,20 @@ type AutoMetadataConfig struct { Fields []AutoMetadataField `json:"fields,omitempty"` } +// MetadataConfigField mirrors one field in the dataset metadata config API. +type MetadataConfigField struct { + Key string `json:"key"` + Type string `json:"type"` + Description *string `json:"description"` + Enum []string `json:"enum"` +} + +// MetadataConfigRequest mirrors PUT /datasets/:dataset_id/metadata/config. +type MetadataConfigRequest struct { + Metadata []MetadataConfigField `json:"metadata"` + BuiltInMetadata []MetadataConfigField `json:"built_in_metadata"` +} + // CreateDatasetRequest represents the request for creating a dataset. type CreateDatasetRequest struct { Name string `json:"name" binding:"required"` @@ -562,6 +582,71 @@ func (s *DatasetService) GetDataset(datasetID, userID string) (map[string]interf return data, common.CodeSuccess, nil } +// GetMetadataConfig gets the auto-metadata configuration for a dataset. +func (s *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) { + kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenantID) + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + return nil, common.CodeServerError, errors.New("Database operation failed") + } + if kb == nil { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + + return map[string]interface{}{ + "metadata": parserConfigValueOrEmptyList(kb.ParserConfig, "metadata"), + "built_in_metadata": parserConfigValueOrEmptyList(kb.ParserConfig, "built_in_metadata"), + }, common.CodeSuccess, nil +} + +// UpdateMetadataConfig updates the auto-metadata configuration for a dataset. +func (s *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) { + datasetID = strings.TrimSpace(datasetID) + tenantID = strings.TrimSpace(tenantID) + + kb, err := s.kbDAO.GetByIDAndTenantID(datasetID, tenantID) + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + return nil, common.CodeServerError, errors.New("Database operation failed") + } + if kb == nil { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + + if req == nil { + req = &MetadataConfigRequest{} + } + + metadata, err := normalizeMetadataConfigFields(req.Metadata, "metadata") + if err != nil { + return nil, common.CodeDataError, err + } + builtInMetadata, err := normalizeMetadataConfigFields(req.BuiltInMetadata, "built_in_metadata") + if err != nil { + return nil, common.CodeDataError, err + } + + parserConfig := kb.ParserConfig + if parserConfig == nil { + parserConfig = entity.JSONMap{} + } + parserConfig["metadata"] = metadata + parserConfig["built_in_metadata"] = builtInMetadata + + if err := s.kbDAO.UpdateByID(kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { + return nil, common.CodeServerError, errors.New("Update auto-metadata error.(Database error)") + } + + return map[string]interface{}{ + "metadata": metadata, + "built_in_metadata": builtInMetadata, + }, common.CodeSuccess, nil +} + // Accessible checks if a knowledge base is accessible by a user func (s *DatasetService) Accessible(kbID, userID string) bool { return s.kbDAO.Accessible(kbID, userID) @@ -955,6 +1040,50 @@ func applyAutoMetadataConfig(parserConfig map[string]interface{}, config *AutoMe return parserConfig } +func parserConfigValueOrEmptyList(parserConfig map[string]interface{}, key string) interface{} { + if parserConfig == nil { + return []interface{}{} + } + + value, ok := parserConfig[key] + if !ok || value == nil { + return []interface{}{} + } + + return value +} + +func normalizeMetadataConfigFields(fields []MetadataConfigField, fieldName string) ([]map[string]interface{}, error) { + normalizedFields := make([]map[string]interface{}, 0, len(fields)) + for i, field := range fields { + key := strings.TrimSpace(field.Key) + if key == "" { + return nil, fmt.Errorf("%s[%d].key is required", fieldName, i) + } + if len(key) > 255 { + return nil, fmt.Errorf("%s[%d].key should have at most 255 characters", fieldName, i) + } + + fieldType := strings.TrimSpace(field.Type) + if _, ok := datasetAllowedMetadataTypes[fieldType]; !ok { + return nil, fmt.Errorf("%s[%d].type should be one of 'string', 'list', 'time' or 'number'", fieldName, i) + } + + if field.Description != nil && len(*field.Description) > 65535 { + return nil, fmt.Errorf("%s[%d].description should have at most 65535 characters", fieldName, i) + } + + normalizedFields = append(normalizedFields, map[string]interface{}{ + "key": key, + "type": fieldType, + "description": field.Description, + "enum": field.Enum, + }) + } + + return normalizedFields, nil +} + func datasetListItemToMap(kb *entity.KnowledgebaseListItem) map[string]interface{} { item := map[string]interface{}{ "id": kb.ID,