diff --git a/internal/dao/tenant_model.go b/internal/dao/tenant_model.go index fd69c3ca4..7df8f809f 100644 --- a/internal/dao/tenant_model.go +++ b/internal/dao/tenant_model.go @@ -18,6 +18,8 @@ package dao import ( "ragflow/internal/entity" + + "gorm.io/gorm" ) // TenantModelDAO tenant model data access object @@ -32,6 +34,21 @@ func (dao *TenantModelDAO) Create(instance *entity.TenantModel) error { return DB.Create(instance).Error } +func (dao *TenantModelDAO) CreateBatch(models []*entity.TenantModel) error { + if len(models) == 0 { + return nil + } + + return DB.Transaction(func(tx *gorm.DB) error { + for _, model := range models { + if err := tx.Create(model).Error; err != nil { + return err + } + } + return nil + }) +} + func (dao *TenantModelDAO) DeleteByModelID(modelID string) (int64, error) { result := DB.Unscoped().Where("id = ?", modelID).Delete(&entity.TenantModel{}) return result.RowsAffected, result.Error diff --git a/internal/entity/models/moonshot.go b/internal/entity/models/moonshot.go index bea74095c..67052ada6 100644 --- a/internal/entity/models/moonshot.go +++ b/internal/entity/models/moonshot.go @@ -184,10 +184,7 @@ func (m *MoonshotModel) ChatWithMessages(modelName string, messages []Message, a var reasonContent string if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") - } + reasonContent, _ = messageMap["reasoning_content"].(string) // if first char of reasonContent is \n remove the \n if reasonContent != "" && reasonContent[0] == '\n' { reasonContent = reasonContent[1:] diff --git a/internal/handler/providers.go b/internal/handler/providers.go index ef5fc0066..e9aab9343 100644 --- a/internal/handler/providers.go +++ b/internal/handler/providers.go @@ -772,7 +772,7 @@ func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) { }) } -func prepareAddCustomModelRequest(req *service.AddCustomModelRequest, providerName, instanceName string) error { +func prepareProviderInstance(providerName, instanceName, reqProviderName, reqInstanceName string) error { if providerName == "" { return errors.New("Provider name is required") } @@ -781,20 +781,34 @@ func prepareAddCustomModelRequest(req *service.AddCustomModelRequest, providerNa return errors.New("Instance name is required") } - if req.ProviderName != "" && !strings.EqualFold(req.ProviderName, providerName) { + if reqProviderName != "" && !strings.EqualFold(reqProviderName, providerName) { return errors.New("Provider name does not match path") } - if req.InstanceName != "" && !strings.EqualFold(req.InstanceName, instanceName) { + if reqInstanceName != "" && !strings.EqualFold(reqInstanceName, instanceName) { return errors.New("Instance name does not match path") } - if req.ModelName == "" { - return errors.New("Model name is required") + return nil +} + +func prepareAddModelRequest(req *service.AddModelRequest, providerName, instanceName string) error { + if err := prepareProviderInstance(providerName, instanceName, req.ProviderName, req.InstanceName); err != nil { + return err } - if len(req.ModelTypes) == 0 { - return errors.New("Model type is required") + if len(req.Models) == 0 { + return errors.New("Models are required") + } + + for _, model := range req.Models { + if model.ModelName == "" { + return errors.New("Model name is required") + } + + if len(model.ModelTypes) == 0 { + return errors.New("Model type is required") + } } req.ProviderName = providerName @@ -802,8 +816,8 @@ func prepareAddCustomModelRequest(req *service.AddCustomModelRequest, providerNa return nil } -func (h *ProviderHandler) AddCustomModel(c *gin.Context) { - var req service.AddCustomModelRequest +func (h *ProviderHandler) AddModel(c *gin.Context) { + var req service.AddModelRequest if err := c.ShouldBindJSON(&req); err != nil { println("JSON bind error: %v (type: %T)", err, err) c.JSON(http.StatusBadRequest, gin.H{ @@ -813,7 +827,7 @@ func (h *ProviderHandler) AddCustomModel(c *gin.Context) { return } - if err := prepareAddCustomModelRequest(&req, c.Param("provider_name"), c.Param("instance_name")); err != nil { + if err := prepareAddModelRequest(&req, c.Param("provider_name"), c.Param("instance_name")); err != nil { c.JSON(http.StatusBadRequest, gin.H{ "code": common.CodeBadRequest, "message": err.Error(), @@ -823,19 +837,19 @@ func (h *ProviderHandler) AddCustomModel(c *gin.Context) { userID := c.GetString("user_id") - errorCode, err := h.modelProviderService.AddCustomModel(&req, userID) + code, err := h.modelProviderService.AddModel(&req, userID) if err != nil { c.JSON(http.StatusOK, gin.H{ - "code": errorCode, + "code": code, "message": err.Error(), }) return } c.JSON(http.StatusOK, gin.H{ - "code": common.CodeSuccess, + "code": code, + "message": "success", }) - } type DropInstanceModelRequest struct { diff --git a/internal/router/router.go b/internal/router/router.go index bc56eb5aa..6a285a898 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -327,7 +327,7 @@ func (r *Router) Setup(engine *gin.Engine) { provider.DELETE("/:provider_name/instances", r.providerHandler.DropProviderInstance) provider.GET("/:provider_name/instances/:instance_name/models", r.providerHandler.ListInstanceModels) provider.PATCH("/:provider_name/instances/:instance_name/models/*model_name", r.providerHandler.EnableOrDisableModel) - provider.POST("/:provider_name/instances/:instance_name/models", r.providerHandler.AddCustomModel) + provider.POST("/:provider_name/instances/:instance_name/models", r.providerHandler.AddModel) provider.DELETE("/:provider_name/instances/:instance_name/models", r.providerHandler.DropInstanceModels) v1.POST("/chat/completions", r.providerHandler.ChatToModel) v1.POST("/embeddings", r.providerHandler.EmbedText) diff --git a/internal/service/model_service.go b/internal/service/model_service.go index ce618978b..3d488eccc 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -25,6 +25,8 @@ import ( "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" "strings" + + "gorm.io/gorm" ) // parseModelName parses a composite model name in format "model@instance@provider" or "model@provider" @@ -1961,36 +1963,37 @@ func (m *ModelProviderService) GetChatModel(tenantID, compositeModelName string) return modelModule.NewChatModel(driver, &modelName, apiConfig), nil } -type AddCustomModelRequest struct { - ProviderName string `json:"provider_name"` - InstanceName string `json:"instance_name"` - ModelName string `json:"model_name"` - ModelTypes []string `json:"model_types"` - MaxTokens int `json:"max_tokens"` - Thinking *bool `json:"thinking"` +type AddModelRequest struct { + ProviderName string `json:"provider_name"` + InstanceName string `json:"instance_name"` + Models []ModelRequest `json:"models"` } -func (m *ModelProviderService) AddCustomModel(request *AddCustomModelRequest, userID string) (common.ErrorCode, error) { +type ModelRequest struct { + ModelName string `json:"model_name"` + ModelTypes []string `json:"model_types"` + MaxTokens int `json:"max_tokens"` + Thinking *bool `json:"thinking"` +} + +func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) (common.ErrorCode, error) { if request == nil { return common.CodeBadRequest, errors.New("request is required") } - if len(request.ModelTypes) == 0 { - return common.CodeBadRequest, errors.New("model type is required") + if len(request.Models) == 0 { + return common.CodeBadRequest, errors.New("models is required") } - // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") if err != nil { return common.CodeServerError, err } - if len(tenants) == 0 { return common.CodeNotFound, errors.New("user has no tenants") } tenantID := tenants[0].TenantID - // Check if provider exists provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, request.ProviderName) if err != nil { return common.CodeServerError, err @@ -2001,41 +2004,67 @@ func (m *ModelProviderService) AddCustomModel(request *AddCustomModelRequest, us return common.CodeServerError, err } - _, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, request.ModelName) - if err == nil { - return common.CodeConflict, errors.New("model already exists") + seen := make(map[string]struct{}) + models := make([]*entity.TenantModel, 0, len(request.Models)) + + for _, model := range request.Models { + modelName := strings.TrimSpace(model.ModelName) + if len(model.ModelTypes) == 0 { + return common.CodeBadRequest, errors.New("model types is required") + } + modelType := strings.TrimSpace(model.ModelTypes[0]) + + if modelName == "" { + return common.CodeBadRequest, errors.New("model name is required") + } + if modelType == "" { + return common.CodeBadRequest, errors.New("model type is required") + } + + duplicateKey := strings.ToLower(modelName) + if _, ok := seen[duplicateKey]; ok { + return common.CodeConflict, fmt.Errorf("duplicate model in request: %s", modelName) + } + seen[duplicateKey] = struct{}{} + + _, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName) + if err == nil { + return common.CodeConflict, fmt.Errorf("model already exists: %s", modelName) + } + if !errors.Is(err, gorm.ErrRecordNotFound) { + return common.CodeServerError, err + } + + modelID, err := generateUUID1Hex() + if err != nil { + return common.CodeServerError, errors.New("fail to get UUID") + } + + extra := map[string]interface{}{ + "max_tokens": model.MaxTokens, + "model_types": []string{modelType}, + } + if model.Thinking != nil { + extra["thinking"] = *model.Thinking + } + + extraByte, err := json.Marshal(extra) + if err != nil { + return common.CodeServerError, errors.New("fail to marshal extra") + } + + models = append(models, &entity.TenantModel{ + ID: modelID, + ModelName: modelName, + ModelType: modelType, + ProviderID: provider.ID, + InstanceID: instance.ID, + Status: "active", + Extra: string(extraByte), + }) } - modelID, err := generateUUID1Hex() - if err != nil { - return common.CodeServerError, errors.New("fail to get UUID") - } - - extra := make(map[string]interface{}) - extra["max_tokens"] = request.MaxTokens - if request.Thinking != nil { - extra["thinking"] = *request.Thinking - } - extra["model_types"] = request.ModelTypes - // convert extra to string - extraByte, err := json.Marshal(extra) - if err != nil { - return common.CodeServerError, errors.New("fail to marshal extra") - } - extraStr := string(extraByte) - - model := &entity.TenantModel{ - ID: modelID, - ModelName: request.ModelName, - ModelType: request.ModelTypes[0], - ProviderID: provider.ID, - InstanceID: instance.ID, - Status: "active", - Extra: extraStr, - } - - err = m.modelDAO.Create(model) - if err != nil { + if err := m.modelDAO.CreateBatch(models); err != nil { return common.CodeServerError, err }