124 lines
3.1 KiB
Go
124 lines
3.1 KiB
Go
package providers
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/chenxin/hunjing/mcp/web-search/internal/config"
|
|
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
|
)
|
|
|
|
// Registry 管理多 provider 注册和查询。
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
providers map[models.ProviderKind]SearchProvider
|
|
}
|
|
|
|
func NewRegistry() *Registry {
|
|
return &Registry{
|
|
providers: make(map[models.ProviderKind]SearchProvider),
|
|
}
|
|
}
|
|
|
|
// Register 注册 provider。
|
|
func (r *Registry) Register(p SearchProvider) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.providers[p.Name()] = p
|
|
}
|
|
|
|
// Get 按名称获取 provider。
|
|
func (r *Registry) Get(name models.ProviderKind) (SearchProvider, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
p, ok := r.providers[name]
|
|
return p, ok
|
|
}
|
|
|
|
// All 返回所有已注册 provider。
|
|
func (r *Registry) All() []SearchProvider {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
out := make([]SearchProvider, 0, len(r.providers))
|
|
for _, p := range r.providers {
|
|
out = append(out, p)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Names 返回所有已注册的 provider 名称。
|
|
func (r *Registry) Names() []models.ProviderKind {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
out := make([]models.ProviderKind, 0, len(r.providers))
|
|
for name := range r.providers {
|
|
out = append(out, name)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// SetupFromConfig 根据配置创建并注册所有启用的 provider。
|
|
func SetupFromConfig(cfg *config.Config) (*Registry, error) {
|
|
reg := NewRegistry()
|
|
|
|
if cfg.Providers.Tavily.Enabled {
|
|
if len(cfg.Providers.Tavily.APIKeys) == 0 {
|
|
slog.Warn("tavily: enabled but no API keys configured, skipping")
|
|
} else {
|
|
reg.Register(NewTavilyProvider(cfg.Providers.Tavily.Timeout, cfg.Providers.Tavily.MaxResults))
|
|
}
|
|
}
|
|
if cfg.Providers.Brave.Enabled {
|
|
if len(cfg.Providers.Brave.APIKeys) == 0 {
|
|
slog.Warn("brave: enabled but no API keys configured, skipping")
|
|
} else {
|
|
reg.Register(NewBraveProvider(cfg.Providers.Brave.Timeout, cfg.Providers.Brave.MaxResults))
|
|
}
|
|
}
|
|
|
|
if len(reg.Names()) == 0 {
|
|
return nil, fmt.Errorf("no providers enabled or configured")
|
|
}
|
|
return reg, nil
|
|
}
|
|
|
|
// SetupFromConfigWithOverrides 从配置创建注册中心，并允许使用显式 provider 配置覆盖配置文件。
|
|
// 本函数方便测试和动态配置场景。
|
|
func SetupFromConfigWithOverrides(
|
|
cfg *config.Config,
|
|
overrideTimeout time.Duration,
|
|
overrideMaxResults int,
|
|
) (*Registry, error) {
|
|
reg := NewRegistry()
|
|
|
|
if cfg.Providers.Tavily.Enabled && len(cfg.Providers.Tavily.APIKeys) > 0 {
|
|
t := cfg.Providers.Tavily.Timeout
|
|
m := cfg.Providers.Tavily.MaxResults
|
|
if overrideTimeout > 0 {
|
|
t = overrideTimeout
|
|
}
|
|
if overrideMaxResults > 0 {
|
|
m = overrideMaxResults
|
|
}
|
|
reg.Register(NewTavilyProvider(t, m))
|
|
}
|
|
if cfg.Providers.Brave.Enabled && len(cfg.Providers.Brave.APIKeys) > 0 {
|
|
t := cfg.Providers.Brave.Timeout
|
|
m := cfg.Providers.Brave.MaxResults
|
|
if overrideTimeout > 0 {
|
|
t = overrideTimeout
|
|
}
|
|
if overrideMaxResults > 0 {
|
|
m = overrideMaxResults
|
|
}
|
|
reg.Register(NewBraveProvider(t, m))
|
|
}
|
|
|
|
if len(reg.Names()) == 0 {
|
|
return nil, fmt.Errorf("no providers enabled or configured")
|
|
}
|
|
return reg, nil
|
|
}
|