apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.3.0__py3-none-any.whl

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

@@ -21,14 +21,15 @@ 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
@@ -36,23 +37,12 @@ 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 +51,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 +83,18 @@ 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",
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.embed_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,
@@ -177,16 +119,21 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.values
 
+    @deprecated(
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.generate_content",
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def generative_model_generate_content(
         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 +147,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.
@@ -219,6 +166,11 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.text
 
+    @deprecated(
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.supervised_fine_tuning_train",
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def supervised_fine_tuning_train(
         self,
@@ -228,10 +180,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.
 
@@ -272,12 +224,17 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return sft_tuning_job
 
+    @deprecated(
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.count_tokens",
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def count_tokens(
         self,
         contents: list,
         location: str,
-        pretrained_model: str = "gemini-pro",
+        pretrained_model: str,
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> types_v1beta1.CountTokensResponse:
         """
@@ -287,7 +244,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.
@@ -359,13 +316,18 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return eval_result
 
+    @deprecated(
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.create_cached_content",
+        category=AirflowProviderDeprecationWarning,
+    )
     def create_cached_content(
         self,
         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:
@@ -393,6 +355,11 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.name
 
+    @deprecated(
+        planned_removal_date="January 3, 2026",
+        use_instead="airflow.providers.google.cloud.hooks.gen_ai.generative_model.GenAIGenerativeModelHook.generate_content",
+        category=AirflowProviderDeprecationWarning,
+    )
     def generate_from_cached_content(
         self,
         location: str,
@@ -413,6 +380,9 @@ class GenerativeModelHook(GoogleBaseHook):
         :param generation_config: Optional. Generation configuration settings.
         :param safety_settings: Optional. Per request settings for blocking unsafe content.
         """
+        # During run of the system test it was found out that names from xcom, e.g. 3402922389 can be
+        # treated as int and throw an error TypeError: expected string or bytes-like object, got 'int'
+        cached_content_name = str(cached_content_name)
         vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
 
         cached_context_model = self.get_cached_context_model(cached_content_name=cached_content_name)
@@ -424,3 +394,37 @@ class GenerativeModelHook(GoogleBaseHook):
         )
 
         return response.text
+
+
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.hooks.vertex_ai.experiment_service.ExperimentRunHook",
+    category=AirflowProviderDeprecationWarning,
+)
+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/hyperparameter_tuning_job.py

@@ -34,7 +34,7 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

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

@@ -27,7 +27,7 @@ from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import ModelServiceClient
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 

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

@@ -39,7 +39,7 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from airflow.providers.google.common.hooks.operation_helpers import OperationHelper

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

@@ -0,0 +1,223 @@
+#
+# 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.
+"""This module contains a Google Cloud Vertex AI hook."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import MutableMapping
+from typing import Any
+
+import vertex_ray
+from google.cloud import aiplatform
+from google.cloud.aiplatform.vertex_ray.util import resources
+from google.cloud.aiplatform_v1 import (
+    PersistentResourceServiceClient,
+)
+from proto.marshal.collections.repeated import Repeated
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class RayHook(GoogleBaseHook):
+    """Hook for Google Cloud Vertex AI Ray APIs."""
+
+    def extract_cluster_id(self, cluster_path) -> str:
+        """Extract cluster_id from cluster_path."""
+        cluster_id = PersistentResourceServiceClient.parse_persistent_resource_path(cluster_path)[
+            "persistent_resource"
+        ]
+        return cluster_id
+
+    def serialize_cluster_obj(self, cluster_obj: resources.Cluster) -> dict:
+        """Serialize Cluster dataclass to dict."""
+
+        def __encode_value(value: Any) -> Any:
+            if isinstance(value, (list, Repeated)):
+                return [__encode_value(nested_value) for nested_value in value]
+            if not isinstance(value, dict) and isinstance(value, MutableMapping):
+                return {key: __encode_value(nested_value) for key, nested_value in dict(value).items()}
+            if dataclasses.is_dataclass(value) and not isinstance(value, type):
+                return dataclasses.asdict(value)
+            return value
+
+        return {
+            field.name: __encode_value(getattr(cluster_obj, field.name))
+            for field in dataclasses.fields(cluster_obj)
+        }
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        head_node_type: resources.Resources = resources.Resources(),
+        python_version: str = "3.10",
+        ray_version: str = "2.33",
+        network: str | None = None,
+        service_account: str | None = None,
+        cluster_name: str | None = None,
+        worker_node_types: list[resources.Resources] | None = None,
+        custom_images: resources.NodeImages | None = None,
+        enable_metrics_collection: bool = True,
+        enable_logging: bool = True,
+        psc_interface_config: resources.PscIConfig | None = None,
+        reserved_ip_ranges: list[str] | None = None,
+        labels: dict[str, str] | None = None,
+    ) -> str:
+        """
+        Create a Ray cluster on the Vertex AI.
+
+        :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 head_node_type: The head node resource. Resources.node_count must be 1. If not set, default
+            value of Resources() class will be used.
+        :param python_version: Python version for the ray cluster.
+        :param ray_version: Ray version for the ray cluster. Default is 2.33.0.
+        :param network: Virtual private cloud (VPC) network. For Ray Client, VPC peering is required to
+            connect to the Ray Cluster managed in the Vertex API service. For Ray Job API, VPC network is not
+            required because Ray Cluster connection can be accessed through dashboard address.
+        :param service_account: Service account to be used for running Ray programs on the cluster.
+        :param cluster_name: This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`.
+            The first character cannot be a number or hyphen.
+        :param worker_node_types: The list of Resources of the worker nodes. The same Resources object should
+            not appear multiple times in the list.
+        :param custom_images: The NodeImages which specifies head node and worker nodes images. All the
+            workers will share the same image. If each Resource has a specific custom image, use
+            `Resources.custom_image` for head/worker_node_type(s). Note that configuring
+            `Resources.custom_image` will override `custom_images` here. Allowlist only.
+        :param enable_metrics_collection: Enable Ray metrics collection for visualization.
+        :param enable_logging: Enable exporting Ray logs to Cloud Logging.
+        :param psc_interface_config: PSC-I config.
+        :param reserved_ip_ranges: A list of names for the reserved IP ranges under the VPC network that can
+            be used for this cluster. If set, we will deploy the cluster within the provided IP ranges.
+            Otherwise, the cluster is deployed to any IP ranges under the provided VPC network.
+            Example: ["vertex-ai-ip-range"].
+        :param labels: The labels with user-defined metadata to organize Ray cluster.
+            Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain
+            lowercase letters, numeric characters, underscores and dashes. International characters are allowed.
+            See https://goo.gl/xmQnxf for more information and examples of labels.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        cluster_path = vertex_ray.create_ray_cluster(
+            head_node_type=head_node_type,
+            python_version=python_version,
+            ray_version=ray_version,
+            network=network,
+            service_account=service_account,
+            cluster_name=cluster_name,
+            worker_node_types=worker_node_types,
+            custom_images=custom_images,
+            enable_metrics_collection=enable_metrics_collection,
+            enable_logging=enable_logging,
+            psc_interface_config=psc_interface_config,
+            reserved_ip_ranges=reserved_ip_ranges,
+            labels=labels,
+        )
+        return cluster_path
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def list_ray_clusters(
+        self,
+        project_id: str,
+        location: str,
+    ) -> list[resources.Cluster]:
+        """
+        List Ray clusters under the currently authenticated project.
+
+        :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.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_clusters = vertex_ray.list_ray_clusters()
+        return ray_clusters
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+    ) -> resources.Cluster:
+        """
+        Get Ray cluster.
+
+        :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 cluster_id: Cluster resource ID.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        ray_cluster = vertex_ray.get_ray_cluster(
+            cluster_resource_name=ray_cluster_name,
+        )
+        return ray_cluster
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def update_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+        worker_node_types: list[resources.Resources],
+    ) -> str:
+        """
+        Update Ray cluster (currently support resizing node counts for worker nodes).
+
+        :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 cluster_id: Cluster resource ID.
+        :param worker_node_types: The list of Resources of the resized worker nodes. The same Resources
+            object should not appear multiple times in the list.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        updated_ray_cluster_name = vertex_ray.update_ray_cluster(
+            cluster_resource_name=ray_cluster_name, worker_node_types=worker_node_types
+        )
+        return updated_ray_cluster_name
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_ray_cluster(
+        self,
+        project_id: str,
+        location: str,
+        cluster_id: str,
+    ) -> None:
+        """
+        Delete Ray cluster.
+
+        :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 cluster_id: Cluster resource ID.
+        """
+        aiplatform.init(project=project_id, location=location, credentials=self.get_credentials())
+        ray_cluster_name = PersistentResourceServiceClient.persistent_resource_path(
+            project=project_id,
+            location=location,
+            persistent_resource=cluster_id,
+        )
+        vertex_ray.delete_ray_cluster(cluster_resource_name=ray_cluster_name)

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 (
@@ -36,7 +36,7 @@ from google.cloud.vision_v1 import (
 )
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 

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

@@ -65,7 +65,7 @@ class WorkflowsHook(GoogleBaseHook):
         Create a new workflow.
 
         If a workflow with the specified name already exists in the
-        specified project and location, the long running operation will
+        specified project and location, the long-running operation will
         return [ALREADY_EXISTS][google.rpc.Code.ALREADY_EXISTS] error.
 
         :param workflow: Required. Workflow to be created.

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},
-        )