diff --git a/providers/google/README.rst b/providers/google/README.rst index 0602b6f5b0043..b50742c5c74c3 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.164.0`` +``google-cloud-aiplatform`` ``>=1.164.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"`` @@ -180,9 +180,10 @@ Dependent package Optional dependencies ---------------------- -==================== ==================================================== +==================== ============================================================================================================================================================================================================ Extra Dependencies -==================== ==================================================== +==================== ============================================================================================================================================================================================================ +``evaluation`` ``tqdm>=4.23.0``, ``scikit-learn>=1.6.0``, ``jsonschema>=4.19.1``, ``ruamel.yaml>=0.19.0``, ``pyyaml>=6.0.3``, ``litellm<1.86.0,>=1.83.7;python_version<'3.14'``, ``litellm>=1.93.0;python_version>='3.14'`` ``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'`` @@ -204,7 +205,7 @@ Extra Dependencies ``http`` ``apache-airflow-providers-http`` ``standard`` ``apache-airflow-providers-standard`` ``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0`` -==================== ==================================================== +==================== ============================================================================================================================================================================================================ The changelog for the provider package can be found in the `changelog `_. diff --git a/providers/google/docs/changelog.rst b/providers/google/docs/changelog.rst index c1e4de295ebe1..67c57b6a411ff 100644 --- a/providers/google/docs/changelog.rst +++ b/providers/google/docs/changelog.rst @@ -27,6 +27,21 @@ Changelog --------- +23.0.0 +...... + +Breaking changes +~~~~~~~~~~~~~~~~ + +* The Vertex AI evaluation feature is now optional. ``RunEvaluationOperator`` and + ``GenerativeModelHook.run_evaluation`` / ``GenerativeModelHook.get_eval_task`` now require the + ``evaluation`` extra. Previously ``google-cloud-aiplatform[evaluation]`` was installed + unconditionally, pulling ``litellm``, ``scikit-learn`` and their transitive dependencies into every + Google provider install. To restore the previous behaviour, install + ``pip install 'apache-airflow-providers-google[evaluation]'``. Without the extra the provider + imports and all non-evaluation features keep working; the evaluation code paths raise + ``AirflowOptionalProviderFeatureException`` with this install hint. + 22.4.0 ...... diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst index 516d2d2ae920d..3108805b4c955 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.164.0`` +``google-cloud-aiplatform`` ``>=1.164.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`` ``tqdm>=4.23.0``, ``scikit-learn>=1.6.0``, ``jsonschema>=4.19.1``, ``ruamel.yaml>=0.19.0``, ``pyyaml>=6.0.3``, ``litellm<1.86.0,>=1.83.7;python_version<'3.14'``, ``litellm>=1.93.0;python_version>='3.14'`` ``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'`` @@ -265,7 +266,7 @@ Extra Dependencies ``http`` ``apache-airflow-providers-http`` ``standard`` ``apache-airflow-providers-standard`` ``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0`` -==================== ==================================================== +==================== ============================================================================================================================================================================================================ Downloading official packages ----------------------------- diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst b/providers/google/docs/operators/cloud/vertex_ai.rst index dd69e361de28d..1d32a0acf497c 100644 --- a/providers/google/docs/operators/cloud/vertex_ai.rst +++ b/providers/google/docs/operators/cloud/vertex_ai.rst @@ -638,6 +638,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 7eec07d5df810..a791457b58fa3 100644 --- a/providers/google/pyproject.toml +++ b/providers/google/pyproject.toml @@ -87,7 +87,7 @@ dependencies = [ # Floor raised to 1.164.0: earlier "evaluation" extras cap litellm below the version that # carries the fixes for CVE-2026-35030 and its follow-on advisories on Python 3.14. # - https://github.com/googleapis/python-aiplatform/issues/7057 - "google-cloud-aiplatform[evaluation]>=1.164.0", + "google-cloud-aiplatform>=1.164.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'", @@ -159,6 +159,20 @@ 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" = [ + # The "evaluation" extra of google-cloud-aiplatform limits litellm to range that has security + # vulnerabilities when we want to get 3.14 support. Until this is fixed we need to bump litellm to + # higher version and remove the "evaluation" extra from google-cloud-aiplatform. + # Tracked in: https://github.com/googleapis/python-aiplatform/issues/7057 + "tqdm>=4.23.0", + "scikit-learn>=1.6.0", + "jsonschema>=4.19.1", + "ruamel.yaml>=0.19.0", + "pyyaml>=6.0.3", + "litellm<1.86.0,>=1.83.7;python_version<'3.14'", + "litellm>=1.93.0;python_version>='3.14'", + # End of evolution extra workaround +] "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..0dd8277e1b040 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,44 @@ 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 + +# ``vertexai.preview.evaluation`` pulls in litellm and scikit-learn, so it ships in the optional +# "evaluation" extra rather than in the base provider. Importing this module must stay safe without +# it: the failure is remembered here and only re-raised from the evaluation code paths. +_evaluation_import_error: ImportError | None = None + +if not TYPE_CHECKING: + # ImportError rather than ModuleNotFoundError: a partially installed evaluation stack should get + # the same install hint instead of a raw traceback. + try: + from vertexai.preview.evaluation import EvalTask + except ImportError as e: + _evaluation_import_error = e + + +def _raise_if_evaluation_unavailable() -> None: + """Raise with an install hint if the optional ``evaluation`` extra is not installed.""" + if _evaluation_import_error is None: + return + 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]" + ) from _evaluation_import_error + class GenerativeModelHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Generative Model APIs.""" @@ -64,6 +91,7 @@ def get_eval_task( experiment: str, ) -> EvalTask: """Return an EvalTask object.""" + _raise_if_evaluation_unavailable() eval_task = EvalTask( dataset=dataset, metrics=metrics, @@ -115,6 +143,7 @@ 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. """ + _raise_if_evaluation_unavailable() 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..3f350f943974c 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,21 @@ 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: + # Only when the optional "evaluation" extra is absent: stand in for the metric constants so the + # tests below still run against the real hook/operator. With the extra installed the real values + # are used. Built in one expression so a single ignore covers it -- mypy sees a class here. + MetricPromptTemplateExamples = mock.MagicMock( # type: ignore[assignment,misc] + Pointwise=mock.MagicMock( + SUMMARIZATION_QUALITY="summarization_quality", + GROUNDEDNESS="groundedness", + VERBOSITY="verbosity", + INSTRUCTION_FOLLOWING="instruction_following", + ) + ) from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import ( GenerativeModelHook, @@ -143,6 +157,7 @@ def setup_method(self): self.hook = GenerativeModelHook(gcp_conn_id=TEST_GCP_CONN_ID) self.hook.get_credentials = self.dummy_get_credentials + @mock.patch(GENERATIVE_MODEL_STRING.format("_evaluation_import_error"), None) @mock.patch(GENERATIVE_MODEL_STRING.format("GenerativeModelHook.get_generative_model")) @mock.patch(GENERATIVE_MODEL_STRING.format("GenerativeModelHook.get_eval_task")) def test_run_evaluation(self, mock_eval_task, mock_model) -> None: 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..26f4ad6c44b74 --- /dev/null +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py @@ -0,0 +1,84 @@ +# +# 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_ERRORS = ( + pytest.param( + ImportError("No module named 'vertexai.preview.evaluation'"), + id="missing-evaluation-module", + ), + pytest.param( + ImportError("No module named 'sklearn'"), + id="missing-sklearn", + ), +) + + +@pytest.fixture +def hook() -> GenerativeModelHook: + with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id): + return GenerativeModelHook() + + +@pytest.mark.parametrize("import_error", MISSING_EVALUATION_ERRORS) +def test_get_eval_task_raises_optional_provider_feature_exception_without_evaluation_extra( + hook: GenerativeModelHook, + import_error: ImportError, +): + with mock.patch(f"{HOOK_MODULE}._evaluation_import_error", import_error): + with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX): + hook.get_eval_task(dataset={}, metrics=[], experiment="test-experiment") + + +@pytest.mark.parametrize("import_error", MISSING_EVALUATION_ERRORS) +def test_run_evaluation_raises_optional_provider_feature_exception_without_evaluation_extra( + hook: GenerativeModelHook, + import_error: ImportError, +): + with mock.patch(f"{HOOK_MODULE}._evaluation_import_error", import_error): + 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..c17bbe09e1ba5 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,21 @@ 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: + # Only when the optional "evaluation" extra is absent: stand in for the metric constants so the + # tests below still run against the real hook/operator. With the extra installed the real values + # are used. Built in one expression so a single ignore covers it -- mypy sees a class here. + MetricPromptTemplateExamples = mock.MagicMock( # type: ignore[assignment,misc] + Pointwise=mock.MagicMock( + SUMMARIZATION_QUALITY="summarization_quality", + GROUNDEDNESS="groundedness", + VERBOSITY="verbosity", + 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 d4f97bd3b9923..38a403f150fa9 100644 --- a/uv.lock +++ b/uv.lock @@ -5516,7 +5516,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" }, @@ -5588,6 +5588,16 @@ cncf-kubernetes = [ common-messaging = [ { name = "apache-airflow-providers-common-messaging" }, ] +evaluation = [ + { name = "jsonschema" }, + { name = "litellm", version = "1.85.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "litellm", version = "1.96.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, + { name = "pyyaml" }, + { name = "ruamel-yaml" }, + { name = "scikit-learn", version = "1.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tqdm" }, +] fab = [ { name = "apache-airflow-providers-fab" }, ] @@ -5710,7 +5720,7 @@ 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.164.0" }, + { name = "google-cloud-aiplatform", specifier = ">=1.164.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" }, @@ -5755,6 +5765,9 @@ requires-dist = [ { name = "grpcio-gcp", specifier = ">=0.2.2" }, { name = "httpx", specifier = ">=0.25.0" }, { name = "immutabledict", specifier = ">=4.2.0" }, + { name = "jsonschema", marker = "extra == 'evaluation'", specifier = ">=4.19.1" }, + { name = "litellm", marker = "python_full_version >= '3.14' and extra == 'evaluation'", specifier = ">=1.93.0" }, + { name = "litellm", marker = "python_full_version < '3.14' and extra == 'evaluation'", specifier = ">=1.83.7,<1.86.0" }, { name = "looker-sdk", specifier = ">=22.4.0,!=24.18.0" }, { name = "pandas", marker = "python_full_version < '3.13'", specifier = ">=2.1.2" }, { name = "pandas", marker = "python_full_version == '3.13.*'", specifier = ">=2.2.3" }, @@ -5766,15 +5779,19 @@ requires-dist = [ { name = "pyarrow", marker = "python_full_version >= '3.14'", specifier = ">=22.0.0" }, { name = "pydantic", marker = "python_full_version >= '3.14'", specifier = ">=2.13.1" }, { name = "python-slugify", specifier = ">=7.0.0" }, + { name = "pyyaml", marker = "extra == 'evaluation'", specifier = ">=6.0.3" }, { name = "ray", extras = ["default"], marker = "python_full_version < '3.13'", specifier = ">=2.42.0" }, { name = "ray", extras = ["default"], marker = "python_full_version == '3.13.*'", specifier = ">=2.49.0" }, { name = "ray", extras = ["default"], marker = "python_full_version == '3.14.*'", specifier = ">=2.55.0" }, + { name = "ruamel-yaml", marker = "extra == 'evaluation'", specifier = ">=0.19.0" }, + { name = "scikit-learn", marker = "extra == 'evaluation'", specifier = ">=1.6.0" }, { name = "sqlalchemy-bigquery", specifier = ">=1.2.1" }, { name = "sqlalchemy-spanner", specifier = ">=1.6.2" }, { name = "tenacity", specifier = ">=8.3.0" }, + { name = "tqdm", marker = "extra == 'evaluation'", specifier = ">=4.23.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 = [