apache-airflow-providers-google 10.16.0rc1__py3-none-any.whl → 10.17.0rc1__py3-none-any.whl

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

@@ -20,6 +20,7 @@
     ImportSshPublicKeyResponse
     oslogin
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -23,6 +23,7 @@ This module contains a Google Pub/Sub Hook.
     MessageStoragePolicy
     ReceivedMessage
 """
+
 from __future__ import annotations
 
 import warnings

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Secrets Manager service."""
+
 from __future__ import annotations
 
 from typing import Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Spanner Hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Callable, NamedTuple, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Speech Hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Stackdriver operators."""
+
 from __future__ import annotations
 
 import contextlib

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Text to Speech Hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Translate Hook."""
+
 from __future__ import annotations
 
 from typing import Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vertex AI hook."""
+
 from __future__ import annotations
 
 import warnings

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

@@ -16,21 +16,24 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vertex AI hook."""
+
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import BatchPredictionJob, Model, explain
-from google.cloud.aiplatform_v1 import JobServiceClient
+from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
 from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
-    from google.api_core.retry import Retry
+    from google.api_core.retry import AsyncRetry, Retry
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager
 
 
@@ -241,6 +244,159 @@ class BatchPredictionJobHook(GoogleBaseHook):
         )
         return self._batch_prediction_job
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def submit_batch_prediction_job(
+        self,
+        project_id: str,
+        region: str,
+        job_display_name: str,
+        model_name: str | Model,
+        instances_format: str = "jsonl",
+        predictions_format: str = "jsonl",
+        gcs_source: str | Sequence[str] | None = None,
+        bigquery_source: str | None = None,
+        gcs_destination_prefix: str | None = None,
+        bigquery_destination_prefix: str | None = None,
+        model_parameters: dict | None = None,
+        machine_type: str | None = None,
+        accelerator_type: str | None = None,
+        accelerator_count: int | None = None,
+        starting_replica_count: int | None = None,
+        max_replica_count: int | None = None,
+        generate_explanation: bool | None = False,
+        explanation_metadata: explain.ExplanationMetadata | None = None,
+        explanation_parameters: explain.ExplanationParameters | None = None,
+        labels: dict[str, str] | None = None,
+        encryption_spec_key_name: str | None = None,
+        create_request_timeout: float | None = None,
+        batch_size: int | None = None,
+    ) -> BatchPredictionJob:
+        """
+        Create a batch prediction job.
+
+        :param project_id: Required. Project to run training in.
+        :param region: Required. Location to run training in.
+        :param job_display_name: Required. The user-defined name of the BatchPredictionJob. The name can be
+            up to 128 characters long and can be consist of any UTF-8 characters.
+        :param model_name: Required. A fully-qualified model resource name or model ID.
+        :param instances_format: Required. The format in which instances are provided. Must be one of the
+            formats listed in `Model.supported_input_storage_formats`. Default is "jsonl" when using
+            `gcs_source`. If a `bigquery_source` is provided, this is overridden to "bigquery".
+        :param predictions_format: Required. The format in which Vertex AI outputs the predictions, must be
+            one of the formats specified in `Model.supported_output_storage_formats`. Default is "jsonl" when
+            using `gcs_destination_prefix`. If a `bigquery_destination_prefix` is provided, this is
+            overridden to "bigquery".
+        :param gcs_source: Google Cloud Storage URI(-s) to your instances to run batch prediction on. They
+            must match `instances_format`. May contain wildcards. For more information on wildcards, see
+            https://cloud.google.com/storage/docs/gsutil/addlhelp/WildcardNames.
+        :param bigquery_source: BigQuery URI to a table, up to 2000 characters long.
+            For example: `bq://projectId.bqDatasetId.bqTableId`
+        :param gcs_destination_prefix: The Google Cloud Storage location of the directory where the output is
+            to be written to. In the given directory a new directory is created. Its name is
+            ``prediction-<model-display-name>-<job-create-time>``, where timestamp is in
+            YYYY-MM-DDThh:mm:ss.sssZ ISO-8601 format. Inside of it files ``predictions_0001.<extension>``,
+            ``predictions_0002.<extension>``, ..., ``predictions_N.<extension>`` are created where
+            ``<extension>`` depends on chosen ``predictions_format``, and N may equal 0001 and depends on the
+            total number of successfully predicted instances. If the Model has both ``instance`` and
+            ``prediction`` schemata defined then each such file contains predictions as per the
+            ``predictions_format``. If prediction for any instance failed (partially or completely), then an
+            additional ``errors_0001.<extension>``, ``errors_0002.<extension>``,..., ``errors_N.<extension>``
+            files are created (N depends on total number of failed predictions). These files contain the
+            failed instances, as per their schema, followed by an additional ``error`` field which as value
+            has ```google.rpc.Status`` <Status>`__ containing only ``code`` and ``message`` fields.
+        :param bigquery_destination_prefix: The BigQuery project location where the output is to be written
+            to. In the given project a new dataset is created with name
+            ``prediction_<model-display-name>_<job-create-time>`` where is made BigQuery-dataset-name
+            compatible (for example, most special characters become underscores), and timestamp is in
+            YYYY_MM_DDThh_mm_ss_sssZ "based on ISO-8601" format. In the dataset two tables will be created,
+            ``predictions``, and ``errors``. If the Model has both ``instance`` and ``prediction`` schemata
+            defined then the tables have columns as follows: The ``predictions`` table contains instances for
+            which the prediction succeeded, it has columns as per a concatenation of the Model's instance and
+            prediction schemata. The ``errors`` table contains rows for which the prediction has failed, it
+            has instance columns, as per the instance schema, followed by a single "errors" column, which as
+            values has ```google.rpc.Status`` <Status>`__ represented as a STRUCT, and containing only
+            ``code`` and ``message``.
+        :param model_parameters: The parameters that govern the predictions. The schema of the parameters may
+            be specified via the Model's `parameters_schema_uri`.
+        :param machine_type: The type of machine for running batch prediction on dedicated resources. Not
+            specifying machine type will result in batch prediction job being run with automatic resources.
+        :param accelerator_type: The type of accelerator(s) that may be attached to the machine as per
+            `accelerator_count`. Only used if `machine_type` is set.
+        :param accelerator_count: The number of accelerators to attach to the `machine_type`. Only used if
+            `machine_type` is set.
+        :param starting_replica_count: The number of machine replicas used at the start of the batch
+            operation. If not set, Vertex AI decides starting number, not greater than `max_replica_count`.
+            Only used if `machine_type` is set.
+        :param max_replica_count: The maximum number of machine replicas the batch operation may be scaled
+            to. Only used if `machine_type` is set. Default is 10.
+        :param generate_explanation: Optional. Generate explanation along with the batch prediction results.
+            This will cause the batch prediction output to include explanations based on the
+            `prediction_format`:
+            - `bigquery`: output includes a column named `explanation`. The value is a struct that conforms
+            to the [aiplatform.gapic.Explanation] object.
+            - `jsonl`: The JSON objects on each line include an additional entry keyed `explanation`. The
+            value of the entry is a JSON object that conforms to the [aiplatform.gapic.Explanation] object.
+            - `csv`: Generating explanations for CSV format is not supported.
+        :param explanation_metadata: Optional. Explanation metadata configuration for this
+            BatchPredictionJob. Can be specified only if `generate_explanation` is set to `True`.
+            This value overrides the value of `Model.explanation_metadata`. All fields of
+            `explanation_metadata` are optional in the request. If a field of the `explanation_metadata`
+            object is not populated, the corresponding field of the `Model.explanation_metadata` object is
+            inherited. For more details, see `Ref docs <http://tinyurl.com/1igh60kt>`
+        :param explanation_parameters: Optional. Parameters to configure explaining for Model's predictions.
+            Can be specified only if `generate_explanation` is set to `True`.
+            This value overrides the value of `Model.explanation_parameters`. All fields of
+            `explanation_parameters` are optional in the request. If a field of the `explanation_parameters`
+            object is not populated, the corresponding field of the `Model.explanation_parameters` object is
+            inherited. For more details, see `Ref docs <http://tinyurl.com/1an4zake>`
+        :param labels: Optional. The labels with user-defined metadata to organize your BatchPredictionJobs.
+            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.
+        :param encryption_spec_key_name: Optional. The Cloud KMS resource identifier of the customer managed
+            encryption key used to protect the job. Has the form:
+            ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. The key needs to be
+            in the same region as where the compute resource is created.
+            If this is set, then all resources created by the BatchPredictionJob will be encrypted with the
+            provided encryption key.
+            Overrides encryption_spec_key_name set in aiplatform.init.
+        :param create_request_timeout: Optional. The timeout for the create request in seconds.
+        :param batch_size: Optional. The number of the records (e.g. instances)
+            of the operation given in each batch
+            to a machine replica. Machine type, and size of a single record should be considered
+            when setting this parameter, higher value speeds up the batch operation's execution,
+            but too high value will result in a whole batch not fitting in a machine's memory,
+            and the whole operation will fail.
+            The default value is same as in the aiplatform's BatchPredictionJob.
+        """
+        self._batch_prediction_job = BatchPredictionJob.submit(
+            job_display_name=job_display_name,
+            model_name=model_name,
+            instances_format=instances_format,
+            predictions_format=predictions_format,
+            gcs_source=gcs_source,
+            bigquery_source=bigquery_source,
+            gcs_destination_prefix=gcs_destination_prefix,
+            bigquery_destination_prefix=bigquery_destination_prefix,
+            model_parameters=model_parameters,
+            machine_type=machine_type,
+            accelerator_type=accelerator_type,
+            accelerator_count=accelerator_count,
+            starting_replica_count=starting_replica_count,
+            max_replica_count=max_replica_count,
+            generate_explanation=generate_explanation,
+            explanation_metadata=explanation_metadata,
+            explanation_parameters=explanation_parameters,
+            labels=labels,
+            project=project_id,
+            location=region,
+            credentials=self.get_credentials(),
+            encryption_spec_key_name=encryption_spec_key_name,
+            create_request_timeout=create_request_timeout,
+            batch_size=batch_size,
+        )
+        return self._batch_prediction_job
+
     @GoogleBaseHook.fallback_to_default_project_id
     def delete_batch_prediction_job(
         self,
@@ -358,3 +514,99 @@ class BatchPredictionJobHook(GoogleBaseHook):
             metadata=metadata,
         )
         return result
+
+
+class BatchPredictionJobAsyncHook(GoogleBaseAsyncHook):
+    """Hook for Google Cloud Vertex AI Batch Prediction Job Async APIs."""
+
+    sync_hook_class = BatchPredictionJobHook
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            gcp_conn_id=gcp_conn_id,
+            impersonation_chain=impersonation_chain,
+            **kwargs,
+        )
+
+    async def get_job_service_client(self, region: str | None = None) -> JobServiceAsyncClient:
+        """Return JobServiceAsyncClient object."""
+        endpoint = f"{region}-aiplatform.googleapis.com:443" if region and region != "global" else None
+        return JobServiceAsyncClient(
+            credentials=(await self.get_sync_hook()).get_credentials(),
+            client_info=CLIENT_INFO,
+            client_options=ClientOptions(api_endpoint=endpoint),
+        )
+
+    async def get_batch_prediction_job(
+        self,
+        project_id: str,
+        location: str,
+        job_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> types.BatchPredictionJob:
+        """Retrieve a batch prediction tuning job.
+
+        :param project_id: Required. The ID of the Google Cloud project that the job belongs to.
+        :param location: Required. The ID of the Google Cloud region that the job belongs to.
+        :param job_id: Required. The hyperparameter tuning job id.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+        """
+        client: JobServiceAsyncClient = await self.get_job_service_client(region=location)
+        job_name = client.batch_prediction_job_path(project_id, location, job_id)
+
+        result = await client.get_batch_prediction_job(
+            request={
+                "name": job_name,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    async def wait_batch_prediction_job(
+        self,
+        project_id: str,
+        location: str,
+        job_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        poll_interval: int = 10,
+    ) -> types.BatchPredictionJob:
+        statuses_complete = {
+            JobState.JOB_STATE_CANCELLED,
+            JobState.JOB_STATE_FAILED,
+            JobState.JOB_STATE_PAUSED,
+            JobState.JOB_STATE_SUCCEEDED,
+        }
+        while True:
+            try:
+                self.log.info("Requesting batch prediction tuning job with id %s", job_id)
+                job: types.BatchPredictionJob = await self.get_batch_prediction_job(
+                    project_id=project_id,
+                    location=location,
+                    job_id=job_id,
+                    retry=retry,
+                    timeout=timeout,
+                    metadata=metadata,
+                )
+            except Exception as ex:
+                self.log.exception("Exception occurred while requesting job %s", job_id)
+                raise AirflowException(ex)
+
+            self.log.info("Status of the batch prediction job %s is %s", job.name, job.state.name)
+            if job.state in statuses_complete:
+                return job
+
+            self.log.info("Sleeping for %s seconds.", poll_interval)
+            await asyncio.sleep(poll_interval)

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vertex AI hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vertex AI hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vertex AI hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -0,0 +1,197 @@
+#
+# 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 Generative Model hook."""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import vertexai
+from vertexai.generative_models import GenerativeModel, Part
+from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
+
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+
+class GenerativeModelHook(GoogleBaseHook):
+    """Hook for Google Cloud Vertex AI Generative Model APIs."""
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        if kwargs.get("delegate_to") is not None:
+            raise RuntimeError(
+                "The `delegate_to` parameter has been deprecated before and finally removed in this version"
+                " of Google Provider. You MUST convert it to `impersonate_chain`"
+            )
+        super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs)
+
+    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)
+        return model
+
+    def get_generative_model(self, pretrained_model: str) -> GenerativeModel:
+        """Return a Generative Model object."""
+        model = GenerativeModel(pretrained_model)
+        return model
+
+    def get_generative_model_part(self, content_gcs_path: str, content_mime_type: str | None = None) -> Part:
+        """Return a Generative Model Part object."""
+        part = Part.from_uri(content_gcs_path, mime_type=content_mime_type)
+        return part
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def prompt_language_model(
+        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 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.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        """
+        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 generate_text_embeddings(
+        self,
+        prompt: str,
+        pretrained_model: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> list:
+        """
+        Use the Vertex AI PaLM API to generate text embeddings.
+
+        :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 generating text embeddings.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+        model = self.get_text_embedding_model(pretrained_model)
+
+        response = model.get_embeddings([prompt])[0]  # single prompt
+
+        return response.values
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def prompt_multimodal_model(
+        self,
+        prompt: str,
+        location: str,
+        pretrained_model: str = "gemini-pro",
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Use the Vertex AI Gemini Pro foundation model to generate natural language text.
+
+        :param prompt: Required. Inputs or queries that a user or a program gives
+            to the Multi-modal model, in order to elicit a specific response.
+        :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
+            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 location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+
+        model = self.get_generative_model(pretrained_model)
+        response = model.generate_content(prompt)
+
+        return response.text
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def prompt_multimodal_model_with_media(
+        self,
+        prompt: str,
+        location: str,
+        media_gcs_path: str,
+        mime_type: str,
+        pretrained_model: str = "gemini-pro-vision",
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Use the Vertex AI Gemini Pro foundation model to generate natural language text.
+
+        :param prompt: Required. Inputs or queries that a user or a program gives
+            to the Multi-modal model, in order to elicit a specific response.
+        :param pretrained_model: By default uses the pre-trained model `gemini-pro-vision`,
+            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 media_gcs_path: A GCS path to a content file such as an image or a video.
+            Can be passed to the multi-modal model as part of the prompt. Used with vision models.
+        :param mime_type: Validates the media type presented by the file in the media_gcs_path.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+
+        model = self.get_generative_model(pretrained_model)
+        part = self.get_generative_model_part(media_gcs_path, mime_type)
+        response = model.generate_content([prompt, part])
+
+        return response.text

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

@@ -22,10 +22,10 @@ This module contains a Google Cloud Vertex AI hook.
 
     JobServiceAsyncClient
 """
+
 from __future__ import annotations
 
 import asyncio
-from functools import lru_cache
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.client_options import ClientOptions
@@ -35,12 +35,11 @@ from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient,
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
-    from google.api_core.retry import Retry
-    from google.api_core.retry_async import AsyncRetry
+    from google.api_core.retry import AsyncRetry, Retry
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager
 
 
@@ -431,9 +430,11 @@ class HyperparameterTuningJobHook(GoogleBaseHook):
         return result
 
 
-class HyperparameterTuningJobAsyncHook(GoogleBaseHook):
+class HyperparameterTuningJobAsyncHook(GoogleBaseAsyncHook):
     """Async hook for Google Cloud Vertex AI Hyperparameter Tuning Job APIs."""
 
+    sync_hook_class = HyperparameterTuningJobHook
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",
@@ -446,8 +447,7 @@ class HyperparameterTuningJobAsyncHook(GoogleBaseHook):
             **kwargs,
         )
 
-    @lru_cache
-    def get_job_service_client(self, region: str | None = None) -> JobServiceAsyncClient:
+    async def get_job_service_client(self, region: str | None = None) -> JobServiceAsyncClient:
         """
         Retrieve Vertex AI async client.
 
@@ -455,7 +455,7 @@ class HyperparameterTuningJobAsyncHook(GoogleBaseHook):
         """
         endpoint = f"{region}-aiplatform.googleapis.com:443" if region and region != "global" else None
         return JobServiceAsyncClient(
-            credentials=self.get_credentials(),
+            credentials=(await self.get_sync_hook()).get_credentials(),
             client_info=CLIENT_INFO,
             client_options=ClientOptions(api_endpoint=endpoint),
         )
@@ -479,7 +479,7 @@ class HyperparameterTuningJobAsyncHook(GoogleBaseHook):
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
         """
-        client: JobServiceAsyncClient = self.get_job_service_client(region=location)
+        client: JobServiceAsyncClient = await self.get_job_service_client(region=location)
         job_name = client.hyperparameter_tuning_job_path(project_id, location, job_id)
 
         result = await client.get_hyperparameter_tuning_job(