Create find_image_online.py

tools/find_image_online.py

@@ -0,0 +1,91 @@
+from typing import Any, Optional
+import re
+from smolagents.tools import Tool
+from smolagents.agent_types import AgentImage
+import requests
+from io import BytesIO
+import os
+import tempfile
+
+class FindImageOnlineTool(Tool):
+    name = "find_image_online"
+    description = "Searches for images online based on a query and returns an image that matches the description."
+    inputs = {'query': {'type': 'string', 'description': 'The search query for the image you want to find.'}}
+    output_type = "image"
+
+    def __init__(self, web_search_tool=None, visit_webpage_tool=None):
+        super().__init__()
+        self.web_search_tool = web_search_tool
+        self.visit_webpage_tool = visit_webpage_tool
+        self.is_initialized = True
+
+    def extract_image_urls(self, markdown_content):
+        # Extract image URLs from markdown using regex
+        # Look for standard markdown image patterns ![alt](url)
+        md_image_pattern = r'!\[.*?\]\((https?://[^)]+\.(jpg|jpeg|png|gif|webp))\)'
+        md_images = re.findall(md_image_pattern, markdown_content)
+        
+        # Also look for direct URLs that end with image extensions
+        direct_url_pattern = r'(https?://[^\s)]+\.(jpg|jpeg|png|gif|webp))'
+        direct_urls = re.findall(direct_url_pattern, markdown_content)
+        
+        # Combine and deduplicate results
+        image_urls = [url for url, _ in md_images] + [url for url, _ in direct_urls]
+        return list(set(image_urls))
+
+    def download_image(self, url):
+        try:
+            response = requests.get(url, stream=True, timeout=10)
+            response.raise_for_status()
+            
+            # Create a temporary file with appropriate extension
+            ext = os.path.splitext(url)[1]
+            if not ext or ext not in ['.jpg', '.jpeg', '.png', '.gif', '.webp']:
+                ext = '.jpg'  # Default extension
+                
+            temp_file = tempfile.NamedTemporaryFile(suffix=ext, delete=False)
+            temp_file.write(response.content)
+            temp_file.close()
+            
+            return temp_file.name, url
+        except Exception as e:
+            print(f"Error downloading image from {url}: {str(e)}")
+            return None, url
+
+    def forward(self, query: str) -> Any:
+        if not self.web_search_tool or not self.visit_webpage_tool:
+            return "Error: Web search and visit webpage tools must be provided."
+        
+        try:
+            # Step 1: Search for the query + "image"
+            search_query = f"{query} image"
+            search_results = self.web_search_tool.forward(search_query)
+            
+            # Step 2: Extract URLs from search results
+            url_pattern = r'\((https?://[^)]+)\)'
+            urls = re.findall(url_pattern, search_results)
+            
+            # Step 3: Visit each page and look for images
+            for url in urls[:3]:  # Limit to first 3 results for efficiency
+                try:
+                    page_content = self.visit_webpage_tool.forward(url)
+                    image_urls = self.extract_image_urls(page_content)
+                    
+                    # Step 4: Download the first valid image found
+                    for img_url in image_urls[:5]:  # Try up to 5 images per page
+                        img_path, source_url = self.download_image(img_url)
+                        if img_path:
+                            # Return both the image and the source information
+                            return {
+                                "image": AgentImage(img_path),
+                                "source_url": source_url,
+                                "page_url": url,
+                                "query": query
+                            }
+                except Exception as e:
+                    continue  # Try the next URL if this one fails
+            
+            return f"Could not find a suitable image for '{query}'. Please try a different query."
+        
+        except Exception as e:
+            return f"Error finding image: {str(e)}"