# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import os from typing import Any from openai import AsyncOpenAI, OpenAI from openai.types import CreateEmbeddingResponse from haystack import component, default_from_dict, default_to_dict from haystack.utils import Secret from haystack.utils.http_client import init_http_client @component class OpenAITextEmbedder: """ Embeds strings using OpenAI models. You can use it to embed user query and send it to an embedding Retriever. ### Usage example ```python from haystack.components.embedders import OpenAITextEmbedder text_to_embed = "I love pizza!" text_embedder = OpenAITextEmbedder() print(text_embedder.run(text_to_embed)) # {'embedding': [0.017020374536514282, -0.023255806416273117, ...], # 'meta': {'model': 'text-embedding-ada-002-v2', # 'usage': {'prompt_tokens': 4, 'total_tokens': 4}}} ``` """ def __init__( self, api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), model: str = "text-embedding-ada-002", dimensions: int | None = None, api_base_url: str | None = None, organization: str | None = None, prefix: str = "", suffix: str = "", timeout: float | None = None, max_retries: int | None = None, http_client_kwargs: dict[str, Any] | None = None, ) -> None: """ Creates an OpenAITextEmbedder component. Before initializing the component, you can set the 'OPENAI_TIMEOUT' and 'OPENAI_MAX_RETRIES' environment variables to override the `timeout` and `max_retries` parameters respectively in the OpenAI client. :param api_key: The OpenAI API key. You can set it with an environment variable `OPENAI_API_KEY`, or pass with this parameter during initialization. :param model: The name of the model to use for calculating embeddings. The default model is `text-embedding-ada-002`. :param dimensions: The number of dimensions of the resulting embeddings. Only `text-embedding-3` and later models support this parameter. :param api_base_url: Overrides default base URL for all HTTP requests. :param organization: Your organization ID. See OpenAI's [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization) for more information. :param prefix: A string to add at the beginning of each text to embed. :param suffix: A string to add at the end of each text to embed. :param timeout: Timeout for OpenAI client calls. If not set, it defaults to either the `OPENAI_TIMEOUT` environment variable, or 30 seconds. :param max_retries: Maximum number of retries to contact OpenAI after an internal error. If not set, it defaults to either the `OPENAI_MAX_RETRIES` environment variable, or set to 5. :param http_client_kwargs: A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`. For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client). """ self.model = model self.dimensions = dimensions self.api_base_url = api_base_url self.organization = organization self.prefix = prefix self.suffix = suffix self.api_key = api_key self.timeout = timeout self.max_retries = max_retries self.http_client_kwargs = http_client_kwargs self.client: OpenAI | None = None self.async_client: AsyncOpenAI | None = None def _client_kwargs(self) -> dict[str, Any]: timeout = self.timeout if self.timeout is not None else float(os.environ.get("OPENAI_TIMEOUT", "30.0")) max_retries = ( self.max_retries if self.max_retries is not None else int(os.environ.get("OPENAI_MAX_RETRIES", "5")) ) return { "api_key": self.api_key.resolve_value(), "organization": self.organization, "base_url": self.api_base_url, "timeout": timeout, "max_retries": max_retries, } def warm_up(self) -> None: """ Initializes the synchronous OpenAI client. """ if self.client is None: self.client = OpenAI( http_client=init_http_client(self.http_client_kwargs, async_client=False), **self._client_kwargs() ) async def warm_up_async(self) -> None: # noqa: RUF029 """ Initializes the asynchronous OpenAI client on the serving event loop. """ if self.async_client is None: self.async_client = AsyncOpenAI( http_client=init_http_client(self.http_client_kwargs, async_client=True), **self._client_kwargs() ) def close(self) -> None: """ Releases the synchronous OpenAI client. """ if self.client is not None: self.client.close() self.client = None async def close_async(self) -> None: """ Releases the asynchronous OpenAI client. """ if self.async_client is not None: await self.async_client.close() self.async_client = None def _get_telemetry_data(self) -> dict[str, Any]: """ Data that is sent to Posthog for usage analytics. """ return {"model": self.model} def to_dict(self) -> dict[str, Any]: """ Serializes the component to a dictionary. :returns: Dictionary with serialized data. """ return default_to_dict( self, api_key=self.api_key, model=self.model, dimensions=self.dimensions, api_base_url=self.api_base_url, organization=self.organization, prefix=self.prefix, suffix=self.suffix, timeout=self.timeout, max_retries=self.max_retries, http_client_kwargs=self.http_client_kwargs, ) @classmethod def from_dict(cls, data: dict[str, Any]) -> "OpenAITextEmbedder": """ Deserializes the component from a dictionary. :param data: Dictionary to deserialize from. :returns: Deserialized component. """ return default_from_dict(cls, data) def _prepare_input(self, text: str) -> dict[str, Any]: if not isinstance(text, str): raise TypeError( "OpenAITextEmbedder expects a string as an input." "In case you want to embed a list of Documents, please use the OpenAIDocumentEmbedder." ) text_to_embed = self.prefix + text + self.suffix kwargs: dict[str, Any] = {"model": self.model, "input": text_to_embed, "encoding_format": "float"} if self.dimensions is not None: kwargs["dimensions"] = self.dimensions return kwargs def _prepare_output(self, result: CreateEmbeddingResponse) -> dict[str, Any]: return {"embedding": result.data[0].embedding, "meta": {"model": result.model, "usage": dict(result.usage)}} @component.output_types(embedding=list[float], meta=dict[str, Any]) def run(self, text: str) -> dict[str, Any]: """ Embeds a single string. :param text: Text to embed. :returns: A dictionary with the following keys: - `embedding`: The embedding of the input text. - `meta`: Information about the usage of the model. """ self.warm_up() create_kwargs = self._prepare_input(text=text) assert self.client is not None # mypy: client is built by warm_up above response = self.client.embeddings.create(**create_kwargs) return self._prepare_output(result=response) @component.output_types(embedding=list[float], meta=dict[str, Any]) async def run_async(self, text: str) -> dict[str, Any]: """ Asynchronously embed a single string. This is the asynchronous version of the `run` method. It has the same parameters and return values but can be used with `await` in async code. :param text: Text to embed. :returns: A dictionary with the following keys: - `embedding`: The embedding of the input text. - `meta`: Information about the usage of the model. """ await self.warm_up_async() create_kwargs = self._prepare_input(text=text) assert self.async_client is not None # mypy: async_client is built by warm_up_async above response = await self.async_client.embeddings.create(**create_kwargs) return self._prepare_output(result=response)