better image getter

app.py

@@ -12,6 +12,7 @@ from tools.visit_webpage import VisitWebpageTool
 from tools.web_search import DuckDuckGoSearchTool
 from tools.html_to_wp_blocks import HTMLToWPBlocksTool
 from tools.wordpress_media import WordPressMediaTool
+from tools.image_handler import ImageHandlerTool
 
 from Gradio_UI import GradioUI
 
@@ -62,8 +63,8 @@ def quick_research(query: str, max_results: int = 3) -> str:
     try:
         # Add travel-specific terms to the query
         travel_query = f"{query} travel guide tourist attractions things to do"
-        search_results = search_tool.search(
-            travel_query, max_results=max_results)
+        search_results = search_tool.forward(
+            query=travel_query, max_results=max_results)
 
         if not search_results:
             return "No travel information found for the given destination."
@@ -152,6 +153,17 @@ model = HfApiModel(
 image_generation_tool = load_tool(
     "agents-course/text-to-image", trust_remote_code=True)
 
+# Create temp directory for images
+image_temp_dir = os.path.join(os.getcwd(), 'temp_images')
+os.makedirs(image_temp_dir, exist_ok=True)
+
+# Initialize tools
+image_handler = ImageHandlerTool(
+    web_search_tool=web_search,
+    image_gen_tool=image_generation_tool,
+    temp_dir=image_temp_dir
+)
+
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
@@ -166,7 +178,7 @@ agent = CodeAgent(
     tools=[
         final_answer,
         # get_current_time_in_timezone,
-        image_generation_tool,
+        image_handler,
         quick_research,
         # get_weather,
         wordpress_post,
@@ -185,32 +197,24 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
-# Update system prompt to include WordPress capabilities
+# Update system prompt to include image handling capabilities
 prompt_templates["system_prompt"] += """
+You can now handle images in blog posts more effectively:
+- image_handler: Gets images from web search or generates them if needed
+- wordpress_media: Uploads media files to WordPress
+- html_to_blocks: Converts HTML content to WordPress blocks, including image blocks
 
-You are also capable of managing a WordPress blog through the following tools:
-- wordpress_post: Publishes posts to WordPress
-- wordpress_media: Uploads media files to WordPress and returns attachment URLs
-- blog_generator: Generates AI-written blog posts
-- html_to_blocks: Converts HTML content to WordPress Gutenberg blocks format
-
-Always check credentials before attempting to post content.
-
-Example WordPress workflow:
-1. Set credentials (first time only)
-2. Generate blog content
-3. Generate or prepare media files
-4. Upload media files to WordPress using wordpress_media tool
-5. Convert HTML content to WordPress blocks format, including media blocks
-6. Publish blocks-formatted content to WordPress
+When creating posts with images:
+1. Use image_handler to get images (it will try web search first, then generation)
+2. Upload obtained images using wordpress_media
+3. Include the WordPress media URLs in your HTML content
+4. Convert to blocks and publish
 
 Remember to:
-- ALWAYS convert HTML content to WordPress blocks format before publishing
-- Upload any media files before referencing them in posts
-- Validate WordPress credentials before posting
-- Generate high-quality, relevant content
-- Handle errors gracefully
-- Provide clear status updates to the user
+- Handle image search/generation errors gracefully
+- Provide appropriate alt text and titles for images
+- Include image attribution if using web images
+- Generate images as a fallback when web search fails
 """
 
 GradioUI(agent).launch()

tools/image_handler.py

@@ -0,0 +1,147 @@
+import os
+import tempfile
+import requests
+from typing import Dict, List, Optional
+from smolagents.tools import Tool
+import time
+
+
+class ImageHandlerTool(Tool):
+    name = "image_handler"
+    description = "Gets or generates images for a given topic, with fallback options"
+    inputs = {
+        'query': {'type': 'string', 'description': 'The topic to get images for'},
+        'num_images': {
+            'type': 'integer',
+            'description': 'Number of images to get/generate',
+            'nullable': True
+        },
+        'style': {
+            'type': 'string',
+            'description': 'Style for generated images (e.g., "photo", "artistic")',
+            'nullable': True
+        }
+    }
+    output_type = "object"
+
+    def __init__(self, web_search_tool, image_gen_tool, temp_dir=None):
+        super().__init__()
+        self.web_search = web_search_tool
+        self.image_gen = image_gen_tool
+        self.temp_dir = temp_dir or tempfile.gettempdir()
+
+    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)
+            response.raise_for_status()
+
+            # Ensure temp directory exists
+            os.makedirs(self.temp_dir, exist_ok=True)
+
+            # Save the image
+            file_path = os.path.join(self.temp_dir, filename)
+            with open(file_path, 'wb') as f:
+                f.write(response.content)
+            return file_path
+        except Exception as e:
+            print(f"Failed to download image from {url}: {str(e)}")
+            return None
+
+    def _try_web_search(self, query: str, num_images: int) -> List[Dict]:
+        """Attempts to find images via web search"""
+        results = []
+        try:
+            # Try different search queries
+            search_queries = [
+                f"{query} high quality photo",
+                f"{query} professional photograph",
+                f"{query} travel photo"
+            ]
+
+            for search_query in search_queries:
+                if len(results) >= num_images:
+                    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']
+                            })
+
+        except Exception as e:
+            print(f"Web search failed: {str(e)}")
+
+        return results
+
+    def _generate_images(self, query: str, num_images: int, style: str = "photo") -> List[Dict]:
+        """Generates images using the image generation tool"""
+        results = []
+        try:
+            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
+                    })
+
+        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:
+        """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
+        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)
+
+        # If we don't have enough images, try generation
+        if len(all_results) < num_images:
+            remaining = num_images - len(all_results)
+            generated = self._generate_images(query, remaining, style)
+            all_results.extend(generated)
+
+        if not all_results:
+            return {
+                "status": "error",
+                "message": "Failed to get any images"
+            }
+
+        return {
+            "status": "success",
+            "images": all_results,
+            "total": len(all_results)
+        }