diff --git a/internal/admin/service.go b/internal/admin/service.go index d4df67bc5..c3a76e370 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -449,10 +449,10 @@ func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) { // Determine API key and base URL based on model type var apiKey, apiBase string switch llm.ModelType { - case string(entity.ModelTypeChat): + case entity.ModelTypeChat.String(): apiKey = factoryConfig.APIKey apiBase = factoryConfig.BaseURL - case string(entity.ModelTypeEmbedding): + case entity.ModelTypeEmbedding.String(): apiKey = cfg.UserDefaultLLM.DefaultModels.EmbeddingModel.APIKey apiBase = cfg.UserDefaultLLM.DefaultModels.EmbeddingModel.BaseURL if apiKey == "" { @@ -461,7 +461,7 @@ func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) { if apiBase == "" { apiBase = factoryConfig.BaseURL } - case string(entity.ModelTypeRerank): + case entity.ModelTypeRerank.String(): apiKey = cfg.UserDefaultLLM.DefaultModels.RerankModel.APIKey apiBase = cfg.UserDefaultLLM.DefaultModels.RerankModel.BaseURL if apiKey == "" { @@ -470,7 +470,7 @@ func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) { if apiBase == "" { apiBase = factoryConfig.BaseURL } - case string(entity.ModelTypeSpeech2Text): + case entity.ModelTypeSpeech2Text.String(): apiKey = cfg.UserDefaultLLM.DefaultModels.ASRModel.APIKey apiBase = cfg.UserDefaultLLM.DefaultModels.ASRModel.BaseURL if apiKey == "" { @@ -479,7 +479,7 @@ func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) { if apiBase == "" { apiBase = factoryConfig.BaseURL } - case string(entity.ModelTypeImage2Text): + case entity.ModelTypeImage2Text.String(): apiKey = cfg.UserDefaultLLM.DefaultModels.Image2TextModel.APIKey apiBase = cfg.UserDefaultLLM.DefaultModels.Image2TextModel.BaseURL if apiKey == "" { diff --git a/internal/dao/tenant_llm.go b/internal/dao/tenant_llm.go index 63ef0eece..4325039c9 100644 --- a/internal/dao/tenant_llm.go +++ b/internal/dao/tenant_llm.go @@ -52,7 +52,8 @@ func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string, // GetByTenantNameAndType get tenant LLM by tenant ID, model name, and model type func (dao *TenantLLMDAO) GetByTenantNameAndType(tenantID, modelName string, modelType entity.ModelType) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM - err := DB.Where("tenant_id = ? AND llm_name = ? AND model_type = ?", tenantID, modelName, modelType).First(&tenantLLM).Error + // tenant_llm.model_type is a VARCHAR column, so convert to string. + err := DB.Where("tenant_id = ? AND llm_name = ? AND model_type = ?", tenantID, modelName, modelType.String()).First(&tenantLLM).Error if err != nil { return nil, err } @@ -62,7 +63,8 @@ func (dao *TenantLLMDAO) GetByTenantNameAndType(tenantID, modelName string, mode // GetByTenantAndType get tenant LLM by tenant ID and model type func (dao *TenantLLMDAO) GetByTenantAndType(tenantID string, modelType entity.ModelType) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM - err := DB.Where("tenant_id = ? AND model_type = ?", tenantID, modelType).First(&tenantLLM).Error + // tenant_llm.model_type is a VARCHAR column, so convert to string. + err := DB.Where("tenant_id = ? AND model_type = ?", tenantID, modelType.String()).First(&tenantLLM).Error if err != nil { return nil, err } @@ -72,7 +74,8 @@ func (dao *TenantLLMDAO) GetByTenantAndType(tenantID string, modelType entity.Mo // GetByTenantAndFactory get tenant LLM by tenant ID, model type and factory func (dao *TenantLLMDAO) GetByTenantAndFactory(tenantID string, modelType entity.ModelType, factory string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM - err := DB.Where("tenant_id = ? AND model_type = ? AND llm_factory = ?", tenantID, modelType, factory).First(&tenantLLM).Error + // tenant_llm.model_type is a VARCHAR column, so convert to string. + err := DB.Where("tenant_id = ? AND model_type = ? AND llm_factory = ?", tenantID, modelType.String(), factory).First(&tenantLLM).Error if err != nil { return nil, err } diff --git a/internal/dao/tenant_model.go b/internal/dao/tenant_model.go index d9572e4db..2fb235ef9 100644 --- a/internal/dao/tenant_model.go +++ b/internal/dao/tenant_model.go @@ -102,9 +102,11 @@ func (dao *TenantModelDAO) GetModelsByProviderIDAndInstanceIDAndModelName(provid return models, nil } -func (dao *TenantModelDAO) GetByProviderIDAndInstanceIDAndModelTypeAndModelName(providerID, instanceID, modelType, modelName string) (*entity.TenantModel, error) { +func (dao *TenantModelDAO) GetByProviderIDAndInstanceIDAndModelTypeAndModelName(providerID, instanceID string, modelType int, modelName string) (*entity.TenantModel, error) { var model entity.TenantModel - err := DB.Where("provider_id = ? AND instance_id = ? AND model_type = ? AND model_name = ?", providerID, instanceID, modelType, modelName).First(&model).Error + // Use bitwise AND to match Python's bin_and(model_type) > 0 pattern. + // A model_type value of 0 (unknown type) matches no row. + err := DB.Where("provider_id = ? AND instance_id = ? AND model_type & ? > 0 AND model_name = ?", providerID, instanceID, modelType, modelName).First(&model).Error if err != nil { return nil, err } diff --git a/internal/dao/tenant_model_test.go b/internal/dao/tenant_model_test.go index 07b33c5a7..45f1125c4 100644 --- a/internal/dao/tenant_model_test.go +++ b/internal/dao/tenant_model_test.go @@ -55,8 +55,8 @@ func TestTenantModelDAODeleteByModelIDAndScopeDeletesOnlyMatchingModel(t *testin db := setupTenantModelDAOTestDB(t) useTenantModelDAOTestDB(t, db) - seedTenantModel(t, db, &entity.TenantModel{ID: "model-delete", ModelName: "m", ModelType: "chat", ProviderID: "provider-1", InstanceID: "instance-1", Status: "active"}) - seedTenantModel(t, db, &entity.TenantModel{ID: "model-keep", ModelName: "m", ModelType: "chat", ProviderID: "provider-1", InstanceID: "instance-2", Status: "active"}) + seedTenantModel(t, db, &entity.TenantModel{ID: "model-delete", ModelName: "m", ModelType: int(entity.ModelTypeChat), ProviderID: "provider-1", InstanceID: "instance-1", Status: "active"}) + seedTenantModel(t, db, &entity.TenantModel{ID: "model-keep", ModelName: "m", ModelType: int(entity.ModelTypeChat), ProviderID: "provider-1", InstanceID: "instance-2", Status: "active"}) rows, err := NewTenantModelDAO().DeleteByModelIDAndProviderIDAndInstanceID("model-delete", "provider-1", "instance-1") if err != nil { @@ -85,7 +85,7 @@ func TestTenantModelDAOUpdateStatusByIDAndScope(t *testing.T) { db := setupTenantModelDAOTestDB(t) useTenantModelDAOTestDB(t, db) - seedTenantModel(t, db, &entity.TenantModel{ID: "model-status", ModelName: "m", ModelType: "chat", ProviderID: "provider-1", InstanceID: "instance-1", Status: "active"}) + seedTenantModel(t, db, &entity.TenantModel{ID: "model-status", ModelName: "m", ModelType: int(entity.ModelTypeChat), ProviderID: "provider-1", InstanceID: "instance-1", Status: "active"}) rows, err := NewTenantModelDAO().UpdateStatusByIDAndScope("model-status", "provider-1", "instance-1", "inactive") if err != nil { diff --git a/internal/entity/chat.go b/internal/entity/chat.go index eeb6b2639..b6ae04fd8 100644 --- a/internal/entity/chat.go +++ b/internal/entity/chat.go @@ -27,7 +27,7 @@ type Chat struct { Icon *string `gorm:"column:icon;type:longtext" json:"icon,omitempty"` Language *string `gorm:"column:language;size:32;index" json:"language,omitempty"` LLMID string `gorm:"column:llm_id;size:128;not null" json:"llm_id"` - TenantLLMID *int64 `gorm:"column:tenant_llm_id;index" json:"tenant_llm_id,omitempty"` + TenantLLMID *string `gorm:"column:tenant_llm_id;size:32;index" json:"tenant_llm_id,omitempty"` LLMSetting JSONMap `gorm:"column:llm_setting;type:longtext;not null" json:"llm_setting"` PromptType string `gorm:"column:prompt_type;size:16;not null;default:'simple';index" json:"prompt_type"` PromptConfig JSONMap `gorm:"column:prompt_config;type:longtext;not null" json:"prompt_config"` @@ -38,7 +38,7 @@ type Chat struct { TopK int64 `gorm:"column:top_k;default:1024" json:"top_k"` DoRefer string `gorm:"column:do_refer;size:1;not null;default:1" json:"do_refer"` RerankID string `gorm:"column:rerank_id;size:128;not null;default:''" json:"rerank_id"` - TenantRerankID *int64 `gorm:"column:tenant_rerank_id;index" json:"tenant_rerank_id,omitempty"` + TenantRerankID *string `gorm:"column:tenant_rerank_id;size:32;index" json:"tenant_rerank_id,omitempty"` KBIDs JSONSlice `gorm:"column:kb_ids;type:longtext;not null" json:"kb_ids"` Status *string `gorm:"column:status;size:1;index" json:"status,omitempty"` BaseModel diff --git a/internal/entity/kb.go b/internal/entity/kb.go index 9424e8581..7a0fe6500 100644 --- a/internal/entity/kb.go +++ b/internal/entity/kb.go @@ -104,7 +104,7 @@ type Knowledgebase struct { Language *string `gorm:"column:language;size:32;index" json:"language,omitempty"` Description *string `gorm:"column:description;type:longtext" json:"description,omitempty"` EmbdID string `gorm:"column:embd_id;size:128;not null;index" json:"embd_id"` - TenantEmbdID *int64 `gorm:"column:tenant_embd_id;index" json:"tenant_embd_id,omitempty"` + TenantEmbdID *string `gorm:"column:tenant_embd_id;size:32;index" json:"tenant_embd_id,omitempty"` Permission string `gorm:"column:permission;size:16;not null;default:me;index" json:"permission"` CreatedBy string `gorm:"column:created_by;size:32;not null;index" json:"created_by"` DocNum int64 `gorm:"column:doc_num;default:0;index" json:"doc_num"` diff --git a/internal/entity/memory.go b/internal/entity/memory.go index 48f4e55a8..3ee294f9a 100644 --- a/internal/entity/memory.go +++ b/internal/entity/memory.go @@ -25,9 +25,9 @@ type Memory struct { MemoryType int64 `gorm:"column:memory_type;default:1;index" json:"memory_type"` StorageType string `gorm:"column:storage_type;size:32;not null;default:table;index" json:"storage_type"` EmbdID string `gorm:"column:embd_id;size:128;not null" json:"embd_id"` - TenantEmbdID *int64 `gorm:"column:tenant_embd_id;index" json:"tenant_embd_id,omitempty"` + TenantEmbdID *string `gorm:"column:tenant_embd_id;size:32;index" json:"tenant_embd_id,omitempty"` LLMID string `gorm:"column:llm_id;size:128;not null" json:"llm_id"` - TenantLLMID *int64 `gorm:"column:tenant_llm_id;index" json:"tenant_llm_id,omitempty"` + TenantLLMID *string `gorm:"column:tenant_llm_id;size:32;index" json:"tenant_llm_id,omitempty"` Permissions string `gorm:"column:permissions;size:16;not null;default:me;index" json:"permissions"` Description *string `gorm:"column:description;type:longtext" json:"description,omitempty"` MemorySize int64 `gorm:"column:memory_size;default:5242880;not null" json:"memory_size"` diff --git a/internal/entity/skill_search.go b/internal/entity/skill_search.go index fb191f188..5b71fcf69 100644 --- a/internal/entity/skill_search.go +++ b/internal/entity/skill_search.go @@ -53,7 +53,7 @@ type SkillSearchConfig struct { SimilarityThreshold float64 `gorm:"column:similarity_threshold;default:0.2" json:"similarity_threshold"` FieldConfig JSONMap `gorm:"column:field_config;type:json" json:"field_config"` RerankID *string `gorm:"column:rerank_id;size:128" json:"rerank_id,omitempty"` - TenantRerankID *int64 `gorm:"column:tenant_rerank_id" json:"tenant_rerank_id,omitempty"` + TenantRerankID *string `gorm:"column:tenant_rerank_id;size:32" json:"tenant_rerank_id,omitempty"` TopK int64 `gorm:"column:top_k;default:10" json:"top_k"` IndexVersion string `gorm:"column:index_version;size:32;default:'1.0.0'" json:"index_version"` BaseModel diff --git a/internal/entity/tenant.go b/internal/entity/tenant.go index c17d74469..9cbd324b6 100644 --- a/internal/entity/tenant.go +++ b/internal/entity/tenant.go @@ -22,20 +22,20 @@ type Tenant struct { Name *string `gorm:"column:name;size:100;index" json:"name,omitempty"` PublicKey *string `gorm:"column:public_key;size:255;index" json:"public_key,omitempty"` LLMID string `gorm:"column:llm_id;size:128;not null;index" json:"llm_id"` - TenantLLMID *int64 `gorm:"column:tenant_llm_id;index" json:"tenant_llm_id,omitempty"` + TenantLLMID *string `gorm:"column:tenant_llm_id;size:32;index" json:"tenant_llm_id,omitempty"` EmbdID string `gorm:"column:embd_id;size:128;not null;index" json:"embd_id"` - TenantEmbdID *int64 `gorm:"column:tenant_embd_id;index" json:"tenant_embd_id,omitempty"` + TenantEmbdID *string `gorm:"column:tenant_embd_id;size:32;index" json:"tenant_embd_id,omitempty"` ASRID string `gorm:"column:asr_id;size:128;not null;index" json:"asr_id"` - TenantASRID *int64 `gorm:"column:tenant_asr_id;index" json:"tenant_asr_id,omitempty"` + TenantASRID *string `gorm:"column:tenant_asr_id;size:32;index" json:"tenant_asr_id,omitempty"` Img2TxtID string `gorm:"column:img2txt_id;size:128;not null;index" json:"img2txt_id"` - TenantImg2TxtID *int64 `gorm:"column:tenant_img2txt_id;index" json:"tenant_img2txt_id,omitempty"` + TenantImg2TxtID *string `gorm:"column:tenant_img2txt_id;size:32;index" json:"tenant_img2txt_id,omitempty"` RerankID string `gorm:"column:rerank_id;size:128;not null;index" json:"rerank_id"` - TenantRerankID *int64 `gorm:"column:tenant_rerank_id;index" json:"tenant_rerank_id,omitempty"` + TenantRerankID *string `gorm:"column:tenant_rerank_id;size:32;index" json:"tenant_rerank_id,omitempty"` TTSID string `gorm:"column:tts_id;size:256;index" json:"tts_id,omitempty"` - TenantTTSID *int64 `gorm:"column:tenant_tts_id;index" json:"tenant_tts_id,omitempty"` + TenantTTSID *string `gorm:"column:tenant_tts_id;size:32;index" json:"tenant_tts_id,omitempty"` ParserIDs string `gorm:"column:parser_ids;size:256;not null;index" json:"parser_ids"` OCRID string `gorm:"column:ocr_id;size:256" json:"ocr_id,omitempty"` - TenantOCRID *int64 `gorm:"column:tenant_ocr_id" json:"tenant_ocr_id,omitempty"` + TenantOCRID *string `gorm:"column:tenant_ocr_id;size:32" json:"tenant_ocr_id,omitempty"` Credit int64 `gorm:"column:credit;default:512;index" json:"credit"` Status *string `gorm:"column:status;size:1;index" json:"status,omitempty"` BaseModel diff --git a/internal/entity/tenant_model.go b/internal/entity/tenant_model.go index 7a34ed195..4593fee9b 100644 --- a/internal/entity/tenant_model.go +++ b/internal/entity/tenant_model.go @@ -22,7 +22,7 @@ type TenantModel struct { ModelName string `gorm:"column:model_name;size:128" json:"model_name"` ProviderID string `gorm:"column:provider_id;size:32;not null" json:"provider_id"` InstanceID string `gorm:"column:instance_id;size:32;not null;index" json:"instance_id"` - ModelType string `gorm:"column:model_type;size:32;not null" json:"model_type"` + ModelType int `gorm:"column:model_type;size:32;not null" json:"model_type"` Status string `gorm:"column:status;size:32;default:'active'" json:"status"` Extra string `gorm:"column:extra;size:1024;default:'{}'" json:"extra"` BaseModel diff --git a/internal/entity/types.go b/internal/entity/types.go index 21ef40cae..c408d37ff 100644 --- a/internal/entity/types.go +++ b/internal/entity/types.go @@ -16,22 +16,126 @@ package entity -// ModelType represents the type of model -type ModelType string +import "strings" + +// ModelType represents the type of model as a bitmask, matching Python's +// ModelTypeBinary enum (common/constants.py). +// Bit flags (LSB->MSB): 1=chat, 2=embedding, 4=speech2text, 8=image2text, 16=rerank, 32=tts, 64=ocr. +// A model can belong to multiple types simultaneously. +type ModelType int const ( - // ModelTypeChat chat model - ModelTypeChat ModelType = "chat" - // ModelTypeEmbedding embedding model - ModelTypeEmbedding ModelType = "embedding" - // ModelTypeSpeech2Text speech to text model - ModelTypeSpeech2Text ModelType = "speech2text" - // ModelTypeImage2Text image to text model - ModelTypeImage2Text ModelType = "image2text" - // ModelTypeRerank rerank model - ModelTypeRerank ModelType = "rerank" - // ModelTypeTTS text to speech model - ModelTypeTTS ModelType = "tts" - // ModelTypeOCR optical character recognition model - ModelTypeOCR ModelType = "ocr" + // ModelTypeChat chat model (1 << 0) + ModelTypeChat ModelType = 1 << iota + // ModelTypeEmbedding embedding model (1 << 1) + ModelTypeEmbedding + // ModelTypeSpeech2Text speech to text model (1 << 2) + ModelTypeSpeech2Text + // ModelTypeImage2Text image to text model (1 << 3) + ModelTypeImage2Text + // ModelTypeRerank rerank model (1 << 4) + ModelTypeRerank + // ModelTypeTTS text to speech model (1 << 5) + ModelTypeTTS + // ModelTypeOCR optical character recognition model (1 << 6) + ModelTypeOCR ) + +// Has returns true if the bitmask contains the given model type. +func (mt ModelType) Has(target ModelType) bool { + return int(mt)&int(target) != 0 +} + +// String returns the string representation of the model type. +// For a single-bit value returns the corresponding name. +// For multi-bit values returns comma-separated names (e.g. "chat,embedding"). +func (mt ModelType) String() string { + if mt == 0 { + return "" + } + parts := make([]string, 0, 4) + if mt.Has(ModelTypeChat) { + parts = append(parts, "chat") + } + if mt.Has(ModelTypeEmbedding) { + parts = append(parts, "embedding") + } + if mt.Has(ModelTypeSpeech2Text) { + parts = append(parts, "speech2text") + } + if mt.Has(ModelTypeImage2Text) { + parts = append(parts, "image2text") + } + if mt.Has(ModelTypeRerank) { + parts = append(parts, "rerank") + } + if mt.Has(ModelTypeTTS) { + parts = append(parts, "tts") + } + if mt.Has(ModelTypeOCR) { + parts = append(parts, "ocr") + } + return strings.Join(parts, ",") +} + +// HumanReadable returns the string names of all model types in this bitmask. +func (mt ModelType) HumanReadable() []string { + if mt == 0 { + return nil + } + parts := make([]string, 0, 4) + if mt.Has(ModelTypeChat) { + parts = append(parts, "chat") + } + if mt.Has(ModelTypeEmbedding) { + parts = append(parts, "embedding") + } + if mt.Has(ModelTypeSpeech2Text) { + parts = append(parts, "speech2text") + } + if mt.Has(ModelTypeImage2Text) { + parts = append(parts, "image2text") + } + if mt.Has(ModelTypeRerank) { + parts = append(parts, "rerank") + } + if mt.Has(ModelTypeTTS) { + parts = append(parts, "tts") + } + if mt.Has(ModelTypeOCR) { + parts = append(parts, "ocr") + } + return parts +} + +// ModelTypeFromString converts a string to a ModelType. +func ModelTypeFromString(s string) ModelType { + switch s { + case "chat": + return ModelTypeChat + case "embedding": + return ModelTypeEmbedding + case "speech2text": + return ModelTypeSpeech2Text + case "image2text": + return ModelTypeImage2Text + case "rerank": + return ModelTypeRerank + case "tts": + return ModelTypeTTS + case "ocr": + return ModelTypeOCR + default: + return 0 + } +} + +// ModelTypeFromStrings converts multiple strings to a combined bitmask ModelType. +// This matches Python's calculate_model_type() which ORs together multiple type values. +func ModelTypeFromStrings(types []string) ModelType { + var result ModelType + for _, t := range types { + result |= ModelTypeFromString(t) + } + return result +} diff --git a/internal/handler/providers_test.go b/internal/handler/providers_test.go index 8f65e293f..58e523f21 100644 --- a/internal/handler/providers_test.go +++ b/internal/handler/providers_test.go @@ -48,7 +48,7 @@ func seedProviderHandlerModel(t *testing.T, db *gorm.DB) { &entity.UserTenant{ID: "user-tenant-1", UserID: "user-1", TenantID: "tenant-1", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, &entity.TenantModelProvider{ID: "provider-1", TenantID: "tenant-1", ProviderName: "OpenAI"}, &entity.TenantModelInstance{ID: "instance-1", ProviderID: "provider-1", InstanceName: "default", APIKey: "sk-test", Status: "active", Extra: "{}"}, - &entity.TenantModel{ID: "model-1", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: "chat", Status: "active"}, + &entity.TenantModel{ID: "model-1", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: int(entity.ModelTypeChat), Status: "active"}, } for _, row := range rows { if err := db.Create(row).Error; err != nil { diff --git a/internal/ingestion/component/pdf_vision_dispatch.go b/internal/ingestion/component/pdf_vision_dispatch.go index 860d8c318..d8fe713ba 100644 --- a/internal/ingestion/component/pdf_vision_dispatch.go +++ b/internal/ingestion/component/pdf_vision_dispatch.go @@ -332,7 +332,7 @@ func resolveModelConfigFromProviderInstance(tenantID string, modelType entity.Mo baseURL := extra["base_url"] modelObj, modelErr := modelDAO.GetByProviderIDAndInstanceIDAndModelTypeAndModelName( - provider.ID, instance.ID, string(modelType), pureModelName, + provider.ID, instance.ID, int(modelType), pureModelName, ) switch { case modelErr == nil: diff --git a/internal/service/chat.go b/internal/service/chat.go index e304c7d60..f8ab3cd9f 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -306,19 +306,19 @@ func validateCreateLLMID(llmID, tenantID string, llmSetting map[string]interface modelType := entity.ModelTypeChat switch confModelType := llmSetting["model_type"].(type) { case string: - if confModelType == string(entity.ModelTypeImage2Text) { + if confModelType == entity.ModelTypeImage2Text.String() { modelType = entity.ModelTypeImage2Text } case []interface{}: for _, item := range confModelType { - if item == string(entity.ModelTypeImage2Text) { + if item == entity.ModelTypeImage2Text.String() { modelType = entity.ModelTypeImage2Text break } } case []string: for _, item := range confModelType { - if item == string(entity.ModelTypeImage2Text) { + if item == entity.ModelTypeImage2Text.String() { modelType = entity.ModelTypeImage2Text break } @@ -924,12 +924,12 @@ func (s *ChatService) validateRESTLLMID(llmID, tenantID string, llmSetting map[s if rawModelType, ok := llmSetting["model_type"]; ok { switch typedModelType := rawModelType.(type) { case string: - if typedModelType == string(entity.ModelTypeImage2Text) { + if typedModelType == entity.ModelTypeImage2Text.String() { modelType = entity.ModelTypeImage2Text } case []interface{}: for _, item := range typedModelType { - if fmt.Sprint(item) == string(entity.ModelTypeImage2Text) { + if fmt.Sprint(item) == entity.ModelTypeImage2Text.String() { modelType = entity.ModelTypeImage2Text break } diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index bdb451926..1bed09756 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -308,7 +308,7 @@ func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID s tenantSvc := service.NewTenantService() modelProviderSvc := service.NewModelProviderService() var err error - llmModelName, err = tenantSvc.GetDefaultModelName(tenantIDs[0], "chat") + llmModelName, err = tenantSvc.GetDefaultModelName(tenantIDs[0], entity.ModelTypeChat) if err != nil || llmModelName == "" { common.Warn("Failed to get default chat model name for LLM transformations", zap.Error(err)) } else { @@ -361,12 +361,8 @@ func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID s modelProviderSvc := service.NewModelProviderService() var embeddingModel *models.EmbeddingModel var embdID string - if kbRecords[0].TenantEmbdID != nil && *kbRecords[0].TenantEmbdID > 0 { - _, embdID, err = dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), *kbRecords[0].TenantEmbdID) - if err != nil { - return nil, fmt.Errorf("failed to get embedding model by tenant_embd_id: %w", err) - } - driver, modelName, apiConfig, maxTokens, getErr := modelProviderSvc.GetModelConfigFromProviderInstance(tenantIDs[0], entity.ModelTypeEmbedding, embdID) + if kbRecords[0].TenantEmbdID != nil && *kbRecords[0].TenantEmbdID != "" { + driver, modelName, apiConfig, maxTokens, getErr := modelProviderSvc.GetModelConfigByID(tenantIDs[0], *kbRecords[0].TenantEmbdID) if getErr != nil { return nil, fmt.Errorf("failed to get embedding model by tenant_embd_id: %w", getErr) } @@ -404,37 +400,20 @@ func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID s // Get rerank model if RerankID is specified var rerankModel *models.RerankModel - var rerankCompositeName string if req.TenantRerankID != nil && *req.TenantRerankID != "" { - tenantRerankIDInt, parseErr := strconv.ParseInt(*req.TenantRerankID, 10, 64) - if parseErr != nil { - return nil, fmt.Errorf("invalid tenant_rerank_id: %w", parseErr) - } - _, rerankCompositeName, err = dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), tenantRerankIDInt) - if err != nil { - return nil, fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", err) - } - } else if req.RerankID != nil && *req.RerankID != "" { - rerankCompositeName = *req.RerankID - if _, _, _, _, getErr := modelProviderSvc.GetModelConfigFromProviderInstance(tenantIDs[0], entity.ModelTypeRerank, rerankCompositeName); getErr != nil { - _, rerankCompositeName, err = dao.LookupTenantLLMByName(dao.NewTenantLLMDAO(), tenantIDs[0], *req.RerankID, entity.ModelTypeRerank) - if err != nil { - return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", getErr) - } - } - } - if rerankCompositeName != "" { - driver, mdlName, apiConfig, _, getErr := modelProviderSvc.GetModelConfigFromProviderInstance(tenantIDs[0], entity.ModelTypeRerank, rerankCompositeName) + driver, mdlName, apiConfig, _, getErr := modelProviderSvc.GetModelConfigByID(tenantIDs[0], *req.TenantRerankID) if getErr != nil { - return nil, fmt.Errorf("failed to get rerank model: %w", getErr) + return nil, fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", getErr) } rerankModel = models.NewRerankModel(driver, &mdlName, apiConfig) - } - - if rerankModel != nil { - common.Info("Fetched rerank model", - zap.String("tenantID", tenantIDs[0]), - zap.String("rerankCompositeName", rerankCompositeName)) + } else if req.RerankID != nil && *req.RerankID != "" { + rerankCompositeName := *req.RerankID + driver, mdlName, apiConfig, _, getErr := modelProviderSvc.GetModelConfigFromProviderInstance(tenantIDs[0], entity.ModelTypeRerank, rerankCompositeName) + if getErr != nil { + rerankModel = nil + } else { + rerankModel = models.NewRerankModel(driver, &mdlName, apiConfig) + } } retrievalReq := &nlp.RetrievalRequest{ @@ -484,8 +463,8 @@ func (s *ChunkService) RetrievalTest(req *service.RetrievalTestRequest, userID s } func knowledgebaseEmbeddingKey(kb *entity.Knowledgebase, tenantID string) string { - if kb.TenantEmbdID != nil && *kb.TenantEmbdID > 0 { - return fmt.Sprintf("tenant:%d", *kb.TenantEmbdID) + if kb.TenantEmbdID != nil && *kb.TenantEmbdID != "" { + return fmt.Sprintf("tenant:%s", *kb.TenantEmbdID) } if kb.EmbdID == "" { return fmt.Sprintf("default:%s", tenantID) diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index a18e62485..49a6f2114 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -78,7 +78,7 @@ func TestHydrateChunkVectors_NoDim(t *testing.T) { } func TestKnowledgebaseEmbeddingKey(t *testing.T) { - tenantEmbdID := int64(42) + tenantEmbdID := "42" tests := []struct { name string @@ -111,7 +111,7 @@ func TestKnowledgebaseEmbeddingKey(t *testing.T) { name: "ignores non-positive tenant embedding id", kb: &entity.Knowledgebase{ EmbdID: "shared-model", - TenantEmbdID: new(int64), + TenantEmbdID: func() *string { s := ""; return &s }(), }, want: "embd:shared-model", }, diff --git a/internal/service/dataset_update_test.go b/internal/service/dataset_update_test.go index a3910c7e9..911afd57e 100644 --- a/internal/service/dataset_update_test.go +++ b/internal/service/dataset_update_test.go @@ -271,7 +271,7 @@ func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing. insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") insertDatasetUpdateModelProvider(t, "provider-1", "tenant-1", "ZHIPU-AI") insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test") - insertDatasetUpdateTenantModel(t, "model-1", "provider-1", "instance-1", "embedding-2", string(entity.ModelTypeEmbedding)) + insertDatasetUpdateTenantModel(t, "model-1", "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding)) embeddingModel := "embedding-2@test@ZHIPU-AI" result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", UpdateDatasetRequest{ @@ -408,7 +408,7 @@ func insertDatasetUpdateModelInstance(t *testing.T, id, providerID, instanceName } } -func insertDatasetUpdateTenantModel(t *testing.T, id, providerID, instanceID, modelName, modelType string) { +func insertDatasetUpdateTenantModel(t *testing.T, id, providerID, instanceID, modelName string, modelType int) { t.Helper() model := &entity.TenantModel{ diff --git a/internal/service/memory.go b/internal/service/memory.go index 51250113f..ae194921b 100644 --- a/internal/service/memory.go +++ b/internal/service/memory.go @@ -414,16 +414,12 @@ func (s *MemoryService) CreateMemory(tenantID string, req *CreateMemoryRequest) SystemPrompt: &systemPrompt, } - // Convert tenant model IDs from string to int64 for database + // Attach tenant model IDs (string form) to memory if req.TenantEmbdID != nil { - if embdID, err := strconv.ParseInt(*req.TenantEmbdID, 10, 64); err == nil { - memory.TenantEmbdID = &embdID - } + memory.TenantEmbdID = req.TenantEmbdID } if req.TenantLLMID != nil { - if llmID, err := strconv.ParseInt(*req.TenantLLMID, 10, 64); err == nil { - memory.TenantLLMID = &llmID - } + memory.TenantLLMID = req.TenantLLMID } if err := s.memoryDAO.Create(memory); err != nil { return nil, errors.New("could not create new memory") @@ -488,15 +484,11 @@ func (s *MemoryService) UpdateMemory(tenantID string, memoryID string, req *Upda } if req.TenantLLMID != nil { - if llmID, err := strconv.ParseInt(*req.TenantLLMID, 10, 64); err == nil { - updateDict["tenant_llm_id"] = llmID - } + updateDict["tenant_llm_id"] = *req.TenantLLMID } if req.TenantEmbdID != nil { - if embdID, err := strconv.ParseInt(*req.TenantEmbdID, 10, 64); err == nil { - updateDict["tenant_embd_id"] = embdID - } + updateDict["tenant_embd_id"] = *req.TenantEmbdID } if req.MemoryType != nil && len(req.MemoryType) > 0 { @@ -600,11 +592,11 @@ func (s *MemoryService) UpdateMemory(tenantID string, memoryID string, req *Upda filteredUpdateDict[field] = value } case "tenant_llm_id": - if currentMemory.TenantLLMID == nil || *currentMemory.TenantLLMID != value.(int64) { + if currentMemory.TenantLLMID == nil || *currentMemory.TenantLLMID != fmt.Sprint(value) { filteredUpdateDict[field] = value } case "tenant_embd_id": - if currentMemory.TenantEmbdID == nil || *currentMemory.TenantEmbdID != value.(int64) { + if currentMemory.TenantEmbdID == nil || *currentMemory.TenantEmbdID != fmt.Sprint(value) { filteredUpdateDict[field] = value } case "memory_type": diff --git a/internal/service/model_service.go b/internal/service/model_service.go index b3fa5e356..0b4f0c90a 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -868,14 +868,14 @@ func (m *ModelProviderService) ListTenantAddedModels(userID, modelTypeFilter str if err != nil { return nil, common.CodeServerError, err } - activeByKey := make(map[string][]string) - inactiveByKey := make(map[string][]string) + activeByKey := make(map[string]int) + inactiveByKey := make(map[string]int) for _, rec := range modelRecords { key := rec.ProviderID + "@" + rec.InstanceID + "@" + rec.ModelName if rec.Status == "inactive" { - inactiveByKey[key] = append(inactiveByKey[key], rec.ModelType) + inactiveByKey[key] |= rec.ModelType } else { - activeByKey[key] = append(activeByKey[key], rec.ModelType) + activeByKey[key] |= rec.ModelType } } @@ -923,26 +923,15 @@ func (m *ModelProviderService) ListTenantAddedModels(userID, modelTypeFilter str } for _, inst := range factoryInstances { key := p.ID + "@" + inst.ID + "@" + llm.Name - // Set-based merge: factory types ∪ active overrides \ inactive overrides. - mergedSet := make(map[string]struct{}, len(llm.ModelTypes)+len(activeByKey[key])) - for _, t := range llm.ModelTypes { - mergedSet[t] = struct{}{} - } - for _, t := range activeByKey[key] { - mergedSet[t] = struct{}{} - } - for _, t := range inactiveByKey[key] { - delete(mergedSet, t) - } - if len(mergedSet) == 0 { + // 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 } - merged := make([]string, 0, len(mergedSet)) - for t := range mergedSet { - merged = append(merged, t) - } added = append(added, map[string]interface{}{ - "model_type": merged, + "model_type": merged.HumanReadable(), "name": llm.Name, "provider_id": inst.ProviderID, "provider_name": p.ProviderName, @@ -1232,7 +1221,7 @@ func (m *ModelProviderService) UpdateModelStatus(providerName, instanceName, mod model = &entity.TenantModel{ ID: modelID, ModelName: modelName, - ModelType: modelSchema.ModelTypes[0], + ModelType: int(entity.ModelTypeFromString(modelSchema.ModelTypes[0])), ProviderID: provider.ID, InstanceID: instance.ID, Status: status, @@ -1529,7 +1518,7 @@ func (m *ModelProviderService) ChatToModelWithMessages(providerName, instanceNam } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "chat" && info.ModelEntity.ModelType != "vision" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeChat) && !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeImage2Text) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a chat or multimodal model", resolvedModelName, resolvedProviderName)) } @@ -1592,7 +1581,7 @@ func (m *ModelProviderService) ChatToModelStreamWithSender(providerName, instanc } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "chat" && info.ModelEntity.ModelType != "vision" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeChat) && !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeImage2Text) { return common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a chat or multimodal model", resolvedModelName, resolvedProviderName)) } @@ -1679,7 +1668,7 @@ func (m *ModelProviderService) EmbedText(providerName, instanceName, modelName, } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "embedding" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeEmbedding) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is an embedding model", resolvedModelName, resolvedProviderName)) } @@ -1745,7 +1734,7 @@ func (m *ModelProviderService) RerankDocument(providerName, instanceName, modelN } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "rerank" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeRerank) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a rerank model", resolvedModelName, resolvedProviderName)) } @@ -1799,7 +1788,7 @@ func (m *ModelProviderService) TranscribeAudio(providerName, instanceName, model } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "asr" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeSpeech2Text) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is an ASR model", *modelName, *providerName)) } @@ -1856,7 +1845,7 @@ func (m *ModelProviderService) TranscribeAudioStream(providerName, instanceName, } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "asr" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeSpeech2Text) { return common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is an ASR model", *modelName, *providerName)) } @@ -1909,7 +1898,7 @@ func (m *ModelProviderService) AudioSpeech(providerName, instanceName, modelName } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "tts" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeTTS) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a TTS model", *modelName, *providerName)) } @@ -1965,7 +1954,7 @@ func (m *ModelProviderService) AudioSpeechStream(providerName, instanceName, mod } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "tts" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeTTS) { return common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a TTS model", *modelName, *providerName)) } @@ -2017,7 +2006,7 @@ func (m *ModelProviderService) OCRFile(providerName, instanceName, modelName, mo } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "ocr" { + if !entity.ModelType(info.ModelEntity.ModelType).Has(entity.ModelTypeOCR) { return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is an OCR model", *modelName, *providerName)) } @@ -2073,9 +2062,8 @@ func (m *ModelProviderService) ParseFile(providerName, instanceName, modelName, } else { // model entity exists if info.ModelEntity.Status == "active" { - if info.ModelEntity.ModelType != "doc_parse" { - return nil, common.CodeNotFound, errors.New(fmt.Sprintf("expect model %s@%s is a ParseFile model", *modelName, *providerName)) - } + // Note: ParseFile model type is not in the ModelType enum; skip entity-type check + // and rely on the factory ModelTypeMap check in the nil-entity branch. modelDriver, err = newModelDriverForBaseURL(info.ProviderInfo.ModelDriver, *providerName, *info.APIConfig.Region, *info.APIConfig.BaseURL) if err != nil { @@ -2148,6 +2136,91 @@ func (m *ModelProviderService) GetTenantDefaultModelByType(tenantID string, mode return m.GetModelConfigFromProviderInstance(tenantID, modelType, modelName) } +// GetModelConfigByID returns model driver and API config for a tenant_model row by its ID. +func (m *ModelProviderService) GetModelConfigByID(userID, modelID string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + common.Debug("GetModelConfigByID", + zap.String("userID", userID), + zap.String("modelID", modelID)) + + modelEntity, err := m.modelDAO.GetByID(modelID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, "", nil, 0, fmt.Errorf("tenant model id=%s not found", modelID) + } + return nil, "", nil, 0, err + } + if modelEntity.Status != "active" { + return nil, "", nil, 0, fmt.Errorf("tenant model id=%s is disabled", modelID) + } + + providerEntity, err := m.modelProviderDAO.GetByID(modelEntity.ProviderID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, "", nil, 0, fmt.Errorf("provider id=%s not found for model id=%s", modelEntity.ProviderID, modelID) + } + return nil, "", nil, 0, err + } + if providerEntity == nil { + return nil, "", nil, 0, fmt.Errorf("provider id=%s not found for model id=%s", modelEntity.ProviderID, modelID) + } + + if providerEntity.TenantID != userID { + userTenants, terr := NewUserTenantService().GetUserTenantRelationByUserID(userID) + if terr != nil { + return nil, "", nil, 0, terr + } + allowed := false + for _, rel := range userTenants { + if rel != nil && rel.TenantID == providerEntity.TenantID { + allowed = true + break + } + } + if !allowed { + return nil, "", nil, 0, fmt.Errorf("tenant %s has no access to provider owned by tenant %s", userID, providerEntity.TenantID) + } + } + + instanceEntity, err := m.modelInstanceDAO.GetByID(modelEntity.InstanceID) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, "", nil, 0, fmt.Errorf("instance id=%s not found for model id=%s", modelEntity.InstanceID, modelID) + } + return nil, "", nil, 0, err + } + + apiKey := instanceEntity.APIKey + var extra map[string]string + if err := json.Unmarshal([]byte(instanceEntity.Extra), &extra); err != nil { + return nil, "", nil, 0, err + } + region := extra["region"] + baseURL := extra["base_url"] + + providerInfo := dao.GetModelProviderManager().FindProvider(providerEntity.ProviderName) + if providerInfo == nil { + return nil, "", nil, 0, fmt.Errorf("provider %q driver not found", providerEntity.ProviderName) + } + modelDriver, err := newModelDriverForBaseURL(providerInfo.ModelDriver, providerEntity.ProviderName, region, baseURL) + if err != nil { + return nil, "", nil, 0, err + } + + maxTokens := 0 + if mi, _ := dao.GetModelProviderManager().GetModelByName(providerEntity.ProviderName, modelEntity.ModelName); mi != nil { + if mi.MaxTokens != nil { + maxTokens = *mi.MaxTokens + } + } + maxTokens, err = maxTokensFromTenantModelExtra(modelEntity, maxTokens) + if err != nil { + return nil, "", nil, 0, err + } + + apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL} + return modelDriver, modelEntity.ModelName, apiConfig, maxTokens, nil +} + // GetModelTypeByName returns the list of model types the given model is enrolled as. func (m *ModelProviderService) GetModelTypeByName(tenantID, modelName string) ([]entity.ModelType, error) { common.Debug("GetModelTypeByName", @@ -2210,7 +2283,7 @@ func (m *ModelProviderService) GetModelTypeByName(tenantID, modelName string) ([ if len(targetProvider.Models[i].ModelTypes) == 0 { return nil, fmt.Errorf("model %q has no model_types in factory catalog", pureModelName) } - return []entity.ModelType{entity.ModelType(targetProvider.Models[i].ModelTypes[0])}, nil + return []entity.ModelType{entity.ModelTypeFromString(targetProvider.Models[i].ModelTypes[0])}, nil } } return nil, fmt.Errorf("model %q not found for model %q", pureModelName, modelName) @@ -2270,13 +2343,30 @@ func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) 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") + + // 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) @@ -2308,7 +2398,6 @@ func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) } extra := map[string]interface{}{ "max_tokens": model.MaxTokens, - "model_types": []string{modelType}, "max_dimension": model.MaxDimension, "dimensions": model.Dimensions, } @@ -2325,7 +2414,7 @@ func (m *ModelProviderService) AddModel(request *AddModelRequest, userID string) models = append(models, &entity.TenantModel{ ID: modelID, ModelName: modelName, - ModelType: modelType, + ModelType: int(combinedType), ProviderID: provider.ID, InstanceID: instance.ID, Status: "active", @@ -2351,13 +2440,13 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(tenantID strin common.Debug("GetModelConfigFromProviderInstance", zap.String("tenantID", tenantID), zap.String("modelName", modelName), - zap.String("modelType", string(modelType))) + zap.String("modelType", modelType.String())) // TEI builtin embedding short-circuit if modelType == entity.ModelTypeEmbedding && strings.Contains(os.Getenv("COMPOSE_PROFILES"), "tei-") { teiModel := os.Getenv("TEI_MODEL") teiBaseURL := os.Getenv("TEI_BASE_URL") - + // First try exact match: handles bare model IDs like "model@q8_0" // where '@' is part of the model name itself. if modelName == teiModel { @@ -2368,7 +2457,7 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(tenantID strin apiConfig := &modelModule.APIConfig{ApiKey: nil, Region: nil, BaseURL: &teiBaseURL} return builtinDriver, modelName, apiConfig, 0, nil } - + // Then try right-anchored parsing for explicit "model@Builtin" or // "model@instance@Builtin" composite keys. teiPure, _, teiProvider := splitRightAnchoredModelName(modelName) @@ -2460,7 +2549,7 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(tenantID strin // Direct model lookup modelObj, modelErr := m.modelDAO.GetByProviderIDAndInstanceIDAndModelTypeAndModelName( - provider.ID, instance.ID, string(modelType), pureModelName, + provider.ID, instance.ID, int(modelType), pureModelName, ) switch { case modelErr == nil: @@ -2498,7 +2587,7 @@ func (m *ModelProviderService) GetModelConfigFromProviderInstance(tenantID strin zap.String("tenantID", tenantID), zap.String("providerName", providerName), zap.String("instanceName", instanceName), - zap.String("modelType", string(modelType)), + zap.String("modelType", modelType.String()), zap.String("modelName", pureModelName)) default: // Surface unexpected DAO errors (e.g. transient DB failure) instead of diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index c6e83069d..bbad3960c 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -146,7 +146,7 @@ func seedModelProviderServiceScope(t *testing.T, db *gorm.DB) { &entity.UserTenant{ID: "user-tenant-1", UserID: "user-1", TenantID: "tenant-1", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, &entity.TenantModelProvider{ID: "provider-1", TenantID: "tenant-1", ProviderName: "OpenAI"}, &entity.TenantModelInstance{ID: "instance-1", ProviderID: "provider-1", InstanceName: "default", APIKey: "sk-test", Status: "active", Extra: "{}"}, - &entity.TenantModel{ID: "model-1", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: "chat", Status: "active"}, + &entity.TenantModel{ID: "model-1", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: int(entity.ModelTypeChat), Status: "active"}, } for _, row := range rows { if err := db.Create(row).Error; err != nil { @@ -177,6 +177,26 @@ func TestModelProviderServiceUpdateModelStatusByID(t *testing.T) { } } +func TestModelProviderServiceGetModelConfigByID(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + seedModelProviderServiceScope(t, db) + + driver, modelName, apiConfig, _, err := NewModelProviderService().GetModelConfigByID("user-1", "model-1") + if err != nil { + t.Fatalf("GetModelConfigByID() error = %v", err) + } + if driver == nil { + t.Fatal("GetModelConfigByID() returned nil driver") + } + if modelName != "gpt-test" { + t.Fatalf("modelName = %q, want %q", modelName, "gpt-test") + } + if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey != "sk-test" { + t.Fatalf("apiConfig.ApiKey = %v, want %q", apiConfig.ApiKey, "sk-test") + } +} + func TestModelProviderServiceUpdateModelStatusRejectsInvalidStatus(t *testing.T) { code, err := NewModelProviderService().UpdateModelStatus("OpenAI", "default", "", "user-1", "model-1", "disabled") if err == nil { diff --git a/internal/service/tenant_test.go b/internal/service/tenant_test.go index 70eb2a06e..fa51e7974 100644 --- a/internal/service/tenant_test.go +++ b/internal/service/tenant_test.go @@ -228,7 +228,7 @@ func TestSetTenantDefaultModels_WithModelID(t *testing.T) { ModelName: "gpt-4o", ProviderID: providerID, InstanceID: "instance-dummy", - ModelType: "chat", + ModelType: int(entity.ModelTypeChat), Status: "active", }).Error if err != nil { diff --git a/internal/service/user.go b/internal/service/user.go index 2a2f835e5..8cb057a5b 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -270,11 +270,11 @@ func (s *UserService) getInitTenantLLM(userID string) ([]*entity.TenantLLM, erro } modelConfigs := map[string]server.ModelConfig{ - string(entity.ModelTypeChat): cfg.UserDefaultLLM.DefaultModels.ChatModel, - string(entity.ModelTypeEmbedding): cfg.UserDefaultLLM.DefaultModels.EmbeddingModel, - string(entity.ModelTypeSpeech2Text): cfg.UserDefaultLLM.DefaultModels.ASRModel, - string(entity.ModelTypeImage2Text): cfg.UserDefaultLLM.DefaultModels.Image2TextModel, - string(entity.ModelTypeRerank): cfg.UserDefaultLLM.DefaultModels.RerankModel, + entity.ModelTypeChat.String(): cfg.UserDefaultLLM.DefaultModels.ChatModel, + entity.ModelTypeEmbedding.String(): cfg.UserDefaultLLM.DefaultModels.EmbeddingModel, + entity.ModelTypeSpeech2Text.String(): cfg.UserDefaultLLM.DefaultModels.ASRModel, + entity.ModelTypeImage2Text.String(): cfg.UserDefaultLLM.DefaultModels.Image2TextModel, + entity.ModelTypeRerank.String(): cfg.UserDefaultLLM.DefaultModels.RerankModel, } seenFactories := make(map[string]bool)