Coverage for haystack/components/embedders/openai_text_embedder.py: 88%
73 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
8from openai import AsyncOpenAI, OpenAI
9from openai.types import CreateEmbeddingResponse
11from haystack import component, default_from_dict, default_to_dict
12from haystack.utils import Secret
13from haystack.utils.http_client import init_http_client
16@component
17class OpenAITextEmbedder:
18 """
19 Embeds strings using OpenAI models.
21 You can use it to embed user query and send it to an embedding Retriever.
23 ### Usage example
24 <!-- test-ignore -->
25 ```python
26 from haystack.components.embedders import OpenAITextEmbedder
28 text_to_embed = "I love pizza!"
29 text_embedder = OpenAITextEmbedder()
31 print(text_embedder.run(text_to_embed))
33 # {'embedding': [0.017020374536514282, -0.023255806416273117, ...],
34 # 'meta': {'model': 'text-embedding-ada-002-v2',
35 # 'usage': {'prompt_tokens': 4, 'total_tokens': 4}}}
36 ```
37 """
39 def __init__(
40 self,
41 api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
42 model: str = "text-embedding-ada-002",
43 dimensions: int | None = None,
44 api_base_url: str | None = None,
45 organization: str | None = None,
46 prefix: str = "",
47 suffix: str = "",
48 timeout: float | None = None,
49 max_retries: int | None = None,
50 http_client_kwargs: dict[str, Any] | None = None,
51 ) -> None:
52 """
53 Creates an OpenAITextEmbedder component.
55 Before initializing the component, you can set the 'OPENAI_TIMEOUT' and 'OPENAI_MAX_RETRIES'
56 environment variables to override the `timeout` and `max_retries` parameters respectively
57 in the OpenAI client.
59 :param api_key:
60 The OpenAI API key.
61 You can set it with an environment variable `OPENAI_API_KEY`, or pass with this parameter
62 during initialization.
63 :param model:
64 The name of the model to use for calculating embeddings.
65 The default model is `text-embedding-ada-002`.
66 :param dimensions:
67 The number of dimensions of the resulting embeddings. Only `text-embedding-3` and
68 later models support this parameter.
69 :param api_base_url:
70 Overrides default base URL for all HTTP requests.
71 :param organization:
72 Your organization ID. See OpenAI's
73 [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
74 for more information.
75 :param prefix:
76 A string to add at the beginning of each text to embed.
77 :param suffix:
78 A string to add at the end of each text to embed.
79 :param timeout:
80 Timeout for OpenAI client calls. If not set, it defaults to either the
81 `OPENAI_TIMEOUT` environment variable, or 30 seconds.
82 :param max_retries:
83 Maximum number of retries to contact OpenAI after an internal error.
84 If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5.
85 :param http_client_kwargs:
86 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
87 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
88 """
89 self.model = model
90 self.dimensions = dimensions
91 self.api_base_url = api_base_url
92 self.organization = organization
93 self.prefix = prefix
94 self.suffix = suffix
95 self.api_key = api_key
96 self.timeout = timeout
97 self.max_retries = max_retries
98 self.http_client_kwargs = http_client_kwargs
100 self.client: OpenAI | None = None
101 self.async_client: AsyncOpenAI | None = None
103 def _client_kwargs(self) -> dict[str, Any]:
104 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
105 max_retries = (
106 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
107 )
108 return {
109 "api_key": self.api_key.resolve_value(),
110 "organization": self.organization,
111 "base_url": self.api_base_url,
112 "timeout": timeout,
113 "max_retries": max_retries,
114 }
116 def warm_up(self) -> None:
117 """
118 Initializes the synchronous OpenAI client.
119 """
120 if self.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=False)
124 self.client = OpenAI(
125 http_client=http_client, # type: ignore[arg-type]
126 **self._client_kwargs(),
127 )
129 async def warm_up_async(self) -> None: # noqa: RUF029
130 """
131 Initializes the asynchronous OpenAI client on the serving event loop.
132 """
133 if self.async_client is None:
134 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
135 # https://github.com/openai/openai-python/blob/main/httpx2.md
136 http_client = init_http_client(self.http_client_kwargs, async_client=True)
137 self.async_client = AsyncOpenAI(
138 http_client=http_client, # type: ignore[arg-type]
139 **self._client_kwargs(),
140 )
142 def close(self) -> None:
143 """
144 Releases the synchronous OpenAI client.
145 """
146 if self.client is not None:
147 self.client.close()
148 self.client = None
150 async def close_async(self) -> None:
151 """
152 Releases the asynchronous OpenAI client.
153 """
154 if self.async_client is not None:
155 await self.async_client.close()
156 self.async_client = None
158 def _get_telemetry_data(self) -> dict[str, Any]:
159 """
160 Data that is sent to Posthog for usage analytics.
161 """
162 return {"model": self.model}
164 def to_dict(self) -> dict[str, Any]:
165 """
166 Serializes the component to a dictionary.
168 :returns:
169 Dictionary with serialized data.
170 """
171 return default_to_dict(
172 self,
173 api_key=self.api_key,
174 model=self.model,
175 dimensions=self.dimensions,
176 api_base_url=self.api_base_url,
177 organization=self.organization,
178 prefix=self.prefix,
179 suffix=self.suffix,
180 timeout=self.timeout,
181 max_retries=self.max_retries,
182 http_client_kwargs=self.http_client_kwargs,
183 )
185 @classmethod
186 def from_dict(cls, data: dict[str, Any]) -> "OpenAITextEmbedder":
187 """
188 Deserializes the component from a dictionary.
190 :param data:
191 Dictionary to deserialize from.
192 :returns:
193 Deserialized component.
194 """
195 return default_from_dict(cls, data)
197 def _prepare_input(self, text: str) -> dict[str, Any]:
198 if not isinstance(text, str):
199 raise TypeError(
200 "OpenAITextEmbedder expects a string as an input."
201 "In case you want to embed a list of Documents, please use the OpenAIDocumentEmbedder."
202 )
204 text_to_embed = self.prefix + text + self.suffix
206 kwargs: dict[str, Any] = {"model": self.model, "input": text_to_embed, "encoding_format": "float"}
207 if self.dimensions is not None:
208 kwargs["dimensions"] = self.dimensions
209 return kwargs
211 def _prepare_output(self, result: CreateEmbeddingResponse) -> dict[str, Any]:
212 return {"embedding": result.data[0].embedding, "meta": {"model": result.model, "usage": dict(result.usage)}}
214 @component.output_types(embedding=list[float], meta=dict[str, Any])
215 def run(self, text: str) -> dict[str, Any]:
216 """
217 Embeds a single string.
219 :param text:
220 Text to embed.
222 :returns:
223 A dictionary with the following keys:
224 - `embedding`: The embedding of the input text.
225 - `meta`: Information about the usage of the model.
226 """
227 self.warm_up()
228 create_kwargs = self._prepare_input(text=text)
229 assert self.client is not None # mypy: client is built by warm_up above
230 response = self.client.embeddings.create(**create_kwargs)
231 return self._prepare_output(result=response)
233 @component.output_types(embedding=list[float], meta=dict[str, Any])
234 async def run_async(self, text: str) -> dict[str, Any]:
235 """
236 Asynchronously embed a single string.
238 This is the asynchronous version of the `run` method. It has the same parameters and return values
239 but can be used with `await` in async code.
241 :param text:
242 Text to embed.
244 :returns:
245 A dictionary with the following keys:
246 - `embedding`: The embedding of the input text.
247 - `meta`: Information about the usage of the model.
248 """
249 await self.warm_up_async()
250 create_kwargs = self._prepare_input(text=text)
251 assert self.async_client is not None # mypy: async_client is built by warm_up_async above
252 response = await self.async_client.embeddings.create(**create_kwargs)
253 return self._prepare_output(result=response)