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

zenml/models/v2/core/model_version_artifact.py

@@ -13,51 +13,46 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and artifacts."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Type, TypeVar, Union
 from uuid import UUID
 
-from pydantic import Field, validator
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from pydantic import ConfigDict, Field
 
 from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
-    BaseResponse,
-    BaseResponseBody,
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
+    BaseRequest,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
-from zenml.models.v2.base.filter import StrFilter
-from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
-    WorkspaceScopedRequest,
-)
+from zenml.models.v2.base.filter import BaseFilter, StrFilter
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
 
 
-class ModelVersionArtifactRequest(WorkspaceScopedRequest):
+class ModelVersionArtifactRequest(BaseRequest):
     """Request model for links between model versions and artifacts."""
 
-    model: UUID
     model_version: UUID
     artifact_version: UUID
-    is_model_artifact: bool = False
-    is_deployment_artifact: bool = False
-
-    @validator("is_deployment_artifact")
-    def _validate_is_endpoint_artifact(
-        cls, is_deployment_artifact: bool, values: Dict[str, Any]
-    ) -> bool:
-        is_model_artifact = values.get("is_model_artifact", False)
-        if is_model_artifact and is_deployment_artifact:
-            raise ValueError(
-                "Artifact cannot be a model artifact and deployment artifact "
-                "at the same time."
-            )
-        return is_deployment_artifact
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
 
 # ------------------ Update Model ------------------
@@ -67,31 +62,34 @@ class ModelVersionArtifactRequest(WorkspaceScopedRequest):
 # ------------------ Response Model ------------------
 
 
-class ModelVersionArtifactResponseBody(BaseResponseBody):
+class ModelVersionArtifactResponseBody(BaseDatedResponseBody):
     """Response body for links between model versions and artifacts."""
 
-    model: UUID
     model_version: UUID
     artifact_version: "ArtifactVersionResponse"
-    is_model_artifact: bool = False
-    is_deployment_artifact: bool = False
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class ModelVersionArtifactResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the model version artifact entity."""
 
 
 class ModelVersionArtifactResponse(
-    BaseResponse[ModelVersionArtifactResponseBody, BaseResponseMetadata]
+    BaseIdentifiedResponse[
+        ModelVersionArtifactResponseBody,
+        BaseResponseMetadata,
+        ModelVersionArtifactResponseResources,
+    ]
 ):
     """Response model for links between model versions and artifacts."""
 
-    # Body and metadata properties
-    @property
-    def model(self) -> UUID:
-        """The `model` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().model
-
     @property
     def model_version(self) -> UUID:
         """The `model_version` property.
@@ -110,68 +108,43 @@ class ModelVersionArtifactResponse(
         """
         return self.get_body().artifact_version
 
-    @property
-    def is_model_artifact(self) -> bool:
-        """The `is_model_artifact` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().is_model_artifact
-
-    @property
-    def is_deployment_artifact(self) -> bool:
-        """The `is_deployment_artifact` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().is_deployment_artifact
-
 
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionArtifactFilter(WorkspaceScopedFilter):
+class ModelVersionArtifactFilter(BaseFilter):
     """Model version pipeline run links filter model."""
 
     # Artifact name and type are not DB fields and need to be handled separately
     FILTER_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *BaseFilter.FILTER_EXCLUDE_FIELDS,
         "artifact_name",
         "only_data_artifacts",
         "only_model_artifacts",
         "only_deployment_artifacts",
         "has_custom_name",
+        "user",
     ]
     CLI_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
+        *BaseFilter.CLI_EXCLUDE_FIELDS,
         "only_data_artifacts",
         "only_model_artifacts",
         "only_deployment_artifacts",
         "has_custom_name",
-        "model_id",
         "model_version_id",
-        "user_id",
-        "workspace_id",
         "updated",
         "id",
     ]
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The workspace of the Model Version"
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The user of the Model Version"
-    )
-    model_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model ID"
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model version ID"
+        default=None,
+        description="Filter by model version ID",
+        union_mode="left_to_right",
     )
     artifact_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by artifact ID"
+        default=None,
+        description="Filter by artifact ID",
+        union_mode="left_to_right",
     )
     artifact_name: Optional[str] = Field(
         default=None,
@@ -181,25 +154,39 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
     only_model_artifacts: Optional[bool] = False
     only_deployment_artifacts: Optional[bool] = False
     has_custom_name: Optional[bool] = None
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the artifact.",
+    )
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        self, table: Type["AnySchema"]
+    ) -> List[Union["ColumnElement[bool]"]]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
-        from sqlalchemy import and_
+        from sqlmodel import and_, col
 
-        from zenml.zen_stores.schemas.artifact_schemas import (
+        from zenml.zen_stores.schemas import (
             ArtifactSchema,
             ArtifactVersionSchema,
-        )
-        from zenml.zen_stores.schemas.model_schemas import (
             ModelVersionArtifactSchema,
+            UserSchema,
         )
 
         if self.artifact_name:
@@ -209,7 +196,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
                 column="name",
                 value=value,
             )
-            artifact_name_filter = and_(  # type: ignore[type-var]
+            artifact_name_filter = and_(
                 ModelVersionArtifactSchema.artifact_version_id
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
@@ -219,25 +206,32 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
 
         if self.only_data_artifacts:
             data_artifact_filter = and_(
-                ModelVersionArtifactSchema.is_model_artifact.is_(False),  # type: ignore[attr-defined]
-                ModelVersionArtifactSchema.is_deployment_artifact.is_(False),  # type: ignore[attr-defined]
+                ModelVersionArtifactSchema.artifact_version_id
+                == ArtifactVersionSchema.id,
+                col(ArtifactVersionSchema.type).not_in(
+                    ["ServiceArtifact", "ModelArtifact"]
+                ),
             )
             custom_filters.append(data_artifact_filter)
 
         if self.only_model_artifacts:
             model_artifact_filter = and_(
-                ModelVersionArtifactSchema.is_model_artifact.is_(True),  # type: ignore[attr-defined]
+                ModelVersionArtifactSchema.artifact_version_id
+                == ArtifactVersionSchema.id,
+                ArtifactVersionSchema.type == "ModelArtifact",
             )
             custom_filters.append(model_artifact_filter)
 
         if self.only_deployment_artifacts:
             deployment_artifact_filter = and_(
-                ModelVersionArtifactSchema.is_deployment_artifact.is_(True),  # type: ignore[attr-defined]
+                ModelVersionArtifactSchema.artifact_version_id
+                == ArtifactVersionSchema.id,
+                ArtifactVersionSchema.type == "ServiceArtifact",
             )
             custom_filters.append(deployment_artifact_filter)
 
         if self.has_custom_name is not None:
-            custom_name_filter = and_(  # type: ignore[type-var]
+            custom_name_filter = and_(
                 ModelVersionArtifactSchema.artifact_version_id
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
@@ -245,4 +239,17 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
             )
             custom_filters.append(custom_name_filter)
 
+        if self.user:
+            user_filter = and_(
+                ModelVersionArtifactSchema.artifact_version_id
+                == ArtifactVersionSchema.id,
+                ArtifactVersionSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
+                ),
+            )
+            custom_filters.append(user_filter)
+
         return custom_filters

zenml/models/v2/core/model_version_pipeline_run.py

@@ -13,37 +13,46 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and pipeline runs."""
 
-from typing import TYPE_CHECKING, Any, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Type, TypeVar, Union
 from uuid import UUID
 
-from pydantic import Field
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from pydantic import ConfigDict, Field
+from sqlalchemy.sql.elements import ColumnElement
 
 from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
-    BaseResponse,
-    BaseResponseBody,
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
+    BaseRequest,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
-from zenml.models.v2.base.filter import StrFilter
-from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
-    WorkspaceScopedRequest,
-)
+from zenml.models.v2.base.filter import BaseFilter, StrFilter
+from zenml.models.v2.core.pipeline_run import PipelineRunResponse
 
 if TYPE_CHECKING:
-    from zenml.models.v2.core.pipeline_run import PipelineRunResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
 
 # ------------------ Request Model ------------------
 
 
-class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
+class ModelVersionPipelineRunRequest(BaseRequest):
     """Request model for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
     pipeline_run: UUID
 
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
 
 # ------------------ Update Model ------------------
 
@@ -52,29 +61,34 @@ class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
 # ------------------ Response Model ------------------
 
 
-class ModelVersionPipelineRunResponseBody(BaseResponseBody):
+class ModelVersionPipelineRunResponseBody(BaseDatedResponseBody):
     """Response body for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
-    pipeline_run: "PipelineRunResponse"
+    pipeline_run: PipelineRunResponse
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class ModelVersionPipelineRunResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the model version pipeline run entity."""
 
 
 class ModelVersionPipelineRunResponse(
-    BaseResponse[ModelVersionPipelineRunResponseBody, BaseResponseMetadata]
+    BaseIdentifiedResponse[
+        ModelVersionPipelineRunResponseBody,
+        BaseResponseMetadata,
+        ModelVersionPipelineRunResponseResources,
+    ]
 ):
     """Response model for links between model versions and pipeline runs."""
 
-    # Body and metadata properties
-    @property
-    def model(self) -> UUID:
-        """The `model` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().model
-
     @property
     def model_version(self) -> UUID:
         """The `model_version` property.
@@ -97,61 +111,67 @@ class ModelVersionPipelineRunResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
+class ModelVersionPipelineRunFilter(BaseFilter):
     """Model version pipeline run links filter model."""
 
-    # Pipeline run name is not a DB field and needs to be handled separately
     FILTER_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *BaseFilter.FILTER_EXCLUDE_FIELDS,
         "pipeline_run_name",
+        "user",
     ]
     CLI_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
-        "model_id",
+        *BaseFilter.CLI_EXCLUDE_FIELDS,
         "model_version_id",
-        "user_id",
-        "workspace_id",
         "updated",
         "id",
     ]
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The workspace of the Model Version"
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The user of the Model Version"
-    )
-    model_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model ID"
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model version ID"
+        default=None,
+        description="Filter by model version ID",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by pipeline run ID"
+        default=None,
+        description="Filter by pipeline run ID",
+        union_mode="left_to_right",
     )
     pipeline_run_name: Optional[str] = Field(
         default=None,
         description="Name of the pipeline run",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the pipeline run.",
+    )
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
-        from sqlalchemy import and_
+        from sqlmodel import and_
 
-        from zenml.zen_stores.schemas.model_schemas import (
+        from zenml.zen_stores.schemas import (
             ModelVersionPipelineRunSchema,
-        )
-        from zenml.zen_stores.schemas.pipeline_run_schemas import (
             PipelineRunSchema,
+            UserSchema,
         )
 
         if self.pipeline_run_name:
@@ -163,11 +183,24 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
                 column="name",
                 value=value,
             )
-            pipeline_run_name_filter = and_(  # type: ignore[type-var]
+            pipeline_run_name_filter = and_(
                 ModelVersionPipelineRunSchema.pipeline_run_id
                 == PipelineRunSchema.id,
                 filter_.generate_query_conditions(PipelineRunSchema),
             )
             custom_filters.append(pipeline_run_name_filter)
 
+        if self.user:
+            user_filter = and_(
+                ModelVersionPipelineRunSchema.pipeline_run_id
+                == PipelineRunSchema.id,
+                PipelineRunSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
+                ),
+            )
+            custom_filters.append(user_filter)
+
         return custom_filters