feat: Add filter to Model Registry list_versions API.

PiperOrigin-RevId: 501046171

Author: vertex-sdk-bot

google/cloud/aiplatform/models.py

@@ -4796,9 +4796,22 @@ def get_model(
 
     def list_versions(
         self,
+        filter: Optional[str] = None,
     ) -> List[VersionInfo]:
         """Lists the versions and version info of a model.
 
+        Args:
+            filter (str):
+                Optional. An expression for filtering the results of the request.
+                For field names both snake_case and camelCase are supported.
+                -  `labels` supports general map functions that is:
+                    -  `labels.key=value` - key:value equality
+                    -  `labels.key:* or labels:key - key existence
+                    -  A key including a space must be quoted.
+                        `labels."a key"`.
+                Some examples:
+                -  `labels.myKey="myValue"`
+
         Returns:
             List[VersionInfo]:
                 A list of VersionInfo, each containing
@@ -4807,8 +4820,13 @@ def list_versions(
 
         _LOGGER.info(f"Getting versions for {self.model_resource_name}")
 
-        page_result = self.client.list_model_versions(
+        request = gca_model_service_compat.ListModelVersionsRequest(
             name=self.model_resource_name,
+            filter=filter,
+        )
+
+        page_result = self.client.list_model_versions(
+            request=request,
         )
 
         versions = [

tests/unit/aiplatform/test_models.py

@@ -255,7 +255,9 @@
 _TEST_VERSION_ID = "2"
 _TEST_VERSION_ALIAS_1 = "myalias"
 _TEST_VERSION_ALIAS_2 = "youralias"
-_TEST_MODEL_VERSION_DESCRIPTION = "My version description"
+_TEST_MODEL_VERSION_DESCRIPTION_1 = "My version 1 description"
+_TEST_MODEL_VERSION_DESCRIPTION_2 = "My version 2 description"
+_TEST_MODEL_VERSION_DESCRIPTION_3 = "My version 3 description"
 
 _TEST_MODEL_VERSIONS_LIST = [
     gca_model.Model(
@@ -265,7 +267,7 @@
         display_name=_TEST_MODEL_NAME,
         name=f"{_TEST_MODEL_PARENT}@1",
         version_aliases=["default"],
-        version_description=_TEST_MODEL_VERSION_DESCRIPTION,
+        version_description=_TEST_MODEL_VERSION_DESCRIPTION_1,
     ),
     gca_model.Model(
         version_id="2",
@@ -274,7 +276,7 @@
         display_name=_TEST_MODEL_NAME,
         name=f"{_TEST_MODEL_PARENT}@2",
         version_aliases=[_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2],
-        version_description=_TEST_MODEL_VERSION_DESCRIPTION,
+        version_description=_TEST_MODEL_VERSION_DESCRIPTION_2,
     ),
     gca_model.Model(
         version_id="3",
@@ -283,9 +285,11 @@
         display_name=_TEST_MODEL_NAME,
         name=f"{_TEST_MODEL_PARENT}@3",
         version_aliases=[],
-        version_description=_TEST_MODEL_VERSION_DESCRIPTION,
+        version_description=_TEST_MODEL_VERSION_DESCRIPTION_3,
+        labels=_TEST_LABEL,
     ),
 ]
+_TEST_MODEL_VERSIONS_WITH_FILTER_LIST = [_TEST_MODEL_VERSIONS_LIST[2]]
 
 _TEST_MODELS_LIST = _TEST_MODEL_VERSIONS_LIST + [
     gca_model.Model(
@@ -295,7 +299,7 @@
         display_name=_TEST_MODEL_NAME_ALT,
         name=_TEST_MODEL_PARENT_ALT,
         version_aliases=["default"],
-        version_description=_TEST_MODEL_VERSION_DESCRIPTION,
+        version_description=_TEST_MODEL_VERSION_DESCRIPTION_1,
     ),
 ]
 
@@ -306,7 +310,7 @@
     display_name=_TEST_MODEL_NAME,
     name=f"{_TEST_MODEL_PARENT}@{_TEST_VERSION_ID}",
     version_aliases=[_TEST_VERSION_ALIAS_1, _TEST_VERSION_ALIAS_2],
-    version_description=_TEST_MODEL_VERSION_DESCRIPTION,
+    version_description=_TEST_MODEL_VERSION_DESCRIPTION_2,
 )
 
 _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_ID}"
@@ -683,6 +687,15 @@ def list_model_versions_mock():
         yield list_model_versions_mock
 
 
+@pytest.fixture
+def list_model_versions_with_filter_mock():
+    with mock.patch.object(
+        model_service_client.ModelServiceClient, "list_model_versions"
+    ) as list_model_versions_mock:
+        list_model_versions_mock.return_value = _TEST_MODEL_VERSIONS_WITH_FILTER_LIST
+        yield list_model_versions_mock
+
+
 @pytest.fixture
 def list_models_mock():
     with mock.patch.object(
@@ -2514,7 +2527,7 @@ def test_init_with_version_in_resource_name(self, get_model_with_version):
         assert model.display_name == _TEST_MODEL_NAME
         assert model.resource_name == _TEST_MODEL_PARENT
         assert model.version_id == _TEST_VERSION_ID
-        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION
+        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION_2
         # The Model yielded from upload should not have a version in resource name
         assert "@" not in model.resource_name
         # The Model yielded from upload SHOULD have a version in the versioned resource name
@@ -2527,7 +2540,7 @@ def test_init_with_version_arg(self, get_model_with_version):
         assert model.display_name == _TEST_MODEL_NAME
         assert model.resource_name == _TEST_MODEL_PARENT
         assert model.version_id == _TEST_VERSION_ID
-        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION
+        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION_2
         # The Model yielded from upload should not have a version in resource name
         assert "@" not in model.resource_name
         # The Model yielded from upload SHOULD have a version in the versioned resource name
@@ -2584,7 +2597,7 @@ def test_upload_new_version(
             "upload_request_timeout": None,
             "model_id": _TEST_ID,
             "parent_model": parent,
-            "version_description": _TEST_MODEL_VERSION_DESCRIPTION,
+            "version_description": _TEST_MODEL_VERSION_DESCRIPTION_2,
             "version_aliases": aliases,
             "is_default_version": default,
         }
@@ -2610,7 +2623,7 @@ def test_upload_new_version(
         assert upload_model_request.model.version_aliases == goal
         assert (
             upload_model_request.model.version_description
-            == _TEST_MODEL_VERSION_DESCRIPTION
+            == _TEST_MODEL_VERSION_DESCRIPTION_2
         )
         assert upload_model_request.parent_model == _TEST_MODEL_PARENT
         assert upload_model_request.model_id == _TEST_ID
@@ -2622,7 +2635,7 @@ def test_get_model_instance_from_registry(self, get_model_with_version):
         assert model.display_name == _TEST_MODEL_NAME
         assert model.resource_name == _TEST_MODEL_PARENT
         assert model.version_id == _TEST_VERSION_ID
-        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION
+        assert model.version_description == _TEST_MODEL_VERSION_DESCRIPTION_2
 
     def test_list_versions(self, list_model_versions_mock, get_model_with_version):
         my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION)
@@ -2643,6 +2656,28 @@ def test_list_versions(self, list_model_versions_mock, get_model_with_version):
             assert model.name.startswith(ver.model_resource_name)
             assert model.name.endswith(ver.version_id)
 
+    def test_list_versions_with_filter(
+        self, list_model_versions_with_filter_mock, get_model_with_version
+    ):
+        my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION)
+        versions = my_model.versioning_registry.list_versions(
+            filter='labels.team="experimentation"'
+        )
+
+        assert len(versions) == len(_TEST_MODEL_VERSIONS_WITH_FILTER_LIST)
+
+        ver = versions[0]
+        model = _TEST_MODEL_VERSIONS_WITH_FILTER_LIST[0]
+        assert ver.version_id == "3"
+        assert ver.version_create_time == model.version_create_time
+        assert ver.version_update_time == model.version_update_time
+        assert ver.model_display_name == model.display_name
+        assert ver.version_aliases == model.version_aliases
+        assert ver.version_description == model.version_description
+
+        assert model.name.startswith(ver.model_resource_name)
+        assert model.name.endswith(ver.version_id)
+
     def test_get_version_info(self, get_model_with_version):
         my_model = models.Model(_TEST_MODEL_NAME, _TEST_PROJECT, _TEST_LOCATION)
         ver = my_model.versioning_registry.get_version_info("2")