image_search_tool.py

import requests
from bs4 import BeautifulSoup
from PIL import Image
from io import BytesIO
import json
import os
from langchain.tools import tool
# Ensure that this tool returns the correct type (e.g., a string path or a list of strings)
@tool
def image_search_tool(query: str, num_images: int = 5) -> list:
    """
    A tool to search for images based on the given query and download them.
    Parameters:
    - query (str): The search query for finding images.
    - num_images (int): The number of images to retrieve (default is 5).
    Returns:
    - list: The paths where the downloaded images are saved.
    """
    url = 'https://www.bing.com/images/search'
    headers = {
        'User-Agent': 'Mozilla/5.0'
    }
    params = {'q': query, 'count': num_images}
    response = requests.get(url, headers=headers, params=params)

    if response.status_code != 200:
        return [f"Failed to retrieve results: {response.status_code}"]

    soup = BeautifulSoup(response.text, 'html.parser')
    results = []
    
    for img_tag in soup.find_all('a', class_='iusc')[:num_images]:
        m = img_tag.get('m')
        m_json = json.loads(m)
        img_url = m_json.get('murl')
        results.append(img_url)
    
    image_paths = []
    folder = 'images'
    if not os.path.exists(folder):
        os.makedirs(folder)

    for i, img_url in enumerate(results, 1):
        try:
            img_response = requests.get(img_url)
            img = Image.open(BytesIO(img_response.content))
            img_path = os.path.join(folder, f"image_{i}.jpg")
            img.save(img_path)
            image_paths.append(img_path)
        except Exception as e:
            return [f"Failed to download {img_url}: {e}"]

    return image_paths