zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/models/v2/base/base.py

@@ -17,8 +17,7 @@ from datetime import datetime
 from typing import Any, Dict, Generic, Optional, TypeVar
 from uuid import UUID
 
-from pydantic import Field, SecretStr
-from pydantic.generics import GenericModel
+from pydantic import ConfigDict, Field
 
 from zenml.analytics.models import AnalyticsTrackedModelMixin
 from zenml.enums import ResponseUpdateStrategy
@@ -35,26 +34,16 @@ class BaseZenModel(YAMLSerializationMixin, AnalyticsTrackedModelMixin):
     """Base model class for all ZenML models.
 
     This class is used as a base class for all ZenML models. It provides
-    functionality for tracking analytics events and proper encoding of
-    SecretStr values.
+    functionality for tracking analytics events.
     """
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # This is needed to allow the REST client and server to unpack SecretStr
-        # values correctly.
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value()
-            if v is not None
-            else None
-        }
-
-        # Allow extras on all models to support forwards and backwards
+    model_config = ConfigDict(
+        # Ignore extras on all models to support forwards and backwards
         # compatibility (e.g. new fields in newer versions of ZenML servers
-        # are allowed to be present in older versions of ZenML clients and
-        # vice versa).
-        extra = "allow"
+        # are allowed to be passed to older versions of ZenML clients and
+        # vice versa but will be ignored).
+        extra="ignore",
+    )
 
 
 # -------------------- Request Model --------------------
@@ -67,26 +56,28 @@ class BaseRequest(BaseZenModel):
     """
 
 
-# -------------------- Response Model --------------------
+# -------------------- Update Model --------------------
 
 
-class BaseResponseBody(BaseZenModel):
-    """Base body model.
+class BaseUpdate(BaseZenModel):
+    """Base update model.
 
-    Used as a base class for all body models associated with responses.
-    Features a creation and update timestamp.
+    Used as a base class for all update models.
     """
 
-    created: Optional[datetime] = Field(
-        title="The timestamp when this resource was created.",
-        default=None,
-    )
-    updated: Optional[datetime] = Field(
-        title="The timestamp when this resource was last updated.",
-        default=None,
+    model_config = ConfigDict(
+        # Ignore extras on all update models.
+        extra="ignore",
     )
 
 
+# -------------------- Response Model --------------------
+
+
+class BaseResponseBody(BaseZenModel):
+    """Base body model."""
+
+
 class BaseResponseMetadata(BaseZenModel):
     """Base metadata model.
 
@@ -94,63 +85,42 @@ class BaseResponseMetadata(BaseZenModel):
     """
 
 
+class BaseResponseResources(BaseZenModel):
+    """Base resources model.
+
+    Used as a base class for all resource models associated with responses.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+
 AnyBody = TypeVar("AnyBody", bound=BaseResponseBody)
 AnyMetadata = TypeVar("AnyMetadata", bound=BaseResponseMetadata)
+AnyResources = TypeVar("AnyResources", bound=BaseResponseResources)
 
 
-class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
-    """Base domain model."""
-
-    id: UUID = Field(title="The unique resource id.")
-    permission_denied: bool = False
+class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
+    """Base domain model for all responses."""
 
     # Body and metadata pair
-    body: Optional["AnyBody"] = Field(title="The body of the resource.")
+    body: Optional["AnyBody"] = Field(
+        default=None, title="The body of the resource."
+    )
     metadata: Optional["AnyMetadata"] = Field(
-        title="The metadata related to this resource."
+        default=None, title="The metadata related to this resource."
+    )
+    resources: Optional["AnyResources"] = Field(
+        default=None, title="The resources related to this resource."
     )
 
-    _response_update_strategy: (
-        ResponseUpdateStrategy
-    ) = ResponseUpdateStrategy.ALLOW
+    _response_update_strategy: ResponseUpdateStrategy = (
+        ResponseUpdateStrategy.ALLOW
+    )
     _warn_on_response_updates: bool = True
 
-    def get_hydrated_version(self) -> "BaseResponse[AnyBody, AnyMetadata]":
-        """Abstract method to fetch the hydrated version of the model.
-
-        Raises:
-            NotImplementedError: in case the method is not implemented.
-        """
-        raise NotImplementedError(
-            "Please implement a `get_hydrated_version` method before "
-            "using/hydrating the model."
-        )
-
-    # Helper functions
-    def __hash__(self) -> int:
-        """Implementation of hash magic method.
-
-        Returns:
-            Hash of the UUID.
-        """
-        return hash((type(self),) + tuple([self.id]))
-
-    def __eq__(self, other: Any) -> bool:
-        """Implementation of equality magic method.
-
-        Args:
-            other: The other object to compare to.
-
-        Returns:
-            True if the other object is of the same type and has the same UUID.
-        """
-        if isinstance(other, type(self)):
-            return self.id == other.id
-        else:
-            return False
-
     def _validate_hydrated_version(
-        self, hydrated_model: "BaseResponse[AnyBody, AnyMetadata]"
+        self,
+        hydrated_model: "BaseResponse[AnyBody, AnyMetadata, AnyResources]",
     ) -> None:
         """Helper method to validate the values within the hydrated version.
 
@@ -168,14 +138,8 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
                 "The hydrated model does not have a metadata field."
             )
 
-        # Check if the ID is the same
-        if self.id != hydrated_model.id:
-            raise HydrationError(
-                "The hydrated version of the model does not have the same id."
-            )
-
         # Check if the name has changed
-        if "name" in self.__fields__:
+        if "name" in self.model_fields:
             original_name = getattr(self, "name")
             hydrated_name = getattr(hydrated_model, "name")
 
@@ -213,7 +177,7 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
                     )
 
         # Check all the fields in the body
-        for field in self.get_body().__fields__:
+        for field in self.get_body().model_fields:
             original_value = getattr(self.get_body(), field)
             hydrated_value = getattr(hydrated_model.get_body(), field)
 
@@ -251,27 +215,39 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
                         f"`{hydrated_value}`"
                     )
 
-    def get_body(self) -> AnyBody:
+    def hydrate(self) -> None:
+        """Hydrate the response."""
+        hydrated_version = self.get_hydrated_version()
+        self._validate_hydrated_version(hydrated_version)
+
+        self.resources = hydrated_version.resources
+        self.metadata = hydrated_version.metadata
+
+    def get_hydrated_version(
+        self,
+    ) -> "BaseResponse[AnyBody, AnyMetadata, AnyResources]":
+        """Abstract method to fetch the hydrated version of the model.
+
+        Raises:
+            NotImplementedError: in case the method is not implemented.
+        """
+        raise NotImplementedError(
+            "Please implement a `get_hydrated_version` method before "
+            "using/hydrating the model."
+        )
+
+    def get_body(self) -> "AnyBody":
         """Fetch the body of the entity.
 
         Returns:
             The body field of the response.
 
         Raises:
-            IllegalOperationError: If the user lacks permission to access the
-                entity represented by this response.
             RuntimeError: If the body was not included in the response.
         """
-        if self.permission_denied:
-            raise IllegalOperationError(
-                f"Missing permissions to access {type(self).__name__} with "
-                f"ID {self.id}."
-            )
-
         if not self.body:
             raise RuntimeError(
-                f"Missing response body for {type(self).__name__} with ID "
-                f"{self.id}."
+                f"Missing response body for {type(self).__name__}."
             )
 
         return self.body
@@ -281,27 +257,27 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
 
         Returns:
             The metadata field of the response.
-
-        Raises:
-            IllegalOperationError: If the user lacks permission to access this
-                entity represented by this response.
         """
-        if self.permission_denied:
-            raise IllegalOperationError(
-                f"Missing permissions to access {type(self).__name__} with "
-                f"ID {self.id}."
-            )
-
         if self.metadata is None:
             # If the metadata is not there, check the class first.
-            metadata_type = self.__fields__["metadata"].type_
+            metadata_annotation = self.model_fields["metadata"].annotation
+            assert metadata_annotation is not None, (
+                "For each response model, an annotated metadata"
+                "field should exist."
+            )
 
-            if len(metadata_type.__fields__):
+            # metadata is defined as:
+            #   metadata: Optional[....ResponseMetadata] = Field(default=None)
+            # We need to find the actual class inside the Optional annotation.
+            from zenml.utils.typing_utils import get_args
+
+            metadata_type = get_args(metadata_annotation)[0]
+            assert issubclass(metadata_type, BaseResponseMetadata)
+
+            if len(metadata_type.model_fields):
                 # If the metadata class defines any fields, fetch the metadata
                 # through the hydrated version.
-                hydrated_version = self.get_hydrated_version()
-                self._validate_hydrated_version(hydrated_version)
-                self.metadata = hydrated_version.metadata
+                self.hydrate()
             else:
                 # Otherwise, use the metadata class to create an empty metadata
                 # object.
@@ -311,6 +287,172 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
 
         return self.metadata
 
+    def get_resources(self) -> "AnyResources":
+        """Fetch the resources related to this entity.
+
+        Returns:
+            The resources field of the response.
+
+        Raises:
+            RuntimeError: If the resources field was not included in the response.
+        """
+        if self.resources is None:
+            # If the resources are not there, check the class first.
+            resources_annotation = self.model_fields["resources"].annotation
+            assert resources_annotation is not None, (
+                "For each response model, an annotated resources"
+                "field should exist."
+            )
+
+            # resources is defined as:
+            #   resources: Optional[....ResponseResources] = Field(default=None)
+            # We need to find the actual class inside the Optional annotation.
+            from zenml.utils.typing_utils import get_args
+
+            resources_type = get_args(resources_annotation)[0]
+            assert issubclass(resources_type, BaseResponseResources)
+
+            if len(resources_type.model_fields):
+                # If the resources class defines any fields, fetch the resources
+                # through the hydrated version.
+                self.hydrate()
+            else:
+                # Otherwise, use the resources class to create an empty
+                # resources object.
+                self.resources = resources_type()
+
+        if self.resources is None:
+            raise RuntimeError(
+                f"Missing response resources for {type(self).__name__}."
+            )
+
+        return self.resources
+
+
+class BaseDatedResponseBody(BaseResponseBody):
+    """Base body model for entities that track a creation and update timestamp.
+
+    Used as a base class for all body models associated with responses.
+    Features a creation and update timestamp.
+    """
+
+    created: datetime = Field(
+        title="The timestamp when this resource was created."
+    )
+    updated: datetime = Field(
+        title="The timestamp when this resource was last updated."
+    )
+
+
+AnyDatedBody = TypeVar("AnyDatedBody", bound=BaseDatedResponseBody)
+
+
+class BaseIdentifiedResponse(
+    BaseResponse[AnyDatedBody, AnyMetadata, AnyResources],
+    Generic[AnyDatedBody, AnyMetadata, AnyResources],
+):
+    """Base domain model for resources with DB representation."""
+
+    id: UUID = Field(title="The unique resource id.")
+
+    permission_denied: bool = False
+
+    # Helper functions
+    def __hash__(self) -> int:
+        """Implementation of hash magic method.
+
+        Returns:
+            Hash of the UUID.
+        """
+        return hash((type(self),) + tuple([self.id]))
+
+    def __eq__(self, other: Any) -> bool:
+        """Implementation of equality magic method.
+
+        Args:
+            other: The other object to compare to.
+
+        Returns:
+            True if the other object is of the same type and has the same UUID.
+        """
+        if isinstance(other, type(self)):
+            return self.id == other.id
+        else:
+            return False
+
+    def _validate_hydrated_version(
+        self,
+        hydrated_model: "BaseResponse[AnyDatedBody, AnyMetadata, AnyResources]",
+    ) -> None:
+        """Helper method to validate the values within the hydrated version.
+
+        Args:
+            hydrated_model: the hydrated version of the model.
+
+        Raises:
+            HydrationError: if the hydrated version has different values set
+                for either the name of the body fields and the
+                _method_body_mutation is set to ResponseBodyUpdate.DENY.
+        """
+        super()._validate_hydrated_version(hydrated_model)
+
+        assert isinstance(hydrated_model, type(self))
+
+        # Check if the ID is the same
+        if self.id != hydrated_model.id:
+            raise HydrationError(
+                "The hydrated version of the model does not have the same id."
+            )
+
+    def get_hydrated_version(
+        self,
+    ) -> "BaseIdentifiedResponse[AnyDatedBody, AnyMetadata, AnyResources]":
+        """Abstract method to fetch the hydrated version of the model.
+
+        Raises:
+            NotImplementedError: in case the method is not implemented.
+        """
+        raise NotImplementedError(
+            "Please implement a `get_hydrated_version` method before "
+            "using/hydrating the model."
+        )
+
+    def get_body(self) -> "AnyDatedBody":
+        """Fetch the body of the entity.
+
+        Returns:
+            The body field of the response.
+
+        Raises:
+            IllegalOperationError: If the user lacks permission to access the
+                entity represented by this response.
+        """
+        if self.permission_denied:
+            raise IllegalOperationError(
+                f"Missing permissions to access {type(self).__name__} with "
+                f"ID {self.id}."
+            )
+
+        return super().get_body()
+
+    def get_metadata(self) -> "AnyMetadata":
+        """Fetch the metadata of the entity.
+
+        Returns:
+            The metadata field of the response.
+
+        Raises:
+            IllegalOperationError: If the user lacks permission to access this
+                entity represented by this response.
+        """
+        if self.permission_denied:
+            raise IllegalOperationError(
+                f"Missing permissions to access {type(self).__name__} with "
+                f"ID {self.id}."
+            )
+
+        return super().get_metadata()
+
     # Analytics
     def get_analytics_metadata(self) -> Dict[str, Any]:
         """Fetches the analytics metadata for base response models.
@@ -324,7 +466,7 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
 
     # Body and metadata properties
     @property
-    def created(self) -> Optional[datetime]:
+    def created(self) -> datetime:
         """The `created` property.
 
         Returns:
@@ -333,7 +475,7 @@ class BaseResponse(GenericModel, Generic[AnyBody, AnyMetadata], BaseZenModel):
         return self.get_body().created
 
     @property
-    def updated(self) -> Optional[datetime]:
+    def updated(self) -> datetime:
         """The `updated` property.
 
         Returns:

zenml/models/v2/base/base_plugin_flavor.py

@@ -0,0 +1,76 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Plugin flavor model definitions."""
+
+from typing import Generic, TypeVar
+
+from pydantic import ConfigDict, Field
+
+from zenml.enums import PluginSubType, PluginType
+from zenml.models.v2.base.base import (
+    BaseResponse,
+    BaseResponseBody,
+    BaseResponseMetadata,
+    BaseResponseResources,
+)
+
+
+class BasePluginResponseBody(BaseResponseBody):
+    """Response body for plugins."""
+
+
+class BasePluginResponseMetadata(BaseResponseMetadata):
+    """Response metadata for plugins."""
+
+
+class BasePluginResponseResources(BaseResponseResources):
+    """Response resources for plugins."""
+
+
+AnyPluginBody = TypeVar("AnyPluginBody", bound=BasePluginResponseBody)
+AnyPluginMetadata = TypeVar(
+    "AnyPluginMetadata", bound=BasePluginResponseMetadata
+)
+AnyPluginResources = TypeVar(
+    "AnyPluginResources", bound=BasePluginResponseResources
+)
+
+
+class BasePluginFlavorResponse(
+    BaseResponse[AnyPluginBody, AnyPluginMetadata, AnyPluginResources],
+    Generic[AnyPluginBody, AnyPluginMetadata, AnyPluginResources],
+):
+    """Base response for all Plugin Flavors."""
+
+    name: str = Field(title="Name of the flavor.")
+    type: PluginType = Field(title="Type of the plugin.")
+    subtype: PluginSubType = Field(title="Subtype of the plugin.")
+    model_config = ConfigDict(extra="ignore")
+
+    def get_hydrated_version(
+        self,
+    ) -> "BasePluginFlavorResponse[AnyPluginBody, AnyPluginMetadata, AnyPluginResources]":
+        """Abstract method to fetch the hydrated version of the model.
+
+        Returns:
+            Hydrated version of the PluginFlavorResponse
+        """
+        # TODO: shouldn't this call the Zen store ? The client should not have
+        #  to know about the plugin flavor registry
+        from zenml.zen_server.utils import plugin_flavor_registry
+
+        plugin_flavor = plugin_flavor_registry().get_flavor_class(
+            name=self.name, _type=self.type, subtype=self.subtype
+        )
+        return plugin_flavor.get_flavor_response_model(hydrate=True)