airflow.providers.google.cloud.triggers.gen_ai

Classes

GenAIGeminiCreateBatchJobTrigger

Trigger that creates Gemini Batch Job and waiting for execution.

GenAIGeminiCreateEmbeddingsBatchJobTrigger

Trigger that creates Gemini Embeddings Batch Job and waiting for execution.

Module Contents

class airflow.providers.google.cloud.triggers.gen_ai.GenAIGeminiCreateBatchJobTrigger(project_id, location, model, input_source, gemini_api_key, create_batch_job_config=None, results_folder=None, gcp_conn_id='google_cloud_default', impersonation_chain=None, retrieve_result=False, polling_interval=30)[source]

Bases: airflow.triggers.base.BaseTrigger

Trigger that creates Gemini Batch Job and waiting for execution.

Parameters:
  • project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.

  • location (str) – Required. The ID of the Google Cloud location that the service belongs to.

  • model (str) – Required. The name of the publisher model to use for Batch job.

  • gemini_api_key (str) – Required. Key to interact with Gemini Batch API.

  • input_source (list | str) – Required. Source of requests, could be inline requests or file name.

  • results_folder (str | None) – Optional. Path to a folder on local machine where file with results will be saved.

  • create_batch_job_config (google.genai.types.CreateBatchJobConfig | dict | None) – Optional. Config for batch job creation.

  • retrieve_result (bool) – Optional. Push the result to XCom. If the input_source is inline, this pushes the execution result. If a file name is specified, this pushes the output file path.

  • polling_interval (int) – Optional. The interval, in seconds, to poll the job status.

  • gcp_conn_id (str) – The connection ID to use connecting to Google Cloud.

  • impersonation_chain (str | collections.abc.Sequence[str] | None) – 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).

project_id[source]
location[source]
gcp_conn_id = 'google_cloud_default'[source]
impersonation_chain = None[source]
model[source]
input_source[source]
create_batch_job_config = None[source]
gemini_api_key[source]
retrieve_result = False[source]
polling_interval = 30[source]
results_folder = None[source]
serialize()[source]

Serialize class arguments and classpath.

async run()[source]

Loop until the job reaches successful final or error state.

Yields a TriggerEvent with success status, if the job reaches successful state.

Yields a TriggerEvent with error status, if the client returns an unexpected terminal job status or any exception is raised while looping.

In any other case the Trigger will wait for a specified amount of time stored in self.polling_interval variable.

class airflow.providers.google.cloud.triggers.gen_ai.GenAIGeminiCreateEmbeddingsBatchJobTrigger(project_id, location, model, input_source, gemini_api_key, create_embeddings_config=None, results_folder=None, gcp_conn_id='google_cloud_default', impersonation_chain=None, retrieve_result=False, polling_interval=30)[source]

Bases: airflow.triggers.base.BaseTrigger

Trigger that creates Gemini Embeddings Batch Job and waiting for execution.

Parameters:
  • project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.

  • location (str) – Required. The ID of the Google Cloud location that the service belongs to.

  • model (str) – Required. The name of the publisher model to use for Batch job.

  • gemini_api_key (str) – Required. Key to interact with Gemini Batch API.

  • input_source (dict | str) – Required. Source of requests, could be inline requests or file name.

  • results_folder (str | None) – Optional. Path to a folder on local machine where file with results will be saved.

  • create_embeddings_config (google.genai.types.CreateBatchJobConfig | dict | None) – Optional. Config for batch job creation.

  • retrieve_result (bool) – Optional. Push the result to XCom. If the input_source is inline, this pushes the execution result. If a file name is specified, this pushes the output file path.

  • polling_interval (int) – Optional. The interval, in seconds, to poll the job status.

  • gcp_conn_id (str) – The connection ID to use connecting to Google Cloud.

  • impersonation_chain (str | collections.abc.Sequence[str] | None) – 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).

project_id[source]
location[source]
gcp_conn_id = 'google_cloud_default'[source]
impersonation_chain = None[source]
model[source]
input_source[source]
create_embeddings_config = None[source]
gemini_api_key[source]
retrieve_result = False[source]
polling_interval = 30[source]
results_folder = None[source]
serialize()[source]

Serialize class arguments and classpath.

async run()[source]

Loop until the job reaches successful final or error state.

Yields a TriggerEvent with success status, if the job reaches successful state.

Yields a TriggerEvent with error status, if the client returns an unexpected terminal job status or any exception is raised while looping.

In any other case the Trigger will wait for a specified amount of time stored in self.polling_interval variable.

Was this entry helpful?