zenml-nightly 0.68.1.dev20241103__tar.gz → 0.68.1.dev20241106__tar.gz

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.68.1.dev20241103
+Version: 0.68.1.dev20241106
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.68.1.dev20241103"
+version = "0.68.1.dev20241106"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]
@@ -475,5 +475,6 @@ module = [
     "prodigy.components.*",
     "prodigy.components.db.*",
     "transformers.*",
+    "langchain_community.*",
 ]
 ignore_missing_imports = true

zenml_nightly-0.68.1.dev20241106/src/zenml/VERSION

@@ -0,0 +1 @@
+0.68.1.dev20241106

zenml_nightly-0.68.1.dev20241103/src/zenml/artifacts/load_directory_materializer.py → zenml_nightly-0.68.1.dev20241106/src/zenml/artifacts/preexisting_data_materializer.py

@@ -14,7 +14,6 @@
 """Only-load materializer for directories."""
 
 import os
-import tempfile
 from pathlib import Path
 from typing import Any, ClassVar, Tuple, Type
 
@@ -46,14 +45,14 @@ class PreexistingDataMaterializer(BaseMaterializer):
         Returns:
             Path to the local directory that contains the artifact files.
         """
-        directory = tempfile.mkdtemp(prefix="zenml-artifact")
-        if fileio.isdir(self.uri):
-            self._copy_directory(src=self.uri, dst=directory)
-            return Path(directory)
-        else:
-            dst = os.path.join(directory, os.path.split(self.uri)[-1])
-            fileio.copy(src=self.uri, dst=dst)
-            return Path(dst)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            if fileio.isdir(self.uri):
+                self._copy_directory(src=self.uri, dst=temp_dir)
+                return Path(temp_dir)
+            else:
+                dst = os.path.join(temp_dir, os.path.split(self.uri)[-1])
+                fileio.copy(src=self.uri, dst=dst)
+                return Path(dst)
 
     def save(self, data: Any) -> None:
         """Store the directory in the artifact store.

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/artifacts/utils.py

@@ -31,7 +31,7 @@ from typing import (
 from uuid import UUID, uuid4
 
 from zenml.artifacts.artifact_config import ArtifactConfig
-from zenml.artifacts.load_directory_materializer import (
+from zenml.artifacts.preexisting_data_materializer import (
     PreexistingDataMaterializer,
 )
 from zenml.client import Client

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/__init__.py

@@ -45,6 +45,7 @@ from zenml.integrations.kubernetes import KubernetesIntegration  # noqa
 from zenml.integrations.label_studio import LabelStudioIntegration  # noqa
 from zenml.integrations.langchain import LangchainIntegration  # noqa
 from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
+
 # from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
@@ -52,7 +53,7 @@ from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
 from zenml.integrations.numpy import NumpyIntegration  # noqa
 from zenml.integrations.openai import OpenAIIntegration  # noqa
 from zenml.integrations.pandas import PandasIntegration  # noqa
-from zenml.integrations.pigeon import PigeonIntegration # noqa
+from zenml.integrations.pigeon import PigeonIntegration  # noqa
 from zenml.integrations.pillow import PillowIntegration  # noqa
 from zenml.integrations.polars import PolarsIntegration  # noqa
 from zenml.integrations.prodigy import ProdigyIntegration  # noqa
@@ -78,3 +79,4 @@ from zenml.integrations.tensorflow import TensorflowIntegration  # noqa
 from zenml.integrations.wandb import WandbIntegration  # noqa
 from zenml.integrations.whylogs import WhylogsIntegration  # noqa
 from zenml.integrations.xgboost import XgboostIntegration  # noqa
+from zenml.integrations.vllm import VLLMIntegration  # noqa

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/bentoml/materializers/bentoml_bento_materializer.py

@@ -14,7 +14,6 @@
 """Materializer for BentoML Bento objects."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import bentoml
@@ -23,7 +22,6 @@ from bentoml.exceptions import BentoMLException
 
 from zenml.enums import ArtifactType
 from zenml.integrations.bentoml.constants import DEFAULT_BENTO_FILENAME
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
@@ -49,23 +47,21 @@ class BentoMaterializer(BaseMaterializer):
         Returns:
             An bento.Bento object.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the Bento from the temporary directory
-        imported_bento = Bento.import_from(
-            os.path.join(temp_dir.name, DEFAULT_BENTO_FILENAME)
-        )
-
-        # Try save the Bento to the local BentoML store
-        try:
-            _ = bentoml.get(imported_bento.tag)
-        except BentoMLException:
-            imported_bento.save()
-        return imported_bento
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            # Copy from artifact store to temporary directory
+            io_utils.copy_dir(self.uri, temp_dir)
+
+            # Load the Bento from the temporary directory
+            imported_bento = Bento.import_from(
+                os.path.join(temp_dir, DEFAULT_BENTO_FILENAME)
+            )
+
+            # Try save the Bento to the local BentoML store
+            try:
+                _ = bentoml.get(imported_bento.tag)
+            except BentoMLException:
+                imported_bento.save()
+            return imported_bento
 
     def save(self, bento: bento.Bento) -> None:
         """Write to artifact store.
@@ -73,18 +69,10 @@ class BentoMaterializer(BaseMaterializer):
         Args:
             bento: An bento.Bento object.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
-        temp_bento_path = os.path.join(temp_dir.name, DEFAULT_BENTO_FILENAME)
-
-        # save the image in a temporary directory
-        bentoml.export_bento(bento.tag, temp_bento_path)
-
-        # copy the saved image to the artifact store
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_bento_path = os.path.join(temp_dir, DEFAULT_BENTO_FILENAME)
+            bentoml.export_bento(bento.tag, temp_bento_path)
+            io_utils.copy_dir(temp_dir, self.uri)
 
     def extract_metadata(
         self, bento: bento.Bento

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/constants.py

@@ -76,4 +76,5 @@ WANDB = "wandb"
 VERTEX = "vertex"
 XGBOOST = "xgboost"
 VAULT = "vault"
+VLLM = "vllm"
 LIGHTNING = "lightning"

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py

@@ -15,7 +15,6 @@
 
 import os
 from collections import defaultdict
-from tempfile import TemporaryDirectory, mkdtemp
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -88,12 +87,12 @@ class HFDatasetMaterializer(BaseMaterializer):
         Returns:
             The dataset read from the specified dir.
         """
-        temp_dir = mkdtemp()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_DATASET_DIR),
-            temp_dir,
-        )
-        return load_from_disk(temp_dir)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_DATASET_DIR),
+                temp_dir,
+            )
+            return load_from_disk(temp_dir)
 
     def save(self, ds: Union[Dataset, DatasetDict]) -> None:
         """Writes a Dataset to the specified dir.
@@ -101,16 +100,13 @@ class HFDatasetMaterializer(BaseMaterializer):
         Args:
             ds: The Dataset to write.
         """
-        temp_dir = TemporaryDirectory()
-        path = os.path.join(temp_dir.name, DEFAULT_DATASET_DIR)
-        try:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            path = os.path.join(temp_dir, DEFAULT_DATASET_DIR)
             ds.save_to_disk(path)
             io_utils.copy_dir(
                 path,
                 os.path.join(self.uri, DEFAULT_DATASET_DIR),
             )
-        finally:
-            fileio.rmtree(temp_dir.name)
 
     def extract_metadata(
         self, ds: Union[Dataset, DatasetDict]

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/huggingface/materializers/huggingface_pt_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFPTModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: PreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFPTModelMaterializer(BaseMaterializer):
         Args:
             model: The Torch Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_PT_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: PreTrainedModel

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/huggingface/materializers/huggingface_t5_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface t5 materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Type, Union
 
 from transformers import (
@@ -52,8 +51,7 @@ class HFT5Materializer(BaseMaterializer):
             ValueError: Unsupported data type used
         """
         filepath = self.uri
-
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Copy files from artifact store to temporary directory
             for file in fileio.listdir(filepath):
                 src = os.path.join(filepath, file)
@@ -86,8 +84,7 @@ class HFT5Materializer(BaseMaterializer):
         Args:
             obj: A T5ForConditionalGeneration model or T5Tokenizer.
         """
-        # Create a temporary directory
-        with tempfile.TemporaryDirectory(prefix="zenml-temp-") as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             # Save the model or tokenizer
             obj.save_pretrained(temp_dir)
 

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/huggingface/materializers/huggingface_tf_model_materializer.py

@@ -15,7 +15,6 @@
 
 import importlib
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Dict, Tuple, Type
 
 from transformers import (
@@ -46,17 +45,17 @@ class HFTFModelMaterializer(BaseMaterializer):
         Returns:
             The model read from the specified dir.
         """
-        temp_dir = TemporaryDirectory()
-        io_utils.copy_dir(
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir.name
-        )
-
-        config = AutoConfig.from_pretrained(temp_dir.name)
-        architecture = "TF" + config.architectures[0]
-        model_cls = getattr(
-            importlib.import_module("transformers"), architecture
-        )
-        return model_cls.from_pretrained(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            io_utils.copy_dir(
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR), temp_dir
+            )
+
+            config = AutoConfig.from_pretrained(temp_dir)
+            architecture = "TF" + config.architectures[0]
+            model_cls = getattr(
+                importlib.import_module("transformers"), architecture
+            )
+            return model_cls.from_pretrained(temp_dir)
 
     def save(self, model: TFPreTrainedModel) -> None:
         """Writes a Model to the specified dir.
@@ -64,12 +63,12 @@ class HFTFModelMaterializer(BaseMaterializer):
         Args:
             model: The TF Model to write.
         """
-        temp_dir = TemporaryDirectory()
-        model.save_pretrained(temp_dir.name)
-        io_utils.copy_dir(
-            temp_dir.name,
-            os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
-        )
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            model.save_pretrained(temp_dir)
+            io_utils.copy_dir(
+                temp_dir,
+                os.path.join(self.uri, DEFAULT_TF_MODEL_DIR),
+            )
 
     def extract_metadata(
         self, model: TFPreTrainedModel

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/huggingface/materializers/huggingface_tokenizer_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the Huggingface tokenizer materializer."""
 
 import os
-from tempfile import TemporaryDirectory
 from typing import Any, ClassVar, Tuple, Type
 
 from transformers import AutoTokenizer
@@ -46,7 +45,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Returns:
             The tokenizer read from the specified dir.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             io_utils.copy_dir(
                 os.path.join(self.uri, DEFAULT_TOKENIZER_DIR), temp_dir
             )
@@ -58,7 +57,7 @@ class HFTokenizerMaterializer(BaseMaterializer):
         Args:
             tokenizer: The HFTokenizer to write.
         """
-        with TemporaryDirectory() as temp_dir:
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
             tokenizer.save_pretrained(temp_dir)
             io_utils.copy_dir(
                 temp_dir,

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/langchain/__init__.py

@@ -26,7 +26,8 @@ class LangchainIntegration(Integration):
 
     NAME = LANGCHAIN
     REQUIREMENTS = [
-        "langchain==0.0.325",
+        "langchain==0.3.0",
+        "langchain-community",
         "pyyaml>=6.0.1",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/langchain/materializers/openai_embedding_materializer.py

@@ -24,11 +24,37 @@ from zenml.materializers.cloudpickle_materializer import (
 if TYPE_CHECKING and sys.version_info < (3, 8):
     OpenAIEmbeddings = Any
 else:
-    from langchain.embeddings import OpenAIEmbeddings
+    from langchain_community.embeddings import (
+        OpenAIEmbeddings,
+    )
 
 
 class LangchainOpenaiEmbeddingMaterializer(CloudpickleMaterializer):
-    """Handle langchain OpenAI embedding objects."""
+    """Materializer for Langchain OpenAI Embeddings."""
 
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL
     ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (OpenAIEmbeddings,)
+
+    def save(self, embeddings: Any) -> None:
+        """Saves the embeddings model after clearing non-picklable clients.
+
+        Args:
+            embeddings: The embeddings model to save.
+        """
+        # Clear the clients which will be recreated on load
+        embeddings.client = None
+        embeddings.async_client = None
+
+        # Use the parent class's save implementation which uses cloudpickle
+        super().save(embeddings)
+
+    def load(self, data_type: Type[Any]) -> Any:
+        """Loads the embeddings model and lets it recreate clients when needed.
+
+        Args:
+            data_type: The type of the data to load.
+
+        Returns:
+            The loaded embeddings model.
+        """
+        return super().load(data_type)

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/lightgbm/materializers/lightgbm_booster_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM booster materializer."""
 
 import os
-import tempfile
 from typing import Any, ClassVar, Tuple, Type
 
 import lightgbm as lgb
@@ -42,18 +41,13 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             A lightgbm Booster object.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        booster = lgb.Booster(model_file=temp_file)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir)
-        return booster
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            booster = lgb.Booster(model_file=temp_file)
+            return booster
 
     def save(self, booster: lgb.Booster) -> None:
         """Creates a JSON serialization for a lightgbm Booster model.
@@ -62,8 +56,7 @@ class LightGBMBoosterMaterializer(BaseMaterializer):
             booster: A lightgbm Booster model.
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_path = os.path.join(tmp_dir, "model.txt")
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            tmp_path = os.path.join(temp_dir, "model.txt")
             booster.save_model(tmp_path)
             fileio.copy(tmp_path, filepath)

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/lightgbm/materializers/lightgbm_dataset_materializer.py

@@ -14,7 +14,6 @@
 """Implementation of the LightGBM materializer."""
 
 import os
-import tempfile
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import lightgbm as lgb
@@ -46,16 +45,14 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Create a temporary folder
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+        with self.get_temporary_directory(delete_at_exit=False) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
 
-        # Copy from artifact store to temporary file
-        fileio.copy(filepath, temp_file)
-        matrix = lgb.Dataset(temp_file, free_raw_data=False)
+            # Copy from artifact store to temporary file
+            fileio.copy(filepath, temp_file)
+            matrix = lgb.Dataset(temp_file, free_raw_data=False)
 
-        # No clean up this time because matrix is lazy loaded
-        return matrix
+            return matrix
 
     def save(self, matrix: lgb.Dataset) -> None:
         """Creates a binary serialization for a lightgbm.Dataset object.
@@ -65,14 +62,12 @@ class LightGBMDatasetMaterializer(BaseMaterializer):
         """
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
 
-        # Make a temporary phantom artifact
-        temp_dir = tempfile.mkdtemp(prefix="zenml-temp-")
-        temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
-        matrix.save_binary(temp_file)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            temp_file = os.path.join(str(temp_dir), DEFAULT_FILENAME)
+            matrix.save_binary(temp_file)
 
-        # Copy it into artifact store
-        fileio.copy(temp_file, filepath)
-        fileio.rmtree(temp_dir)
+            # Copy it into artifact store
+            fileio.copy(temp_file, filepath)
 
     def extract_metadata(
         self, matrix: lgb.Dataset

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/openai/__init__.py

@@ -21,7 +21,7 @@ class OpenAIIntegration(Integration):
     """Definition of OpenAI integration for ZenML."""
 
     NAME = OPEN_AI
-    REQUIREMENTS = ["openai>=0.27.0,<1.0.0"]
+    REQUIREMENTS = ["openai>=1.0.0"]
 
 
 OpenAIIntegration.check_installation()

{zenml_nightly-0.68.1.dev20241103 → zenml_nightly-0.68.1.dev20241106}/src/zenml/integrations/openai/hooks/open_ai_failure_hook.py

@@ -15,8 +15,9 @@
 
 import io
 import sys
+from typing import Optional
 
-import openai
+from openai import OpenAI
 from rich.console import Console
 
 from zenml import get_step_context
@@ -38,6 +39,8 @@ def openai_alerter_failure_hook_helper(
     Args:
         exception: The exception that was raised.
         model_name: The OpenAI model to use for the chatbot.
+
+    This implementation uses the OpenAI v1 SDK with automatic retries and backoff.
     """
     client = Client()
     context = get_step_context()
@@ -47,12 +50,15 @@ def openai_alerter_failure_hook_helper(
         openai_secret = client.get_secret(
             "openai", allow_partial_name_match=False
         )
-        openai_api_key = openai_secret.secret_values.get("api_key")
+        openai_api_key: Optional[str] = openai_secret.secret_values.get(
+            "api_key"
+        )
     except (KeyError, NotImplementedError):
         openai_api_key = None
 
     alerter = client.active_stack.alerter
     if alerter and openai_api_key:
+        # Capture rich traceback
         output_captured = io.StringIO()
         original_stdout = sys.stdout
         sys.stdout = output_captured
@@ -62,25 +68,44 @@ def openai_alerter_failure_hook_helper(
         sys.stdout = original_stdout
         rich_traceback = output_captured.getvalue()
 
-        response = openai.ChatCompletion.create(  # type: ignore
+        # Initialize OpenAI client with timeout and retry settings
+        openai_client = OpenAI(
+            api_key=openai_api_key,
+            max_retries=3,  # Will retry 3 times with exponential backoff
+            timeout=60.0,  # 60 second timeout
+        )
+
+        # Create chat completion using the new client pattern
+        response = openai_client.chat.completions.create(
             model=model_name,
             messages=[
                 {
                     "role": "user",
-                    "content": f"This is an error message (following an exception of type '{type(exception)}') I encountered while executing a ZenML step. Please suggest ways I might fix the problem. Feel free to give code snippets as examples, and note that your response will be piped to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. Thank you!",
+                    "content": f"This is an error message (following an exception of type '{type(exception)}') "
+                    f"I encountered while executing a ZenML step. Please suggest ways I might fix the problem. "
+                    f"Feel free to give code snippets as examples, and note that your response will be piped "
+                    f"to a Slack bot so make sure the formatting is appropriate: {exception} -- {rich_traceback}. "
+                    f"Thank you!",
                 }
             ],
         )
-        suggestion = response["choices"][0]["message"]["content"]
-        message = "*Failure Hook Notification! Step failed!*" + "\n\n"
-        message += f"Run name: `{context.pipeline_run.name}`" + "\n"
-        message += f"Step name: `{context.step_run.name}`" + "\n"
-        message += f"Parameters: `{context.step_run.config.parameters}`" + "\n"
-        message += f"Exception: `({type(exception)}) {exception}`" + "\n\n"
-        message += (
-            f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`"
-            + "\n"
+
+        suggestion = response.choices[0].message.content
+
+        # Format the alert message
+        message = "\n".join(
+            [
+                "*Failure Hook Notification! Step failed!*",
+                "",
+                f"Run name: `{context.pipeline_run.name}`",
+                f"Step name: `{context.step_run.name}`",
+                f"Parameters: `{context.step_run.config.parameters}`",
+                f"Exception: `({type(exception)}) {exception}`",
+                "",
+                f"*OpenAI ChatGPT's suggestion (model = `{model_name}`) on how to fix it:*\n `{suggestion}`",
+            ]
         )
+
         alerter.post(message)
     elif not openai_api_key:
         logger.warning(
@@ -111,4 +136,4 @@ def openai_gpt4_alerter_failure_hook(
     Args:
         exception: The exception that was raised.
     """
-    openai_alerter_failure_hook_helper(exception, "gpt-4")
+    openai_alerter_failure_hook_helper(exception, "gpt-4o")