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

google/genai/tests/transformers/test_t_tool.py

@@ -0,0 +1,157 @@
+# 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_tool."""
+
+import typing
+from typing import Any
+import pytest
+from ... import _transformers as t
+from ... import client as google_genai_client_module
+from ... import types
+
+_is_mcp_imported = False
+if typing.TYPE_CHECKING:
+  from mcp import types as mcp_types
+
+  _is_mcp_imported = True
+else:
+  try:
+    from mcp import types as mcp_types
+
+    _is_mcp_imported = True
+  except ImportError:
+    _is_mcp_imported = False
+
+
+@pytest.fixture
+def client(use_vertex):
+  if use_vertex:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, project='test-project', location='test-location'
+    )
+  else:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, api_key='test-api-key'
+    )
+
+
+@pytest.mark.usefixtures('client')
+def test_none(client):
+  assert t.t_tool(client, None) == None
+
+
+@pytest.mark.usefixtures('client')
+def test_function(client):
+  def test_func(arg1: str, arg2: int):
+    pass
+
+  assert t.t_tool(client, test_func) == types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='test_func',
+              parameters=types.Schema(
+                  type='OBJECT',
+                  properties={
+                      'arg1': types.Schema(type='STRING'),
+                      'arg2': types.Schema(type='INTEGER'),
+                  },
+                  required=['arg1', 'arg2'],
+              ),
+          )
+      ]
+  )
+
+
+@pytest.mark.usefixtures('client')
+def test_dictionary(client):
+  assert t.t_tool(
+      client,
+      {
+          'function_declarations': [{
+              'name': 'tool',
+              'description': 'tool-description',
+              'parameters': {'type': 'OBJECT', 'properties': {}},
+          }]
+      },
+  ) == types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='tool',
+              description='tool-description',
+              parameters=types.Schema(
+                  type='OBJECT',
+                  properties={},
+              ),
+          )
+      ]
+  )
+
+
+@pytest.mark.usefixtures('client')
+def test_mcp_tool(client):
+  if not _is_mcp_imported:
+    return
+
+  mcp_tool = mcp_types.Tool(
+      name='tool',
+      description='tool-description',
+      inputSchema={
+          'type': 'object',
+          'properties': {
+              'key1': {
+                  'type': 'string',
+              },
+              'key2': {
+                  'type': 'number',
+              },
+          },
+      },
+  )
+  assert t.t_tool(client, mcp_tool) == types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='tool',
+              description='tool-description',
+              parameters=types.Schema(
+                  type='OBJECT',
+                  properties={
+                      'key1': types.Schema(type='STRING'),
+                      'key2': types.Schema(type='NUMBER'),
+                  },
+              ),
+          )
+      ]
+  )
+
+
+@pytest.mark.usefixtures('client')
+def test_tool(client):
+  tool = types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='tool',
+              description='tool-description',
+              parameters=types.Schema(
+                  type='OBJECT',
+                  properties={
+                      'key1': types.Schema(type='STRING'),
+                      'key2': types.Schema(type='NUMBER'),
+                  },
+              ),
+          )
+      ]
+  )
+  assert t.t_tool(client, tool) == tool

google/genai/tests/transformers/test_t_tools.py

@@ -0,0 +1,195 @@
+# 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_tools."""
+
+import typing
+from typing import Any
+import pytest
+from ... import _transformers as t
+from ... import client as google_genai_client_module
+from ... import types
+
+_is_mcp_imported = False
+if typing.TYPE_CHECKING:
+  from mcp import ClientSession as McpClientSession
+  from mcp import types as mcp_types
+
+  _is_mcp_imported = True
+else:
+  try:
+    from mcp import ClientSession as McpClientSession
+    from mcp import types as mcp_types
+
+    _is_mcp_imported = True
+  except ImportError:
+    _is_mcp_imported = False
+
+
+@pytest.fixture
+def client(use_vertex):
+  if use_vertex:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, project='test-project', location='test-location'
+    )
+  else:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, api_key='test-api-key'
+    )
+
+
+@pytest.mark.usefixtures('client')
+def test_empty(client):
+  assert t.t_tools(client, []) == []
+
+
+@pytest.mark.usefixtures('client')
+def test_function(client):
+  def test_func(arg1: str, arg2: int):
+    pass
+
+  assert t.t_tools(client, [test_func]) == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='test_func',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={
+                          'arg1': types.Schema(type='STRING'),
+                          'arg2': types.Schema(type='INTEGER'),
+                      },
+                      required=['arg1', 'arg2'],
+                  ),
+              )
+          ]
+      )
+  ]
+
+
+@pytest.mark.usefixtures('client')
+def test_tool(client):
+  tool = types.Tool(
+      function_declarations=[
+          types.FunctionDeclaration(
+              name='tool',
+              description='tool-description',
+              parameters=types.Schema(
+                  type='OBJECT',
+                  properties={
+                      'key1': types.Schema(type='STRING'),
+                      'key2': types.Schema(type='NUMBER'),
+                  },
+              ),
+          )
+      ]
+  )
+  assert t.t_tools(client, [tool]) == [tool]
+
+
+@pytest.mark.usefixtures('client')
+def test_mcp_tool(client):
+  if not _is_mcp_imported:
+    return
+
+  mcp_tool = mcp_types.Tool(
+      name='tool',
+      description='tool-description',
+      inputSchema={
+          'type': 'object',
+          'properties': {
+              'key1': {
+                  'type': 'string',
+              },
+              'key2': {
+                  'type': 'number',
+              },
+          },
+      },
+  )
+  assert t.t_tools(client, [mcp_tool]) == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={
+                          'key1': types.Schema(type='STRING'),
+                          'key2': types.Schema(type='NUMBER'),
+                      },
+                  ),
+              )
+          ]
+      )
+  ]
+
+
+@pytest.mark.usefixtures('client')
+def test_multiple_tools(client):
+  if not _is_mcp_imported:
+    return
+
+  mcp_tool1 = mcp_types.Tool(
+      name='tool1',
+      description='tool1-description',
+      inputSchema={
+          'type': 'object',
+          'properties': {
+              'key1': {
+                  'type': 'string',
+              },
+          },
+      },
+  )
+  mcp_tool2 = mcp_types.Tool(
+      name='tool2',
+      description='tool2-description',
+      inputSchema={
+          'type': 'object',
+          'properties': {
+              'key1': {
+                  'type': 'number',
+              },
+          },
+      },
+  )
+  assert t.t_tools(client, [mcp_tool1, mcp_tool2]) == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool1',
+                  description='tool1-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={
+                          'key1': types.Schema(type='STRING'),
+                      },
+                  ),
+              ),
+              types.FunctionDeclaration(
+                  name='tool2',
+                  description='tool2-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={
+                          'key1': types.Schema(type='NUMBER'),
+                      },
+                  ),
+              ),
+          ]
+      ),
+  ]

google/genai/tests/tunings/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+#
+
+

google/genai/tests/tunings/test_cancel.py

@@ -0,0 +1,39 @@
+# 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 tunings.cancel()."""
+
+from ... import types as genai_types
+from .. import pytest_helper
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name="test_cancel_vertex",
+        parameters=genai_types._CancelTuningJobParameters(
+            name="projects/801452371447/locations/us-central1/tuningJobs/2470687341547618304"
+        ),
+        exception_if_mldev="Not Found",
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="tunings.cancel",
+    test_table=test_table,
+)
+
+pytest_plugins = ("pytest_asyncio",)

google/genai/tests/tunings/test_end_to_end.py

@@ -0,0 +1,106 @@
+# 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 tunings.tune()."""
+
+import time
+from ... import _replay_api_client
+from ... import types as genai_types
+from .. import pytest_helper
+
+
+test_table: list[pytest_helper.TestTableItem] = []
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="tunings.tune",
+    test_table=test_table,
+)
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+def test_tune_until_success(client):
+  if client._api_client.vertexai:
+    job = client.tunings.tune(
+        base_model="gemini-2.5-flash-lite",
+        training_dataset=genai_types.TuningDataset(
+            gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
+        ),
+        config=genai_types.CreateTuningJobConfig(
+            epoch_count=1,
+        ),
+    )
+  else:
+    # Remove GenAI SDK test since it is deprecated:
+    # https://ai.google.dev/gemini-api/docs/model-tuning
+    return
+
+  while not job.has_ended:
+    # Skipping the sleep for when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(300)
+    job = client.tunings.get(name=job.name)
+
+  assert job.has_ended
+  assert job.has_succeeded
+
+
+def test_continuous_tuning(client):
+  if not client._api_client.vertexai:
+    return
+
+  job = client.tunings.tune(
+      base_model="gemini-2.5-flash",
+      training_dataset=genai_types.TuningDataset(
+          gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
+      ),
+      config=genai_types.CreateTuningJobConfig(
+          epoch_count=1,
+      ),
+  )
+
+  while not job.has_ended:
+    # Skipping the sleep for when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(300)
+    job = client.tunings.get(name=job.name)
+
+  assert job.has_ended
+  assert job.has_succeeded
+
+  continuous_job = client.tunings.tune(
+      base_model=job.tuned_model.model,
+      training_dataset=genai_types.TuningDataset(
+          gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
+      ),
+      config=genai_types.CreateTuningJobConfig(
+          tuned_model_display_name="continuous tuning job",
+          epoch_count=1,
+      )
+  )
+
+  while not continuous_job.has_ended:
+    # Skipping the sleep for when in replay mode.
+    if client._api_client._mode not in ("replay", "auto"):
+      time.sleep(300)
+    continuous_job = client.tunings.get(name=continuous_job.name)
+
+  assert continuous_job.has_ended
+  assert continuous_job.has_succeeded
+
+

google/genai/tests/tunings/test_get.py

@@ -0,0 +1,67 @@
+# 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 tunings.get()."""
+
+from ... import types as genai_types
+from .. import pytest_helper
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name="test_vertexai",
+        parameters=genai_types._GetTuningJobParameters(
+            name="projects/801452371447/locations/us-central1/tuningJobs/4303478340632707072"
+        ),
+        exception_if_mldev="Not Found",
+    ),
+    pytest_helper.TestTableItem(
+        name="test_vertexai_with_pretuned_model",
+        parameters=genai_types._GetTuningJobParameters(
+            name="projects/801452371447/locations/us-central1/tuningJobs/8520221517529743360",
+        ),
+        exception_if_mldev="Not Found",
+    ),
+    pytest_helper.TestTableItem(
+        name="test_mldev",
+        parameters=genai_types._GetTuningJobParameters(
+            name="tunedModels/generate-num-1896"
+        ),
+        exception_if_vertex="Not Found",
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="tunings.get",
+    test_table=test_table,
+)
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+def test_helper_properties(client):
+  if client._api_client.vertexai:
+    job = client.tunings.get(
+        name="projects/801452371447/locations/us-central1/tuningJobs/4303478340632707072",
+    )
+  else:
+    job = client.tunings.get(
+        name="tunedModels/testdatasetexamples-model-j0fpgpaksvri"
+    )
+
+  assert job.has_ended
+  assert job.has_succeeded

google/genai/tests/tunings/test_list.py

@@ -0,0 +1,75 @@
+# 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 tunings.list()."""
+
+import pytest
+
+from ... import types as genai_types
+from .. import pytest_helper
+
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_default',
+        parameters=genai_types._ListTuningJobsParameters(),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_config',
+        parameters=genai_types._ListTuningJobsParameters(
+            config=genai_types.ListTuningJobsConfig(page_size=2)
+        ),
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='tunings.list',
+    test_table=test_table,
+)
+
+pytest_plugins = ('pytest_asyncio',)
+
+
+def test_pager(client):
+  tuning_jobs = client.tunings.list(config={'page_size': 2})
+  assert 'content-type' in tuning_jobs.sdk_http_response.headers
+  assert tuning_jobs.name == 'tuning_jobs'
+  assert tuning_jobs.page_size == 2
+  assert len(tuning_jobs) <= 2
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  for _ in tuning_jobs:
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    tuning_jobs.next_page()
+
+
+@pytest.mark.asyncio
+async def test_async_pager(client):
+  tuning_jobs = await client.aio.tunings.list(config={'page_size': 2})
+
+  assert 'Content-Type' in tuning_jobs.sdk_http_response.headers
+  assert tuning_jobs.name == 'tuning_jobs'
+  assert tuning_jobs.page_size == 2
+  assert len(tuning_jobs) <= 2
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  async for _ in tuning_jobs:
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    await tuning_jobs.next_page()