Skip to content
This repository has been archived by the owner on Mar 1, 2024. It is now read-only.

feat: OpenAI Image Generation Tool #628

Merged
merged 10 commits into from
Nov 18, 2023
Prev Previous commit
Next Next commit
chore: fix tests and lint
EmanuelCampos committed Nov 11, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit 2c098d762a66a51d35c2f9c388017eff9434f5be
26 changes: 15 additions & 11 deletions llama_hub/tools/openai/image_generation/base.py
Original file line number Diff line number Diff line change
@@ -7,11 +7,9 @@
from typing import Optional
from llama_index.tools.tool_spec.base import BaseToolSpec

from PIL import Image
from io import BytesIO

DEFAULT_CACHE_DIR = "../../../img_cache"


class OpenAIImageGenerationToolSpec(BaseToolSpec):
"""OpenAI Image Generation tool spec."""

@@ -28,14 +26,20 @@ def __init__(self, api_key: str, cache_dir: Optional[str] = None) -> None:
"""Initialize with parameters."""
self.client = OpenAI(api_key=api_key)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

api_key should probably be optional, since it could be in os.environ

self.cache_dir = cache_dir or DEFAULT_CACHE_DIR



def get_cache_dir(self):
return self.cache_dir

def save_base64_image(self, base64_str, image_name):
try:
from PIL import Image
from io import BytesIO
except ImportError:
raise ImportError(
"Please install Pillow with `pip install Pillow` to use this tool"
)
cache_dir = self.cache_dir

# Create cache directory if it doesn't exist
if not os.path.exists(cache_dir):
os.makedirs(cache_dir)
@@ -76,11 +80,11 @@ def image_generation(
n=num_images,
response_format="b64_json",
)

image_bytes = response.data[0].b64_json

filename = f"{time.time()}.jpg"
saved_image_path = self.save_base64_image(image_bytes, filename);

saved_image_path = self.save_base64_image(image_bytes, filename)

return saved_image_path