refactor: extract OpenAI client lifecycle into mixin

haystack/components/_openai_client_mixin.py

@@ -0,0 +1,94 @@
+# 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 haystack.utils.http_client import init_http_client
+
+
+class OpenAIClientMixin:
+    """
+    Mixin providing OpenAI client lifecycle management.
+
+    Supplies ``_client_kwargs``, ``warm_up``, ``warm_up_async``, ``close``
+    and ``close_async`` so that every OpenAI-backed component shares a single
+    implementation of these methods.
+
+    Subclasses must set the following attributes **before** any mixin method
+    is called (typically at the end of ``__init__``):
+
+    * ``api_key`` – a :class:`~haystack.utils.Secret`
+    * ``organization`` – ``str | None``
+    * ``api_base_url`` – ``str | None``
+    * ``timeout`` – ``float | None``
+    * ``max_retries`` – ``int | None``
+    * ``http_client_kwargs`` – ``dict[str, Any] | None``
+    * ``client`` – initialised to ``None``
+    * ``async_client`` – initialised to ``None``
+    """
+
+    # Declared here so that mypy knows the mixin expects these on *self*.
+    api_key: Any
+    organization: str | None
+    api_base_url: str | None
+    timeout: float | None
+    max_retries: int | None
+    http_client_kwargs: dict[str, Any] | None
+    client: OpenAI | None
+    async_client: AsyncOpenAI | None
+
+    def _client_kwargs(self) -> dict[str, Any]:
+        """Build keyword arguments for the OpenAI client constructors."""
+        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 hasattr(self, "_warm_up_tools"):
+            self._warm_up_tools()
+        if self.client is None:
+            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
+            # https://github.com/openai/openai-python/blob/main/httpx2.md
+            http_client = init_http_client(self.http_client_kwargs, async_client=False)
+            self.client = OpenAI(
+                http_client=http_client,  # type: ignore[arg-type]
+                **self._client_kwargs(),
+            )
+
+    async def warm_up_async(self) -> None:  # noqa: RUF029
+        """Initializes the asynchronous OpenAI client on the serving event loop."""
+        if hasattr(self, "_warm_up_tools"):
+            self._warm_up_tools()
+        if self.async_client is None:
+            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
+            # https://github.com/openai/openai-python/blob/main/httpx2.md
+            http_client = init_http_client(self.http_client_kwargs, async_client=True)
+            self.async_client = AsyncOpenAI(
+                http_client=http_client,  # type: ignore[arg-type]
+                **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

haystack/components/embedders/azure_document_embedder.py

@@ -123,7 +123,7 @@ def __init__(  # noqa: PLR0913, PLR0917 (too-many-arguments, too-many-positional
         if api_key is None and azure_ad_token is None:
             raise ValueError("Please provide an API key or an Azure Active Directory token.")
 
-        self.api_key = api_key  # type: ignore[assignment] # mypy does not understand that api_key can be None
+        self.api_key = api_key
         self.azure_ad_token = azure_ad_token
         self.api_version = api_version
         self.azure_endpoint = azure_endpoint

haystack/components/embedders/azure_text_embedder.py

@@ -109,7 +109,7 @@ def __init__(  # noqa: PLR0913
         if api_key is None and azure_ad_token is None:
             raise ValueError("Please provide an API key or an Azure Active Directory token.")
 
-        self.api_key = api_key  # type: ignore[assignment] # mypy does not understand that api_key can be None
+        self.api_key = api_key
         self.azure_ad_token = azure_ad_token
         self.api_version = api_version
         self.azure_endpoint = azure_endpoint

haystack/components/embedders/openai_document_embedder.py

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import os
 from dataclasses import replace
 from typing import Any
 
@@ -12,14 +11,14 @@
 from tqdm.asyncio import tqdm as async_tqdm
 
 from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack.components._openai_client_mixin import OpenAIClientMixin
 from haystack.utils import Secret
-from haystack.utils.http_client import init_http_client
 
 logger = logging.getLogger(__name__)
 
 
 @component
-class OpenAIDocumentEmbedder:
+class OpenAIDocumentEmbedder(OpenAIClientMixin):
     """
     Computes document embeddings using OpenAI models.
 
@@ -125,61 +124,6 @@ def __init__(  # noqa: PLR0913, PLR0917 (too-many-arguments, too-many-positional
         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:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=False)
-            self.client = OpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **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:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=True)
-            self.async_client = AsyncOpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **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.

haystack/components/embedders/openai_text_embedder.py

@@ -2,19 +2,18 @@
 #
 # 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.components._openai_client_mixin import OpenAIClientMixin
 from haystack.utils import Secret
-from haystack.utils.http_client import init_http_client
 
 
 @component
-class OpenAITextEmbedder:
+class OpenAITextEmbedder(OpenAIClientMixin):
     """
     Embeds strings using OpenAI models.
 
@@ -100,61 +99,6 @@ def __init__(
         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:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=False)
-            self.client = OpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **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:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=True)
-            self.async_client = AsyncOpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **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.

haystack/components/generators/chat/azure.py

@@ -214,7 +214,7 @@ def __init__(
 
         # The check above makes mypy incorrectly infer that api_key is never None,
         # which propagates the incorrect type.
-        self.api_key = api_key  # type: ignore
+        self.api_key = api_key
         self.azure_ad_token = azure_ad_token
         self.generation_kwargs = generation_kwargs or {}
         self.streaming_callback = streaming_callback

haystack/components/generators/chat/openai.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import json
-import os
 from datetime import datetime
 from typing import Any, ClassVar
 
@@ -23,6 +22,7 @@
 from pydantic import BaseModel
 
 from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.components._openai_client_mixin import OpenAIClientMixin
 from haystack.components.generators.utils import (
     _convert_streaming_chunks_to_chat_message,
     _normalize_messages,
@@ -49,13 +49,12 @@
     warm_up_tools,
 )
 from haystack.utils import Secret, deserialize_callable, serialize_callable
-from haystack.utils.http_client import init_http_client
 
 logger = logging.getLogger(__name__)
 
 
 @component
-class OpenAIChatGenerator:
+class OpenAIChatGenerator(OpenAIClientMixin):
     """
     Completes chats using OpenAI's large language models (LLMs).
 
@@ -213,68 +212,11 @@ def __init__(
         self.async_client: AsyncOpenAI | None = None
         self._tools_warmed_up = False
 
-    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_tools(self) -> None:
         if not self._tools_warmed_up:
             warm_up_tools(self.tools)
             self._tools_warmed_up = True
 
-    def warm_up(self) -> None:
-        """
-        Warm up the tools and initialize the synchronous OpenAI client.
-        """
-        self._warm_up_tools()
-        if self.client is None:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=False)
-            self.client = OpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **self._client_kwargs(),
-            )
-
-    async def warm_up_async(self) -> None:  # noqa: RUF029
-        """
-        Warm up the tools and initialize the asynchronous OpenAI client on the serving event loop.
-        """
-        self._warm_up_tools()
-        if self.async_client is None:
-            # openai>=3 annotates http_client as httpx2, but legacy httpx clients are supported at runtime.
-            # https://github.com/openai/openai-python/blob/main/httpx2.md
-            http_client = init_http_client(self.http_client_kwargs, async_client=True)
-            self.async_client = AsyncOpenAI(
-                http_client=http_client,  # type: ignore[arg-type]
-                **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.