Update agent.py

agent.py

@@ -1,4 +1,6 @@
 from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 
 
 
@@ -126,6 +128,77 @@ class modulo(Tool):
     def forward(self, a: int, b: int):
         return a % b
 
+class WikipediaSearchTool(tool):
+    name = "wikipedia_search_tool"
+    
+    description = """
+        Search Wikipedia for a query and return top 2 results.
+    
+        Args:
+            query: the search query.
+    """
+    
+    inputs = {
+        "query":{
+            "type":"string",
+            "description":"the search query"
+        }
+    }
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        documents = WikipediaLoader(query=query, load_max_docs=2).load()
+        condensed_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 documents
+        ])
+        return {"wikipedia_results": condensed_docs}
+
+class TavilySearchTool(tool):
+    name = "tavily_search_tool"
+    
+    description = """
+        Search the internet using the browser Tavily and return the top 3 results.
+    
+        Args:
+            query: the search query.
+    """
+    
+    inputs = {
+        "query":{
+            "type":"string",
+            "description":"the search query"
+        }
+    }
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        documents = TavilySearchResults(max_results=3).invoke(query=query)
+        condensed_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 documents
+        ])
+        return {"web_search_results": condensed_docs}
+
+
+@tool
+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}
+
+
+
 tools=[
     add(),
     subtract(),