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

zenml/models/v2/core/user.py

@@ -16,6 +16,7 @@
 from secrets import token_hex
 from typing import (
     TYPE_CHECKING,
+    AbstractSet,
     Any,
     ClassVar,
     Dict,
@@ -26,48 +27,32 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import Field, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
+    BaseZenModel,
 )
 from zenml.models.v2.base.filter import AnyQuery, BaseFilter
-from zenml.models.v2.base.update import update_model
 
 if TYPE_CHECKING:
     from passlib.context import CryptContext
 
     from zenml.models.v2.base.filter import AnySchema
 
-# ------------------ Request Model ------------------
-
+# ------------------ Base Model ------------------
 
-class UserRequest(BaseRequest):
-    """Request model for users."""
 
-    # Analytics fields for user request models
-    ANALYTICS_FIELDS: ClassVar[List[str]] = [
-        "name",
-        "full_name",
-        "active",
-        "email_opted_in",
-    ]
+class UserBase(BaseModel):
+    """Base model for users."""
 
     # Fields
-    name: str = Field(
-        title="The unique username for the account.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
-    full_name: str = Field(
-        default="",
-        title="The full name for the account owner. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
+
     email: Optional[str] = Field(
         default=None,
         title="The email address associated with the account.",
@@ -80,12 +65,6 @@ class UserRequest(BaseRequest):
         description="`null` if not answered, `true` if agreed, "
         "`false` if skipped.",
     )
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     password: Optional[str] = Field(
         default=None,
         title="A password for the user.",
@@ -98,17 +77,10 @@ class UserRequest(BaseRequest):
         default=None,
         title="The external user ID associated with the account.",
     )
-    active: bool = Field(default=False, title="Whether the account is active.")
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Validate attributes when assigning them
-        validate_assignment = True
-
-        # Forbid extra attributes to prevent unexpected behavior
-        extra = "forbid"
-        underscore_attrs_are_private = True
+    user_metadata: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The metadata associated with the user.",
+    )
 
     @classmethod
     def _get_crypt_context(cls) -> "CryptContext":
@@ -164,19 +136,77 @@ class UserRequest(BaseRequest):
         return self.activation_token
 
 
+# ------------------ Request Model ------------------
+
+
+class UserRequest(UserBase, BaseRequest):
+    """Request model for users."""
+
+    # Analytics fields for user request models
+    ANALYTICS_FIELDS: ClassVar[List[str]] = [
+        "name",
+        "full_name",
+        "active",
+        "email_opted_in",
+    ]
+
+    name: str = Field(
+        title="The unique username for the account.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    full_name: str = Field(
+        default="",
+        title="The full name for the account owner. Only relevant for user "
+        "accounts.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    is_admin: bool = Field(
+        title="Whether the account is an administrator.",
+    )
+    active: bool = Field(default=False, title="Whether the account is active.")
+
+    model_config = ConfigDict(
+        # Validate attributes when assigning them
+        validate_assignment=True,
+        # Forbid extra attributes to prevent unexpected behavior
+        extra="forbid",
+    )
+
+
 # ------------------ Update Model ------------------
 
 
-@update_model
-class UserUpdate(UserRequest):
+class UserUpdate(UserBase, BaseZenModel):
     """Update model for users."""
 
-    @root_validator
-    def user_email_updates(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        """Validate that the UserUpdateModel conforms to the email-opt-in-flow.
+    name: Optional[str] = Field(
+        title="The unique username for the account.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    full_name: Optional[str] = Field(
+        default=None,
+        title="The full name for the account owner. Only relevant for user "
+        "accounts.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    is_admin: Optional[bool] = Field(
+        default=None,
+        title="Whether the account is an administrator.",
+    )
+    active: Optional[bool] = Field(
+        default=None, title="Whether the account is active."
+    )
+    old_password: Optional[str] = Field(
+        default=None,
+        title="The previous password for the user. Only relevant for user "
+        "accounts. Required when updating the password.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
 
-        Args:
-            values: The values to validate.
+    @model_validator(mode="after")
+    def user_email_updates(self) -> "UserUpdate":
+        """Validate that the UserUpdateModel conforms to the email-opt-in-flow.
 
         Returns:
             The validated values.
@@ -187,24 +217,40 @@ class UserUpdate(UserRequest):
         """
         # When someone sets the email, or updates the email and hasn't
         #  before explicitly opted out, they are opted in
-        if values["email"] is not None:
-            if values["email_opted_in"] is None:
-                values["email_opted_in"] = True
+        if self.email is not None:
+            if self.email_opted_in is None:
+                self.email_opted_in = True
 
         # It should not be possible to do opt in without an email
-        if values["email_opted_in"] is True:
-            if values["email"] is None:
+        if self.email_opted_in is True:
+            if self.email is None:
                 raise ValueError(
                     "Please provide an email, when you are opting-in with "
                     "your email."
                 )
-        return values
+        return self
+
+    def create_copy(self, exclude: AbstractSet[str]) -> "UserUpdate":
+        """Create a copy of the current instance.
+
+        Args:
+            exclude: Fields to exclude from the copy.
+
+        Returns:
+            A copy of the current instance.
+        """
+        return UserUpdate(
+            **self.model_dump(
+                exclude=set(exclude),
+                exclude_unset=True,
+            )
+        )
 
 
 # ------------------ Response Model ------------------
 
 
-class UserResponseBody(BaseResponseBody):
+class UserResponseBody(BaseDatedResponseBody):
     """Response body for users."""
 
     active: bool = Field(default=False, title="Whether the account is active.")
@@ -230,6 +276,9 @@ class UserResponseBody(BaseResponseBody):
     is_service_account: bool = Field(
         title="Indicates whether this is a service account or a user account."
     )
+    is_admin: bool = Field(
+        title="Whether the account is an administrator.",
+    )
 
 
 class UserResponseMetadata(BaseResponseMetadata):
@@ -241,20 +290,26 @@ class UserResponseMetadata(BaseResponseMetadata):
         "for user accounts.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     external_user_id: Optional[UUID] = Field(
         default=None,
         title="The external user ID associated with the account. Only relevant "
         "for user accounts.",
     )
+    user_metadata: Dict[str, Any] = Field(
+        default={},
+        title="The metadata associated with the user.",
+    )
+
 
+class UserResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the user entity."""
 
-class UserResponse(BaseResponse[UserResponseBody, UserResponseMetadata]):
+
+class UserResponse(
+    BaseIdentifiedResponse[
+        UserResponseBody, UserResponseMetadata, UserResponseResources
+    ]
+):
     """Response model for user and service accounts.
 
     This returns the activation_token that is required for the
@@ -331,6 +386,15 @@ class UserResponse(BaseResponse[UserResponseBody, UserResponseMetadata]):
         """
         return self.get_body().is_service_account
 
+    @property
+    def is_admin(self) -> bool:
+        """The `is_admin` property.
+
+        Returns:
+            Whether the user is an admin.
+        """
+        return self.get_body().is_admin
+
     @property
     def email(self) -> Optional[str]:
         """The `email` property.
@@ -341,22 +405,22 @@ class UserResponse(BaseResponse[UserResponseBody, UserResponseMetadata]):
         return self.get_metadata().email
 
     @property
-    def hub_token(self) -> Optional[str]:
-        """The `hub_token` property.
+    def external_user_id(self) -> Optional[UUID]:
+        """The `external_user_id` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().hub_token
+        return self.get_metadata().external_user_id
 
     @property
-    def external_user_id(self) -> Optional[UUID]:
-        """The `external_user_id` property.
+    def user_metadata(self) -> Dict[str, Any]:
+        """The `user_metadata` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().external_user_id
+        return self.get_metadata().user_metadata
 
     # Helper methods
     @classmethod
@@ -392,14 +456,17 @@ class UserFilter(BaseFilter):
     active: Optional[Union[bool, str]] = Field(
         default=None,
         description="Whether the user is active",
+        union_mode="left_to_right",
     )
     email_opted_in: Optional[Union[bool, str]] = Field(
         default=None,
         description="Whether the user has opted in to emails",
+        union_mode="left_to_right",
     )
     external_user_id: Optional[Union[UUID, str]] = Field(
         default=None,
         title="The external user ID associated with the account.",
+        union_mode="left_to_right",
     )
 
     def apply_filter(

zenml/models/v2/core/workspace.py

@@ -19,13 +19,14 @@ from pydantic import Field
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
+    BaseUpdate,
 )
 from zenml.models.v2.base.filter import BaseFilter
-from zenml.models.v2.base.update import update_model
 
 # ------------------ Request Model ------------------
 
@@ -47,15 +48,25 @@ class WorkspaceRequest(BaseRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class WorkspaceUpdate(WorkspaceRequest):
+class WorkspaceUpdate(BaseUpdate):
     """Update model for workspaces."""
 
+    name: Optional[str] = Field(
+        title="The unique name of the workspace.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    description: Optional[str] = Field(
+        title="The description of the workspace.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 
 
-class WorkspaceResponseBody(BaseResponseBody):
+class WorkspaceResponseBody(BaseDatedResponseBody):
     """Response body for workspaces."""
 
 
@@ -69,8 +80,16 @@ class WorkspaceResponseMetadata(BaseResponseMetadata):
     )
 
 
+class WorkspaceResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the workspace entity."""
+
+
 class WorkspaceResponse(
-    BaseResponse[WorkspaceResponseBody, WorkspaceResponseMetadata]
+    BaseIdentifiedResponse[
+        WorkspaceResponseBody,
+        WorkspaceResponseMetadata,
+        WorkspaceResponseResources,
+    ]
 ):
     """Response model for workspaces."""
 

zenml/models/v2/misc/auth_models.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """Models representing OAuth2 requests and responses."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import OAuthGrantTypes
 
@@ -29,6 +29,7 @@ class OAuthDeviceAuthorizationRequest(BaseModel):
     """OAuth2 device authorization grant request."""
 
     client_id: UUID
+    device_id: Optional[UUID] = None
 
 
 class OAuthDeviceVerificationRequest(BaseModel):
@@ -119,6 +120,14 @@ class OAuthTokenResponse(BaseModel):
     expires_in: Optional[int] = None
     refresh_token: Optional[str] = None
     scope: Optional[str] = None
+    cookie_name: Optional[str] = None
+    device_id: Optional[UUID] = None
+    device_metadata: Optional[Dict[str, Any]] = None
+
+    model_config = ConfigDict(
+        # Allow extra attributes to allow compatibility with different versions
+        extra="allow",
+    )
 
 
 class OAuthRedirectResponse(BaseModel):

zenml/models/v2/misc/build_item.py

@@ -33,13 +33,13 @@ class BuildItem(BaseModel):
 
     image: str = Field(title="The image name or digest.")
     dockerfile: Optional[str] = Field(
-        title="The dockerfile used to build the image."
+        default=None, title="The dockerfile used to build the image."
     )
     requirements: Optional[str] = Field(
-        title="The pip requirements installed in the image."
+        default=None, title="The pip requirements installed in the image."
     )
     settings_checksum: Optional[str] = Field(
-        title="The checksum of the build settings."
+        default=None, title="The checksum of the build settings."
     )
     contains_code: bool = Field(
         default=True, title="Whether the image contains user files."

zenml/models/v2/misc/external_user.py

@@ -16,7 +16,7 @@
 from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ExternalUserModel(BaseModel):
@@ -25,9 +25,6 @@ class ExternalUserModel(BaseModel):
     id: UUID
     email: str
     name: Optional[str] = None
+    is_admin: bool = False
 
-    class Config:
-        """Pydantic configuration."""
-
-        # ignore arbitrary fields
-        extra = "ignore"
+    model_config = ConfigDict(extra="ignore")

zenml/models/v2/misc/info_models.py

@@ -0,0 +1,78 @@
+#  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.
+"""Models representing full stack requests."""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+from zenml.enums import StackComponentType
+
+if TYPE_CHECKING:
+    from zenml.models.v2.core.component import ComponentResponse
+
+
+class ServiceConnectorInfo(BaseModel):
+    """Information about the service connector when creating a full stack."""
+
+    type: str
+    auth_method: str
+    configuration: Dict[str, Any] = {}
+
+
+class ComponentInfo(BaseModel):
+    """Information about each stack components when creating a full stack."""
+
+    flavor: str
+    service_connector_index: Optional[int] = Field(
+        default=None,
+        title="The id of the service connector from the list "
+        "`service_connectors`.",
+        description="The id of the service connector from the list "
+        "`service_connectors` from `FullStackRequest`.",
+    )
+    service_connector_resource_id: Optional[str] = None
+    configuration: Dict[str, Any] = {}
+
+
+class ResourcesInfo(BaseModel):
+    """Information about the resources needed for CLI and UI."""
+
+    flavor: str
+    flavor_display_name: str
+    required_configuration: Dict[str, str] = {}
+    use_resource_value_as_fixed_config: bool = False
+
+    accessible_by_service_connector: List[str]
+    connected_through_service_connector: List["ComponentResponse"]
+
+    @model_validator(mode="after")
+    def _validate_resource_info(self) -> "ResourcesInfo":
+        if (
+            self.use_resource_value_as_fixed_config
+            and len(self.required_configuration) > 1
+        ):
+            raise ValueError(
+                "Cannot use resource value as fixed config if more than "
+                "one required configuration key is provided."
+            )
+        return self
+
+
+class ServiceConnectorResourcesInfo(BaseModel):
+    """Information about the service connector resources needed for CLI and UI."""
+
+    connector_type: str
+
+    components_resources_info: Dict[StackComponentType, List[ResourcesInfo]]

zenml/models/v2/misc/loaded_visualization.py

@@ -15,7 +15,7 @@
 
 from typing import Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from zenml.enums import VisualizationType
 
@@ -24,4 +24,4 @@ class LoadedVisualization(BaseModel):
     """Model for loaded visualizations."""
 
     type: VisualizationType
-    value: Union[str, bytes]
+    value: Union[str, bytes] = Field(union_mode="left_to_right")

zenml/models/v2/misc/run_metadata.py

@@ -0,0 +1,38 @@
+#  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.
+"""Utility classes for modeling run metadata."""
+
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import MetadataResourceTypes
+from zenml.metadata.metadata_types import MetadataType
+
+
+class RunMetadataResource(BaseModel):
+    """Utility class to help identify resources to tag metadata to."""
+
+    id: UUID = Field(title="The ID of the resource.")
+    type: MetadataResourceTypes = Field(title="The type of the resource.")
+
+
+class RunMetadataEntry(BaseModel):
+    """Utility class to sort/list run metadata entries."""
+
+    value: MetadataType = Field(title="The value for the run metadata entry")
+    created: datetime = Field(
+        title="The timestamp when this resource was created."
+    )

zenml/models/v2/misc/server_models.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Model definitions for ZenML servers."""
 
+from datetime import datetime
+from typing import Dict, Optional
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field
@@ -50,10 +52,16 @@ class ServerModel(BaseModel):
 
     id: UUID = Field(default_factory=uuid4, title="The unique server id.")
 
+    name: Optional[str] = Field(None, title="The name of the ZenML server.")
+
     version: str = Field(
         title="The ZenML version that the server is running.",
     )
 
+    active: bool = Field(
+        True, title="Flag to indicate whether the server is active."
+    )
+
     debug: bool = Field(
         False, title="Flag to indicate whether ZenML is running on debug mode."
     )
@@ -73,6 +81,31 @@ class ServerModel(BaseModel):
     auth_scheme: AuthScheme = Field(
         title="The authentication scheme that the server is using.",
     )
+    server_url: str = Field(
+        "",
+        title="The URL where the ZenML server API is reachable. If not "
+        "specified, the clients will use the same URL used to connect them to "
+        "the ZenML server.",
+    )
+    dashboard_url: str = Field(
+        "",
+        title="The URL where the ZenML dashboard is reachable. If "
+        "not specified, the `server_url` value will be used instead.",
+    )
+    analytics_enabled: bool = Field(
+        default=True,  # We set a default for migrations from < 0.57.0
+        title="Enable server-side analytics.",
+    )
+
+    metadata: Dict[str, str] = Field(
+        {},
+        title="The metadata associated with the server.",
+    )
+
+    last_user_activity: Optional[datetime] = Field(
+        None,
+        title="Timestamp of latest user activity traced on the server.",
+    )
 
     def is_local(self) -> bool:
         """Return whether the server is running locally.
@@ -85,3 +118,26 @@ class ServerModel(BaseModel):
         # Local ZenML servers are identifiable by the fact that their
         # server ID is the same as the local client (user) ID.
         return self.id == GlobalConfiguration().user_id
+
+
+class ServerLoadInfo(BaseModel):
+    """Domain model for ZenML server load information."""
+
+    threads: int = Field(
+        title="Number of threads that the server is currently using."
+    )
+
+    db_connections_total: int = Field(
+        title="Total number of database connections (active and idle) that the "
+        "server currently has established."
+    )
+
+    db_connections_active: int = Field(
+        title="Number of database connections that the server is currently "
+        "actively using to make queries or transactions."
+    )
+
+    db_connections_overflow: int = Field(
+        title="Number of overflow database connections that the server is "
+        "currently actively using to make queries or transactions."
+    )