feat(web_search): add load balance and failover for api keys

This commit is contained in:
stark 2026-03-03 20:59:55 +08:00
parent f18d661454
commit 98e5ff81d8
8 changed files with 82 additions and 44 deletions

View file

@ -216,7 +216,10 @@
"web": { "web": {
"brave": { "brave": {
"enabled": false, "enabled": false,
"api_keys": "YOUR_BRAVE_API_KEY", "api_key": "YOUR_BRAVE_API_KEY",
"api_keys": [
"YOUR_BRAVE_API_KEY"
],
"max_results": 5 "max_results": 5
}, },
"duckduckgo": { "duckduckgo": {
@ -225,7 +228,10 @@
}, },
"perplexity": { "perplexity": {
"enabled": false, "enabled": false,
"api_keys": "pplx-xxx", "api_key": "pplx-xxx",
"api_keys": [
"pplx-xxx"
],
"max_results": 5 "max_results": 5
}, },
"proxy": "" "proxy": ""

View file

@ -100,16 +100,16 @@ func registerSharedTools(
// Web tools // Web tools
searchTool, err := tools.NewWebSearchTool(tools.WebSearchToolOptions{ searchTool, err := tools.NewWebSearchTool(tools.WebSearchToolOptions{
BraveAPIKeys: cfg.Tools.Web.Brave.APIKeys, BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys),
BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, BraveMaxResults: cfg.Tools.Web.Brave.MaxResults,
BraveEnabled: cfg.Tools.Web.Brave.Enabled, BraveEnabled: cfg.Tools.Web.Brave.Enabled,
TavilyAPIKeys: cfg.Tools.Web.Tavily.APIKeys, TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys),
TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL,
TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults,
TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, TavilyEnabled: cfg.Tools.Web.Tavily.Enabled,
DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults,
DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled,
PerplexityAPIKeys: cfg.Tools.Web.Perplexity.APIKeys, PerplexityAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Perplexity.APIKey, cfg.Tools.Web.Perplexity.APIKeys),
PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults,
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
Proxy: cfg.Tools.Web.Proxy, Proxy: cfg.Tools.Web.Proxy,

View file

@ -4,6 +4,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"os" "os"
"strings"
"sync/atomic" "sync/atomic"
"github.com/caarlos0/env/v11" "github.com/caarlos0/env/v11"
@ -494,16 +495,18 @@ type GatewayConfig struct {
} }
type BraveConfig struct { type BraveConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_BRAVE_ENABLED"` Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_BRAVE_ENABLED"`
APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEYS"` APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEY"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_BRAVE_MAX_RESULTS"` APIKeys []string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_BRAVE_MAX_RESULTS"`
} }
type TavilyConfig struct { type TavilyConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_TAVILY_ENABLED"` Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_TAVILY_ENABLED"`
APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEYS"` APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEY"`
BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_TAVILY_BASE_URL"` APIKeys []string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"` BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_TAVILY_BASE_URL"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"`
} }
type DuckDuckGoConfig struct { type DuckDuckGoConfig struct {
@ -512,9 +515,10 @@ type DuckDuckGoConfig struct {
} }
type PerplexityConfig struct { type PerplexityConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_ENABLED"` Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_ENABLED"`
APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEYS"` APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEY"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"` APIKeys []string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"`
} }
type WebToolsConfig struct { type WebToolsConfig struct {
@ -779,3 +783,26 @@ func (c *Config) ValidateModelList() error {
} }
return nil return nil
} }
func MergeAPIKeys(apiKey string, apiKeys []string) []string {
seen := make(map[string]struct{})
var all []string
if k := strings.TrimSpace(apiKey); k != "" {
if _, exists := seen[k]; !exists {
seen[k] = struct{}{}
all = append(all, k)
}
}
for _, k := range apiKeys {
if trimmed := strings.TrimSpace(k); trimmed != "" {
if _, exists := seen[trimmed]; !exists {
seen[trimmed] = struct{}{}
all = append(all, trimmed)
}
}
}
return all
}

View file

@ -319,7 +319,14 @@ func DefaultConfig() *Config {
FetchLimitBytes: 10 * 1024 * 1024, // 10MB by default FetchLimitBytes: 10 * 1024 * 1024, // 10MB by default
Brave: BraveConfig{ Brave: BraveConfig{
Enabled: false, Enabled: false,
APIKeys: "", APIKey: "",
APIKeys: []string{"YOUR_BRAVE_API_KEY"},
MaxResults: 5,
},
Tavily: TavilyConfig{
Enabled: false,
APIKey: "",
APIKeys: []string{"YOUR_TAVILY_API_KEY"},
MaxResults: 5, MaxResults: 5,
}, },
DuckDuckGo: DuckDuckGoConfig{ DuckDuckGo: DuckDuckGoConfig{
@ -328,7 +335,8 @@ func DefaultConfig() *Config {
}, },
Perplexity: PerplexityConfig{ Perplexity: PerplexityConfig{
Enabled: false, Enabled: false,
APIKeys: "", APIKey: "",
APIKeys: []string{"YOUR_PERPLEXITY_API_KEY"},
MaxResults: 5, MaxResults: 5,
}, },
}, },

View file

@ -719,12 +719,14 @@ type WebToolsConfig struct {
type BraveConfig struct { type BraveConfig struct {
Enabled bool `json:"enabled"` Enabled bool `json:"enabled"`
APIKey string `json:"api_key"` APIKey string `json:"api_key"`
APIKeys []string `json:"api_keys"`
MaxResults int `json:"max_results"` MaxResults int `json:"max_results"`
} }
type TavilyConfig struct { type TavilyConfig struct {
Enabled bool `json:"enabled"` Enabled bool `json:"enabled"`
APIKey string `json:"api_key"` APIKey string `json:"api_key"`
APIKeys []string `json:"api_keys"`
BaseURL string `json:"base_url"` BaseURL string `json:"base_url"`
MaxResults int `json:"max_results"` MaxResults int `json:"max_results"`
} }
@ -737,6 +739,7 @@ type DuckDuckGoConfig struct {
type PerplexityConfig struct { type PerplexityConfig struct {
Enabled bool `json:"enabled"` Enabled bool `json:"enabled"`
APIKey string `json:"api_key"` APIKey string `json:"api_key"`
APIKeys []string `json:"api_keys"`
MaxResults int `json:"max_results"` MaxResults int `json:"max_results"`
} }
@ -1043,12 +1046,13 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig {
Web: config.WebToolsConfig{ Web: config.WebToolsConfig{
Brave: config.BraveConfig{ Brave: config.BraveConfig{
Enabled: c.Web.Brave.Enabled, Enabled: c.Web.Brave.Enabled,
APIKeys: c.Web.Brave.APIKey, APIKey: c.Web.Brave.APIKey,
APIKeys: c.Web.Brave.APIKeys,
MaxResults: c.Web.Brave.MaxResults, MaxResults: c.Web.Brave.MaxResults,
}, },
Tavily: config.TavilyConfig{ Tavily: config.TavilyConfig{
Enabled: c.Web.Tavily.Enabled, Enabled: c.Web.Tavily.Enabled,
APIKeys: c.Web.Tavily.APIKey, APIKey: c.Web.Tavily.APIKey,
BaseURL: c.Web.Tavily.BaseURL, BaseURL: c.Web.Tavily.BaseURL,
MaxResults: c.Web.Tavily.MaxResults, MaxResults: c.Web.Tavily.MaxResults,
}, },
@ -1058,7 +1062,7 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig {
}, },
Perplexity: config.PerplexityConfig{ Perplexity: config.PerplexityConfig{
Enabled: c.Web.Perplexity.Enabled, Enabled: c.Web.Perplexity.Enabled,
APIKeys: c.Web.Perplexity.APIKey, APIKey: c.Web.Perplexity.APIKey,
MaxResults: c.Web.Perplexity.MaxResults, MaxResults: c.Web.Perplexity.MaxResults,
}, },
Proxy: c.Web.Proxy, Proxy: c.Web.Proxy,

View file

@ -82,13 +82,7 @@ type APIKeyPool struct {
current uint32 current uint32
} }
func NewAPIKeyPool(keysStr string) *APIKeyPool { func NewAPIKeyPool(keys []string) *APIKeyPool {
var keys []string
for _, k := range strings.Split(keysStr, ",") {
if trimmed := strings.TrimSpace(k); trimmed != "" {
keys = append(keys, trimmed)
}
}
return &APIKeyPool{ return &APIKeyPool{
keys: keys, keys: keys,
} }
@ -301,7 +295,6 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
} }
} }
return strings.Join(lines, "\n"), nil return strings.Join(lines, "\n"), nil
} }
@ -499,16 +492,16 @@ type WebSearchTool struct {
} }
type WebSearchToolOptions struct { type WebSearchToolOptions struct {
BraveAPIKeys string BraveAPIKeys []string
BraveMaxResults int BraveMaxResults int
BraveEnabled bool BraveEnabled bool
TavilyAPIKeys string TavilyAPIKeys []string
TavilyBaseURL string TavilyBaseURL string
TavilyMaxResults int TavilyMaxResults int
TavilyEnabled bool TavilyEnabled bool
DuckDuckGoMaxResults int DuckDuckGoMaxResults int
DuckDuckGoEnabled bool DuckDuckGoEnabled bool
PerplexityAPIKeys string PerplexityAPIKeys []string
PerplexityMaxResults int PerplexityMaxResults int
PerplexityEnabled bool PerplexityEnabled bool
Proxy string Proxy string
@ -519,7 +512,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
maxResults := 5 maxResults := 5
// Priority: Perplexity > Brave > Tavily > DuckDuckGo // Priority: Perplexity > Brave > Tavily > DuckDuckGo
if opts.PerplexityEnabled && opts.PerplexityAPIKeys != "" { if opts.PerplexityEnabled && len(opts.PerplexityAPIKeys) > 0 {
client, err := createHTTPClient(opts.Proxy, perplexityTimeout) client, err := createHTTPClient(opts.Proxy, perplexityTimeout)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err)
@ -528,7 +521,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if opts.PerplexityMaxResults > 0 { if opts.PerplexityMaxResults > 0 {
maxResults = opts.PerplexityMaxResults maxResults = opts.PerplexityMaxResults
} }
} else if opts.BraveEnabled && opts.BraveAPIKeys != "" { } else if opts.BraveEnabled && len(opts.BraveAPIKeys) > 0 {
client, err := createHTTPClient(opts.Proxy, searchTimeout) client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err)
@ -537,7 +530,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if opts.BraveMaxResults > 0 { if opts.BraveMaxResults > 0 {
maxResults = opts.BraveMaxResults maxResults = opts.BraveMaxResults
} }
} else if opts.TavilyEnabled && opts.TavilyAPIKeys != "" { } else if opts.TavilyEnabled && len(opts.TavilyAPIKeys) > 0 {
client, err := createHTTPClient(opts.Proxy, searchTimeout) client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err)

View file

@ -249,7 +249,7 @@ func TestWebFetchTool_PayloadTooLarge(t *testing.T) {
// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing // TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing
func TestWebTool_WebSearch_NoApiKey(t *testing.T) { func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: ""}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: nil})
if err != nil { if err != nil {
t.Fatalf("Unexpected error: %v", err) t.Fatalf("Unexpected error: %v", err)
} }
@ -269,7 +269,7 @@ func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
// TestWebTool_WebSearch_MissingQuery verifies error handling for missing query // TestWebTool_WebSearch_MissingQuery verifies error handling for missing query
func TestWebTool_WebSearch_MissingQuery(t *testing.T) { func TestWebTool_WebSearch_MissingQuery(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: "test-key", BraveMaxResults: 5}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: []string{"test-key"}, BraveMaxResults: 5})
if err != nil { if err != nil {
t.Fatalf("Unexpected error: %v", err) t.Fatalf("Unexpected error: %v", err)
} }
@ -553,7 +553,7 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("perplexity", func(t *testing.T) { t.Run("perplexity", func(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
PerplexityEnabled: true, PerplexityEnabled: true,
PerplexityAPIKeys: "k", PerplexityAPIKeys: []string{"k"},
PerplexityMaxResults: 3, PerplexityMaxResults: 3,
Proxy: "http://127.0.0.1:7890", Proxy: "http://127.0.0.1:7890",
}) })
@ -572,7 +572,7 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("brave", func(t *testing.T) { t.Run("brave", func(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
BraveEnabled: true, BraveEnabled: true,
BraveAPIKeys: "k", BraveAPIKeys: []string{"k"},
BraveMaxResults: 3, BraveMaxResults: 3,
Proxy: "http://127.0.0.1:7890", Proxy: "http://127.0.0.1:7890",
}) })
@ -650,7 +650,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true, TavilyEnabled: true,
TavilyAPIKeys: "test-key", TavilyAPIKeys: []string{"test-key"},
TavilyBaseURL: server.URL, TavilyBaseURL: server.URL,
TavilyMaxResults: 5, TavilyMaxResults: 5,
}) })
@ -683,7 +683,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
} }
func TestAPIKeyPool(t *testing.T) { func TestAPIKeyPool(t *testing.T) {
pool := NewAPIKeyPool("key1, key2 ,key3") pool := NewAPIKeyPool([]string{"key1", " key2 ", "key3"})
if len(pool.keys) != 3 { if len(pool.keys) != 3 {
t.Fatalf("expected 3 keys, got %d", len(pool.keys)) t.Fatalf("expected 3 keys, got %d", len(pool.keys))
} }
@ -705,12 +705,12 @@ func TestAPIKeyPool(t *testing.T) {
t.Errorf("expected key1, got %s", k) t.Errorf("expected key1, got %s", k)
} }
emptyPool := NewAPIKeyPool(" ") emptyPool := NewAPIKeyPool([]string{" "})
if k := emptyPool.Get(); k != "" { if k := emptyPool.Get(); k != "" {
t.Errorf("expected empty string, got %s", k) t.Errorf("expected empty string, got %s", k)
} }
singlePool := NewAPIKeyPool("single") singlePool := NewAPIKeyPool([]string{"single"})
if k := singlePool.Get(); k != "single" { if k := singlePool.Get(); k != "single" {
t.Errorf("expected single, got %s", k) t.Errorf("expected single, got %s", k)
} }
@ -757,7 +757,7 @@ func TestWebTool_TavilySearch_Failover(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true, TavilyEnabled: true,
TavilyAPIKeys: "key1, key2", TavilyAPIKeys: []string{"key1", "key2"},
TavilyBaseURL: server.URL, TavilyBaseURL: server.URL,
TavilyMaxResults: 5, TavilyMaxResults: 5,
}) })