Update agent.py

agent.py

@@ -1,6 +1,9 @@
 import json
 import os
-from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, tool, InferenceClientModel
+from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, tool, InferenceClientModel, Tool
+from duckduckgo_search import DDGS
+import requests
+from bs4 import BeautifulSoup
 
 #========Tools========#
 @tool
@@ -52,6 +55,41 @@ def modulus(a: int, b: int) -> int:
 
 search_tool = DuckDuckGoSearchTool()
 web_search = VisitWebpageTool()
+
+def process_tool_output(tool_output):
+    """Convert tool output to a string."""
+    if isinstance(tool_output, list):
+        return '\n'.join([item.get('title', str(item)) for item in tool_output if item]) or 'No results found'
+    elif isinstance(tool_output, dict):
+        return tool_output.get('title', str(tool_output))
+    return str(tool_output)
+
+class CustomDuckDuckGoSearchTool(Tool):
+    def __call__(self, query: str) -> str:
+        """Perform a DuckDuckGo search and return a string result."""
+        try:
+            with DDGS() as ddgs:
+                results = list(ddgs.text(query, max_results=3))  # Convert generator to list
+            return process_tool_output(results)
+        except Exception as e:
+            print(f"Error in DuckDuckGoSearchTool: {e}")
+            return f"Search error: {e}"
+
+class CustomVisitWebpageTool(Tool):
+    def __call__(self, url: str) -> str:
+        """Visit a webpage and extract content as a string."""
+        try:
+            response = requests.get(url, timeout=10)
+            response.raise_for_status()
+            soup = BeautifulSoup(response.text, 'html.parser')
+            # Extract text from the webpage
+            text = soup.get_text(separator=' ', strip=True)
+            return text[:1000]  # Limit output length to avoid overload
+        except Exception as e:
+            print(f"Error in VisitWebpageTool: {e}")
+            return f"Webpage access error: {e}"
+
+
 #=======Agent========#
 hf_token = os.environ.get("HF_TOKEN")
 model = InferenceClientModel(token = hf_token, model_id = 'meta-llama/Llama-2-7b-chat-hf', provider = "auto")
@@ -63,7 +101,9 @@ agent = CodeAgent(
     divide,
     modulus,
     search_tool,
-    web_search],
+    web_search,
+    CustomDuckDuckGoSearchTool,
+    CustomVisitWebpageTool],
     model = model,
     add_base_tools=True,
     planning_interval=3