apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.1.0__py3-none-any.whl

airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py

@@ -21,38 +21,26 @@ from __future__ import annotations
 
 import time
 from datetime import timedelta
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Literal
 
 import vertexai
+from google.cloud import aiplatform
 from vertexai.generative_models import GenerativeModel
-from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
+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 vertexai.preview.generative_models import GenerativeModel as preview_generative_model
 from vertexai.preview.tuning import sft
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:
-    from google.cloud.aiplatform_v1 import types as types_v1
     from google.cloud.aiplatform_v1beta1 import types as types_v1beta1
 
 
 class GenerativeModelHook(GoogleBaseHook):
     """Hook for Google Cloud Vertex AI Generative Model APIs."""
 
-    @deprecated(
-        planned_removal_date="April 09, 2025",
-        use_instead="GenerativeModelHook.get_generative_model",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_text_generation_model(self, pretrained_model: str):
-        """Return a Model Garden Model object based on Text Generation."""
-        model = TextGenerationModel.from_pretrained(pretrained_model)
-        return model
-
     def get_text_embedding_model(self, pretrained_model: str):
         """Return a Model Garden Model object based on Text Embedding."""
         model = TextEmbeddingModel.from_pretrained(pretrained_model)
@@ -61,7 +49,7 @@ class GenerativeModelHook(GoogleBaseHook):
     def get_generative_model(
         self,
         pretrained_model: str,
-        system_instruction: str | None = None,
+        system_instruction: Any | None = None,
         generation_config: dict | None = None,
         safety_settings: dict | None = None,
         tools: list | None = None,
@@ -93,66 +81,13 @@ class GenerativeModelHook(GoogleBaseHook):
     def get_cached_context_model(
         self,
         cached_content_name: str,
-    ) -> preview_generative_model:
+    ) -> Any:
         """Return a Generative Model with Cached Context."""
         cached_content = CachedContent(cached_content_name=cached_content_name)
 
-        cached_context_model = preview_generative_model.from_cached_content(cached_content)
+        cached_context_model = preview_generative_model.GenerativeModel.from_cached_content(cached_content)
         return cached_context_model
 
-    @deprecated(
-        planned_removal_date="April 09, 2025",
-        use_instead="GenerativeModelHook.generative_model_generate_content",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def text_generation_model_predict(
-        self,
-        prompt: str,
-        pretrained_model: str,
-        temperature: float,
-        max_output_tokens: int,
-        top_p: float,
-        top_k: int,
-        location: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> str:
-        """
-        Use the Vertex AI PaLM API to generate natural language text.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
-        :param prompt: Required. Inputs or queries that a user or a program gives
-            to the Vertex AI PaLM API, in order to elicit a specific response.
-        :param pretrained_model: A pre-trained model optimized for performing natural
-            language tasks such as classification, summarization, extraction, content
-            creation, and ideation.
-        :param temperature: Temperature controls the degree of randomness in token
-            selection.
-        :param max_output_tokens: Token limit determines the maximum amount of text
-            output.
-        :param top_p: Tokens are selected from most probable to least until the sum
-            of their probabilities equals the top_p value. Defaults to 0.8.
-        :param top_k: A top_k of 1 means the selected token is the most probable
-            among all tokens.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-
-        parameters = {
-            "temperature": temperature,
-            "max_output_tokens": max_output_tokens,
-            "top_p": top_p,
-            "top_k": top_k,
-        }
-
-        model = self.get_text_generation_model(pretrained_model)
-
-        response = model.predict(
-            prompt=prompt,
-            **parameters,
-        )
-        return response.text
-
     @GoogleBaseHook.fallback_to_default_project_id
     def text_embedding_model_get_embeddings(
         self,
@@ -182,11 +117,11 @@ class GenerativeModelHook(GoogleBaseHook):
         self,
         contents: list,
         location: str,
+        pretrained_model: str,
         tools: list | None = None,
         generation_config: dict | None = None,
         safety_settings: dict | None = None,
         system_instruction: str | None = None,
-        pretrained_model: str = "gemini-pro",
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> str:
         """
@@ -200,7 +135,7 @@ class GenerativeModelHook(GoogleBaseHook):
         :param safety_settings: Optional. Per request settings for blocking unsafe content.
         :param tools: Optional. A list of tools available to the model during evaluation, such as a data store.
         :param system_instruction: Optional. An instruction given to the model to guide its behavior.
-        :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
+        :param pretrained_model: Required. Model,
             supporting prompts with text-only input, including natural language
             tasks, multi-turn text and code chat, and code generation. It can
             output text and code.
@@ -228,10 +163,10 @@ class GenerativeModelHook(GoogleBaseHook):
         tuned_model_display_name: str | None = None,
         validation_dataset: str | None = None,
         epochs: int | None = None,
-        adapter_size: int | None = None,
+        adapter_size: Literal[1, 4, 8, 16] | None = None,
         learning_rate_multiplier: float | None = None,
         project_id: str = PROVIDE_PROJECT_ID,
-    ) -> types_v1.TuningJob:
+    ) -> Any:
         """
         Use the Supervised Fine Tuning API to create a tuning job.
 
@@ -277,7 +212,7 @@ class GenerativeModelHook(GoogleBaseHook):
         self,
         contents: list,
         location: str,
-        pretrained_model: str = "gemini-pro",
+        pretrained_model: str,
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> types_v1beta1.CountTokensResponse:
         """
@@ -287,7 +222,7 @@ class GenerativeModelHook(GoogleBaseHook):
         :param location: Required. The ID of the Google Cloud location that the service belongs to.
         :param contents: Required. The multi-part content of a message that a user or a program
             gives to the generative model, in order to elicit a specific response.
-        :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
+        :param pretrained_model: Required. Model,
             supporting prompts with text-only input, including natural language
             tasks, multi-turn text and code chat, and code generation. It can
             output text and code.
@@ -364,8 +299,8 @@ class GenerativeModelHook(GoogleBaseHook):
         model_name: str,
         location: str,
         ttl_hours: float = 1,
-        system_instruction: str | None = None,
-        contents: list | None = None,
+        system_instruction: Any | None = None,
+        contents: list[Any] | None = None,
         display_name: str | None = None,
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> str:
@@ -424,3 +359,32 @@ class GenerativeModelHook(GoogleBaseHook):
         )
 
         return response.text
+
+
+class ExperimentRunHook(GoogleBaseHook):
+    """Use the Vertex AI SDK for Python to create and manage your experiment runs."""
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_experiment_run(
+        self,
+        experiment_run_name: str,
+        experiment_name: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        delete_backing_tensorboard_run: bool = False,
+    ) -> None:
+        """
+        Delete experiment run from the experiment.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param experiment_name: Required. The name of the evaluation experiment.
+        :param experiment_run_name: Required. The specific run name or ID for this experiment.
+        :param delete_backing_tensorboard_run: Whether to delete the backing Vertex AI TensorBoard run
+            that stores time series metrics for this run.
+        """
+        self.log.info("Next experiment run will be deleted: %s", experiment_run_name)
+        experiment_run = aiplatform.ExperimentRun(
+            run_name=experiment_run_name, experiment=experiment_name, project=project_id, location=location
+        )
+        experiment_run.delete(delete_backing_tensorboard_run=delete_backing_tensorboard_run)

airflow/providers/google/cloud/hooks/vertex_ai/ray.py

@@ -22,8 +22,17 @@ from __future__ import annotations
 import dataclasses
 from typing import Any
 
-import vertex_ray
-from google._upb._message import ScalarMapContainer
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+
+try:
+    import vertex_ray
+    from google._upb._message import ScalarMapContainer  # type: ignore[attr-defined]
+except ImportError:
+    # Fallback for environments where the upb module is not available.
+    raise AirflowOptionalProviderFeatureException(
+        "google._upb._message.ScalarMapContainer is not available. "
+        "Please install the ray package to use this feature."
+    )
 from google.cloud import aiplatform
 from google.cloud.aiplatform.vertex_ray.util import resources
 from google.cloud.aiplatform_v1 import (

airflow/providers/google/cloud/hooks/vision.py

@@ -19,10 +19,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from copy import deepcopy
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.vision_v1 import (

airflow/providers/google/cloud/links/alloy_db.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 ALLOY_DB_BASE_LINK = "/alloydb"
 ALLOY_DB_CLUSTER_LINK = (
     ALLOY_DB_BASE_LINK + "/locations/{location_id}/clusters/{cluster_id}?project={project_id}"
@@ -44,20 +38,6 @@ class AlloyDBClusterLink(BaseGoogleLink):
     key = "alloy_db_cluster"
     format_str = ALLOY_DB_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        cluster_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBClusterLink.key,
-            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
-        )
-
 
 class AlloyDBUsersLink(BaseGoogleLink):
     """Helper class for constructing AlloyDB users Link."""
@@ -66,20 +46,6 @@ class AlloyDBUsersLink(BaseGoogleLink):
     key = "alloy_db_users"
     format_str = ALLOY_DB_USERS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        cluster_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBUsersLink.key,
-            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
-        )
-
 
 class AlloyDBBackupsLink(BaseGoogleLink):
     """Helper class for constructing AlloyDB backups Link."""
@@ -87,15 +53,3 @@ class AlloyDBBackupsLink(BaseGoogleLink):
     name = "AlloyDB Backups"
     key = "alloy_db_backups"
     format_str = ALLOY_DB_BACKUPS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=AlloyDBBackupsLink.key,
-            value={"project_id": project_id},
-        )

airflow/providers/google/cloud/links/base.py

@@ -19,19 +19,23 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.providers.google.version_compat import (
+    AIRFLOW_V_3_0_PLUS,
+    BaseOperator,
+    BaseOperatorLink,
+    BaseSensorOperator,
+)
 
 if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
     from airflow.sdk.execution_time.xcom import XCom
 else:
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
     from airflow.models.xcom import XCom  # type: ignore[no-redef]
 
+if TYPE_CHECKING:
+    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+    from airflow.utils.context import Context
+
 BASE_LINK = "https://console.cloud.google.com"
 
 
@@ -39,6 +43,12 @@ class BaseGoogleLink(BaseOperatorLink):
     """
     Base class for all Google links.
 
+    When you inherit this class in a Link class;
+      - You can call the persist method to push data to the XCom to use it later in the get_link method.
+      - If you have an operator which inherit the GoogleCloudBaseOperator or BaseSensorOperator
+        You can define extra_links_params method in the operator to pass the operator properties
+        to the get_link method.
+
     :meta private:
     """
 
@@ -46,15 +56,69 @@ class BaseGoogleLink(BaseOperatorLink):
     key: ClassVar[str]
     format_str: ClassVar[str]
 
+    @property
+    def xcom_key(self) -> str:
+        # NOTE: in Airflow 3 we need to have xcom_key property in the Link class.
+        # Since we have the key property already, this is just a proxy property method to use same
+        # key as in Airflow 2.
+        return self.key
+
+    @classmethod
+    def persist(cls, context: Context, **value):
+        """
+        Push arguments to the XCom to use later for link formatting at the `get_link` method.
+
+        Note: for Airflow 2 we need to call this function with context variable only
+        where we have the extra_links_params property method defined
+        """
+        params = {}
+        # TODO: remove after Airflow v2 support dropped
+        if not AIRFLOW_V_3_0_PLUS:
+            common_params = getattr(context["task"], "extra_links_params", None)
+            if common_params:
+                params.update(common_params)
+
+        context["ti"].xcom_push(
+            key=cls.key,
+            value={
+                **params,
+                **value,
+            },
+        )
+
+    def get_config(self, operator, ti_key):
+        conf = {}
+        conf.update(getattr(operator, "extra_links_params", {}))
+        conf.update(XCom.get_value(key=self.key, ti_key=ti_key) or {})
+
+        # if the config did not define, return None to stop URL formatting
+        if not conf:
+            return None
+
+        # Add a default value for the 'namespace' parameter for backward compatibility.
+        # This is for datafusion
+        conf.setdefault("namespace", "default")
+        return conf
+
     def get_link(
         self,
         operator: BaseOperator,
         *,
         ti_key: TaskInstanceKey,
     ) -> str:
-        conf = XCom.get_value(key=self.key, ti_key=ti_key)
+        if TYPE_CHECKING:
+            assert isinstance(operator, (GoogleCloudBaseOperator, BaseSensorOperator))
+
+        conf = self.get_config(operator, ti_key)
         if not conf:
             return ""
-        if self.format_str.startswith("http"):
-            return self.format_str.format(**conf)
-        return BASE_LINK + self.format_str.format(**conf)
+        return self._format_link(**conf)
+
+    def _format_link(self, **kwargs):
+        try:
+            formatted_str = self.format_str.format(**kwargs)
+            if formatted_str.startswith("http"):
+                return formatted_str
+            return BASE_LINK + formatted_str
+        except KeyError:
+            return ""

airflow/providers/google/cloud/links/bigquery.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 BIGQUERY_BASE_LINK = "/bigquery"
 BIGQUERY_DATASET_LINK = (
     BIGQUERY_BASE_LINK + "?referrer=search&project={project_id}&d={dataset_id}&p={project_id}&page=dataset"
@@ -47,19 +41,6 @@ class BigQueryDatasetLink(BaseGoogleLink):
     key = "bigquery_dataset"
     format_str = BIGQUERY_DATASET_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        dataset_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDatasetLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id},
-        )
-
 
 class BigQueryTableLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Table Link."""
@@ -68,20 +49,6 @@ class BigQueryTableLink(BaseGoogleLink):
     key = "bigquery_table"
     format_str = BIGQUERY_TABLE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        table_id: str,
-        dataset_id: str | None = None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryTableLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id, "table_id": table_id},
-        )
-
 
 class BigQueryJobDetailLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Job Detail Link."""
@@ -89,17 +56,3 @@ class BigQueryJobDetailLink(BaseGoogleLink):
     name = "BigQuery Job Detail"
     key = "bigquery_job_detail"
     format_str = BIGQUERY_JOB_DETAIL_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        location: str,
-        job_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryJobDetailLink.key,
-            value={"project_id": project_id, "location": location, "job_id": job_id},
-        )

airflow/providers/google/cloud/links/bigquery_dts.py

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 BIGQUERY_BASE_LINK = "/bigquery/transfers"
 BIGQUERY_DTS_LINK = BIGQUERY_BASE_LINK + "/locations/{region}/configs/{config_id}/runs?project={project_id}"
 
@@ -37,17 +31,3 @@ class BigQueryDataTransferConfigLink(BaseGoogleLink):
     name = "BigQuery Data Transfer Config"
     key = "bigquery_dts_config"
     format_str = BIGQUERY_DTS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        region: str,
-        config_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDataTransferConfigLink.key,
-            value={"project_id": project_id, "region": region, "config_id": config_id},
-        )

airflow/providers/google/cloud/links/bigtable.py

@@ -16,13 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 BIGTABLE_BASE_LINK = "/bigtable"
 BIGTABLE_INSTANCE_LINK = BIGTABLE_BASE_LINK + "/instances/{instance_id}/overview?project={project_id}"
 BIGTABLE_CLUSTER_LINK = (
@@ -38,20 +33,6 @@ class BigtableInstanceLink(BaseGoogleLink):
     key = "instance_key"
     format_str = BIGTABLE_INSTANCE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableInstanceLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableClusterLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Cluster link."""
@@ -60,21 +41,6 @@ class BigtableClusterLink(BaseGoogleLink):
     key = "cluster_key"
     format_str = BIGTABLE_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableClusterLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "cluster_id": task_instance.cluster_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableTablesLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Tables link."""
@@ -82,17 +48,3 @@ class BigtableTablesLink(BaseGoogleLink):
     name = "Bigtable Tables"
     key = "tables_key"
     format_str = BIGTABLE_TABLES_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableTablesLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )