from smolagents import Tool
from typing import Any, Optional

class SimpleTool(Tool):
    name = "pollinations_generate"
    description = "Generate an image from a prompt"
    inputs = {'prompt': {'type': 'string', 'description': 'The prompt to generate the image from'}, 'width': {'type': 'integer', 'nullable': True, 'description': 'image width'}, 'height': {'type': 'integer', 'nullable': True, 'description': 'image height'}, 'seed': {'type': 'integer', 'nullable': True, 'description': 'A decimal integer  Each seed generates a new image variation Returns generated image as a PILL image'}}
    output_type = "object"

    def forward(self, prompt: str, width: int = 1024, height: int = 1024, seed: int = 0) -> 'PIL.Image.Image':
      """
      Generate an image from a prompt
         Args:
             prompt : The prompt to generate the image from
             width : image width
             height : image height
             seed : A decimal integer  Each seed generates a new image variation
          Returns
            generated image as a PILL image

      """
      import PIL
      from PIL import Image
      import requests
      import random
      from io import BytesIO
      if seed ==0:
        seed=random.randint(1,8*1024**2)
      else:
        random.seed(seed)

      url = f"https://pollinations.ai/p/{prompt}"
      params={"width":width,"height":height,"seed":seed}
      response = requests.get(url)
      return Image.open(requests.get(url, stream=True,params=params).raw)