diff --git a/providers/google/README.rst b/providers/google/README.rst index bd1c048e1a7e1..9bf97e7a945c5 100644 --- a/providers/google/README.rst +++ b/providers/google/README.rst @@ -77,7 +77,7 @@ PIP package Version required ``google-auth`` ``>=2.29.0`` ``google-auth-httplib2`` ``>=0.0.1`` ``google-genai`` ``>=2.8.0`` -``google-cloud-aiplatform[evaluation]`` ``>=1.155.0`` +``google-cloud-aiplatform`` ``>=1.155.0`` ``ray[default]`` ``>=2.42.0; python_version < "3.13"`` ``ray[default]`` ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"`` ``ray[default]`` ``>=2.55.0; python_version >= "3.14" and python_version < "3.15"`` diff --git a/providers/google/docs/changelog.rst b/providers/google/docs/changelog.rst index 8d3603296bd97..d77f0e09a2bd9 100644 --- a/providers/google/docs/changelog.rst +++ b/providers/google/docs/changelog.rst @@ -44,6 +44,16 @@ Changelog because a team name may itself contain ``--`` and the resulting name would be ambiguous. See :doc:`/secrets-backends/google-cloud-secret-manager-backend` for the full convention. +Breaking changes +~~~~~~~~~~~~~~~~ + +* The Evaluation feature of Vertex AI is now optional. ``RunEvaluationOperator`` + and ``GenerativeModelHook.run_evaluation`` require + ``apache-airflow-providers-google[evaluation]`` extra. Previously + ``google-cloud-aiplatform[evaluation]`` was installed unconditionally, + pulling ``litellm`` and ``scikit-learn`` for all provider users. To + restore old behavior, install with ``pip install apache-airflow-providers-google[evaluation]``. + Features ~~~~~~~~ diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst index 05c764f3c1cce..f8a9e91a5823b 100644 --- a/providers/google/docs/index.rst +++ b/providers/google/docs/index.rst @@ -130,7 +130,7 @@ PIP package Version required ``google-auth`` ``>=2.29.0`` ``google-auth-httplib2`` ``>=0.0.1`` ``google-genai`` ``>=2.8.0`` -``google-cloud-aiplatform[evaluation]`` ``>=1.155.0`` +``google-cloud-aiplatform`` ``>=1.155.0`` ``ray[default]`` ``>=2.42.0; python_version < "3.13"`` ``ray[default]`` ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"`` ``ray[default]`` ``>=2.55.0; python_version >= "3.14" and python_version < "3.15"`` @@ -238,12 +238,13 @@ Install them when installing from PyPI. For example: .. code-block:: bash - pip install apache-airflow-providers-google[cncf.kubernetes] + pip install apache-airflow-providers-google[evaluation] ==================== ==================================================== Extra Dependencies ==================== ==================================================== +``evaluation`` ``google-cloud-aiplatform[evaluation]>=1.155.0`` ``cncf.kubernetes`` ``apache-airflow-providers-cncf-kubernetes>=10.1.0`` ``fab`` ``apache-airflow-providers-fab>=2.0.0`` ``leveldb`` ``plyvel>=1.5.1; python_version < '3.13'`` diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst b/providers/google/docs/operators/cloud/vertex_ai.rst index f11b8787ef684..41df2c3c4f4b8 100644 --- a/providers/google/docs/operators/cloud/vertex_ai.rst +++ b/providers/google/docs/operators/cloud/vertex_ai.rst @@ -627,6 +627,9 @@ To evaluate a model you can use :class:`~airflow.providers.google.cloud.operators.vertex_ai.generative_model.RunEvaluationOperator`. The operator returns the evaluation summary metrics in :ref:`XCom ` under ``summary_metrics`` key. +Vertex AI evaluation requires the ``evaluation`` extra. Install it with +``pip install apache-airflow-providers-google[evaluation]``. + .. exampleinclude:: /../../google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py :language: python :dedent: 4 diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml index e750891904d1a..43cb82f0729ad 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -84,7 +84,7 @@ dependencies = [ # google-cloud-aiplatform doesn't install ray for python 3.12 (issue: https://github.com/googleapis/python-aiplatform/issues/5252). # Temporarily lock in ray 2.42.0 which is compatible with python 3.12 until linked issue is solved. # Remove the ray dependency as well as google-cloud-bigquery-storage once linked issue is fixed - "google-cloud-aiplatform[evaluation]>=1.155.0", + "google-cloud-aiplatform>=1.155.0", "ray[default]>=2.42.0;python_version<'3.13'", "ray[default]>=2.49.0;python_version>='3.13' and python_version <'3.14'", "ray[default]>=2.55.0;python_version>='3.14' and python_version <'3.15'", @@ -165,6 +165,9 @@ dependencies = [ # The optional dependencies should be modified in place in the generated file # Any change in the dependencies is preserved when the file is regenerated [project.optional-dependencies] +"evaluation" = [ + "google-cloud-aiplatform[evaluation]>=1.155.0", +] "cncf.kubernetes" = [ "apache-airflow-providers-cncf-kubernetes>=10.1.0", ] diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py index 06854d313248a..506afab210e71 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py @@ -19,17 +19,31 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any import vertexai from vertexai.generative_models import GenerativeModel from vertexai.language_models import TextEmbeddingModel from vertexai.preview import generative_models as preview_generative_model from vertexai.preview.caching import CachedContent -from vertexai.preview.evaluation import EvalResult, EvalTask +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from vertexai.preview.evaluation import EvalResult, EvalTask + +_evaluation_import_error: ImportError | None = None + +if not TYPE_CHECKING: + try: + from vertexai.preview.evaluation import EvalResult, EvalTask + except ImportError as e: + _evaluation_import_error = e + # Runtime fallback: guard checks _evaluation_import_error and raises before using these. + EvalResult = Any + EvalTask = Any + class GenerativeModelHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Generative Model APIs.""" @@ -64,6 +78,12 @@ def get_eval_task( experiment: str, ) -> EvalTask: """Return an EvalTask object.""" + if _evaluation_import_error: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + f"Original error: {_evaluation_import_error}. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) eval_task = EvalTask( dataset=dataset, metrics=metrics, @@ -115,6 +135,12 @@ def run_evaluation( :param system_instruction: Optional. An instruction given to the model to guide its behavior. :param tools: Optional. A list of tools available to the model during evaluation, such as a data store. """ + if _evaluation_import_error: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + f"Original error: {_evaluation_import_error}. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) vertexai.init(project=project_id, location=location, credentials=self.get_credentials()) model = self.get_generative_model( diff --git a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py index c9a6f9ffcafa4..31764d07327cb 100644 --- a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py +++ b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py @@ -42,6 +42,7 @@ ) from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException from airflow.providers.google.cloud.operators.gen_ai import ( GenAICountTokensOperator, GenAICreateCachedContentOperator, @@ -169,13 +170,14 @@ def _get_actual_models(key) -> dict[str, str]: def _get_metrics(): - """ - Lazily import and return the metrics list. - - This avoids slow imports during DAG parsing by deferring the import - until the operator is actually created. - """ - from vertexai.preview.evaluation import MetricPromptTemplateExamples + """Return metrics without importing optional evaluation dependencies during module import.""" + try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples + except ImportError as e: + raise AirflowOptionalProviderFeatureException( + "The 'evaluation' extra is required for Vertex AI evaluation. " + "Install with: pip install apache-airflow-providers-google[evaluation]" + ) from e return [ MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY, diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py index 3146992dba599..752f4413b2603 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py @@ -25,7 +25,15 @@ pytest.importorskip("google.cloud.aiplatform_v1") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Part, Tool, grounding -from vertexai.preview.evaluation import MetricPromptTemplateExamples + +try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples +except ImportError: + MetricPromptTemplateExamples = mock.MagicMock() + MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY = "summarization_quality" + MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS = "groundedness" + MetricPromptTemplateExamples.Pointwise.VERBOSITY = "verbosity" + MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING = "instruction_following" from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import ( GenerativeModelHook, diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py new file mode 100644 index 0000000000000..139760c91845c --- /dev/null +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py @@ -0,0 +1,71 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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 +# +# http://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. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +from unit.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id + +pytest.importorskip("google.cloud.aiplatform_v1") + +from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import ( + GenerativeModelHook, +) + +HOOK_MODULE = "airflow.providers.google.cloud.hooks.vertex_ai.generative_model" +BASE_HOOK_INIT = "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__" +INSTALL_EXTRA_REGEX = r"apache-airflow-providers-google\[evaluation\]" + +# Patching the guard variable (instead of reloading the hook module with the evaluation import +# blocked) keeps the module and class objects stable for other test files in the same session. +MISSING_EVALUATION_IMPORT = ImportError("No module named 'vertexai.preview.evaluation'") + + +@pytest.fixture +def hook() -> GenerativeModelHook: + with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): + return GenerativeModelHook() + + +@mock.patch(f"{HOOK_MODULE}._evaluation_import_error", MISSING_EVALUATION_IMPORT) +def test_get_eval_task_raises_optional_provider_feature_exception_without_evaluation_extra( + hook: GenerativeModelHook, +): + with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): + hook.get_eval_task(dataset={}, metrics=[], experiment="test-experiment") + + +@mock.patch(f"{HOOK_MODULE}._evaluation_import_error", MISSING_EVALUATION_IMPORT) +def test_run_evaluation_raises_optional_provider_feature_exception_without_evaluation_extra( + hook: GenerativeModelHook, +): + with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): + hook.run_evaluation( + project_id="test-project", + location="us-central1", + pretrained_model="gemini-pro", + eval_dataset={}, + metrics=[], + experiment_name="test-experiment", + experiment_run_name="test-run", + prompt_template="{prompt}", + ) diff --git a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py index e9bd014eac794..5cd5eebe6124d 100644 --- a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py +++ b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py @@ -25,7 +25,15 @@ pytest.importorskip("google.cloud.aiplatform_v1beta1") vertexai = pytest.importorskip("vertexai.generative_models") from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding -from vertexai.preview.evaluation import MetricPromptTemplateExamples + +try: + from vertexai.preview.evaluation import MetricPromptTemplateExamples +except ImportError: + MetricPromptTemplateExamples = mock.MagicMock() + MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY = "summarization_quality" + MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS = "groundedness" + MetricPromptTemplateExamples.Pointwise.VERBOSITY = "verbosity" + MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING = "instruction_following" from airflow.providers.google.cloud.operators.vertex_ai.generative_model import ( RunEvaluationOperator, diff --git a/uv.lock b/uv.lock index 3b52d8d26c2f2..09fc6ab614252 100644 --- a/uv.lock +++ b/uv.lock @@ -5520,7 +5520,7 @@ dependencies = [ { name = "google-api-python-client" }, { name = "google-auth" }, { name = "google-auth-httplib2" }, - { name = "google-cloud-aiplatform", extra = ["evaluation"] }, + { name = "google-cloud-aiplatform" }, { name = "google-cloud-alloydb" }, { name = "google-cloud-automl" }, { name = "google-cloud-batch" }, @@ -5591,6 +5591,9 @@ cncf-kubernetes = [ common-messaging = [ { name = "apache-airflow-providers-common-messaging" }, ] +evaluation = [ + { name = "google-cloud-aiplatform", extra = ["evaluation"] }, +] fab = [ { name = "apache-airflow-providers-fab" }, ] @@ -5713,7 +5716,8 @@ requires-dist = [ { name = "google-api-python-client", specifier = ">=2.0.2" }, { name = "google-auth", specifier = ">=2.29.0" }, { name = "google-auth-httplib2", specifier = ">=0.0.1" }, - { name = "google-cloud-aiplatform", extras = ["evaluation"], specifier = ">=1.155.0" }, + { name = "google-cloud-aiplatform", specifier = ">=1.155.0" }, + { name = "google-cloud-aiplatform", extras = ["evaluation"], marker = "extra == 'evaluation'", specifier = ">=1.155.0" }, { name = "google-cloud-alloydb", specifier = ">=0.4.0" }, { name = "google-cloud-automl", specifier = ">=2.12.0" }, { name = "google-cloud-batch", specifier = ">=0.13.0" }, @@ -5777,7 +5781,7 @@ requires-dist = [ { name = "tenacity", specifier = ">=8.3.0" }, { name = "types-protobuf", specifier = ">=5.27.0,!=5.29.1.20250402" }, ] -provides-extras = ["cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-messaging"] +provides-extras = ["evaluation", "cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-messaging"] [package.metadata.requires-dev] dev = [