From c6865fe852f4e163767b78b6df72a06b5fdc6204 Mon Sep 17 00:00:00 2001 From: Vidish <57653368+ulolol@users.noreply.github.com> Date: Sun, 22 Feb 2026 22:00:14 +0530 Subject: [PATCH] feat: integrate Tavily search (#340) * feat: integrate Tavily search * fix: set include_raw_content to false in Tavily search as wealready get relevant data inside content * refactor: update Go type declarations to `any`, apply formatting fixes. --- README.ja.md | 29 ++++++++++-- README.md | 9 +++- README.zh.md | 27 ++++++----- pkg/config/config.go | 8 ++++ pkg/tools/registry_test.go | 20 ++++---- pkg/tools/web.go | 97 +++++++++++++++++++++++++++++++++++++- pkg/tools/web_test.go | 72 ++++++++++++++++++++++++++++ 7 files changed, 232 insertions(+), 30 deletions(-) diff --git a/README.ja.md b/README.ja.md index bb0bdfb28..3506c77c2 100644 --- a/README.ja.md +++ b/README.ja.md @@ -162,7 +162,7 @@ docker compose --profile gateway up -d > [!TIP] > `~/.picoclaw/config.json` に API キーを設定してください。 > API キーの取得先: [OpenRouter](https://openrouter.ai/keys) (LLM) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) (LLM) -> Web 検索は **任意** です - 無料の [Brave Search API](https://brave.com/search/api) (月 2000 クエリ無料) +> Web 検索は **任意** です - 無料の [Tavily API](https://tavily.com) (月 1000 クエリ無料) または [Brave Search API](https://brave.com/search/api) (月 2000 クエリ無料) **1. 初期化** @@ -193,14 +193,34 @@ picoclaw onboard "token": "YOUR_TELEGRAM_BOT_TOKEN", "allow_from": [] } + }, + "tools": { + "web": { + "search": { + "api_key": "YOUR_BRAVE_API_KEY", + "max_results": 5 + }, + "tavily": { + "enabled": false, + "api_key": "YOUR_TAVILY_API_KEY", + "max_results": 5 + } + }, + "cron": { + "exec_timeout_minutes": 5 + } + }, + "heartbeat": { + "enabled": true, + "interval": 30 } } ``` **3. API キーの取得** -- **LLM プロバイダー**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) · [Qwen](https://dashscope.console.aliyun.com) -- **Web 検索**（任意）: [Brave Search](https://brave.com/search/api) - 無料枠あり（月 2000 リクエスト） +- **LLM プロバイダー**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) +- **Web 検索**（任意）: [Tavily](https://tavily.com) - AI エージェント向けに最適化 (月 1000 リクエスト) · [Brave Search](https://brave.com/search/api) - 無料枠あり（月 2000 リクエスト） > **注意**: 完全な設定テンプレートは `config.example.json` を参照してください。 @@ -985,7 +1005,7 @@ Discord: https://discord.gg/V4sAZ9XWpN 検索 API キーをまだ設定していない場合、これは正常です。PicoClaw は手動検索用の便利なリンクを提供します。 Web 検索を有効にするには： -1. [https://brave.com/search/api](https://brave.com/search/api) で無料の API キーを取得（月 2000 クエリ無料） +1. [https://tavily.com](https://tavily.com) (月 1000 クエリ無料) または [https://brave.com/search/api](https://brave.com/search/api) で無料の API キーを取得（月 2000 クエリ無料） 2. `~/.picoclaw/config.json` に追加： ```json { @@ -1023,5 +1043,6 @@ Web 検索を有効にするには： | **Zhipu** | 月 200K トークン | 中国ユーザー向け最適 | | **Qwen** | 無料枠あり | 通義千問 (Qwen) | | **Brave Search** | 月 2000 クエリ | Web 検索機能 | +| **Tavily** | 月 1000 クエリ | AI エージェント検索最適化 | | **Groq** | 無料枠あり | 高速推論（Llama, Mixtral） | | **Cerebras** | 無料枠あり | 高速推論（Llama, Qwen など） | diff --git a/README.md b/README.md index de6fd87ea..825f57340 100644 --- a/README.md +++ b/README.md @@ -200,7 +200,7 @@ docker compose --profile gateway up -d > [!TIP] > Set your API key in `~/.picoclaw/config.json`. > Get API keys: [OpenRouter](https://openrouter.ai/keys) (LLM) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) (LLM) -> Web search is **optional** - get free [Brave Search API](https://brave.com/search/api) (2000 free queries/month) or use built-in auto fallback. +> Web Search is **optional** - get free [Tavily API](https://tavily.com) (1000 free queries/month) or [Brave Search API](https://brave.com/search/api) (2000 free queries/month) or use built-in auto fallback. **1. Initialize** @@ -240,6 +240,11 @@ picoclaw onboard "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 }, + "tavily": { + "enabled": false, + "api_key": "YOUR_TAVILY_API_KEY", + "max_results": 5 + }, "duckduckgo": { "enabled": true, "max_results": 5 @@ -254,7 +259,7 @@ picoclaw onboard **3. Get API Keys** * **LLM Provider**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) -* **Web Search** (optional): [Brave Search](https://brave.com/search/api) - Free tier available (2000 requests/month) +* **Web Search** (optional): [Tavily](https://tavily.com) - Optimized for AI Agents (1000 requests/month) · [Brave Search](https://brave.com/search/api) - Free tier available (2000 requests/month) > **Note**: See `config.example.json` for a complete configuration template. diff --git a/README.zh.md b/README.zh.md index 4d739c5eb..fd188567d 100644 --- a/README.zh.md +++ b/README.zh.md @@ -205,7 +205,7 @@ docker compose --profile gateway up -d > [!TIP] > 在 `~/.picoclaw/config.json` 中设置您的 API Key。 > 获取 API Key: [OpenRouter](https://openrouter.ai/keys) (LLM) · [Zhipu (智谱)](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) (LLM) -> 网络搜索是 **可选的** - 获取免费的 [Brave Search API](https://brave.com/search/api) (每月 2000 次免费查询) +> 网络搜索是 **可选的** - 获取免费的 [Tavily API](https://tavily.com) (每月 1000 次免费查询) 或 [Brave Search API](https://brave.com/search/api) (每月 2000 次免费查询) **1. 初始化 (Initialize)** @@ -246,8 +246,9 @@ picoclaw onboard "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 }, - "duckduckgo": { - "enabled": true, + "tavily": { + "enabled": false, + "api_key": "YOUR_TAVILY_API_KEY", "max_results": 5 } }, @@ -262,8 +263,8 @@ picoclaw onboard **3. 获取 API Key** -- **LLM 提供商**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) -- **网络搜索** (可选): [Brave Search](https://brave.com/search/api) - 提供免费层级 (2000 请求/月) +* **LLM 提供商**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) +* **网络搜索** (可选): [Tavily](https://tavily.com) - 专为 AI Agent 优化 (1000 请求/月) · [Brave Search](https://brave.com/search/api) - 提供免费层级 (2000 请求/月) > **注意**: 完整的配置模板请参考 `config.example.json`。 @@ -771,7 +772,7 @@ Discord: [https://discord.gg/V4sAZ9XWpN](https://discord.gg/V4sAZ9XWpN) 启用网络搜索： -1. 在 [https://brave.com/search/api](https://brave.com/search/api) 获取免费 API Key (每月 2000 次免费查询) +1. 在 [https://tavily.com](https://tavily.com) (1000 次免费) 或 [https://brave.com/search/api](https://brave.com/search/api) 获取免费 API Key (2000 次免费) 2. 添加到 `~/.picoclaw/config.json`: ```json @@ -804,10 +805,10 @@ Discord: [https://discord.gg/V4sAZ9XWpN](https://discord.gg/V4sAZ9XWpN) ## 📝 API Key 对比 -| 服务 | 免费层级 | 适用场景 | -| ---------------- | -------------- | ----------------------------- | -| **OpenRouter** | 200K tokens/月 | 多模型聚合 (Claude, GPT-4 等) | -| **智谱 (Zhipu)** | 200K tokens/月 | 最适合中国用户 | -| **Brave Search** | 2000 次查询/月 | 网络搜索功能 | -| **Groq** | 提供免费层级 | 极速推理 (Llama, Mixtral) | -| **Cerebras** | 提供免费层级 | 极速推理 (Llama, Qwen 等) | +| 服务 | 免费层级 | 适用场景 | +| --- | --- | --- | +| **OpenRouter** | 200K tokens/月 | 多模型聚合 (Claude, GPT-4 等) | +| **智谱 (Zhipu)** | 200K tokens/月 | 最适合中国用户 | +| **Brave Search** | 2000 次查询/月 | 网络搜索功能 | +| **Tavily** | 1000 次查询/月 | AI Agent 搜索优化 | +| **Groq** | 提供免费层级 | 极速推理 (Llama, Mixtral) | diff --git a/pkg/config/config.go b/pkg/config/config.go index 20556011a..036021e49 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -418,6 +418,13 @@ type BraveConfig struct { 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"` + BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_TAVILY_BASE_URL"` + MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_TAVILY_MAX_RESULTS"` +} + type DuckDuckGoConfig struct { Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_DUCKDUCKGO_ENABLED"` MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_DUCKDUCKGO_MAX_RESULTS"` @@ -431,6 +438,7 @@ type PerplexityConfig struct { type WebToolsConfig struct { Brave BraveConfig `json:"brave"` + Tavily TavilyConfig `json:"tavily"` DuckDuckGo DuckDuckGoConfig `json:"duckduckgo"` Perplexity PerplexityConfig `json:"perplexity"` } diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index 33978e543..8ae13b20c 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -14,14 +14,14 @@ import ( type mockRegistryTool struct { name string desc string - params map[string]interface{} + params map[string]any result *ToolResult } -func (m *mockRegistryTool) Name() string { return m.name } -func (m *mockRegistryTool) Description() string { return m.desc } -func (m *mockRegistryTool) Parameters() map[string]interface{} { return m.params } -func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]interface{}) *ToolResult { +func (m *mockRegistryTool) Name() string { return m.name } +func (m *mockRegistryTool) Description() string { return m.desc } +func (m *mockRegistryTool) Parameters() map[string]any { return m.params } +func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult { return m.result } @@ -51,7 +51,7 @@ func newMockTool(name, desc string) *mockRegistryTool { return &mockRegistryTool{ name: name, desc: desc, - params: map[string]interface{}{"type": "object"}, + params: map[string]any{"type": "object"}, result: SilentResult("ok"), } } @@ -109,7 +109,7 @@ func TestToolRegistry_Execute_Success(t *testing.T) { r.Register(&mockRegistryTool{ name: "greet", desc: "says hello", - params: map[string]interface{}{}, + params: map[string]any{}, result: SilentResult("hello"), }) @@ -203,7 +203,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { if defs[0]["type"] != "function" { t.Errorf("expected type 'function', got %v", defs[0]["type"]) } - fn, ok := defs[0]["function"].(map[string]interface{}) + fn, ok := defs[0]["function"].(map[string]any) if !ok { t.Fatal("expected 'function' key to be a map") } @@ -217,7 +217,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { func TestToolRegistry_ToProviderDefs(t *testing.T) { r := NewToolRegistry() - params := map[string]interface{}{"type": "object", "properties": map[string]interface{}{}} + params := map[string]any{"type": "object", "properties": map[string]any{}} r.Register(&mockRegistryTool{ name: "beta", desc: "tool B", @@ -310,7 +310,7 @@ func TestToolToSchema(t *testing.T) { if schema["type"] != "function" { t.Errorf("expected type 'function', got %v", schema["type"]) } - fn, ok := schema["function"].(map[string]interface{}) + fn, ok := schema["function"].(map[string]any) if !ok { t.Fatal("expected 'function' to be a map") } diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 301e00daf..059437889 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -1,6 +1,7 @@ package tools import ( + "bytes" "context" "encoding/json" "fmt" @@ -84,6 +85,88 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in return strings.Join(lines, "\n"), nil } +type TavilySearchProvider struct { + apiKey string + baseURL string +} + +func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { + searchURL := p.baseURL + if searchURL == "" { + 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, + } + + 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) + + client := &http.Client{Timeout: 10 * time.Second} + resp, err := 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 + } + 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 +} + type DuckDuckGoSearchProvider struct{} func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { @@ -256,6 +339,10 @@ type WebSearchToolOptions struct { BraveAPIKey string BraveMaxResults int BraveEnabled bool + TavilyAPIKey string + TavilyBaseURL string + TavilyMaxResults int + TavilyEnabled bool DuckDuckGoMaxResults int DuckDuckGoEnabled bool PerplexityAPIKey string @@ -267,7 +354,7 @@ func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool { var provider SearchProvider maxResults := 5 - // Priority: Perplexity > Brave > DuckDuckGo + // Priority: Perplexity > Brave > Tavily > DuckDuckGo if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey} if opts.PerplexityMaxResults > 0 { @@ -278,6 +365,14 @@ func NewWebSearchTool(opts WebSearchToolOptions) *WebSearchTool { if opts.BraveMaxResults > 0 { maxResults = opts.BraveMaxResults } + } else if opts.TavilyEnabled && opts.TavilyAPIKey != "" { + provider = &TavilySearchProvider{ + apiKey: opts.TavilyAPIKey, + baseURL: opts.TavilyBaseURL, + } + if opts.TavilyMaxResults > 0 { + maxResults = opts.TavilyMaxResults + } } else if opts.DuckDuckGoEnabled { provider = &DuckDuckGoSearchProvider{} if opts.DuckDuckGoMaxResults > 0 { diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index d999d8958..75e0d8d16 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -333,3 +333,75 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) { t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) } } + +// TestWebTool_TavilySearch_Success verifies successful Tavily search +func TestWebTool_TavilySearch_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) + } + + // Verify payload + var payload map[string]any + json.NewDecoder(r.Body).Decode(&payload) + if payload["api_key"] != "test-key" { + t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) + } + if payload["query"] != "test query" { + t.Errorf("Expected query 'test query', got %v", payload["query"]) + } + + // Return mock response + response := map[string]any{ + "results": []map[string]any{ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "content": "Content for result 1", + }, + { + "title": "Test Result 2", + "url": "https://example.com/2", + "content": "Content for result 2", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + tool := NewWebSearchTool(WebSearchToolOptions{ + TavilyEnabled: true, + TavilyAPIKey: "test-key", + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, + }) + + ctx := context.Background() + args := map[string]any{ + "query": "test query", + } + + result := tool.Execute(ctx, args) + + // Success should not be an error + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + + // ForUser should contain result titles and URLs + if !strings.Contains(result.ForUser, "Test Result 1") || + !strings.Contains(result.ForUser, "https://example.com/1") { + t.Errorf("Expected results in output, got: %s", result.ForUser) + } + + // Should mention via Tavily + if !strings.Contains(result.ForUser, "via Tavily") { + t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) + } +}