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,

156
pkg/tools/web.go
View file

@ -11,6 +11,7 @@ import (
"net/url"
"regexp"
"strings"
"sync/atomic"
"time"
)
@ -76,12 +77,43 @@ 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
keyPool *APIKeyPool
proxy string
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",
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)
if err != nil {
return "", fmt.Errorf("failed to create request: %w", err)
}
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)
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)
resp.Body.Close()
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 {
@ -121,7 +176,6 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
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)
}
@ -143,10 +197,13 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
}
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,8 +215,17 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
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{
"api_key": p.apiKey,
"api_key": apiKey,
"query": query,
"search_depth": "advanced",
"include_answer": false,
@ -183,17 +249,27 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
resp, err := p.client.Do(req)
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)
resp.Body.Close()
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 {
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 {
@ -225,7 +301,11 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
}
}
return strings.Join(lines, "\n"), nil
}
return "", fmt.Errorf("all api keys failed, last error: %w", lastErr)
}
type DuckDuckGoSearchProvider struct {
@ -320,7 +400,7 @@ func stripTags(content string) string {
}
type PerplexitySearchProvider struct {
apiKey string
keyPool *APIKeyPool
proxy string
client *http.Client
}
@ -328,6 +408,15 @@ type PerplexitySearchProvider struct {
func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
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{
"model": "sonar",
"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("Authorization", "Bearer "+p.apiKey)
req.Header.Set("Authorization", "Bearer "+apiKey)
req.Header.Set("User-Agent", userAgent)
resp, err := p.client.Do(req)
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)
resp.Body.Close()
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 {
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 {
@ -389,6 +488,9 @@ 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.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)
}
}