Update agent.py

agent.py

@@ -43,6 +43,22 @@ def GoogleSearchTool(query: str) -> str:
     except Exception as e:
         return f"Error performing Google search: {str(e)}"
 
+@tool
+def search_arxiv(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query.
+     Returns:
+        str: Formatted search results
+    """
+    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}
 
 
 class VisitWebpageTool(Tool):
@@ -235,19 +251,6 @@ class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
     def forward(self, query: str) -> str:
         return super().forward(query)
 
-@tool
-def search_arxiv(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}
 
 
 
@@ -287,10 +290,10 @@ class MagAgent:
                 ExcelReaderTool(),
                 VisitWebpageTool(),
                 PythonCodeReaderTool(),
-                search_arxiv(),
+                search_arxiv,
 #                PNG2FENTool,
 #                ChessEngineTool(),
-#                GoogleSearchTool,
+                GoogleSearchTool,
 #                ImageAnalysisTool,
             ],
             verbosity_level=2,