zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/models/v2/core/artifact_version.py

@@ -24,12 +24,13 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
-from zenml.config.source import Source, convert_source_validator
+from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
+from zenml.model.model import Model
 from zenml.models.v2.base.filter import StrFilter
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
@@ -43,9 +44,8 @@ from zenml.models.v2.core.artifact import ArtifactResponse
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+    from sqlalchemy.sql.elements import ColumnElement
 
-    from zenml.model.model import Model
     from zenml.models.v2.core.artifact_visualization import (
         ArtifactVisualizationRequest,
         ArtifactVisualizationResponse,
@@ -68,8 +68,7 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
         title="ID of the artifact to which this version belongs.",
     )
     version: Union[str, int] = Field(
-        title="Version of the artifact.",
-        max_length=STR_FIELD_MAX_LENGTH,
+        title="Version of the artifact.", union_mode="left_to_right"
     )
     has_custom_name: bool = Field(
         title="Whether the name is custom (True) or auto-generated (False).",
@@ -83,10 +82,10 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     uri: str = Field(
         title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
     )
-    materializer: Source = Field(
+    materializer: SourceWithValidator = Field(
         title="Materializer class to use for this artifact.",
     )
-    data_type: Source = Field(
+    data_type: SourceWithValidator = Field(
         title="Data type of the artifact.",
     )
     tags: Optional[List[str]] = Field(
@@ -98,7 +97,26 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
         default=None, title="Visualizations of the artifact."
     )
 
-    _convert_source = convert_source_validator("materializer", "data_type")
+    @field_validator("version")
+    @classmethod
+    def str_field_max_length_check(cls, value: Any) -> Any:
+        """Checks if the length of the value exceeds the maximum str length.
+
+        Args:
+            value: the value set in the field
+
+        Returns:
+            the value itself.
+
+        Raises:
+            AssertionError: if the length of the field is longer than the
+                maximum threshold.
+        """
+        assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
+            "The length of the value for this field can not "
+            f"exceed {STR_FIELD_MAX_LENGTH}"
+        )
+        return value
 
 
 # ------------------ Update Model ------------------
@@ -121,28 +139,45 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
     artifact: ArtifactResponse = Field(
         title="Artifact to which this version belongs."
     )
-    version: Union[str, int] = Field(
-        title="Version of the artifact.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
+    version: str = Field(title="Version of the artifact.")
     uri: str = Field(
         title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
     )
     type: ArtifactType = Field(title="Type of the artifact.")
-    materializer: Source = Field(
+    materializer: SourceWithValidator = Field(
         title="Materializer class to use for this artifact.",
     )
-    data_type: Source = Field(
+    data_type: SourceWithValidator = Field(
         title="Data type of the artifact.",
     )
     tags: List[TagResponse] = Field(
         title="Tags associated with the model",
     )
     producer_pipeline_run_id: Optional[UUID] = Field(
-        title="The ID of the pipeline run that generated this artifact version."
+        title="The ID of the pipeline run that generated this artifact version.",
+        default=None,
     )
 
-    _convert_source = convert_source_validator("materializer", "data_type")
+    @field_validator("version")
+    @classmethod
+    def str_field_max_length_check(cls, value: Any) -> Any:
+        """Checks if the length of the value exceeds the maximum str length.
+
+        Args:
+            value: the value set in the field
+
+        Returns:
+            the value itself.
+
+        Raises:
+            AssertionError: if the length of the field is longer than the
+                maximum threshold.
+        """
+        assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
+            "The length of the value for this field can not "
+            f"exceed {STR_FIELD_MAX_LENGTH}"
+        )
+        return value
 
 
 class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -396,7 +431,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of artifact versions."""
 
     # `name` and `only_unused` refer to properties related to other entities
-    #  rather than a field in the db, hence they needs to be handled
+    #  rather than a field in the db, hence they need to be handled
     #  explicitly
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
@@ -407,6 +442,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     artifact_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="ID of the artifact to which this version belongs.",
+        union_mode="left_to_right",
     )
     name: Optional[str] = Field(
         default=None,
@@ -419,6 +455,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     version_number: Optional[Union[int, str]] = Field(
         default=None,
         description="Version of the artifact if it is an integer",
+        union_mode="left_to_right",
     )
     uri: Optional[str] = Field(
         default=None,
@@ -437,13 +474,19 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="Datatype of the artifact",
     )
     artifact_store_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Artifact store for this artifact"
+        default=None,
+        description="Artifact store for this artifact",
+        union_mode="left_to_right",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace for this artifact"
+        default=None,
+        description="Workspace for this artifact",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that produced this artifact"
+        default=None,
+        description="User that produced this artifact",
+        union_mode="left_to_right",
     )
     only_unused: Optional[bool] = Field(
         default=False, description="Filter only for unused artifacts"
@@ -453,9 +496,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="Filter only artifacts with/without custom names.",
     )
 
-    def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+    def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
         """Get custom filters.
 
         Returns:
@@ -463,8 +504,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         """
         custom_filters = super().get_custom_filters()
 
-        from sqlalchemy import and_
-        from sqlmodel import select
+        from sqlmodel import and_, select
 
         from zenml.zen_stores.schemas.artifact_schemas import (
             ArtifactSchema,
@@ -482,7 +522,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
                 column="name",
                 value=value,
             )
-            artifact_name_filter = and_(  # type: ignore[type-var]
+            artifact_name_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                 filter_.generate_query_conditions(ArtifactSchema),
             )
@@ -500,7 +540,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             custom_filters.append(unused_filter)
 
         if self.has_custom_name is not None:
-            custom_name_filter = and_(  # type: ignore[type-var]
+            custom_name_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                 ArtifactSchema.has_custom_name == self.has_custom_name,
             )
@@ -520,9 +560,9 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
     """
 
     id: Optional[UUID] = None  # type: ignore[assignment]
-    _lazy_load_name: Optional[str] = None
-    _lazy_load_version: Optional[str] = None
-    _lazy_load_model: "Model"
+    lazy_load_name: Optional[str] = None
+    lazy_load_version: Optional[str] = None
+    lazy_load_model: Model
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.
@@ -552,13 +592,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
 
         return RunMetadataLazyGetter(  # type: ignore[return-value]
-            self._lazy_load_model,
-            self._lazy_load_name,
-            self._lazy_load_version,
+            self.lazy_load_model,
+            self.lazy_load_name,
+            self.lazy_load_version,
         )
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Allow extras to include the lazy load attributes
-        extra = "allow"

zenml/models/v2/core/code_repository.py

@@ -20,6 +20,7 @@ from pydantic import Field
 
 from zenml.config.source import Source
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -28,9 +29,6 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import (
-    update_model,
-)
 
 # ------------------ Request Model ------------------
 
@@ -48,21 +46,45 @@ class CodeRepositoryRequest(WorkspaceScopedRequest):
     source: Source = Field(description="The code repository source.")
     logo_url: Optional[str] = Field(
         description="Optional URL of a logo (png, jpg or svg) for the "
-        "code repository."
+        "code repository.",
+        default=None,
     )
     description: Optional[str] = Field(
         description="Code repository description.",
         max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
     )
 
 
 # ------------------ Update Model ------------------
 
 
-@update_model
-class CodeRepositoryUpdate(CodeRepositoryRequest):
+class CodeRepositoryUpdate(BaseUpdate):
     """Update model for code repositories."""
 
+    name: Optional[str] = Field(
+        title="The name of the code repository.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    config: Optional[Dict[str, Any]] = Field(
+        description="Configuration for the code repository.",
+        default=None,
+    )
+    source: Optional[Source] = Field(
+        description="The code repository source.", default=None
+    )
+    logo_url: Optional[str] = Field(
+        description="Optional URL of a logo (png, jpg or svg) for the "
+        "code repository.",
+        default=None,
+    )
+    description: Optional[str] = Field(
+        description="Code repository description.",
+        max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 
@@ -165,10 +187,15 @@ class CodeRepositoryFilter(WorkspaceScopedFilter):
 
     name: Optional[str] = Field(
         description="Name of the code repository.",
+        default=None,
     )
-    workspace_id: Union[UUID, str, None] = Field(
-        description="Workspace of the code repository."
+    workspace_id: Optional[Union[UUID, str]] = Field(
+        description="Workspace of the code repository.",
+        default=None,
+        union_mode="left_to_right",
     )
-    user_id: Union[UUID, str, None] = Field(
-        description="User that created the code repository."
+    user_id: Optional[Union[UUID, str]] = Field(
+        description="User that created the code repository.",
+        default=None,
+        union_mode="left_to_right",
     )

zenml/models/v2/core/component.py

@@ -25,11 +25,11 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import LogicalOperators, StackComponentType
-from zenml.models.v2.base.internal import server_owned_request_model
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -38,11 +38,10 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.update import update_model
 from zenml.utils import secret_utils
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+    from sqlalchemy.sql.elements import ColumnElement
     from sqlmodel import SQLModel
 
     from zenml.models.v2.core.service_connector import (
@@ -103,7 +102,8 @@ class ComponentRequest(ComponentBase, WorkspaceScopedRequest):
         title="The service connector linked to this stack component.",
     )
 
-    @validator("name")
+    @field_validator("name")
+    @classmethod
     def name_cant_be_a_secret_reference(cls, name: str) -> str:
         """Validator to ensure that the given name is not a secret reference.
 
@@ -124,22 +124,63 @@ class ComponentRequest(ComponentBase, WorkspaceScopedRequest):
         return name
 
 
-@server_owned_request_model
 class InternalComponentRequest(ComponentRequest):
     """Internal component request model."""
 
-    pass
+    user: Optional[UUID] = Field(  # type: ignore[assignment]
+        title="The id of the user that created this resource.",
+        default=None,
+    )
 
 
 # ------------------ Update Model ------------------
 
 
-@update_model
-class ComponentUpdate(ComponentRequest):
+class ComponentUpdate(BaseUpdate):
     """Update model for stack components."""
 
+    ANALYTICS_FIELDS: ClassVar[List[str]] = ["type", "flavor"]
+
+    name: Optional[str] = Field(
+        title="The name of the stack component.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    type: Optional[StackComponentType] = Field(
+        title="The type of the stack component.",
+        default=None,
+    )
+    flavor: Optional[str] = Field(
+        title="The flavor of the stack component.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    configuration: Optional[Dict[str, Any]] = Field(
+        title="The stack component configuration.",
+        default=None,
+    )
+    connector_resource_id: Optional[str] = Field(
+        description="The ID of a specific resource instance to "
+        "gain access to through the connector",
+        default=None,
+    )
+    labels: Optional[Dict[str, Any]] = Field(
+        title="The stack component labels.",
+        default=None,
+    )
+    component_spec_path: Optional[str] = Field(
+        title="The path to the component spec used for mlstacks deployments.",
+        default=None,
+    )
+    connector: Optional[UUID] = Field(
+        title="The service connector linked to this stack component.",
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
+
+
 class ComponentResponseBody(WorkspaceScopedResponseBody):
     """Response body for components."""
 
@@ -311,16 +352,24 @@ class ComponentFilter(WorkspaceScopedFilter):
         description="Type of the stack component",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the stack component"
+        default=None,
+        description="Workspace of the stack component",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the stack component"
+        default=None,
+        description="User of the stack component",
+        union_mode="left_to_right",
     )
     connector_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Connector linked to the stack component"
+        default=None,
+        description="Connector linked to the stack component",
+        union_mode="left_to_right",
     )
     stack_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Stack of the stack component"
+        default=None,
+        description="Stack of the stack component",
+        union_mode="left_to_right",
     )
 
     def set_scope_type(self, component_type: str) -> None:
@@ -333,7 +382,7 @@ class ComponentFilter(WorkspaceScopedFilter):
 
     def generate_filter(
         self, table: Type["SQLModel"]
-    ) -> Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]:
+    ) -> Union["ColumnElement[bool]"]:
         """Generate the filter for the query.
 
         Stack components can be scoped by type to narrow the search.
@@ -344,7 +393,7 @@ class ComponentFilter(WorkspaceScopedFilter):
         Returns:
             The filter expression for the query.
         """
-        from sqlalchemy import and_, or_
+        from sqlmodel import and_, or_
 
         from zenml.zen_stores.schemas import (
             StackComponentSchema,
@@ -361,7 +410,7 @@ class ComponentFilter(WorkspaceScopedFilter):
                 or_ if self.logical_operator == LogicalOperators.OR else and_
             )
 
-            stack_filter = and_(  # type: ignore[type-var]
+            stack_filter = and_(
                 StackCompositionSchema.stack_id == self.stack_id,
                 StackCompositionSchema.component_id == StackComponentSchema.id,
             )

zenml/models/v2/core/device.py

@@ -440,22 +440,32 @@ class OAuthDeviceFilter(UserScopedFilter):
     """Model to enable advanced filtering of OAuth2 devices."""
 
     expires: Optional[Union[datetime, str, None]] = Field(
-        default=None, description="The expiration date of the OAuth2 device."
+        default=None,
+        description="The expiration date of the OAuth2 device.",
+        union_mode="left_to_right",
     )
     client_id: Union[UUID, str, None] = Field(
-        default=None, description="The client ID of the OAuth2 device."
+        default=None,
+        description="The client ID of the OAuth2 device.",
+        union_mode="left_to_right",
     )
     status: Union[OAuthDeviceStatus, str, None] = Field(
-        default=None, description="The status of the OAuth2 device."
+        default=None,
+        description="The status of the OAuth2 device.",
+        union_mode="left_to_right",
     )
     trusted_device: Union[bool, str, None] = Field(
         default=None,
         description="Whether the OAuth2 device was marked as trusted.",
+        union_mode="left_to_right",
     )
     failed_auth_attempts: Union[int, str, None] = Field(
         default=None,
         description="The number of failed authentication attempts.",
+        union_mode="left_to_right",
     )
     last_login: Optional[Union[datetime, str, None]] = Field(
-        default=None, description="The date of the last successful login."
+        default=None,
+        description="The date of the last successful login.",
+        union_mode="left_to_right",
     )

zenml/models/v2/core/event_source.py

@@ -48,7 +48,6 @@ class EventSourceRequest(WorkspaceScopedRequest):
     )
     plugin_subtype: PluginSubType = Field(
         title="The plugin subtype of the event source.",
-        max_length=STR_FIELD_MAX_LENGTH,
     )
     description: str = Field(
         default="",
@@ -118,7 +117,6 @@ class EventSourceResponseBody(WorkspaceScopedResponseBody):
     )
     plugin_subtype: PluginSubType = Field(
         title="The plugin subtype of the event source.",
-        max_length=STR_FIELD_MAX_LENGTH,
     )
     is_active: bool = Field(
         title="Whether the event source is active.",
@@ -240,6 +238,7 @@ class EventSourceFilter(WorkspaceScopedFilter):
         description="Flavor of the event source",
     )
     plugin_subtype: Optional[str] = Field(
+        default=None,
         title="The plugin sub type of the event source.",
         max_length=STR_FIELD_MAX_LENGTH,
     )

zenml/models/v2/core/flavor.py

@@ -20,7 +20,7 @@ from pydantic import Field
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import StackComponentType
-from zenml.models.v2.base.internal import server_owned_request_model
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     UserScopedRequest,
     UserScopedResponse,
@@ -29,7 +29,6 @@ from zenml.models.v2.base.scoped import (
     UserScopedResponseResources,
     WorkspaceScopedFilter,
 )
-from zenml.models.v2.base.update import update_model
 
 if TYPE_CHECKING:
     from zenml.models import (
@@ -104,20 +103,83 @@ class FlavorRequest(UserScopedRequest):
     )
 
 
-@server_owned_request_model
 class InternalFlavorRequest(FlavorRequest):
     """Internal flavor request model."""
 
-    pass
+    user: Optional[UUID] = Field(  # type: ignore[assignment]
+        title="The id of the user that created this resource.",
+        default=None,
+    )
 
 
 # ------------------ Update Model ------------------
 
 
-@update_model
-class FlavorUpdate(FlavorRequest):
+class FlavorUpdate(BaseUpdate):
     """Update model for flavors."""
 
+    name: Optional[str] = Field(
+        title="The name of the Flavor.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    type: Optional[StackComponentType] = Field(
+        title="The type of the Flavor.", default=None
+    )
+    config_schema: Optional[Dict[str, Any]] = Field(
+        title="The JSON schema of this flavor's corresponding configuration.",
+        default=None,
+    )
+    connector_type: Optional[str] = Field(
+        title="The type of the connector that this flavor uses.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    connector_resource_type: Optional[str] = Field(
+        title="The resource type of the connector that this flavor uses.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    connector_resource_id_attr: Optional[str] = Field(
+        title="The name of an attribute in the stack component configuration "
+        "that plays the role of resource ID when linked to a service "
+        "connector.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    source: Optional[str] = Field(
+        title="The path to the module which contains this Flavor.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    integration: Optional[str] = Field(
+        title="The name of the integration that the Flavor belongs to.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    logo_url: Optional[str] = Field(
+        title="Optionally, a url pointing to a png,"
+        "svg or jpg can be attached.",
+        default=None,
+    )
+    docs_url: Optional[str] = Field(
+        title="Optionally, a url pointing to docs, within docs.zenml.io.",
+        default=None,
+    )
+    sdk_docs_url: Optional[str] = Field(
+        title="Optionally, a url pointing to SDK docs,"
+        "within sdkdocs.zenml.io.",
+        default=None,
+    )
+    is_custom: Optional[bool] = Field(
+        title="Whether or not this flavor is a custom, user created flavor.",
+        default=None,
+    )
+    workspace: Optional[UUID] = Field(
+        title="The workspace to which this resource belongs.",
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 
@@ -367,8 +429,12 @@ class FlavorFilter(WorkspaceScopedFilter):
         description="Integration associated with the flavor",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the stack"
+        default=None,
+        description="Workspace of the stack",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the stack"
+        default=None,
+        description="User of the stack",
+        union_mode="left_to_right",
     )