// // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // package dao import ( "fmt" "ragflow/internal/entity" ) // TenantLLMDAO tenant LLM data access object type TenantLLMDAO struct{} // NewTenantLLMDAO create tenant LLM DAO func NewTenantLLMDAO() *TenantLLMDAO { return &TenantLLMDAO{} } // GetByID get tenant LLM by primary key ID func (dao *TenantLLMDAO) GetByID(id int64) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM err := DB.Where("id = ?", id).First(&tenantLLM).Error if err != nil { return nil, err } return &tenantLLM, nil } // GetByTenantAndModelName get tenant LLM by tenant ID and model name func (dao *TenantLLMDAO) GetByTenantAndModelName(tenantID, providerName string, modelName string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM err := DB.Where("tenant_id = ? AND llm_factory = ? AND llm_name = ?", tenantID, providerName, modelName).First(&tenantLLM).Error if err != nil { return nil, err } return &tenantLLM, nil } // 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 if err != nil { return nil, err } return &tenantLLM, nil } // 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 if err != nil { return nil, err } return &tenantLLM, nil } // 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 if err != nil { return nil, err } return &tenantLLM, nil } // ListByTenant list all tenant LLMs for a tenant func (dao *TenantLLMDAO) ListByTenant(tenantID string) ([]entity.TenantLLM, error) { var tenantLLMs []entity.TenantLLM err := DB.Where("tenant_id = ?", tenantID).Find(&tenantLLMs).Error if err != nil { return nil, err } return tenantLLMs, nil } // GetByTenantFactoryAndModelName get tenant LLM by tenant ID, factory and model name func (dao *TenantLLMDAO) GetByTenantFactoryAndModelName(tenantID, factory, modelName string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM err := DB.Where("tenant_id = ? AND llm_factory = ? AND llm_name = ?", tenantID, factory, modelName).First(&tenantLLM).Error if err != nil { return nil, err } return &tenantLLM, nil } // Create create a new tenant LLM record func (dao *TenantLLMDAO) Create(tenantLLM *entity.TenantLLM) error { return DB.Create(tenantLLM).Error } // Update update an existing tenant LLM record func (dao *TenantLLMDAO) Update(tenantLLM *entity.TenantLLM) error { return DB.Save(tenantLLM).Error } // Delete delete a tenant LLM record by tenant ID, factory and model name func (dao *TenantLLMDAO) Delete(tenantID, factory, modelName string) error { return DB.Where("tenant_id = ? AND llm_factory = ? AND llm_name = ?", tenantID, factory, modelName).Delete(&entity.TenantLLM{}).Error } // GetMyLLMs get tenant LLMs with factory details func (dao *TenantLLMDAO) GetMyLLMs(tenantID string) ([]entity.MyLLM, error) { var myLLMs []entity.MyLLM err := DB.Table("tenant_llm tl"). Select("tl.id, tl.llm_factory, lf.logo, lf.tags, tl.model_type, tl.llm_name, tl.used_tokens, tl.status"). Joins("JOIN llm_factories lf ON tl.llm_factory = lf.name"). Where("tl.tenant_id = ? AND tl.api_key IS NOT NULL", tenantID). Find(&myLLMs).Error if err != nil { return nil, err } return myLLMs, nil } // ListValidByTenant lists valid tenant LLMs for a tenant func (dao *TenantLLMDAO) ListValidByTenant(tenantID string) ([]*entity.TenantLLM, error) { var tenantLLMs []*entity.TenantLLM err := DB.Where("tenant_id = ? AND api_key IS NOT NULL AND api_key != ? AND status = ?", tenantID, "", "1").Find(&tenantLLMs).Error if err != nil { return nil, err } return tenantLLMs, nil } // ListAllByTenant lists all tenant LLMs for a tenant func (dao *TenantLLMDAO) ListAllByTenant(tenantID string) ([]*entity.TenantLLM, error) { var tenantLLMs []*entity.TenantLLM err := DB.Where("tenant_id = ?", tenantID).Find(&tenantLLMs).Error if err != nil { return nil, err } return tenantLLMs, nil } // InsertMany inserts multiple tenant LLM records func (dao *TenantLLMDAO) InsertMany(tenantLLMs []*entity.TenantLLM) error { if len(tenantLLMs) == 0 { return nil } return DB.Create(&tenantLLMs).Error } // DeleteByTenantID deletes all tenant LLM records by tenant ID (hard delete) func (dao *TenantLLMDAO) DeleteByTenantID(tenantID string) (int64, error) { result := DB.Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.TenantLLM{}) return result.RowsAffected, result.Error } // splitModelNameAndFactory splits model name and factory from combined format // This matches Python's split_model_name_and_factory logic // // Parameters: // - modelName: The model name which can be in format "ModelName" or "ModelName@Factory" // // Returns: // - string: The model name without factory prefix // - string: The factory name (empty string if not specified) // // Example: // // modelName, factory := splitModelNameAndFactory("gpt-4") // // Returns: "gpt-4", "" // // modelName, factory := splitModelNameAndFactory("gpt-4@OpenAI") // // Returns: "gpt-4", "OpenAI" func splitModelNameAndFactory(modelName string) (string, string) { // Split by "@" separator // Handle cases like "model@factory" or "model@sub@factory" lastAtIndex := -1 for i := len(modelName) - 1; i >= 0; i-- { if modelName[i] == '@' { lastAtIndex = i break } } // No "@" found, return original name if lastAtIndex == -1 { return modelName, "" } // Split into model name and potential factory modelNamePart := modelName[:lastAtIndex] factory := modelName[lastAtIndex+1:] // Validate if factory exists in llm_factories table // This matches Python's logic of checking against model providers var factoryCount int64 DB.Model(&entity.LLMFactories{}).Where("name = ?", factory).Count(&factoryCount) // If factory doesn't exist in database, treat the whole string as model name if factoryCount == 0 { return modelName, "" } return modelNamePart, factory } // GetByTenantIDAndLLMName gets tenant LLM by tenant ID and LLM name // This is used to resolve tenant_llm_id from llm_id // It supports both simple model names and factory-prefixed names (e.g., "gpt-4@OpenAI") // // Parameters: // - tenantID: The tenant identifier // - llmName: The LLM model name (can include factory prefix like "OpenAI@gpt-4") // // Returns: // - *model.TenantLLM: The tenant LLM record // - error: Error if not found // // Example: // // // Simple model name // tenantLLM, err := dao.GetByTenantIDAndLLMName("tenant123", "gpt-4") // // // Model name with factory prefix // tenantLLM, err := dao.GetByTenantIDAndLLMName("tenant123", "gpt-4@OpenAI") func (dao *TenantLLMDAO) GetByTenantIDAndLLMName(tenantID string, llmName string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM // Split model name and factory from the combined format modelName, factory := splitModelNameAndFactory(llmName) // First attempt: try to find with model name only err := DB.Where("tenant_id = ? AND llm_name = ?", tenantID, modelName).First(&tenantLLM).Error if err == nil { return &tenantLLM, nil } // Second attempt: if factory is specified, try with both model name and factory if factory != "" { err = DB.Where("tenant_id = ? AND llm_name = ? AND llm_factory = ?", tenantID, modelName, factory).First(&tenantLLM).Error if err == nil { return &tenantLLM, nil } // Special handling for LocalAI and HuggingFace (matching Python logic) // These factories append "___FactoryName" to the model name if factory == "LocalAI" || factory == "HuggingFace" || factory == "OpenAI-API-Compatible" { specialModelName := modelName + "___" + factory err = DB.Where("tenant_id = ? AND llm_name = ?", tenantID, specialModelName).First(&tenantLLM).Error if err == nil { return &tenantLLM, nil } } } // Return the last error (record not found) return nil, err } // GetByTenantIDLLMNameAndFactory gets tenant LLM by tenant ID, LLM name and factory // This is used when model name includes factory suffix (e.g., "model@factory") // // Parameters: // - tenantID: The tenant identifier // - llmName: The LLM model name // - factory: The LLM factory name // // Returns: // - *model.TenantLLM: The tenant LLM record // - error: Error if not found // // Example: // // tenantLLM, err := dao.GetByTenantIDLLMNameAndFactory("tenant123", "gpt-4", "OpenAI") func (dao *TenantLLMDAO) GetByTenantIDLLMNameAndFactory(tenantID, llmName, factory string) (*entity.TenantLLM, error) { var tenantLLM entity.TenantLLM err := DB.Where("tenant_id = ? AND llm_name = ? AND llm_factory = ?", tenantID, llmName, factory).First(&tenantLLM).Error if err != nil { return nil, err } return &tenantLLM, nil } // LookupTenantLLMByID looks up a TenantLLM record by ID and returns the record plus composite model name. func LookupTenantLLMByID(tenantLLMDao *TenantLLMDAO, id int64) (*entity.TenantLLM, string, error) { tenantLLM, err := tenantLLMDao.GetByID(id) if err != nil { return nil, "", fmt.Errorf("failed to get tenant_llm by id %d: %w", id, err) } if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { return nil, "", fmt.Errorf("tenant_llm record not found for id %d", id) } compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) return tenantLLM, compositeName, nil } // LookupTenantLLMByName looks up a TenantLLM record by tenant name and model type. func LookupTenantLLMByName(tenantLLMDao *TenantLLMDAO, tenantID, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) { // Parse factory from name if present (e.g., "model@Factory") modelName, factory := splitModelNameAndFactory(name) // If factory is found, use factory-based lookup if factory != "" { return LookupTenantLLMByFactory(tenantLLMDao, tenantID, factory, modelName, modelType) } tenantLLM, err := tenantLLMDao.GetByTenantNameAndType(tenantID, modelName, modelType) if err != nil { return nil, "", fmt.Errorf("failed to get tenant_llm by name %s: %w", name, err) } if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { return nil, "", fmt.Errorf("tenant_llm record not found for name %s", name) } compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) return tenantLLM, compositeName, nil } // LookupTenantLLMByFactory looks up a TenantLLM record by tenant, factory, and model name. func LookupTenantLLMByFactory(tenantLLMDao *TenantLLMDAO, tenantID, factory, name string, modelType entity.ModelType) (*entity.TenantLLM, string, error) { tenantLLM, err := tenantLLMDao.GetByTenantFactoryAndModelName(tenantID, factory, name) if err != nil { return nil, "", fmt.Errorf("failed to get tenant_llm by factory %s and name %s: %w", factory, name, err) } if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" { return nil, "", fmt.Errorf("tenant_llm record not found for factory %s and name %s", factory, name) } compositeName := fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory) return tenantLLM, compositeName, nil }