opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

0
cmd/picoclaw/__debug_bin1689969562 Normal file
View file

4
config/config.example.json
View file

@ -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": ""

6
pkg/agent/loop.go
View file

@ -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,

6
pkg/config/config.go
View file

@ -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"`
}

2
pkg/config/config_test.go
View file

@ -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 {

4
pkg/config/defaults.go
View file

@ -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,
},
},

6
pkg/migrate/sources/openclaw/openclaw_config.go
View file

@ -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,

464
pkg/tools/web.go
View file

@ -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,

108
pkg/tools/web_test.go
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
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)
}
}