chnjinlei/hive
1
0
Fork 0
Code Issues Pull Requests Actions 39 Packages Projects Releases Wiki Activity

refactor(tools): add multi-provider support to web_search tool (#795)

Browse Source 
* 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 <abdamus@hilti.com>
This commit is contained in:
vrijmetse
2026-01-27 21:46:41 +07:00
committed by GitHub
parent 9d39c09e27
commit a59d6ac6db
6 changed files with 355 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/README.md
+6 -2
View File
@@ -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 |
tools/src/aden_tools/credentials/search.py
+19 -10
View File
@@ -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",
),
}
tools/src/aden_tools/tools/__init__.py
+8 -5
View File
@@ -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
tools/src/aden_tools/tools/web_search_tool/README.md
+40 -7
View File
@@ -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: <error>` - Connection or DNS issues
tools/src/aden_tools/tools/web_search_tool/web_search_tool.py
+175 -66
View File
@@ -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"}
tools/tests/tools/test_web_search_tool.py
+107 -13
View File
@@ -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)