updatetool

tools.py

@@ -17,6 +17,61 @@ from huggingface_hub import list_models
 from pathlib import Path
 import json 
 
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    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 {"wiki_results": formatted_search_docs}
+
+wiki_search_tool = Tool(
+    name="wiki_search",
+    func=wiki_search,
+    description="Searches for and retrieves information on Wikipedia using an inputted query."
+)
+
+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}
+web_search_tool = Tool(
+    name="web_search",
+    func=web_search,
+    description="Searches for and retrieves information online using an inputted query using Tavily."
+)
+
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+arvix_search_tool = Tool(
+    name="arvix_search",
+    func=arvix_search,
+    description="Searches for and retrieves information online using an inputted query on the website arvix, that is a database storing many factual scientific papers."
+)
 
 def extract_image_text(img_path: str) -> str:
     """