zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-06 19:38:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat[Go]: implement add multi_models (#15563)

Browse Source
This commit is contained in:
Haruko386
2026-06-03 15:26:46 +08:00
committed by GitHub
parent c0e00a7f6e
commit 473d06d1ad
5 changed files with 122 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
internal/dao/tenant_model.go
View File

@@ -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

5
internal/entity/models/moonshot.go
View File

@@ -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:]

42
internal/handler/providers.go
View File

@@ -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 {

2
internal/router/router.go
View File

@@ -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)

121
internal/service/model_service.go
View File

@@ -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
}