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

This module contains Google Vertex AI Generative AI operators.

Module Contents

Classes

PromptLanguageModelOperator

Uses the Vertex AI PaLM API to generate natural language text.

GenerateTextEmbeddingsOperator

Uses the Vertex AI PaLM API to generate natural language text.

PromptMultimodalModelOperator

Use the Vertex AI Gemini Pro foundation model to generate natural language text.

PromptMultimodalModelWithMediaOperator

Use the Vertex AI Gemini Pro foundation model to generate natural language text.

class airflow.providers.google.cloud.operators.vertex_ai.generative_model.PromptLanguageModelOperator(*, project_id, location, prompt, pretrained_model='text-bison', temperature=0.0, max_output_tokens=256, top_p=0.8, top_k=40, gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.cloud.operators.cloud_base.GoogleCloudBaseOperator

Uses the Vertex AI PaLM API to generate natural language text.

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.

  • prompt (str) – Required. Inputs or queries that a user or a program gives to the Vertex AI PaLM API, in order to elicit a specific response.

  • pretrained_model (str) – By default uses the pre-trained model text-bison, optimized for performing natural language tasks such as classification, summarization, extraction, content creation, and ideation.

  • temperature (float) – Temperature controls the degree of randomness in token selection. Defaults to 0.0.

  • max_output_tokens (int) – Token limit determines the maximum amount of text output. Defaults to 256.

  • top_p (float) – Tokens are selected from most probable to least until the sum of their probabilities equals the top_p value. Defaults to 0.8.

  • top_k (int) – A top_k of 1 means the selected token is the most probable among all tokens. Defaults to 0.4.

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

  • impersonation_chain (str | 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).

execute(context)[source]

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.generative_model.GenerateTextEmbeddingsOperator(*, project_id, location, prompt, pretrained_model='textembedding-gecko', gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.cloud.operators.cloud_base.GoogleCloudBaseOperator

Uses the Vertex AI PaLM API to generate natural language text.

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.

  • prompt (str) – Required. Inputs or queries that a user or a program gives to the Vertex AI PaLM API, in order to elicit a specific response.

  • pretrained_model (str) – By default uses the pre-trained model textembedding-gecko, optimized for performing text embeddings.

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

  • impersonation_chain (str | 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).

execute(context)[source]

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.generative_model.PromptMultimodalModelOperator(*, project_id, location, prompt, pretrained_model='gemini-pro', gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.cloud.operators.cloud_base.GoogleCloudBaseOperator

Use the Vertex AI Gemini Pro foundation model to generate natural language text.

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.

  • prompt (str) – Required. Inputs or queries that a user or a program gives to the Multi-modal model, in order to elicit a specific response.

  • pretrained_model (str) – 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.

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

  • impersonation_chain (str | 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).

execute(context)[source]

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

class airflow.providers.google.cloud.operators.vertex_ai.generative_model.PromptMultimodalModelWithMediaOperator(*, project_id, location, prompt, media_gcs_path, mime_type, pretrained_model='gemini-pro-vision', gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.cloud.operators.cloud_base.GoogleCloudBaseOperator

Use the Vertex AI Gemini Pro foundation model to generate natural language text.

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.

  • prompt (str) – Required. Inputs or queries that a user or a program gives to the Multi-modal model, in order to elicit a specific response.

  • pretrained_model (str) – 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.

  • media_gcs_path (str) – A GCS path to a media 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.

  • mime_type (str) – Validates the media type presented by the file in the media_gcs_path.

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

  • impersonation_chain (str | 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).

execute(context)[source]

Derive when creating an operator.

Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

Was this entry helpful?