update tavily

agent.py

@@ -12,7 +12,11 @@ from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.tools import Tool
-from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_tavily import TavilySearch
+from langchain_tavily import TavilySearch, TavilyExtract
+
+# from langchain_community.tools.tavily_search import TavilySearchResults
+
 from langchain_experimental.utilities import PythonREPL
 import assemblyai as aai
 
@@ -26,6 +30,16 @@ repl_tool = Tool(
     func=PythonREPL().run,
 )
 
+# Initialize Tavily Search Tool
+tavily_search_tool = TavilySearch(
+    max_results=5,
+    topic="general",
+)
+
+# Initialize Tavily Extract Tool
+tavily_extract_tool = TavilyExtract()
+
+
 @tool
 def read_excel_file(excel_file_path: str) -> str:
     """Read the content of excel file.
@@ -80,25 +94,18 @@ def wiki_search(query: str) -> str:
 
 # @tool
 # def web_search(query: str) -> str:
-#     """Search DuckDuckGo for a query and return maximum 5 results.
+#     """Search Tavily for a query and return maximum 3 results.
     
 #     Args:
 #         query: The search query."""
-#     search_docs = DuckDuckGoSearchResults(max_results=5).invoke(query)
-#     return {"web_results": search_docs}
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"web_results": formatted_search_docs}
+#     search_docs = TavilySearchResults(max_results=5).invoke(query=query)
+#     formatted_search_docs = "\n\n---\n\n".join(
+#         [
+#             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+#             for doc in search_docs
+#         ])
+#     print({"web_results": formatted_search_docs})
+#     return {"web_results": formatted_search_docs}
 
     
 system_prompt = """
@@ -113,7 +120,8 @@ sys_msg = SystemMessage(content=system_prompt)
 
 tools = [
     wiki_search,
-    web_search,
+    tavily_search_tool,
+    tavily_extract_tool,
     repl_tool,
     read_excel_file,
     transcribe_audio