more reliable way of getting images

tools/image_handler.py

@@ -18,7 +18,12 @@ class ImageHandlerTool(Tool):
         },
         'style': {
             'type': 'string',
-            'description': 'Style for generated images (e.g., "photo", "artistic")',
+            'description': 'Style for generated images (e.g., "photo", "artistic", "realistic")',
+            'nullable': True
+        },
+        'skip_web_search': {
+            'type': 'boolean',
+            'description': 'Whether to skip web search and go straight to generation',
             'nullable': True
         }
     }
@@ -33,9 +38,17 @@ class ImageHandlerTool(Tool):
     def _download_image(self, url: str, filename: str) -> Optional[str]:
         """Downloads an image from a URL and saves it to a temporary file"""
         try:
-            response = requests.get(url, timeout=10)
+            headers = {
+                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+            }
+            response = requests.get(url, timeout=10, headers=headers)
             response.raise_for_status()
 
+            # Check if response is actually an image
+            content_type = response.headers.get('content-type', '')
+            if not content_type.startswith('image/'):
+                return None
+
             # Ensure temp directory exists
             os.makedirs(self.temp_dir, exist_ok=True)
 
@@ -52,11 +65,12 @@ class ImageHandlerTool(Tool):
         """Attempts to find images via web search"""
         results = []
         try:
-            # Try different search queries
+            # Try different search queries with better targeting
             search_queries = [
-                f"{query} high quality photo",
-                f"{query} professional photograph",
-                f"{query} travel photo"
+                f"{query} high resolution photo",
+                f"{query} professional photography",
+                f"{query} best pictures",
+                f"{query} travel photography"
             ]
 
             for search_query in search_queries:
@@ -64,23 +78,33 @@ class ImageHandlerTool(Tool):
                     break
 
                 time.sleep(2)  # Rate limiting
-                search_results = self.web_search.forward(
-                    query=search_query, max_results=num_images)
-
-                for idx, result in enumerate(search_results):
-                    if len(results) >= num_images:
-                        break
-
-                    if 'image_url' in result:
-                        filename = f"{query.replace(' ', '_')}_{idx}.jpg"
-                        file_path = self._download_image(
-                            result['image_url'], filename)
-                        if file_path:
-                            results.append({
-                                'file_path': file_path,
-                                'source': 'web',
-                                'url': result['image_url']
-                            })
+                try:
+                    search_results = self.web_search.forward(
+                        query=search_query, max_results=num_images)
+
+                    if isinstance(search_results, str):  # Handle string responses
+                        continue
+
+                    for idx, result in enumerate(search_results):
+                        if len(results) >= num_images:
+                            break
+
+                        # Try both image_url and direct URL fields
+                        url = result.get('image_url') or result.get('url')
+                        if url and url.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')):
+                            filename = f"{query.replace(' ', '_')}_{idx}.jpg"
+                            file_path = self._download_image(url, filename)
+                            if file_path:
+                                results.append({
+                                    'file_path': file_path,
+                                    'source': 'web',
+                                    'url': url,
+                                    'title': result.get('title', ''),
+                                    'attribution': result.get('source', '')
+                                })
+                except Exception as search_error:
+                    print(f"Search query failed: {str(search_error)}")
+                    continue
 
         except Exception as e:
             print(f"Web search failed: {str(e)}")
@@ -91,42 +115,61 @@ class ImageHandlerTool(Tool):
         """Generates images using the image generation tool"""
         results = []
         try:
+            # Enhanced prompts for better generation
+            prompts = [
+                f"Generate a {style} style image of {query}, high quality, detailed",
+                f"Create a {style} representation of {query}, professional quality",
+                f"Make a {style} image showing {query}, realistic and clear"
+            ]
+
             for idx in range(num_images):
-                prompt = f"Generate a {style} of {query}"
-                response = self.image_gen.forward(prompt=prompt)
-
-                if isinstance(response, dict) and 'image_path' in response:
-                    results.append({
-                        'file_path': response['image_path'],
-                        'source': 'generated',
-                        'prompt': prompt
-                    })
-                elif isinstance(response, str) and os.path.exists(response):
-                    results.append({
-                        'file_path': response,
-                        'source': 'generated',
-                        'prompt': prompt
-                    })
+                prompt = prompts[idx % len(prompts)].replace(
+                    query, query + f" {idx+1}")
+                try:
+                    response = self.image_gen.forward(prompt=prompt)
+
+                    if isinstance(response, dict) and 'image_path' in response:
+                        results.append({
+                            'file_path': response['image_path'],
+                            'source': 'generated',
+                            'prompt': prompt,
+                            'style': style
+                        })
+                    elif isinstance(response, str) and os.path.exists(response):
+                        results.append({
+                            'file_path': response,
+                            'source': 'generated',
+                            'prompt': prompt,
+                            'style': style
+                        })
+                except Exception as gen_error:
+                    print(
+                        f"Failed to generate image {idx+1}: {str(gen_error)}")
+                    continue
+
+                time.sleep(1)  # Brief pause between generations
 
         except Exception as e:
             print(f"Image generation failed: {str(e)}")
 
         return results
 
-    def forward(self, query: str, num_images: int = 2, style: str = "photo") -> Dict:
+    def forward(self, query: str, num_images: int = 2, style: str = "photo", skip_web_search: bool = False) -> Dict:
         """Gets or generates images for the query
         Args:
             query: What to get images of
             num_images: How many images to get
             style: Style for generated images
+            skip_web_search: Whether to skip web search
         Returns:
             Dict containing results and status
         """
         all_results = []
 
-        # First try web search
-        web_results = self._try_web_search(query, num_images)
-        all_results.extend(web_results)
+        # Try web search first unless skipped
+        if not skip_web_search:
+            web_results = self._try_web_search(query, num_images)
+            all_results.extend(web_results)
 
         # If we don't have enough images, try generation
         if len(all_results) < num_images:
@@ -143,5 +186,9 @@ class ImageHandlerTool(Tool):
         return {
             "status": "success",
             "images": all_results,
-            "total": len(all_results)
+            "total": len(all_results),
+            "sources": {
+                "web": len([img for img in all_results if img['source'] == 'web']),
+                "generated": len([img for img in all_results if img['source'] == 'generated'])
+            }
         }

tools/visit_webpage.py

@@ -1,42 +1,40 @@
+import re
 from typing import Any, Optional
 from smolagents.tools import Tool
 import requests
 import markdownify
 import smolagents
 
+
 class VisitWebpageTool(Tool):
     name = "visit_webpage"
     description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
-    inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
+    inputs = {'url': {'type': 'string',
+                      'description': 'The url of the webpage to visit.'}}
     output_type = "string"
 
     def forward(self, url: str) -> str:
         try:
-            import requests
-            from markdownify import markdownify
-            from requests.exceptions import RequestException
-
-            from smolagents.utils import truncate_content
-        except ImportError as e:
-            raise ImportError(
-                "You must install packages `markdownify` and `requests` to run this tool: for instance run `pip install markdownify requests`."
-            ) from e
-        try:
+            # Add headers to mimic a browser
+            headers = {
+                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+            }
+
             # Send a GET request to the URL with a 20-second timeout
-            response = requests.get(url, timeout=20)
-            response.raise_for_status()  # Raise an exception for bad status codes
+            response = requests.get(url, timeout=20, headers=headers)
+            response.raise_for_status()
 
             # Convert the HTML content to Markdown
-            markdown_content = markdownify(response.text).strip()
+            markdown_content = markdownify.markdownify(response.text).strip()
 
             # Remove multiple line breaks
             markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
 
-            return truncate_content(markdown_content, 10000)
+            return markdown_content[:10000]  # Limit content length
 
         except requests.exceptions.Timeout:
             return "The request timed out. Please try again later or check the URL."
-        except RequestException as e:
+        except requests.exceptions.RequestException as e:
             return f"Error fetching the webpage: {str(e)}"
         except Exception as e:
             return f"An unexpected error occurred: {str(e)}"