apache-airflow-providers-google 18.0.0rc1__py3-none-any.whl → 18.1.0__py3-none-any.whl

airflow/providers/google/cloud/operators/datafusion.py

@@ -40,6 +40,7 @@ from airflow.providers.google.cloud.utils.helpers import resource_path_to_dict
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -777,6 +778,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         self.pipeline_timeout = pipeline_timeout
         self.deferrable = deferrable
         self.poll_interval = poll_interval
+        self.pipeline_id: str | None = None
 
         if success_states:
             self.success_states = success_states
@@ -796,14 +798,14 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
             project_id=self.project_id,
         )
         api_url = instance["apiEndpoint"]
-        pipeline_id = hook.start_pipeline(
+        self.pipeline_id = hook.start_pipeline(
             pipeline_name=self.pipeline_name,
             pipeline_type=self.pipeline_type,
             instance_url=api_url,
             namespace=self.namespace,
             runtime_args=self.runtime_args,
         )
-        self.log.info("Pipeline %s submitted successfully.", pipeline_id)
+        self.log.info("Pipeline %s submitted successfully.", self.pipeline_id)
 
         DataFusionPipelineLink.persist(
             context=context,
@@ -824,7 +826,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
                     namespace=self.namespace,
                     pipeline_name=self.pipeline_name,
                     pipeline_type=self.pipeline_type.value,
-                    pipeline_id=pipeline_id,
+                    pipeline_id=self.pipeline_id,
                     poll_interval=self.poll_interval,
                     gcp_conn_id=self.gcp_conn_id,
                     impersonation_chain=self.impersonation_chain,
@@ -834,19 +836,21 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         else:
             if not self.asynchronous:
                 # when NOT using asynchronous mode it will just wait for pipeline to finish and print message
-                self.log.info("Waiting when pipeline %s will be in one of the success states", pipeline_id)
+                self.log.info(
+                    "Waiting when pipeline %s will be in one of the success states", self.pipeline_id
+                )
                 hook.wait_for_pipeline_state(
                     success_states=self.success_states,
-                    pipeline_id=pipeline_id,
+                    pipeline_id=self.pipeline_id,
                     pipeline_name=self.pipeline_name,
                     pipeline_type=self.pipeline_type,
                     namespace=self.namespace,
                     instance_url=api_url,
                     timeout=self.pipeline_timeout,
                 )
-                self.log.info("Pipeline %s discovered success state.", pipeline_id)
+                self.log.info("Pipeline %s discovered success state.", self.pipeline_id)
             #  otherwise, return pipeline_id so that sensor can use it later to check the pipeline state
-        return pipeline_id
+        return self.pipeline_id
 
     def execute_complete(self, context: Context, event: dict[str, Any]):
         """
@@ -863,6 +867,31 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         )
         return event["pipeline_id"]
 
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """Build and return OpenLineage facets and datasets for the completed pipeline start."""
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.facets import DataFusionRunFacet
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        pipeline_resource = f"{self.project_id}:{self.location}:{self.instance_name}:{self.pipeline_name}"
+
+        inputs = [Dataset(namespace="datafusion", name=pipeline_resource)]
+
+        if self.pipeline_id:
+            output_name = f"{pipeline_resource}:{self.pipeline_id}"
+        else:
+            output_name = f"{pipeline_resource}:unknown"
+        outputs = [Dataset(namespace="datafusion", name=output_name)]
+
+        run_facets = {
+            "dataFusionRun": DataFusionRunFacet(
+                runId=self.pipeline_id,
+                runtimeArgs=self.runtime_args,
+            )
+        }
+
+        return OperatorLineage(inputs=inputs, outputs=outputs, run_facets=run_facets, job_facets={})
+
 
 class CloudDataFusionStopPipelineOperator(GoogleCloudBaseOperator):
     """

airflow/providers/google/cloud/operators/gen_ai.py

@@ -0,0 +1,389 @@
+#
+# 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 Google Gen AI operators."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.google.cloud.hooks.gen_ai import (
+    GenAIGenerativeModelHook,
+)
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from google.genai.types import (
+        ContentListUnion,
+        ContentListUnionDict,
+        CountTokensConfigOrDict,
+        CreateCachedContentConfigOrDict,
+        CreateTuningJobConfigOrDict,
+        EmbedContentConfigOrDict,
+        GenerateContentConfig,
+        TuningDatasetOrDict,
+    )
+
+    from airflow.utils.context import Context
+
+
+class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator):
+    """
+    Uses the Gemini AI Embeddings API to generate embeddings for words, phrases, sentences, and code.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to (templated).
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to (templated).
+    :param model: Required. The name of the model to use for content generation,
+        which can be a text-only or multimodal model. For example, `gemini-pro` or
+        `gemini-pro-vision`.
+    :param contents: Optional. The contents to use for embedding.
+    :param config: Optional. Configuration for embeddings.
+    :param gcp_conn_id: Optional. The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional. Service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("location", "project_id", "impersonation_chain", "contents", "model", "config")
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        model: str,
+        contents: ContentListUnion | ContentListUnionDict | list[str],
+        config: EmbedContentConfigOrDict | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.contents = contents
+        self.config = config
+        self.model = model
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenAIGenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Generating text embeddings...")
+        response = self.hook.embed_content(
+            project_id=self.project_id,
+            location=self.location,
+            contents=self.contents,
+            model=self.model,
+            config=self.config,
+        )
+
+        self.log.info("Model response: %s", response)
+        context["ti"].xcom_push(key="model_response", value=response)
+
+        return response
+
+
+class GenAIGenerateContentOperator(GoogleCloudBaseOperator):
+    """
+    Generate a model response based on given configuration. Input capabilities differ between models, including tuned models.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to (templated).
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to (templated).
+    :param model: Required. The name of the model to use for content generation,
+        which can be a text-only or multimodal model. For example, `gemini-pro` or
+        `gemini-pro-vision`.
+    :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 generation_config: Optional. Generation configuration settings.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = (
+        "generation_config",
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "contents",
+        "model",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        contents: ContentListUnionDict,
+        model: str,
+        generation_config: GenerateContentConfig | dict[str, Any] | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.contents = contents
+        self.generation_config = generation_config
+        self.model = model
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenAIGenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.generate_content(
+            project_id=self.project_id,
+            location=self.location,
+            model=self.model,
+            contents=self.contents,
+            generation_config=self.generation_config,
+        )
+
+        self.log.info("Created Content: %s", response)
+        context["ti"].xcom_push(key="model_response", value=response)
+
+        return response
+
+
+class GenAISupervisedFineTuningTrainOperator(GoogleCloudBaseOperator):
+    """
+    Create a tuning job to adapt model behavior with a labeled dataset.
+
+    :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 source_model: Required. A pre-trained model optimized for performing natural
+        language tasks such as classification, summarization, extraction, content
+        creation, and ideation.
+    :param training_dataset: Required. Cloud Storage URI of your training dataset. The dataset
+        must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples.
+    :param tuning_job_config: Optional. Configuration of the Tuning job to be created.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "training_dataset",
+        "tuning_job_config",
+        "source_model",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        source_model: str,
+        training_dataset: TuningDatasetOrDict,
+        tuning_job_config: CreateTuningJobConfigOrDict | dict[str, Any] | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.source_model = source_model
+        self.training_dataset = training_dataset
+        self.tuning_job_config = tuning_job_config
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenAIGenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.supervised_fine_tuning_train(
+            project_id=self.project_id,
+            location=self.location,
+            source_model=self.source_model,
+            training_dataset=self.training_dataset,
+            tuning_job_config=self.tuning_job_config,
+        )
+
+        self.log.info("Tuned Model Name: %s", response.tuned_model.model)  # type: ignore[union-attr,arg-type]
+        self.log.info("Tuned Model EndpointName: %s", response.tuned_model.endpoint)  # type: ignore[union-attr,arg-type]
+
+        context["ti"].xcom_push(key="tuned_model_name", value=response.tuned_model.model)  # type: ignore[union-attr,arg-type]
+        context["ti"].xcom_push(key="tuned_model_endpoint_name", value=response.tuned_model.endpoint)  # type: ignore[union-attr,arg-type]
+
+        result = {
+            "tuned_model_name": response.tuned_model.model,  # type: ignore[union-attr,arg-type]
+            "tuned_model_endpoint_name": response.tuned_model.endpoint,  # type: ignore[union-attr,arg-type]
+        }
+
+        return result
+
+
+class GenAICountTokensOperator(GoogleCloudBaseOperator):
+    """
+    Use Count Tokens API to calculate the number of input tokens before sending a request to Gemini API.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to (templated).
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to (templated).
+    :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 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.
+    :param config: Optional. Configuration for Count Tokens.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("location", "project_id", "impersonation_chain", "contents", "model", "config")
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        contents: ContentListUnion | ContentListUnionDict,
+        model: str,
+        config: CountTokensConfigOrDict | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.contents = contents
+        self.model = model
+        self.config = config
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenAIGenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.count_tokens(
+            project_id=self.project_id,
+            location=self.location,
+            contents=self.contents,
+            model=self.model,
+            config=self.config,
+        )
+
+        self.log.info("Total tokens: %s", response.total_tokens)
+        context["ti"].xcom_push(key="total_tokens", value=response.total_tokens)
+
+
+class GenAICreateCachedContentOperator(GoogleCloudBaseOperator):
+    """
+    Create CachedContent resource to reduce the cost of requests that contain repeat content with high input token counts.
+
+    :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 model: Required. The name of the publisher model to use for cached content.
+    :param cached_content_config: Optional. Configuration of the Cached Content.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = ("location", "project_id", "impersonation_chain", "model", "cached_content_config")
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        model: str,
+        cached_content_config: CreateCachedContentConfigOrDict | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.project_id = project_id
+        self.location = location
+        self.model = model
+        self.cached_content_config = cached_content_config
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenAIGenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        cached_content_name = self.hook.create_cached_content(
+            project_id=self.project_id,
+            location=self.location,
+            model=self.model,
+            cached_content_config=self.cached_content_config,
+        )
+
+        self.log.info("Cached Content Name: %s", cached_content_name)
+        context["ti"].xcom_push(key="cached_content", value=cached_content_name)
+
+        return cached_content_name

airflow/providers/google/cloud/operators/spanner.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
@@ -29,6 +30,7 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -254,6 +256,13 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @cached_property
+    def hook(self) -> SpannerHook:
+        return SpannerHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
     def _validate_inputs(self) -> None:
         if self.project_id == "":
             raise AirflowException("The required parameter 'project_id' is empty")
@@ -265,10 +274,6 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
             raise AirflowException("The required parameter 'query' is empty")
 
     def execute(self, context: Context):
-        hook = SpannerHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
         if isinstance(self.query, str):
             queries = [x.strip() for x in self.query.split(";")]
             self.sanitize_queries(queries)
@@ -281,7 +286,7 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
             self.database_id,
         )
         self.log.info("Executing queries: %s", queries)
-        result_rows_count_per_query = hook.execute_dml(
+        result_rows_count_per_query = self.hook.execute_dml(
             project_id=self.project_id,
             instance_id=self.instance_id,
             database_id=self.database_id,
@@ -291,7 +296,7 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
             context=context,
             instance_id=self.instance_id,
             database_id=self.database_id,
-            project_id=self.project_id or hook.project_id,
+            project_id=self.project_id or self.hook.project_id,
         )
         return result_rows_count_per_query
 
@@ -305,6 +310,17 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         if queries and queries[-1] == "":
             queries.pop()
 
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """Build a generic OpenLineage facet, aligned with SQL-based operators."""
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+
+        return get_openlineage_facets_with_sql(
+            hook=self.hook,
+            sql=self.query,
+            conn_id=self.gcp_conn_id,
+            database=self.database_id,
+        )
+
 
 class SpannerDeployDatabaseInstanceOperator(GoogleCloudBaseOperator):
     """

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

@@ -29,6 +29,7 @@ from google.cloud.aiplatform import datasets
 from google.cloud.aiplatform.models import Model
 from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
@@ -36,6 +37,7 @@ from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAITrainingPipelinesLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.deprecated import deprecated
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -473,6 +475,11 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="March 24, 2026",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Video Training job."""
 

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

@@ -36,6 +36,11 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAIGenerateEmbeddingsOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class TextEmbeddingModelGetEmbeddingsOperator(GoogleCloudBaseOperator):
     """
     Uses the Vertex AI Embeddings API to generate embeddings based on prompt.
@@ -99,6 +104,11 @@ class TextEmbeddingModelGetEmbeddingsOperator(GoogleCloudBaseOperator):
         return response
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAIGenerateContentOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator):
     """
     Use the Vertex AI Gemini Pro foundation model to generate content.
@@ -178,6 +188,11 @@ class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator):
         return response
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAISupervisedFineTuningTrainOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator):
     """
     Use the Supervised Fine Tuning API to create a tuning job.
@@ -280,6 +295,11 @@ class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator):
         return result
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAICountTokensOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class CountTokensOperator(GoogleCloudBaseOperator):
     """
     Use the Vertex AI Count Tokens API to calculate the number of input tokens before sending a request to the Gemini API.
@@ -443,6 +463,11 @@ class RunEvaluationOperator(GoogleCloudBaseOperator):
         return response.summary_metrics
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAICreateCachedContentOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class CreateCachedContentOperator(GoogleCloudBaseOperator):
     """
     Create CachedContent to reduce the cost of requests that contain repeat content with high input token counts.
@@ -522,6 +547,11 @@ class CreateCachedContentOperator(GoogleCloudBaseOperator):
         return cached_content_name
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.gen_ai.generative_model.GenAIGenerateContentOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class GenerateFromCachedContentOperator(GoogleCloudBaseOperator):
     """
     Generate a response from CachedContent.