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

google/genai/batches.py

@@ -1013,6 +1013,11 @@ def _GenerateContentConfig_to_mldev(
         getv(from_object, ['enable_enhanced_civic_answers']),
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    raise ValueError(
+        'model_armor_config parameter is not supported in Gemini API.'
+    )
+
   return to_object
 
 
@@ -1203,6 +1208,9 @@ def _InlinedResponse_from_mldev(
         ),
     )
 
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
   if getv(from_object, ['error']) is not None:
     setv(to_object, ['error'], getv(from_object, ['error']))
 

google/genai/files.py

@@ -101,6 +101,17 @@ def _GetFileParameters_to_mldev(
   return to_object
 
 
+def _InternalRegisterFilesParameters_to_mldev(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['uris']) is not None:
+    setv(to_object, ['uris'], getv(from_object, ['uris']))
+
+  return to_object
+
+
 def _ListFilesConfig_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -152,17 +163,6 @@ def _ListFilesResponse_from_mldev(
   return to_object
 
 
-def _RegisterFilesParameters_to_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['uris']) is not None:
-    setv(to_object, ['uris'], getv(from_object, ['uris']))
-
-  return to_object
-
-
 def _RegisterFilesResponse_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -438,7 +438,7 @@ class Files(_api_module.BaseModule):
       uris: list[str],
       config: Optional[types.RegisterFilesConfigOrDict] = None,
   ) -> types.RegisterFilesResponse:
-    parameter_model = types._RegisterFilesParameters(
+    parameter_model = types._InternalRegisterFilesParameters(
         uris=uris,
         config=config,
     )
@@ -449,7 +449,7 @@ class Files(_api_module.BaseModule):
           'This method is only supported in the Gemini Developer client.'
       )
     else:
-      request_dict = _RegisterFilesParameters_to_mldev(parameter_model)
+      request_dict = _InternalRegisterFilesParameters_to_mldev(parameter_model)
       request_url_dict = request_dict.get('_url')
       if request_url_dict:
         path = 'files:register'.format_map(request_url_dict)
@@ -977,7 +977,7 @@ class AsyncFiles(_api_module.BaseModule):
       uris: list[str],
       config: Optional[types.RegisterFilesConfigOrDict] = None,
   ) -> types.RegisterFilesResponse:
-    parameter_model = types._RegisterFilesParameters(
+    parameter_model = types._InternalRegisterFilesParameters(
         uris=uris,
         config=config,
     )
@@ -988,7 +988,7 @@ class AsyncFiles(_api_module.BaseModule):
           'This method is only supported in the Gemini Developer client.'
       )
     else:
-      request_dict = _RegisterFilesParameters_to_mldev(parameter_model)
+      request_dict = _InternalRegisterFilesParameters_to_mldev(parameter_model)
       request_url_dict = request_dict.get('_url')
       if request_url_dict:
         path = 'files:register'.format_map(request_url_dict)

google/genai/models.py

@@ -1117,6 +1117,11 @@ def _GenerateContentConfig_to_mldev(
         getv(from_object, ['enable_enhanced_civic_answers']),
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    raise ValueError(
+        'model_armor_config parameter is not supported in Gemini API.'
+    )
+
   return to_object
 
 
@@ -1279,6 +1284,13 @@ def _GenerateContentConfig_to_vertex(
         'enable_enhanced_civic_answers parameter is not supported in Vertex AI.'
     )
 
+  if getv(from_object, ['model_armor_config']) is not None:
+    setv(
+        parent_object,
+        ['modelArmorConfig'],
+        getv(from_object, ['model_armor_config']),
+    )
+
   return to_object
 
 

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/models/test_generate_content.py

@@ -546,6 +546,22 @@ test_table: list[pytest_helper.TestTableItem] = [
         ),
         exception_if_vertex='not supported',
     ),
+    pytest_helper.TestTableItem(
+        name='test_model_armor_config',
+        parameters=types._GenerateContentParameters(
+            model=GEMINI_FLASH_LATEST,
+            contents=t.t_contents('What is your name?'),
+            config={
+                'model_armor_config': {
+                    'prompt_template_name': '',
+                    'response_template_name': '',
+                    # Intentionally left blank just to test that the SDK doesn't
+                    # throw an exception.
+                },
+            },
+        ),
+        exception_if_mldev='not supported',
+    ),
 ]
 
 pytestmark = pytest_helper.setup(

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(