airflow.providers.google.cloud.hooks.gen_ai¶

This module contains a Google Cloud GenAI Generative Model hook.

Classes¶

GenAIGenerativeModelHook

Class for Google Cloud Generative AI Vertex AI hook.

Module Contents¶

class airflow.providers.google.cloud.hooks.gen_ai.GenAIGenerativeModelHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]¶

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Class for Google Cloud Generative AI Vertex AI hook.

get_genai_client(project_id, location)[source]¶

embed_content(model, location, contents, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Generate embeddings for words, phrases, sentences, and code.

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 model to use.
contents (google.genai.types.ContentListUnion | google.genai.types.ContentListUnionDict | list[str]) – Optional. The contents to use for embedding.
config (google.genai.types.EmbedContentConfigOrDict | None) – Optional. Configuration for embeddings.

generate_content(location, model, contents, generation_config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Make an API request to generate content using a model.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
model (str) – Required. The model to use.
contents (google.genai.types.ContentListUnionDict) – 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.
generation_config (google.genai.types.GenerateContentConfig | None) – Optional. Generation configuration settings.

supervised_fine_tuning_train(source_model, location, training_dataset, tuning_job_config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Create a tuning job to adapt model behavior with a labeled dataset.

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.
source_model (str) – Required. A pre-trained model optimized for performing natural language tasks such as classification, summarization, extraction, content creation, and ideation.
train_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.
tuning_job_config (google.genai.types.CreateTuningJobConfigOrDict | dict[str, Any] | None) – Optional. Configuration of the Tuning job to be created.

count_tokens(location, model, contents, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Use Count Tokens API to calculate the number of input tokens before sending a request to Gemini API.

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.
contents (google.genai.types.ContentListUnion | google.genai.types.ContentListUnionDict) – 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.
model (str) – 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.
config (google.genai.types.CountTokensConfigOrDict | None) – Optional. Configuration for Count Tokens.

create_cached_content(model, location, cached_content_config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Create CachedContent to reduce the cost of requests containing repeat content.

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 cached content.
cached_content_config (google.genai.types.CreateCachedContentConfigOrDict | None) – Optional. Configuration of the Cached Content.