Update tools.py

tools.py

@@ -1,12 +1,13 @@
 # Tools extracted from one of leader repos > https://huggingface.co/spaces/baixianger/RobotPai/blob/main/agent.py
 
-from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_tavily import TavilySearch
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.document_loaders import ArxivLoader
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain.tools.retriever import create_retriever_tool
+from langchain_core.messages import ToolMessage
 
 from typing import Dict
 
@@ -78,13 +79,13 @@ def wiki_search(query: str) -> Dict[str, list]:
     return {"wiki_results": formatted_search_docs}
 
 @tool(parse_docstring=True)
-def web_search(query: str) -> Dict[str, list]:
-    """Search in the web with Tavily for a query and return maximum 3 results.
+def web_search(query: str) -> ToolMessage:
+    """Search in the web with Tavily for a query and return maximum 4 results.
     
     Args:
         query: The search query.
     """
-    search_docs = TavilySearchResults(max_results=4, include_images=False).invoke({"query": query})
+    return TavilySearch(max_results=4, include_images=False).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>'