feat: Implement check_license_agreement_status and accept_model_license_agreement for Model Garden OpenModel.

PiperOrigin-RevId: 753629633

Author: haomengchao

tests/unit/vertexai/model_garden/test_model_garden.py

@@ -41,6 +41,7 @@
 
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"
+_TEST_PROJECT_NUMBER = "1234567890"
 
 _TEST_MODEL_FULL_RESOURCE_NAME = (
  "publishers/google/models/paligemma@paligemma-224-float32"
@@ -398,6 +399,40 @@ def list_publisher_models_mock():
  yield list_publisher_models
 
 
+@pytest.fixture
+def check_license_agreement_status_mock():
+ """Mocks the check_license_agreement_status method."""
+ with mock.patch.object(
+ model_garden_service.ModelGardenServiceClient,
+ "check_publisher_model_eula_acceptance",
+ ) as check_license_agreement_status:
+ check_license_agreement_status.return_value = (
+ types.PublisherModelEulaAcceptance(
+ project_number=_TEST_PROJECT_NUMBER,
+ publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+ publisher_model_eula_acked=True,
+ )
+ )
+ yield check_license_agreement_status
+
+
+@pytest.fixture
+def accept_model_license_agreement_mock():
+ """Mocks the accept_model_license_agreement method."""
+ with mock.patch.object(
+ model_garden_service.ModelGardenServiceClient,
+ "accept_publisher_model_eula",
+ ) as accept_model_license_agreement:
+ accept_model_license_agreement.return_value = (
+ types.PublisherModelEulaAcceptance(
+ project_number=_TEST_PROJECT_NUMBER,
+ publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+ publisher_model_eula_acked=True,
+ )
+ )
+ yield accept_model_license_agreement
+
+
 @pytest.mark.usefixtures(
  "google_auth_mock",
  "deploy_mock",
@@ -406,6 +441,8 @@ def list_publisher_models_mock():
  "export_publisher_model_mock",
  "batch_prediction_mock",
  "complete_bq_uri_mock",
+ "check_license_agreement_status_mock",
+ "accept_model_license_agreement_mock",
 )
 class TestModelGarden:
  """Test cases for ModelGarden class."""
@@ -999,3 +1036,43 @@ def test_batch_prediction_success(self, batch_prediction_mock):
  batch_prediction_job=expected_gapic_batch_prediction_job,
  timeout=None,
  )
+
+ def test_check_license_agreement_status_success(
+ self, check_license_agreement_status_mock
+ ):
+ """Tests checking EULA acceptance for a model."""
+ aiplatform.init(
+ project=_TEST_PROJECT,
+ location=_TEST_LOCATION,
+ )
+ model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
+ eula_acceptance = model.check_license_agreement_status()
+ check_license_agreement_status_mock.assert_called_once_with(
+ types.CheckPublisherModelEulaAcceptanceRequest(
+ parent=f"projects/{_TEST_PROJECT}",
+ publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+ )
+ )
+ assert eula_acceptance
+
+ def test_accept_model_license_agreement_success(
+ self, accept_model_license_agreement_mock
+ ):
+ """Tests accepting EULA for a model."""
+ aiplatform.init(
+ project=_TEST_PROJECT,
+ location=_TEST_LOCATION,
+ )
+ model = model_garden.OpenModel(model_name=_TEST_MODEL_FULL_RESOURCE_NAME)
+ eula_acceptance = model.accept_model_license_agreement()
+ accept_model_license_agreement_mock.assert_called_once_with(
+ types.AcceptPublisherModelEulaRequest(
+ parent=f"projects/{_TEST_PROJECT}",
+ publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+ )
+ )
+ assert eula_acceptance == types.PublisherModelEulaAcceptance(
+ project_number=_TEST_PROJECT_NUMBER,
+ publisher_model=_TEST_MODEL_FULL_RESOURCE_NAME,
+ publisher_model_eula_acked=True,
+ )

vertexai/model_garden/_model_garden.py

@@ -715,3 +715,44 @@ def batch_predict(
  starting_replica_count=starting_replica_count,
  max_replica_count=max_replica_count,
  )
+
+ def check_license_agreement_status(self) -> bool:
+ """Check whether the project has accepted the license agreement of the model.
+
+ EULA (End User License Agreement) is a legal document that the user must
+ accept before using the model. For Models having license restrictions,
+ the user must accept the EULA before using the model. You can check the
+ details of the License in Model Garden.
+
+ Returns:
+ bool : True if the project has accepted the End User License
+ Agreement, False otherwise.
+ """
+ request = types.CheckPublisherModelEulaAcceptanceRequest(
+ parent=f"projects/{self._project}",
+ publisher_model=self._publisher_model_name,
+ )
+ response = self._model_garden_client.check_publisher_model_eula_acceptance(
+ request
+ )
+ return response.publisher_model_eula_acked
+
+ def accept_model_license_agreement(
+ self,
+ ) -> types.model_garden_service.PublisherModelEulaAcceptance:
+ """Accepts the EULA(End User License Agreement) of the model for the project.
+
+ For Models having license restrictions, the user must accept the EULA
+ before using the model. Calling this method will mark the EULA as accepted
+ for the project.
+
+ Returns:
+ types.model_garden_service.PublisherModelEulaAcceptance:
+ The response of the accept_eula call, containing project number,
+ model name and acceptance status.
+ """
+ request = types.AcceptPublisherModelEulaRequest(
+ parent=f"projects/{self._project}",
+ publisher_model=self._publisher_model_name,
+ )
+ return self._model_garden_client.accept_publisher_model_eula(request)