Coverage for haystack/components/generators/openai_image_generator.py: 100%
80 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import os
6from typing import Any, Literal
8from openai import AsyncOpenAI, OpenAI
9from openai.types.image import Image
11from haystack import component, default_from_dict, default_to_dict, logging
12from haystack.utils import Secret
13from haystack.utils.http_client import init_http_client
15logger = logging.getLogger(__name__)
18@component
19class OpenAIImageGenerator:
20 """
21 Generates images using OpenAI's image generation models such as `gpt-image-2`.
23 For details on OpenAI API parameters, see
24 [OpenAI documentation](https://developers.openai.com/api/reference/resources/images/methods/generate).
26 ### Usage example
27 ```python
28 from haystack.components.generators import OpenAIImageGenerator
29 image_generator = OpenAIImageGenerator()
30 response = image_generator.run("Show me a picture of a black cat.")
31 print(response)
32 ```
33 """
35 def __init__(
36 self,
37 model: str = "gpt-image-2",
38 quality: Literal["auto", "high", "medium", "low"] = "auto",
39 size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] = "1024x1024",
40 response_format: Literal["b64_json"] = "b64_json",
41 api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
42 api_base_url: str | None = None,
43 organization: str | None = None,
44 timeout: float | None = None,
45 max_retries: int | None = None,
46 http_client_kwargs: dict[str, Any] | None = None,
47 ) -> None:
48 """
49 Creates an instance of OpenAIImageGenerator. Unless specified otherwise in `model`, uses OpenAI's gpt-image-2.
51 :param model: The model to use for image generation. Model names can be found in the
52 [OpenAI documentation](https://developers.openai.com/api/docs/models/all).
53 :param quality: The quality of the generated image. Can be "auto", "high", "medium", or "low".
54 :param size: The size of the generated images. One of 1024x1024, 1024x1536, 1536x1024, or "auto".
55 `gpt-image-2` also supports arbitrary sizes. You can find more information about supported sizes in
56 the [OpenAI documentation](https://developers.openai.com/api/reference/resources/images/methods/generate).
57 :param response_format: This parameter is ignored and only kept for backward compatibility.
58 :param api_key: The OpenAI API key to connect to OpenAI.
59 :param api_base_url: An optional base URL.
60 :param organization: The Organization ID, defaults to `None`.
61 :param timeout:
62 Timeout for OpenAI Client calls. If not set, it is inferred from the `OPENAI_TIMEOUT` environment variable
63 or set to 30.
64 :param max_retries:
65 Maximum retries to establish contact with OpenAI if it returns an internal error. If not set, it is inferred
66 from the `OPENAI_MAX_RETRIES` environment variable or set to 5.
67 :param http_client_kwargs:
68 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
69 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
70 """
71 self.model = model
72 if quality not in ["auto", "high", "medium", "low"]:
73 logger.warning("Invalid quality: {quality}. Defaulting to 'auto'.", quality=quality)
74 quality = "auto"
75 self.quality = quality
76 self.size = size
77 if response_format != "b64_json":
78 logger.warning("response_format is ignored. A base64-encoded image will be returned.")
79 self.api_key = api_key
80 self.api_base_url = api_base_url
81 self.organization = organization
83 self.timeout = timeout
84 self.max_retries = max_retries
85 self.http_client_kwargs = http_client_kwargs
87 self.client: OpenAI | None = None
88 self.async_client: AsyncOpenAI | None = None
90 def _client_kwargs(self) -> dict[str, Any]:
91 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
92 max_retries = (
93 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
94 )
95 return {
96 "api_key": self.api_key.resolve_value(),
97 "organization": self.organization,
98 "base_url": self.api_base_url,
99 "timeout": timeout,
100 "max_retries": max_retries,
101 }
103 def warm_up(self) -> None:
104 """
105 Initializes the synchronous OpenAI client.
106 """
107 if self.client is None:
108 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
109 # https://github.com/openai/openai-python/blob/main/httpx2.md
110 http_client = init_http_client(self.http_client_kwargs, async_client=False)
111 self.client = OpenAI(
112 http_client=http_client, # type: ignore[arg-type]
113 **self._client_kwargs(),
114 )
116 async def warm_up_async(self) -> None: # noqa: RUF029
117 """
118 Initializes the asynchronous OpenAI client on the serving event loop.
119 """
120 if self.async_client is None:
121 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
122 # https://github.com/openai/openai-python/blob/main/httpx2.md
123 http_client = init_http_client(self.http_client_kwargs, async_client=True)
124 self.async_client = AsyncOpenAI(
125 http_client=http_client, # type: ignore[arg-type]
126 **self._client_kwargs(),
127 )
129 def close(self) -> None:
130 """
131 Releases the synchronous OpenAI client.
132 """
133 if self.client is not None:
134 self.client.close()
135 self.client = None
137 async def close_async(self) -> None:
138 """
139 Releases the asynchronous OpenAI client.
140 """
141 if self.async_client is not None:
142 await self.async_client.close()
143 self.async_client = None
145 @component.output_types(images=list[str], revised_prompt=str)
146 def run(
147 self,
148 prompt: str,
149 size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] | None = None,
150 quality: Literal["auto", "high", "medium", "low"] | None = None,
151 response_format: Literal["b64_json"] | None = None, # noqa: ARG002
152 ) -> dict[str, Any]:
153 """
154 Invokes the image generation inference based on the provided prompt and generation parameters.
156 :param prompt: The prompt to generate the image.
157 :param size: If provided, overrides the size provided during initialization.
158 :param quality: If provided, overrides the quality provided during initialization.
159 :param response_format: This parameter is ignored and only kept for backward compatibility.
161 :returns:
162 A dictionary containing the generated list of images as base64 encoded JSON strings and the revised prompt.
163 The revised prompt is the prompt that was used to generate the image, if there was any revision
164 to the prompt made by OpenAI.
165 """
166 self.warm_up()
168 # at this point the client is initialized, but mypy doesn't know that
169 assert self.client is not None
171 size = size or self.size
172 quality = quality or self.quality
173 response = self.client.images.generate(model=self.model, prompt=prompt, size=size, quality=quality, n=1)
174 image_str = ""
175 revised_prompt = ""
176 if response.data is not None:
177 image: Image = response.data[0]
178 image_str = image.b64_json or ""
179 revised_prompt = image.revised_prompt or ""
181 return {"images": [image_str], "revised_prompt": revised_prompt}
183 @component.output_types(images=list[str], revised_prompt=str)
184 async def run_async(
185 self,
186 prompt: str,
187 size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] | None = None,
188 quality: Literal["auto", "high", "medium", "low"] | None = None,
189 response_format: Literal["b64_json"] | None = None, # noqa: ARG002
190 ) -> dict[str, Any]:
191 """
192 Asynchronously invokes the image generation inference based on the provided prompt and generation parameters.
194 This is the asynchronous version of the `run` method. It has the same parameters and return values
195 but can be used with `await` in an async code.
197 :param prompt: The prompt to generate the image.
198 :param size: If provided, overrides the size provided during initialization.
199 :param quality: If provided, overrides the quality provided during initialization.
200 :param response_format: This parameter is ignored and only kept for backward compatibility.
202 :returns:
203 A dictionary containing the generated list of images as base64 encoded JSON strings and the revised prompt.
204 The revised prompt is the prompt that was used to generate the image, if there was any revision
205 to the prompt made by OpenAI.
206 """
207 await self.warm_up_async()
209 # at this point the client is initialized, but mypy doesn't know that
210 assert self.async_client is not None
212 size = size or self.size
213 quality = quality or self.quality
214 response = await self.async_client.images.generate(
215 model=self.model, prompt=prompt, size=size, quality=quality, n=1
216 )
217 image_str = ""
218 revised_prompt = ""
219 if response.data is not None:
220 image: Image = response.data[0]
221 image_str = image.b64_json or ""
222 revised_prompt = image.revised_prompt or ""
224 return {"images": [image_str], "revised_prompt": revised_prompt}
226 def to_dict(self) -> dict[str, Any]:
227 """
228 Serialize this component to a dictionary.
230 :returns:
231 The serialized component as a dictionary.
232 """
233 return default_to_dict(
234 self,
235 model=self.model,
236 quality=self.quality,
237 size=self.size,
238 api_key=self.api_key,
239 api_base_url=self.api_base_url,
240 organization=self.organization,
241 http_client_kwargs=self.http_client_kwargs,
242 )
244 @classmethod
245 def from_dict(cls, data: dict[str, Any]) -> "OpenAIImageGenerator":
246 """
247 Deserialize this component from a dictionary.
249 :param data:
250 The dictionary representation of this component.
251 :returns:
252 The deserialized component instance.
253 """
254 return default_from_dict(cls, data)