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-03 20:59:55 +08:00
parent f18d661454
commit 98e5ff81d8
8 changed files with 82 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
cmd/picoclaw/__debug_bin1689969562
View file

10
config/config.example.json
View file

@ -216,7 +216,10 @@
"web": {
"brave": {
"enabled": false,
"api_keys": "YOUR_BRAVE_API_KEY",
"api_key": "YOUR_BRAVE_API_KEY",
"api_keys": [
"YOUR_BRAVE_API_KEY"
],
"max_results": 5
},
"duckduckgo": {
@ -225,7 +228,10 @@
},
"perplexity": {
"enabled": false,
"api_keys": "pplx-xxx",
"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{
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,
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,
TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults,
TavilyEnabled: cfg.Tools.Web.Tavily.Enabled,
DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults,
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,
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
Proxy: cfg.Tools.Web.Proxy,

33
pkg/config/config.go
View file

@ -4,6 +4,7 @@ import (
"encoding/json"
"fmt"
"os"
"strings"
"sync/atomic"
"github.com/caarlos0/env/v11"
@ -495,13 +496,15 @@ type GatewayConfig struct {
type BraveConfig struct {
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"`
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"`
APIKeys string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_TAVILY_API_KEYS"`
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 +516,8 @@ type DuckDuckGoConfig struct {
type PerplexityConfig struct {
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"`
APIKeys []string `json:"api_keys" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_API_KEYS"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"`
}
@ -779,3 +783,26 @@ func (c *Config) ValidateModelList() error {
}
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
}

12
pkg/config/defaults.go
View file

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

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

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

21
pkg/tools/web.go
View file

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

18
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, BraveAPIKeys: ""})
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: nil})
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, BraveAPIKeys: "test-key", BraveMaxResults: 5})
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: []string{"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,
PerplexityAPIKeys: "k",
PerplexityAPIKeys: []string{"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,
BraveAPIKeys: "k",
BraveAPIKeys: []string{"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,
TavilyAPIKeys: "test-key",
TavilyAPIKeys: []string{"test-key"},
TavilyBaseURL: server.URL,
TavilyMaxResults: 5,
})
@ -683,7 +683,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
}
func TestAPIKeyPool(t *testing.T) {
pool := NewAPIKeyPool("key1, key2 ,key3")
pool := NewAPIKeyPool([]string{"key1", " key2 ", "key3"})
if len(pool.keys) != 3 {
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)
}
emptyPool := NewAPIKeyPool(" ")
emptyPool := NewAPIKeyPool([]string{" "})
if k := emptyPool.Get(); k != "" {
t.Errorf("expected empty string, got %s", k)
}
singlePool := NewAPIKeyPool("single")
singlePool := NewAPIKeyPool([]string{"single"})
if k := singlePool.Get(); k != "single" {
t.Errorf("expected single, got %s", k)
}
@ -757,7 +757,7 @@ func TestWebTool_TavilySearch_Failover(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true,
TavilyAPIKeys: "key1, key2",
TavilyAPIKeys: []string{"key1", "key2"},
TavilyBaseURL: server.URL,
TavilyMaxResults: 5,
})