chore: import upstream snapshot with attribution

haystack/components/embedders/openai_text_embedder.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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
+    <!-- test-ignore -->
+    ```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)