diff --git a/haystack/components/_openai_client_mixin.py b/haystack/components/_openai_client_mixin.py new file mode 100644 index 00000000000..37beeb789d0 --- /dev/null +++ b/haystack/components/_openai_client_mixin.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# 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 diff --git a/haystack/components/embedders/azure_document_embedder.py b/haystack/components/embedders/azure_document_embedder.py index 38ebc03838e..15835d4df5f 100644 --- a/haystack/components/embedders/azure_document_embedder.py +++ b/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 diff --git a/haystack/components/embedders/azure_text_embedder.py b/haystack/components/embedders/azure_text_embedder.py index 76d4c2a1b70..ef18647a41d 100644 --- a/haystack/components/embedders/azure_text_embedder.py +++ b/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 diff --git a/haystack/components/embedders/openai_document_embedder.py b/haystack/components/embedders/openai_document_embedder.py index e9911c36c5d..f4be76fcd51 100644 --- a/haystack/components/embedders/openai_document_embedder.py +++ b/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. diff --git a/haystack/components/embedders/openai_text_embedder.py b/haystack/components/embedders/openai_text_embedder.py index 22d0b75a96f..97654f961c4 100644 --- a/haystack/components/embedders/openai_text_embedder.py +++ b/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. diff --git a/haystack/components/generators/chat/azure.py b/haystack/components/generators/chat/azure.py index f5907fcdbe3..ef52c2639d5 100644 --- a/haystack/components/generators/chat/azure.py +++ b/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 diff --git a/haystack/components/generators/chat/openai.py b/haystack/components/generators/chat/openai.py index 4cf1823df2a..b3fb37594e0 100644 --- a/haystack/components/generators/chat/openai.py +++ b/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. diff --git a/haystack/components/generators/chat/openai_responses.py b/haystack/components/generators/chat/openai_responses.py index 45c11b7c7aa..83169d64059 100644 --- a/haystack/components/generators/chat/openai_responses.py +++ b/haystack/components/generators/chat/openai_responses.py @@ -13,6 +13,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 _normalize_messages, _serialize_object from haystack.dataclasses import ( ChatMessage, @@ -38,7 +39,6 @@ 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__) @@ -64,7 +64,7 @@ def _get_response_finish_reason(response: Response | ParsedResponse) -> FinishRe @component -class OpenAIResponsesChatGenerator: +class OpenAIResponsesChatGenerator(OpenAIClientMixin): """ Completes chats using OpenAI's Responses API. @@ -255,50 +255,6 @@ def _warm_up_tools(self) -> None: warm_up_tools(self.tools) # type: ignore[arg-type] 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. diff --git a/haystack/components/generators/openai_image_generator.py b/haystack/components/generators/openai_image_generator.py index 1dc47dc79c0..bd14e7c14f4 100644 --- a/haystack/components/generators/openai_image_generator.py +++ b/haystack/components/generators/openai_image_generator.py @@ -2,21 +2,20 @@ # # SPDX-License-Identifier: Apache-2.0 -import os from typing import Any, Literal from openai import AsyncOpenAI, OpenAI from openai.types.image import Image from haystack import 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 OpenAIImageGenerator: +class OpenAIImageGenerator(OpenAIClientMixin): """ Generates images using OpenAI's image generation models such as `gpt-image-2`. @@ -87,61 +86,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 - @component.output_types(images=list[str], revised_prompt=str) def run( self, diff --git a/releasenotes/notes/Extract-OpenAI-client-lifecycle-into-mixin-dd46dd928ad3a106.yaml b/releasenotes/notes/Extract-OpenAI-client-lifecycle-into-mixin-dd46dd928ad3a106.yaml new file mode 100644 index 00000000000..6b63debac12 --- /dev/null +++ b/releasenotes/notes/Extract-OpenAI-client-lifecycle-into-mixin-dd46dd928ad3a106.yaml @@ -0,0 +1,4 @@ +--- +enhancements: + - | + Extracted duplicated OpenAI client lifecycle management logic (``_client_kwargs``, ``warm_up``, ``warm_up_async``, ``close``, ``close_async``) into a new ``OpenAIClientMixin`` class to unify behavior and reduce boilerplate across OpenAI and Azure components. diff --git a/test/components/embedders/test_openai_document_embedder.py b/test/components/embedders/test_openai_document_embedder.py index 84cb5753e88..44d998edd91 100644 --- a/test/components/embedders/test_openai_document_embedder.py +++ b/test/components/embedders/test_openai_document_embedder.py @@ -9,7 +9,7 @@ import pytest from openai import APIError -import haystack.components.embedders.openai_document_embedder as openai_document_embedder_module +import haystack.components._openai_client_mixin as openai_client_mixin_module from haystack import Document from haystack.components.embedders.openai_document_embedder import OpenAIDocumentEmbedder from haystack.utils.auth import Secret @@ -347,8 +347,8 @@ def mock_openai_clients(monkeypatch): sync_cls = MagicMock(name="OpenAI") async_cls = MagicMock(name="AsyncOpenAI") async_cls.return_value.close = AsyncMock() - monkeypatch.setattr(openai_document_embedder_module, "OpenAI", sync_cls) - monkeypatch.setattr(openai_document_embedder_module, "AsyncOpenAI", async_cls) + monkeypatch.setattr(openai_client_mixin_module, "OpenAI", sync_cls) + monkeypatch.setattr(openai_client_mixin_module, "AsyncOpenAI", async_cls) return sync_cls, async_cls diff --git a/test/components/embedders/test_openai_text_embedder.py b/test/components/embedders/test_openai_text_embedder.py index d9283434b17..4656992df1b 100644 --- a/test/components/embedders/test_openai_text_embedder.py +++ b/test/components/embedders/test_openai_text_embedder.py @@ -10,7 +10,7 @@ from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage -import haystack.components.embedders.openai_text_embedder as openai_text_embedder_module +import haystack.components._openai_client_mixin as openai_client_mixin_module from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder from haystack.utils.auth import Secret @@ -232,8 +232,8 @@ def mock_openai_clients(monkeypatch): sync_cls = MagicMock(name="OpenAI") async_cls = MagicMock(name="AsyncOpenAI") async_cls.return_value.close = AsyncMock() - monkeypatch.setattr(openai_text_embedder_module, "OpenAI", sync_cls) - monkeypatch.setattr(openai_text_embedder_module, "AsyncOpenAI", async_cls) + monkeypatch.setattr(openai_client_mixin_module, "OpenAI", sync_cls) + monkeypatch.setattr(openai_client_mixin_module, "AsyncOpenAI", async_cls) return sync_cls, async_cls diff --git a/test/components/generators/chat/test_openai.py b/test/components/generators/chat/test_openai.py index c33c1a1ed9c..4fabb22ad4c 100644 --- a/test/components/generators/chat/test_openai.py +++ b/test/components/generators/chat/test_openai.py @@ -31,7 +31,7 @@ from openai.types.completion_usage import CompletionTokensDetails, CompletionUsage, PromptTokensDetails from pydantic import BaseModel -import haystack.components.generators.chat.openai as openai_chat_module +import haystack.components._openai_client_mixin as openai_client_mixin_module from haystack import component from haystack.components.generators.chat.openai import ( OpenAIChatGenerator, @@ -1755,8 +1755,8 @@ def mock_openai_clients(monkeypatch): sync_cls = MagicMock(name="OpenAI") async_cls = MagicMock(name="AsyncOpenAI") async_cls.return_value.close = AsyncMock() - monkeypatch.setattr(openai_chat_module, "OpenAI", sync_cls) - monkeypatch.setattr(openai_chat_module, "AsyncOpenAI", async_cls) + monkeypatch.setattr(openai_client_mixin_module, "OpenAI", sync_cls) + monkeypatch.setattr(openai_client_mixin_module, "AsyncOpenAI", async_cls) return sync_cls, async_cls diff --git a/test/components/generators/chat/test_openai_responses.py b/test/components/generators/chat/test_openai_responses.py index 16b8d09cb2e..26a47baa7ee 100644 --- a/test/components/generators/chat/test_openai_responses.py +++ b/test/components/generators/chat/test_openai_responses.py @@ -12,7 +12,7 @@ from openai import AsyncOpenAI, OpenAIError from pydantic import BaseModel -import haystack.components.generators.chat.openai_responses as openai_responses_module +import haystack.components._openai_client_mixin as openai_client_mixin_module from haystack import component from haystack.components.agents import Agent from haystack.components.generators.chat.openai_responses import OpenAIResponsesChatGenerator @@ -363,8 +363,8 @@ def mock_openai_clients(monkeypatch): sync_cls = MagicMock(name="OpenAI") async_cls = MagicMock(name="AsyncOpenAI") async_cls.return_value.close = AsyncMock() - monkeypatch.setattr(openai_responses_module, "OpenAI", sync_cls) - monkeypatch.setattr(openai_responses_module, "AsyncOpenAI", async_cls) + monkeypatch.setattr(openai_client_mixin_module, "OpenAI", sync_cls) + monkeypatch.setattr(openai_client_mixin_module, "AsyncOpenAI", async_cls) return sync_cls, async_cls diff --git a/test/components/generators/test_openai_image_generator.py b/test/components/generators/test_openai_image_generator.py index 2c9998a446f..738e336062e 100644 --- a/test/components/generators/test_openai_image_generator.py +++ b/test/components/generators/test_openai_image_generator.py @@ -11,7 +11,7 @@ from openai.types import ImagesResponse from openai.types.image import Image -import haystack.components.generators.openai_image_generator as openai_image_generator_module +import haystack.components._openai_client_mixin as openai_client_mixin_module from haystack.components.generators.openai_image_generator import OpenAIImageGenerator from haystack.utils import Secret @@ -257,8 +257,8 @@ def mock_openai_clients(monkeypatch): sync_cls = MagicMock(name="OpenAI") async_cls = MagicMock(name="AsyncOpenAI") async_cls.return_value.close = AsyncMock() - monkeypatch.setattr(openai_image_generator_module, "OpenAI", sync_cls) - monkeypatch.setattr(openai_image_generator_module, "AsyncOpenAI", async_cls) + monkeypatch.setattr(openai_client_mixin_module, "OpenAI", sync_cls) + monkeypatch.setattr(openai_client_mixin_module, "AsyncOpenAI", async_cls) return sync_cls, async_cls