From 48f4fbc41d5edb7780e098e5e5c70c0ad1428664 Mon Sep 17 00:00:00 2001
From: user <user@localhost>
Date: Mon, 23 Dec 2024 16:48:47 -0500
Subject: [PATCH] Add Tavily tool.

---
 ra_aid/tools/__init__.py          |  2 ++
 ra_aid/tools/web_search_tavily.py | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+)
 create mode 100644 ra_aid/tools/web_search_tavily.py

diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py
index 832dee0..ccd1334 100644
--- a/ra_aid/tools/__init__.py
+++ b/ra_aid/tools/__init__.py
@@ -1,4 +1,5 @@
 from .shell import run_shell_command
+from .web_search_tavily import web_search_tavily
 from .research import monorepo_detected, existing_project_detected, ui_detected
 from .human import ask_human
 from .programmer import run_programming_task
@@ -20,6 +21,7 @@ __all__ = [
     'ask_expert',
     'delete_key_facts',
     'delete_key_snippets',
+    'web_search_tavily',
     'deregister_related_files',
     'emit_expert_context', 
     'emit_key_facts',
diff --git a/ra_aid/tools/web_search_tavily.py b/ra_aid/tools/web_search_tavily.py
new file mode 100644
index 0000000..98a98b6
--- /dev/null
+++ b/ra_aid/tools/web_search_tavily.py
@@ -0,0 +1,19 @@
+import os
+from typing import Dict
+from tavily import TavilyClient
+from langchain_core.tools import tool
+
+@tool
+def web_search_tavily(query: str) -> Dict:
+    """
+    Perform a web search using Tavily API.
+    
+    Args:
+        query: The search query string
+        
+    Returns:
+        Dict containing search results from Tavily
+    """
+    client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+    search_result = client.search(query=query)
+    return search_result