google-genai 1.54.0__py3-none-any.whl → 1.55.0__py3-none-any.whl

google/genai/tests/transformers/test_t_batch.py

@@ -0,0 +1,286 @@
+# Copyright 2025 Google LLC
+#
+# Licensed 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.
+#
+
+"""Tests for t_batch."""
+
+import pytest
+from ... import _transformers as t
+from ... import client as google_genai_client_module
+from ... import types
+
+
+@pytest.fixture
+def vertex_client():
+  yield google_genai_client_module.Client(
+      vertexai=True, project='test-project', location='test-location'
+  )
+
+
+@pytest.fixture
+def mldev_client():
+  yield google_genai_client_module.Client(
+      vertexai=False, api_key='test-api-key'
+  )
+
+
+@pytest.mark.usefixtures('vertex_client', 'mldev_client')
+class TestBatchJobSource:
+
+  # Test cases for src as str
+  @pytest.mark.parametrize(
+      'src_str, expected_gcs_uri, expected_bigquery_uri, expected_file_name,'
+      ' expected_format',
+      [
+          (
+              'gs://bucket/path/to/data.jsonl',
+              ['gs://bucket/path/to/data.jsonl'],
+              None,
+              None,
+              'jsonl',
+          ),
+          (
+              'bq://project.dataset.table',
+              None,
+              'bq://project.dataset.table',
+              None,
+              'bigquery',
+          ),
+          ('files/data.csv', None, None, 'files/data.csv', None),
+      ],
+  )
+  def test_batch_job_source_str(
+      self,
+      mldev_client,
+      src_str,
+      expected_gcs_uri,
+      expected_bigquery_uri,
+      expected_file_name,
+      expected_format,
+  ):
+    result = t.t_batch_job_source(mldev_client, src_str)
+    expected = types.BatchJobSource(
+        gcs_uri=expected_gcs_uri,
+        bigquery_uri=expected_bigquery_uri,
+        file_name=expected_file_name,
+        format=expected_format,
+    )
+    assert result == expected
+
+  def test_batch_job_source_str_unsupported(self, mldev_client):
+    with pytest.raises(ValueError, match='Unsupported source'):
+      t.t_batch_job_source(mldev_client, 'http://example.com/data')
+    with pytest.raises(ValueError, match='Unsupported source'):
+      t.t_batch_job_source(mldev_client, 'invalid-path')
+
+  # Test cases for src as list
+  def test_batch_job_source_list_empty(self, mldev_client):
+    result = t.t_batch_job_source(mldev_client, [])
+    expected = types.BatchJobSource(inlined_requests=[])
+    assert result == expected
+
+  def test_batch_job_source_list_with_items(self, mldev_client):
+    inlined_data = [
+        types.InlinedRequest(
+            contents=[types.Content(parts=[types.Part(text='item1')])]
+        ),
+        types.InlinedRequest(
+            contents=[types.Content(parts=[types.Part(text='item2')])]
+        ),
+    ]
+    result = t.t_batch_job_source(mldev_client, inlined_data)
+    expected = types.BatchJobSource(inlined_requests=inlined_data)
+    assert result == expected
+
+  # Test cases for src as dict
+  def test_batch_job_source_dict(self, vertex_client, mldev_client):
+    src_dict = {'gcs_uri': ['gs://test/file.jsonl'], 'format': 'jsonl'}
+    result = t.t_batch_job_source(vertex_client, src_dict)
+    expected = types.BatchJobSource(
+        gcs_uri=['gs://test/file.jsonl'], format='jsonl'
+    )
+    assert result == expected
+
+    src_dict_inlined = {
+        'inlined_requests': [{
+            'contents': [{
+                'parts': [{
+                    'text': 'Hello!',
+                }],
+                'role': 'user',
+            }],
+        }]
+    }
+    result_inlined = t.t_batch_job_source(mldev_client, src_dict_inlined)
+    expected_inlined = types.BatchJobSource(inlined_requests=[{
+            'contents': [{
+                'parts': [{
+                    'text': 'Hello!',
+                }],
+                'role': 'user',
+            }],
+        }])
+    assert result_inlined == expected_inlined
+
+  # Test cases for src as types.BatchJobSource (mocked MockBatchJobSource)
+
+  def test_batch_job_source_mldev_valid_file_name(self, mldev_client):
+    src_obj = types.BatchJobSource(file_name='files/my_data.csv')
+    result = t.t_batch_job_source(mldev_client, src_obj)
+    assert result is src_obj
+
+  def test_batch_job_source_mldev_invalid_both_set(self, mldev_client):
+    src_obj = types.BatchJobSource(
+        inlined_requests=[types.InlinedRequest()],
+        file_name='files/data.csv',
+    )
+    with pytest.raises(
+        ValueError,
+        match='`inlined_requests`, `file_name`,',
+    ):
+      t.t_batch_job_source(mldev_client, src_obj)
+
+  def test_batch_job_source_mldev_invalid_neither_set(self, mldev_client):
+    src_obj = types.BatchJobSource(
+        gcs_uri=['gs://temp']
+    )
+    with pytest.raises(
+        ValueError,
+        match='`inlined_requests`, `file_name`,',
+    ):
+      t.t_batch_job_source(mldev_client, src_obj)
+
+  # client.vertexai = True (Vertex AI)
+  def test_batch_job_source_vertexai_valid_gcs(self, vertex_client):
+    src_obj = types.BatchJobSource(gcs_uri=['gs://vertex-bucket/data.jsonl'])
+    result = t.t_batch_job_source(vertex_client, src_obj)
+    assert result is src_obj
+
+  def test_batch_job_source_vertexai_valid_bigquery(self, vertex_client):
+    src_obj = types.BatchJobSource(bigquery_uri='bq://project.dataset.table')
+    result = t.t_batch_job_source(vertex_client, src_obj)
+    assert result is src_obj
+
+  def test_batch_job_source_vertexai_valid_both(self, vertex_client):
+    src_obj = types.BatchJobSource(
+        gcs_uri=['gs://vertex-bucket/data.jsonl'],
+        bigquery_uri='bq://project.dataset.table',
+    )
+    with pytest.raises(ValueError, match='`gcs_uri` or `bigquery_uri`'):
+      t.t_batch_job_source(vertex_client, src_obj)
+
+  def test_batch_job_source_vertexai_invalid_neither_set(self, vertex_client):
+    src_obj = types.BatchJobSource(
+        file_name='files/data.csv'
+    )
+    with pytest.raises(ValueError, match='`gcs_uri` or `bigquery_uri`'):
+      t.t_batch_job_source(vertex_client, src_obj)
+
+
+class TestBatchJobDestination:
+
+  @pytest.mark.parametrize(
+      'dest_str, expected_gcs_uri, expected_bigquery_uri, expected_format',
+      [
+          (
+              'gs://bucket/path/to/output',
+              'gs://bucket/path/to/output',
+              None,
+              'jsonl',
+          ),
+          (
+              'bq://project.dataset.output_table',
+              None,
+              'bq://project.dataset.output_table',
+              'bigquery',
+          ),
+      ],
+  )
+  def test_valid_destinations(
+      self, dest_str, expected_gcs_uri, expected_bigquery_uri, expected_format
+  ):
+    result = t.t_batch_job_destination(dest_str)
+    expected = types.BatchJobDestination(
+        gcs_uri=expected_gcs_uri,
+        bigquery_uri=expected_bigquery_uri,
+        format=expected_format,
+    )
+    assert result == expected
+
+  def test_unsupported_destination(self):
+    with pytest.raises(ValueError, match='Unsupported destination'):
+      t.t_batch_job_destination('local/path/output.jsonl')
+    with pytest.raises(ValueError, match='Unsupported destination'):
+      t.t_batch_job_destination('http://some.url')
+
+
+class TestBatchJobName:
+
+  def test_mldev_valid_name(self, mldev_client):
+    name = 'batches/my-job-123'
+    assert t.t_batch_job_name(mldev_client, name) == 'my-job-123'
+
+  @pytest.mark.parametrize(
+      'invalid_name',
+      [
+          'my-job-123',
+          'batches/my-job/suffix',
+          'batches/',
+          'batches/my-job-123/',
+      ],
+  )
+  def test_mldev_invalid_name(self, mldev_client, invalid_name):
+    with pytest.raises(ValueError, match='Invalid batch job name'):
+      t.t_batch_job_name(mldev_client, invalid_name)
+
+  def test_vertexai_valid_full_path(self, vertex_client):
+    name = 'projects/123/locations/us-central1/batchPredictionJobs/my-vertex-job-456'
+    assert t.t_batch_job_name(vertex_client, name) == 'my-vertex-job-456'
+
+  def test_vertexai_valid_digit_name(self, vertex_client):
+    name = '9876543210'
+    assert t.t_batch_job_name(vertex_client, name) == '9876543210'
+
+  @pytest.mark.parametrize(
+      'invalid_name',
+      [
+          'my-vertex-job-456',
+          'projects/123/locations/us-central1/job/my-vertex-job-456',
+          'projects/123/locations/us-central1/batchPredictionJobs/',
+          'not-a-number-or-path',
+          'batchPredictionJobs/my-job',
+      ],
+  )
+  def test_vertexai_invalid_name(self, vertex_client, invalid_name):
+    with pytest.raises(ValueError, match='Invalid batch job name'):
+      t.t_batch_job_name(vertex_client, invalid_name)
+
+
+class TestJobState:
+
+  @pytest.mark.parametrize(
+      'input_state, expected_output_state',
+      [
+          ('BATCH_STATE_UNSPECIFIED', 'JOB_STATE_UNSPECIFIED'),
+          ('BATCH_STATE_PENDING', 'JOB_STATE_PENDING'),
+          ('BATCH_STATE_SUCCEEDED', 'JOB_STATE_SUCCEEDED'),
+          ('BATCH_STATE_FAILED', 'JOB_STATE_FAILED'),
+          ('BATCH_STATE_CANCELLED', 'JOB_STATE_CANCELLED'),
+          ('BATCH_STATE_EXPIRED', 'JOB_STATE_EXPIRED'),
+          ('BATCH_STATE_RUNNING', 'JOB_STATE_RUNNING'),
+          ('BATCH_STATE_FOOBAR', 'BATCH_STATE_FOOBAR'),
+      ],
+  )
+  def test_job_state_mapping(self, input_state, expected_output_state):
+    assert t.t_job_state(input_state) == expected_output_state

google/genai/tests/transformers/test_t_content.py

@@ -0,0 +1,160 @@
+# Copyright 2025 Google LLC
+#
+# Licensed 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.
+#
+
+"""Tests for t_content."""
+
+import pytest
+import pydantic
+
+from ... import _transformers as t
+from ... import types
+
+
+def test_none():
+  with pytest.raises(ValueError):
+    t.t_content( None)
+
+
+def test_content():
+  assert t.t_content(
+      types.Content(parts=[types.Part(text='test')])
+  ) == types.Content(parts=[types.Part(text='test')])
+
+
+def test_content_dict():
+  assert t.t_content(
+      {'role': 'user', 'parts': [{'text': 'test'}]}
+  ) == types.Content(parts=[types.Part(text='test')], role='user')
+
+
+def test_content_dict_invalid():
+  with pytest.raises(pydantic.ValidationError):
+    t.t_content({'invalid_key': 'test'})
+
+
+def test_text_part_dict():
+  assert t.t_content({'text': 'test'}) == types.UserContent(
+      parts=[types.Part(text='test')]
+  )
+
+
+def test_function_call_part_dict():
+  assert t.t_content(
+      {'function_call': {'name': 'test_func', 'args': {'arg1': 'value1'}}}
+  ) == (
+      types.ModelContent(
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      name='test_func', args={'arg1': 'value1'}
+                  )
+              )
+          ]
+      )
+  )
+
+
+def test_text_part():
+  assert t.t_content(types.Part(text='test')) == types.UserContent(
+      parts=[types.Part(text='test')]
+  )
+
+
+def test_function_call_part():
+  assert t.t_content(
+      types.Part(
+          function_call=types.FunctionCall(
+              name='test_func', args={'arg1': 'value1'}
+          )
+      ),
+  ) == (
+      types.ModelContent(
+          parts=[
+              types.Part(
+                  function_call=types.FunctionCall(
+                      name='test_func', args={'arg1': 'value1'}
+                  )
+              )
+          ]
+      )
+  )
+
+
+def test_string():
+  assert t.t_content('test') == types.UserContent(
+      parts=[types.Part(text='test')]
+  )
+
+
+def test_file():
+  assert t.t_content(
+      types.File(uri='gs://test', mime_type='image/png')
+  ) == types.UserContent(
+      parts=[
+          types.Part(
+              file_data=types.FileData(
+                  file_uri='gs://test', mime_type='image/png'
+              )
+          )
+      ]
+  )
+
+
+def test_file_no_uri():
+  with pytest.raises(ValueError):
+    t.t_content(types.File(mime_type='image/png'))
+
+
+def test_file_no_mime_type():
+  with pytest.raises(ValueError):
+    t.t_content(types.File(uri='gs://test'))
+
+
+def test_file_dict():
+  assert t.t_content({'file_uri': 'gs://test', 'mime_type': 'image/png'}) == types.UserContent(
+          parts=[
+              types.Part(
+                  file_data=types.FileData(
+                      file_uri='gs://test', mime_type='image/png'
+                  )
+              )
+          ]
+      )
+
+def test_int():
+  try:
+    t.t_content(1)
+  except ValueError as e:
+    assert 'Unsupported content part type: <class \'int\'>' in str(e)
+
+
+@pytest.mark.parametrize(
+    'contents',
+    [
+        types.Content(parts=[types.Part(inline_data=types.Blob(data=b'test'))]),
+        {'parts': [{'inline_data': {'data': b'test'}}]},
+        [
+            types.Content(
+                parts=[types.Part(inline_data=types.Blob(data=b'test'))]
+            )
+        ],
+        [{'parts': [{'inline_data': {'data': b'test'}}]}],
+    ],
+)
+def test_t_contents_strict_content(contents):
+  result = t.t_contents_strict(contents)
+  assert result == [
+      types.Content(parts=[types.Part(inline_data=types.Blob(data=b'test'))])
+  ]