From 2665746d5b7c181cd51bf970e9d88fd6e1a3b98d Mon Sep 17 00:00:00 2001
From: Ayush Agrawal <ayushagra@google.com>
Date: Fri, 20 Feb 2026 09:27:16 -0800
Subject: [PATCH] chore: update deprecated embedding model id
 (text-embedding-004 to gemini-embedding-001)

PiperOrigin-RevId: 872936859
---
 .../tests/batches/test_create_with_bigquery.py    |  6 +++---
 .../genai/tests/batches/test_create_with_gcs.py   |  6 +++---
 google/genai/tests/models/test_embed_content.py   | 15 ++++-----------
 google/genai/tests/shared/models/test_embed.py    |  5 ++++-
 4 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/google/genai/tests/batches/test_create_with_bigquery.py b/google/genai/tests/batches/test_create_with_bigquery.py
index 10c36e672..37dec9800 100644
--- a/google/genai/tests/batches/test_create_with_bigquery.py
+++ b/google/genai/tests/batches/test_create_with_bigquery.py
@@ -24,8 +24,8 @@
 
 _GEMINI_MODEL = 'gemini-2.5-flash'
 _GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
-_EMBEDDING_MODEL = 'text-embedding-004'
-_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/text-embedding-004'
+_EMBEDDING_MODEL = 'gemini-embedding-001'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
 _DISPLAY_NAME = 'test_batch'
 _GENERATE_CONTENT_BQ_INPUT_FILE = (
     'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_requests'
@@ -56,7 +56,7 @@
     pytest_helper.TestTableItem(
         name='test_union_embedding_with_bigquery',
         parameters=types._CreateBatchJobParameters(
-            model=_EMBEDDING_MODEL_FULL_NAME,
+            model='publishers/google/models/text-embedding-004',
             src=_EMBEDDING_BQ_INPUT_FILE,
             config={
                 'display_name': _DISPLAY_NAME,
diff --git a/google/genai/tests/batches/test_create_with_gcs.py b/google/genai/tests/batches/test_create_with_gcs.py
index 5517b0219..cb8934668 100644
--- a/google/genai/tests/batches/test_create_with_gcs.py
+++ b/google/genai/tests/batches/test_create_with_gcs.py
@@ -24,8 +24,8 @@
 
 _GEMINI_MODEL = 'gemini-2.5-flash'
 _GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
-_EMBEDDING_MODEL = 'text-embedding-004'
-_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/text-embedding-004'
+_EMBEDDING_MODEL = 'gemini-embedding-001'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
 _DISPLAY_NAME = 'test_batch'
 
 _GENERATE_CONTENT_GCS_INPUT_FILE = (
@@ -53,7 +53,7 @@
     pytest_helper.TestTableItem(
         name='test_union_embedding_with_gcs',
         parameters=types._CreateBatchJobParameters(
-            model=_EMBEDDING_MODEL,
+            model='text-embedding-004',
             src=_EMBEDDING_GCS_INPUT_FILE,
             config={
                 'display_name': _DISPLAY_NAME,
diff --git a/google/genai/tests/models/test_embed_content.py b/google/genai/tests/models/test_embed_content.py
index ed853a22f..8dd6cd130 100644
--- a/google/genai/tests/models/test_embed_content.py
+++ b/google/genai/tests/models/test_embed_content.py
@@ -33,17 +33,10 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
     return file.read()
 
 test_table: list[pytest_helper.TestTableItem] = [
-    pytest_helper.TestTableItem(
-        name='test_single_text',
-        parameters=types.EmbedContentParameters(
-            model='text-embedding-004',
-            contents=t.t_contents('What is your name?'),
-        ),
-    ),
     pytest_helper.TestTableItem(
         name='test_multi_texts_with_config',
         parameters=types.EmbedContentParameters(
-            model='text-embedding-004',
+            model='gemini-embedding-001',
             contents=[
                 t.t_content('What is your name?'),
                 t.t_content('I am a model.'),
@@ -61,7 +54,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
     pytest_helper.TestTableItem(
         name='test_single_text_with_mime_type_not_supported_in_mldev',
         parameters=types.EmbedContentParameters(
-            model='text-embedding-004',
+            model='gemini-embedding-001',
             contents=t.t_contents('What is your name?'),
             config={
                 'output_dimensionality': 10,
@@ -73,7 +66,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
     pytest_helper.TestTableItem(
         name='test_single_text_with_auto_truncate_not_supported_in_mldev',
         parameters=types.EmbedContentParameters(
-            model='text-embedding-004',
+            model='gemini-embedding-001',
             contents=t.t_contents('What is your name?'),
             config={
                 'output_dimensionality': 10,
@@ -201,7 +194,7 @@ def _get_bytes_from_file(relative_path: str) -> bytes:
 @pytest.mark.asyncio
 async def test_async(client):
   response = await client.aio.models.embed_content(
-      model='text-embedding-004',
+      model='gemini-embedding-001',
       contents='What is your name?',
       config={'output_dimensionality': 10},
   )
diff --git a/google/genai/tests/shared/models/test_embed.py b/google/genai/tests/shared/models/test_embed.py
index 3e1a7e018..394388f81 100644
--- a/google/genai/tests/shared/models/test_embed.py
+++ b/google/genai/tests/shared/models/test_embed.py
@@ -24,8 +24,11 @@
     pytest_helper.TestTableItem(
         name='test_embed',
         parameters=genai_types.EmbedContentParameters(
-            model='text-embedding-004',
+            model='gemini-embedding-001',
             contents='Hello world!',
+            config={
+                'output_dimensionality': 10,
+            },
         ),
     ),
 ]