From 6c242ef3bbfe722d4159e60c931d20e2a38a0570 Mon Sep 17 00:00:00 2001
From: siddas27 <siddharthdas8@gmail.com>
Date: Sat, 30 Nov 2024 14:04:06 -0600
Subject: [PATCH] add brave search tool

---
 .../tools/brave_search_tool/README.md         | 30 +++++++
 .../brave_search_tool/brave_search_tool.py    | 82 +++++++++++++++++++
 tests/tools/brave_search_tool_test.py         | 13 +++
 3 files changed, 125 insertions(+)
 create mode 100644 src/crewai_tools/tools/brave_search_tool/README.md
 create mode 100644 src/crewai_tools/tools/brave_search_tool/brave_search_tool.py
 create mode 100644 tests/tools/brave_search_tool_test.py

diff --git a/src/crewai_tools/tools/brave_search_tool/README.md b/src/crewai_tools/tools/brave_search_tool/README.md
new file mode 100644
index 000000000..a66210491
--- /dev/null
+++ b/src/crewai_tools/tools/brave_search_tool/README.md
@@ -0,0 +1,30 @@
+# BraveSearchTool Documentation
+
+## Description
+This tool is designed to perform a web search for a specified query from a text's content across the internet. It utilizes the Brave Web Search API, which is a REST API to query Brave Search and get back search results from the web. The following sections describe how to curate requests, including parameters and headers, to Brave Web Search API and get a JSON response back.
+
+## Installation
+To incorporate this tool into your project, follow the installation instructions below:
+```shell
+pip install 'crewai[tools]'
+```
+
+## Example
+The following example demonstrates how to initialize the tool and execute a search with a given query:
+
+```python
+from crewai_tools import BraveSearchTool
+
+# Initialize the tool for internet searching capabilities
+tool = BraveSearchTool()
+```
+
+## Steps to Get Started
+To effectively use the `BraveSearchTool`, follow these steps:
+
+1. **Package Installation**: Confirm that the `crewai[tools]` package is installed in your Python environment.
+2. **API Key Acquisition**: Acquire a API key [here](https://api.search.brave.com/app/keys).
+3. **Environment Configuration**: Store your obtained API key in an environment variable named `BRAVE_API_KEY` to facilitate its use by the tool.
+
+## Conclusion
+By integrating the `BraveSearchTool` into Python projects, users gain the ability to conduct real-time, relevant searches across the internet directly from their applications. By adhering to the setup and usage guidelines provided, incorporating this tool into projects is streamlined and straightforward.
diff --git a/src/crewai_tools/tools/brave_search_tool/brave_search_tool.py b/src/crewai_tools/tools/brave_search_tool/brave_search_tool.py
new file mode 100644
index 000000000..54f546f1e
--- /dev/null
+++ b/src/crewai_tools/tools/brave_search_tool/brave_search_tool.py
@@ -0,0 +1,82 @@
+import datetime
+import os
+from typing import Any, Optional, Type
+
+import requests
+from pydantic import BaseModel, Field
+
+from crewai_tools.tools.base_tool import BaseTool
+
+
+def _save_results_to_file(content: str) -> None:
+    """Saves the search results to a file."""
+    filename = f"search_results_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt"
+    with open(filename, "w") as file:
+        file.write(content)
+    print(f"Results saved to {filename}")
+
+
+class BraveSearchToolSchema(BaseModel):
+    """Input for BraveSearchTool."""
+
+    search_query: str = Field(
+        ..., description="Mandatory search query you want to use to search the internet"
+    )
+
+
+class BraveSearchTool(BaseTool):
+    name: str = "Search the internet"
+    description: str = (
+        "A tool that can be used to search the internet with a search_query."
+    )
+    args_schema: Type[BaseModel] = BraveSearchToolSchema
+    search_url: str = "https://api.search.brave.com/res/v1/web/search"
+    country: Optional[str] = ""
+    n_results: int = 10
+    save_file: bool = False
+
+    def _run(
+        self,
+        **kwargs: Any,
+    ) -> Any:
+        search_query = kwargs.get("search_query") or kwargs.get("query")
+        save_file = kwargs.get("save_file", self.save_file)
+        n_results = kwargs.get("n_results", self.n_results)
+
+        payload = {"q": search_query, "count": n_results}
+
+        if self.country != "":
+            payload["country"] = self.country
+
+        headers = {
+            "X-Subscription-Token": os.environ["BRAVE_API_KEY"],
+            "Accept": "application/json",
+        }
+
+        response = requests.get(self.search_url, headers=headers, params=payload)
+        results = response.json()
+
+        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:
+                    continue
+
+            content = "\n".join(string)
+            if save_file:
+                _save_results_to_file(content)
+            return f"\nSearch results: {content}\n"
+        else:
+            return results
diff --git a/tests/tools/brave_search_tool_test.py b/tests/tools/brave_search_tool_test.py
new file mode 100644
index 000000000..16c1bcb92
--- /dev/null
+++ b/tests/tools/brave_search_tool_test.py
@@ -0,0 +1,13 @@
+from crewai_tools.tools.brave_search_tool.brave_search_tool import BraveSearchTool
+
+
+def test_brave_tool():
+    tool = BraveSearchTool(
+        n_results=2,
+    )
+
+    print(tool.run(search_query="ChatGPT"))
+
+
+if __name__ == "__main__":
+    test_brave_tool()