From b343c71b9bcecf1b77aa8a98eb79e0bf5d7897a4 Mon Sep 17 00:00:00 2001
From: Eduardo Chiarotti <dudumelgaco@hotmail.com>
Date: Tue, 30 Jul 2024 22:17:23 -0300
Subject: [PATCH] feat: Add Dall-E tool to generate images

---
 src/crewai_tools/tools/dalle_tool/README.MD   |  0
 .../tools/dalle_tool/dalle_tool.py            | 48 +++++++++++++++++++
 2 files changed, 48 insertions(+)
 create mode 100644 src/crewai_tools/tools/dalle_tool/README.MD
 create mode 100644 src/crewai_tools/tools/dalle_tool/dalle_tool.py

diff --git a/src/crewai_tools/tools/dalle_tool/README.MD b/src/crewai_tools/tools/dalle_tool/README.MD
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/crewai_tools/tools/dalle_tool/dalle_tool.py b/src/crewai_tools/tools/dalle_tool/dalle_tool.py
new file mode 100644
index 000000000..a4f7738a6
--- /dev/null
+++ b/src/crewai_tools/tools/dalle_tool/dalle_tool.py
@@ -0,0 +1,48 @@
+import json
+from typing import Type
+
+from crewai_tools.tools.base_tool import BaseTool
+from openai import OpenAI
+from pydantic.v1 import BaseModel
+
+
+class ImagePromptSchema(BaseModel):
+    """Input for Dall-E Tool."""
+
+    image_description: str = "Description of the image to be generated by Dall-E."
+
+
+class DallETool(BaseTool):
+    name: str = "Dall-E Tool"
+    description: str = "Generates images using OpenAI's Dall-E model."
+    args_schema: Type[BaseModel] = ImagePromptSchema
+
+    model: str = "dall-e-3"
+    size: str = "1024x1024"
+    quality: str = "standard"
+    n: int = 1
+
+    def _run(self, **kwargs) -> str:
+        client = OpenAI()
+
+        image_description = kwargs.get("image_description")
+
+        if not image_description:
+            return "Image description is required."
+
+        response = client.images.generate(
+            model=self.model,
+            prompt=image_description,
+            size=self.size,
+            quality=self.quality,
+            n=self.n,
+        )
+
+        image_data = json.dumps(
+            {
+                "image_url": response.data[0].url,
+                "image_description": response.data[0].revised_prompt,
+            }
+        )
+
+        return image_data