diff --git a/internal/dao/connector.go b/internal/dao/connector.go index 0961a1240..a5a4491ea 100644 --- a/internal/dao/connector.go +++ b/internal/dao/connector.go @@ -18,6 +18,7 @@ package dao import ( "errors" + "ragflow/internal/common" "ragflow/internal/entity" "time" @@ -83,6 +84,85 @@ func (dao *ConnectorDAO) ListByDatasetID(datasetID string) ([]*ConnectorDatasetL return connectors, nil } +// DatasetConnectorLink is the connector relation payload accepted by dataset update. +type DatasetConnectorLink struct { + ID string + AutoParse string +} + +// LinkDatasetConnectors syncs connector2kb rows for a dataset. +func (dao *ConnectorDAO) LinkDatasetConnectors(kbID string, connectors []DatasetConnectorLink) error { + return DB.Transaction(func(tx *gorm.DB) error { + var existing []entity.Connector2Kb + if err := tx.Where("kb_id = ?", kbID).Find(&existing).Error; err != nil { + return err + } + + oldConnectorIDs := make(map[string]entity.Connector2Kb, len(existing)) + for _, row := range existing { + oldConnectorIDs[row.ConnectorID] = row + } + + nextConnectorIDs := make(map[string]struct{}, len(connectors)) + for _, connector := range connectors { + nextConnectorIDs[connector.ID] = struct{}{} + autoParse := connector.AutoParse + if autoParse == "" { + autoParse = "1" + } + + if _, ok := oldConnectorIDs[connector.ID]; ok { + if err := tx.Model(&entity.Connector2Kb{}). + Where("connector_id = ? AND kb_id = ?", connector.ID, kbID). + Update("auto_parse", autoParse).Error; err != nil { + return err + } + continue + } + + if err := tx.Create(&entity.Connector2Kb{ + ID: common.GenerateUUID(), + ConnectorID: connector.ID, + KbID: kbID, + AutoParse: autoParse, + }).Error; err != nil { + return err + } + + if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { + return err + } + + var fullConnector entity.Connector + if err := tx.Where("id = ?", connector.ID).First(&fullConnector).Error; err != nil { + return err + } + if connectorConfigBool(fullConnector.Config, "sync_deleted_files") { + if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { + return err + } + } + } + + for connectorID := range oldConnectorIDs { + if _, ok := nextConnectorIDs[connectorID]; ok { + continue + } + if err := tx.Where("kb_id = ? AND connector_id = ?", kbID, connectorID). + Delete(&entity.Connector2Kb{}).Error; err != nil { + return err + } + if err := tx.Model(&entity.SyncLogs{}). + Where("connector_id = ? AND kb_id = ? AND status IN ?", connectorID, kbID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). + Update("status", string(entity.TaskStatusCancel)).Error; err != nil { + return err + } + } + + return nil + }) +} + // GetByID get connector by ID func (dao *ConnectorDAO) GetByID(id string) (*entity.Connector, error) { var connector entity.Connector diff --git a/internal/dao/database.go b/internal/dao/database.go index 5142b8866..366b6b424 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -19,10 +19,13 @@ package dao import ( "fmt" "log" + "os" + "path/filepath" "ragflow/internal/common" "ragflow/internal/entity" "ragflow/internal/entity/models" "strings" + "sync" "time" "ragflow/internal/server" @@ -36,6 +39,7 @@ import ( var DB *gorm.DB var modelProviderManager *models.ProviderManager +var modelProviderManagerMu sync.Mutex // LLMFactoryConfig represents a single LLM factory configuration type LLMFactoryConfig struct { @@ -188,9 +192,44 @@ func GetDB() *gorm.DB { // GetModelProviderManager get database instance func GetModelProviderManager() *models.ProviderManager { + if modelProviderManager != nil { + return modelProviderManager + } + + modelProviderManagerMu.Lock() + defer modelProviderManagerMu.Unlock() + if modelProviderManager != nil { + return modelProviderManager + } + if existing := models.GetProviderManager(); existing != nil { + modelProviderManager = existing + return modelProviderManager + } + modelConfigDir, err := findModelConfigDir() + if err != nil { + log.Fatal("Failed to locate model providers:", err) + } + if err := models.InitProviderManager(modelConfigDir); err != nil { + log.Fatal("Failed to load model providers:", err) + } + modelProviderManager = models.GetProviderManager() return modelProviderManager } +func findModelConfigDir() (string, error) { + candidates := []string{ + "conf/models", + filepath.Join("..", "..", "conf", "models"), + filepath.Join("..", "..", "..", "conf", "models"), + } + for _, candidate := range candidates { + if info, err := os.Stat(candidate); err == nil && info.IsDir() { + return candidate, nil + } + } + return "", fmt.Errorf("conf/models not found") +} + // autoMigrateSafely runs AutoMigrate and ignores duplicate index errors // This handles cases where indexes already exist (e.g., created by Python backend) func autoMigrateSafely(db *gorm.DB, model interface{}) error { diff --git a/internal/handler/datasets.go b/internal/handler/datasets.go index 346aded37..d430741a8 100644 --- a/internal/handler/datasets.go +++ b/internal/handler/datasets.go @@ -171,6 +171,48 @@ func (h *DatasetsHandler) GetDataset(c *gin.Context) { }) } +// UpdateDataset Update a dataset. +func (h *DatasetsHandler) UpdateDataset(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + userID := strings.TrimSpace(user.ID) + if userID == "" { + jsonError(c, common.CodeDataError, "user id is required") + return + } + + datasetID := strings.TrimSpace(c.Param("dataset_id")) + if datasetID == "" { + jsonError(c, common.CodeBadRequest, "dataset id is required") + return + } + + var req service.UpdateDatasetRequest + if err := c.ShouldBindJSON(&req); err != nil { + jsonError(c, common.CodeDataError, err.Error()) + return + } + + result, code, err := h.datasetsService.UpdateDataset(datasetID, userID, req) + if err != nil { + jsonError(c, code, err.Error()) + return + } + if code != common.CodeSuccess { + jsonError(c, code, "dataset updated failed") + return + } + + c.JSON(http.StatusOK, gin.H{ + "code": common.CodeSuccess, + "data": result, + }) +} + func (h *DatasetsHandler) GetMetadataConfig(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { diff --git a/internal/router/router.go b/internal/router/router.go index ca0131a6c..29f6898ca 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -252,6 +252,7 @@ func (r *Router) Setup(engine *gin.Engine) { { datasets.GET("", r.datasetsHandler.ListDatasets) datasets.GET("/:dataset_id", r.datasetsHandler.GetDataset) + datasets.PUT("/:dataset_id", r.datasetsHandler.UpdateDataset) datasets.GET("/:dataset_id/graph", r.datasetsHandler.GetKnowledgeGraph) datasets.DELETE("/:dataset_id/tags", r.datasetsHandler.RemoveTags) datasets.DELETE("/:dataset_id/graph", r.datasetsHandler.DeleteKnowledgeGraph) @@ -324,32 +325,32 @@ func (r *Router) Setup(engine *gin.Engine) { commitFolders.GET("/:folder_id/changes", r.fileCommitHandler.GetUncommittedChanges) } - // /workspace/{workspace_id}/commits — alias for /folders/ (workspace_id == folder_id) - commitWorkspace := v1.Group("/workspace") - { - commitWorkspace.POST("/:folder_id/commits", r.fileCommitHandler.CreateCommit) - commitWorkspace.GET("/:folder_id/commits", r.fileCommitHandler.ListCommits) - commitWorkspace.GET("/:folder_id/commits/diff", r.fileCommitHandler.DiffCommits) - commitWorkspace.GET("/:folder_id/commits/:commit_id", r.fileCommitHandler.GetCommit) - commitWorkspace.GET("/:folder_id/commits/:commit_id/files", r.fileCommitHandler.ListCommitFiles) - commitWorkspace.GET("/:folder_id/commits/:commit_id/tree", r.fileCommitHandler.GetCommitTree) - commitWorkspace.GET("/:folder_id/commits/:commit_id/files/:file_id/content", r.fileCommitHandler.GetCommitFileContent) - commitWorkspace.GET("/:folder_id/changes", r.fileCommitHandler.GetUncommittedChanges) - } + // /workspace/{workspace_id}/commits — alias for /folders/ (workspace_id == folder_id) + commitWorkspace := v1.Group("/workspace") + { + commitWorkspace.POST("/:folder_id/commits", r.fileCommitHandler.CreateCommit) + commitWorkspace.GET("/:folder_id/commits", r.fileCommitHandler.ListCommits) + commitWorkspace.GET("/:folder_id/commits/diff", r.fileCommitHandler.DiffCommits) + commitWorkspace.GET("/:folder_id/commits/:commit_id", r.fileCommitHandler.GetCommit) + commitWorkspace.GET("/:folder_id/commits/:commit_id/files", r.fileCommitHandler.ListCommitFiles) + commitWorkspace.GET("/:folder_id/commits/:commit_id/tree", r.fileCommitHandler.GetCommitTree) + commitWorkspace.GET("/:folder_id/commits/:commit_id/files/:file_id/content", r.fileCommitHandler.GetCommitFileContent) + commitWorkspace.GET("/:folder_id/changes", r.fileCommitHandler.GetUncommittedChanges) + } - // /datasets/{dataset_id}/commits — resolve dataset_id → folder_id via middleware - commitDatasets := v1.Group("/datasets/:dataset_id") - commitDatasets.Use(handler.CommitFolderResolver(r.fileCommitHandler, "datasets", "dataset_id")) - { - commitDatasets.POST("/commits", r.fileCommitHandler.CreateCommit) - commitDatasets.GET("/commits", r.fileCommitHandler.ListCommits) - commitDatasets.GET("/commits/diff", r.fileCommitHandler.DiffCommits) - commitDatasets.GET("/commits/:commit_id", r.fileCommitHandler.GetCommit) - commitDatasets.GET("/commits/:commit_id/files", r.fileCommitHandler.ListCommitFiles) - commitDatasets.GET("/commits/:commit_id/tree", r.fileCommitHandler.GetCommitTree) - commitDatasets.GET("/commits/:commit_id/files/:file_id/content", r.fileCommitHandler.GetCommitFileContent) - commitDatasets.GET("/changes", r.fileCommitHandler.GetUncommittedChanges) - } + // /datasets/{dataset_id}/commits — resolve dataset_id → folder_id via middleware + commitDatasets := v1.Group("/datasets/:dataset_id") + commitDatasets.Use(handler.CommitFolderResolver(r.fileCommitHandler, "datasets", "dataset_id")) + { + commitDatasets.POST("/commits", r.fileCommitHandler.CreateCommit) + commitDatasets.GET("/commits", r.fileCommitHandler.ListCommits) + commitDatasets.GET("/commits/diff", r.fileCommitHandler.DiffCommits) + commitDatasets.GET("/commits/:commit_id", r.fileCommitHandler.GetCommit) + commitDatasets.GET("/commits/:commit_id/files", r.fileCommitHandler.ListCommitFiles) + commitDatasets.GET("/commits/:commit_id/tree", r.fileCommitHandler.GetCommitTree) + commitDatasets.GET("/commits/:commit_id/files/:file_id/content", r.fileCommitHandler.GetCommitFileContent) + commitDatasets.GET("/changes", r.fileCommitHandler.GetUncommittedChanges) + } // Author routes authors := v1.Group("/authors") @@ -372,6 +373,7 @@ func (r *Router) Setup(engine *gin.Engine) { message := v1.Group("/messages") { message.DELETE("/:memory_message", r.memoryHandler.ForgetMessage) + message.PUT("/:memory_message", r.memoryHandler.UpdateMessage) } // Skill search routes diff --git a/internal/router/router_test.go b/internal/router/router_test.go index 5e6c420af..414ac4283 100644 --- a/internal/router/router_test.go +++ b/internal/router/router_test.go @@ -6,6 +6,8 @@ import ( "testing" "github.com/gin-gonic/gin" + + "ragflow/internal/handler" ) func TestConnectorRoutesDoNotConflictWithOAuthCallbacks(t *testing.T) { @@ -69,3 +71,25 @@ func TestConnectorRoutesDoNotConflictWithOAuthCallbacks(t *testing.T) { }) } } + +func TestRouterSetupRegistersUpdateDatasetRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + + engine := gin.New() + r := &Router{ + authHandler: handler.NewAuthHandler(), + datasetsHandler: handler.NewDatasetsHandler(nil, nil), + } + r.Setup(engine) + + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/api/v1/datasets/dataset-1", nil) + engine.ServeHTTP(resp, req) + + if resp.Code == http.StatusNotFound { + t.Fatalf("PUT /api/v1/datasets/:dataset_id returned 404; UpdateDataset route is not registered") + } + if resp.Code != http.StatusUnauthorized { + t.Fatalf("status=%d body=%s; want auth middleware to handle registered UpdateDataset route", resp.Code, resp.Body.String()) + } +} diff --git a/internal/service/dataset.go b/internal/service/dataset.go index 65f80bc52..d7c6c870f 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -1022,6 +1022,366 @@ func (s *DatasetService) GetDataset(datasetID, userID string) (map[string]interf return data, common.CodeSuccess, nil } +type DatasetConnectorRequest struct { + ID string `json:"id"` + AutoParse string `json:"auto_parse,omitempty"` +} + +type UpdateDatasetRequest struct { + Name *string `json:"name,omitempty"` + Avatar *string `json:"avatar,omitempty"` + Description *string `json:"description,omitempty"` + Language *string `json:"language,omitempty"` + Connectors *[]DatasetConnectorRequest `json:"connectors,omitempty"` + EmbdID *string `json:"embd_id,omitempty"` + EmbeddingModel *string `json:"embedding_model,omitempty"` + Permission *string `json:"permission,omitempty"` + ParserID *string `json:"parser_id,omitempty"` + ChunkMethod *string `json:"chunk_method,omitempty"` + Pagerank *int64 `json:"pagerank,omitempty"` + ParserConfig map[string]interface{} `json:"parser_config,omitempty"` + PipelineID *string `json:"pipeline_id,omitempty"` + AutoMetadataConfig *AutoMetadataConfig `json:"auto_metadata_config,omitempty"` + Ext map[string]interface{} `json:"ext,omitempty"` +} + +// UpdateDataset Update a dataset +func (s *DatasetService) UpdateDataset(datasetID, tenantID string, req UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) { + kb, err := s.kbDAO.GetByID(datasetID) + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, errors.New("Dataset not found") + } + return nil, common.CodeServerError, errors.New("Database operation failed") + } + + if kb == nil || kb.TenantID != tenantID { + return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID) + } + + connectorsProvided := req.Connectors != nil + connectors := make([]DatasetConnectorRequest, 0) + if req.Connectors != nil { + connectors = *req.Connectors + } + + updates := make(map[string]interface{}) + extUpdates := normalizeDatasetUpdateExt(req.Ext) + + if req.Name != nil { + name := strings.TrimSpace(*req.Name) + if name == "" { + return nil, common.CodeDataError, errors.New("String should have at least 1 character") + } + if len(name) > 128 { + return nil, common.CodeDataError, errors.New("String should have at most 128 characters") + } + updates["name"] = name + } + if req.Avatar != nil { + if len(*req.Avatar) > 65535 { + return nil, common.CodeDataError, errors.New("String should have at most 65535 characters") + } + if err := validateDatasetAvatar(*req.Avatar); err != nil { + return nil, common.CodeDataError, err + } + updates["avatar"] = *req.Avatar + } + if req.Description != nil { + if len(*req.Description) > 65535 { + return nil, common.CodeDataError, errors.New("String should have at most 65535 characters") + } + updates["description"] = *req.Description + } + if req.Language != nil { + language := strings.TrimSpace(*req.Language) + if len(language) > 32 { + return nil, common.CodeDataError, errors.New("String should have at most 32 characters") + } + updates["language"] = language + } + if req.Permission != nil { + permission := strings.TrimSpace(*req.Permission) + if permission != "me" && permission != "team" { + return nil, common.CodeDataError, errors.New("Input should be 'me' or 'team'") + } + updates["permission"] = permission + } + if req.PipelineID != nil { + pipelineID, err := normalizeDatasetPipelineID(*req.PipelineID) + if err != nil { + return nil, common.CodeDataError, err + } + if pipelineID != nil { + updates["pipeline_id"] = *pipelineID + } + } + + for key, value := range extUpdates { + if _, exists := updates[key]; !exists { + updates[key] = value + } + } + + parserID, parserIDProvided, err := datasetUpdateParserID(req) + if err != nil { + return nil, common.CodeDataError, err + } + if !parserIDProvided { + if extParserID, ok := updates["parser_id"]; ok { + parserIDValue, ok := extParserID.(string) + if !ok { + return nil, common.CodeDataError, errors.New(datasetChunkMethodErrorMessage) + } + parserID = strings.TrimSpace(parserIDValue) + if err := validateDatasetChunkMethod(parserID); err != nil { + return nil, common.CodeDataError, err + } + parserIDProvided = true + } + } + if parserIDProvided { + updates["parser_id"] = parserID + } + + embdID, embdIDProvided, err := datasetUpdateEmbeddingID(req) + if err != nil { + return nil, common.CodeDataError, err + } + if !embdIDProvided { + if extEmbdID, ok := updates["embd_id"]; ok { + embdIDValue, ok := extEmbdID.(string) + if !ok { + return nil, common.CodeDataError, errors.New("Embedding model identifier must follow @ format") + } + embdID = strings.TrimSpace(embdIDValue) + if embdID != "" { + if err := validateDatasetEmbeddingModel(embdID); err != nil { + return nil, common.CodeDataError, err + } + } + embdIDProvided = true + } + } + if embdIDProvided { + if embdID == "" { + embdID = kb.EmbdID + } + ok, message := s.verifyEmbeddingAvailability(embdID, tenantID) + if !ok { + return nil, common.CodeDataError, errors.New(message) + } + updates["embd_id"] = embdID + } + + if req.AutoMetadataConfig != nil { + req.ParserConfig = applyAutoMetadataConfig(req.ParserConfig, req.AutoMetadataConfig) + } + if req.ParserConfig != nil { + if err := validateDatasetParserConfigSize(req.ParserConfig); err != nil { + return nil, common.CodeDataError, err + } + if len(req.ParserConfig) > 0 { + parserConfig := normalizeDatasetUpdateParserConfig(req.ParserConfig) + updates["parser_config"] = entity.JSONMap(common.DeepMergeMaps(kb.ParserConfig, parserConfig)) + } + } + + if req.Pagerank != nil && *req.Pagerank != kb.Pagerank { + if *req.Pagerank < 0 || *req.Pagerank > 100 { + return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") + } + if s.engineType == server.EngineInfinity { + return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch") + } + indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) + if *req.Pagerank > 0 { + err = s.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"kb_id": kb.ID}, map[string]interface{}{common.PAGERANK_FLD: *req.Pagerank}, indexName, kb.ID) + } else { + err = s.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, indexName, kb.ID) + } + if err != nil { + return nil, common.CodeServerError, err + } + updates["pagerank"] = *req.Pagerank + } + + if parserIDProvided && parserID != kb.ParserID { + if _, ok := updates["parser_config"]; !ok { + updates["parser_config"] = entity.JSONMap(common.GetParserConfig(parserID, nil)) + } + } + if kb.PipelineID != nil && parserIDProvided { + if _, ok := updates["pipeline_id"]; !ok { + updates["pipeline_id"] = "" + } + } + + if nameValue, ok := updates["name"].(string); ok && strings.ToLower(nameValue) != strings.ToLower(kb.Name) { + existing, lookupErr := s.kbDAO.GetByName(nameValue, tenantID) + if lookupErr != nil && !dao.IsNotFoundErr(lookupErr) { + return nil, common.CodeServerError, errors.New("Database operation failed") + } + if existing != nil { + return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", nameValue) + } + } + + if len(updates) == 0 && !connectorsProvided { + return nil, common.CodeDataError, errors.New("No properties were modified") + } + + if len(updates) > 0 { + if err = s.kbDAO.UpdateByID(kb.ID, updates); err != nil { + return nil, common.CodeServerError, errors.New("Update dataset error.(Database error)") + } + } + + if connectorsProvided { + connectorLinks := make([]dao.DatasetConnectorLink, 0, len(connectors)) + for _, connector := range connectors { + connectorID := strings.TrimSpace(connector.ID) + if connectorID == "" { + return nil, common.CodeDataError, errors.New("connector id is required") + } + connectorLinks = append(connectorLinks, dao.DatasetConnectorLink{ + ID: connectorID, + AutoParse: connector.AutoParse, + }) + } + if err = s.connectorDAO.LinkDatasetConnectors(kb.ID, connectorLinks); err != nil { + return nil, common.CodeServerError, errors.New("Database operation failed") + } + } + + updatedKB, err := s.kbDAO.GetByID(kb.ID) + if err != nil { + return nil, common.CodeDataError, errors.New("Dataset updated failed") + } + + data := datasetToMap(updatedKB) + linkedConnectors, err := s.connectorDAO.ListByDatasetID(kb.ID) + if err != nil { + return nil, common.CodeServerError, errors.New("Database operation failed") + } + data["connectors"] = linkedConnectors + return data, common.CodeSuccess, nil +} + +func datasetUpdateParserID(req UpdateDatasetRequest) (string, bool, error) { + parserID := "" + provided := false + if req.ParserID != nil { + parserID = strings.TrimSpace(*req.ParserID) + provided = true + } + if req.ChunkMethod != nil { + parserID = strings.TrimSpace(*req.ChunkMethod) + provided = true + } + if !provided { + return "", false, nil + } + if err := validateDatasetChunkMethod(parserID); err != nil { + return "", true, err + } + return parserID, true, nil +} + +func datasetUpdateEmbeddingID(req UpdateDatasetRequest) (string, bool, error) { + embdID := "" + provided := false + if req.EmbdID != nil { + embdID = strings.TrimSpace(*req.EmbdID) + provided = true + } + if req.EmbeddingModel != nil { + embdID = strings.TrimSpace(*req.EmbeddingModel) + provided = true + } + if !provided { + return "", false, nil + } + if embdID != "" { + if err := validateDatasetEmbeddingModel(embdID); err != nil { + return "", true, err + } + } + return embdID, true, nil +} + +func normalizeDatasetUpdateExt(ext map[string]interface{}) map[string]interface{} { + if ext == nil { + return nil + } + + updates := make(map[string]interface{}, len(ext)) + for key, value := range ext { + switch key { + case "embedding_model": + updates["embd_id"] = value + case "chunk_method": + updates["parser_id"] = value + case "connectors", "auto_metadata_config", "ext", "parse_type": + continue + default: + updates[key] = value + } + } + return updates +} + +func normalizeDatasetUpdateParserConfig(parserConfig map[string]interface{}) map[string]interface{} { + normalized := common.DeepMergeMaps(nil, parserConfig) + parentChild, _ := normalized["parent_child"].(map[string]interface{}) + if parentChild == nil { + parentChild = map[string]interface{}{} + } + + if datasetBoolValue(parentChild["use_parent_child"]) { + childrenDelimiter, ok := parentChild["children_delimiter"] + if !ok { + childrenDelimiter = "\n" + } + normalized["children_delimiter"] = childrenDelimiter + enableChildren, ok := parentChild["use_parent_child"] + if !ok { + enableChildren = true + } + normalized["enable_children"] = enableChildren + } else { + normalized["children_delimiter"] = "" + normalized["enable_children"] = false + normalized["parent_child"] = map[string]interface{}{} + } + + if extFields, ok := normalized["ext"].(map[string]interface{}); ok { + delete(normalized, "ext") + for key, value := range extFields { + normalized[key] = value + } + } + + return normalized +} + +func datasetBoolValue(value interface{}) bool { + switch typedValue := value.(type) { + case bool: + return typedValue + case string: + return typedValue == "1" || strings.EqualFold(typedValue, "true") + case int: + return typedValue != 0 + case int64: + return typedValue != 0 + case float64: + return typedValue != 0 + default: + return false + } +} + // 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) @@ -1424,32 +1784,11 @@ func normalizeDatasetUUID1(id string) (string, error) { } func (s *DatasetService) verifyEmbeddingAvailability(embdID string, tenantID string) (bool, string) { - modelName, _, provider, err := parseModelName(embdID) + _, _, _, _, err := NewModelProviderService().GetModelConfigFromProviderInstance(tenantID, entity.ModelTypeEmbedding, embdID) if err != nil { - return false, "Embedding model identifier must follow @ format" + return false, err.Error() } - - if provider == "Builtin" { - return true, "" - } - - tenantLLMs, err := s.tenantLLMDAO.ListValidByTenant(tenantID) - if err != nil { - return false, "Database operation failed" - } - - for _, tenantLLM := range tenantLLMs { - if tenantLLM == nil || tenantLLM.LLMName == nil || tenantLLM.ModelType == nil { - continue - } - if *tenantLLM.LLMName == modelName && - tenantLLM.LLMFactory == provider && - *tenantLLM.ModelType == string(entity.ModelTypeEmbedding) { - return true, "" - } - } - - return false, fmt.Sprintf("Unauthorized model: <%s>", embdID) + return true, "" } func applyAutoMetadataConfig(parserConfig map[string]interface{}, config *AutoMetadataConfig) map[string]interface{} { diff --git a/internal/service/dataset_update_test.go b/internal/service/dataset_update_test.go new file mode 100644 index 000000000..0377da1b3 --- /dev/null +++ b/internal/service/dataset_update_test.go @@ -0,0 +1,400 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "strings" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + + "gorm.io/gorm" +) + +func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + name := " Renamed Dataset " + description := "updated description" + language := "English" + permission := string(entity.TenantPermissionTeam) + chunkMethod := string(entity.ParserTypeBook) + embeddingModel := "BAAI/bge-large-zh-v1.5@Builtin" + pipelineID := "ABCDEF0123456789ABCDEF0123456789" + + result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{ + Name: &name, + Description: &description, + Language: &language, + Permission: &permission, + ChunkMethod: &chunkMethod, + EmbeddingModel: &embeddingModel, + PipelineID: &pipelineID, + ParserConfig: map[string]interface{}{ + "parent_child": map[string]interface{}{ + "use_parent_child": true, + }, + "ext": map[string]interface{}{ + "delimiter": "\n\n", + }, + }, + }) + if err != nil { + t.Fatalf("UpdateDataset failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if result["name"] != "Renamed Dataset" { + t.Fatalf("expected trimmed name in response, got %#v", result["name"]) + } + if result["chunk_method"] != chunkMethod { + t.Fatalf("expected chunk method %q, got %#v", chunkMethod, result["chunk_method"]) + } + if result["embedding_model"] != embeddingModel { + t.Fatalf("expected embedding model %q, got %#v", embeddingModel, result["embedding_model"]) + } + if result["pipeline_id"] != strings.ToLower(pipelineID) { + t.Fatalf("expected normalized pipeline id, got %#v", result["pipeline_id"]) + } + if connectors, ok := result["connectors"].([]*dao.ConnectorDatasetListItem); !ok || len(connectors) != 0 { + t.Fatalf("expected empty connector list, got %#v", result["connectors"]) + } + + persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get updated kb: %v", err) + } + if persisted.Name != "Renamed Dataset" { + t.Fatalf("expected persisted trimmed name, got %q", persisted.Name) + } + if persisted.Description == nil || *persisted.Description != description { + t.Fatalf("expected persisted description, got %#v", persisted.Description) + } + if persisted.Permission != permission { + t.Fatalf("expected persisted permission %q, got %q", permission, persisted.Permission) + } + if persisted.ParserID != chunkMethod { + t.Fatalf("expected parser id %q, got %q", chunkMethod, persisted.ParserID) + } + if persisted.EmbdID != embeddingModel { + t.Fatalf("expected embd id %q, got %q", embeddingModel, persisted.EmbdID) + } + if persisted.PipelineID == nil || *persisted.PipelineID != strings.ToLower(pipelineID) { + t.Fatalf("expected normalized pipeline id persisted, got %#v", persisted.PipelineID) + } + if persisted.ParserConfig["chunk_token_num"] != float64(128) { + t.Fatalf("expected existing parser_config value preserved, got %#v", persisted.ParserConfig) + } + if persisted.ParserConfig["enable_children"] != true { + t.Fatalf("expected enable_children normalized, got %#v", persisted.ParserConfig) + } + if persisted.ParserConfig["children_delimiter"] != "\n" { + t.Fatalf("expected default children delimiter, got %#v", persisted.ParserConfig["children_delimiter"]) + } + if persisted.ParserConfig["delimiter"] != "\n\n" { + t.Fatalf("expected ext parser config fields flattened, got %#v", persisted.ParserConfig) + } +} + +func TestDatasetServiceUpdateDatasetRejectsMissingDataset(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + + name := "Renamed" + _, code, err := testDatasetUpdateService(t).UpdateDataset("missing-kb", "tenant-1", UpdateDatasetRequest{Name: &name}) + if err == nil { + t.Fatal("expected missing dataset error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if err.Error() != "Dataset not found" { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceUpdateDatasetRejectsNonOwner(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + name := "Renamed" + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-2", UpdateDatasetRequest{Name: &name}) + if err == nil { + t.Fatal("expected permission error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if !strings.Contains(err.Error(), "lacks permission") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceUpdateDatasetValidatesName(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + name := " " + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{Name: &name}) + if err == nil { + t.Fatal("expected name validation error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if err.Error() != "String should have at least 1 character" { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceUpdateDatasetRejectsDuplicateName(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Existing") + + name := "Existing" + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{Name: &name}) + if err == nil { + t.Fatal("expected duplicate name error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if !strings.Contains(err.Error(), "already exists") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceUpdateDatasetRejectsNoPropertiesModified(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{}) + if err == nil { + t.Fatal("expected no-op update error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if err.Error() != "No properties were modified" { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestDatasetServiceUpdateDatasetLinksConnectors(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + insertDatasetUpdateConnector(t, "connector-1", "tenant-1") + + autoParse := "0" + result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{ + Connectors: &[]DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}}, + }) + if err != nil { + t.Fatalf("UpdateDataset failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + connectors, ok := result["connectors"].([]*dao.ConnectorDatasetListItem) + if !ok { + t.Fatalf("expected connector list, got %#v", result["connectors"]) + } + if len(connectors) != 1 { + t.Fatalf("expected one connector, got %d", len(connectors)) + } + if connectors[0].ID != "connector-1" || connectors[0].AutoParse != autoParse { + t.Fatalf("unexpected connector: %#v", connectors[0]) + } + + var link entity.Connector2Kb + if err := dao.DB.Where("kb_id = ? AND connector_id = ?", "kb-1", "connector-1").First(&link).Error; err != nil { + t.Fatalf("expected connector link: %v", err) + } + if link.AutoParse != autoParse { + t.Fatalf("expected auto_parse %q, got %q", autoParse, link.AutoParse) + } +} + +func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + insertDatasetUpdateModelProvider(t, "provider-1", "tenant-1", "ZHIPU-AI") + insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test") + insertDatasetUpdateTenantModel(t, "model-1", "provider-1", "instance-1", "embedding-2", string(entity.ModelTypeEmbedding)) + + embeddingModel := "embedding-2@test@ZHIPU-AI" + result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{ + EmbeddingModel: &embeddingModel, + }) + if err != nil { + t.Fatalf("UpdateDataset failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if result["embedding_model"] != embeddingModel { + t.Fatalf("expected embedding model %q, got %#v", embeddingModel, result["embedding_model"]) + } + + persisted, err := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if err != nil { + t.Fatalf("get updated kb: %v", err) + } + if persisted.EmbdID != embeddingModel { + t.Fatalf("expected persisted embedding model %q, got %q", embeddingModel, persisted.EmbdID) + } +} + +func TestDatasetServiceUpdateDatasetRejectsEmptyConnectorID(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + connectors := []DatasetConnectorRequest{{ID: " "}} + _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{ + Connectors: &connectors, + }) + if err == nil { + t.Fatal("expected connector validation error") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if err.Error() != "connector id is required" { + t.Fatalf("unexpected error: %v", err) + } +} + +func setupDatasetUpdateTestDB(t *testing.T) *gorm.DB { + t.Helper() + + db := setupServiceTestDB(t) + if err := db.AutoMigrate( + &entity.Connector{}, + &entity.Connector2Kb{}, + &entity.SyncLogs{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, + ); err != nil { + t.Fatalf("failed to migrate dataset update tables: %v", err) + } + return db +} + +func testDatasetUpdateService(t *testing.T) *DatasetService { + t.Helper() + + return &DatasetService{ + kbDAO: dao.NewKnowledgebaseDAO(), + connectorDAO: dao.NewConnectorDAO(), + } +} + +func insertDatasetUpdateKB(t *testing.T, id, tenantID, name string) { + t.Helper() + + kb := &entity.Knowledgebase{ + ID: id, + TenantID: tenantID, + Name: name, + EmbdID: "BAAI/bge-large-zh-v1.5@Builtin", + CreatedBy: tenantID, + Permission: string(entity.TenantPermissionMe), + ParserID: string(entity.ParserTypeNaive), + ParserConfig: entity.JSONMap{"chunk_token_num": float64(128)}, + Status: sptr(string(entity.StatusValid)), + } + if err := dao.DB.Create(kb).Error; err != nil { + t.Fatalf("insert test kb: %v", err) + } +} + +func insertDatasetUpdateConnector(t *testing.T, id, tenantID string) { + t.Helper() + + connector := &entity.Connector{ + ID: id, + TenantID: tenantID, + Name: "Test Connector", + Source: "google_drive", + InputType: "oauth", + Config: entity.JSONMap{"sync_deleted_files": false}, + Status: string(entity.TaskStatusDone), + } + if err := dao.DB.Create(connector).Error; err != nil { + t.Fatalf("insert test connector: %v", err) + } +} + +func insertDatasetUpdateModelProvider(t *testing.T, id, tenantID, providerName string) { + t.Helper() + + provider := &entity.TenantModelProvider{ + ID: id, + TenantID: tenantID, + ProviderName: providerName, + } + if err := dao.DB.Create(provider).Error; err != nil { + t.Fatalf("insert test model provider: %v", err) + } +} + +func insertDatasetUpdateModelInstance(t *testing.T, id, providerID, instanceName string) { + t.Helper() + + instance := &entity.TenantModelInstance{ + ID: id, + ProviderID: providerID, + InstanceName: instanceName, + APIKey: "test-api-key", + Status: "active", + Extra: "{}", + } + if err := dao.DB.Create(instance).Error; err != nil { + t.Fatalf("insert test model instance: %v", err) + } +} + +func insertDatasetUpdateTenantModel(t *testing.T, id, providerID, instanceID, modelName, modelType string) { + t.Helper() + + model := &entity.TenantModel{ + ID: id, + ProviderID: providerID, + InstanceID: instanceID, + ModelName: modelName, + ModelType: modelType, + Status: "active", + Extra: "{}", + } + if err := dao.DB.Create(model).Error; err != nil { + t.Fatalf("insert test tenant model: %v", err) + } +}