This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 671ad9dceb8 Fix models selection in genai tests (#72323)
671ad9dceb8 is described below
commit 671ad9dceb87a3417de8fc4525e7e84fcab6640b
Author: Marcin <[email protected]>
AuthorDate: Wed Sep 9 00:12:48 2026 +0200
Fix models selection in genai tests (#72323)
---
.../gen_ai/example_gen_ai_gemini_batch_api.py | 16 +--
.../gen_ai/example_gen_ai_generative_model.py | 116 ++-------------------
.../example_gen_ai_generative_model_tuning.py | 84 +++------------
.../tests/system/google/cloud/gen_ai/llm_models.py | 46 ++++++++
.../gemini_batch_embeddings_requests.jsonl | 4 +-
5 files changed, 78 insertions(+), 188 deletions(-)
diff --git
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_gemini_batch_api.py
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_gemini_batch_api.py
index 56db1256a65..12718f85e2b 100644
---
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_gemini_batch_api.py
+++
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_gemini_batch_api.py
@@ -55,10 +55,14 @@ from airflow.providers.google.cloud.operators.gen_ai import
(
from airflow.providers.google.common.utils.get_secret import get_secret
from airflow.providers.standard.operators.bash import BashOperator
+from system.google.cloud.gen_ai import llm_models
+
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
REGION = "us-central1"
DAG_ID = "gen_ai_gemini_batch_api"
+TEXT_EMBEDDING_MODEL = llm_models.get_text_embedding_gemini_model()
+DEFAULT_MODEL = llm_models.get_default_gemini_model()
GEMINI_API_KEY = "api_key"
@@ -178,7 +182,7 @@ with DAG(
task_id="create_batch_job_using_inlined_requests_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-3-pro-preview",
+ model=DEFAULT_MODEL,
gemini_api_key=GEMINI_XCOM_API_KEY,
create_batch_job_config={
"display_name": "inlined-requests-batch-job",
@@ -194,7 +198,7 @@ with DAG(
task_id="create_batch_job_using_inlined_requests_deferrable_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-3-pro-preview",
+ model=DEFAULT_MODEL,
gemini_api_key=GEMINI_XCOM_API_KEY,
create_batch_job_config={
"display_name": "deferrable-inlined-requests-batch-job",
@@ -210,7 +214,7 @@ with DAG(
task_id="create_batch_job_using_file_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-3-pro-preview",
+ model=DEFAULT_MODEL,
gemini_api_key=GEMINI_XCOM_API_KEY,
create_batch_job_config={
"display_name": "file-upload-batch-job",
@@ -227,7 +231,7 @@ with DAG(
task_id="create_embeddings_job_using_inlined_requests_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-embedding-001",
+ model=TEXT_EMBEDDING_MODEL,
wait_until_complete=False,
gemini_api_key=GEMINI_XCOM_API_KEY,
create_embeddings_config={
@@ -242,7 +246,7 @@ with DAG(
task_id="create_embeddings_job_using_file_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-embedding-001",
+ model=TEXT_EMBEDDING_MODEL,
wait_until_complete=False,
gemini_api_key=GEMINI_XCOM_API_KEY,
create_embeddings_config={
@@ -257,7 +261,7 @@ with DAG(
task_id="create_embeddings_job_using_file_deferrable_task",
project_id=PROJECT_ID,
location=REGION,
- model="gemini-embedding-001",
+ model=TEXT_EMBEDDING_MODEL,
retrieve_result=True,
deferrable=True,
gemini_api_key=GEMINI_XCOM_API_KEY,
diff --git
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
index c9a6f9ffcaf..85ed7a671ec 100644
---
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
+++
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
@@ -25,13 +25,6 @@ from __future__ import annotations
import os
from datetime import datetime
-import requests
-
-try:
- from airflow.sdk import task
-except ImportError:
- # Airflow 2 path
- from airflow.decorators import task # type: ignore[attr-defined,no-redef]
from google.genai.types import (
Content,
CreateCachedContentConfig,
@@ -54,92 +47,17 @@ from
airflow.providers.google.cloud.operators.vertex_ai.experiment_service impor
DeleteExperimentRunOperator,
)
from airflow.providers.google.cloud.operators.vertex_ai.generative_model
import RunEvaluationOperator
-from airflow.providers.google.common.utils.get_secret import get_secret
-
-
-def _get_actual_models(key) -> dict[str, str]:
- models: dict[str, str] = {
- "multimodal": "",
- "text-embedding": "",
- "cached-model": "",
- }
- try:
- response = requests.get(
- "https://generativelanguage.googleapis.com/v1beta/models";,
- {"key": key},
- timeout=10,
- )
- response.raise_for_status()
- available_models = response.json()
- except requests.exceptions.RequestException as e:
- print(f"Error fetching models from API: {e}")
- return models
-
- for model in available_models.get("models", []):
- try:
- model_name = model["name"].split("/")[-1]
- splited_model_name = model_name.split("-")
- if not models["text-embedding"] and ("gemini" in model_name and
"embedding" in model_name):
- models["text-embedding"] = model_name
- elif (
- models["text-embedding"]
- and ("text" in model_name and "embedding" in model_name)
- and int(splited_model_name[-1]) >
int(models["text-embedding"].split("-")[-1])
- ):
- models["text-embedding"] = model_name
- elif ("vision" not in model_name or "image" in model_name) and (
- "flash" in model_name or "pro" in model_name
- ):
- if not models["multimodal"] and "pro" in model_name:
- models["multimodal"] = model_name
- elif (
- models["multimodal"]
- and "pro" in model_name
- and float(models["multimodal"].split("-")[1]) <
float(splited_model_name[1])
- ):
- models["multimodal"] = model_name
- elif (
- models["multimodal"]
- and "pro" in model_name
- and (
- float(models["multimodal"].split("-")[1]) ==
float(splited_model_name[1])
- and int(splited_model_name[-1]) >
int(models["multimodal"].split("-")[-1])
- )
- ):
- models["multimodal"] = model_name
- if "createCachedContent" in
model["supportedGenerationMethods"]:
- if not models["cached-model"]:
- models["cached-model"] = model_name
- elif models["cached-model"] and
float(models["cached-model"].split("-")[1]) < float(
- splited_model_name[1]
- ):
- models["cached-model"] = model_name
- elif (
- models["cached-model"]
- and float(models["cached-model"].split("-")[1]) ==
float(splited_model_name[1])
- and int(splited_model_name[-1]) >
int(models["cached-model"].split("-")[-1])
- ):
- models["cached-model"] = model_name
- except (ValueError, IndexError) as e:
- print(f"Could not parse model name '{model.get('name')}'.
Skipping. Error: {e}")
- continue
- if not any(models.values()):
- raise ValueError(f"Some of the models not found {models}")
- return models
+from system.google.cloud.gen_ai import llm_models
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
-GEMINI_API_KEY = "api_key"
-MODELS = "{{ task_instance.xcom_pull('get_actual_models') }}"
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
DAG_ID = "gen_ai_generative_model_dag"
REGION = "us-central1"
PROMPT = "In 10 words or less, why is Apache Airflow amazing?"
CONTENTS = [PROMPT]
-TEXT_EMBEDDING_MODEL = "{{
task_instance.xcom_pull('get_actual_models')['text-embedding'] }}"
-MULTIMODAL_MODEL = "{{
task_instance.xcom_pull('get_actual_models')['multimodal'] }}"
-MEDIA_GCS_PATH =
"gs://download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg"
-MIME_TYPE = "image/jpeg"
+TEXT_EMBEDDING_MODEL = llm_models.get_text_embedding_gemini_model()
+MULTIMODAL_MODEL = llm_models.get_default_gemini_model()
TOOLS = [Tool(google_search=GoogleSearch())]
GENERATION_CONFIG_CREATE_CONTENT = GenerateContentConfig(
max_output_tokens=256,
@@ -194,7 +112,7 @@ EXPERIMENT_NAME =
f"eval-test-experiment-airflow-operator-{ENV_ID}".replace("_",
EXPERIMENT_RUN_NAME =
f"eval-experiment-airflow-operator-run-{ENV_ID}".replace("_", "-")
PROMPT_TEMPLATE = "{instruction}. Article: {context}. Summary:"
-CACHED_MODEL = "{{
task_instance.xcom_pull('get_actual_models')['cached-model'] }}"
+CACHED_MODEL = llm_models.get_default_gemini_model()
CACHED_SYSTEM_INSTRUCTION = """
You are an expert researcher. You always stick to the facts in the sources
provided, and never make up new facts.
Now look at these research papers, and answer the following questions.
@@ -228,19 +146,6 @@ with DAG(
tags=["example", "gen_ai", "generative_model"],
render_template_as_native_obj=True,
) as dag:
-
- @task
- def get_gemini_api_key():
- return get_secret(GEMINI_API_KEY)
-
- get_gemini_api_key_task = get_gemini_api_key()
-
- @task
- def get_actual_models(key):
- return _get_actual_models(key)
-
- get_actual_models_task = get_actual_models(get_gemini_api_key_task)
-
# [START how_to_cloud_gen_ai_generate_embeddings_task]
generate_embeddings_task = GenAIGenerateEmbeddingsOperator(
task_id="generate_embeddings_task",
@@ -330,16 +235,9 @@ with DAG(
model=CACHED_MODEL,
)
# [END how_to_cloud_gen_ai_generate_from_cached_content_operator]
- get_gemini_api_key_task >> get_actual_models_task
- get_actual_models_task >> [generate_embeddings_task, count_tokens_task,
generate_content_task]
- get_actual_models_task >> create_cached_content_task >>
generate_from_cached_content_task
- (
- get_actual_models_task
- >> create_experiment_task
- >> run_evaluation_task
- >> delete_experiment_run_task
- >> delete_experiment_task
- )
+ [generate_embeddings_task, count_tokens_task, generate_content_task]
+ create_cached_content_task >> generate_from_cached_content_task
+ (create_experiment_task >> run_evaluation_task >>
delete_experiment_run_task >> delete_experiment_task)
from tests_common.test_utils.watcher import watcher
diff --git
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model_tuning.py
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model_tuning.py
index e779c35fbd0..171f535db85 100644
---
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model_tuning.py
+++
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model_tuning.py
@@ -26,16 +26,9 @@ import os
from datetime import datetime
from pathlib import Path
-import requests
-
from airflow.providers.google.cloud.operators.gcs import
GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.transfers.local_to_gcs import
LocalFilesystemToGCSOperator
-try:
- from airflow.sdk import task
-except ImportError:
- # Airflow 2 path
- from airflow.decorators import task # type: ignore[attr-defined,no-redef]
try:
from airflow.sdk import TriggerRule
except ImportError:
@@ -47,59 +40,28 @@ from airflow.models.dag import DAG
from airflow.providers.google.cloud.operators.gen_ai import (
GenAISupervisedFineTuningTrainOperator,
)
-from airflow.providers.google.common.utils.get_secret import get_secret
-
-
-def _get_actual_model(key) -> str:
- source_model: str | None = None
- try:
- response =
requests.get("https://generativelanguage.googleapis.com/v1/models";, {"key":
key})
- response.raise_for_status()
- available_models = response.json()
- except requests.exceptions.RequestException as e:
- print(f"Error fetching models from API: {e}")
- return ""
- for model in available_models.get("models", []):
- try:
- model_name = model["name"].split("/")[-1]
- splited_model_name = model_name.split("-")
- if not source_model and "flash" in model_name:
- source_model = model_name
- elif (
- source_model
- and "flash" in model_name
- and float(source_model.split("-")[1]) <
float(splited_model_name[1])
- ):
- source_model = model_name
- elif (
- source_model
- and "flash" in model_name
- and (
- float(source_model.split("-")[1]) ==
float(splited_model_name[1])
- and int(splited_model_name[-1]) >
int(source_model.split("-")[-1])
- )
- ):
- source_model = model_name
- except (ValueError, IndexError) as e:
- print(f"Could not parse model name '{model.get('name')}'.
Skipping. Error: {e}")
- continue
- if not source_model:
- raise ValueError("Source model not found")
- return source_model
+from system.google.cloud.gen_ai import llm_models
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
DAG_ID = "gen_ai_generative_model_tuning_dag"
REGION = "us-central1"
-GEMINI_API_KEY = "api_key"
-SOURCE_MODEL = "{{ task_instance.xcom_pull('get_actual_model') }}"
+SOURCE_MODEL = llm_models.get_sft_enabled_gemini_model()
TRAIN_DATASET = TuningDataset(
gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-1_5/text/sft_train_data.jsonl",
)
TUNED_MODEL_DISPLAY_NAME = "my_tuned_gemini_model"
-TUNING_JOB_CONFIG = {"tuned_model_display_name": TUNED_MODEL_DISPLAY_NAME}
+TUNING_JOB_CONFIG = {
+ "tuned_model_display_name": TUNED_MODEL_DISPLAY_NAME,
+ "epoch_count": 1,
+ "export_last_checkpoint_only": True,
+}
TUNED_VIDEO_MODEL_DISPLAY_NAME = "my_tuned_gemini_video_model"
-TUNING_JOB_VIDEO_MODEL_CONFIG = {"tuned_model_display_name":
TUNED_VIDEO_MODEL_DISPLAY_NAME}
+TUNING_JOB_VIDEO_MODEL_CONFIG = {
+ "tuned_model_display_name": TUNED_VIDEO_MODEL_DISPLAY_NAME,
+ "epoch_count": 1,
+ "export_last_checkpoint_only": True,
+}
BUCKET_NAME = f"bucket_tuning_dag_{PROJECT_ID}"
FILE_NAME = "video_tuning_dataset.jsonl"
@@ -116,19 +78,6 @@ with DAG(
tags=["example", "vertex_ai", "generative_model"],
render_template_as_native_obj=True,
) as dag:
-
- @task
- def get_gemini_api_key():
- return get_secret(GEMINI_API_KEY)
-
- get_gemini_api_key_task = get_gemini_api_key()
-
- @task
- def get_actual_model(key):
- return _get_actual_model(key)
-
- get_actual_model_task = get_actual_model(get_gemini_api_key_task)
-
create_bucket = GCSCreateBucketOperator(
task_id="create_bucket",
bucket_name=BUCKET_NAME,
@@ -168,14 +117,7 @@ with DAG(
delete_bucket.trigger_rule = TriggerRule.ALL_DONE
- (
- get_gemini_api_key_task
- >> get_actual_model_task
- >> create_bucket
- >> upload_file
- >> [sft_train_task, sft_video_task]
- >> delete_bucket
- )
+ (create_bucket >> upload_file >> [sft_train_task, sft_video_task] >>
delete_bucket)
from tests_common.test_utils.watcher import watcher
diff --git a/providers/google/tests/system/google/cloud/gen_ai/llm_models.py
b/providers/google/tests/system/google/cloud/gen_ai/llm_models.py
new file mode 100644
index 00000000000..caa5a4be5dd
--- /dev/null
+++ b/providers/google/tests/system/google/cloud/gen_ai/llm_models.py
@@ -0,0 +1,46 @@
+#
+# 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.
+from __future__ import annotations
+
+import os
+
+DEFAULT_GEMINI_MODEL = "gemini-3.7-flash"
+SFT_ENABLED_GEMINI_MODEL = "gemini-3.5-flash"
+TEXT_EMBEDDING_GEMINI_MODEL = "gemini-embedding-2"
+
+
+def _get_model_env_variable(name: str, default: str) -> str:
+ value = os.environ.get(name, default)
+ if not value:
+ raise ValueError(f"Environment variable {name} must not be empty.")
+ return value
+
+
+# it should be a general purpose model, multimodal & cached
+def get_default_gemini_model() -> str:
+ return _get_model_env_variable("DEFAULT_GEMINI_MODEL",
DEFAULT_GEMINI_MODEL)
+
+
+# sft enabled model
+def get_sft_enabled_gemini_model() -> str:
+ return _get_model_env_variable("SFT_ENABLED_GEMINI_MODEL",
SFT_ENABLED_GEMINI_MODEL)
+
+
+# embedding model with text support
+def get_text_embedding_gemini_model() -> str:
+ return _get_model_env_variable("TEXT_EMBEDDING_GEMINI_MODEL",
TEXT_EMBEDDING_GEMINI_MODEL)
diff --git
a/providers/google/tests/system/google/cloud/gen_ai/resources/gemini_batch_embeddings_requests.jsonl
b/providers/google/tests/system/google/cloud/gen_ai/resources/gemini_batch_embeddings_requests.jsonl
index b898b469662..a3e26fa37a2 100644
---
a/providers/google/tests/system/google/cloud/gen_ai/resources/gemini_batch_embeddings_requests.jsonl
+++
b/providers/google/tests/system/google/cloud/gen_ai/resources/gemini_batch_embeddings_requests.jsonl
@@ -1,2 +1,2 @@
-{"key": "request_1", "request": {"output_dimensionality": 3, "content":
{"parts": [{"text": "1"}]}}}
-{"key": "request_2", "request": {"output_dimensionality": 4, "content":
{"parts": [{"text": "2"}]}}}
+{"key": "request_1", "request": {"output_dimensionality": 128, "content":
{"parts": [{"text": "1"}]}}}
+{"key": "request_2", "request": {"output_dimensionality": 256, "content":
{"parts": [{"text": "2"}]}}}
