google-genai 1.60.0__py3-none-any.whl → 1.62.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/errors/test_api_error.py

@@ -23,6 +23,7 @@ import pickle
 
 import httpx
 import pytest
+import websockets
 
 from ... import errors
 
@@ -257,6 +258,43 @@ def test_constructor_message_not_present():
   }
 
 
+def test_constructor_with_websocket_connection_closed_error():
+  actual_error = errors.APIError(
+      1007,
+      'At most one response modality can be specified in the setup request.'
+      ' To enable simultaneous transcription and audio output,',
+      None,
+  )
+  assert actual_error.code == 1007
+  assert (
+      actual_error.details
+      == 'At most one response modality can be specified in the setup request.'
+      ' To enable simultaneous transcription and audio output,',
+  )
+  assert actual_error.status == None
+  assert actual_error.message == None
+
+
+def test_raise_for_websocket_connection_closed_error():
+  try:
+    errors.APIError.raise_error(
+        1007,
+        'At most one response modality can be specified in the setup request.'
+        ' To enable simultaneous transcription and audio output,',
+        None,
+    )
+  except errors.APIError as actual_error:
+    assert actual_error.code == 1007
+    assert (
+        actual_error.details
+        == 'At most one response modality can be specified in the setup'
+        ' request.'
+        ' To enable simultaneous transcription and audio output,'
+    )
+    assert actual_error.status == None
+    assert actual_error.message == None
+
+
 def test_raise_for_response_code_exist_json_decoder_error():
   class FakeResponse(httpx.Response):
 

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(