From a59d6ac6db4a41a367b27272d5ab343b042cbf08 Mon Sep 17 00:00:00 2001 From: vrijmetse <56536692+vrijmetse@users.noreply.github.com> Date: Tue, 27 Jan 2026 21:46:41 +0700 Subject: [PATCH] refactor(tools): add multi-provider support to web_search tool (#795) * feat(tools): add Google Custom Search as alternative to Brave Search Adds google_search tool using Google Custom Search API as an alternative to the existing web_search tool (Brave Search). Changes: - Add google_search_tool with full implementation - Register Google credentials (GOOGLE_API_KEY, GOOGLE_CSE_ID) - Register tool in tools/__init__.py - Add README with setup instructions Closes #793 * test(tools): add unit tests for google_search tool Adds 7 tests mirroring web_search_tool test patterns: - Missing API key error handling - Missing CSE ID error handling - Empty query validation - Long query validation - num_results clamping - Default parameters - Custom language/country parameters All tests pass. * refactor(tools): add multi-provider support to web_search tool BREAKING CHANGE: None - backward compatible. Brave remains default. - Add Google Custom Search as alternative provider in web_search - Add 'provider' parameter: 'auto' (default), 'google', 'brave' - Auto mode tries Brave first for backward compatibility - Remove separate google_search_tool (consolidated into web_search) - Update tests to cover multi-provider functionality (13 tests) - Update README documentation Users with BRAVE_SEARCH_API_KEY: No changes needed Users with GOOGLE_API_KEY + GOOGLE_CSE_ID: Can use provider='google' Users with both: Brave preferred by default, use provider='google' to force Closes #793 * feat(tools): fixed readme --------- Co-authored-by: Mustafa Abdat --- tools/README.md | 8 +- tools/src/aden_tools/credentials/search.py | 29 ++- tools/src/aden_tools/tools/__init__.py | 13 +- .../tools/web_search_tool/README.md | 47 +++- .../tools/web_search_tool/web_search_tool.py | 241 +++++++++++++----- tools/tests/tools/test_web_search_tool.py | 120 ++++++++- 6 files changed, 355 insertions(+), 103 deletions(-) diff --git a/tools/README.md b/tools/README.md index 98d30ad1..9f5b9787 100644 --- a/tools/README.md +++ b/tools/README.md @@ -25,7 +25,11 @@ cp .env.example .env | Variable | Required For | Get Key | | ---------------------- | ----------------------------- | ------------------------------------------------------- | | `ANTHROPIC_API_KEY` | MCP server startup, LLM nodes | [console.anthropic.com](https://console.anthropic.com/) | -| `BRAVE_SEARCH_API_KEY` | `web_search` tool | [brave.com/search/api](https://brave.com/search/api/) | +| `BRAVE_SEARCH_API_KEY` | `web_search` tool (Brave) | [brave.com/search/api](https://brave.com/search/api/) | +| `GOOGLE_API_KEY` | `web_search` tool (Google) | [console.cloud.google.com](https://console.cloud.google.com/) | +| `GOOGLE_CSE_ID` | `web_search` tool (Google) | [programmablesearchengine.google.com](https://programmablesearchengine.google.com/) | + +> **Note:** `web_search` supports multiple providers. Set either Brave OR Google credentials. Brave is preferred for backward compatibility. Alternatively, export as environment variables: @@ -68,7 +72,7 @@ python mcp_server.py | `apply_patch` | Apply unified patches to files | | `grep_search` | Search file contents with regex | | `execute_command_tool` | Execute shell commands | -| `web_search` | Search the web using Brave Search API | +| `web_search` | Search the web (Google or Brave, auto-detected) | | `web_scrape` | Scrape and extract content from webpages | | `pdf_read` | Read and extract text from PDF files | diff --git a/tools/src/aden_tools/credentials/search.py b/tools/src/aden_tools/credentials/search.py index e70406b0..0f180c74 100644 --- a/tools/src/aden_tools/credentials/search.py +++ b/tools/src/aden_tools/credentials/search.py @@ -3,6 +3,7 @@ Search tool credentials. Contains credentials for search providers like Brave Search, Google, Bing, etc. """ + from .base import CredentialSpec SEARCH_CREDENTIALS = { @@ -15,14 +16,22 @@ SEARCH_CREDENTIALS = { help_url="https://brave.com/search/api/", description="API key for Brave Search", ), - # Future search providers: - # "google_search": CredentialSpec( - # env_var="GOOGLE_SEARCH_API_KEY", - # tools=["google_search"], - # node_types=[], - # required=True, - # startup_required=False, - # help_url="https://developers.google.com/custom-search/v1/overview", - # description="API key for Google Custom Search", - # ), + "google_search": CredentialSpec( + env_var="GOOGLE_API_KEY", + tools=["google_search"], + node_types=[], + required=True, + startup_required=False, + help_url="https://console.cloud.google.com/", + description="API key for Google Custom Search", + ), + "google_cse": CredentialSpec( + env_var="GOOGLE_CSE_ID", + tools=["google_search"], + node_types=[], + required=True, + startup_required=False, + help_url="https://programmablesearchengine.google.com/", + description="Google Custom Search Engine ID", + ), } diff --git a/tools/src/aden_tools/tools/__init__.py b/tools/src/aden_tools/tools/__init__.py index 73169c07..6ab64e73 100644 --- a/tools/src/aden_tools/tools/__init__.py +++ b/tools/src/aden_tools/tools/__init__.py @@ -10,6 +10,7 @@ Usage: credentials = CredentialManager() register_all_tools(mcp, credentials=credentials) """ + from typing import List, Optional, TYPE_CHECKING from fastmcp import FastMCP @@ -27,11 +28,15 @@ from .pdf_read_tool import register_tools as register_pdf_read from .file_system_toolkits.view_file import register_tools as register_view_file from .file_system_toolkits.write_to_file import register_tools as register_write_to_file from .file_system_toolkits.list_dir import register_tools as register_list_dir -from .file_system_toolkits.replace_file_content import register_tools as register_replace_file_content +from .file_system_toolkits.replace_file_content import ( + register_tools as register_replace_file_content, +) from .file_system_toolkits.apply_diff import register_tools as register_apply_diff from .file_system_toolkits.apply_patch import register_tools as register_apply_patch from .file_system_toolkits.grep_search import register_tools as register_grep_search -from .file_system_toolkits.execute_command_tool import register_tools as register_execute_command +from .file_system_toolkits.execute_command_tool import ( + register_tools as register_execute_command, +) from .csv_tool import register_tools as register_csv @@ -56,9 +61,7 @@ def register_all_tools( register_pdf_read(mcp) # Tools that need credentials (pass credentials if provided) - # web_search handles both credential sources internally: - # - If credentials provided: uses credentials.get("brave_search") - # - If credentials is None: falls back to os.getenv("BRAVE_SEARCH_API_KEY") + # web_search supports multiple providers (Google, Brave) with auto-detection register_web_search(mcp, credentials=credentials) # Register file system toolkits diff --git a/tools/src/aden_tools/tools/web_search_tool/README.md b/tools/src/aden_tools/tools/web_search_tool/README.md index 7344962e..cf4a39c8 100644 --- a/tools/src/aden_tools/tools/web_search_tool/README.md +++ b/tools/src/aden_tools/tools/web_search_tool/README.md @@ -1,31 +1,64 @@ # Web Search Tool -Search the web using the Brave Search API. +Search the web using multiple providers with automatic detection. ## Description Returns titles, URLs, and snippets for search results. Use when you need current information, research topics, or find websites. +Supports multiple search providers: +- **Brave Search API** (default, for backward compatibility) +- **Google Custom Search API** (fallback) + ## Arguments | Argument | Type | Required | Default | Description | |----------|------|----------|---------|-------------| | `query` | str | Yes | - | The search query (1-500 chars) | -| `num_results` | int | No | `10` | Number of results to return (1-20) | -| `country` | str | No | `us` | Country code for localized results (us, uk, de, etc.) | +| `num_results` | int | No | `10` | Number of results (1-10 for Google, 1-20 for Brave) | +| `country` | str | No | `us` | Country code for localized results | +| `language` | str | No | `en` | Language code (Google only) | +| `provider` | str | No | `auto` | Provider: "auto", "google", or "brave" | ## Environment Variables +Set credentials for at least one provider: + +### Option 1: Google Custom Search +| Variable | Required | Description | +|----------|----------|-------------| +| `GOOGLE_API_KEY` | Yes | API key from [Google Cloud Console](https://console.cloud.google.com/) | +| `GOOGLE_CSE_ID` | Yes | Search Engine ID from [Programmable Search Engine](https://programmablesearchengine.google.com/) | + +### Option 2: Brave Search | Variable | Required | Description | |----------|----------|-------------| | `BRAVE_SEARCH_API_KEY` | Yes | API key from [Brave Search API](https://brave.com/search/api/) | +## Provider Selection + +- `provider="auto"` (default): Uses Brave if available, otherwise Google (backward compatible) +- `provider="brave"`: Force Brave Search +- `provider="google"`: Force Google Custom Search + +## Example Usage + +```python +# Auto-detect provider based on available credentials +result = web_search(query="climate change effects") + +# Force specific provider +result = web_search(query="python tutorial", provider="google") +result = web_search(query="local news", provider="brave", country="id") +``` + ## Error Handling Returns error dicts for common issues: -- `BRAVE_SEARCH_API_KEY environment variable not set` - Missing API key +- `No search credentials configured` - No API keys set +- `Google credentials not configured` - Missing Google keys when provider="google" +- `Brave credentials not configured` - Missing Brave key when provider="brave" - `Query must be 1-500 characters` - Empty or too long query -- `Invalid API key` - API key rejected (HTTP 401) -- `Rate limit exceeded. Try again later.` - Too many requests (HTTP 429) +- `Invalid API key` - API key rejected +- `Rate limit exceeded` - Too many requests - `Search request timed out` - Request exceeded 30s timeout -- `Network error: ` - Connection or DNS issues diff --git a/tools/src/aden_tools/tools/web_search_tool/web_search_tool.py b/tools/src/aden_tools/tools/web_search_tool/web_search_tool.py index 47b82518..146c2785 100644 --- a/tools/src/aden_tools/tools/web_search_tool/web_search_tool.py +++ b/tools/src/aden_tools/tools/web_search_tool/web_search_tool.py @@ -1,13 +1,17 @@ """ -Web Search Tool - Search the web using Brave Search API. +Web Search Tool - Search the web using multiple providers. -Requires BRAVE_SEARCH_API_KEY environment variable. -Returns search results with titles, URLs, and snippets. +Supports: +- Google Custom Search API (GOOGLE_API_KEY + GOOGLE_CSE_ID) +- Brave Search API (BRAVE_SEARCH_API_KEY) + +Auto-detection: If provider="auto", tries Brave first (backward compatible), then Google. """ + from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Literal, Optional import httpx from fastmcp import FastMCP @@ -22,88 +26,193 @@ def register_tools( ) -> None: """Register web search tools with the MCP server.""" + def _search_google( + query: str, + num_results: int, + country: str, + language: str, + api_key: str, + cse_id: str, + ) -> dict: + """Execute search using Google Custom Search API.""" + response = httpx.get( + "https://www.googleapis.com/customsearch/v1", + params={ + "key": api_key, + "cx": cse_id, + "q": query, + "num": min(num_results, 10), + "lr": f"lang_{language}", + "gl": country, + }, + timeout=30.0, + ) + + if response.status_code == 401: + return {"error": "Invalid Google API key"} + elif response.status_code == 403: + return {"error": "Google API key not authorized or quota exceeded"} + elif response.status_code == 429: + return {"error": "Google rate limit exceeded. Try again later."} + elif response.status_code != 200: + return {"error": f"Google API request failed: HTTP {response.status_code}"} + + data = response.json() + results = [] + for item in data.get("items", [])[:num_results]: + results.append( + { + "title": item.get("title", ""), + "url": item.get("link", ""), + "snippet": item.get("snippet", ""), + } + ) + + return { + "query": query, + "results": results, + "total": len(results), + "provider": "google", + } + + def _search_brave( + query: str, + num_results: int, + country: str, + api_key: str, + ) -> dict: + """Execute search using Brave Search API.""" + response = httpx.get( + "https://api.search.brave.com/res/v1/web/search", + params={ + "q": query, + "count": min(num_results, 20), + "country": country, + }, + headers={ + "X-Subscription-Token": api_key, + "Accept": "application/json", + }, + timeout=30.0, + ) + + if response.status_code == 401: + return {"error": "Invalid Brave API key"} + elif response.status_code == 429: + return {"error": "Brave rate limit exceeded. Try again later."} + elif response.status_code != 200: + return {"error": f"Brave API request failed: HTTP {response.status_code}"} + + data = response.json() + results = [] + for item in data.get("web", {}).get("results", [])[:num_results]: + results.append( + { + "title": item.get("title", ""), + "url": item.get("url", ""), + "snippet": item.get("description", ""), + } + ) + + return { + "query": query, + "results": results, + "total": len(results), + "provider": "brave", + } + + def _get_credentials() -> dict: + """Get available search credentials.""" + if credentials is not None: + return { + "google_api_key": credentials.get("google_search"), + "google_cse_id": credentials.get("google_cse"), + "brave_api_key": credentials.get("brave_search"), + } + return { + "google_api_key": os.getenv("GOOGLE_API_KEY"), + "google_cse_id": os.getenv("GOOGLE_CSE_ID"), + "brave_api_key": os.getenv("BRAVE_SEARCH_API_KEY"), + } + @mcp.tool() def web_search( query: str, num_results: int = 10, country: str = "us", + language: str = "en", + provider: Literal["auto", "google", "brave"] = "auto", ) -> dict: """ - Search the web for information using Brave Search API. + Search the web for information. - Returns titles, URLs, and snippets. Use when you need current - information, research, or to find websites. - - Requires BRAVE_SEARCH_API_KEY environment variable. + Supports multiple search providers: + - "auto": Tries Brave first (backward compatible), then Google + - "google": Use Google Custom Search API (requires GOOGLE_API_KEY + GOOGLE_CSE_ID) + - "brave": Use Brave Search API (requires BRAVE_SEARCH_API_KEY) Args: query: The search query (1-500 chars) - num_results: Number of results to return (1-20) - country: Country code for localized results (us, uk, de, etc.) + num_results: Number of results to return (1-20 for Brave, 1-10 for Google) + country: Country code for localized results (us, id, uk, de, etc.) + language: Language code for results (en, id, etc.) - Google only + provider: Search provider to use ("auto", "google", "brave") Returns: - Dict with search results or error dict + Dict with search results, total count, and provider used """ - # Get API key - use CredentialManager if provided, fallback to direct env - if credentials is not None: - api_key = credentials.get("brave_search") - else: - # Backward compatibility: direct env access - api_key = os.getenv("BRAVE_SEARCH_API_KEY") - - if not api_key: - return { - "error": "BRAVE_SEARCH_API_KEY environment variable not set", - "help": "Get an API key at https://brave.com/search/api/", - } - - # Validate inputs if not query or len(query) > 500: return {"error": "Query must be 1-500 characters"} - if num_results < 1 or num_results > 20: - num_results = max(1, min(20, num_results)) + + creds = _get_credentials() + google_available = creds["google_api_key"] and creds["google_cse_id"] + brave_available = bool(creds["brave_api_key"]) try: - # Make request to Brave Search API - response = httpx.get( - "https://api.search.brave.com/res/v1/web/search", - params={ - "q": query, - "count": num_results, - "country": country, - }, - headers={ - "X-Subscription-Token": api_key, - "Accept": "application/json", - }, - timeout=30.0, - ) + if provider == "google": + if not google_available: + return { + "error": "Google credentials not configured", + "help": "Set GOOGLE_API_KEY and GOOGLE_CSE_ID environment variables", + } + return _search_google( + query, + num_results, + country, + language, + creds["google_api_key"], + creds["google_cse_id"], + ) - if response.status_code == 401: - return {"error": "Invalid API key"} - elif response.status_code == 429: - return {"error": "Rate limit exceeded. Try again later."} - elif response.status_code != 200: - return {"error": f"API request failed: HTTP {response.status_code}"} + elif provider == "brave": + if not brave_available: + return { + "error": "Brave credentials not configured", + "help": "Set BRAVE_SEARCH_API_KEY environment variable", + } + return _search_brave( + query, num_results, country, creds["brave_api_key"] + ) - data = response.json() - - # Extract results - results = [] - web_results = data.get("web", {}).get("results", []) - - for item in web_results[:num_results]: - results.append({ - "title": item.get("title", ""), - "url": item.get("url", ""), - "snippet": item.get("description", ""), - }) - - return { - "query": query, - "results": results, - "total": len(results), - } + else: # auto - try Brave first for backward compatibility + if brave_available: + return _search_brave( + query, num_results, country, creds["brave_api_key"] + ) + elif google_available: + return _search_google( + query, + num_results, + country, + language, + creds["google_api_key"], + creds["google_cse_id"], + ) + else: + return { + "error": "No search credentials configured", + "help": "Set either GOOGLE_API_KEY+GOOGLE_CSE_ID or BRAVE_SEARCH_API_KEY", + } except httpx.TimeoutException: return {"error": "Search request timed out"} diff --git a/tools/tests/tools/test_web_search_tool.py b/tools/tests/tools/test_web_search_tool.py index 8e50c48f..d15c570f 100644 --- a/tools/tests/tools/test_web_search_tool.py +++ b/tools/tests/tools/test_web_search_tool.py @@ -1,4 +1,5 @@ -"""Tests for web_search tool (FastMCP).""" +"""Tests for web_search tool with multi-provider support (FastMCP).""" + import pytest from fastmcp import FastMCP @@ -15,14 +16,16 @@ def web_search_fn(mcp: FastMCP): class TestWebSearchTool: """Tests for web_search tool.""" - def test_search_missing_api_key(self, web_search_fn, monkeypatch): - """Search without API key returns helpful error.""" + def test_no_credentials_returns_error(self, web_search_fn, monkeypatch): + """Search without any credentials returns helpful error.""" monkeypatch.delenv("BRAVE_SEARCH_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_CSE_ID", raising=False) result = web_search_fn(query="test query") assert "error" in result - assert "BRAVE_SEARCH_API_KEY" in result["error"] + assert "No search credentials configured" in result["error"] assert "help" in result def test_empty_query_returns_error(self, web_search_fn, monkeypatch): @@ -32,7 +35,9 @@ class TestWebSearchTool: result = web_search_fn(query="") assert "error" in result - assert "1-500" in result["error"].lower() or "character" in result["error"].lower() + assert ( + "1-500" in result["error"].lower() or "character" in result["error"].lower() + ) def test_long_query_returns_error(self, web_search_fn, monkeypatch): """Query exceeding 500 chars returns error.""" @@ -42,16 +47,105 @@ class TestWebSearchTool: assert "error" in result - def test_num_results_clamped_to_valid_range(self, web_search_fn, monkeypatch): - """num_results outside 1-20 is clamped (not error).""" + +class TestBraveProvider: + """Tests for Brave Search provider.""" + + def test_brave_missing_api_key(self, web_search_fn, monkeypatch): + """Brave provider without API key returns error.""" + monkeypatch.delenv("BRAVE_SEARCH_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + + result = web_search_fn(query="test", provider="brave") + + assert "error" in result + assert "Brave credentials not configured" in result["error"] + + def test_brave_explicit_provider(self, web_search_fn, monkeypatch): + """Brave provider can be explicitly selected.""" + monkeypatch.setenv("BRAVE_SEARCH_API_KEY", "test-key") + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + + result = web_search_fn(query="test", provider="brave") + assert isinstance(result, dict) + + +class TestGoogleProvider: + """Tests for Google Custom Search provider.""" + + def test_google_missing_api_key(self, web_search_fn, monkeypatch): + """Google provider without API key returns error.""" + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_CSE_ID", raising=False) + + result = web_search_fn(query="test", provider="google") + + assert "error" in result + assert "Google credentials not configured" in result["error"] + + def test_google_missing_cse_id(self, web_search_fn, monkeypatch): + """Google provider with API key but missing CSE ID returns error.""" + monkeypatch.setenv("GOOGLE_API_KEY", "test-key") + monkeypatch.delenv("GOOGLE_CSE_ID", raising=False) + + result = web_search_fn(query="test", provider="google") + + assert "error" in result + assert "Google credentials not configured" in result["error"] + + def test_google_explicit_provider(self, web_search_fn, monkeypatch): + """Google provider can be explicitly selected.""" + monkeypatch.setenv("GOOGLE_API_KEY", "test-key") + monkeypatch.setenv("GOOGLE_CSE_ID", "test-cse-id") + + result = web_search_fn(query="test", provider="google") + assert isinstance(result, dict) + + +class TestAutoProvider: + """Tests for auto provider selection.""" + + def test_auto_prefers_brave_for_backward_compatibility( + self, web_search_fn, monkeypatch + ): + """Auto mode uses Brave first for backward compatibility.""" + monkeypatch.setenv("GOOGLE_API_KEY", "test-google-key") + monkeypatch.setenv("GOOGLE_CSE_ID", "test-cse-id") + monkeypatch.setenv("BRAVE_SEARCH_API_KEY", "test-brave-key") + + result = web_search_fn(query="test", provider="auto") + assert isinstance(result, dict) + + def test_auto_falls_back_to_google(self, web_search_fn, monkeypatch): + """Auto mode falls back to Google when Brave not available.""" + monkeypatch.setenv("GOOGLE_API_KEY", "test-google-key") + monkeypatch.setenv("GOOGLE_CSE_ID", "test-cse-id") + monkeypatch.delenv("BRAVE_SEARCH_API_KEY", raising=False) + + result = web_search_fn(query="test", provider="auto") + assert isinstance(result, dict) + + def test_default_provider_is_auto(self, web_search_fn, monkeypatch): + """Default provider is auto.""" monkeypatch.setenv("BRAVE_SEARCH_API_KEY", "test-key") - # Test that the function handles out-of-range values gracefully - # The implementation clamps values, so we just verify it doesn't crash - # (actual API call would fail with invalid key, but that's expected) - result = web_search_fn(query="test", num_results=0) - # Should either clamp or error - both are acceptable + result = web_search_fn(query="test") assert isinstance(result, dict) - result = web_search_fn(query="test", num_results=100) + +class TestParameters: + """Tests for tool parameters.""" + + def test_custom_language_and_country(self, web_search_fn, monkeypatch): + """Custom language and country parameters are accepted.""" + monkeypatch.setenv("BRAVE_SEARCH_API_KEY", "test-key") + + result = web_search_fn(query="test", language="id", country="id") + assert isinstance(result, dict) + + def test_num_results_parameter(self, web_search_fn, monkeypatch): + """num_results parameter is accepted.""" + monkeypatch.setenv("BRAVE_SEARCH_API_KEY", "test-key") + + result = web_search_fn(query="test", num_results=5) assert isinstance(result, dict)