13972131258

Committed 20 Mar 2025 02:43PM UTC coverage: 90.021% (-0.03%) from 90.054%

Build # 13972131258

Build Type

Pull #9069

github

Committed by

web-flow

Commit Message

Merge 8371761b0 into 67ab3788e

Pull Request Pull Request #9069: refactor!: `ChatMessage` serialization-deserialization updates

Run Details

9833 of 10923 relevant lines covered (90.02%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.5

haystack/components/generators/openai_dalle.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import os
from typing import Any, Dict, List, Literal, Optional

from openai import OpenAI
from openai.types.image import Image

from haystack import component, default_from_dict, default_to_dict
from haystack.utils import Secret, deserialize_secrets_inplace


@component
class DALLEImageGenerator:
    """
    Generates images using OpenAI's DALL-E model.

    For details on OpenAI API parameters, see
    [OpenAI documentation](https://platform.openai.com/docs/api-reference/images/create).

    ### Usage example

    ```python
    from haystack.components.generators import DALLEImageGenerator
    image_generator = DALLEImageGenerator()
    response = image_generator.run("Show me a picture of a black cat.")
    print(response)
    ```
    """

    def __init__(  # pylint: disable=too-many-positional-arguments
        self,
        model: str = "dall-e-3",
        quality: Literal["standard", "hd"] = "standard",
        size: Literal["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"] = "1024x1024",
        response_format: Literal["url", "b64_json"] = "url",
        api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
        api_base_url: Optional[str] = None,
        organization: Optional[str] = None,
        timeout: Optional[float] = None,
        max_retries: Optional[int] = None,
    ):
        """
        Creates an instance of DALLEImageGenerator. Unless specified otherwise in `model`, uses OpenAI's dall-e-3.

        :param model: The model to use for image generation. Can be "dall-e-2" or "dall-e-3".
        :param quality: The quality of the generated image. Can be "standard" or "hd".
        :param size: The size of the generated images.
            Must be one of 256x256, 512x512, or 1024x1024 for dall-e-2.
            Must be one of 1024x1024, 1792x1024, or 1024x1792 for dall-e-3 models.
        :param response_format: The format of the response. Can be "url" or "b64_json".
        :param api_key: The OpenAI API key to connect to OpenAI.
        :param api_base_url: An optional base URL.
        :param organization: The Organization ID, defaults to `None`.
        :param timeout:
            Timeout for OpenAI Client calls. If not set, it is inferred from the `OPENAI_TIMEOUT` environment variable
            or set to 30.
        :param max_retries:
            Maximum retries to establish contact with OpenAI if it returns an internal error. If not set, it is inferred
            from the `OPENAI_MAX_RETRIES` environment variable or set to 5.
        """
        self.model = model
        self.quality = quality
        self.size = size
        self.response_format = response_format
        self.api_key = api_key
        self.api_base_url = api_base_url
        self.organization = organization

        self.timeout = timeout or float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
        self.max_retries = max_retries or int(os.environ.get("OPENAI_MAX_RETRIES", "5"))

        self.client: Optional[OpenAI] = None

    def warm_up(self) -> None:
        """
        Warm up the OpenAI client.
        """
        if self.client is None:
            self.client = OpenAI(
                api_key=self.api_key.resolve_value(),
                organization=self.organization,
                base_url=self.api_base_url,
                timeout=self.timeout,
                max_retries=self.max_retries,
            )

    @component.output_types(images=List[str], revised_prompt=str)
    def run(
        self,
        prompt: str,
        size: Optional[Literal["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"]] = None,
        quality: Optional[Literal["standard", "hd"]] = None,
        response_format: Optional[Optional[Literal["url", "b64_json"]]] = None,
    ):
        """
        Invokes the image generation inference based on the provided prompt and generation parameters.

        :param prompt: The prompt to generate the image.
        :param size: If provided, overrides the size provided during initialization.
        :param quality: If provided, overrides the quality provided during initialization.
        :param response_format: If provided, overrides the response format provided during initialization.

        :returns:
            A dictionary containing the generated list of images and the revised prompt.
            Depending on the `response_format` parameter, the list of images can be URLs or base64 encoded JSON strings.
            The revised prompt is the prompt that was used to generate the image, if there was any revision
            to the prompt made by OpenAI.
        """
        if self.client is None:
            raise RuntimeError(
                "The component DALLEImageGenerator wasn't warmed up. Run 'warm_up()' before calling 'run()'."
            )

        size = size or self.size
        quality = quality or self.quality
        response_format = response_format or self.response_format
        response = self.client.images.generate(
            model=self.model, prompt=prompt, size=size, quality=quality, response_format=response_format, n=1
        )
        image: Image = response.data[0]
        image_str = image.url or image.b64_json or ""
        return {"images": [image_str], "revised_prompt": image.revised_prompt or ""}

    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize this component to a dictionary.

        :returns:
            The serialized component as a dictionary.
        """
        return default_to_dict(  # type: ignore
            self,
            model=self.model,
            quality=self.quality,
            size=self.size,
            response_format=self.response_format,
            api_key=self.api_key.to_dict(),
            api_base_url=self.api_base_url,
            organization=self.organization,
        )

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DALLEImageGenerator":
        """
        Deserialize this component from a dictionary.

        :param data:
            The dictionary representation of this component.
        :returns:
            The deserialized component instance.
        """
        init_params = data.get("init_parameters", {})
        deserialize_secrets_inplace(init_params, keys=["api_key"])
        return default_from_dict(cls, data)  # type: ignore

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import os	1✔
6	from typing import Any, Dict, List, Literal, Optional	1✔
7
8	from openai import OpenAI	1✔
9	from openai.types.image import Image	1✔
10
11	from haystack import component, default_from_dict, default_to_dict	1✔
12	from haystack.utils import Secret, deserialize_secrets_inplace	1✔
13
14
15	@component	1✔
16	class DALLEImageGenerator:	1✔
17	"""
18	Generates images using OpenAI's DALL-E model.
19
20	For details on OpenAI API parameters, see
21	[OpenAI documentation](https://platform.openai.com/docs/api-reference/images/create).
22
23	### Usage example
24
25	```python
26	from haystack.components.generators import DALLEImageGenerator
27	image_generator = DALLEImageGenerator()
28	response = image_generator.run("Show me a picture of a black cat.")
29	print(response)
30	```
31	"""
32
33	def __init__( # pylint: disable=too-many-positional-arguments	1✔
34	self,
35	model: str = "dall-e-3",
36	quality: Literal["standard", "hd"] = "standard",
37	size: Literal["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"] = "1024x1024",
38	response_format: Literal["url", "b64_json"] = "url",
39	api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
40	api_base_url: Optional[str] = None,
41	organization: Optional[str] = None,
42	timeout: Optional[float] = None,
43	max_retries: Optional[int] = None,
44	):
45	"""
46	Creates an instance of DALLEImageGenerator. Unless specified otherwise in `model`, uses OpenAI's dall-e-3.
47
48	:param model: The model to use for image generation. Can be "dall-e-2" or "dall-e-3".
49	:param quality: The quality of the generated image. Can be "standard" or "hd".
50	:param size: The size of the generated images.
51	Must be one of 256x256, 512x512, or 1024x1024 for dall-e-2.
52	Must be one of 1024x1024, 1792x1024, or 1024x1792 for dall-e-3 models.
53	:param response_format: The format of the response. Can be "url" or "b64_json".
54	:param api_key: The OpenAI API key to connect to OpenAI.
55	:param api_base_url: An optional base URL.
56	:param organization: The Organization ID, defaults to `None`.
57	:param timeout:
58	Timeout for OpenAI Client calls. If not set, it is inferred from the `OPENAI_TIMEOUT` environment variable
59	or set to 30.
60	:param max_retries:
61	Maximum retries to establish contact with OpenAI if it returns an internal error. If not set, it is inferred
62	from the `OPENAI_MAX_RETRIES` environment variable or set to 5.
63	"""
64	self.model = model	1✔
65	self.quality = quality	1✔
66	self.size = size	1✔
67	self.response_format = response_format	1✔
68	self.api_key = api_key	1✔
69	self.api_base_url = api_base_url	1✔
70	self.organization = organization	1✔
71
72	self.timeout = timeout or float(os.environ.get("OPENAI_TIMEOUT", "30.0"))	1✔
73	self.max_retries = max_retries or int(os.environ.get("OPENAI_MAX_RETRIES", "5"))	1✔
74
75	self.client: Optional[OpenAI] = None	1✔
76
77	def warm_up(self) -> None:	1✔
78	"""
79	Warm up the OpenAI client.
80	"""
81	if self.client is None:	1✔
82	self.client = OpenAI(	1✔
83	api_key=self.api_key.resolve_value(),
84	organization=self.organization,
85	base_url=self.api_base_url,
86	timeout=self.timeout,
87	max_retries=self.max_retries,
88	)
89
90	@component.output_types(images=List[str], revised_prompt=str)	1✔
91	def run(	1✔
92	self,
93	prompt: str,
94	size: Optional[Literal["256x256", "512x512", "1024x1024", "1792x1024", "1024x1792"]] = None,
95	quality: Optional[Literal["standard", "hd"]] = None,
96	response_format: Optional[Optional[Literal["url", "b64_json"]]] = None,
97	):
98	"""
99	Invokes the image generation inference based on the provided prompt and generation parameters.
100
101	:param prompt: The prompt to generate the image.
102	:param size: If provided, overrides the size provided during initialization.
103	:param quality: If provided, overrides the quality provided during initialization.
104	:param response_format: If provided, overrides the response format provided during initialization.
105
106	:returns:
107	A dictionary containing the generated list of images and the revised prompt.
108	Depending on the `response_format` parameter, the list of images can be URLs or base64 encoded JSON strings.
109	The revised prompt is the prompt that was used to generate the image, if there was any revision
110	to the prompt made by OpenAI.
111	"""
112	if self.client is None:	1✔
113	raise RuntimeError(	×
114	"The component DALLEImageGenerator wasn't warmed up. Run 'warm_up()' before calling 'run()'."
115	)
116
117	size = size or self.size	1✔
118	quality = quality or self.quality	1✔
119	response_format = response_format or self.response_format	1✔
120	response = self.client.images.generate(	1✔
121	model=self.model, prompt=prompt, size=size, quality=quality, response_format=response_format, n=1
122	)
123	image: Image = response.data[0]	1✔
124	image_str = image.url or image.b64_json or ""	1✔
125	return {"images": [image_str], "revised_prompt": image.revised_prompt or ""}	1✔
126
127	def to_dict(self) -> Dict[str, Any]:	1✔
128	"""
129	Serialize this component to a dictionary.
130
131	:returns:
132	The serialized component as a dictionary.
133	"""
134	return default_to_dict( # type: ignore	1✔
135	self,
136	model=self.model,
137	quality=self.quality,
138	size=self.size,
139	response_format=self.response_format,
140	api_key=self.api_key.to_dict(),
141	api_base_url=self.api_base_url,
142	organization=self.organization,
143	)
144
145	@classmethod	1✔
146	def from_dict(cls, data: Dict[str, Any]) -> "DALLEImageGenerator":	1✔
147	"""
148	Deserialize this component from a dictionary.
149
150	:param data:
151	The dictionary representation of this component.
152	:returns:
153	The deserialized component instance.
154	"""
155	init_params = data.get("init_parameters", {})	1✔
156	deserialize_secrets_inplace(init_params, keys=["api_key"])	1✔
157	return default_from_dict(cls, data) # type: ignore	1✔

deepset-ai / haystack / 13972131258

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous