diff --git a/prem_utils/connectors/openai.py b/prem_utils/connectors/openai.py index bcc5c54..5c26670 100644 --- a/prem_utils/connectors/openai.py +++ b/prem_utils/connectors/openai.py @@ -305,10 +305,17 @@ def generate_image( n: int = 1, quality: str = "standard", style: str = "vivid", + response_format: str = "url", ): try: response = self.client.images.generate( - model=model, prompt=prompt, n=n, size=size, quality=quality, style=style + model=model, + prompt=prompt, + n=n, + size=size, + quality=quality, + style=style, + response_format=response_format, ) except ( NotFoundError, @@ -326,4 +333,4 @@ def generate_image( ) as error: custom_exception = self.exception_mapping.get(type(error), errors.PremProviderError) raise custom_exception(error, provider="openai", model=model, provider_message=str(error)) - return {"created": None, "data": [{"url": image.url} for image in response.data]} + return {"created": None, "data": [{"url": image.url, "b64_json": image.b64_json} for image in response.data]}