refactor: Moves the LLM factory initialization logic to the dao package. (#13476)

### What problem does this PR solve?

refactor: Moves the LLM factory initialization logic to the `dao`
package.

Removes the `init_data` package and integrates the LLM factory
initialization functionality into the `dao` package.
Adds a `utility` package to provide general utility functions.
Updates `server_main.go` to use the new initialization path.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Co-authored-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
chanx
2026-03-10 10:35:55 +08:00
committed by GitHub
parent 88a40b95a2
commit 02108772d8
4 changed files with 204 additions and 163 deletions

View File

@ -17,21 +17,50 @@
package dao
import (
"encoding/json"
"fmt"
"log"
"os"
"path/filepath"
"time"
"ragflow/internal/logger"
"ragflow/internal/model"
"ragflow/internal/server"
"time"
"ragflow/internal/utility"
gormLogger "gorm.io/gorm/logger"
"gorm.io/driver/mysql"
"gorm.io/gorm"
"ragflow/internal/logger"
)
var DB *gorm.DB
// LLMFactoryConfig represents a single LLM factory configuration
type LLMFactoryConfig struct {
Name string `json:"name"`
Logo string `json:"logo"`
Tags string `json:"tags"`
Status string `json:"status"`
Rank string `json:"rank"`
LLM []LLMConfig `json:"llm"`
}
// LLMConfig represents a single LLM model configuration
type LLMConfig struct {
LLMName string `json:"llm_name"`
Tags string `json:"tags"`
MaxTokens int64 `json:"max_tokens"`
ModelType string `json:"model_type"`
IsTools bool `json:"is_tools"`
}
// LLMFactoriesFile represents the structure of llm_factories.json
type LLMFactoriesFile struct {
FactoryLLMInfos []LLMFactoryConfig `json:"factory_llm_infos"`
}
// InitDB initialize database connection
func InitDB() error {
cfg := server.GetConfig()
@ -132,3 +161,93 @@ func InitDB() error {
func GetDB() *gorm.DB {
return DB
}
// InitLLMFactory initializes LLM factories and models from JSON file.
// It reads the llm_factories.json configuration file and populates the database
// with LLM factory and model information. If a factory or model already exists,
// it will be updated with the new configuration.
//
// Returns:
// - error: An error if the initialization fails, nil otherwise.
func InitLLMFactory() error {
configPath := filepath.Join(utility.GetProjectBaseDirectory(), "conf", "llm_factories.json")
data, err := os.ReadFile(configPath)
if err != nil {
return fmt.Errorf("failed to read llm_factories.json: %w", err)
}
var fileData LLMFactoriesFile
if err := json.Unmarshal(data, &fileData); err != nil {
return fmt.Errorf("failed to parse llm_factories.json: %w", err)
}
db := DB
for _, factory := range fileData.FactoryLLMInfos {
status := factory.Status
if status == "" {
status = "1"
}
llmFactory := &model.LLMFactories{
Name: factory.Name,
Logo: utility.StringPtr(factory.Logo),
Tags: factory.Tags,
Rank: utility.ParseInt64(factory.Rank),
Status: &status,
}
var existingFactory model.LLMFactories
result := db.Where("name = ?", factory.Name).First(&existingFactory)
if result.Error != nil {
if err := db.Create(llmFactory).Error; err != nil {
log.Printf("Failed to create LLM factory %s: %v", factory.Name, err)
continue
}
} else {
if err := db.Model(&model.LLMFactories{}).Where("name = ?", factory.Name).Updates(map[string]interface{}{
"logo": llmFactory.Logo,
"tags": llmFactory.Tags,
"rank": llmFactory.Rank,
"status": llmFactory.Status,
}).Error; err != nil {
log.Printf("Failed to update LLM factory %s: %v", factory.Name, err)
}
}
for _, llm := range factory.LLM {
llmStatus := "1"
llmModel := &model.LLM{
LLMName: llm.LLMName,
ModelType: llm.ModelType,
FID: factory.Name,
MaxTokens: llm.MaxTokens,
Tags: llm.Tags,
IsTools: llm.IsTools,
Status: &llmStatus,
}
var existingLLM model.LLM
result := db.Where("llm_name = ? AND fid = ?", llm.LLMName, factory.Name).First(&existingLLM)
if result.Error != nil {
if err := db.Create(llmModel).Error; err != nil {
log.Printf("Failed to create LLM %s/%s: %v", factory.Name, llm.LLMName, err)
}
} else {
if err := db.Model(&model.LLM{}).Where("llm_name = ? AND fid = ?", llm.LLMName, factory.Name).Updates(map[string]interface{}{
"model_type": llmModel.ModelType,
"max_tokens": llmModel.MaxTokens,
"tags": llmModel.Tags,
"is_tools": llmModel.IsTools,
"status": llmModel.Status,
}).Error; err != nil {
log.Printf("Failed to update LLM %s/%s: %v", factory.Name, llm.LLMName, err)
}
}
}
}
log.Println("LLM factories initialized successfully")
return nil
}