Update agent.py

agent.py

@@ -4,6 +4,8 @@ from typing import List, Dict, Any, Optional
 import tempfile
 import pandas as pd
 import numpy as np
+import requests
+from pathlib import Path
 
 """Langraph"""
 from langgraph.graph import START, StateGraph, MessagesState
@@ -29,8 +31,17 @@ def duckduckgo_search(query: str) -> str:
     Args:
         query: The search query."""
     try:
-        # Initialize DuckDuckGo search
-        ddgs = DDGS()
+        # Initialize DuckDuckGo search with appropriate user agent
+        ddgs_user_agent_header = "/".join(
+            [
+                "Mozilla",
+                "5.0 (Windows NT 10.0; Win64; x64; rv:124.0) Gecko",
+                "20100101 Firefox",
+                "124.0",
+            ]
+        )
+        
+        ddgs = DDGS(headers={"User-Agent": ddgs_user_agent_header})
         
         # Perform the search and get the first 5 results
         results = list(ddgs.text(query, max_results=5))
@@ -52,6 +63,57 @@ def duckduckgo_search(query: str) -> str:
     except Exception as e:
         return {"search_results": f"Error performing search: {str(e)}"}
 
+@tool
+def download_image(url: str, filename: Optional[str] = None) -> str:
+    """Download an image from a URL and save it to a file.
+    Args:
+        url: The URL of the image to download.
+        filename: Optional; The name to save the file as. If not provided, a name will be derived from the URL.
+    """
+    try:
+        # Send a GET request to the URL
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise an exception for HTTP errors
+        
+        # Determine the file extension from the content type
+        content_type = response.headers.get('Content-Type', '')
+        ext = '.jpg'  # Default extension
+        if 'png' in content_type:
+            ext = '.png'
+        elif 'jpeg' in content_type or 'jpg' in content_type:
+            ext = '.jpg'
+        elif 'gif' in content_type:
+            ext = '.gif'
+        elif 'webp' in content_type:
+            ext = '.webp'
+        
+        # Create a filename if not provided
+        if not filename:
+            # Use the last part of the URL or a timestamp
+            url_path = Path(url.split('?')[0])
+            url_filename = url_path.name
+            if url_filename:
+                filename = url_filename
+                if not Path(filename).suffix:
+                    filename = f"{filename}{ext}"
+            else:
+                import time
+                filename = f"image_{int(time.time())}{ext}"
+        
+        # Ensure the file has an extension
+        if not Path(filename).suffix:
+            filename = f"{filename}{ext}"
+        
+        # Save the file
+        filepath = os.path.join(tempfile.gettempdir(), filename)
+        with open(filepath, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        
+        return f"Image downloaded and saved to {filepath}"
+    except Exception as e:
+        return f"Error downloading image: {str(e)}"
+
 ### =============== DOCUMENT PROCESSING TOOLS =============== ###
 # File handling still requires external tools
 
@@ -87,6 +149,7 @@ sys_msg = SystemMessage(content=system_prompt)
 tools = [
     duckduckgo_search,
     save_and_read_file,
+    download_image,
 ]