diff --git a/internal/entity/models/google.go b/internal/entity/models/google.go index dc74fbc618..0ef6112735 100644 --- a/internal/entity/models/google.go +++ b/internal/entity/models/google.go @@ -22,6 +22,7 @@ import ( "fmt" "math" "ragflow/internal/common" + "ragflow/internal/entity" "strings" "google.golang.org/genai" @@ -44,12 +45,62 @@ func collectGoogleModelNames(ctx context.Context, listPage func(context.Context, models = append(models, page.items...) if page.nextPageToken == "" { - return ParseListModel(ModelList{Models: models}), nil + return finalizeGoogleModelList(ParseListModel(ModelList{Models: models})), nil } pageToken = page.nextPageToken } } +// googleUsableActions lists the Gemini supportedActions RAGFlow can serve: +// generateContent covers chat/vision/tts, embedContent covers embedding. +var googleUsableActions = []string{"generateContent", "embedContent", "batchEmbedContents"} + +// googleSupportsUsableAction reports whether the model supports at least one +// action RAGFlow can use. Models limited to other actions (image/video/music +// generation, question answering, etc.) are filtered out while listing so +// the list never offers models whose model type is unknown or unsupported. +func googleSupportsUsableAction(actions []string) bool { + for _, action := range actions { + for _, usable := range googleUsableActions { + if action == usable { + return true + } + } + } + return false +} + +// finalizeGoogleModelList resolves model types for listed models and filters +// out unknown or unsupported model_type values at list-generation time, +// rather than letting them flow into AddModel where they would be stored as +// model_type = 0 and repaired later. Catalog types win; models missing from +// the static catalog fall back to name-hint inference; models that still +// have no supported type are dropped. +func finalizeGoogleModelList(list []ListModelResponse) []ListModelResponse { + if list == nil { + return nil + } + filtered := make([]ListModelResponse, 0, len(list)) + for _, item := range list { + types := item.ModelTypes + if len(types) == 0 { + types = InferModelTypes(item.Name) + } + supported := make([]string, 0, len(types)) + for _, t := range types { + if entity.ModelTypeFromString(t) != 0 { + supported = append(supported, t) + } + } + if len(supported) == 0 { + continue + } + item.ModelTypes = supported + filtered = append(filtered, item) + } + return filtered +} + var googleListModels = func(ctx context.Context, config *genai.ClientConfig) ([]ListModelResponse, error) { client, err := genai.NewClient(ctx, config) if err != nil { @@ -64,14 +115,25 @@ var googleListModels = func(ctx context.Context, config *genai.ClientConfig) ([] var modelNames []ModelListItem for _, m := range models.Items { - modelName := strings.TrimSpace(m.DisplayName) + // Skip models limited to actions RAGFlow cannot serve + // (e.g. imagen/veo generation-only models); see + // finalizeGoogleModelList. + if len(m.SupportedActions) > 0 && !googleSupportsUsableAction(m.SupportedActions) { + continue + } + // Use the API model ID ("models/gemini-2.5-flash" → + // "gemini-2.5-flash") so listed models match the static + // catalog (model types / max_tokens) and are directly + // usable in chat requests. Display names ("Gemini 2.5 + // Flash") are not accepted by the Gemini API. + modelName := strings.TrimSpace(strings.TrimPrefix(m.Name, "models/")) if modelName == "" { - modelName = strings.TrimSpace(m.Name) + modelName = strings.TrimSpace(m.DisplayName) } if modelName != "" { modelNames = append(modelNames, ModelListItem{ ID: modelName, - OwnedBy: "Google", + OwnedBy: "Gemini", }) } } diff --git a/internal/entity/models/google_test.go b/internal/entity/models/google_test.go index 530fb223cf..6cac27b269 100644 --- a/internal/entity/models/google_test.go +++ b/internal/entity/models/google_test.go @@ -392,7 +392,10 @@ func TestCollectGoogleModelNamesPaginates(t *testing.T) { t.Fatalf("expected no error, got %v", err) } - expectedModels := []ListModelResponse{{Name: "Gemini 2.5 Flash"}, {Name: "Gemini 2.5 Pro"}} + expectedModels := []ListModelResponse{ + {Name: "Gemini 2.5 Flash", ModelTypes: []string{"chat"}}, + {Name: "Gemini 2.5 Pro", ModelTypes: []string{"chat"}}, + } if !reflect.DeepEqual(models, expectedModels) { t.Fatalf("expected models %v, got %v", expectedModels, models) } @@ -430,6 +433,57 @@ func TestCollectGoogleModelNamesReturnsPageError(t *testing.T) { } } +func TestFinalizeGoogleModelListFiltersUnknownModelTypes(t *testing.T) { + list := []ListModelResponse{ + {Name: "gemini-2.5-pro"}, // not in catalog: inferred + {Name: "gemini-embedding-001"}, // not in catalog: inferred + {Name: "custom", ModelTypes: []string{"chat", "image-gen"}}, // unsupported value stripped + {Name: "broken", ModelTypes: []string{"image-gen"}}, // no supported type: dropped + } + + got := finalizeGoogleModelList(list) + + expected := []ListModelResponse{ + {Name: "gemini-2.5-pro", ModelTypes: []string{"chat"}}, + {Name: "gemini-embedding-001", ModelTypes: []string{"embedding"}}, + {Name: "custom", ModelTypes: []string{"chat"}}, + } + if !reflect.DeepEqual(got, expected) { + t.Fatalf("expected models %v, got %v", expected, got) + } +} + +func TestFinalizeGoogleModelListPreservesNil(t *testing.T) { + if got := finalizeGoogleModelList(nil); got != nil { + t.Fatalf("expected nil, got %v", got) + } +} + +func TestGoogleSupportsUsableAction(t *testing.T) { + usable := [][]string{ + {"generateContent", "countTokens"}, + {"embedContent"}, + {"batchEmbedContents"}, + } + for _, actions := range usable { + if !googleSupportsUsableAction(actions) { + t.Fatalf("expected actions %v to be usable", actions) + } + } + unusable := [][]string{ + nil, + {"predict"}, // imagen-style image generation + {"predictLongRunning"}, // veo-style video generation + {"generateAnswer"}, // aqa-style question answering + {"createCachedContent"}, // cache-only entry + } + for _, actions := range unusable { + if googleSupportsUsableAction(actions) { + t.Fatalf("expected actions %v to be filtered out", actions) + } + } +} + func TestGoogleGenerateContentConfigConvertsTools(t *testing.T) { toolChoice := "required" cfg, err := googleGenerateContentConfig(&ChatConfig{ diff --git a/internal/handler/providers.go b/internal/handler/providers.go index b3592961dd..9869a0c004 100644 --- a/internal/handler/providers.go +++ b/internal/handler/providers.go @@ -727,11 +727,6 @@ func (h *ProviderHandler) AddModel(c *gin.Context) { return } - if len(req.ModelTypes) == 0 { - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_type is required") - return - } - userID := c.GetString("user_id") ctx := c.Request.Context() diff --git a/internal/service/model_service.go b/internal/service/model_service.go index e420d157d2..e917650b57 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3703,10 +3703,6 @@ func (m *ModelProviderService) AddModel(ctx context.Context, request *AddModelRe return common.CodeBadRequest, errors.New("model_name is required") } - if len(request.ModelTypes) == 0 { - return common.CodeBadRequest, errors.New("model_type is required") - } - tenants, err := m.userTenantDAO.GetByUserIDAndRole(ctx, dao.DB, userID, "owner") if err != nil { return common.CodeServerError, err @@ -3741,18 +3737,15 @@ func (m *ModelProviderService) AddModel(ctx context.Context, request *AddModelRe return common.CodeServerError, err } - // Compute model type bitmask. + // Compute model type bitmask. Matches Python's calculate_model_type: + // empty and unrecognized type names are ignored, not rejected. combinedType := entity.ModelType(0) for _, rawType := range request.ModelTypes { mt := strings.TrimSpace(rawType) if mt == "" { continue } - t := entity.ModelTypeFromString(mt) - if t == 0 { - return common.CodeBadRequest, fmt.Errorf("invalid model type: %s", mt) - } - combinedType |= t + combinedType |= entity.ModelTypeFromString(mt) } maxTokens := request.MaxTokens