adds additional search params

Introduces support for additional Brave Search API web-search parameters.

Author: jonathansampson

lib/crewai-tools/src/crewai_tools/tools/brave_search_tool/brave_search_tool.py

@@ -1,10 +1,12 @@
 from datetime import datetime
+import json
 import os
 import time
-from typing import Any, ClassVar
+from typing import Annotated, Any, ClassVar, Literal
 
 from crewai.tools import BaseTool, EnvVar
 from pydantic import BaseModel, Field
+from pydantic.types import StringConstraints
 import requests
 
 
@@ -15,14 +17,60 @@ def _save_results_to_file(content: str) -> None:
         file.write(content)
 
 
+FreshnessPreset = Literal["pd", "pw", "pm", "py"]
+FreshnessRange = Annotated[
+    str, StringConstraints(pattern=r"^\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2}$")
+]
+Freshness = FreshnessPreset | FreshnessRange
+SafeSearch = Literal["off", "moderate", "strict"]
+
+
 class BraveSearchToolSchema(BaseModel):
-    """Input for BraveSearchTool."""
+    """Input for BraveSearchTool"""
 
-    search_query: str = Field(
-        ..., description="Mandatory search query you want to use to search the internet"
+    query: str = Field(..., description="Search query to perform")
+    country: str | None = Field(
+        default=None,
+        description="Country code for geo-targeting (e.g., 'US', 'BR').",
+    )
+    search_language: str | None = Field(
+        default=None,
+        description="Language code for the search results (e.g., 'en', 'es').",
+    )
+    count: int | None = Field(
+        default=None,
+        description="The maximum number of results to return. Actual number may be less.",
+    )
+    offset: int | None = Field(
+        default=None, description="Skip the first N result sets/pages. Max is 9."
+    )
+    safesearch: SafeSearch | None = Field(
+        default=None,
+        description="Filter out explicit content. Options: off/moderate/strict",
+    )
+    spellcheck: bool | None = Field(
+        default=None,
+        description="Attempt to correct spelling errors in the search query.",
+    )
+    freshness: Freshness | None = Field(
+        default=None,
+        description="Enforce freshness of results. Options: pd/pw/pm/py, or YYYY-MM-DDtoYYYY-MM-DD",
+    )
+    text_decorations: bool | None = Field(
+        default=None,
+        description="Include markup to highlight search terms in the results.",
+    )
+    extra_snippets: bool | None = Field(
+        default=None,
+        description="Include up to 5 text snippets for each page if possible.",
+    )
+    operators: bool | None = Field(
+        default=None,
+        description="Whether to apply search operators (e.g., site:example.com).",
     )
 
 
+# TODO: Extend support to additional endpoints (e.g., /images, /news, etc.)
 class BraveSearchTool(BaseTool):
     """BraveSearchTool - A tool for performing web searches using the Brave Search API.
 
@@ -35,13 +83,13 @@ class BraveSearchTool(BaseTool):
         - python-dotenv (for API key management)
     """
 
-    name: str = "Brave Web Search the internet"
+    name: str = "Brave Search"
     description: str = (
-        "A tool that can be used to search the internet with a search_query."
+        "A tool that performs web searches using the Brave Search API. "
+        "Results are returned as structured JSON data."
     )
     args_schema: type[BaseModel] = BraveSearchToolSchema
     search_url: str = "https://api.search.brave.com/res/v1/web/search"
-    country: str | None = ""
     n_results: int = 10
     save_file: bool = False
     _last_request_time: ClassVar[float] = 0
@@ -73,19 +121,64 @@ def _run(
                 self._min_request_interval - (current_time - self._last_request_time)
             )
         BraveSearchTool._last_request_time = time.time()
+
+        # Construct and send the request
         try:
-            search_query = kwargs.get("search_query") or kwargs.get("query")
-            if not search_query:
-                raise ValueError("Search query is required")
+            # Maintain both "search_query" and "query" for backwards compatibility
+            query = kwargs.get("search_query") or kwargs.get("query")
+            if not query:
+                raise ValueError("Query is required")
+
+            payload = {"q": query}
+
+            if country := kwargs.get("country"):
+                payload["country"] = country
+
+            if search_language := kwargs.get("search_language"):
+                payload["search_language"] = search_language
+
+            # Fallback to deprecated n_results parameter if no count is provided
+            count = kwargs.get("count")
+            if count is not None:
+                payload["count"] = count
+            else:
+                payload["count"] = self.n_results
+
+            # Offset may be 0, so avoid truthiness check
+            offset = kwargs.get("offset")
+            if offset is not None:
+                payload["offset"] = offset
+
+            if safesearch := kwargs.get("safesearch"):
+                payload["safesearch"] = safesearch
 
             save_file = kwargs.get("save_file", self.save_file)
-            n_results = kwargs.get("n_results", self.n_results)
+            if freshness := kwargs.get("freshness"):
+                payload["freshness"] = freshness
+
+            # Boolean parameters
+            spellcheck = kwargs.get("spellcheck")
+            if spellcheck is not None:
+                payload["spellcheck"] = spellcheck
+
+            text_decorations = kwargs.get("text_decorations")
+            if text_decorations is not None:
+                payload["text_decorations"] = text_decorations
 
-            payload = {"q": search_query, "count": n_results}
+            extra_snippets = kwargs.get("extra_snippets")
+            if extra_snippets is not None:
+                payload["extra_snippets"] = extra_snippets
 
-            if self.country != "":
-                payload["country"] = self.country
+            operators = kwargs.get("operators")
+            if operators is not None:
+                payload["operators"] = operators
 
+            # Limit the result types to "web" since there is presently no
+            # handling of other types like "discussions", "faq", "infobox",
+            # "news", "videos", or "locations".
+            payload["result_filter"] = "web"
+
+            # Setup Request Headers
             headers = {
                 "X-Subscription-Token": os.environ["BRAVE_API_KEY"],
                 "Accept": "application/json",
@@ -97,25 +190,32 @@ def _run(
             response.raise_for_status()  # Handle non-200 responses
             results = response.json()
 
+            # TODO: Handle other result types like "discussions", "faq", etc.
+            web_results_items = []
             if "web" in results:
-                results = results["web"]["results"]
-                string = []
-                for result in results:
-                    try:
-                        string.append(
-                            "\n".join(
-                                [
-                                    f"Title: {result['title']}",
-                                    f"Link: {result['url']}",
-                                    f"Snippet: {result['description']}",
-                                    "---",
-                                ]
-                            )
-                        )
-                    except KeyError:  # noqa: PERF203
+                web_results = results["web"]["results"]
+
+                for result in web_results:
+                    url = result.get("url")
+                    title = result.get("title")
+                    # If, for whatever reason, this entry does not have a title
+                    # or url, skip it.
+                    if not url or not title:
                         continue
-
-            content = "\n".join(string)
+                    item = {
+                        "url": url,
+                        "title": title,
+                    }
+                    description = result.get("description")
+                    if description:
+                        item["description"] = description
+                    snippets = result.get("extra_snippets")
+                    if snippets:
+                        item["snippets"] = snippets
+
+                    web_results_items.append(item)
+
+            content = json.dumps(web_results_items)
         except requests.RequestException as e:
             return f"Error performing search: {e!s}"
         except KeyError as e:

lib/crewai-tools/tests/tools/brave_search_tool_test.py

@@ -30,16 +30,30 @@ def test_brave_tool_search(mock_get, brave_tool):
     }
     mock_get.return_value.json.return_value = mock_response
 
-    result = brave_tool.run(search_query="test")
+    result = brave_tool.run(query="test")
     assert "Test Title" in result
     assert "http://test.com" in result
 
 
-def test_brave_tool():
-    tool = BraveSearchTool(
-        n_results=2,
-    )
-    tool.run(search_query="ChatGPT")
+@patch("requests.get")
+def test_brave_tool(mock_get):
+    mock_response = {
+        "web": {
+            "results": [
+                {
+                    "title": "Brave Browser",
+                    "url": "https://brave.com",
+                    "description": "Brave Browser description",
+                }
+            ]
+        }
+    }
+    mock_get.return_value.json.return_value = mock_response
+    
+    tool = BraveSearchTool(n_results=2)
+    result = tool.run(query="Brave Browser")
+    assert "Brave Browser" in result
+    assert "brave.com" in result
 
 
 if __name__ == "__main__":