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.

GenAIGeminiAPIHook

Class for Google Cloud Generative AI Gemini Developer API 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.

class airflow.providers.google.cloud.hooks.gen_ai.GenAIGeminiAPIHook(gemini_api_key, **kwargs)[source]

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

Class for Google Cloud Generative AI Gemini Developer API hook.

gemini_api_key[source]
get_genai_client()[source]
get_batch_job(job_name)[source]

Get batch job using Gemini Batch API.

Parameters:

job_name (str) – Required. Batch job name.

list_batch_jobs(list_batch_jobs_config=None)[source]

Get list of batch jobs using Gemini Batch API.

Parameters:

list_batch_jobs_config (google.genai.types.ListBatchJobsConfig | dict | None) – Optional. Configuration of returned iterator.

create_batch_job(model, source, create_batch_job_config=None)[source]

Create batch job using Gemini Batch API to process large-scale, non-urgent tasks.

Parameters:

  • model (str) – Required. Gemini model name to process requests.

  • source (list | str) – Required. Requests that will be sent to chosen model. Can be in format of Inline requests or file name.

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

delete_batch_job(job_name)[source]

Delete batch job using Gemini Batch API.

Parameters:

job_name (str) – Required. Batch job name.

cancel_batch_job(job_name)[source]

Cancel batch job using Gemini Batch API.

Parameters:

job_name (str) – Required. Batch job name.

create_embeddings(model, source, create_embeddings_config=None)[source]

Create batch job for embeddings using Gemini Batch API to process large-scale, non-urgent tasks.

Parameters:

  • model (str) – Required. Gemini model name to process requests.

  • source (dict | str) – Required. Requests that will be sent to chosen model. Can be in format of Inline requests or file name.

  • create_embeddings_config (google.genai.types.CreateBatchJobConfig | dict | None) – Optional. Configuration parameters for embeddings batch job.

upload_file(path_to_file, upload_file_config=None)[source]

Upload file for batch job or embeddings batch job using Gemini Files API.

Parameters:

  • path_to_file (str) – Required. Path to file on local filesystem.

  • upload_file_config (dict | None) – Optional. Configuration for file upload.

get_file(file_name)[source]

Get file’s metadata for batch job or embeddings batch job using Gemini Files API.

Parameters:

file_name (str) – Required. Name of the file in Gemini Files API.

download_file(file_name)[source]

Download file for batch job or embeddings batch job using Gemini Files API.

Parameters:

file_name (str) – Required. Name of the file in Gemini Files API.

list_files()[source]

List files for stored in Gemini Files API.

delete_file(file_name)[source]

Delete file from Gemini Files API storage.

Parameters:

file_name (str) – Required. Name of the file in Gemini Files API.

Was this entry helpful?