zenml-nightly 0.58.2.dev20240617__py3-none-any.whl → 0.58.2.dev20240619__py3-none-any.whl

zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py

@@ -100,15 +100,16 @@ def resolve_duplicate_names() -> None:
     """Resolve duplicate names for shareable entities."""
     connection = op.get_bind()
 
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
     meta.reflect(
+        bind=op.get_bind(),
         only=(
             "stack",
             "stack_component",
             "stack_composition",
             "service_connector",
             "workspace",
-        )
+        ),
     )
 
     stack_table = sa.Table("stack", meta)

zenml/zen_stores/migrations/versions/7834208cc3f6_artifact_project_scoping.py

@@ -49,15 +49,16 @@ def upgrade() -> None:
     # ------------
 
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
     meta.reflect(
+        bind=op.get_bind(),
         only=(
             "pipeline_run",
             "step_run",
             "artifact",
             "step_run_output_artifact",
             "workspace",
-        )
+        ),
     )
     pipeline_runs = sa.Table("pipeline_run", meta)
     step_runs = sa.Table("step_run", meta)
@@ -67,12 +68,12 @@ def upgrade() -> None:
 
     # Use the first workspace as the default workspace in case an artifact or
     # step run has no associated pipeline run or producer step run.
-    default_workspace = conn.execute(select([workspace])).first()
+    default_workspace = conn.execute(select(workspace)).first()
 
     # For each step run, set user and project according to its pipeline run
-    for step_run in conn.execute(select([step_runs])).all():
+    for step_run in conn.execute(select(step_runs)).all():
         pipeline_run = conn.execute(
-            select([pipeline_runs]).where(
+            select(pipeline_runs).where(
                 pipeline_runs.c.id == step_run.pipeline_run_id
             )
         ).first()
@@ -90,9 +91,9 @@ def upgrade() -> None:
         )
 
     # For each artifact, set user and project according to its producer step run
-    for artifact in conn.execute(select([artifacts])).all():
+    for artifact in conn.execute(select(artifacts)).all():
         producer_step_run = conn.execute(
-            select([step_runs])
+            select(step_runs)
             .where(step_runs.c.status == "completed")
             .where(step_runs.c.id == step_run_output_artifacts.c.step_id)
             .where(step_run_output_artifacts.c.artifact_id == artifact.id)

zenml/zen_stores/migrations/versions/7b651bf6822e_track_secrets_in_db.py

@@ -58,7 +58,7 @@ class BaseSecretsStoreBackend(BaseModel):
     """Base class for accessing and listing ZenML secrets."""
 
     server_id: UUID
-    client: Any
+    client: Any = None
 
     def _get_secret_metadata(
         self,
@@ -603,15 +603,17 @@ def upgrade() -> None:
     # Transfer secrets from the external secrets store to the db
 
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("secret", "user", "workspace", "identity"))
+    meta = sa.MetaData()
+    meta.reflect(
+        only=("secret", "user", "workspace", "identity"), bind=op.get_bind()
+    )
     secrets = sa.Table("secret", meta)
     users = sa.Table("user", meta)
     workspaces = sa.Table("workspace", meta)
     identity = sa.Table("identity", meta)
 
     # Extract the ZenML deployment ID from the identity table
-    server_id = UUID(conn.execute(identity.select()).first()["id"])  # type: ignore[index]
+    server_id = UUID(conn.execute(sa.select(identity.c.id)).scalar_one())
 
     # Initialize the secrets store backend based on the backend client extracted
     # from the secrets store.

zenml/zen_stores/migrations/versions/7e4a481d17f7_add_identity_table.py

@@ -27,11 +27,11 @@ def upgrade() -> None:
         sa.PrimaryKeyConstraint("id"),
     )
     # get metadata from current connection
-    meta = sa.MetaData(bind=op.get_bind())
+    meta = sa.MetaData()
 
     # pass in tuple with tables we want to reflect, otherwise whole database
     #  will get reflected
-    meta.reflect(only=("identity",))
+    meta.reflect(only=("identity",), bind=op.get_bind())
     # Prefill the identity table with a single row that contains the deployment
     # id extracted from the global configuration
     op.bulk_insert(

zenml/zen_stores/migrations/versions/7f603e583dd7_fixed_migration.py

@@ -44,7 +44,7 @@ def upgrade() -> None:
                 sa.Column(
                     "step_configuration",
                     sa.String(length=16777215).with_variant(
-                        mysql.MEDIUMTEXT(), "mysql"
+                        mysql.MEDIUMTEXT, "mysql"
                     ),
                     nullable=True,
                 )

zenml/zen_stores/migrations/versions/a39c4184c8ce_remove_secrets_manager_flavors.py

@@ -19,8 +19,8 @@ depends_on = None
 def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("flavor",))
+    meta = sa.MetaData()
+    meta.reflect(only=("flavor",), bind=op.get_bind())
     flavors = sa.Table("flavor", meta)
 
     # Remove all secrets manager flavors

zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py

@@ -46,13 +46,13 @@ def upgrade() -> None:
 
     # Find all unique artifact names and create a new artifact entry for each
     conn = op.get_bind()
-    meta = sa.MetaData(bind=op.get_bind())
-    meta.reflect(only=("artifact", "artifact_version"))
+    meta = sa.MetaData()
+    meta.reflect(only=("artifact", "artifact_version"), bind=op.get_bind())
     artifacts = sa.Table("artifact", meta)
     artifact_versions = sa.Table("artifact_version", meta)
     artifact_names = conn.execute(
         sa.select(
-            [artifact_versions.c.name, artifact_versions.c.has_custom_name]
+            artifact_versions.c.name, artifact_versions.c.has_custom_name
         ).distinct()
     ).all()
     unique_artifact_names = {
@@ -72,7 +72,7 @@ def upgrade() -> None:
     # Set artifact_id column in artifact_version
     conn.execute(
         artifact_versions.update().values(
-            artifact_id=sa.select([artifacts.c.id])
+            artifact_id=sa.select(artifacts.c.id)
             .where(artifacts.c.name == artifact_versions.c.name)
             .scalar_subquery()
         )

zenml/zen_stores/migrations/versions/alembic_start.py

@@ -22,7 +22,7 @@ def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     # If the tables already exist, skip this migration.
     conn = op.get_bind()
-    inspector = Inspector.from_engine(conn)
+    inspector = Inspector.from_engine(conn.engine)
     tables = inspector.get_table_names()
     if "projectschema" in tables:
         return

zenml/zen_stores/migrations/versions/fbd7f18ced1e_increase_step_run_field_lengths.py

@@ -25,7 +25,7 @@ def upgrade() -> None:
             "parameters",
             existing_type=sa.TEXT(),
             type_=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             existing_nullable=False,
         )
@@ -33,7 +33,7 @@ def upgrade() -> None:
             "step_configuration",
             existing_type=sa.TEXT(),
             type_=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             existing_nullable=False,
         )
@@ -48,7 +48,7 @@ def downgrade() -> None:
         batch_op.alter_column(
             "step_configuration",
             existing_type=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             type_=sa.TEXT(),
             existing_nullable=False,
@@ -56,7 +56,7 @@ def downgrade() -> None:
         batch_op.alter_column(
             "parameters",
             existing_type=sa.String(length=16777215).with_variant(
-                mysql.MEDIUMTEXT(), "mysql"
+                mysql.MEDIUMTEXT, "mysql"
             ),
             type_=sa.TEXT(),
             existing_nullable=False,

zenml/zen_stores/rest_zen_store.py

@@ -32,12 +32,19 @@ from uuid import UUID
 
 import requests
 import urllib3
-from pydantic import BaseModel, root_validator, validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
 from requests.adapters import HTTPAdapter, Retry
 
 import zenml
 from zenml.analytics import source_context
 from zenml.config.global_config import GlobalConfiguration
+from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.store_config import StoreConfiguration
 from zenml.constants import (
     API,
@@ -97,9 +104,7 @@ from zenml.enums import (
     OAuthGrantTypes,
     StoreType,
 )
-from zenml.exceptions import (
-    AuthorizationException,
-)
+from zenml.exceptions import AuthorizationException, MethodNotAllowedError
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.models import (
@@ -274,16 +279,13 @@ class RestZenStoreConfiguration(StoreConfiguration):
     password: Optional[str] = None
     api_key: Optional[str] = None
     api_token: Optional[str] = None
-    verify_ssl: Union[bool, str] = True
+    verify_ssl: Union[bool, str] = Field(True, union_mode="left_to_right")
     http_timeout: int = DEFAULT_HTTP_TIMEOUT
 
-    @root_validator
-    def validate_credentials(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def validate_credentials(self) -> "RestZenStoreConfiguration":
         """Validates the credentials provided in the values dictionary.
 
-        Args:
-            values: A dictionary containing the values to be validated.
-
         Raises:
             ValueError: If neither api_token nor username nor api_key is set.
 
@@ -292,18 +294,15 @@ class RestZenStoreConfiguration(StoreConfiguration):
         """
         # Check if the values dictionary contains either an API token, an API
         # key or a username as non-empty strings.
-        if (
-            values.get("api_token")
-            or values.get("username")
-            or values.get("api_key")
-        ):
-            return values
+        if self.api_token or self.username or self.api_key:
+            return self
         raise ValueError(
             "Neither api_token nor username nor api_key is set in the "
             "store config."
         )
 
-    @validator("url")
+    @field_validator("url")
+    @classmethod
     def validate_url(cls, url: str) -> str:
         """Validates that the URL is a well-formed REST store URL.
 
@@ -331,7 +330,8 @@ class RestZenStoreConfiguration(StoreConfiguration):
 
         return url
 
-    @validator("verify_ssl")
+    @field_validator("verify_ssl")
+    @classmethod
     def validate_verify_ssl(
         cls, verify_ssl: Union[bool, str]
     ) -> Union[bool, str]:
@@ -387,15 +387,14 @@ class RestZenStoreConfiguration(StoreConfiguration):
             with open(self.verify_ssl, "r") as f:
                 self.verify_ssl = f.read()
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Don't validate attributes when assigning them. This is necessary
         # because the `verify_ssl` attribute can be expanded to the contents
         # of the certificate file.
-        validate_assignment = False
+        validate_assignment=False,
         # Forbid extra attributes set in the class.
-        extra = "forbid"
+        extra="forbid",
+    )
 
 
 class RestZenStore(BaseZenStore):
@@ -440,7 +439,7 @@ class RestZenStore(BaseZenStore):
             Information about the server.
         """
         body = self.get(INFO)
-        return ServerModel.parse_obj(body)
+        return ServerModel.model_validate(body)
 
     def get_deployment_id(self) -> UUID:
         """Get the ID of the deployment.
@@ -465,7 +464,7 @@ class RestZenStore(BaseZenStore):
             The server settings.
         """
         response_body = self.get(SERVER_SETTINGS, params={"hydrate": hydrate})
-        return ServerSettingsResponse.parse_obj(response_body)
+        return ServerSettingsResponse.model_validate(response_body)
 
     def update_server_settings(
         self, settings_update: ServerSettingsUpdate
@@ -479,7 +478,7 @@ class RestZenStore(BaseZenStore):
             The updated server settings.
         """
         response_body = self.put(SERVER_SETTINGS, body=settings_update)
-        return ServerSettingsResponse.parse_obj(response_body)
+        return ServerSettingsResponse.model_validate(response_body)
 
     # ----------------------------- API Keys -----------------------------
 
@@ -612,7 +611,7 @@ class RestZenStore(BaseZenStore):
             f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}/{str(api_key_name_or_id)}{API_KEY_ROTATE}",
             body=rotate_request,
         )
-        return APIKeyResponse.parse_obj(response_body)
+        return APIKeyResponse.model_validate(response_body)
 
     def delete_api_key(
         self,
@@ -1415,9 +1414,34 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_BUILDS,
         )
 
-        # ----------------------
-        # Pipeline Deployments
-        # ----------------------
+    def run_build(
+        self,
+        build_id: UUID,
+        run_configuration: Optional[PipelineRunConfiguration] = None,
+    ) -> PipelineRunResponse:
+        """Run a pipeline from a build.
+
+        Args:
+            build_id: The ID of the build to run.
+            run_configuration: Configuration for the run.
+
+        Raises:
+            RuntimeError: If the server does not support running a build.
+
+        Returns:
+            Model of the pipeline run.
+        """
+        run_configuration = run_configuration or PipelineRunConfiguration()
+        try:
+            response_body = self.post(
+                f"{PIPELINE_BUILDS}/{build_id}/run", body=run_configuration
+            )
+        except MethodNotAllowedError as e:
+            raise RuntimeError(
+                "Running a build is not supported for this server."
+            ) from e
+
+        return PipelineRunResponse.parse_obj(response_body)
 
     # -------------------------- Pipeline Deployments --------------------------
 
@@ -1493,6 +1517,37 @@ class RestZenStore(BaseZenStore):
             route=PIPELINE_DEPLOYMENTS,
         )
 
+    def run_deployment(
+        self,
+        deployment_id: UUID,
+        run_configuration: Optional[PipelineRunConfiguration] = None,
+    ) -> PipelineRunResponse:
+        """Run a pipeline from a deployment.
+
+        Args:
+            deployment_id: The ID of the deployment to run.
+            run_configuration: Configuration for the run.
+
+        Raises:
+            RuntimeError: If the server does not support running a deployment.
+
+        Returns:
+            Model of the pipeline run.
+        """
+        run_configuration = run_configuration or PipelineRunConfiguration()
+
+        try:
+            response_body = self.post(
+                f"{PIPELINE_DEPLOYMENTS}/{deployment_id}/run",
+                body=run_configuration,
+            )
+        except MethodNotAllowedError as e:
+            raise RuntimeError(
+                "Running a deployment is not supported for this server."
+            ) from e
+
+        return PipelineRunResponse.parse_obj(response_body)
+
     # -------------------- Event Sources  --------------------
 
     def create_event_source(
@@ -1721,7 +1776,7 @@ class RestZenStore(BaseZenStore):
         result: List[RunMetadataResponse] = []
         if isinstance(response_body, list):
             for metadata in response_body or []:
-                result.append(RunMetadataResponse.parse_obj(metadata))
+                result.append(RunMetadataResponse.model_validate(metadata))
         return result
 
     def get_run_metadata(
@@ -2315,7 +2370,9 @@ class RestZenStore(BaseZenStore):
             params={"list_resources": list_resources},
         )
 
-        resources = ServiceConnectorResourcesModel.parse_obj(response_body)
+        resources = ServiceConnectorResourcesModel.model_validate(
+            response_body
+        )
         self._populate_connector_type(resources)
         return resources
 
@@ -2350,7 +2407,9 @@ class RestZenStore(BaseZenStore):
             params=params,
         )
 
-        resources = ServiceConnectorResourcesModel.parse_obj(response_body)
+        resources = ServiceConnectorResourcesModel.model_validate(
+            response_body
+        )
         self._populate_connector_type(resources)
         return resources
 
@@ -2381,7 +2440,7 @@ class RestZenStore(BaseZenStore):
             params=params,
         )
 
-        connector = ServiceConnectorResponse.parse_obj(response_body)
+        connector = ServiceConnectorResponse.model_validate(response_body)
         self._populate_connector_type(connector)
         return connector
 
@@ -2418,7 +2477,7 @@ class RestZenStore(BaseZenStore):
 
         assert isinstance(response_body, list)
         resource_list = [
-            ServiceConnectorResourcesModel.parse_obj(item)
+            ServiceConnectorResourcesModel.model_validate(item)
             for item in response_body
         ]
 
@@ -2491,7 +2550,8 @@ class RestZenStore(BaseZenStore):
 
         assert isinstance(response_body, list)
         remote_connector_types = [
-            ServiceConnectorTypeModel.parse_obj(item) for item in response_body
+            ServiceConnectorTypeModel.model_validate(item)
+            for item in response_body
         ]
 
         # Mark the remote connector types as being only remotely available
@@ -2547,7 +2607,7 @@ class RestZenStore(BaseZenStore):
             response_body = self.get(
                 f"{SERVICE_CONNECTOR_TYPES}/{connector_type}",
             )
-            remote_connector_type = ServiceConnectorTypeModel.parse_obj(
+            remote_connector_type = ServiceConnectorTypeModel.model_validate(
                 response_body
             )
             if local_connector_type:
@@ -2935,7 +2995,7 @@ class RestZenStore(BaseZenStore):
             )
         else:
             body = self.get(CURRENT_USER, params={"hydrate": hydrate})
-            return UserResponse.parse_obj(body)
+            return UserResponse.model_validate(body)
 
     def list_users(
         self,
@@ -2994,7 +3054,7 @@ class RestZenStore(BaseZenStore):
             f"{USERS}/{str(user_name_or_id)}{DEACTIVATE}",
         )
 
-        return UserResponse.parse_obj(response_body)
+        return UserResponse.model_validate(response_body)
 
     def delete_user(self, user_name_or_id: Union[str, UUID]) -> None:
         """Deletes a user.
@@ -3924,7 +3984,7 @@ class RestZenStore(BaseZenStore):
         return self._request(
             "POST",
             self.url + API + VERSION_1 + path,
-            data=body.json(),
+            data=body.model_dump_json(),
             params=params,
             **kwargs,
         )
@@ -3948,7 +4008,7 @@ class RestZenStore(BaseZenStore):
             The response body.
         """
         logger.debug(f"Sending PUT request to {path}...")
-        data = body.json(exclude_unset=True) if body else None
+        data = body.model_dump_json(exclude_unset=True) if body else None
         return self._request(
             "PUT",
             self.url + API + VERSION_1 + path,
@@ -3978,7 +4038,7 @@ class RestZenStore(BaseZenStore):
         """
         response_body = self.post(f"{route}", body=resource, params=params)
 
-        return response_model.parse_obj(response_body)
+        return response_model.model_validate(response_body)
 
     def _create_workspace_scoped_resource(
         self,
@@ -4054,7 +4114,7 @@ class RestZenStore(BaseZenStore):
                 f"response from the {route}{GET_OR_CREATE} endpoint but got "
                 f"{type(was_created)} instead."
             )
-        return response_model.parse_obj(model_json), was_created
+        return response_model.model_validate(model_json), was_created
 
     def _get_or_create_workspace_scoped_resource(
         self,
@@ -4102,7 +4162,7 @@ class RestZenStore(BaseZenStore):
             The retrieved resource.
         """
         body = self.get(f"{route}/{str(resource_id)}", params=params)
-        return response_model.parse_obj(body)
+        return response_model.model_validate(body)
 
     def _list_paginated_resources(
         self,
@@ -4127,17 +4187,17 @@ class RestZenStore(BaseZenStore):
         """
         # leave out filter params that are not supplied
         params = params or {}
-        params.update(filter_model.dict(exclude_none=True))
+        params.update(filter_model.model_dump(exclude_none=True))
         body = self.get(f"{route}", params=params)
         if not isinstance(body, dict):
             raise ValueError(
                 f"Bad API Response. Expected list, got {type(body)}"
             )
         # The initial page of items will be of type BaseResponseModel
-        page_of_items: Page[AnyResponse] = Page.parse_obj(body)
+        page_of_items: Page[AnyResponse] = Page.model_validate(body)
         # So these items will be parsed into their correct types like here
         page_of_items.items = [
-            response_model.parse_obj(generic_item)
+            response_model.model_validate(generic_item)
             for generic_item in body["items"]
         ]
         return page_of_items
@@ -4168,7 +4228,7 @@ class RestZenStore(BaseZenStore):
             raise ValueError(
                 f"Bad API Response. Expected list, got {type(body)}"
             )
-        return [response_model.parse_obj(entry) for entry in body]
+        return [response_model.model_validate(entry) for entry in body]
 
     def _update_resource(
         self,
@@ -4195,7 +4255,7 @@ class RestZenStore(BaseZenStore):
             f"{route}/{str(resource_id)}", body=resource_update, params=params
         )
 
-        return response_model.parse_obj(response_body)
+        return response_model.model_validate(response_body)
 
     def _delete_resource(
         self, resource_id: Union[str, UUID], route: str

zenml/zen_stores/schemas/api_key_schemas.py

@@ -188,7 +188,7 @@ class APIKeySchema(NamedSchema, table=True):
         Returns:
             The updated `APIKeySchema`.
         """
-        for field, value in update.dict(exclude_none=True).items():
+        for field, value in update.model_dump(exclude_none=True).items():
             if hasattr(self, field):
                 setattr(self, field, value)
 

zenml/zen_stores/schemas/artifact_schemas.py

@@ -171,7 +171,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
     # Fields
     version: str
     version_number: Optional[int]
-    type: ArtifactType
+    type: str
     uri: str = Field(sa_column=Column(TEXT, nullable=False))
     materializer: str = Field(sa_column=Column(TEXT, nullable=False))
     data_type: str = Field(sa_column=Column(TEXT, nullable=False))
@@ -277,10 +277,10 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             artifact_store_id=artifact_version_request.artifact_store_id,
             workspace_id=artifact_version_request.workspace,
             user_id=artifact_version_request.user,
-            type=artifact_version_request.type,
+            type=artifact_version_request.type.value,
             uri=artifact_version_request.uri,
-            materializer=artifact_version_request.materializer.json(),
-            data_type=artifact_version_request.data_type.json(),
+            materializer=artifact_version_request.materializer.model_dump_json(),
+            data_type=artifact_version_request.data_type.model_dump_json(),
         )
 
     def to_model(
@@ -302,13 +302,13 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             The created `ArtifactVersionResponse`.
         """
         try:
-            materializer = Source.parse_raw(self.materializer)
+            materializer = Source.model_validate_json(self.materializer)
         except ValidationError:
             # This is an old source which was an importable source path
             materializer = Source.from_import_path(self.materializer)
 
         try:
-            data_type = Source.parse_raw(self.data_type)
+            data_type = Source.model_validate_json(self.data_type)
         except ValidationError:
             # This is an old source which was an importable source path
             data_type = Source.from_import_path(self.data_type)
@@ -331,10 +331,10 @@ class ArtifactVersionSchema(BaseSchema, table=True):
         # Create the body of the model
         body = ArtifactVersionResponseBody(
             artifact=self.artifact.to_model(),
-            version=self.version_number or self.version,
+            version=self.version or str(self.version_number),
             user=self.user.to_model() if self.user else None,
             uri=self.uri,
-            type=self.type,
+            type=ArtifactType(self.type),
             materializer=materializer,
             data_type=data_type,
             created=self.created,

zenml/zen_stores/schemas/artifact_visualization_schemas.py

@@ -37,7 +37,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True):
     __tablename__ = "artifact_visualization"
 
     # Fields
-    type: VisualizationType
+    type: str
     uri: str = Field(sa_column=Column(TEXT, nullable=False))
 
     # Foreign Keys
@@ -71,7 +71,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True):
             The `ArtifactVisualizationSchema`.
         """
         return cls(
-            type=artifact_visualization_request.type,
+            type=artifact_visualization_request.type.value,
             uri=artifact_visualization_request.uri,
             artifact_version_id=artifact_version_id,
         )
@@ -95,7 +95,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True):
             The `Visualization`.
         """
         body = ArtifactVisualizationResponseBody(
-            type=self.type,
+            type=VisualizationType(self.type),
             uri=self.uri,
             created=self.created,
             updated=self.updated,