feat: support arxiv & brave search

src/tools/__init__.py

@@ -1,12 +1,19 @@
 from .crawl import crawl_tool
 from .python_repl import python_repl_tool
-from .search import tavily_search_tool, duckduckgo_search_tool
+from .search import (
+    tavily_search_tool,
+    duckduckgo_search_tool,
+    brave_search_tool,
+    arxiv_search_tool,
+)
 from src.config import SELECTED_SEARCH_ENGINE, SearchEngine
 
 # Map search engine names to their respective tools
 search_tool_mappings = {
     SearchEngine.TAVILY.value: tavily_search_tool,
     SearchEngine.DUCKDUCKGO.value: duckduckgo_search_tool,
+    SearchEngine.BRAVE_SEARCH.value: brave_search_tool,
+    SearchEngine.ARXIV.value: arxiv_search_tool,
 }
 
 web_search_tool = search_tool_mappings.get(SELECTED_SEARCH_ENGINE, tavily_search_tool)

src/tools/search.py

@@ -1,6 +1,10 @@
 import logging
+import os
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_community.tools import BraveSearch
+from langchain_community.tools.arxiv import ArxivQueryRun
+from langchain_community.utilities import ArxivAPIWrapper, BraveSearchWrapper
 from src.config import SEARCH_MAX_RESULTS
 from .decorators import create_logged_tool
 
@@ -16,3 +20,22 @@ LoggedDuckDuckGoSearch = create_logged_tool(DuckDuckGoSearchResults)
 duckduckgo_search_tool = LoggedDuckDuckGoSearch(
     name="web_search", max_results=SEARCH_MAX_RESULTS
 )
+
+LoggedBraveSearch = create_logged_tool(BraveSearch)
+brave_search_tool = LoggedBraveSearch(
+    name="web_search",
+    search_wrapper=BraveSearchWrapper(
+        api_key=os.getenv("BRAVE_SEARCH_API_KEY", ""),
+        search_kwargs={"count": SEARCH_MAX_RESULTS},
+    ),
+)
+
+LoggedArxivSearch = create_logged_tool(ArxivQueryRun)
+arxiv_search_tool = LoggedArxivSearch(
+    name="web_search",
+    api_wrapper=ArxivAPIWrapper(
+        top_k_results=SEARCH_MAX_RESULTS,
+        load_max_docs=SEARCH_MAX_RESULTS,
+        load_all_available_meta=True,
+    ),
+)