Update agent.py

agent.py

@@ -53,7 +53,7 @@ class VisitWebpageTool(Tool):
 
     def forward(self, url: str) -> str:
         try:
-            response = requests.get(url, timeout=30)
+            response = requests.get(url, timeout=50)
             response.raise_for_status()
             markdown_content = markdownify(response.text).strip()
             markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
@@ -89,7 +89,7 @@ class DownloadTaskAttachmentTool(Tool):
                 while not self.rate_limiter.consume(1):
                     print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
                     time.sleep(4)  # Assuming 15 RPM
-            response = requests.get(file_url, stream=True, timeout=25)
+            response = requests.get(file_url, stream=True, timeout=50)
             response.raise_for_status()
             
             # Determine file extension based on Content-Type
@@ -235,6 +235,21 @@ 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}
+
+
 
 class MagAgent:
     def __init__(self, rate_limiter: Optional[Limiter] = None):
@@ -272,6 +287,7 @@ class MagAgent:
                 ExcelReaderTool(),
                 VisitWebpageTool(),
                 PythonCodeReaderTool(),
+                search_arxiv(),
 #                PNG2FENTool,
 #                ChessEngineTool(),
 #                GoogleSearchTool,