diff --git a/Dockerfile b/Dockerfile index 6e5d87f945..7200d95a1a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -157,6 +157,8 @@ RUN --mount=type=cache,id=ragflow_apt,target=/var/cache/apt,sharing=locked \ fi || \ { echo "Failed to install ODBC driver"; exit 1; } + + # Add dependencies of selenium RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/chrome-linux64-121-0-6167-85,target=/chrome-linux64.zip \ unzip /chrome-linux64.zip && \ diff --git a/api/apps/restful_apis/provider_api.py b/api/apps/restful_apis/provider_api.py index 6010e144c7..422e4a2956 100644 --- a/api/apps/restful_apis/provider_api.py +++ b/api/apps/restful_apis/provider_api.py @@ -641,9 +641,7 @@ async def update_provider_instance(tenant_id: str = None, provider_id_or_name: s verify = data.get("verify", True) try: - success, msg = await provider_api_service.update_provider_instance( - tenant_id, provider_id_or_name, instance_id_or_name, instance_name, api_key, base_url, region, model_info, verify - ) + success, msg = await provider_api_service.update_provider_instance(tenant_id, provider_id_or_name, instance_id_or_name, instance_name, api_key, base_url, region, model_info, verify) if success: return get_result(message=msg) else: @@ -961,13 +959,12 @@ async def delete_models_from_instance(tenant_id: str, provider_id_or_name: str, return get_error_data_result(message="Internal server error") - @manager.route("/providers//instances//models/", methods=["PATCH"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs async def alter_model(tenant_id: str = None, provider_id_or_name: str = None, instance_id_or_name: str = None, model_name: str = None): """ - Enable or disable a model, or update max_tokens + Alter a model's status, max_tokens, model_type, or extra fields. --- tags: - Providers @@ -996,35 +993,40 @@ async def alter_model(tenant_id: str = None, provider_id_or_name: str = None, in description: Bearer token for authentication. - in: body name: body - description: Model status update. + description: Model update parameters. required: true schema: type: object - required: - - status properties: status: type: string enum: ["active", "inactive"] description: Model status. + max_tokens: + type: integer + description: Maximum number of tokens. + model_type: + type: string + description: Model type (e.g. "chat", "embedding"). + extra: + type: object + description: Additional model configuration. responses: 200: - description: Model status updated. + description: Model updated. schema: type: object """ data = await request.get_json() - if not data or ("status" not in data and "max_tokens" not in data): - return get_error_argument_result(message="status or max_tokens required.") + if not data or not any(k in data for k in ("status", "max_tokens", "model_type", "extra")): + return get_error_argument_result(message="At least one of status, max_tokens, model_type, or extra is required.") update_dict = {k: data[k] for k in ["status", "max_tokens", "model_type", "extra"] if k in data} if update_dict.get("status") and update_dict["status"] not in ("active", "inactive"): return get_error_argument_result(message="status must be 'active' or 'inactive'") try: - success, msg = provider_api_service.update_model( - tenant_id, provider_id_or_name, instance_id_or_name, model_name, update_dict - ) + success, msg = provider_api_service.update_model(tenant_id, provider_id_or_name, instance_id_or_name, model_name, update_dict) if success: return get_result(message=msg) else: diff --git a/api/utils/model_utils.py b/api/utils/model_utils.py index f2a462de3b..f27a08044f 100644 --- a/api/utils/model_utils.py +++ b/api/utils/model_utils.py @@ -28,7 +28,7 @@ _LEGACY_MODEL_TYPE_ALIASES = { } -def calculate_model_type(model_type_name_list: List[str]|str) -> int: +def calculate_model_type(model_type_name_list: List[str] | str) -> int: model_type = 0 if isinstance(model_type_name_list, str): model_type_name_list = [model_type_name_list] diff --git a/common/constants.py b/common/constants.py index b171263b91..66965af7b9 100644 --- a/common/constants.py +++ b/common/constants.py @@ -94,7 +94,7 @@ class LLMType(StrEnum): class ModelTypeBinary(Enum): - CHAT = 0b0000001 # 1 << 0 = 1 + CHAT = 0b0000001 # 1 << 0 = 1 EMBEDDING = 0b0000010 # 1 << 1 = 2 ASR = 0b0000100 # 1 << 2 = 4 VISION = 0b0001000 # 1 << 3 = 8 diff --git a/internal/agent/component/categorize_test.go b/internal/agent/component/categorize_test.go index 57f80a00e9..fcb75bc1eb 100644 --- a/internal/agent/component/categorize_test.go +++ b/internal/agent/component/categorize_test.go @@ -199,46 +199,6 @@ func TestSplitCompositeLLMID(t *testing.T) { } } -func TestCategorize_ResolvesTenantLLMCredentials(t *testing.T) { - db := setupComponentTestDB(t) - pushComponentDB(t, db) - apiKey := "tenant-llm-key" - apiBase := "https://tenant-llm.example" - modelName := "Qwen/Qwen3-8B" - if err := db.Create(&entity.TenantLLM{ - TenantID: "tenant-1", - LLMFactory: "SILICONFLOW", - LLMName: &modelName, - APIKey: &apiKey, - APIBase: &apiBase, - Status: "1", - }).Error; err != nil { - t.Fatalf("create tenant_llm: %v", err) - } - - stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "support", Model: "stub"}} - withStubInvoker(t, stub) - - c := NewCategorizeComponent(CategorizeParam{ - ModelID: "Qwen/Qwen3-8B@default@SILICONFLOW", - Categories: []string{"sales", "support"}, - DefaultCategory: "support", - }) - _, err := c.Invoke(stateWithTenant("tenant-1"), map[string]any{}) - if err != nil { - t.Fatalf("Invoke: %v", err) - } - if stub.captured == nil { - t.Fatal("invoker not called") - } - if stub.captured.APIKey != apiKey { - t.Fatalf("APIKey=%q, want %q", stub.captured.APIKey, apiKey) - } - if stub.captured.BaseURL != apiBase { - t.Fatalf("BaseURL=%q, want %q", stub.captured.BaseURL, apiBase) - } -} - func TestCategorize_ResolvesTenantModelInstanceCredentials(t *testing.T) { db := setupComponentTestDB(t) pushComponentDB(t, db) @@ -387,7 +347,6 @@ func setupComponentTestDB(t *testing.T) *gorm.DB { t.Fatalf("open sqlite: %v", err) } if err := db.AutoMigrate( - &entity.TenantLLM{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}, diff --git a/internal/binding/cpp/rag_analyzer.cpp b/internal/binding/cpp/rag_analyzer.cpp index a4a3bcc92d..f9d52e7336 100644 --- a/internal/binding/cpp/rag_analyzer.cpp +++ b/internal/binding/cpp/rag_analyzer.cpp @@ -733,6 +733,12 @@ int32_t RAGAnalyzer::Load() { int32_t freq = std::stoi(results[1]); freq = int32_t(std::log(float(freq) / DENOMINATOR) + 0.5); int32_t pos_idx = pos_table_->GetPOSIndex(results[2]); + // If the POS tag is not in the POS table (e.g. "eng" for + // English words), default to index 0. Using -1 would cause + // Encode() to produce a negative int32_t which the darts + // trie rejects. + if (pos_idx < 0) + pos_idx = 0; int value = Encode(freq, pos_idx); trie_->Add(results[0], value); std::string rkey = RKey(results[0]); diff --git a/internal/entity/models/model.go b/internal/entity/models/model.go index 96dcef6423..c13ff52de6 100644 --- a/internal/entity/models/model.go +++ b/internal/entity/models/model.go @@ -169,6 +169,7 @@ type Model struct { MaxDimension *int `json:"max_dimension"` // used by embedding models Dimensions []int `json:"dimensions"` Alias []string `json:"alias"` + Rank *int `json:"rank"` ModelTypeMap map[string]bool } @@ -488,7 +489,7 @@ func (pm *ProviderManager) ListModels(providerName string) ([]map[string]interfa modelData := map[string]interface{}{ "name": model.Name, "max_tokens": model.MaxTokens, - "model_type": model.ModelTypes, + "model_types": model.ModelTypes, "max_dimension": model.MaxDimension, "dimensions": model.Dimensions, } @@ -513,7 +514,7 @@ func (pm *ProviderManager) GetModelByName(providerName, modelName string) (*Mode if provider == nil { return nil, fmt.Errorf("provider '%s' not found", providerName) } - model := pm.findModel(provider, modelName) + model := pm.FindModel(provider, modelName) if model == nil { return nil, fmt.Errorf("model '%s' not found", modelName) } @@ -525,7 +526,7 @@ func (pm *ProviderManager) GetModelUrl(providerName, modelName, modelType string if provider == nil { return nil, nil, fmt.Errorf("provider '%s' not found", providerName) } - model := pm.findModel(provider, modelName) + model := pm.FindModel(provider, modelName) if model == nil { return nil, nil, fmt.Errorf("model '%s' not found", modelName) } @@ -568,7 +569,7 @@ func (pm *ProviderManager) SearchModelInfo(providerName, modelName string, filte return resp } - model := pm.findModel(provider, modelName) + model := pm.FindModel(provider, modelName) if model == nil { resp.Code = 404 resp.Message = fmt.Sprintf("Model '%s' not found in provider '%s'", modelName, providerName) @@ -794,7 +795,7 @@ func (pm *ProviderManager) FindProvider(name string) *Provider { } // Helper: Find model by name -func (pm *ProviderManager) findModel(provider *Provider, modelName string) *Model { +func (pm *ProviderManager) FindModel(provider *Provider, modelName string) *Model { for i := range provider.Models { if strings.EqualFold(provider.Models[i].Name, modelName) { return provider.Models[i] diff --git a/internal/handler/providers.go b/internal/handler/providers.go index 29092d87ec..6aeb8600bc 100644 --- a/internal/handler/providers.go +++ b/internal/handler/providers.go @@ -244,16 +244,15 @@ func (h *ProviderHandler) ShowProviderInstance(c *gin.Context) { return } - instanceName := c.Param("instance_name") - if instanceName == "" { + instanceIDOrName := c.Param("instance_name") + if instanceIDOrName == "" { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Instance name is required") return } userID := c.GetString("user_id") - // Get tenant ID from user - instance, errorCode, err := h.modelProviderService.ShowProviderInstance(providerName, instanceName, userID) + instance, errorCode, err := h.modelProviderService.ShowProviderInstance(providerName, instanceIDOrName, userID) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -516,12 +515,15 @@ func (h *ProviderHandler) ListInstanceModels(c *gin.Context) { common.SuccessWithData(c, modelInstances, "success") } -type EnableOrDisableModelRequest struct { - ModelID string `json:"model_id"` - Status string `json:"status"` +type AlterModelRequest struct { + ModelID string `json:"model_id"` + Status string `json:"status"` + MaxTokens int `json:"max_tokens"` + ModelType interface{} `json:"model_type"` + Extra interface{} `json:"extra"` } -func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) { +func (h *ProviderHandler) AlterModel(c *gin.Context) { providerName := c.Param("provider_name") if providerName == "" { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required") @@ -534,29 +536,43 @@ func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) { return } - var req EnableOrDisableModelRequest + var req AlterModelRequest if err := c.ShouldBindJSON(&req); err != nil { - println("JSON bind error: %v (type: %T)", err, err) common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } userID := c.GetString("user_id") - modelID := strings.TrimSpace(req.ModelID) modelName := strings.TrimPrefix(c.Param("model_name"), "/") modelName = strings.TrimSpace(modelName) + modelID := strings.TrimSpace(req.ModelID) + if modelName == "" && modelID == "" { common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_name or model_id is required") return } status := strings.TrimSpace(req.Status) - if status != "active" && status != "inactive" { - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "Status must be active or inactive") + if status != "" && status != "active" && status != "inactive" { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "status must be 'active' or 'inactive'") return } - code, err := h.modelProviderService.UpdateModelStatus(providerName, instanceName, modelName, userID, modelID, status) + updateDict := make(map[string]interface{}) + if status != "" { + updateDict["status"] = status + } + if req.MaxTokens > 0 { + updateDict["max_tokens"] = req.MaxTokens + } + if req.ModelType != nil { + updateDict["model_type"] = req.ModelType + } + if req.Extra != nil { + updateDict["extra"] = req.Extra + } + + code, err := h.modelProviderService.AlterModel(providerName, instanceName, modelName, userID, modelID, updateDict) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -585,30 +601,6 @@ func prepareProviderInstance(providerName, instanceName, reqProviderName, reqIns 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.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 - req.InstanceName = instanceName - return nil -} - func (h *ProviderHandler) AddModel(c *gin.Context) { var req service.AddModelRequest if err := c.ShouldBindJSON(&req); err != nil { @@ -617,8 +609,21 @@ func (h *ProviderHandler) AddModel(c *gin.Context) { return } - if err := prepareAddModelRequest(&req, c.Param("provider_name"), c.Param("instance_name")); err != nil { - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeDataError, nil, err.Error()) + req.ProviderName = c.Param("provider_name") + req.InstanceName = c.Param("instance_name") + + if req.ProviderName == "" || req.InstanceName == "" { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "provider_name and instance_name are required") + return + } + + if req.ModelName == "" { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_name is required") + return + } + + if len(req.ModelTypes) == 0 { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_type is required") return } @@ -634,8 +639,7 @@ func (h *ProviderHandler) AddModel(c *gin.Context) { } type DropInstanceModelRequest struct { - ModelIDs []string `json:"model_ids"` - Models []string `json:"models"` + ModelNames []string `json:"model_name"` } func (h *ProviderHandler) DropInstanceModels(c *gin.Context) { @@ -655,14 +659,14 @@ func (h *ProviderHandler) DropInstanceModels(c *gin.Context) { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } - if len(req.ModelIDs) == 0 && len(req.Models) == 0 { - common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_ids or models is required") + if len(req.ModelNames) == 0 { + common.ResponseWithHttpCodeData(c, http.StatusBadRequest, common.CodeBadRequest, nil, "model_name is required") return } userID := c.GetString("user_id") - code, err := h.modelProviderService.DropInstanceModels(providerName, instanceName, userID, req.ModelIDs, req.Models) + code, err := h.modelProviderService.DropInstanceModels(providerName, instanceName, userID, req.ModelNames) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/providers_test.go b/internal/handler/providers_test.go index 58e523f21a..c19f14c61d 100644 --- a/internal/handler/providers_test.go +++ b/internal/handler/providers_test.go @@ -84,7 +84,7 @@ func decodeProviderHandlerResponse(t *testing.T, recorder *httptest.ResponseReco return body } -func TestProviderHandlerEnableOrDisableModelRejectsMissingModelSelector(t *testing.T) { +func TestProviderHandlerAlterModelRejectsMissingModelSelector(t *testing.T) { ctx, recorder := newProviderHandlerRequest( t, map[string]interface{}{"status": "active"}, @@ -92,7 +92,7 @@ func TestProviderHandlerEnableOrDisableModelRejectsMissingModelSelector(t *testi gin.Param{Key: "instance_name", Value: "default"}, ) - NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx) + NewProviderHandler(nil, service.NewModelProviderService()).AlterModel(ctx) if recorder.Code != http.StatusBadRequest { t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadRequest, recorder.Body.String()) @@ -103,7 +103,7 @@ func TestProviderHandlerEnableOrDisableModelRejectsMissingModelSelector(t *testi } } -func TestProviderHandlerEnableOrDisableModelRejectsInvalidStatus(t *testing.T) { +func TestProviderHandlerAlterModelRejectsInvalidStatus(t *testing.T) { ctx, recorder := newProviderHandlerRequest( t, map[string]interface{}{"status": "disabled"}, @@ -112,7 +112,7 @@ func TestProviderHandlerEnableOrDisableModelRejectsInvalidStatus(t *testing.T) { gin.Param{Key: "model_name", Value: "gpt-test"}, ) - NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx) + NewProviderHandler(nil, service.NewModelProviderService()).AlterModel(ctx) if recorder.Code != http.StatusBadRequest { t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadRequest, recorder.Body.String()) @@ -123,7 +123,7 @@ func TestProviderHandlerEnableOrDisableModelRejectsInvalidStatus(t *testing.T) { } } -func TestProviderHandlerEnableOrDisableModelUpdatesStatus(t *testing.T) { +func TestProviderHandlerAlterModelUpdatesStatus(t *testing.T) { db := setupProviderHandlerTestDB(t) useProviderHandlerTestDB(t, db) seedProviderHandlerModel(t, db) @@ -136,7 +136,7 @@ func TestProviderHandlerEnableOrDisableModelUpdatesStatus(t *testing.T) { gin.Param{Key: "model_name", Value: "gpt-test"}, ) - NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx) + NewProviderHandler(nil, service.NewModelProviderService()).AlterModel(ctx) if recorder.Code != http.StatusOK { t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String()) diff --git a/internal/router/router.go b/internal/router/router.go index f4c750e4c8..ae3ae6a29e 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -539,7 +539,7 @@ func (r *Router) Setup(engine *gin.Engine) { provider.PUT("/:provider_name/instances/:instance_name", r.providerHandler.AlterProviderInstance) 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.PATCH("/:provider_name/instances/:instance_name/models/*model_name", r.providerHandler.AlterModel) 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/to_model", r.providerHandler.ChatToModel) diff --git a/internal/service/agent_run_e2e_test.go b/internal/service/agent_run_e2e_test.go index 11d9f75c13..6f739605ee 100644 --- a/internal/service/agent_run_e2e_test.go +++ b/internal/service/agent_run_e2e_test.go @@ -158,6 +158,11 @@ func TestRunAgent_RealCanvas_BeginMessage(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -240,6 +245,11 @@ func TestRunAgent_RealCanvas_WaitForUserResume(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -367,6 +377,11 @@ func TestRunAgent_RealCanvas_WaitForUserResume_EventSemantics(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -466,6 +481,11 @@ func TestRunAgent_RealCanvas_GroupedParallelOuterFollower(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -594,6 +614,11 @@ func TestRunAgent_AllFixture_LoopInterruptResume(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -698,6 +723,11 @@ func TestRunAgent_AllFixture_LoopInterruptResume_MultiTurn(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -786,6 +816,11 @@ func TestRunAgent_AllFixture_IterationFormatsItems(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -882,6 +917,11 @@ func TestRunAgent_AllFixture_VarAssigner(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -942,6 +982,11 @@ func TestRunAgent_AllFixture_DataOps(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -1018,6 +1063,11 @@ func TestRunAgent_RealCanvas_CompileFails(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -1090,6 +1140,11 @@ func TestRunAgent_AllFixture_CategorizeResume(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -1217,6 +1272,11 @@ func TestRunAgent_RealCanvas_InvokeFails(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } @@ -1319,6 +1379,11 @@ func TestRunAgent_RunTracker_AttachCheckpoint_CallSequence(t *testing.T) { &entity.UserTenant{}, &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.APIToken{}, + &entity.API4Conversation{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } diff --git a/internal/service/agent_test.go b/internal/service/agent_test.go index 727dbed9c6..80848f5e7e 100644 --- a/internal/service/agent_test.go +++ b/internal/service/agent_test.go @@ -220,7 +220,9 @@ func TestGetVersion_WrongCanvas(t *testing.T) { t.Helper() if err := testDB.AutoMigrate( + &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.UserTenant{}, ); err != nil { t.Fatalf("failed to migrate: %v", err) } @@ -248,7 +250,9 @@ func TestGetVersion_NotFound(t *testing.T) { t.Helper() if err := testDB.AutoMigrate( + &entity.UserCanvas{}, &entity.UserCanvasVersion{}, + &entity.UserTenant{}, ); err != nil { t.Fatalf("failed to migrate: %v", err) } diff --git a/internal/service/model_service.go b/internal/service/model_service.go index 1b1a8ff2d0..f3a4ea4976 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -25,6 +25,7 @@ import ( "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" "ragflow/internal/utility" + "sort" "strings" "go.uber.org/zap" @@ -387,8 +388,6 @@ func (m *ModelProviderService) ListSupportedModels(providerName, instanceName, u return result, nil } -// CreateInstanceModelInfo mirrors Python's model_info dict in -// create_provider_instance. See api/apps/services/provider_api_service.py:434. type CreateInstanceModelInfo struct { ModelName string `json:"model_name"` ModelTypes []string `json:"model_type"` @@ -396,8 +395,16 @@ type CreateInstanceModelInfo struct { Extra map[string]interface{} `json:"extra"` } -func (m *ModelProviderService) CreateProviderInstance(providerName, instanceName, apiKey, baseURL, region, userID string, modelInfo []CreateInstanceModelInfo) (common.ErrorCode, error) { - providerName = strings.TrimSpace(providerName) +func (m *ModelProviderService) getProviderByIDOrName(tenantID, providerIDOrName string) (*entity.TenantModelProvider, error) { + provider, err := m.modelProviderDAO.GetByID(providerIDOrName) + if err == nil && provider.TenantID == tenantID { + return provider, nil + } + return m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, strings.TrimSpace(providerIDOrName)) +} + +func (m *ModelProviderService) CreateProviderInstance(providerIDOrName, instanceName, apiKey, baseURL, region, userID string, modelInfo []CreateInstanceModelInfo) (common.ErrorCode, error) { + providerIDOrName = strings.TrimSpace(providerIDOrName) // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") @@ -411,15 +418,15 @@ func (m *ModelProviderService) CreateProviderInstance(providerName, instanceName tenantID := tenants[0].TenantID - // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) if err != nil { - return common.CodeServerError, err + return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) } + providerName := provider.ProviderName // Normalize api_key: VLLM with empty api_key defaults to "x". // Mirrors Python's _normalize_provider_api_key. - if strings.EqualFold(providerName, "vllm") && apiKey == "" { + if strings.EqualFold(providerName, "VLLM") && apiKey == "" { apiKey = "x" } @@ -466,16 +473,54 @@ func (m *ModelProviderService) CreateProviderInstance(providerName, instanceName return common.CodeServerError, err } } + } else { + // model_info not provided — add all factory default models. + // Mirrors Python's create_provider_instance + // (api/apps/services/provider_api_service.py:506-531). + targetFactoryName := providerName + if region == "intl" && strings.EqualFold(providerName, "siliconflow") { + targetFactoryName = "siliconflow_intl" + } + factoryProvider := dao.GetModelProviderManager().FindProvider(targetFactoryName) + if factoryProvider != nil { + for _, llm := range factoryProvider.Models { + verifyStatus := modelVerifyResult[llm.Name] + if verifyStatus == "" { + verifyStatus = entity.ModelVerifyUnknown + } + extra := map[string]interface{}{ + "verify": verifyStatus, + } + if llm.Tools != nil { + extra["is_tools"] = llm.Tools.Support + } + if llm.Thinking != nil { + extra["thinking"] = llm.Thinking.DefaultValue + } + if err := m.addModelToInstance(tenantID, providerName, instanceName, CreateInstanceModelInfo{ + ModelName: llm.Name, + ModelTypes: llm.ModelTypes, + MaxTokens: func() int { + if llm.MaxTokens != nil { + return *llm.MaxTokens + } + return 8192 + }(), + Extra: extra, + }); err != nil { + return common.CodeServerError, err + } + } + } } return common.CodeSuccess, nil } // CreateNameOnlyProviderInstance creates a provider instance with only a name, -// skipping API key validation and model creation. Mirrors Python's -// create_name_only_provider_instance in provider_api_service.py:536. -func (m *ModelProviderService) CreateNameOnlyProviderInstance(providerName, instanceName, userID string) (common.ErrorCode, error) { - providerName = strings.TrimSpace(providerName) +// skipping API key validation and model creation. +func (m *ModelProviderService) CreateNameOnlyProviderInstance(providerIDOrName, instanceName, userID string) (common.ErrorCode, error) { + providerIDOrName = strings.TrimSpace(providerIDOrName) if instanceName == "default" { return common.CodeBadRequest, errors.New("instance name cannot be 'default'") @@ -490,9 +535,9 @@ func (m *ModelProviderService) CreateNameOnlyProviderInstance(providerName, inst } tenantID := tenants[0].TenantID - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) if err != nil { - return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) + return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) } instanceID := utility.GenerateToken() @@ -515,8 +560,7 @@ func (m *ModelProviderService) CreateNameOnlyProviderInstance(providerName, inst // verifyProviderAPIKey verifies the API key against the provider by calling // the driver's CheckConnection. It returns a map from model name to verify -// status (success/fail/unknown). Mirrors Python's verify_api_key in -// provider_api_service.py:596. +// status (success/fail/unknown). func (m *ModelProviderService) verifyProviderAPIKey(providerName, apiKey, region, baseURL string, modelInfo []CreateInstanceModelInfo) map[string]string { result := make(map[string]string) @@ -567,7 +611,6 @@ func (m *ModelProviderService) verifyProviderAPIKey(providerName, apiKey, region } // addModelToInstance creates a single model under the given provider instance. -// Mirrors Python's add_model_to_instance in provider_api_service.py:969. func (m *ModelProviderService) addModelToInstance(tenantID, providerName, instanceName string, model CreateInstanceModelInfo) error { provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) if err != nil { @@ -626,8 +669,8 @@ func (m *ModelProviderService) addModelToInstance(tenantID, providerName, instan return nil } -func (m *ModelProviderService) ListProviderInstances(providerName, userID string) ([]map[string]interface{}, common.ErrorCode, error) { - providerName = strings.TrimSpace(providerName) +func (m *ModelProviderService) ListProviderInstances(providerIDOrName, userID string) ([]map[string]interface{}, common.ErrorCode, error) { + providerIDOrName = strings.TrimSpace(providerIDOrName) // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") @@ -641,20 +684,10 @@ func (m *ModelProviderService) ListProviderInstances(providerName, userID string tenantID := tenants[0].TenantID - // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + // Check if provider exists — try by ID first, then by name. + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) if err != nil { - // "provider not connected to this tenant" is a normal, expected state - // (e.g. demo tenant has never added SiliconFlow). Mirrors the Python - // contract in api/apps/services/provider_api_service.py:349-355 which - // returns (False, "No provider found for provider ''") on this - // path. The REST layer maps that to get_error_data_result with - // code=RetCode.DATA_ERROR (=102), so the Go port must do the same — - // NOT a 500 server error. - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeDataError, fmt.Errorf("No provider found for provider '%s'", providerName) - } - return nil, common.CodeServerError, err + return nil, common.CodeDataError, fmt.Errorf("No provider found for provider '%s'", providerIDOrName) } // Check if provider exists @@ -670,36 +703,35 @@ func (m *ModelProviderService) ListProviderInstances(providerName, userID string // crash on a freshly created tenant. result := make([]map[string]interface{}, 0, len(instances)) for _, instance := range instances { - // convert instance.Extra (JSON string) to map - var extra map[string]string - err = json.Unmarshal([]byte(instance.Extra), &extra) - if err != nil { - return nil, common.CodeServerError, err + // Parse extra to extract region + var extraFields map[string]string + if instance.Extra != "" { + if err := json.Unmarshal([]byte(instance.Extra), &extraFields); err != nil { + return nil, common.CodeServerError, err + } + } + if extraFields == nil { + extraFields = make(map[string]string) } - // Emit snake_case keys to match the IProviderInstance TypeScript - // contract (web/src/interfaces/database/llm.ts) and the Python - // `TenantModelInstanceService.query(joinedload(...))` response - // shape. The Go port previously emitted camelCase - // (`instanceName`/`providerID`/`apiKey`), which broke the - // `instance.instance_name` reads in - // web/src/pages/user-setting/setting-model/components/used-model.tsx - // and made `useFetchInstanceModels(providerName, - // instance.instance_name)` hit `/api/v1/providers/

/instances/undefined/models`. + // Emit snake_case keys to match the Python + // list_provider_instances response shape + // (provider_api_service.py:563). result = append(result, map[string]interface{}{ "id": instance.ID, "instance_name": instance.InstanceName, "provider_id": instance.ProviderID, - "api_key": instance.APIKey, + "region": extraFields["region"], "status": instance.Status, - "extra": instance.Extra, }) } return result, common.CodeSuccess, nil } -func (m *ModelProviderService) ShowProviderInstance(providerName, instanceName, userID string) (map[string]interface{}, common.ErrorCode, error) { +func (m *ModelProviderService) ShowProviderInstance(providerName, instanceIDOrName, userID string) (map[string]interface{}, common.ErrorCode, error) { + providerName = strings.TrimSpace(providerName) + providerName = strings.ToLower(providerName) // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") @@ -713,37 +745,50 @@ func (m *ModelProviderService) ShowProviderInstance(providerName, instanceName, tenantID := tenants[0].TenantID - // Check if provider exists + // Check if provider exists. Mirrors Python's: + // provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_id(tenant_id, provider_id_or_name) + // if not provider_obj: + // provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, provider_id_or_name) provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) if err != nil { - return nil, common.CodeServerError, err + return nil, common.CodeDataError, fmt.Errorf("No provider found for provider '%s'", providerName) } - instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName) + // Find the instance — try by ID first, then by name. + // Mirrors Python's: + // _, instance_obj = TenantModelInstanceService.get_by_id(instance_id_or_name) + // if instance_obj and instance_obj.provider_id != provider_id: instance_obj = None + // if not instance_obj: + // instance_obj = TenantModelInstanceService.get_by_provider_id_and_instance_name(provider_id, instance_id_or_name) + var instance *entity.TenantModelInstance + instance, err = m.modelInstanceDAO.GetByID(instanceIDOrName) + if err != nil || instance.ProviderID != provider.ID { + instance, err = m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceIDOrName) + } if err != nil { - return nil, common.CodeServerError, err + return nil, common.CodeDataError, fmt.Errorf("No instance found for provider '%s' and instance '%s'", providerName, instanceIDOrName) } - // convert instance.Extra (JSON string) to map - var extra map[string]string - err = json.Unmarshal([]byte(instance.Extra), &extra) - if err != nil { - return nil, common.CodeServerError, err + // Parse extra fields. Mirrors Python's: + // extra_fields = json.loads(instance_obj.extra) if instance_obj.extra else {} + var extraFields map[string]string + if instance.Extra != "" { + if err := json.Unmarshal([]byte(instance.Extra), &extraFields); err != nil { + return nil, common.CodeServerError, err + } + } + if extraFields == nil { + extraFields = make(map[string]string) } - // Emit snake_case keys to match the IProviderInstance TypeScript - // contract — see ListProviderInstances above. The previous shape - // mixed conventions (`apikey` lowercase, `instanceName`/`providerID` - // camelCase) and broke the front-end's `instance.api_key` / - // `instance.instance_name` reads. result := map[string]interface{}{ "id": instance.ID, "instance_name": instance.InstanceName, - "provider_id": instance.ProviderID, - "status": instance.Status, + "provider_id": provider.ID, + "region": extraFields["region"], + "base_url": extraFields["base_url"], "api_key": instance.APIKey, - "region": extra["region"], - "base_url": extra["base_url"], + "status": instance.Status, } return result, common.CodeSuccess, nil @@ -1067,6 +1112,16 @@ func (m *ModelProviderService) ListTenantAddedModels(userID, ownerTenantID, mode modelTypeFilter = strings.ToLower(strings.TrimSpace(modelTypeFilter)) } + // Mirror Python's ensure_*_from_env calls. + _ = m.ensureMineruFromEnv(tenantID) + _ = m.ensurePaddleOCREnabledFromEnv(tenantID) + _ = m.ensureOpenDataLoaderFromEnv(tenantID) + + var modelTypeFilterBin entity.ModelType + if modelTypeFilter != "" { + modelTypeFilterBin = entity.ModelTypeFromString(modelTypeFilter) + } + providers, err := m.modelProviderDAO.GetByTenantID(tenantID) if err != nil { return nil, common.CodeServerError, err @@ -1097,105 +1152,127 @@ func (m *ModelProviderService) ListTenantAddedModels(userID, ownerTenantID, mode instanceInfoByID[inst.ID] = inst } - // Per-tenant enable/disable overrides. In the Go port this is - // typically empty (no writers), but we still honor active/inactive - // rows for correctness and parity. + // Fetch tenant model records and filter by model_type if needed. modelRecords, err := m.modelDAO.GetModelsByProviderIDsAndInstanceIDs(providerIDs, instanceIDs) if err != nil { return nil, common.CodeServerError, err } - activeByKey := make(map[string]int) - inactiveByKey := make(map[string]int) - activeModelIDByKey := make(map[string]string) - for _, rec := range modelRecords { - key := rec.ProviderID + "@" + rec.InstanceID + "@" + rec.ModelName - if rec.Status == "inactive" { - inactiveByKey[key] |= rec.ModelType - } else { - activeByKey[key] |= rec.ModelType - if activeModelIDByKey[key] == "" { - activeModelIDByKey[key] = rec.ID + + var targetRecords []*entity.TenantModel + if modelTypeFilterBin != 0 { + for _, rec := range modelRecords { + if entity.ModelType(rec.ModelType)&modelTypeFilterBin != 0 { + targetRecords = append(targetRecords, rec) } } + } else { + targetRecords = modelRecords } - // Group instances by provider_name for the outer loop. - instancesByProviderName := make(map[string][]*entity.TenantModelInstance) - for _, inst := range instances { - p, ok := providerInfoByID[inst.ProviderID] - if !ok { - continue - } - instancesByProviderName[p.ProviderName] = append(instancesByProviderName[p.ProviderName], inst) - } - + // Build model rank map from factory catalog (mirrors Python's model_rank_map). providerManager := dao.GetModelProviderManager() - added := make([]map[string]interface{}, 0) - - // factory rank is not present in the Go entity.Provider struct, so we - // follow Python's stable ordering intent (factory rank desc, then - // provider_name, then instance_name) by simply iterating providers in - // the tenant's own order. With one provider today this is a no-op. - for _, p := range providers { - factory := providerManager.FindProvider(p.ProviderName) - if factory == nil { - // Factory not in the static catalog. The tenant has linked - // a provider we have no model list for. Skip — there is - // nothing to expose. - continue - } - factoryInstances := instancesByProviderName[p.ProviderName] - if len(factoryInstances) == 0 { - continue - } + modelRankMap := make(map[string]int) // key: "providerName@modelName" + factoryRankMapping := make(map[string]int) + for i := range providerManager.Providers { + factory := &providerManager.Providers[i] + factoryRankMapping[factory.Name] = -factory.Rank for _, llm := range factory.Models { - if modelTypeFilter != "" { - match := false - for _, t := range llm.ModelTypes { - if strings.EqualFold(t, modelTypeFilter) { - match = true + rank := 500 + if llm.Rank != nil { + rank = *llm.Rank + } + modelRankMap[factory.Name+"@"+llm.Name] = rank + } + } + + added := make([]map[string]interface{}, 0, len(targetRecords)) + for _, modelRecord := range targetRecords { + provInfo, provOK := providerInfoByID[modelRecord.ProviderID] + instInfo, instOK := instanceInfoByID[modelRecord.InstanceID] + if !provOK || !instOK { + continue + } + rank := modelRankMap[provInfo.ProviderName+"@"+modelRecord.ModelName] + if rank == 0 { + rank = 500 + } + added = append(added, map[string]interface{}{ + "model_id": modelRecord.ID, + "tenant_id": provInfo.TenantID, + "tenant_name": tenantName, + "model_type": entity.ModelType(modelRecord.ModelType).HumanReadable(), + "name": modelRecord.ModelName, + "provider_id": modelRecord.ProviderID, + "provider_name": provInfo.ProviderName, + "instance_id": modelRecord.InstanceID, + "instance_name": instInfo.InstanceName, + "rank": rank, + }) + } + + // Add TEI Builtin embedding model if configured (mirrors Python). + composeProfiles := common.GetEnv(common.EnvComposeProfiles) + teiModel := common.GetEnv(common.EnvTEIModel) + if strings.Contains(composeProfiles, "tei-") && teiModel != "" { + if modelTypeFilter == "" || modelTypeFilter == "embedding" { + teiAlreadyAdded := false + for _, m := range added { + if pn, _ := m["provider_name"].(string); pn == "Builtin" { + if n, _ := m["name"].(string); n == teiModel { + teiAlreadyAdded = true break } } - if !match { - continue - } } - for _, inst := range factoryInstances { - key := p.ID + "@" + inst.ID + "@" + llm.Name - // Bitmask merge: factory types ∪ active overrides \ inactive overrides. - merged := entity.ModelTypeFromStrings(llm.ModelTypes) - merged |= entity.ModelType(activeByKey[key]) - merged &^= entity.ModelType(inactiveByKey[key]) - if merged == 0 { - continue + if !teiAlreadyAdded { + rank := modelRankMap["Builtin@"+teiModel] + if rank == 0 { + rank = 500 } added = append(added, map[string]interface{}{ - "model_id": activeModelIDByKey[key], + "model_id": "", "tenant_id": tenant.ID, "tenant_name": tenantName, - "model_type": modelTypesForAPI(merged), - "name": llm.Name, - "provider_id": inst.ProviderID, - "provider_name": p.ProviderName, - "instance_id": inst.ID, - "instance_name": inst.InstanceName, + "model_type": []string{"embedding"}, + "name": teiModel, + "provider_id": "", + "provider_name": "Builtin", + "instance_id": "", + "instance_name": "default", + "rank": rank, }) } } } - return added, common.CodeSuccess, nil -} - -func modelTypesForAPI(modelType entity.ModelType) []string { - types := modelType.HumanReadable() - for i, t := range types { - if t == "image2text" { - types[i] = "vision" + // Sort by factory rank desc, provider_name, instance_name, model rank desc, name. + sort.Slice(added, func(i, j int) bool { + pi, _ := added[i]["provider_name"].(string) + pj, _ := added[j]["provider_name"].(string) + fi := factoryRankMapping[pi] + fj := factoryRankMapping[pj] + if fi != fj { + return fi < fj // factory rank: smaller (more negative) = higher priority } - } - return types + if pi != pj { + return pi < pj + } + ini, _ := added[i]["instance_name"].(string) + inj, _ := added[j]["instance_name"].(string) + if ini != inj { + return ini < inj + } + ri, _ := added[i]["rank"].(int) + rj, _ := added[j]["rank"].(int) + if ri != rj { + return ri > rj // model rank desc + } + ni, _ := added[i]["name"].(string) + nj, _ := added[j]["name"].(string) + return ni < nj + }) + + return added, common.CodeSuccess, nil } func (m *ModelProviderService) resolveModelListTenant(userID, ownerTenantID string) (*entity.Tenant, common.ErrorCode, error) { @@ -1235,8 +1312,18 @@ func (m *ModelProviderService) resolveModelListTenant(userID, ownerTenantID stri return tenant, common.CodeSuccess, nil } -func (m *ModelProviderService) AlterProviderInstance(userID, providerName, instanceIDOrName, newInstanceName, apiKey, baseURL, region string, modelInfo []CreateInstanceModelInfo, verify bool) (common.ErrorCode, error) { - providerName = strings.TrimSpace(providerName) +// ensureMineruFromEnv is a stub that mirrors Python's ensure_mineru_from_env. +// It ensures a MinerU OCR provider instance exists when env vars are configured. +func (m *ModelProviderService) ensureMineruFromEnv(tenantID string) error { return nil } + +// ensurePaddleOCREnabledFromEnv is a stub that mirrors Python's ensure_paddleocr_from_env. +func (m *ModelProviderService) ensurePaddleOCREnabledFromEnv(tenantID string) error { return nil } + +// ensureOpenDataLoaderFromEnv is a stub that mirrors Python's ensure_opendataloader_from_env. +func (m *ModelProviderService) ensureOpenDataLoaderFromEnv(tenantID string) error { return nil } + +func (m *ModelProviderService) AlterProviderInstance(userID, providerIDOrName, instanceIDOrName, newInstanceName, apiKey, baseURL, region string, modelInfo []CreateInstanceModelInfo, verify bool) (common.ErrorCode, error) { + providerIDOrName = strings.TrimSpace(providerIDOrName) tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") if err != nil { @@ -1247,10 +1334,11 @@ func (m *ModelProviderService) AlterProviderInstance(userID, providerName, insta } tenantID := tenants[0].TenantID - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) if err != nil { - return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) + return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) } + providerName := provider.ProviderName // Find the instance — try by ID first, then by name. instance, err := m.modelInstanceDAO.GetByID(instanceIDOrName) @@ -1404,27 +1492,8 @@ func (m *ModelProviderService) AlterProviderInstance(userID, providerName, insta return common.CodeSuccess, nil } -func (m *ModelProviderService) getProviderInstancesByNames(providerID string, instanceNames []string) ([]*entity.TenantModelInstance, []string, error) { - modelInstances := make([]*entity.TenantModelInstance, 0, len(instanceNames)) - missingInstanceNames := make([]string, 0) - - for _, instanceName := range instanceNames { - modelInstance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(providerID, instanceName) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - missingInstanceNames = append(missingInstanceNames, instanceName) - continue - } - return nil, nil, err - } - modelInstances = append(modelInstances, modelInstance) - } - - return modelInstances, missingInstanceNames, nil -} - -func (m *ModelProviderService) DropProviderInstances(providerName, userID string, instanceNames []string) (common.ErrorCode, error) { - if len(instanceNames) == 0 { +func (m *ModelProviderService) DropProviderInstances(providerIDOrName, userID string, instanceIDOrNames []string) (common.ErrorCode, error) { + if len(instanceIDOrNames) == 0 { return common.CodeBadRequest, errors.New("instances is required") } @@ -1440,42 +1509,46 @@ func (m *ModelProviderService) DropProviderInstances(providerName, userID string tenantID := tenants[0].TenantID - // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + // Find provider by ID or name + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return common.CodeNotFound, fmt.Errorf("no provider found for provider %q", providerName) + return common.CodeNotFound, fmt.Errorf("no provider found for provider %q", providerIDOrName) } return common.CodeServerError, err } - modelInstances, missingInstanceNames, err := m.getProviderInstancesByNames(provider.ID, instanceNames) - if err != nil { - return common.CodeServerError, err - } - if len(missingInstanceNames) > 0 { - return common.CodeNotFound, fmt.Errorf("no instance found for provider %q and instances %q", providerName, missingInstanceNames) - } - - for _, modelInstance := range modelInstances { - if _, err = m.modelDAO.DeleteByProviderIDAndInstanceID(provider.ID, modelInstance.ID); err != nil { - return common.CodeServerError, err - } - - count, err := m.modelInstanceDAO.DeleteByProviderIDAndInstanceName(provider.ID, modelInstance.InstanceName) + // Delete each instance by provider_id + instance_name. + // Since instance names are unique within a provider, we can delete directly. + notExistInstances := make([]string, 0) + for _, instanceName := range instanceIDOrNames { + instanceObj, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName) if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + notExistInstances = append(notExistInstances, instanceName) + continue + } return common.CodeServerError, err } - if count == 0 { - return common.CodeNotFound, fmt.Errorf("provider instance %q not found", modelInstance.InstanceName) + + // Delete models under this instance first + if _, err := m.modelDAO.DeleteByProviderIDAndInstanceID(provider.ID, instanceObj.ID); err != nil { + return common.CodeServerError, err } + + // Delete the instance itself + if _, err := m.modelInstanceDAO.DeleteByProviderIDAndInstanceName(provider.ID, instanceName); err != nil { + return common.CodeServerError, err + } + } + if len(notExistInstances) > 0 { + return common.CodeNotFound, fmt.Errorf("no instance found for provider %q and instances %q", providerIDOrName, notExistInstances) } return common.CodeSuccess, nil } -func (m *ModelProviderService) DropInstanceModels(providerName, instanceName, userID string, modelIDs, models []string) (common.ErrorCode, error) { - +func (m *ModelProviderService) DropInstanceModels(providerIDOrName, instanceIDOrName, userID string, modelNames []string) (common.ErrorCode, error) { // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") if err != nil { @@ -1488,51 +1561,54 @@ func (m *ModelProviderService) DropInstanceModels(providerName, instanceName, us tenantID := tenants[0].TenantID - // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + // Get provider by ID or name (matches Python's get_by_tenant_id_and_provider_id → get_by_tenant_id_and_provider_name fallback). + provider, err := m.getProviderByIDOrName(tenantID, providerIDOrName) + if err != nil { + return common.CodeDataError, fmt.Errorf("No provider found for provider '%s'", providerIDOrName) + } + + // Get instance by ID or name (matches Python's get_by_id → get_by_provider_id_and_instance_name fallback). + instance, err := m.modelInstanceDAO.GetByID(instanceIDOrName) + if err != nil || instance.ProviderID != provider.ID { + instance, err = m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceIDOrName) + } + if err != nil { + return common.CodeDataError, fmt.Errorf("No instance found for provider '%s' and instance '%s'", providerIDOrName, instanceIDOrName) + } + + // Fetch all models under this instance and validate all requested names exist. + modelObjs, err := m.modelDAO.GetModelsByInstanceID(instance.ID) if err != nil { return common.CodeServerError, err } - var modelInstance *entity.TenantModelInstance - modelInstance, err = m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName) - if err != nil { + existingNames := make(map[string]struct{}, len(modelObjs)) + for _, obj := range modelObjs { + existingNames[obj.ModelName] = struct{}{} + } + + notExist := make([]string, 0) + for _, name := range modelNames { + if _, ok := existingNames[name]; !ok { + notExist = append(notExist, name) + } + } + if len(notExist) > 0 { + return common.CodeNotFound, fmt.Errorf("Models %v not found for provider '%s' and instance '%s'", notExist, providerIDOrName, instanceIDOrName) + } + + // Collect IDs of models to delete. + idsToDelete := make([]string, 0, len(modelNames)) + for _, obj := range modelObjs { + if _, ok := existingNames[obj.ModelName]; ok { + idsToDelete = append(idsToDelete, obj.ID) + } + } + + if _, err := m.modelDAO.DeleteByIDs(idsToDelete); err != nil { return common.CodeServerError, err } - for _, modelID := range modelIDs { - modelID = strings.TrimSpace(modelID) - if modelID == "" { - return common.CodeBadRequest, errors.New("model ID is required") - } - var count int64 = 0 - count, err = m.modelDAO.DeleteByModelIDAndProviderIDAndInstanceID(modelID, provider.ID, modelInstance.ID) - if err != nil { - return common.CodeServerError, err - } - - if count == 0 { - return common.CodeNotFound, fmt.Errorf("model %s not found", modelID) - } - } - - for _, modelName := range models { - modelName = strings.TrimSpace(modelName) - if modelName == "" { - return common.CodeBadRequest, errors.New("model name is required") - } - // Delete all models of this instance - var count int64 = 0 - count, err = m.modelDAO.DeleteByProviderIDAndInstanceIDAndModelName(provider.ID, modelInstance.ID, modelName) - if err != nil { - return common.CodeServerError, err - } - - if count == 0 { - return common.CodeNotFound, fmt.Errorf("model: %s not found", modelName) - } - } - return common.CodeSuccess, nil } @@ -1549,64 +1625,89 @@ func (m *ModelProviderService) ListInstanceModels(providerName, instanceName, us tenantID := tenants[0].TenantID - // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + // Find provider by ID or name + provider, err := m.getProviderByIDOrName(tenantID, providerName) if err != nil { return nil, err } - // Get instance - instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName) - if err != nil { - return nil, err + // Find instance by ID first, then by name + instance, err := m.modelInstanceDAO.GetByID(instanceName) + if err != nil || instance.ProviderID != provider.ID { + instance, err = m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName) + if err != nil { + return nil, err + } } // Get all models for this instance - disabledModels, err := m.modelDAO.GetModelsByInstanceID(instance.ID) + modelObjs, err := m.modelDAO.GetModelsByInstanceID(instance.ID) if err != nil { return nil, err } - allModels, err := dao.GetModelProviderManager().ListModels(providerName) + modelList := make([]map[string]interface{}, 0, len(modelObjs)) + for _, model := range modelObjs { + modelExtra := make(map[string]interface{}) + if model.Extra != "" { + _ = json.Unmarshal([]byte(model.Extra), &modelExtra) + } - // insert models name into a set - modelNames := make(map[string]bool) - for _, model := range disabledModels { - if model.Status == "active" { - modelData := map[string]interface{}{ - "name": model.ModelName, + maxTokens := 8192 + if v, ok := modelExtra["max_tokens"]; ok { + switch val := v.(type) { + case float64: + maxTokens = int(val) + case int: + maxTokens = val } - allModels = append(allModels, modelData) - } else { - modelNames[model.ModelName] = true } - } - - for _, model := range allModels { - // convert model["name"] to string - modelName := model["name"].(string) - if modelNames[modelName] { - model["status"] = "inactive" - } else { - model["status"] = "active" + verify := "unknown" + if v, ok := modelExtra["verify"].(string); ok { + verify = v } + + features := make([]string, 0) + if v, ok := modelExtra["is_tools"]; ok && isTruthy(v) { + features = append(features, "is_tools") + } + if v, ok := modelExtra["thinking"]; ok && isTruthy(v) { + features = append(features, "thinking") + } + + modelList = append(modelList, map[string]interface{}{ + "name": model.ModelName, + "model_type": entity.ModelType(model.ModelType).HumanReadable(), + "max_tokens": maxTokens, + "status": model.Status, + "verify": verify, + "features": features, + }) } - return allModels, nil + // Sort by name + sort.Slice(modelList, func(i, j int) bool { + return modelList[i]["name"].(string) < modelList[j]["name"].(string) + }) + + return modelList, nil } -func (m *ModelProviderService) UpdateModelStatus(providerName, instanceName, modelName, userID, modelID, status string) (common.ErrorCode, error) { +func (m *ModelProviderService) AlterModel(providerName, instanceName, modelName, userID, modelID string, updateDict map[string]interface{}) (common.ErrorCode, error) { modelName = strings.TrimSpace(modelName) modelID = strings.TrimSpace(modelID) - status = strings.TrimSpace(status) - if status != "active" && status != "inactive" { - return common.CodeBadRequest, errors.New("status must be active or inactive") - } if modelName == "" && modelID == "" { return common.CodeBadRequest, errors.New("model name or model ID is required") } + // Validate status early, before any DB lookup. + if status, ok := updateDict["status"].(string); ok && status != "" { + if status != "active" && status != "inactive" { + return common.CodeBadRequest, errors.New("status must be 'active' or 'inactive'") + } + } + // Get tenant ID from user tenants, err := m.userTenantDAO.GetByUserIDAndRole(userID, "owner") if err != nil { @@ -1620,7 +1721,7 @@ func (m *ModelProviderService) UpdateModelStatus(providerName, instanceName, mod tenantID := tenants[0].TenantID // Check if provider exists - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) + provider, err := m.getProviderByIDOrName(tenantID, providerName) if err != nil { return common.CodeServerError, err } @@ -1631,75 +1732,112 @@ func (m *ModelProviderService) UpdateModelStatus(providerName, instanceName, mod } var model *entity.TenantModel - - if modelID != "" { - model, err = m.modelDAO.GetByID(modelID) - if err != nil || model == nil { - return common.CodeNotFound, errors.New("model not found") - } - - if model.ProviderID != provider.ID || model.InstanceID != instance.ID { - return common.CodeNotFound, errors.New("model not found") - } - - if modelName != "" && model.ModelName != modelName { - return common.CodeBadRequest, errors.New("model ID does not match model name") - } - - count, err := m.modelDAO.UpdateStatusByIDAndScope(modelID, provider.ID, instance.ID, status) - if err != nil { - return common.CodeServerError, err - } - if count == 0 { - return common.CodeNotFound, errors.New("model not found") - } - return common.CodeSuccess, nil - } else { + if modelName != "" { model, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName) if err != nil { if !errors.Is(err, gorm.ErrRecordNotFound) { return common.CodeServerError, err } - - modelID = utility.GenerateToken() - - var modelSchema *modelModule.Model - modelSchema, err = dao.GetModelProviderManager().GetModelByName(providerName, modelName) - if err != nil { - return common.CodeNotFound, errors.New(fmt.Sprintf("provider %s model %s not found", providerName, modelName)) + return common.CodeNotFound, fmt.Errorf("model %q not found for provider %q and instance %q", modelName, providerName, instanceName) + } + if modelID != "" && model.ID != modelID { + return common.CodeBadRequest, errors.New("model ID does not match model name") + } + } else { + model, err = m.modelDAO.GetByID(modelID) + if err != nil { + if !errors.Is(err, gorm.ErrRecordNotFound) { + return common.CodeServerError, err } - - // Get model info from provider - if len(modelSchema.ModelTypes) == 0 { - return common.CodeServerError, fmt.Errorf("provider %s model %s has no model types", providerName, modelName) - } - model = &entity.TenantModel{ - ID: modelID, - ModelName: modelName, - ModelType: int(entity.ModelTypeFromString(modelSchema.ModelTypes[0])), - ProviderID: provider.ID, - InstanceID: instance.ID, - Status: status, - } - err = m.modelDAO.Create(model) - if err != nil { - return common.CodeServerError, errors.New("fail to create model") - } - return common.CodeSuccess, nil + return common.CodeNotFound, fmt.Errorf("model with ID %q not found", modelID) + } + if model.ProviderID != provider.ID || model.InstanceID != instance.ID { + return common.CodeNotFound, fmt.Errorf("model with ID %q does not belong to provider %q and instance %q", modelID, providerName, instanceName) } } - count, err := m.modelDAO.UpdateStatusByIDAndScope(model.ID, provider.ID, instance.ID, status) - if err != nil { - return common.CodeServerError, err + toUpdate := make(map[string]interface{}) + + // Handle status (validation already done above, here we only apply the change) + if status, ok := updateDict["status"].(string); ok && status != "" { + if status != model.Status { + toUpdate["status"] = status + } } - if count == 0 { - return common.CodeNotFound, errors.New("model not found") + + // Handle model_type + if modelTypeVal, ok := updateDict["model_type"]; ok && modelTypeVal != nil { + targetModelType := resolveModelType(modelTypeVal) + if targetModelType != model.ModelType { + toUpdate["model_type"] = targetModelType + } + } + + // Handle extra (max_tokens, extra) + newExtra := make(map[string]interface{}) + if extraVal, ok := updateDict["extra"]; ok && extraVal != nil { + if extraMap, ok := extraVal.(map[string]interface{}); ok { + for k, v := range extraMap { + newExtra[k] = v + } + } + } + if maxTokensVal, ok := updateDict["max_tokens"]; ok { + switch v := maxTokensVal.(type) { + case int: + if v > 0 { + newExtra["max_tokens"] = v + } + case float64: + if v > 0 { + newExtra["max_tokens"] = int(v) + } + } + } + if len(newExtra) > 0 { + dbExtra := make(map[string]interface{}) + if model.Extra != "" { + _ = json.Unmarshal([]byte(model.Extra), &dbExtra) + } + for k, v := range newExtra { + dbExtra[k] = v + } + extraBytes, err := json.Marshal(dbExtra) + if err != nil { + return common.CodeServerError, err + } + toUpdate["extra"] = string(extraBytes) + } + + if len(toUpdate) > 0 { + if err := m.modelDAO.UpdateByID(model.ID, toUpdate); err != nil { + return common.CodeServerError, err + } } return common.CodeSuccess, nil } +// resolveModelType converts a model_type value (string, []string, or float64) to an int bitmask. +func resolveModelType(val interface{}) int { + switch v := val.(type) { + case string: + return int(entity.ModelTypeFromString(v)) + case []interface{}: + types := make([]string, 0, len(v)) + for _, t := range v { + if s, ok := t.(string); ok { + types = append(types, s) + } + } + return int(entity.ModelTypeFromStrings(types)) + case float64: + return int(v) + default: + return 0 + } +} + type ModelInstanceAndProviderInfo struct { ProviderEntity *entity.TenantModelProvider ProviderInfo *modelModule.Provider @@ -2569,9 +2707,12 @@ func (m *ModelProviderService) GetRerankModel(tenantID, compositeModelName strin } type AddModelRequest struct { - ProviderName string `json:"provider_name"` - InstanceName string `json:"instance_name"` - Models []ModelRequest `json:"models"` + ProviderName string `json:"provider_name"` + InstanceName string `json:"instance_name"` + ModelName string `json:"model_name"` + ModelTypes []string `json:"model_type"` + MaxTokens int `json:"max_tokens"` + Extra map[string]interface{} `json:"extra"` } func (m *ModelProviderService) GetTenantDefaultModelByType(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { @@ -2880,8 +3021,14 @@ func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) if request == nil { return common.CodeBadRequest, errors.New("request is required") } - if len(request.Models) == 0 { - return common.CodeBadRequest, errors.New("models is required") + + modelName := strings.TrimSpace(request.ModelName) + if modelName == "" { + 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(userID, "owner") @@ -2894,105 +3041,92 @@ func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) tenantID := tenants[0].TenantID - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, request.ProviderName) + // Get provider by ID or name (matches Python's get_by_tenant_id_and_provider_id → get_by_tenant_id_and_provider_name fallback). + provider, err := m.getProviderByIDOrName(tenantID, request.ProviderName) if err != nil { - return common.CodeServerError, err + return common.CodeDataError, fmt.Errorf("No provider found for provider '%s'", request.ProviderName) } - instance, err := m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, request.InstanceName) + // Get instance by ID or name (matches Python's get_by_id → get_by_provider_id_and_instance_name fallback). + instance, err := m.modelInstanceDAO.GetByID(request.InstanceName) + if err != nil || instance.ProviderID != provider.ID { + instance, err = m.modelInstanceDAO.GetByProviderIDAndInstanceName(provider.ID, request.InstanceName) + } if err != nil { + return common.CodeDataError, fmt.Errorf("No instance found for provider '%s' and instance '%s'", request.ProviderName, request.InstanceName) + } + + // Check for duplicate model. + _, err = m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(provider.ID, instance.ID, modelName) + if err == nil { + return common.CodeConflict, fmt.Errorf("Model '%s' already exists for provider '%s' and instance '%s'", modelName, request.ProviderName, request.InstanceName) + } + if !errors.Is(err, gorm.ErrRecordNotFound) { return common.CodeServerError, err } - 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") + // Compute model type bitmask. + combinedType := entity.ModelType(0) + for _, rawType := range request.ModelTypes { + mt := strings.TrimSpace(rawType) + if mt == "" { + continue } - - if modelName == "" { - return common.CodeBadRequest, errors.New("model name is required") + t := entity.ModelTypeFromString(mt) + if t == 0 { + return common.CodeBadRequest, fmt.Errorf("invalid model type: %s", mt) } - - // Validate each model type string and compute the combined bitmask. - combinedType := entity.ModelType(0) - seenTypes := make(map[string]struct{}, len(model.ModelTypes)) - validTypes := make([]string, 0, len(model.ModelTypes)) - for _, rawType := range model.ModelTypes { - mt := strings.TrimSpace(rawType) - if mt == "" { - return common.CodeBadRequest, errors.New("model type is required") - } - if _, ok := seenTypes[mt]; ok { - continue // skip duplicates - } - seenTypes[mt] = struct{}{} - validTypes = append(validTypes, mt) - t := entity.ModelTypeFromString(mt) - if t == 0 { - return common.CodeBadRequest, fmt.Errorf("invalid model type: %s", mt) - } - combinedType |= t - } - - 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 := utility.GenerateToken() - - if model.MaxDimension < 0 { - return common.CodeBadRequest, errors.New("max_dimension must be non-negative") - } - for _, dimension := range model.Dimensions { - if dimension <= 0 { - return common.CodeBadRequest, errors.New("dimensions must contain positive values") - } - if model.MaxDimension > 0 && dimension > model.MaxDimension { - return common.CodeBadRequest, fmt.Errorf("dimension %d exceeds max_dimension %d", dimension, model.MaxDimension) - } - } - extra := map[string]interface{}{ - "max_tokens": model.MaxTokens, - "max_dimension": model.MaxDimension, - "dimensions": model.Dimensions, - } - if model.Thinking != nil { - extra["thinking"] = *model.Thinking - } - - var extraByte []byte - 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: int(combinedType), - ProviderID: provider.ID, - InstanceID: instance.ID, - Status: "active", - Extra: string(extraByte), - }) + combinedType |= t } - if err = m.modelDAO.CreateBatch(models); err != nil { - return common.CodeServerError, err + maxTokens := request.MaxTokens + if maxTokens <= 0 { + maxTokens = 8192 + } + + // Build extra fields. Mirrors Python's add_model_to_instance. + extraFields := map[string]interface{}{ + "max_tokens": maxTokens, + } + + // Look up factory info to populate is_tools and thinking (matches Python). + targetProvider := dao.GetModelProviderManager().FindProvider(provider.ProviderName) + if targetProvider != nil { + targetModel := dao.GetModelProviderManager().FindModel(targetProvider, modelName) + if targetModel != nil { + if targetModel.Tools != nil && targetModel.Tools.Support { + extraFields["is_tools"] = true + } + if targetModel.Thinking != nil { + extraFields["thinking"] = true + } + } + } + + if request.Extra != nil { + for k, v := range request.Extra { + extraFields[k] = v + } + } + + extraBytes, err := json.Marshal(extraFields) + if err != nil { + return common.CodeServerError, errors.New("fail to marshal extra") + } + + modelID := utility.GenerateToken() + tenantModel := &entity.TenantModel{ + ID: modelID, + ModelName: modelName, + ModelType: int(combinedType), + ProviderID: provider.ID, + InstanceID: instance.ID, + Status: "active", + Extra: string(extraBytes), + } + + if err := m.modelDAO.Create(tenantModel); err != nil { + return common.CodeServerError, fmt.Errorf("fail to create model '%s': %s", modelName, err.Error()) } return common.CodeSuccess, nil diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index 2839ad67d1..0f031667e4 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -155,14 +155,14 @@ func seedModelProviderServiceScope(t *testing.T, db *gorm.DB) { } } -func TestModelProviderServiceUpdateModelStatusByID(t *testing.T) { +func TestModelProviderServiceAlterModelStatusByID(t *testing.T) { db := setupModelProviderServiceTestDB(t) useModelProviderServiceTestDB(t, db) seedModelProviderServiceScope(t, db) - code, err := NewModelProviderService().UpdateModelStatus("OpenAI", "default", "", "user-1", "model-1", "inactive") + code, err := NewModelProviderService().AlterModel("OpenAI", "default", "", "user-1", "model-1", map[string]interface{}{"status": "inactive"}) if err != nil { - t.Fatalf("UpdateModelStatus() error = %v", err) + t.Fatalf("AlterModel() error = %v", err) } if code != common.CodeSuccess { t.Fatalf("code = %v, want %v", code, common.CodeSuccess) @@ -197,23 +197,23 @@ func TestModelProviderServiceGetModelConfigByID(t *testing.T) { } } -func TestModelProviderServiceUpdateModelStatusRejectsInvalidStatus(t *testing.T) { - code, err := NewModelProviderService().UpdateModelStatus("OpenAI", "default", "", "user-1", "model-1", "disabled") +func TestModelProviderServiceAlterModelRejectsInvalidStatus(t *testing.T) { + code, err := NewModelProviderService().AlterModel("OpenAI", "default", "", "user-1", "model-1", map[string]interface{}{"status": "disabled"}) if err == nil { - t.Fatalf("UpdateModelStatus() error = nil, want invalid status error") + t.Fatalf("AlterModel() error = nil, want invalid status error") } if code != common.CodeBadRequest { t.Fatalf("code = %v, want %v", code, common.CodeBadRequest) } - if !strings.Contains(err.Error(), "status must be active or inactive") { + if !strings.Contains(err.Error(), "status must be") { t.Fatalf("error = %v, want status validation message", err) } } -func TestModelProviderServiceUpdateModelStatusRejectsMissingModelSelector(t *testing.T) { - code, err := NewModelProviderService().UpdateModelStatus("OpenAI", "default", "", "user-1", "", "active") +func TestModelProviderServiceAlterModelRejectsMissingModelSelector(t *testing.T) { + code, err := NewModelProviderService().AlterModel("OpenAI", "default", "", "user-1", "", map[string]interface{}{"status": "active"}) if err == nil { - t.Fatalf("UpdateModelStatus() error = nil, want missing model selector error") + t.Fatalf("AlterModel() error = nil, want missing model selector error") } if code != common.CodeBadRequest { t.Fatalf("code = %v, want %v", code, common.CodeBadRequest) @@ -223,7 +223,7 @@ func TestModelProviderServiceUpdateModelStatusRejectsMissingModelSelector(t *tes } } -func TestModelProviderServiceUpdateModelStatusRejectsWrongScopedModelID(t *testing.T) { +func TestModelProviderServiceAlterModelRejectsWrongScopedModelID(t *testing.T) { db := setupModelProviderServiceTestDB(t) useModelProviderServiceTestDB(t, db) seedModelProviderServiceScope(t, db) @@ -231,9 +231,9 @@ func TestModelProviderServiceUpdateModelStatusRejectsWrongScopedModelID(t *testi t.Fatalf("failed to seed second instance: %v", err) } - code, err := NewModelProviderService().UpdateModelStatus("OpenAI", "other", "", "user-1", "model-1", "inactive") + code, err := NewModelProviderService().AlterModel("OpenAI", "other", "", "user-1", "model-1", map[string]interface{}{"status": "inactive"}) if err == nil { - t.Fatalf("UpdateModelStatus() error = nil, want not found error") + t.Fatalf("AlterModel() error = nil, want not found error") } if code != common.CodeNotFound { t.Fatalf("code = %v, want %v", code, common.CodeNotFound) diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index 026c3c83e3..b156dc6774 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -358,6 +358,7 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID if result.Final { finalContent := strings.TrimSpace(result.Answer) fullContent = finalContent + finalReference = []FormattedChunk{} if ref, ok := result.Reference["chunks"]; ok { if chunks, ok := ref.([]map[string]interface{}); ok { finalReference = formatChunks(chunks) @@ -397,6 +398,7 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID } if finalReference == nil && openaiReq.NeedReference { + finalReference = []FormattedChunk{} if ref, ok := lastResult.Reference["chunks"]; ok { if chunks, ok := ref.([]map[string]interface{}); ok { finalReference = formatChunks(chunks) @@ -441,6 +443,7 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID TotalTokens: promptTokens + completionTokens, } if openaiReq.NeedReference { + resp.Reference = []FormattedChunk{} if ref, ok := finalResult.Reference["chunks"]; ok { if chunks, ok := ref.([]map[string]interface{}); ok { resp.Reference = formatChunks(chunks) diff --git a/internal/service/tenant.go b/internal/service/tenant.go index fafcd17bbc..fda50f4066 100644 --- a/internal/service/tenant.go +++ b/internal/service/tenant.go @@ -489,71 +489,103 @@ func (s *TenantService) GetDefaultModelName(tenantID string, modelType entity.Mo } func (s *TenantService) GetModelInfo(tenantID string, defaultModel string, modelType string) (*string, *string, *string, bool, error) { - // normally the model string is: modelName@instanceName@providerName, sometimes it's just modelName@providerName - // for the 1st case, parse defaultChatModel into three parts - defaultChatModelParts := strings.Split(defaultModel, "@") - var providerName *string - var instanceName *string - var modelName *string - if len(defaultChatModelParts) >= 3 { - providerName = &defaultChatModelParts[len(defaultChatModelParts)-1] - instanceName = &defaultChatModelParts[len(defaultChatModelParts)-2] - joinedModelName := strings.Join(defaultChatModelParts[:len(defaultChatModelParts)-2], "@") - modelName = &joinedModelName - - } else if len(defaultChatModelParts) == 2 { - providerName = &defaultChatModelParts[1] - instanceName = new(string) - *instanceName = "default" - modelName = &defaultChatModelParts[0] - } else { - return nil, nil, nil, false, fmt.Errorf("invalid model string: %s", defaultModel) + // Mirror Python's _get_model_info: right-anchored rsplit so that model + // names containing '@' (e.g. LM Studio IDs like + // "text-embedding-nomic-embed-text-v1.5@q8_0") remain intact. + // The composite key is: modelName@instanceName@providerName or + // modelName@providerName. + parts := rsplitN(defaultModel, "@", 2) + var modelName, instanceName, providerName string + switch len(parts) { + case 3: + modelName, instanceName, providerName = parts[0], parts[1], parts[2] + case 2: + modelName, providerName = parts[0], parts[1] + instanceName = "default" + default: + modelName = parts[0] + providerName = "" + instanceName = "default" } - if modelType == "ocr" { - if *providerName == "infiniflow" && *instanceName == "default" && *modelName == "deepdoc" { - return providerName, instanceName, modelName, true, nil - } + // Special case: OCR with infiniflow@default@deepdoc is always enabled. + if modelType == "ocr" && providerName == "infiniflow" && instanceName == "default" && modelName == "deepdoc" { + return &providerName, &instanceName, &modelName, true, nil } - // Check if the provider and instance exists - modelProvider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, *providerName) + // Special case: TEI Builtin embedding model. + composeProfiles := common.GetEnv(common.EnvComposeProfiles) + teiModel := common.GetEnv(common.EnvTEIModel) + if modelType == "embedding" && strings.Contains(composeProfiles, "tei-") && teiModel != "" && + modelName == teiModel && (providerName == "" || providerName == "Builtin") { + return &providerName, &instanceName, &modelName, true, nil + } + + // Check if the provider exists for the tenant. + modelProvider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(tenantID, providerName) if err != nil { return nil, nil, nil, false, err } - modelInstance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(modelProvider.ID, *instanceName) + // Check if the instance exists. + modelInstance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(modelProvider.ID, instanceName) if err != nil { return nil, nil, nil, false, err } - modelSchema, err := dao.GetModelProviderManager().GetModelByName(*providerName, *modelName) + // Validate that the factory model supports this model type. + modelSchema, err := dao.GetModelProviderManager().GetModelByName(providerName, modelName) if err != nil { return nil, nil, nil, false, err } - factoryModelType := factoryModelTypeName(modelType) if !modelSchema.ModelTypeMap[factoryModelType] { - return nil, nil, nil, false, fmt.Errorf("model %s isn't a %s model", *modelName, modelType) + return nil, nil, nil, false, fmt.Errorf("model %s isn't a %s model", modelName, modelType) } - var modelEntity *entity.TenantModel - modelEntity, err = s.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(modelProvider.ID, modelInstance.ID, *modelName) + // Check if the model exists and is active. + modelEntity, err := s.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(modelProvider.ID, modelInstance.ID, modelName) if err != nil { - errString := err.Error() - if !strings.Contains(errString, "record not found") { + if !dao.IsNotFoundErr(err) { return nil, nil, nil, false, err } } if modelEntity == nil { - return nil, nil, nil, false, fmt.Errorf("model %s isn't available", *modelName) + return nil, nil, nil, false, fmt.Errorf("model %s isn't available", modelName) } if modelEntity.Status != "active" { - return nil, nil, nil, false, fmt.Errorf("model %s isn't available", *modelName) + return nil, nil, nil, false, fmt.Errorf("model %s isn't available", modelName) } - return providerName, instanceName, modelName, true, nil + return &providerName, &instanceName, &modelName, true, nil +} +// rsplitN splits s by sep from the right, limiting to n+1 parts (mirrors +// Python's str.rsplit(sep, maxsplit)). +func rsplitN(s, sep string, n int) []string { + if n <= 0 { + return []string{s} + } + result := make([]string, 0, n+1) + remaining := s + for i := 0; i < n; i++ { + idx := strings.LastIndex(remaining, sep) + if idx < 0 { + result = append(result, remaining) + // Reverse the collected parts. + for left, right := 0, len(result)-1; left < right; left, right = left+1, right-1 { + result[left], result[right] = result[right], result[left] + } + return result + } + result = append(result, remaining[idx+len(sep):]) + remaining = remaining[:idx] + } + result = append(result, remaining) + for left, right := 0, len(result)-1; left < right; left, right = left+1, right-1 { + result[left], result[right] = result[right], result[left] + } + return result } func (s *TenantService) ListTenantDefaultModels(userID string) ([]ModelItem, error) { diff --git a/test/testcases/restful_api/test_openai_compatible.py b/test/testcases/restful_api/test_openai_compatible.py index b3dbbbd319..efbbc7656c 100644 --- a/test/testcases/restful_api/test_openai_compatible.py +++ b/test/testcases/restful_api/test_openai_compatible.py @@ -241,4 +241,5 @@ def test_openai_compatible_reference_metadata_fields_filter_accepts_array(rest_c choice_msg = payload["choices"][0]["message"] skip_if_go_proxy_upstream_error(choice_msg) assert "reference" in choice_msg, payload + print(payload) assert isinstance(choice_msg["reference"], list), payload diff --git a/test/unit_test/api/apps/services/test_delete_datasets.py b/test/unit_test/api/apps/services/test_delete_datasets.py index 3e036023a4..b85b4f0fa4 100644 --- a/test/unit_test/api/apps/services/test_delete_datasets.py +++ b/test/unit_test/api/apps/services/test_delete_datasets.py @@ -27,14 +27,14 @@ import pytest pytestmark = pytest.mark.p2 -class _StubModelTypeBinary(IntEnum): - CHAT = 1 - EMBEDDING = 2 - ASR = 4 - VISION = 8 - RERANK = 16 - TTS = 32 - OCR = 64 +class _StubModelTypeBinary(IntEnum): + CHAT = 1 + EMBEDDING = 2 + ASR = 4 + VISION = 8 + RERANK = 16 + TTS = 32 + OCR = 64 def _stub(monkeypatch, name, **attrs):