import os

from PIL import Image
from IPython.display import display

## Setting up the Agent

from llama_index.tools.openai.image_generation import OpenAIImageGenerationToolSpec
from llama_index.agent import ReActAgent
from llama_index.tools import FunctionTool


def show_image(filename: str) -> Image:
    """Display an image based on the filename"""
    img = Image.open(filename)
    return display(img)


image_generation_tool = OpenAIImageGenerationToolSpec(
    api_key=os.environ["OPENAI_API_KEY"]
)
show_image_tool = FunctionTool.from_defaults(fn=show_image)

agent = ReActAgent.from_tools(
    [*image_generation_spec.to_tool_list(), show_image_tool], verbose=True
)

response = agent.query(
    "generate a hacker image with size 1024x1024, use the filename and show the image"
)