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

Merge branch 'adenhq:main' into fix/concurrent-storage-file-locks-leak

Browse Source
This commit is contained in:
Tahir yamin
2026-01-27 20:36:19 +05:00
committed by GitHub
parent 112b1baf2e a59d6ac6db
commit 0381a5c87b
9 changed files with 495 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/setup_mcp.py
+61 -47
View File
@@ -6,11 +6,24 @@ This script installs the framework and configures the MCP server.
"""
import json
import logging
import os
import subprocess
import sys
from pathlib import Path
logger = logging.getLogger(__name__)
def setup_logger():
"""Configure logger for CLI usage with colored output."""
if not logger.handlers:
handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter("%(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
class Colors:
"""ANSI color codes for terminal output."""
@@ -22,19 +35,19 @@ class Colors:
NC = "\033[0m" # No Color
def print_step(message: str, color: str = Colors.YELLOW):
"""Print a colored step message."""
print(f"{color}{message}{Colors.NC}")
def log_step(message: str):
"""Log a colored step message."""
logger.info(f"{Colors.YELLOW}{message}{Colors.NC}")
def print_success(message: str):
"""Print a success message."""
print(f"{Colors.GREEN}✓ {message}{Colors.NC}")
def log_success(message: str):
"""Log a success message."""
logger.info(f"{Colors.GREEN}✓ {message}{Colors.NC}")
def print_error(message: str):
"""Print an error message."""
print(f"{Colors.RED}✗ {message}{Colors.NC}", file=sys.stderr)
def log_error(message: str):
"""Log an error message."""
logger.error(f"{Colors.RED}✗ {message}{Colors.NC}")
def run_command(cmd: list, error_msg: str) -> bool:
@@ -43,52 +56,53 @@ def run_command(cmd: list, error_msg: str) -> bool:
subprocess.run(cmd, check=True, capture_output=True, text=True)
return True
except subprocess.CalledProcessError as e:
print_error(error_msg)
print(f"Error output: {e.stderr}", file=sys.stderr)
log_error(error_msg)
logger.error(f"Error output: {e.stderr}")
return False
def main():
"""Main setup function."""
print("=== Aden Hive Framework MCP Server Setup ===")
print()
setup_logger()
logger.info("=== Aden Hive Framework MCP Server Setup ===")
logger.info("")
# Get script directory
script_dir = Path(__file__).parent.absolute()
os.chdir(script_dir)
# Step 1: Install framework package
print_step("Step 1: Installing framework package...")
log_step("Step 1: Installing framework package...")
if not run_command(
[sys.executable, "-m", "pip", "install", "-e", "."], "Failed to install framework package"
):
sys.exit(1)
print_success("Framework package installed")
print()
log_success("Framework package installed")
logger.info("")
# Step 2: Install MCP dependencies
print_step("Step 2: Installing MCP dependencies...")
log_step("Step 2: Installing MCP dependencies...")
if not run_command(
[sys.executable, "-m", "pip", "install", "mcp", "fastmcp"],
"Failed to install MCP dependencies",
):
sys.exit(1)
print_success("MCP dependencies installed")
print()
log_success("MCP dependencies installed")
logger.info("")
# Step 3: Verify/create MCP configuration
print_step("Step 3: Verifying MCP server configuration...")
log_step("Step 3: Verifying MCP server configuration...")
mcp_config_path = script_dir / ".mcp.json"
if mcp_config_path.exists():
print_success("MCP configuration found at .mcp.json")
print("Configuration:")
log_success("MCP configuration found at .mcp.json")
logger.info("Configuration:")
with open(mcp_config_path) as f:
config = json.load(f)
print(json.dumps(config, indent=2))
logger.info(json.dumps(config, indent=2))
else:
print_error("No .mcp.json found")
print("Creating default MCP configuration...")
log_error("No .mcp.json found")
logger.info("Creating default MCP configuration...")
config = {
"mcpServers": {
@@ -103,11 +117,11 @@ def main():
with open(mcp_config_path, "w") as f:
json.dump(config, f, indent=2)
print_success("Created .mcp.json")
print()
log_success("Created .mcp.json")
logger.info("")
# Step 4: Test MCP server
print_step("Step 4: Testing MCP server...")
log_step("Step 4: Testing MCP server...")
try:
# Try importing the MCP server module
subprocess.run(
@@ -116,27 +130,27 @@ def main():
capture_output=True,
text=True,
)
print_success("MCP server module verified")
log_success("MCP server module verified")
except subprocess.CalledProcessError as e:
print_error("Failed to import MCP server module")
print(f"Error: {e.stderr}", file=sys.stderr)
log_error("Failed to import MCP server module")
logger.error(f"Error: {e.stderr}")
sys.exit(1)
print()
logger.info("")
# Success summary
print(f"{Colors.GREEN}=== Setup Complete ==={Colors.NC}")
print()
print("The MCP server is now ready to use!")
print()
print(f"{Colors.BLUE}To start the MCP server manually:{Colors.NC}")
print(" python -m framework.mcp.agent_builder_server")
print()
print(f"{Colors.BLUE}MCP Configuration location:{Colors.NC}")
print(f" {mcp_config_path}")
print()
print(f"{Colors.BLUE}To use with Claude Desktop or other MCP clients,{Colors.NC}")
print(f"{Colors.BLUE}add the following to your MCP client configuration:{Colors.NC}")
print()
logger.info(f"{Colors.GREEN}=== Setup Complete ==={Colors.NC}")
logger.info("")
logger.info("The MCP server is now ready to use!")
logger.info("")
logger.info(f"{Colors.BLUE}To start the MCP server manually:{Colors.NC}")
logger.info(" python -m framework.mcp.agent_builder_server")
logger.info("")
logger.info(f"{Colors.BLUE}MCP Configuration location:{Colors.NC}")
logger.info(f" {mcp_config_path}")
logger.info("")
logger.info(f"{Colors.BLUE}To use with Claude Desktop or other MCP clients,{Colors.NC}")
logger.info(f"{Colors.BLUE}add the following to your MCP client configuration:{Colors.NC}")
logger.info("")
example_config = {
"mcpServers": {
@@ -147,8 +161,8 @@ def main():
}
}
}
print(json.dumps(example_config, indent=2))
print()
logger.info(json.dumps(example_config, indent=2))
logger.info("")
if __name__ == "__main__":
core/verify_mcp.py
+50 -35
View File
@@ -6,10 +6,23 @@ This script checks if the MCP server is properly installed and configured.
"""
import json
import logging
import subprocess
import sys
from pathlib import Path
logger = logging.getLogger(__name__)
def setup_logger():
"""Configure logger for CLI usage."""
if not logger.handlers:
handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter("%(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
class Colors:
GREEN = "\033[0;32m"
@@ -21,29 +34,31 @@ class Colors:
def check(description: str) -> bool:
"""Print check description and return a context manager for result."""
print(f"Checking {description}...", end=" ")
logger.info(f"Checking {description}... ", extra={"end": ""})
sys.stdout.flush()
return True
def success(msg: str = "OK"):
"""Print success message."""
print(f"{Colors.GREEN}✓ {msg}{Colors.NC}")
"""Log success message."""
logger.info(f"{Colors.GREEN}✓ {msg}{Colors.NC}")
def warning(msg: str):
"""Print warning message."""
print(f"{Colors.YELLOW}⚠ {msg}{Colors.NC}")
"""Log warning message."""
logger.warning(f"{Colors.YELLOW}⚠ {msg}{Colors.NC}")
def error(msg: str):
"""Print error message."""
print(f"{Colors.RED}✗ {msg}{Colors.NC}")
"""Log error message."""
logger.error(f"{Colors.RED}✗ {msg}{Colors.NC}")
def main():
"""Run verification checks."""
print("=== MCP Server Verification ===")
print()
setup_logger()
logger.info("=== MCP Server Verification ===")
logger.info("")
script_dir = Path(__file__).parent.absolute()
all_checks_passed = True
@@ -61,7 +76,7 @@ def main():
success(f"installed at {framework_path}")
except subprocess.CalledProcessError:
error("framework package not found")
print(f" Run: pip install -e {script_dir}")
logger.info(f" Run: pip install -e {script_dir}")
all_checks_passed = False
# Check 2: MCP dependencies
@@ -75,7 +90,7 @@ def main():
if missing_deps:
error(f"missing: {', '.join(missing_deps)}")
print(f" Run: pip install {' '.join(missing_deps)}")
logger.info(f" Run: pip install {' '.join(missing_deps)}")
all_checks_passed = False
else:
success("all installed")
@@ -92,7 +107,7 @@ def main():
success("loads successfully")
except subprocess.CalledProcessError as e:
error("failed to import")
print(f" Error: {e.stderr}")
logger.error(f" Error: {e.stderr}")
all_checks_passed = False
# Check 4: MCP configuration file
@@ -106,9 +121,9 @@ def main():
if "mcpServers" in config and "agent-builder" in config["mcpServers"]:
server_config = config["mcpServers"]["agent-builder"]
success("found and valid")
print(f" Command: {server_config.get('command')}")
print(f" Args: {' '.join(server_config.get('args', []))}")
print(f" CWD: {server_config.get('cwd')}")
logger.info(f" Command: {server_config.get('command')}")
logger.info(f" Args: {' '.join(server_config.get('args', []))}")
logger.info(f" CWD: {server_config.get('cwd')}")
else:
warning("exists but missing agent-builder config")
all_checks_passed = False
@@ -117,8 +132,8 @@ def main():
all_checks_passed = False
else:
warning("not found (optional)")
print(f" Location would be: {mcp_config}")
print(" Run setup_mcp.py to create it")
logger.info(f" Location would be: {mcp_config}")
logger.info(" Run setup_mcp.py to create it")
# Check 5: Framework modules
check("core framework modules")
@@ -168,28 +183,28 @@ def main():
warning("server startup slow (might be OK)")
except subprocess.CalledProcessError as e:
error("server failed to start")
print(f" Error: {e.stderr}")
logger.error(f" Error: {e.stderr}")
all_checks_passed = False
print()
print("=" * 40)
logger.info("")
logger.info("=" * 40)
if all_checks_passed:
print(f"{Colors.GREEN}✓ All checks passed!{Colors.NC}")
print()
print("Your MCP server is ready to use.")
print()
print(f"{Colors.BLUE}To start the server:{Colors.NC}")
print(" python -m framework.mcp.agent_builder_server")
print()
print(f"{Colors.BLUE}To use with Claude Desktop:{Colors.NC}")
print(" Add the configuration from .mcp.json to your")
print(" Claude Desktop MCP settings")
logger.info(f"{Colors.GREEN}✓ All checks passed!{Colors.NC}")
logger.info("")
logger.info("Your MCP server is ready to use.")
logger.info("")
logger.info(f"{Colors.BLUE}To start the server:{Colors.NC}")
logger.info(" python -m framework.mcp.agent_builder_server")
logger.info("")
logger.info(f"{Colors.BLUE}To use with Claude Desktop:{Colors.NC}")
logger.info(" Add the configuration from .mcp.json to your")
logger.info(" Claude Desktop MCP settings")
else:
print(f"{Colors.RED}✗ Some checks failed{Colors.NC}")
print()
print("To fix issues, run:")
print(f" python {script_dir / 'setup_mcp.py'}")
print()
logger.info(f"{Colors.RED}✗ Some checks failed{Colors.NC}")
logger.info("")
logger.info("To fix issues, run:")
logger.info(f" python {script_dir / 'setup_mcp.py'}")
logger.info("")
if __name__ == "__main__":
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/mcp_server.py
+29 -10
View File
@@ -26,27 +26,46 @@ Note:
See aden_tools.credentials for details.
"""
import argparse
import logging
import os
import sys
logger = logging.getLogger(__name__)
def setup_logger():
"""Configure logger for MCP server."""
if not logger.handlers:
# For STDIO mode, log to stderr; for HTTP mode, log to stdout
stream = sys.stderr if "--stdio" in sys.argv else sys.stdout
handler = logging.StreamHandler(stream)
formatter = logging.Formatter("[MCP] %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
setup_logger()
# Suppress FastMCP banner in STDIO mode
if "--stdio" in sys.argv:
# Monkey-patch rich Console to redirect to stderr
import rich.console
_original_console_init = rich.console.Console.__init__
def _patched_console_init(self, *args, **kwargs):
kwargs['file'] = sys.stderr # Force all rich output to stderr
kwargs["file"] = sys.stderr # Force all rich output to stderr
_original_console_init(self, *args, **kwargs)
rich.console.Console.__init__ = _patched_console_init
from fastmcp import FastMCP
from starlette.requests import Request
from starlette.responses import PlainTextResponse
from fastmcp import FastMCP # noqa: E402
from starlette.requests import Request # noqa: E402
from starlette.responses import PlainTextResponse # noqa: E402
from aden_tools.credentials import CredentialManager, CredentialError
from aden_tools.tools import register_all_tools
from aden_tools.credentials import CredentialError, CredentialManager # noqa: E402
from aden_tools.tools import register_all_tools # noqa: E402
# Create credential manager
credentials = CredentialManager()
@@ -54,10 +73,10 @@ credentials = CredentialManager()
# Tier 1: Validate startup-required credentials (if any)
try:
credentials.validate_startup()
print("[MCP] Startup credentials validated")
logger.info("Startup credentials validated")
except CredentialError as e:
# Non-fatal - tools will validate their own credentials when called
print(f"[MCP] Warning: {e}", file=sys.stderr)
logger.warning(str(e))
mcp = FastMCP("tools")
@@ -65,7 +84,7 @@ mcp = FastMCP("tools")
tools = register_all_tools(mcp, credentials=credentials)
# Only print to stdout in HTTP mode (STDIO mode requires clean stdout for JSON-RPC)
if "--stdio" not in sys.argv:
print(f"[MCP] Registered {len(tools)} tools: {tools}")
logger.info(f"Registered {len(tools)} tools: {tools}")
@mcp.custom_route("/health", methods=["GET"])
@@ -105,7 +124,7 @@ def main() -> None:
# STDIO mode: only JSON-RPC messages go to stdout
mcp.run(transport="stdio")
else:
print(f"[MCP] Starting HTTP server on {args.host}:{args.port}")
logger.info(f"Starting HTTP server on {args.host}:{args.port}")
mcp.run(transport="http", host=args.host, port=args.port)
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)