From 15d37ec338bfe8c741adca0e5e9bdf3c9c8a6a3b Mon Sep 17 00:00:00 2001 From: wener Date: Fri, 12 Dec 2025 14:23:55 +0800 Subject: [PATCH] feat: support config.yaml and config.yml loading with reference support --- transports/bifrost-http/lib/config.go | 57 +++++++-- .../bifrost-http/lib/config_yaml_test.go | 111 ++++++++++++++++++ transports/go.mod | 2 +- 3 files changed, 156 insertions(+), 14 deletions(-) create mode 100644 transports/bifrost-http/lib/config_yaml_test.go diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index 13dd6d41d..4dd864da7 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -27,6 +27,7 @@ import ( "github.com/maximhq/bifrost/framework/modelcatalog" "github.com/maximhq/bifrost/framework/vectorstore" "github.com/maximhq/bifrost/plugins/semanticcache" + "gopkg.in/yaml.v3" "gorm.io/gorm" ) @@ -265,30 +266,60 @@ func (c *Config) initializeEncryption(configKey string) error { // - Graceful handling of missing config files func LoadConfig(ctx context.Context, configDirPath string) (*Config, error) { // Initialize separate database connections for optimal performance at scale - configFilePath := filepath.Join(configDirPath, "config.json") configDBPath := filepath.Join(configDirPath, "config.db") logsDBPath := filepath.Join(configDirPath, "logs.db") + // Initialize config config := &Config{ - configPath: configFilePath, + configPath: filepath.Join(configDirPath, "config.json"), EnvKeys: make(map[string][]configstore.EnvKeyInfo), Providers: make(map[schemas.ModelProvider]configstore.ProviderConfig), Plugins: atomic.Pointer[[]schemas.Plugin]{}, } - // Getting absolute path for config file - absConfigFilePath, err := filepath.Abs(configFilePath) - if err != nil { - return nil, fmt.Errorf("failed to get absolute path for config file: %w", err) + + // Check if config file exists + // Check order: config.json -> config.yaml -> config.yml + var data []byte + var err error + var absConfigFilePath string + var loadedFile string + + possibleFiles := []string{"config.json", "config.yaml", "config.yml"} + for _, file := range possibleFiles { + fullPath := filepath.Join(configDirPath, file) + data, err = os.ReadFile(fullPath) + if err == nil { + loadedFile = file + absConfigFilePath, _ = filepath.Abs(fullPath) + config.configPath = fullPath + break + } + if !os.IsNotExist(err) { + // Real error reading file + return nil, fmt.Errorf("failed to read config file %s: %w", file, err) + } } // Check if config file exists - data, err := os.ReadFile(configFilePath) - if err != nil { - // If config file doesn't exist, we will directly use the config store (create one if it doesn't exist) - if os.IsNotExist(err) { - logger.Info("config file not found at path: %s, initializing with default values", absConfigFilePath) - return loadConfigFromDefaults(ctx, config, configDBPath, logsDBPath) + if loadedFile == "" { + // No config file found, use defaults + absConfigFilePath, _ = filepath.Abs(filepath.Join(configDirPath, "config.json")) // Default for logging + logger.Info("config file not found in %s, initializing with default values", configDirPath) + return loadConfigFromDefaults(ctx, config, configDBPath, logsDBPath) + } + + // If it's a YAML file, convert to JSON first + // This handles YAML aliases/anchors by unmarshaling to interface{} first + if strings.HasSuffix(loadedFile, ".yaml") || strings.HasSuffix(loadedFile, ".yml") { + var yamlData interface{} + if err := yaml.Unmarshal(data, &yamlData); err != nil { + return nil, fmt.Errorf("failed to parse yaml config file: %w", err) + } + // Convert back to JSON + jsonData, err := json.Marshal(yamlData) + if err != nil { + return nil, fmt.Errorf("failed to convert yaml config to json: %w", err) } - return nil, fmt.Errorf("failed to read config file: %w", err) + data = jsonData } // If file exists, we will do a quick check if that file includes "$schema":"https://www.getbifrost.ai/schema", If not we will show a warning in a box - yellow color var schema map[string]interface{} diff --git a/transports/bifrost-http/lib/config_yaml_test.go b/transports/bifrost-http/lib/config_yaml_test.go new file mode 100644 index 000000000..a8b9a9221 --- /dev/null +++ b/transports/bifrost-http/lib/config_yaml_test.go @@ -0,0 +1,111 @@ +package lib + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestLoadConfig_YAML(t *testing.T) { + dir := createTempDir(t) + + yamlContent := ` +client: + initial_pool_size: 50 +providers: + openai: + keys: + - &common + name: common-key + value: sk-test-123 + models: + - gpt-4 + - <<: *common + name: specific-key +` + err := os.WriteFile(filepath.Join(dir, "config.yaml"), []byte(yamlContent), 0644) + if err != nil { + t.Fatalf("failed to write yaml file: %v", err) + } + + config, err := LoadConfig(context.Background(), dir) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + if config.ClientConfig.InitialPoolSize != 50 { + t.Errorf("expected InitialPoolSize 50, got %d", config.ClientConfig.InitialPoolSize) + } + + openai, ok := config.Providers["openai"] + if !ok { + t.Fatal("openai provider not found") + } + + if len(openai.Keys) != 2 { + t.Fatalf("expected 2 keys, got %d", len(openai.Keys)) + } + + // Order might vary depending on how YAML unmarshal handles keys map if it were a map, but keys is a list here + // The order in list should be preserved. + + // First key + if openai.Keys[0].Name != "common-key" { + t.Errorf("expected key 0 name common-key, got %s", openai.Keys[0].Name) + } + if openai.Keys[0].Value != "sk-test-123" { + t.Errorf("expected key 0 value sk-test-123, got %s", openai.Keys[0].Value) + } + + // Second key (merged) + if openai.Keys[1].Name != "specific-key" { + t.Errorf("expected key 1 name specific-key, got %s", openai.Keys[1].Name) + } + if openai.Keys[1].Value != "sk-test-123" { + // Should inherit from anchor + t.Errorf("expected key 1 value sk-test-123, got %s", openai.Keys[1].Value) + } +} + +func TestLoadConfig_YAML_Precedence(t *testing.T) { + dir := createTempDir(t) + + // Create config.json + err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"client":{"initial_pool_size": 10}}`), 0644) + if err != nil { + t.Fatalf("failed to write json file: %v", err) + } + + // Create config.yaml + err = os.WriteFile(filepath.Join(dir, "config.yaml"), []byte(` +client: + initial_pool_size: 20 +`), 0644) + if err != nil { + t.Fatalf("failed to write yaml file: %v", err) + } + + config, err := LoadConfig(context.Background(), dir) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + // Should prefer config.json + if config.ClientConfig.InitialPoolSize != 10 { + t.Errorf("expected InitialPoolSize 10 (from json), got %d", config.ClientConfig.InitialPoolSize) + } + + // Remove json and retry + os.Remove(filepath.Join(dir, "config.json")) + + config, err = LoadConfig(context.Background(), dir) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + + // Should fallback to config.yaml + if config.ClientConfig.InitialPoolSize != 20 { + t.Errorf("expected InitialPoolSize 20 (from yaml), got %d", config.ClientConfig.InitialPoolSize) + } +} diff --git a/transports/go.mod b/transports/go.mod index b0ad42b09..bf4c828ec 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -18,6 +18,7 @@ require ( github.com/maximhq/bifrost/plugins/telemetry v1.3.46 github.com/prometheus/client_golang v1.23.0 github.com/valyala/fasthttp v1.67.0 + gopkg.in/yaml.v3 v3.0.1 gorm.io/driver/sqlite v1.6.0 gorm.io/gorm v1.31.1 ) @@ -124,6 +125,5 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20251111163417-95abcf5c77ba // indirect google.golang.org/grpc v1.76.0 // indirect google.golang.org/protobuf v1.36.10 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect gorm.io/driver/postgres v1.6.0 // indirect )