google-genai 1.60.0__py3-none-any.whl → 1.61.0__py3-none-any.whl

google/genai/tests/batches/test_create_with_inlined_requests.py

@@ -42,20 +42,36 @@ _SAFETY_SETTINGS = [
     },
 ]
 
-_INLINED_REQUEST = {
-    'contents': [{
-        'parts': [{
-            'text': 'Hello!',
+_INLINED_REQUESTS = [
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 1? return just the number.',
+            }],
+            'role': 'user',
         }],
-        'role': 'user',
-    }],
-    'metadata': {
-        'key': 'request-1',
+        'metadata': {
+            'key': 'request-1',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
     },
-    'config': {
-        'safety_settings': _SAFETY_SETTINGS,
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 2? return just the number.',
+            }],
+            'role': 'user',
+        }],
+        'metadata': {
+            'key': 'request-2',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
     },
-}
+]
 _INLINED_TEXT_REQUEST_UNION = {
     'contents': [{
         'parts': [{
@@ -154,7 +170,7 @@ test_table: list[pytest_helper.TestTableItem] = [
         name='test_union_with_inlined_request',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src=[_INLINED_REQUEST],
+            src=_INLINED_REQUESTS,
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -166,7 +182,7 @@ test_table: list[pytest_helper.TestTableItem] = [
         name='test_with_inlined_request',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src={'inlined_requests': [_INLINED_REQUEST]},
+            src={'inlined_requests': _INLINED_REQUESTS},
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -177,7 +193,7 @@ test_table: list[pytest_helper.TestTableItem] = [
         name='test_with_inlined_request_config',
         parameters=types._CreateBatchJobParameters(
             model=_MLDEV_GEMINI_MODEL,
-            src={'inlined_requests': [_INLINED_TEXT_REQUEST]},
+            src={'inlined_requests': _INLINED_REQUESTS},
             config={
                 'display_name': _DISPLAY_NAME,
             },
@@ -247,7 +263,7 @@ async def test_async_create(client):
   with pytest_helper.exception_if_vertex(client, ValueError):
     batch_job = await client.aio.batches.create(
         model=_GEMINI_MODEL,
-        src=[_INLINED_REQUEST],
+        src=_INLINED_REQUESTS,
     )
     assert batch_job.name.startswith('batches/')
     assert (

google/genai/tests/batches/test_get.py

@@ -29,7 +29,7 @@ _BATCH_JOB_FULL_RESOURCE_NAME = (
     f'batchPredictionJobs/{_BATCH_JOB_NAME}'
 )
 # MLDev batch operation name.
-_MLDEV_BATCH_OPERATION_NAME = 'batches/0yew7plxupyybd7appsrq5vw7w0lp3l79lab'
+_MLDEV_BATCH_OPERATION_NAME = 'batches/z2p8ksus4lyxt25rntl3fpd67p2niw4hfij5'
 _INVALID_BATCH_JOB_NAME = 'invalid_name'
 
 

google/genai/tests/client/test_client_close.py

@@ -90,7 +90,6 @@ async def test_async_httpx_client_context_manager():
   assert async_client._api_client._async_httpx_client.is_closed
 
 
-@requires_aiohttp
 @pytest.fixture
 def mock_request():
   mock_aiohttp_response = mock.Mock(spec=aiohttp.ClientSession.request)

google/genai/tests/files/test_register_table.py

@@ -42,7 +42,7 @@ def get_headers():
 test_table: list[pytest_helper.TestTableItem] = [
     pytest_helper.TestTableItem(
         name='test_register',
-        parameters=types._RegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
+        parameters=types._InternalRegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
         exception_if_vertex='only supported in the Gemini Developer client',
         skip_in_api_mode=(
             'The files have a TTL, they cannot be reliably retrieved for a long'

google/genai/tests/transformers/test_schema.py

@@ -79,7 +79,6 @@ class CountryInfoWithAnyOf(pydantic.BaseModel):
 
 
 @pytest.fixture
-@pytest.mark.parametrize('use_vertex', [True, False])
 def client(use_vertex):
   if use_vertex:
     yield google_genai_client_module.Client(
@@ -91,6 +90,7 @@ def client(use_vertex):
     )
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_build_schema_for_list_of_pydantic_schema(client):
   """Tests _build_schema() when list[pydantic.BaseModel] is provided to response_schema."""
 
@@ -112,6 +112,7 @@ def test_build_schema_for_list_of_pydantic_schema(client):
     assert list_schema['required'] == list(country_info_fields.keys())
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_build_schema_for_list_of_nested_pydantic_schema(client):
   """Tests _build_schema() when list[pydantic.BaseModel] is provided to response_schema and the pydantic.BaseModel has nested pydantic fields."""
   list_schema = _transformers.t_schema(
@@ -132,6 +133,7 @@ def test_build_schema_for_list_of_nested_pydantic_schema(client):
     assert field_name in currency_info_fields
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_for_pydantic_schema(client):
   """Tests t_schema when pydantic.BaseModel is passed to response_schema."""
   transformed_schema = _transformers.t_schema(client, CountryInfo)
@@ -143,6 +145,7 @@ def test_t_schema_for_pydantic_schema(client):
     )
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_for_list_of_pydantic_schema(client):
   """Tests t_schema when list[pydantic.BaseModel] is passed to response_schema."""
   transformed_schema = _transformers.t_schema(client, list[CountryInfo])
@@ -156,6 +159,7 @@ def test_t_schema_for_list_of_pydantic_schema(client):
     )
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_for_null_fields(client):
   """Tests t_schema when null fields are present."""
   transformed_schema = _transformers.t_schema(client, CountryInfoWithNullFields)
@@ -163,6 +167,7 @@ def test_t_schema_for_null_fields(client):
   assert transformed_schema.properties['population'].nullable
 
 
+#@pytest.mark.parametrize('use_vertex', [True, False])
 def test_schema_with_no_null_fields_is_unchanged():
   """Tests handle_null_fields() doesn't change anything when no null fields are present."""
   test_properties = {
@@ -207,6 +212,7 @@ def test_schema_with_default_value(client):
   assert transformed_schema == expected_schema
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_schema_with_any_of(client):
   transformed_schema = _transformers.t_schema(client, CountryInfoWithAnyOf)
   expected_schema = types.Schema(
@@ -601,6 +607,7 @@ def test_process_schema_order_properties_propagates_into_any_of(
     assert schema == schema_without_property_ordering
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_does_not_change_property_ordering_if_set(client):
   """Tests t_schema doesn't overwrite the property_ordering field if already set."""
 
@@ -612,6 +619,7 @@ def test_t_schema_does_not_change_property_ordering_if_set(client):
   assert transformed_schema.property_ordering == custom_property_ordering
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_sets_property_ordering_for_json_schema(client):
   """Tests t_schema sets the property_ordering field for json schemas."""
 
@@ -629,6 +637,7 @@ def test_t_schema_sets_property_ordering_for_json_schema(client):
   ]
 
 
+@pytest.mark.parametrize('use_vertex', [True, False])
 def test_t_schema_sets_property_ordering_for_schema_type(client):
   """Tests t_schema sets the property_ordering field for Schema types."""
 

google/genai/tests/tunings/test_tune.py

@@ -20,6 +20,12 @@ from ... import types as genai_types
 from .. import pytest_helper
 import pytest
 
+
+VERTEX_HTTP_OPTIONS = {
+    'api_version': 'v1beta1',
+    'base_url': 'https://us-central1-autopush-aiplatform.sandbox.googleapis.com/',
+}
+
 evaluation_config=genai_types.EvaluationConfig(
     metrics=[
         genai_types.Metric(name="bleu", prompt_template="test prompt template")
@@ -158,6 +164,87 @@ test_table: list[pytest_helper.TestTableItem] = [
         ),
         exception_if_mldev="vertex_dataset_resource parameter is not supported in Gemini API.",
     ),
+    pytest_helper.TestTableItem(
+        name="test_tune_distillation",
+        parameters=genai_types.CreateTuningJobParameters(
+            base_model="meta/llama3_1@llama-3.1-8b-instruct",
+            training_dataset=genai_types.TuningDataset(
+                gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl",
+            ),
+            config=genai_types.CreateTuningJobConfig(
+                method="DISTILLATION",
+                base_teacher_model="deepseek-ai/deepseek-v3.1-maas",
+                epoch_count=20,
+                validation_dataset=genai_types.TuningValidationDataset(
+                    gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-val-openai-opposites.jsonl",
+                ),
+                output_uri="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test-folder",
+                http_options=VERTEX_HTTP_OPTIONS,
+            ),
+        ),
+        exception_if_mldev="parameter is not supported in Gemini API.",
+    ),
+    pytest_helper.TestTableItem(
+        name="test_tune_oss_sft",
+        parameters=genai_types.CreateTuningJobParameters(
+            base_model="meta/llama3_1@llama-3.1-8b-instruct",
+            training_dataset=genai_types.TuningDataset(
+                gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl",
+            ),
+            config=genai_types.CreateTuningJobConfig(
+                epoch_count=20,
+                validation_dataset=genai_types.TuningValidationDataset(
+                    gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-val-openai-opposites.jsonl",
+                ),
+                custom_base_model="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test-folder/postprocess/node-0/checkpoints/final",
+                output_uri="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test",
+                http_options=VERTEX_HTTP_OPTIONS,
+            ),
+        ),
+        exception_if_mldev="not supported in Gemini API",
+    ),
+    pytest_helper.TestTableItem(
+        name="test_tune_oss_sft_hyperparams",
+        parameters=genai_types.CreateTuningJobParameters(
+            base_model="meta/llama3_1@llama-3.1-8b-instruct",
+            training_dataset=genai_types.TuningDataset(
+                gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl",
+            ),
+            config=genai_types.CreateTuningJobConfig(
+                epoch_count=20,
+                validation_dataset=genai_types.TuningValidationDataset(
+                    gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-val-openai-opposites.jsonl",
+                ),
+                learning_rate=2.5e-4,
+                tuning_mode="TUNING_MODE_FULL",
+                custom_base_model="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test-folder/postprocess/node-0/checkpoints/final",
+                output_uri="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test",
+                http_options=VERTEX_HTTP_OPTIONS,
+            ),
+        ),
+        exception_if_mldev="not supported in Gemini API",
+    ),
+    pytest_helper.TestTableItem(
+        name="test_tune_oss_distillation",
+        parameters=genai_types.CreateTuningJobParameters(
+            base_model="meta/llama3_1@llama-3.1-8b-instruct",
+            training_dataset=genai_types.TuningDataset(
+                gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl",
+            ),
+            config=genai_types.CreateTuningJobConfig(
+                method="DISTILLATION",
+                base_teacher_model="deepseek-ai/deepseek-v3.1-maas",
+                epoch_count=20,
+                validation_dataset=genai_types.TuningValidationDataset(
+                    gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-val-openai-opposites.jsonl",
+                ),
+                custom_base_model="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test-folder/postprocess/node-0/checkpoints/final",
+                output_uri="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test",
+                http_options=VERTEX_HTTP_OPTIONS,
+            ),
+        ),
+        exception_if_mldev="not supported in Gemini API",
+    ),
 ]
 
 pytestmark = pytest_helper.setup(

google/genai/tunings.py

@@ -188,6 +188,14 @@ def _CreateTuningJobConfig_to_mldev(
   if getv(from_object, ['adapter_size']) is not None:
     raise ValueError('adapter_size parameter is not supported in Gemini API.')
 
+  if getv(from_object, ['tuning_mode']) is not None:
+    raise ValueError('tuning_mode parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['custom_base_model']) is not None:
+    raise ValueError(
+        'custom_base_model parameter is not supported in Gemini API.'
+    )
+
   if getv(from_object, ['batch_size']) is not None:
     setv(
         parent_object,
@@ -213,6 +221,24 @@ def _CreateTuningJobConfig_to_mldev(
   if getv(from_object, ['beta']) is not None:
     raise ValueError('beta parameter is not supported in Gemini API.')
 
+  if getv(from_object, ['base_teacher_model']) is not None:
+    raise ValueError(
+        'base_teacher_model parameter is not supported in Gemini API.'
+    )
+
+  if getv(from_object, ['tuned_teacher_model_source']) is not None:
+    raise ValueError(
+        'tuned_teacher_model_source parameter is not supported in Gemini API.'
+    )
+
+  if getv(from_object, ['sft_loss_weight_multiplier']) is not None:
+    raise ValueError(
+        'sft_loss_weight_multiplier parameter is not supported in Gemini API.'
+    )
+
+  if getv(from_object, ['output_uri']) is not None:
+    raise ValueError('output_uri parameter is not supported in Gemini API.')
+
   return to_object
 
 
@@ -246,6 +272,16 @@ def _CreateTuningJobConfig_to_vertex(
           ),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['validation_dataset']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec'],
+          _TuningValidationDataset_to_vertex(
+              getv(from_object, ['validation_dataset']), to_object, root_object
+          ),
+      )
+
   if getv(from_object, ['tuned_model_display_name']) is not None:
     setv(
         parent_object,
@@ -275,6 +311,14 @@ def _CreateTuningJobConfig_to_vertex(
           getv(from_object, ['epoch_count']),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['epoch_count']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'hyperParameters', 'epochCount'],
+          getv(from_object, ['epoch_count']),
+      )
+
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
     discriminator = 'SUPERVISED_FINE_TUNING'
@@ -298,6 +342,14 @@ def _CreateTuningJobConfig_to_vertex(
           getv(from_object, ['learning_rate_multiplier']),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['learning_rate_multiplier']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'hyperParameters', 'learningRateMultiplier'],
+          getv(from_object, ['learning_rate_multiplier']),
+      )
+
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
     discriminator = 'SUPERVISED_FINE_TUNING'
@@ -317,6 +369,14 @@ def _CreateTuningJobConfig_to_vertex(
           getv(from_object, ['export_last_checkpoint_only']),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['export_last_checkpoint_only']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'exportLastCheckpointOnly'],
+          getv(from_object, ['export_last_checkpoint_only']),
+      )
+
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
     discriminator = 'SUPERVISED_FINE_TUNING'
@@ -336,11 +396,53 @@ def _CreateTuningJobConfig_to_vertex(
           getv(from_object, ['adapter_size']),
       )
 
-  if getv(from_object, ['batch_size']) is not None:
-    raise ValueError('batch_size parameter is not supported in Vertex AI.')
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['adapter_size']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'hyperParameters', 'adapterSize'],
+          getv(from_object, ['adapter_size']),
+      )
 
-  if getv(from_object, ['learning_rate']) is not None:
-    raise ValueError('learning_rate parameter is not supported in Vertex AI.')
+  discriminator = getv(root_object, ['config', 'method'])
+  if discriminator is None:
+    discriminator = 'SUPERVISED_FINE_TUNING'
+  if discriminator == 'SUPERVISED_FINE_TUNING':
+    if getv(from_object, ['tuning_mode']) is not None:
+      setv(
+          parent_object,
+          ['supervisedTuningSpec', 'tuningMode'],
+          getv(from_object, ['tuning_mode']),
+      )
+
+  if getv(from_object, ['custom_base_model']) is not None:
+    setv(
+        parent_object,
+        ['customBaseModel'],
+        getv(from_object, ['custom_base_model']),
+    )
+
+  discriminator = getv(root_object, ['config', 'method'])
+  if discriminator is None:
+    discriminator = 'SUPERVISED_FINE_TUNING'
+  if discriminator == 'SUPERVISED_FINE_TUNING':
+    if getv(from_object, ['batch_size']) is not None:
+      setv(
+          parent_object,
+          ['supervisedTuningSpec', 'hyperParameters', 'batchSize'],
+          getv(from_object, ['batch_size']),
+      )
+
+  discriminator = getv(root_object, ['config', 'method'])
+  if discriminator is None:
+    discriminator = 'SUPERVISED_FINE_TUNING'
+  if discriminator == 'SUPERVISED_FINE_TUNING':
+    if getv(from_object, ['learning_rate']) is not None:
+      setv(
+          parent_object,
+          ['supervisedTuningSpec', 'hyperParameters', 'learningRate'],
+          getv(from_object, ['learning_rate']),
+      )
 
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
@@ -365,6 +467,16 @@ def _CreateTuningJobConfig_to_vertex(
           ),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['evaluation_config']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'evaluationConfig'],
+          _EvaluationConfig_to_vertex(
+              getv(from_object, ['evaluation_config']), to_object, root_object
+          ),
+      )
+
   if getv(from_object, ['labels']) is not None:
     setv(parent_object, ['labels'], getv(from_object, ['labels']))
 
@@ -375,6 +487,30 @@ def _CreateTuningJobConfig_to_vertex(
         getv(from_object, ['beta']),
     )
 
+  if getv(from_object, ['base_teacher_model']) is not None:
+    setv(
+        parent_object,
+        ['distillationSpec', 'baseTeacherModel'],
+        getv(from_object, ['base_teacher_model']),
+    )
+
+  if getv(from_object, ['tuned_teacher_model_source']) is not None:
+    setv(
+        parent_object,
+        ['distillationSpec', 'tunedTeacherModelSource'],
+        getv(from_object, ['tuned_teacher_model_source']),
+    )
+
+  if getv(from_object, ['sft_loss_weight_multiplier']) is not None:
+    setv(
+        parent_object,
+        ['distillationSpec', 'hyperParameters', 'sftLossWeightMultiplier'],
+        getv(from_object, ['sft_loss_weight_multiplier']),
+    )
+
+  if getv(from_object, ['output_uri']) is not None:
+    setv(parent_object, ['outputUri'], getv(from_object, ['output_uri']))
+
   return to_object
 
 
@@ -920,6 +1056,14 @@ def _TuningDataset_to_vertex(
           getv(from_object, ['gcs_uri']),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['gcs_uri']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'promptDatasetUri'],
+          getv(from_object, ['gcs_uri']),
+      )
+
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
     discriminator = 'SUPERVISED_FINE_TUNING'
@@ -939,6 +1083,14 @@ def _TuningDataset_to_vertex(
           getv(from_object, ['vertex_dataset_resource']),
       )
 
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['vertex_dataset_resource']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'promptDatasetUri'],
+          getv(from_object, ['vertex_dataset_resource']),
+      )
+
   if getv(from_object, ['examples']) is not None:
     raise ValueError('examples parameter is not supported in Vertex AI.')
 
@@ -1066,6 +1218,13 @@ def _TuningJob_from_vertex(
         getv(from_object, ['preferenceOptimizationSpec']),
     )
 
+  if getv(from_object, ['distillationSpec']) is not None:
+    setv(
+        to_object,
+        ['distillation_spec'],
+        getv(from_object, ['distillationSpec']),
+    )
+
   if getv(from_object, ['tuningDataStats']) is not None:
     setv(
         to_object, ['tuning_data_stats'], getv(from_object, ['tuningDataStats'])