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

This commit is contained in:
stark 2026-03-02 17:16:24 +08:00
parent d5370c9605
commit f18d661454
9 changed files with 400 additions and 200 deletions

View file

View file

@ -216,7 +216,7 @@
"web": { "web": {
"brave": { "brave": {
"enabled": false, "enabled": false,
"api_key": "YOUR_BRAVE_API_KEY", "api_keys": "YOUR_BRAVE_API_KEY",
"max_results": 5 "max_results": 5
}, },
"duckduckgo": { "duckduckgo": {
@ -225,7 +225,7 @@
}, },
"perplexity": { "perplexity": {
"enabled": false, "enabled": false,
"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{
BraveAPIKey: cfg.Tools.Web.Brave.APIKey, BraveAPIKeys: 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,
TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, TavilyAPIKeys: 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,
PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, PerplexityAPIKeys: 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

@ -495,13 +495,13 @@ 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"`
APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEY"` APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_BRAVE_MAX_RESULTS"` 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"`
APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEY"` APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEYS"`
BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_TAVILY_BASE_URL"` BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_TAVILY_BASE_URL"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"` MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"`
} }
@ -513,7 +513,7 @@ 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"`
APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEY"` APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"` MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"`
} }

View file

@ -293,7 +293,7 @@ func TestDefaultConfig_WebTools(t *testing.T) {
if cfg.Tools.Web.Brave.MaxResults != 5 { if cfg.Tools.Web.Brave.MaxResults != 5 {
t.Error("Expected Brave MaxResults 5, got ", cfg.Tools.Web.Brave.MaxResults) t.Error("Expected Brave MaxResults 5, got ", cfg.Tools.Web.Brave.MaxResults)
} }
if cfg.Tools.Web.Brave.APIKey != "" { if cfg.Tools.Web.Brave.APIKeys != "" {
t.Error("Brave API key should be empty by default") t.Error("Brave API key should be empty by default")
} }
if cfg.Tools.Web.DuckDuckGo.MaxResults != 5 { if cfg.Tools.Web.DuckDuckGo.MaxResults != 5 {

View file

@ -319,7 +319,7 @@ 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,
APIKey: "", APIKeys: "",
MaxResults: 5, MaxResults: 5,
}, },
DuckDuckGo: DuckDuckGoConfig{ DuckDuckGo: DuckDuckGoConfig{
@ -328,7 +328,7 @@ func DefaultConfig() *Config {
}, },
Perplexity: PerplexityConfig{ Perplexity: PerplexityConfig{
Enabled: false, Enabled: false,
APIKey: "", APIKeys: "",
MaxResults: 5, MaxResults: 5,
}, },
}, },

View file

@ -1043,12 +1043,12 @@ 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,
APIKey: c.Web.Brave.APIKey, APIKeys: c.Web.Brave.APIKey,
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,
APIKey: c.Web.Tavily.APIKey, APIKeys: 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 +1058,7 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig {
}, },
Perplexity: config.PerplexityConfig{ Perplexity: config.PerplexityConfig{
Enabled: c.Web.Perplexity.Enabled, Enabled: c.Web.Perplexity.Enabled,
APIKey: c.Web.Perplexity.APIKey, APIKeys: c.Web.Perplexity.APIKey,
MaxResults: c.Web.Perplexity.MaxResults, MaxResults: c.Web.Perplexity.MaxResults,
}, },
Proxy: c.Web.Proxy, Proxy: c.Web.Proxy,

View file

@ -11,6 +11,7 @@ import (
"net/url" "net/url"
"regexp" "regexp"
"strings" "strings"
"sync/atomic"
"time" "time"
) )
@ -76,12 +77,43 @@ func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, err
return client, nil return client, nil
} }
type APIKeyPool struct {
keys []string
current uint32
}
func NewAPIKeyPool(keysStr string) *APIKeyPool {
var keys []string
for _, k := range strings.Split(keysStr, ",") {
if trimmed := strings.TrimSpace(k); trimmed != "" {
keys = append(keys, trimmed)
}
}
return &APIKeyPool{
keys: keys,
}
}
func (p *APIKeyPool) Get() string {
if len(p.keys) == 0 {
return ""
}
if len(p.keys) == 1 {
return p.keys[0]
}
idx := atomic.AddUint32(&p.current, 1) - 1
if idx >= uint32(len(p.keys))-1 {
atomic.CompareAndSwapUint32(&p.current, idx+1, 0)
}
return p.keys[idx%uint32(len(p.keys))]
}
type SearchProvider interface { type SearchProvider interface {
Search(ctx context.Context, query string, count int) (string, error) Search(ctx context.Context, query string, count int) (string, error)
} }
type BraveSearchProvider struct { type BraveSearchProvider struct {
apiKey string keyPool *APIKeyPool
proxy string proxy string
client *http.Client client *http.Client
} }
@ -90,23 +122,46 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d",
url.QueryEscape(query), count) url.QueryEscape(query), count)
var lastErr error
maxAttempts := len(p.keyPool.keys)
if maxAttempts == 0 {
return "", errors.New("no api key available for Brave")
}
for attempt := 0; attempt < maxAttempts; attempt++ {
apiKey := p.keyPool.Get()
req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil)
if err != nil { if err != nil {
return "", fmt.Errorf("failed to create request: %w", err) return "", fmt.Errorf("failed to create request: %w", err)
} }
req.Header.Set("Accept", "application/json") req.Header.Set("Accept", "application/json")
req.Header.Set("X-Subscription-Token", p.apiKey) req.Header.Set("X-Subscription-Token", apiKey)
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) lastErr = fmt.Errorf("request failed: %w", err)
continue
} }
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil { if err != nil {
return "", fmt.Errorf("failed to read response: %w", err) lastErr = fmt.Errorf("failed to read response: %w", err)
continue
}
if resp.StatusCode != http.StatusOK {
lastErr = fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
if resp.StatusCode == http.StatusTooManyRequests ||
resp.StatusCode == http.StatusUnauthorized ||
resp.StatusCode == http.StatusForbidden ||
resp.StatusCode >= 500 {
continue
}
return "", lastErr
} }
var searchResp struct { var searchResp struct {
@ -121,7 +176,6 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
if err := json.Unmarshal(body, &searchResp); err != nil { if err := json.Unmarshal(body, &searchResp); err != nil {
// Log error body for debugging // Log error body for debugging
fmt.Printf("Brave API Error Body: %s\n", string(body))
return "", fmt.Errorf("failed to parse response: %w", err) return "", fmt.Errorf("failed to parse response: %w", err)
} }
@ -145,8 +199,11 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
return strings.Join(lines, "\n"), nil return strings.Join(lines, "\n"), nil
} }
return "", fmt.Errorf("all api keys failed, last error: %w", lastErr)
}
type TavilySearchProvider struct { type TavilySearchProvider struct {
apiKey string keyPool *APIKeyPool
baseURL string baseURL string
proxy string proxy string
client *http.Client client *http.Client
@ -158,8 +215,17 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
searchURL = "https://api.tavily.com/search" searchURL = "https://api.tavily.com/search"
} }
var lastErr error
maxAttempts := len(p.keyPool.keys)
if maxAttempts == 0 {
return "", errors.New("no api key available for Tavily")
}
for attempt := 0; attempt < maxAttempts; attempt++ {
apiKey := p.keyPool.Get()
payload := map[string]any{ payload := map[string]any{
"api_key": p.apiKey, "api_key": apiKey,
"query": query, "query": query,
"search_depth": "advanced", "search_depth": "advanced",
"include_answer": false, "include_answer": false,
@ -183,17 +249,27 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) lastErr = fmt.Errorf("request failed: %w", err)
continue
} }
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil { if err != nil {
return "", fmt.Errorf("failed to read response: %w", err) lastErr = fmt.Errorf("failed to read response: %w", err)
continue
} }
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("tavily api error (status %d): %s", resp.StatusCode, string(body)) lastErr = fmt.Errorf("tavily api error (status %d): %s", resp.StatusCode, string(body))
if resp.StatusCode == http.StatusTooManyRequests ||
resp.StatusCode == http.StatusUnauthorized ||
resp.StatusCode == http.StatusForbidden ||
resp.StatusCode >= 500 {
continue
}
return "", lastErr
} }
var searchResp struct { var searchResp struct {
@ -225,9 +301,13 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
} }
} }
return strings.Join(lines, "\n"), nil return strings.Join(lines, "\n"), nil
} }
return "", fmt.Errorf("all api keys failed, last error: %w", lastErr)
}
type DuckDuckGoSearchProvider struct { type DuckDuckGoSearchProvider struct {
proxy string proxy string
client *http.Client client *http.Client
@ -320,7 +400,7 @@ func stripTags(content string) string {
} }
type PerplexitySearchProvider struct { type PerplexitySearchProvider struct {
apiKey string keyPool *APIKeyPool
proxy string proxy string
client *http.Client client *http.Client
} }
@ -328,6 +408,15 @@ type PerplexitySearchProvider struct {
func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
searchURL := "https://api.perplexity.ai/chat/completions" searchURL := "https://api.perplexity.ai/chat/completions"
var lastErr error
maxAttempts := len(p.keyPool.keys)
if maxAttempts == 0 {
return "", errors.New("no api key available for Perplexity")
}
for attempt := 0; attempt < maxAttempts; attempt++ {
apiKey := p.keyPool.Get()
payload := map[string]any{ payload := map[string]any{
"model": "sonar", "model": "sonar",
"messages": []map[string]string{ "messages": []map[string]string{
@ -354,22 +443,32 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
} }
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("Authorization", "Bearer "+apiKey)
req.Header.Set("User-Agent", userAgent) req.Header.Set("User-Agent", userAgent)
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) lastErr = fmt.Errorf("request failed: %w", err)
continue
} }
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil { if err != nil {
return "", fmt.Errorf("failed to read response: %w", err) lastErr = fmt.Errorf("failed to read response: %w", err)
continue
} }
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("Perplexity API error: %s", string(body)) lastErr = fmt.Errorf("Perplexity API error: %s", string(body))
if resp.StatusCode == http.StatusTooManyRequests ||
resp.StatusCode == http.StatusUnauthorized ||
resp.StatusCode == http.StatusForbidden ||
resp.StatusCode >= 500 {
continue
}
return "", lastErr
} }
var searchResp struct { var searchResp struct {
@ -391,22 +490,25 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil
} }
return "", fmt.Errorf("all api keys failed, last error: %w", lastErr)
}
type WebSearchTool struct { type WebSearchTool struct {
provider SearchProvider provider SearchProvider
maxResults int maxResults int
} }
type WebSearchToolOptions struct { type WebSearchToolOptions struct {
BraveAPIKey string BraveAPIKeys string
BraveMaxResults int BraveMaxResults int
BraveEnabled bool BraveEnabled bool
TavilyAPIKey string TavilyAPIKeys string
TavilyBaseURL string TavilyBaseURL string
TavilyMaxResults int TavilyMaxResults int
TavilyEnabled bool TavilyEnabled bool
DuckDuckGoMaxResults int DuckDuckGoMaxResults int
DuckDuckGoEnabled bool DuckDuckGoEnabled bool
PerplexityAPIKey string PerplexityAPIKeys string
PerplexityMaxResults int PerplexityMaxResults int
PerplexityEnabled bool PerplexityEnabled bool
Proxy string Proxy string
@ -417,31 +519,31 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
maxResults := 5 maxResults := 5
// Priority: Perplexity > Brave > Tavily > DuckDuckGo // Priority: Perplexity > Brave > Tavily > DuckDuckGo
if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { if opts.PerplexityEnabled && opts.PerplexityAPIKeys != "" {
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)
} }
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client} provider = &PerplexitySearchProvider{keyPool: NewAPIKeyPool(opts.PerplexityAPIKeys), proxy: opts.Proxy, client: client}
if opts.PerplexityMaxResults > 0 { if opts.PerplexityMaxResults > 0 {
maxResults = opts.PerplexityMaxResults maxResults = opts.PerplexityMaxResults
} }
} else if opts.BraveEnabled && opts.BraveAPIKey != "" { } else if opts.BraveEnabled && opts.BraveAPIKeys != "" {
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)
} }
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client} provider = &BraveSearchProvider{keyPool: NewAPIKeyPool(opts.BraveAPIKeys), proxy: opts.Proxy, client: client}
if opts.BraveMaxResults > 0 { if opts.BraveMaxResults > 0 {
maxResults = opts.BraveMaxResults maxResults = opts.BraveMaxResults
} }
} else if opts.TavilyEnabled && opts.TavilyAPIKey != "" { } else if opts.TavilyEnabled && opts.TavilyAPIKeys != "" {
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)
} }
provider = &TavilySearchProvider{ provider = &TavilySearchProvider{
apiKey: opts.TavilyAPIKey, keyPool: NewAPIKeyPool(opts.TavilyAPIKeys),
baseURL: opts.TavilyBaseURL, baseURL: opts.TavilyBaseURL,
proxy: opts.Proxy, proxy: opts.Proxy,
client: client, client: client,

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, BraveAPIKey: ""}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: ""})
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, BraveAPIKey: "test-key", BraveMaxResults: 5}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: "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,
PerplexityAPIKey: "k", PerplexityAPIKeys: "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,
BraveAPIKey: "k", BraveAPIKeys: "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,
TavilyAPIKey: "test-key", TavilyAPIKeys: "test-key",
TavilyBaseURL: server.URL, TavilyBaseURL: server.URL,
TavilyMaxResults: 5, TavilyMaxResults: 5,
}) })
@ -681,3 +681,101 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser)
} }
} }
func TestAPIKeyPool(t *testing.T) {
pool := NewAPIKeyPool("key1, key2 ,key3")
if len(pool.keys) != 3 {
t.Fatalf("expected 3 keys, got %d", len(pool.keys))
}
if pool.keys[0] != "key1" || pool.keys[1] != "key2" || pool.keys[2] != "key3" {
t.Fatalf("unexpected keys: %v", pool.keys)
}
// Test Get()
if k := pool.Get(); k != "key1" {
t.Errorf("expected key1, got %s", k)
}
if k := pool.Get(); k != "key2" {
t.Errorf("expected key2, got %s", k)
}
if k := pool.Get(); k != "key3" {
t.Errorf("expected key3, got %s", k)
}
if k := pool.Get(); k != "key1" {
t.Errorf("expected key1, got %s", k)
}
emptyPool := NewAPIKeyPool(" ")
if k := emptyPool.Get(); k != "" {
t.Errorf("expected empty string, got %s", k)
}
singlePool := NewAPIKeyPool("single")
if k := singlePool.Get(); k != "single" {
t.Errorf("expected single, got %s", k)
}
if k := singlePool.Get(); k != "single" {
t.Errorf("expected single, got %s", k)
}
}
func TestWebTool_TavilySearch_Failover(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var payload map[string]any
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
t.Fatalf("failed to decode payload: %v", err)
}
apiKey := payload["api_key"].(string)
if apiKey == "key1" {
w.WriteHeader(http.StatusTooManyRequests)
w.Write([]byte("Rate limited"))
return
}
if apiKey == "key2" {
// Success
response := map[string]any{
"results": []map[string]any{
{
"title": "Success Result",
"url": "https://example.com/success",
"content": "Success content",
},
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(response)
return
}
w.WriteHeader(http.StatusBadRequest)
}))
defer server.Close()
tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true,
TavilyAPIKeys: "key1, key2",
TavilyBaseURL: server.URL,
TavilyMaxResults: 5,
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
ctx := context.Background()
args := map[string]any{
"query": "test query",
}
result := tool.Execute(ctx, args)
if result.IsError {
t.Errorf("Expected success, got Error: %s", result.ForLLM)
}
if !strings.Contains(result.ForUser, "Success Result") {
t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser)
}
}