Update agent.py

agent.py

@@ -1,4 +1,5 @@
-from smolagents import CodeAgent, ToolCallingAgent, LiteLLMModel, tool, Tool, load_tool, DuckDuckGoSearchTool, WikipediaSearchTool
+from smolagents import CodeAgent, ToolCallingAgent, LiteLLMModel, tool, Tool, load_tool, WebSearchTool
+#DuckDuckGoSearchTool, WikipediaSearchTool
 import asyncio
 import os
 import re
@@ -59,8 +60,9 @@ class VisitWebpageTool(Tool):
             response.raise_for_status()
             markdown_content = markdownify(response.text).strip()
             markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
-            from smolagents.utils import truncate_content
-            return truncate_content(markdown_content, 10000)
+            #from smolagents.utils import truncate_content
+            #return truncate_content(markdown_content, 10000)
+            return markdown_content
         except requests.exceptions.HTTPError as e:
             if e.response.status_code == 429:
                 raise  # Retry on 429
@@ -163,15 +165,15 @@ class PythonCodeReaderTool(Tool):
 
 from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
-class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
-    @retry(
-        stop=stop_after_attempt(3),
-        wait=wait_exponential(multiplier=1, min=4, max=10),
-        retry=retry_if_exception_type(Exception)
-    )
+#class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
+#    @retry(
+#        stop=stop_after_attempt(3),
+#        wait=wait_exponential(multiplier=1, min=4, max=10),
+#        retry=retry_if_exception_type(Exception)
+#    )
 
-    def forward(self, query: str) -> str:
-        return super().forward(query)
+#    def forward(self, query: str) -> str:
+#        return super().forward(query)
 
 class MagAgent:
     def __init__(self, rate_limiter: Optional[Limiter] = None):
@@ -241,8 +243,9 @@ class MagAgent:
         web_agent = ToolCallingAgent(
             tools=[
 #                RetryDuckDuckGoSearchTool(),
-                WikipediaSearchTool(),
+#                WikipediaSearchTool(),
 #                SpeechToTextTool(),
+                WebSearchTool(),
                 VisitWebpageTool(),
 #                ExcelReaderTool(),
 #                PythonCodeReaderTool(),
@@ -261,7 +264,7 @@ class MagAgent:
             add_base_tools=True,
             additional_authorized_imports=self.imports,
             verbosity_level=2,
-            max_steps=20
+            max_steps=10
         )
         print("MagAgent initialized.")