From f18d661454c4a934db6de6d3d208f58921a32e49 Mon Sep 17 00:00:00 2001 From: stark Date: Mon, 2 Mar 2026 17:16:24 +0800 Subject: [PATCH] feat(web_search): add load balance and failover for api keys --- cmd/picoclaw/__debug_bin1689969562 | 0 config/config.example.json | 4 +- pkg/agent/loop.go | 6 +- pkg/config/config.go | 6 +- pkg/config/config_test.go | 2 +- pkg/config/defaults.go | 4 +- .../sources/openclaw/openclaw_config.go | 6 +- pkg/tools/web.go | 464 +++++++++++------- pkg/tools/web_test.go | 108 +++- 9 files changed, 400 insertions(+), 200 deletions(-) create mode 100644 cmd/picoclaw/__debug_bin1689969562 diff --git a/cmd/picoclaw/__debug_bin1689969562 b/cmd/picoclaw/__debug_bin1689969562 new file mode 100644 index 000000000..e69de29bb diff --git a/config/config.example.json b/config/config.example.json index d885ef94b..a0e4795e9 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -216,7 +216,7 @@ "web": { "brave": { "enabled": false, - "api_key": "YOUR_BRAVE_API_KEY", + "api_keys": "YOUR_BRAVE_API_KEY", "max_results": 5 }, "duckduckgo": { @@ -225,7 +225,7 @@ }, "perplexity": { "enabled": false, - "api_key": "pplx-xxx", + "api_keys": "pplx-xxx", "max_results": 5 }, "proxy": "" diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 00b0f096a..ea916ff72 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -100,16 +100,16 @@ func registerSharedTools( // Web tools searchTool, err := tools.NewWebSearchTool(tools.WebSearchToolOptions{ - BraveAPIKey: cfg.Tools.Web.Brave.APIKey, + BraveAPIKeys: cfg.Tools.Web.Brave.APIKeys, BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, BraveEnabled: cfg.Tools.Web.Brave.Enabled, - TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, + TavilyAPIKeys: cfg.Tools.Web.Tavily.APIKeys, TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, - PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, + PerplexityAPIKeys: cfg.Tools.Web.Perplexity.APIKeys, PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, Proxy: cfg.Tools.Web.Proxy, diff --git a/pkg/config/config.go b/pkg/config/config.go index 9f4769de4..df0c0ddf5 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -495,13 +495,13 @@ type GatewayConfig struct { type BraveConfig struct { 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"` } type TavilyConfig struct { 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"` MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"` } @@ -513,7 +513,7 @@ type DuckDuckGoConfig struct { type PerplexityConfig struct { 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"` } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 6af7c209e..0ef7dee35 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -293,7 +293,7 @@ func TestDefaultConfig_WebTools(t *testing.T) { if cfg.Tools.Web.Brave.MaxResults != 5 { 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") } if cfg.Tools.Web.DuckDuckGo.MaxResults != 5 { diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 44f4de7e9..9f8095334 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -319,7 +319,7 @@ func DefaultConfig() *Config { FetchLimitBytes: 10 * 1024 * 1024, // 10MB by default Brave: BraveConfig{ Enabled: false, - APIKey: "", + APIKeys: "", MaxResults: 5, }, DuckDuckGo: DuckDuckGoConfig{ @@ -328,7 +328,7 @@ func DefaultConfig() *Config { }, Perplexity: PerplexityConfig{ Enabled: false, - APIKey: "", + APIKeys: "", MaxResults: 5, }, }, diff --git a/pkg/migrate/sources/openclaw/openclaw_config.go b/pkg/migrate/sources/openclaw/openclaw_config.go index 39ad48fad..4d0149336 100644 --- a/pkg/migrate/sources/openclaw/openclaw_config.go +++ b/pkg/migrate/sources/openclaw/openclaw_config.go @@ -1043,12 +1043,12 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig { Web: config.WebToolsConfig{ Brave: config.BraveConfig{ Enabled: c.Web.Brave.Enabled, - APIKey: c.Web.Brave.APIKey, + APIKeys: c.Web.Brave.APIKey, MaxResults: c.Web.Brave.MaxResults, }, Tavily: config.TavilyConfig{ Enabled: c.Web.Tavily.Enabled, - APIKey: c.Web.Tavily.APIKey, + APIKeys: c.Web.Tavily.APIKey, BaseURL: c.Web.Tavily.BaseURL, MaxResults: c.Web.Tavily.MaxResults, }, @@ -1058,7 +1058,7 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig { }, Perplexity: config.PerplexityConfig{ Enabled: c.Web.Perplexity.Enabled, - APIKey: c.Web.Perplexity.APIKey, + APIKeys: c.Web.Perplexity.APIKey, MaxResults: c.Web.Perplexity.MaxResults, }, Proxy: c.Web.Proxy, diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 10498126b..52f172b2d 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -11,6 +11,7 @@ import ( "net/url" "regexp" "strings" + "sync/atomic" "time" ) @@ -76,77 +77,133 @@ func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, err 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 { Search(ctx context.Context, query string, count int) (string, error) } type BraveSearchProvider struct { - apiKey string - proxy string - client *http.Client + keyPool *APIKeyPool + proxy string + client *http.Client } func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", url.QueryEscape(query), count) - req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) + var lastErr error + maxAttempts := len(p.keyPool.keys) + if maxAttempts == 0 { + return "", errors.New("no api key available for Brave") } - req.Header.Set("Accept", "application/json") - req.Header.Set("X-Subscription-Token", p.apiKey) + for attempt := 0; attempt < maxAttempts; attempt++ { + apiKey := p.keyPool.Get() - resp, err := p.client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("failed to read response: %w", err) - } - - var searchResp struct { - Web struct { - Results []struct { - Title string `json:"title"` - URL string `json:"url"` - Description string `json:"description"` - } `json:"results"` - } `json:"web"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - // Log error body for debugging - fmt.Printf("Brave API Error Body: %s\n", string(body)) - return "", fmt.Errorf("failed to parse response: %w", err) - } - - results := searchResp.Web.Results - if len(results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s", query)) - for i, item := range results { - if i >= count { - break + req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) } - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) - if item.Description != "" { - lines = append(lines, fmt.Sprintf(" %s", item.Description)) + + req.Header.Set("Accept", "application/json") + req.Header.Set("X-Subscription-Token", apiKey) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + 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 { + Web struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Description string `json:"description"` + } `json:"results"` + } `json:"web"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + // Log error body for debugging + return "", fmt.Errorf("failed to parse response: %w", err) + } + + results := searchResp.Web.Results + if len(results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s", query)) + for i, item := range results { + if i >= count { + break + } + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) + if item.Description != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Description)) + } + } + + 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 { - apiKey string + keyPool *APIKeyPool baseURL string proxy string client *http.Client @@ -158,74 +215,97 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i searchURL = "https://api.tavily.com/search" } - payload := map[string]any{ - "api_key": p.apiKey, - "query": query, - "search_depth": "advanced", - "include_answer": false, - "include_images": false, - "include_raw_content": false, - "max_results": count, + var lastErr error + maxAttempts := len(p.keyPool.keys) + if maxAttempts == 0 { + return "", errors.New("no api key available for Tavily") } - bodyBytes, err := json.Marshal(payload) - if err != nil { - return "", fmt.Errorf("failed to marshal payload: %w", err) - } + for attempt := 0; attempt < maxAttempts; attempt++ { + apiKey := p.keyPool.Get() - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewBuffer(bodyBytes)) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("User-Agent", userAgent) - - resp, err := p.client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("tavily api error (status %d): %s", resp.StatusCode, string(body)) - } - - var searchResp struct { - Results []struct { - Title string `json:"title"` - URL string `json:"url"` - Content string `json:"content"` - } `json:"results"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - results := searchResp.Results - if len(results) == 0 { - return fmt.Sprintf("No results for: %s", query), nil - } - - var lines []string - lines = append(lines, fmt.Sprintf("Results for: %s (via Tavily)", query)) - for i, item := range results { - if i >= count { - break + payload := map[string]any{ + "api_key": apiKey, + "query": query, + "search_depth": "advanced", + "include_answer": false, + "include_images": false, + "include_raw_content": false, + "max_results": count, } - lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) - if item.Content != "" { - lines = append(lines, fmt.Sprintf(" %s", item.Content)) + + bodyBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal payload: %w", err) } + + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewBuffer(bodyBytes)) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", userAgent) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + lastErr = fmt.Errorf("failed to read response: %w", err) + continue + } + + if resp.StatusCode != http.StatusOK { + 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 { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + results := searchResp.Results + if len(results) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + var lines []string + lines = append(lines, fmt.Sprintf("Results for: %s (via Tavily)", query)) + for i, item := range results { + if i >= count { + break + } + lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.URL)) + if item.Content != "" { + lines = append(lines, fmt.Sprintf(" %s", item.Content)) + } + } + + + 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 { @@ -320,75 +400,97 @@ func stripTags(content string) string { } type PerplexitySearchProvider struct { - apiKey string - proxy string - client *http.Client + keyPool *APIKeyPool + proxy string + client *http.Client } func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { searchURL := "https://api.perplexity.ai/chat/completions" - payload := map[string]any{ - "model": "sonar", - "messages": []map[string]string{ - { - "role": "system", - "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", + 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{ + "model": "sonar", + "messages": []map[string]string{ + { + "role": "system", + "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", + }, + { + "role": "user", + "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), + }, }, - { - "role": "user", - "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), - }, - }, - "max_tokens": 1000, + "max_tokens": 1000, + } + + payloadBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) + if err != nil { + return "", fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("User-Agent", userAgent) + + resp, err := p.client.Do(req) + if err != nil { + lastErr = fmt.Errorf("request failed: %w", err) + continue + } + + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + + if err != nil { + lastErr = fmt.Errorf("failed to read response: %w", err) + continue + } + + if resp.StatusCode != http.StatusOK { + 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 { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + if err := json.Unmarshal(body, &searchResp); err != nil { + return "", fmt.Errorf("failed to parse response: %w", err) + } + + if len(searchResp.Choices) == 0 { + return fmt.Sprintf("No results for: %s", query), nil + } + + return fmt.Sprintf("Results for: %s (via Perplexity)\n%s", query, searchResp.Choices[0].Message.Content), nil } - payloadBytes, err := json.Marshal(payload) - if err != nil { - return "", fmt.Errorf("failed to marshal request: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, "POST", searchURL, strings.NewReader(string(payloadBytes))) - if err != nil { - return "", fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+p.apiKey) - req.Header.Set("User-Agent", userAgent) - - resp, err := p.client.Do(req) - if err != nil { - return "", fmt.Errorf("request failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("Perplexity API error: %s", string(body)) - } - - var searchResp struct { - Choices []struct { - Message struct { - Content string `json:"content"` - } `json:"message"` - } `json:"choices"` - } - - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("failed to parse response: %w", err) - } - - if len(searchResp.Choices) == 0 { - return fmt.Sprintf("No results for: %s", query), 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 { @@ -397,16 +499,16 @@ type WebSearchTool struct { } type WebSearchToolOptions struct { - BraveAPIKey string + BraveAPIKeys string BraveMaxResults int BraveEnabled bool - TavilyAPIKey string + TavilyAPIKeys string TavilyBaseURL string TavilyMaxResults int TavilyEnabled bool DuckDuckGoMaxResults int DuckDuckGoEnabled bool - PerplexityAPIKey string + PerplexityAPIKeys string PerplexityMaxResults int PerplexityEnabled bool Proxy string @@ -417,31 +519,31 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { maxResults := 5 // Priority: Perplexity > Brave > Tavily > DuckDuckGo - if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { + if opts.PerplexityEnabled && opts.PerplexityAPIKeys != "" { client, err := createHTTPClient(opts.Proxy, perplexityTimeout) if err != nil { 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 { maxResults = opts.PerplexityMaxResults } - } else if opts.BraveEnabled && opts.BraveAPIKey != "" { + } else if opts.BraveEnabled && opts.BraveAPIKeys != "" { client, err := createHTTPClient(opts.Proxy, searchTimeout) if err != nil { 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 { maxResults = opts.BraveMaxResults } - } else if opts.TavilyEnabled && opts.TavilyAPIKey != "" { + } else if opts.TavilyEnabled && opts.TavilyAPIKeys != "" { client, err := createHTTPClient(opts.Proxy, searchTimeout) if err != nil { return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) } provider = &TavilySearchProvider{ - apiKey: opts.TavilyAPIKey, + keyPool: NewAPIKeyPool(opts.TavilyAPIKeys), baseURL: opts.TavilyBaseURL, proxy: opts.Proxy, client: client, diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 8a8b88131..acf664069 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -249,7 +249,7 @@ func TestWebFetchTool_PayloadTooLarge(t *testing.T) { // TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing func TestWebTool_WebSearch_NoApiKey(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) + tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: ""}) if err != nil { 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 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 { t.Fatalf("Unexpected error: %v", err) } @@ -553,7 +553,7 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("perplexity", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ PerplexityEnabled: true, - PerplexityAPIKey: "k", + PerplexityAPIKeys: "k", PerplexityMaxResults: 3, Proxy: "http://127.0.0.1:7890", }) @@ -572,7 +572,7 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("brave", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ BraveEnabled: true, - BraveAPIKey: "k", + BraveAPIKeys: "k", BraveMaxResults: 3, Proxy: "http://127.0.0.1:7890", }) @@ -650,7 +650,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ TavilyEnabled: true, - TavilyAPIKey: "test-key", + TavilyAPIKeys: "test-key", TavilyBaseURL: server.URL, TavilyMaxResults: 5, }) @@ -681,3 +681,101 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { 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) + } +}