zenml-nightly 0.75.0.dev20250317__py3-none-any.whl → 0.80.0.dev20250321__py3-none-any.whl

zenml/models/v2/base/filter.py

@@ -37,7 +37,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from sqlalchemy import Float, and_, asc, cast, desc
+from sqlalchemy import Float, and_, asc, cast, desc, or_
 from sqlmodel import SQLModel
 
 from zenml.constants import (
@@ -171,8 +171,6 @@ class BoolFilter(Filter):
 class StrFilter(Filter):
     """Filter for all string fields."""
 
-    json_encode_value: bool = False
-
     ALLOWED_OPS: ClassVar[List[str]] = [
         GenericFilterOps.EQUALS,
         GenericFilterOps.NOT_EQUALS,
@@ -201,6 +199,23 @@ class StrFilter(Filter):
                 raise ValueError(ONEOF_ERROR)
         return self
 
+    def _check_if_column_is_json_encoded(self, column: Any) -> bool:
+        """Check if the column is json encoded.
+
+        Args:
+            column: The column of an SQLModel table on which to filter.
+
+        Returns:
+            True if the column is json encoded, False otherwise.
+        """
+        from zenml.zen_stores.schemas import RunMetadataSchema
+
+        JSON_ENCODED_COLUMNS = [RunMetadataSchema.value]
+
+        if column in JSON_ENCODED_COLUMNS:
+            return True
+        return False
+
     def generate_query_conditions_from_column(self, column: Any) -> Any:
         """Generate query conditions for a string column.
 
@@ -209,70 +224,168 @@ class StrFilter(Filter):
 
         Returns:
             A list of query conditions.
-
-        Raises:
-            ValueError: the comparison of the column to a numeric value fails.
         """
+        # Handle numeric comparisons (GT, LT, GTE, LTE)
         if self.operation in {
             GenericFilterOps.GT,
             GenericFilterOps.LT,
             GenericFilterOps.GTE,
             GenericFilterOps.LTE,
         }:
-            try:
-                numeric_column = cast(column, Float)
-
-                assert self.value is not None
-
-                if self.operation == GenericFilterOps.GT:
-                    return and_(
-                        numeric_column, numeric_column > float(self.value)
-                    )
-                if self.operation == GenericFilterOps.LT:
-                    return and_(
-                        numeric_column, numeric_column < float(self.value)
-                    )
-                if self.operation == GenericFilterOps.GTE:
-                    return and_(
-                        numeric_column, numeric_column >= float(self.value)
-                    )
-                if self.operation == GenericFilterOps.LTE:
-                    return and_(
-                        numeric_column, numeric_column <= float(self.value)
-                    )
-            except Exception as e:
-                raise ValueError(
-                    f"Failed to compare the column '{column}' to the "
-                    f"value '{self.value}' (must be numeric): {e}"
-                )
+            return self._handle_numeric_comparison(column)
+
+        # Handle operations that need special treatment for JSON-encoded columns
+        is_json_encoded = self._check_if_column_is_json_encoded(column)
 
+        # Handle list operations
         if self.operation == GenericFilterOps.ONEOF:
             assert isinstance(self.value, list)
-            # Convert the list of values to a list of json strings
-            json_list = (
-                [json.dumps(v) for v in self.value]
-                if self.json_encode_value
-                else self.value
-            )
-            return column.in_(json_list)
+            return self._handle_oneof(column, is_json_encoded)
 
-        # Don't convert the value to a json string if the operation is contains
-        # because the quotes around strings will mess with the comparison
+        # Handle pattern matching operations
         if self.operation == GenericFilterOps.CONTAINS:
             return column.like(f"%{self.value}%")
 
-        json_value = (
-            json.dumps(self.value) if self.json_encode_value else self.value
-        )
-
         if self.operation == GenericFilterOps.STARTSWITH:
-            return column.startswith(f"{json_value}")
+            return self._handle_startswith(column, is_json_encoded)
+
         if self.operation == GenericFilterOps.ENDSWITH:
-            return column.endswith(f"{json_value}")
+            return self._handle_endswith(column, is_json_encoded)
+
         if self.operation == GenericFilterOps.NOT_EQUALS:
-            return column != json_value
+            return self._handle_not_equals(column, is_json_encoded)
 
-        return column == json_value
+        # Default case (EQUALS)
+        return self._handle_equals(column, is_json_encoded)
+
+    def _handle_numeric_comparison(self, column: Any) -> Any:
+        """Handle numeric comparison operations.
+
+        Args:
+            column: The column to compare.
+
+        Returns:
+            The query condition.
+
+        Raises:
+            ValueError: If the comparison fails.
+        """
+        try:
+            numeric_column = cast(column, Float)
+            assert self.value is not None
+
+            operations = {
+                GenericFilterOps.GT: lambda col, val: and_(
+                    col, col > float(val)
+                ),
+                GenericFilterOps.LT: lambda col, val: and_(
+                    col, col < float(val)
+                ),
+                GenericFilterOps.GTE: lambda col, val: and_(
+                    col, col >= float(val)
+                ),
+                GenericFilterOps.LTE: lambda col, val: and_(
+                    col, col <= float(val)
+                ),
+            }
+
+            return operations[self.operation](numeric_column, self.value)  # type: ignore[no-untyped-call]
+        except Exception as e:
+            raise ValueError(
+                f"Failed to compare the column '{column}' to the "
+                f"value '{self.value}' (must be numeric): {e}"
+            )
+
+    def _handle_oneof(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the ONEOF operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        from sqlalchemy import or_
+
+        conditions = []
+
+        assert isinstance(self.value, list)
+
+        for value in self.value:
+            if is_json_encoded:
+                # For JSON encoded columns, add conditions for both raw and JSON-quoted values
+                conditions.append(column == value)
+                conditions.append(column == f'"{value}"')
+            else:
+                conditions.append(column == value)
+
+        return or_(*conditions)
+
+    def _handle_startswith(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the STARTSWITH operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(
+                column.startswith(self.value),
+                column.startswith(f'"{self.value}'),
+            )
+        else:
+            return column.startswith(self.value)
+
+    def _handle_endswith(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the ENDSWITH operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(
+                column.endswith(self.value), column.endswith(f'{self.value}"')
+            )
+        else:
+            return column.endswith(self.value)
+
+    def _handle_not_equals(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the NOT_EQUALS operation.
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return and_(column != self.value, column != f'"{self.value}"')
+        else:
+            return column != self.value
+
+    def _handle_equals(self, column: Any, is_json_encoded: bool) -> Any:
+        """Handle the EQUALS operation (default).
+
+        Args:
+            column: The column to check.
+            is_json_encoded: Whether the column is JSON encoded.
+
+        Returns:
+            The query condition.
+        """
+        if is_json_encoded:
+            return or_(column == self.value, column == f'"{self.value}"')
+        else:
+            return column == self.value
 
 
 class UUIDFilter(StrFilter):
@@ -751,7 +864,6 @@ class BaseFilter(BaseModel):
         value: Any,
         table: Type[SQLModel],
         column: str,
-        json_encode_value: bool = False,
     ) -> "ColumnElement[bool]":
         """Generate custom filter conditions for a column of a table.
 
@@ -759,7 +871,6 @@ class BaseFilter(BaseModel):
             value: The filter value.
             table: The table which contains the column.
             column: The column name.
-            json_encode_value: Whether to json encode the value.
 
         Returns:
             The query conditions.
@@ -768,9 +879,6 @@ class BaseFilter(BaseModel):
         filter_ = FilterGenerator(table).define_filter(
             column=column, value=value, operator=operator
         )
-        if isinstance(filter_, StrFilter):
-            filter_.json_encode_value = json_encode_value
-
         return filter_.generate_query_conditions(table=table)
 
     @property

zenml/models/v2/base/scoped.py

@@ -426,6 +426,10 @@ class TaggableFilter(BaseFilter):
         *BaseFilter.CUSTOM_SORTING_OPTIONS,
         "tags",
     ]
+    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
+        *BaseFilter.API_MULTI_INPUT_PARAMS,
+        "tags",
+    ]
 
     @model_validator(mode="after")
     def add_tag_to_tags(self) -> "TaggableFilter":
@@ -583,3 +587,131 @@ class TaggableFilter(BaseFilter):
             return query
 
         return super().apply_sorting(query=query, table=table)
+
+
+class RunMetadataFilterMixin(BaseFilter):
+    """Model to enable filtering and sorting by run metadata."""
+
+    run_metadata: Optional[List[str]] = Field(
+        default=None,
+        description="The run_metadata to filter the pipeline runs by.",
+    )
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *BaseFilter.FILTER_EXCLUDE_FIELDS,
+        "run_metadata",
+    ]
+    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
+        *BaseFilter.API_MULTI_INPUT_PARAMS,
+        "run_metadata",
+    ]
+
+    @model_validator(mode="after")
+    def validate_run_metadata_format(self) -> "RunMetadataFilterMixin":
+        """Validates that run_metadata entries are in the correct format.
+
+        Each run_metadata entry must be in one of the following formats:
+        1. "key:value" - Direct equality comparison (key equals value)
+        2. "key:filterop:value" - Where filterop is one of the GenericFilterOps:
+           - equals: Exact match
+           - notequals: Not equal to
+           - contains: String contains value
+           - startswith: String starts with value
+           - endswith: String ends with value
+           - oneof: Value is one of the specified options
+           - gte: Greater than or equal to
+           - gt: Greater than
+           - lte: Less than or equal to
+           - lt: Less than
+           - in: Value is in a list
+
+        Examples:
+        - "status:completed" - Find entries where status equals "completed"
+        - "name:contains:test" - Find entries where name contains "test"
+        - "duration:gt:10" - Find entries where duration is greater than 10
+
+        Returns:
+            self
+
+        Raises:
+            ValueError: If any entry in run_metadata does not contain a colon.
+        """
+        if self.run_metadata:
+            for entry in self.run_metadata:
+                if ":" not in entry:
+                    raise ValueError(
+                        f"Invalid run_metadata entry format: '{entry}'. "
+                        "Entry must be in format 'key:value' for direct "
+                        "equality comparison or 'key:filterop:value' where "
+                        "filterop is one of: equals, notequals, "
+                        f"contains, startswith, endswith, oneof, gte, gt, "
+                        f"lte, lt, in."
+                    )
+        return self
+
+    def get_custom_filters(
+        self, table: Type["AnySchema"]
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom run metadata filters.
+
+        Args:
+            table: The query table.
+
+        Returns:
+            A list of custom filters.
+        """
+        custom_filters = super().get_custom_filters(table)
+
+        if self.run_metadata is not None:
+            from sqlmodel import exists, select
+
+            from zenml.enums import MetadataResourceTypes
+            from zenml.zen_stores.schemas import (
+                ArtifactVersionSchema,
+                ModelVersionSchema,
+                PipelineRunSchema,
+                RunMetadataResourceSchema,
+                RunMetadataSchema,
+                ScheduleSchema,
+                StepRunSchema,
+            )
+
+            resource_type_mapping = {
+                ArtifactVersionSchema: MetadataResourceTypes.ARTIFACT_VERSION,
+                ModelVersionSchema: MetadataResourceTypes.MODEL_VERSION,
+                PipelineRunSchema: MetadataResourceTypes.PIPELINE_RUN,
+                StepRunSchema: MetadataResourceTypes.STEP_RUN,
+                ScheduleSchema: MetadataResourceTypes.SCHEDULE,
+            }
+
+            # Create an EXISTS subquery for each run_metadata filter
+            for entry in self.run_metadata:
+                # Split at the first colon to get the key
+                key, value = entry.split(":", 1)
+
+                # Create an exists subquery
+                exists_subquery = exists(
+                    select(RunMetadataResourceSchema.id)
+                    .join(
+                        RunMetadataSchema,
+                        RunMetadataSchema.id  # type: ignore[arg-type]
+                        == RunMetadataResourceSchema.run_metadata_id,
+                    )
+                    .where(
+                        RunMetadataResourceSchema.resource_id == table.id,
+                        RunMetadataResourceSchema.resource_type
+                        == resource_type_mapping[table].value,
+                        self.generate_custom_query_conditions_for_column(
+                            value=key,
+                            table=RunMetadataSchema,
+                            column="key",
+                        ),
+                        self.generate_custom_query_conditions_for_column(
+                            value=value,
+                            table=RunMetadataSchema,
+                            column="value",
+                        ),
+                    )
+                )
+                custom_filters.append(exists_subquery)
+
+        return custom_filters

zenml/models/v2/core/artifact_version.py

@@ -47,6 +47,7 @@ from zenml.models.v2.base.scoped import (
     ProjectScopedResponseBody,
     ProjectScopedResponseMetadata,
     ProjectScopedResponseResources,
+    RunMetadataFilterMixin,
     TaggableFilter,
 )
 from zenml.models.v2.core.artifact import ArtifactResponse
@@ -470,12 +471,15 @@ class ArtifactVersionResponse(
 # ------------------ Filter Model ------------------
 
 
-class ArtifactVersionFilter(ProjectScopedFilter, TaggableFilter):
+class ArtifactVersionFilter(
+    ProjectScopedFilter, TaggableFilter, RunMetadataFilterMixin
+):
     """Model to enable advanced filtering of artifact versions."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
         *TaggableFilter.FILTER_EXCLUDE_FIELDS,
+        *RunMetadataFilterMixin.FILTER_EXCLUDE_FIELDS,
         "artifact_id",
         "artifact",
         "only_unused",
@@ -483,17 +487,23 @@ class ArtifactVersionFilter(ProjectScopedFilter, TaggableFilter):
         "model",
         "pipeline_run",
         "model_version_id",
-        "run_metadata",
     ]
     CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
         *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
         *TaggableFilter.CUSTOM_SORTING_OPTIONS,
+        *RunMetadataFilterMixin.CUSTOM_SORTING_OPTIONS,
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
         *TaggableFilter.CLI_EXCLUDE_FIELDS,
+        *RunMetadataFilterMixin.CLI_EXCLUDE_FIELDS,
         "artifact_id",
     ]
+    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
+        *ProjectScopedFilter.API_MULTI_INPUT_PARAMS,
+        *TaggableFilter.API_MULTI_INPUT_PARAMS,
+        *RunMetadataFilterMixin.API_MULTI_INPUT_PARAMS,
+    ]
 
     artifact: Optional[Union[UUID, str]] = Field(
         default=None,
@@ -561,10 +571,6 @@ class ArtifactVersionFilter(ProjectScopedFilter, TaggableFilter):
         description="Name/ID of a pipeline run that is associated with this "
         "artifact version.",
     )
-    run_metadata: Optional[Dict[str, Any]] = Field(
-        default=None,
-        description="The run_metadata to filter the artifact versions by.",
-    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -590,8 +596,6 @@ class ArtifactVersionFilter(ProjectScopedFilter, TaggableFilter):
             ModelVersionArtifactSchema,
             ModelVersionSchema,
             PipelineRunSchema,
-            RunMetadataResourceSchema,
-            RunMetadataSchema,
             StepRunInputArtifactSchema,
             StepRunOutputArtifactSchema,
             StepRunSchema,
@@ -674,31 +678,6 @@ class ArtifactVersionFilter(ProjectScopedFilter, TaggableFilter):
             )
             custom_filters.append(pipeline_run_filter)
 
-        if self.run_metadata is not None:
-            from zenml.enums import MetadataResourceTypes
-
-            for key, value in self.run_metadata.items():
-                additional_filter = and_(
-                    RunMetadataResourceSchema.resource_id
-                    == ArtifactVersionSchema.id,
-                    RunMetadataResourceSchema.resource_type
-                    == MetadataResourceTypes.ARTIFACT_VERSION.value,
-                    RunMetadataResourceSchema.run_metadata_id
-                    == RunMetadataSchema.id,
-                    self.generate_custom_query_conditions_for_column(
-                        value=key,
-                        table=RunMetadataSchema,
-                        column="key",
-                    ),
-                    self.generate_custom_query_conditions_for_column(
-                        value=value,
-                        table=RunMetadataSchema,
-                        column="value",
-                        json_encode_value=True,
-                    ),
-                )
-                custom_filters.append(additional_filter)
-
         return custom_filters
 
     @model_validator(mode="after")

zenml/models/v2/core/model_version.py

@@ -40,14 +40,13 @@ from zenml.models.v2.base.scoped import (
     ProjectScopedResponseBody,
     ProjectScopedResponseMetadata,
     ProjectScopedResponseResources,
+    RunMetadataFilterMixin,
     TaggableFilter,
 )
 from zenml.models.v2.core.service import ServiceResponse
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import ColumnElement
-
     from zenml.model.model import Model
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.model import ModelResponse
@@ -566,24 +565,33 @@ class ModelVersionResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionFilter(ProjectScopedFilter, TaggableFilter):
+class ModelVersionFilter(
+    ProjectScopedFilter, TaggableFilter, RunMetadataFilterMixin
+):
     """Filter model for model versions."""
 
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
         *TaggableFilter.FILTER_EXCLUDE_FIELDS,
+        *RunMetadataFilterMixin.FILTER_EXCLUDE_FIELDS,
         "model",
-        "run_metadata",
     ]
     CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
         *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
         *TaggableFilter.CUSTOM_SORTING_OPTIONS,
+        *RunMetadataFilterMixin.CUSTOM_SORTING_OPTIONS,
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
         *TaggableFilter.CLI_EXCLUDE_FIELDS,
+        *RunMetadataFilterMixin.CLI_EXCLUDE_FIELDS,
         "model",
     ]
+    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
+        *ProjectScopedFilter.API_MULTI_INPUT_PARAMS,
+        *TaggableFilter.API_MULTI_INPUT_PARAMS,
+        *RunMetadataFilterMixin.API_MULTI_INPUT_PARAMS,
+    ]
 
     name: Optional[str] = Field(
         default=None,
@@ -598,10 +606,6 @@ class ModelVersionFilter(ProjectScopedFilter, TaggableFilter):
         default=None,
         union_mode="left_to_right",
     )
-    run_metadata: Optional[Dict[str, str]] = Field(
-        default=None,
-        description="The run_metadata to filter the model versions by.",
-    )
     model: Optional[Union[str, UUID]] = Field(
         default=None,
         description="The name or ID of the model which the search is scoped "
@@ -611,48 +615,6 @@ class ModelVersionFilter(ProjectScopedFilter, TaggableFilter):
         union_mode="left_to_right",
     )
 
-    def get_custom_filters(
-        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(table)
-
-        from sqlmodel import and_
-
-        from zenml.zen_stores.schemas import (
-            ModelVersionSchema,
-            RunMetadataResourceSchema,
-            RunMetadataSchema,
-        )
-
-        if self.run_metadata is not None:
-            from zenml.enums import MetadataResourceTypes
-
-            for key, value in self.run_metadata.items():
-                additional_filter = and_(
-                    RunMetadataResourceSchema.resource_id
-                    == ModelVersionSchema.id,
-                    RunMetadataResourceSchema.resource_type
-                    == MetadataResourceTypes.MODEL_VERSION,
-                    RunMetadataResourceSchema.run_metadata_id
-                    == RunMetadataSchema.id,
-                    self.generate_custom_query_conditions_for_column(
-                        value=value,
-                        table=RunMetadataSchema,
-                        column="value",
-                    ),
-                )
-                custom_filters.append(additional_filter)
-
-        return custom_filters
-
     def apply_filter(
         self,
         query: AnyQuery,