zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-09 12:54:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix(go): adapt to new db columns (#16745)

Browse Source
This commit is contained in:
Lynn
2026-07-08 18:02:11 +08:00
committed by GitHub
parent 21266286cb
commit 330033d7c2
22 changed files with 346 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
internal/admin/service.go
View File

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

9
internal/dao/tenant_llm.go
View File

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

6
internal/dao/tenant_model.go
View File

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

6
internal/dao/tenant_model_test.go
View File

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

4
internal/entity/chat.go
View File

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

2
internal/entity/kb.go
View File

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

4
internal/entity/memory.go
View File

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

2
internal/entity/skill_search.go
View File

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

14
internal/entity/tenant.go
View File

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

2
internal/entity/tenant_model.go
View File

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

136
internal/entity/types.go
View File

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

2
internal/handler/providers_test.go
View File

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

2
internal/ingestion/component/pdf_vision_dispatch.go
View File

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

10
internal/service/chat.go
View File

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

51
internal/service/chunk/chunk.go
View File

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

4
internal/service/chunk/chunk_test.go
View File

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

4
internal/service/dataset_update_test.go
View File

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

22
internal/service/memory.go
View File

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

179
internal/service/model_service.go
View File

@@ -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: &region, 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

22
internal/service/model_service_test.go
View File

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

2
internal/service/tenant_test.go
View File

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

10
internal/service/user.go
View File

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