google-genai 1.56.0__py3-none-any.whl → 1.58.0__py3-none-any.whl

google/genai/tests/mcp/test_mcp_to_gemini_tools.py

@@ -0,0 +1,191 @@
+# 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.
+#
+
+from ... import _mcp_utils
+from ... import types
+
+try:
+  from mcp import types as mcp_types
+except ImportError as e:
+  import sys
+
+  if sys.version_info < (3, 10):
+    raise ImportError(
+        'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
+        ' version.'
+    ) from e
+  else:
+    raise e
+
+
+def test_empty_mcp_tools_list():
+  """Test conversion of empty MCP tools list to Gemini tools list."""
+  result = _mcp_utils.mcp_to_gemini_tools([])
+
+  assert result == []
+
+
+def test_unknown_field_conversion():
+  """Test conversion of MCP tools with unknown fields to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'properties': {},
+              'unknown_field': 'unknownField',
+              'unknown_object': {},
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={},
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_items_conversion():
+  """Test conversion of MCP tools with items to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'array',
+              'items': {
+                  'type': 'object',
+                  'properties': {
+                      'key1': {
+                          'type': 'string',
+                      },
+                      'key2': {
+                          'type': 'number',
+                      },
+                  },
+              },
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='ARRAY',
+                      items=types.Schema(
+                          type='OBJECT',
+                          properties={
+                              'key1': types.Schema(type='STRING'),
+                              'key2': types.Schema(type='NUMBER'),
+                          },
+                      ),
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_any_of_conversion():
+  """Test conversion of MCP tools with any_of to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'any_of': [
+                  {
+                      'type': 'string',
+                  },
+                  {
+                      'type': 'number',
+                  },
+              ],
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      any_of=[
+                          types.Schema(type='STRING'),
+                          types.Schema(type='NUMBER'),
+                      ],
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_properties_conversion():
+  """Test conversion of MCP tools with properties to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'properties': {
+                  'key1': {
+                      'type': 'string',
+                  },
+                  'key2': {
+                      'type': 'number',
+                  },
+              },
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      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'),
+                      },
+                  ),
+              ),
+          ],
+      ),
+  ]

google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py

@@ -0,0 +1,201 @@
+# 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.
+#
+
+import _asyncio
+import pytest
+from ... import _extra_utils
+from ... import types
+from ..._adapters import McpToGenAiToolAdapter
+
+try:
+  from mcp import types as mcp_types
+  from mcp import ClientSession as McpClientSession
+except ImportError as e:
+  import sys
+
+  if sys.version_info < (3, 10):
+    raise ImportError(
+        'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
+        ' version.'
+    ) from e
+  else:
+    raise e
+
+
+@pytest.mark.asyncio
+async def test_parse_empty_config_dict():
+  """Test conversion of empty GenerateContentConfigDict to parsed config."""
+  config = {}
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert parsed_config == None
+  assert not mcp_to_genai_tool_adapters
+
+
+@pytest.mark.asyncio
+async def test_parse_empty_config_object():
+  """Test conversion of empty GenerateContentConfig to parsed config."""
+  config = types.GenerateContentConfig()
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert config is not parsed_config  # config is not modified
+  assert not mcp_to_genai_tool_adapters
+
+
+@pytest.mark.asyncio
+async def test_parse_config_object_with_tools():
+  """Test conversion of GenerateContentConfig with tools to parsed config."""
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+
+    async def list_tools(self):
+      return mcp_types.ListToolsResult(
+          tools=[
+              mcp_types.Tool(
+                  name='get_weather',
+                  description='Get the weather in a city.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+              mcp_types.Tool(
+                  name='get_weather_2',
+                  description='Different tool to get the weather.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+          ]
+      )
+
+  mock_session_instance = MockMcpClientSession()
+  config = types.GenerateContentConfig(tools=[mock_session_instance])
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 1
+  assert config.tools[0] is mock_session_instance
+  assert config is not parsed_config  # config is not modified
+  assert len(mcp_to_genai_tool_adapters) == 2
+  assert mcp_to_genai_tool_adapters.keys() == {
+      'get_weather',
+      'get_weather_2',
+  }
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
+  )
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
+  )
+
+
+@pytest.mark.asyncio
+async def test_parse_config_object_with_tools_complex_type():
+  """Test conversion of GenerateContentConfig with tools to parsed config."""
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+      self._future = _asyncio.Future()  # This object cannot be pickled.
+
+    async def list_tools(self):
+      return mcp_types.ListToolsResult(
+          tools=[
+              mcp_types.Tool(
+                  name='get_weather',
+                  description='Get the weather in a city.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+              mcp_types.Tool(
+                  name='get_weather_2',
+                  description='Different tool to get the weather.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+          ]
+      )
+
+  mock_session_instance = MockMcpClientSession()
+  config = types.GenerateContentConfig(tools=[mock_session_instance])
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 1
+  assert config.tools[0] is mock_session_instance
+  assert config is not parsed_config  # config is not modified
+  assert len(mcp_to_genai_tool_adapters) == 2
+  assert mcp_to_genai_tool_adapters.keys() == {
+      'get_weather',
+      'get_weather_2',
+  }
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
+  )
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
+  )
+
+
+@pytest.mark.asyncio
+async def test_parse_config_object_with_non_mcp_tools():
+  """Test conversion of GenerateContentConfig with regular tools to parsed config."""
+
+  config = types.GenerateContentConfig(
+      tools=[
+          types.Tool(
+              function_declarations=[
+                  {'name': 'tool-1', 'description': 'tool-1-description'}
+              ]
+          ),
+          types.Tool(
+              function_declarations=[
+                  {'name': 'tool-2', 'description': 'tool-2-description'}
+              ]
+          ),
+      ]
+  )
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 2
+  assert config is not parsed_config  # config is not modified
+  assert mcp_to_genai_tool_adapters == {}
+  assert parsed_config.tools == [
+      types.Tool(
+          function_declarations=[
+              {'name': 'tool-1', 'description': 'tool-1-description'}
+          ]
+      ),
+      types.Tool(
+          function_declarations=[
+              {'name': 'tool-2', 'description': 'tool-2-description'}
+          ]
+      ),
+  ]

google/genai/tests/mcp/test_parse_config_for_mcp_usage.py

@@ -0,0 +1,130 @@
+# 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.
+#
+
+import re
+import pytest
+from ... import _extra_utils
+from ... import types
+from ..._adapters import McpToGenAiToolAdapter
+
+try:
+  from mcp import types as mcp_types
+  from mcp import ClientSession as McpClientSession
+except ImportError as e:
+  import sys
+
+  if sys.version_info < (3, 10):
+    raise ImportError(
+        'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
+        ' version.'
+    ) from e
+  else:
+    raise e
+
+
+def test_parse_empty_config_dict():
+  """Test conversion of empty GenerateContentConfigDict to parsed config."""
+  config = {}
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert parsed_config == None
+
+
+def test_parse_empty_config_object():
+  """Test conversion of empty GenerateContentConfig to parsed config."""
+  config = types.GenerateContentConfig()
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert config is not parsed_config  # config is not modified
+
+
+def test_parse_config_object_with_tools():
+  """Test conversion of GenerateContentConfig with tools to parsed config."""
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+
+    async def list_tools(self):
+      return mcp_types.ListToolsResult(
+          tools=[
+              mcp_types.Tool(
+                  name='get_weather',
+                  description='Get the weather in a city.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+              mcp_types.Tool(
+                  name='get_weather_2',
+                  description='Different tool to get the weather.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+          ]
+      )
+
+  mock_session_instance = MockMcpClientSession()
+  config = types.GenerateContentConfig(tools=[mock_session_instance])
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert config.http_options is None
+  assert config is not parsed_config  # config is not modified
+  assert re.match(
+      r'mcp_used/\d+\.\d+\.\d+',
+      parsed_config.http_options.headers['x-goog-api-client'],
+  )
+
+
+def test_parse_config_object_with_tools_and_existing_headers():
+  """Test conversion of GenerateContentConfig with tools and existing headers to parsed config."""
+
+  config = types.GenerateContentConfig(
+      tools=[
+          mcp_types.Tool(
+              name='get_weather',
+              description='Get the weather in a city.',
+              inputSchema={
+                  'type': 'object',
+                  'properties': {'location': {'type': 'string'}},
+              },
+          ),
+          mcp_types.Tool(
+              name='get_weather_2',
+              description='Different tool to get the weather.',
+              inputSchema={
+                  'type': 'object',
+                  'properties': {'location': {'type': 'string'}},
+              },
+          ),
+      ],
+      http_options=types.HttpOptions(
+          headers={
+              'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'
+          }
+      ),
+  )
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert not re.match(
+      r'mcp_used/\d+\.\d+\.\d+',
+      config.http_options.headers['x-goog-api-client'],
+  )
+  assert config is not parsed_config  # config is not modified
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      parsed_config.http_options.headers['x-goog-api-client'],
+  )

google/genai/tests/mcp/test_set_mcp_usage_header.py

@@ -0,0 +1,72 @@
+# 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.
+#
+
+from importlib.metadata import version
+import re
+import typing
+from typing import Any
+
+from ... import _mcp_utils
+from ... import types
+
+_is_mcp_imported = False
+if typing.TYPE_CHECKING:
+  import mcp
+
+  _is_mcp_imported = True
+else:
+  try:
+    import mcp
+
+    _is_mcp_imported = True
+  except ImportError:
+    _is_mcp_imported = False
+
+
+def test_set_mcp_usage_header_from_empty_dict():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an empty dict."""
+  headers = {}
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(r'mcp_used/\d+\.\d+\.\d+', headers['x-goog-api-client'])
+
+
+def test_set_mcp_usage_header_with_existing_header():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an existing header."""
+  headers = {'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'}
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      headers['x-goog-api-client'],
+  )
+
+
+def test_set_mcp_usage_header_with_existing_mcp_header():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an existing MCP header."""
+  headers = {
+      'x-goog-api-client': (
+          'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/1.0.0'
+      )
+  }
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      headers['x-goog-api-client'],
+  )

google/genai/tests/models/__init__.py

@@ -0,0 +1,17 @@
+# 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 the Google GenAI SDK's models module."""

google/genai/tests/models/constants.py

@@ -0,0 +1,8 @@
+VERTEX_HTTP_OPTIONS = {
+    'api_version': 'v1beta1',
+    'base_url': 'https://us-central1-aiplatform.googleapis.com/',
+}
+MLDEV_HTTP_OPTIONS = {
+    'api_version': 'v1beta',
+    'base_url': 'https://generativelanguage.googleapis.com/',
+}