hunjing/mcp/web-search/internal/config/config.go

// Package config 管理博闻配置。
package config

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"
)

// Config 顶层配置。
type Config struct {
	Server    ServerConfig    `yaml:"server"`
	Providers ProvidersConfig `yaml:"providers"`
	Balancer  BalancerConfig  `yaml:"balancer"`
	Cache     CacheConfig     `yaml:"cache"`
	Metrics   MetricsConfig   `yaml:"metrics"`
}

// ServerConfig 服务端配置。
type ServerConfig struct {
	Port         int           `yaml:"port"`
	ReadTimeout  time.Duration `yaml:"read_timeout"`
	WriteTimeout time.Duration `yaml:"write_timeout"`
}

// ProvidersConfig 聚合所有搜索提供商配置。
type ProvidersConfig struct {
	Tavily ProviderConfig `yaml:"tavily"`
	Brave  ProviderConfig `yaml:"brave"`
}

// ProviderConfig 单个搜索提供商配置。
type ProviderConfig struct {
	Enabled     bool            `yaml:"enabled"`
	Timeout     time.Duration   `yaml:"timeout"`
	MaxResults  int             `yaml:"max_results"`
	HealthCheck HealthCheckConf `yaml:"health_check"`
	// APIKeys 从环境变量注入，不在 YAML 中
	APIKeys []string `yaml:"-"`
}

// HealthCheckConf 健康检查配置。
type HealthCheckConf struct {
	Interval         time.Duration `yaml:"interval"`
	Timeout          time.Duration `yaml:"timeout"`
	FailureThreshold int           `yaml:"failure_threshold"`
}

// BalancerConfig 负载均衡配置。
type BalancerConfig struct {
	Strategy       string            `yaml:"strategy"`
	RetryCount     int               `yaml:"retry_count"`
	CircuitBreaker CircuitBreakerCfg `yaml:"circuit_breaker"`
}

// CircuitBreakerCfg 熔断器配置。
type CircuitBreakerCfg struct {
	Enabled          bool          `yaml:"enabled"`
	FailureThreshold int           `yaml:"failure_threshold"`
	Cooldown         time.Duration `yaml:"cooldown"`
}

// CacheConfig 缓存配置。
type CacheConfig struct {
	Enabled  bool          `yaml:"enabled"`
	TTL      time.Duration `yaml:"ttl"`
	MaxItems int           `yaml:"max_items"`
	Redis    RedisConfig   `yaml:"redis"`
}

// RedisConfig Redis 缓存配置（二级缓存层，可开关）。
type RedisConfig struct {
	Enabled  bool   `yaml:"enabled"`
	Addr     string `yaml:"addr"`
	Password string `yaml:"password"`
	DB       int    `yaml:"db"`
}

// MetricsConfig 监控指标配置。
type MetricsConfig struct {
	Enabled bool   `yaml:"enabled"`
	Port    int    `yaml:"port"`
	Path    string `yaml:"path"`
}

// envVarKeys 定义各 provider 的 API Key 环境变量名。
var envVarKeys = map[string]string{
	"tavily": "TAVILY_API_KEYS",
	"brave":  "BRAVE_API_KEYS",
}

// Load 加载配置，API Keys 从环境变量读取。
// configPath 为可选 YAML 配置文件路径。
func Load(configPath string) (*Config, error) {
	cfg := &Config{}
	setDefaults(cfg)

	// 从 YAML 文件加载（如存在）
	if configPath != "" {
		if err := loadFromFile(configPath, cfg); err != nil {
			return nil, err
		}
	}

	// 从环境变量加载 API Keys
	cfg.Providers.Tavily.APIKeys = parseEnvKeys(envVarKeys["tavily"])
	cfg.Providers.Brave.APIKeys = parseEnvKeys(envVarKeys["brave"])

	// 从环境变量加载 Redis 配置
	if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ENABLED"); env != "" {
		cfg.Cache.Redis.Enabled, _ = strconv.ParseBool(env)
	}
	if env := os.Getenv("MCP_WEB_SEARCH_REDIS_ADDR"); env != "" {
		cfg.Cache.Redis.Addr = env
	}
	if env := os.Getenv("MCP_WEB_SEARCH_REDIS_PASSWORD"); env != "" {
		cfg.Cache.Redis.Password = env
	}
	if env := os.Getenv("MCP_WEB_SEARCH_REDIS_DB"); env != "" {
		if db, err := strconv.Atoi(env); err == nil {
			cfg.Cache.Redis.DB = db
		}
	}

	return cfg, nil
}

func setDefaults(cfg *Config) {
	cfg.Server.Port = 8080
	cfg.Server.ReadTimeout = 30 * time.Second
	cfg.Server.WriteTimeout = 30 * time.Second
	// Provider 默认启用（可通过 YAML 中 enabled: false 关闭）
	cfg.Providers.Tavily.Enabled = true
	cfg.Providers.Tavily.Timeout = 10 * time.Second
	cfg.Providers.Tavily.MaxResults = 10
	cfg.Providers.Brave.Enabled = true
	cfg.Providers.Brave.Timeout = 10 * time.Second
	cfg.Providers.Brave.MaxResults = 10
	cfg.Balancer.Strategy = "weighted_round_robin"
	cfg.Balancer.RetryCount = 2
	cfg.Balancer.CircuitBreaker.Enabled = true
	cfg.Balancer.CircuitBreaker.FailureThreshold = 5
	cfg.Balancer.CircuitBreaker.Cooldown = 30 * time.Second
	cfg.Cache.Enabled = true
	cfg.Cache.TTL = 5 * time.Minute
	cfg.Cache.MaxItems = 1000
	cfg.Cache.Redis.Enabled = false
	cfg.Cache.Redis.Addr = "127.0.0.1:6379"
	cfg.Cache.Redis.Password = ""
	cfg.Cache.Redis.DB = 0
	cfg.Metrics.Enabled = true
	cfg.Metrics.Port = 9090
	cfg.Metrics.Path = "/metrics"
}

func loadFromFile(path string, cfg *Config) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("read config: %w", err)
	}
	return yamlUnmarshal(data, cfg)
}

func parseEnvKeys(envName string) []string {
	val := os.Getenv(envName)
	if val == "" {
		return nil
	}
	parts := strings.Split(val, ",")
	keys := make([]string, 0, len(parts))
	for _, p := range parts {
		if trimmed := strings.TrimSpace(p); trimmed != "" {
			keys = append(keys, trimmed)
		}
	}
	return keys
}