Make Vertex AI evaluation optional to avoid litellm bloat

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"``

providers/google/docs/changelog.rst

@@ -27,6 +27,19 @@
 Changelog
 ---------
 
+22.3.0
+......
+
+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]``.
+
 22.2.2
 ......
 

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'``

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 <concepts: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

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'",
@@ -156,6 +156,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",
 ]

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
+
+try:
+    from vertexai.preview.evaluation import EvalResult, EvalTask
+except ImportError as e:
+    _evaluation_import_error = e
+    # Fallback for mypy: use Any so that calls like EvalTask(...) don't error at type-check time.
+    # At runtime, guard checks _evaluation_import_error and raises before using these.
+    EvalResult = Any  # type: ignore[no-redef]
+    EvalTask = Any  # type: ignore[no-redef]
+else:
+    _evaluation_import_error = None
+
 
 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(

providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model.py

@@ -23,6 +23,7 @@
 
 # For no Pydantic environment, we need to skip the tests
 pytest.importorskip("google.cloud.aiplatform_v1")
+pytest.importorskip("vertexai.preview.evaluation")
 
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Part, Tool, grounding
 from vertexai.preview.evaluation import MetricPromptTemplateExamples

providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_generative_model_optional_evaluation.py

@@ -0,0 +1,91 @@
+#
+# 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
+
+import importlib
+import sys
+from collections.abc import Iterator
+from types import ModuleType
+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")
+
+EVALUATION_MODULE = "vertexai.preview.evaluation"
+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\]"
+
+_MISSING = object()
+
+
+@pytest.fixture
+def generative_model_hook_module_without_evaluation() -> Iterator[ModuleType]:
+    original_evaluation_module = sys.modules.get(EVALUATION_MODULE, _MISSING)
+    original_hook_module = sys.modules.get(HOOK_MODULE, _MISSING)
+
+    sys.modules[EVALUATION_MODULE] = None
+    hook_module = importlib.import_module(HOOK_MODULE)
+    hook_module = importlib.reload(hook_module)
+
+    try:
+        yield hook_module
+    finally:
+        if original_evaluation_module is _MISSING:
+            sys.modules.pop(EVALUATION_MODULE, None)
+        else:
+            sys.modules[EVALUATION_MODULE] = original_evaluation_module
+
+        if original_hook_module is _MISSING:
+            sys.modules.pop(HOOK_MODULE, None)
+        else:
+            importlib.reload(hook_module)
+
+
+def test_get_eval_task_raises_optional_provider_feature_exception_without_evaluation_extra(
+    generative_model_hook_module_without_evaluation: ModuleType,
+):
+    with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id):
+        hook = generative_model_hook_module_without_evaluation.GenerativeModelHook()
+
+    with pytest.raises(AirflowOptionalProviderFeatureException, match=INSTALL_EXTRA_REGEX):
+        hook.get_eval_task(dataset={}, metrics=[], experiment="test-experiment")
+
+
+def test_run_evaluation_raises_optional_provider_feature_exception_without_evaluation_extra(
+    generative_model_hook_module_without_evaluation: ModuleType,
+):
+    with mock.patch(BASE_HOOK_INIT, new=mock_base_gcp_hook_default_project_id):
+        hook = generative_model_hook_module_without_evaluation.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}",
+        )

providers/google/tests/unit/google/cloud/operators/vertex_ai/test_generative_model.py

@@ -23,6 +23,7 @@
 # For no Pydantic environment, we need to skip the tests
 pytest.importorskip("google.cloud.aiplatform_v1")
 pytest.importorskip("google.cloud.aiplatform_v1beta1")
+pytest.importorskip("vertexai.preview.evaluation")
 vertexai = pytest.importorskip("vertexai.generative_models")
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding
 from vertexai.preview.evaluation import MetricPromptTemplateExamples