Coverage for haystack/components/embedders/openai_document_embedder.py: 70%
145 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 dataclasses import replace
7from typing import Any
9from more_itertools import batched
10from openai import APIError, AsyncOpenAI, OpenAI
11from tqdm import tqdm
12from tqdm.asyncio import tqdm as async_tqdm
14from haystack import Document, component, default_from_dict, default_to_dict, logging
15from haystack.utils import Secret
16from haystack.utils.http_client import init_http_client
18logger = logging.getLogger(__name__)
21@component
22class OpenAIDocumentEmbedder:
23 """
24 Computes document embeddings using OpenAI models.
26 ### Usage example
27 <!-- test-ignore -->
28 ```python
29 from haystack import Document
30 from haystack.components.embedders import OpenAIDocumentEmbedder
32 doc = Document(content="I love pizza!")
33 document_embedder = OpenAIDocumentEmbedder()
34 result = document_embedder.run([doc])
36 print(result['documents'][0].embedding)
38 # [0.017020374536514282, -0.023255806416273117, ...]
39 ```
40 """
42 def __init__( # noqa: PLR0913, PLR0917 (too-many-arguments, too-many-positional-arguments)
43 self,
44 api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
45 model: str = "text-embedding-ada-002",
46 dimensions: int | None = None,
47 api_base_url: str | None = None,
48 organization: str | None = None,
49 prefix: str = "",
50 suffix: str = "",
51 batch_size: int = 32,
52 progress_bar: bool = True,
53 meta_fields_to_embed: list[str] | None = None,
54 embedding_separator: str = "\n",
55 timeout: float | None = None,
56 max_retries: int | None = None,
57 http_client_kwargs: dict[str, Any] | None = None,
58 *,
59 raise_on_failure: bool = False,
60 ) -> None:
61 """
62 Creates an OpenAIDocumentEmbedder component.
64 Before initializing the component, you can set the 'OPENAI_TIMEOUT' and 'OPENAI_MAX_RETRIES'
65 environment variables to override the `timeout` and `max_retries` parameters respectively
66 in the OpenAI client.
68 :param api_key:
69 The OpenAI API key.
70 You can set it with an environment variable `OPENAI_API_KEY`, or pass with this parameter
71 during initialization.
72 :param model:
73 The name of the model to use for calculating embeddings.
74 The default model is `text-embedding-ada-002`.
75 :param dimensions:
76 The number of dimensions of the resulting embeddings. Only `text-embedding-3` and
77 later models support this parameter.
78 :param api_base_url:
79 Overrides the default base URL for all HTTP requests.
80 :param organization:
81 Your OpenAI organization ID. See OpenAI's
82 [Setting Up Your Organization](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
83 for more information.
84 :param prefix:
85 A string to add at the beginning of each text.
86 :param suffix:
87 A string to add at the end of each text.
88 :param batch_size:
89 Number of documents to embed at once.
90 :param progress_bar:
91 If `True`, shows a progress bar when running.
92 :param meta_fields_to_embed:
93 List of metadata fields to embed along with the document text.
94 :param embedding_separator:
95 Separator used to concatenate the metadata fields to the document text.
96 :param timeout:
97 Timeout for OpenAI client calls. If not set, it defaults to either the
98 `OPENAI_TIMEOUT` environment variable, or 30 seconds.
99 :param max_retries:
100 Maximum number of retries to contact OpenAI after an internal error.
101 If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or 5 retries.
102 :param http_client_kwargs:
103 A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
104 For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
105 :param raise_on_failure:
106 Whether to raise an exception if the embedding request fails. If `False`, the component will log the error
107 and continue processing the remaining documents. If `True`, it will raise an exception on failure.
108 """
109 self.api_key = api_key
110 self.model = model
111 self.dimensions = dimensions
112 self.api_base_url = api_base_url
113 self.organization = organization
114 self.prefix = prefix
115 self.suffix = suffix
116 self.batch_size = batch_size
117 self.progress_bar = progress_bar
118 self.meta_fields_to_embed = meta_fields_to_embed or []
119 self.embedding_separator = embedding_separator
120 self.timeout = timeout
121 self.max_retries = max_retries
122 self.http_client_kwargs = http_client_kwargs
123 self.raise_on_failure = raise_on_failure
125 self.client: OpenAI | None = None
126 self.async_client: AsyncOpenAI | None = None
128 def _client_kwargs(self) -> dict[str, Any]:
129 timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0"))
130 max_retries = (
131 self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5"))
132 )
133 return {
134 "api_key": self.api_key.resolve_value(),
135 "organization": self.organization,
136 "base_url": self.api_base_url,
137 "timeout": timeout,
138 "max_retries": max_retries,
139 }
141 def warm_up(self) -> None:
142 """
143 Initializes the synchronous OpenAI client.
144 """
145 if self.client is None:
146 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
147 # https://github.com/openai/openai-python/blob/main/httpx2.md
148 http_client = init_http_client(self.http_client_kwargs, async_client=False)
149 self.client = OpenAI(
150 http_client=http_client, # type: ignore[arg-type]
151 **self._client_kwargs(),
152 )
154 async def warm_up_async(self) -> None: # noqa: RUF029
155 """
156 Initializes the asynchronous OpenAI client on the serving event loop.
157 """
158 if self.async_client is None:
159 # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
160 # https://github.com/openai/openai-python/blob/main/httpx2.md
161 http_client = init_http_client(self.http_client_kwargs, async_client=True)
162 self.async_client = AsyncOpenAI(
163 http_client=http_client, # type: ignore[arg-type]
164 **self._client_kwargs(),
165 )
167 def close(self) -> None:
168 """
169 Releases the synchronous OpenAI client.
170 """
171 if self.client is not None:
172 self.client.close()
173 self.client = None
175 async def close_async(self) -> None:
176 """
177 Releases the asynchronous OpenAI client.
178 """
179 if self.async_client is not None:
180 await self.async_client.close()
181 self.async_client = None
183 def _get_telemetry_data(self) -> dict[str, Any]:
184 """
185 Data that is sent to Posthog for usage analytics.
186 """
187 return {"model": self.model}
189 def to_dict(self) -> dict[str, Any]:
190 """
191 Serializes the component to a dictionary.
193 :returns:
194 Dictionary with serialized data.
195 """
196 return default_to_dict(
197 self,
198 api_key=self.api_key,
199 model=self.model,
200 dimensions=self.dimensions,
201 api_base_url=self.api_base_url,
202 organization=self.organization,
203 prefix=self.prefix,
204 suffix=self.suffix,
205 batch_size=self.batch_size,
206 progress_bar=self.progress_bar,
207 meta_fields_to_embed=self.meta_fields_to_embed,
208 embedding_separator=self.embedding_separator,
209 timeout=self.timeout,
210 max_retries=self.max_retries,
211 http_client_kwargs=self.http_client_kwargs,
212 raise_on_failure=self.raise_on_failure,
213 )
215 @classmethod
216 def from_dict(cls, data: dict[str, Any]) -> "OpenAIDocumentEmbedder":
217 """
218 Deserializes the component from a dictionary.
220 :param data:
221 Dictionary to deserialize from.
222 :returns:
223 Deserialized component.
224 """
225 return default_from_dict(cls, data)
227 def _prepare_texts_to_embed(self, documents: list[Document]) -> dict[str, str]:
228 """
229 Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
230 """
231 texts_to_embed = {}
232 for doc in documents:
233 meta_values_to_embed = [
234 str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
235 ]
237 texts_to_embed[doc.id] = (
238 self.prefix + self.embedding_separator.join(meta_values_to_embed + [doc.content or ""]) + self.suffix
239 )
241 return texts_to_embed
243 def _embed_batch(
244 self, texts_to_embed: dict[str, str], batch_size: int
245 ) -> tuple[dict[str, list[float]], dict[str, Any]]:
246 """
247 Embed a list of texts in batches.
248 """
250 doc_ids_to_embeddings: dict[str, list[float]] = {}
251 meta: dict[str, Any] = {}
252 for batch in tqdm(
253 batched(texts_to_embed.items(), batch_size), disable=not self.progress_bar, desc="Calculating embeddings"
254 ):
255 args: dict[str, Any] = {"model": self.model, "input": [b[1] for b in batch], "encoding_format": "float"}
257 if self.dimensions is not None:
258 args["dimensions"] = self.dimensions
260 try:
261 # this method is invoked after warm_up, so client is not None
262 assert self.client is not None
263 response = self.client.embeddings.create(**args)
264 except APIError as exc:
265 ids = ", ".join(b[0] for b in batch)
266 msg = "Failed embedding of documents {ids} caused by {exc}"
267 logger.exception(msg, ids=ids, exc=exc)
268 if self.raise_on_failure:
269 raise exc
270 continue
272 embeddings = [el.embedding for el in response.data]
273 doc_ids_to_embeddings.update(dict(zip((b[0] for b in batch), embeddings, strict=True)))
275 if "model" not in meta:
276 meta["model"] = response.model
277 if "usage" not in meta:
278 meta["usage"] = dict(response.usage)
279 else:
280 meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
281 meta["usage"]["total_tokens"] += response.usage.total_tokens
283 return doc_ids_to_embeddings, meta
285 async def _embed_batch_async(
286 self, texts_to_embed: dict[str, str], batch_size: int
287 ) -> tuple[dict[str, list[float]], dict[str, Any]]:
288 """
289 Embed a list of texts in batches asynchronously.
290 """
292 doc_ids_to_embeddings: dict[str, list[float]] = {}
293 meta: dict[str, Any] = {}
295 batches = list(batched(texts_to_embed.items(), batch_size))
296 if self.progress_bar:
297 batches = async_tqdm(batches, desc="Calculating embeddings")
299 for batch in batches:
300 args: dict[str, Any] = {"model": self.model, "input": [b[1] for b in batch]}
302 if self.dimensions is not None:
303 args["dimensions"] = self.dimensions
305 try:
306 # this method is invoked after warm_up_async, so async_client is not None
307 assert self.async_client is not None
308 response = await self.async_client.embeddings.create(**args)
309 except APIError as exc:
310 ids = ", ".join(b[0] for b in batch)
311 msg = "Failed embedding of documents {ids} caused by {exc}"
312 logger.exception(msg, ids=ids, exc=exc)
313 if self.raise_on_failure:
314 raise exc
315 continue
317 embeddings = [el.embedding for el in response.data]
318 doc_ids_to_embeddings.update(dict(zip((b[0] for b in batch), embeddings, strict=True)))
320 if "model" not in meta:
321 meta["model"] = response.model
322 if "usage" not in meta:
323 meta["usage"] = dict(response.usage)
324 else:
325 meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
326 meta["usage"]["total_tokens"] += response.usage.total_tokens
328 return doc_ids_to_embeddings, meta
330 @component.output_types(documents=list[Document], meta=dict[str, Any])
331 def run(self, documents: list[Document]) -> dict[str, Any]:
332 """
333 Embeds a list of documents.
335 :param documents:
336 A list of documents to embed.
338 :returns:
339 A dictionary with the following keys:
340 - `documents`: A list of documents with embeddings.
341 - `meta`: Information about the usage of the model.
342 """
343 if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
344 raise TypeError(
345 "OpenAIDocumentEmbedder expects a list of Documents as input."
346 "In case you want to embed a string, please use the OpenAITextEmbedder."
347 )
349 self.warm_up()
351 texts_to_embed = self._prepare_texts_to_embed(documents=documents)
353 doc_ids_to_embeddings, meta = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
355 new_documents = []
356 for doc in documents:
357 if doc.id in doc_ids_to_embeddings:
358 new_documents.append(replace(doc, embedding=doc_ids_to_embeddings[doc.id]))
359 else:
360 new_documents.append(replace(doc))
362 return {"documents": new_documents, "meta": meta}
364 @component.output_types(documents=list[Document], meta=dict[str, Any])
365 async def run_async(self, documents: list[Document]) -> dict[str, Any]:
366 """
367 Embeds a list of documents asynchronously.
369 :param documents:
370 A list of documents to embed.
372 :returns:
373 A dictionary with the following keys:
374 - `documents`: A list of documents with embeddings.
375 - `meta`: Information about the usage of the model.
376 """
377 if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
378 raise TypeError(
379 "OpenAIDocumentEmbedder expects a list of Documents as input. "
380 "In case you want to embed a string, please use the OpenAITextEmbedder."
381 )
383 await self.warm_up_async()
385 texts_to_embed = self._prepare_texts_to_embed(documents=documents)
387 doc_ids_to_embeddings, meta = await self._embed_batch_async(
388 texts_to_embed=texts_to_embed, batch_size=self.batch_size
389 )
391 new_documents = []
392 for doc in documents:
393 if doc.id in doc_ids_to_embeddings:
394 new_documents.append(replace(doc, embedding=doc_ids_to_embeddings[doc.id]))
395 else:
396 new_documents.append(replace(doc))
398 return {"documents": new_documents, "meta": meta}