From 4dd7b165dc54d3ae75367f68d05f9d9951688f54 Mon Sep 17 00:00:00 2001 From: Wanlin Du Date: Fri, 20 Feb 2026 12:32:11 -0800 Subject: [PATCH] feat: enable server side MCP and disable all other AFC when server side MCP is configured. PiperOrigin-RevId: 873027776 --- google/genai/_extra_utils.py | 7 +- google/genai/_live_converters.py | 10 ++ google/genai/_tokens_converters.py | 7 ++ google/genai/batches.py | 7 ++ google/genai/caches.py | 10 ++ google/genai/models.py | 16 ++- ...test_find_afc_incompatible_tool_indexes.py | 32 +++++- google/genai/tests/chats/test_send_message.py | 103 +++++++++++++++++- .../models/test_generate_content_tools.py | 65 +++++++++++ google/genai/tests/types/test_part_type.py | 4 +- google/genai/types.py | 93 ++++++++++++++++ 11 files changed, 343 insertions(+), 11 deletions(-) diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index e0fb9c105..129c05f7d 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -141,9 +141,12 @@ def find_afc_incompatible_tool_indexes( return incompatible_tools_indexes for index, tool in enumerate(config_model.tools): - if isinstance(tool, types.Tool) and tool.function_declarations: + if not isinstance(tool, types.Tool): + continue + if tool.function_declarations: + incompatible_tools_indexes.append(index) + if tool.mcp_servers: incompatible_tools_indexes.append(index) - return incompatible_tools_indexes diff --git a/google/genai/_live_converters.py b/google/genai/_live_converters.py index de3e7e77b..0bb03daf5 100644 --- a/google/genai/_live_converters.py +++ b/google/genai/_live_converters.py @@ -1397,6 +1397,13 @@ def _Tool_to_mldev( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + setv( + to_object, + ['mcpServers'], + [item for item in getv(from_object, ['mcp_servers'])], + ) + return to_object @@ -1450,6 +1457,9 @@ def _Tool_to_vertex( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + raise ValueError('mcp_servers parameter is not supported in Vertex AI.') + return to_object diff --git a/google/genai/_tokens_converters.py b/google/genai/_tokens_converters.py index 82b81d207..6e0345e6e 100644 --- a/google/genai/_tokens_converters.py +++ b/google/genai/_tokens_converters.py @@ -522,4 +522,11 @@ def _Tool_to_mldev( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + setv( + to_object, + ['mcpServers'], + [item for item in getv(from_object, ['mcp_servers'])], + ) + return to_object diff --git a/google/genai/batches.py b/google/genai/batches.py index 21dce5ce9..ef3ccec01 100644 --- a/google/genai/batches.py +++ b/google/genai/batches.py @@ -1498,6 +1498,13 @@ def _Tool_to_mldev( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + setv( + to_object, + ['mcpServers'], + [item for item in getv(from_object, ['mcp_servers'])], + ) + return to_object diff --git a/google/genai/caches.py b/google/genai/caches.py index df8aa9401..03d77f82c 100644 --- a/google/genai/caches.py +++ b/google/genai/caches.py @@ -704,6 +704,13 @@ def _Tool_to_mldev( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + setv( + to_object, + ['mcpServers'], + [item for item in getv(from_object, ['mcp_servers'])], + ) + return to_object @@ -757,6 +764,9 @@ def _Tool_to_vertex( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + raise ValueError('mcp_servers parameter is not supported in Vertex AI.') + return to_object diff --git a/google/genai/models.py b/google/genai/models.py index cd25dc44b..bf5e7d690 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -3770,6 +3770,13 @@ def _Tool_to_mldev( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + setv( + to_object, + ['mcpServers'], + [item for item in getv(from_object, ['mcp_servers'])], + ) + return to_object @@ -3824,6 +3831,9 @@ def _Tool_to_vertex( if getv(from_object, ['url_context']) is not None: setv(to_object, ['urlContext'], getv(from_object, ['url_context'])) + if getv(from_object, ['mcp_servers']) is not None: + raise ValueError('mcp_servers parameter is not supported in Vertex AI.') + return to_object @@ -5600,7 +5610,7 @@ def generate_content( 'Tools at indices [%s] are not compatible with automatic function ' 'calling (AFC). AFC is disabled. If AFC is intended, please ' 'include python callables in the tool list, and do not include ' - 'function declaration in the tool list.', + 'function declaration and MCP server in the tool list.', indices_str, ) return self._generate_content( @@ -7469,7 +7479,7 @@ async def generate_content( 'Tools at indices [%s] are not compatible with automatic function ' 'calling (AFC). AFC is disabled. If AFC is intended, please ' 'include python callables in the tool list, and do not include ' - 'function declaration in the tool list.', + 'function declaration and MCP server in the tool list.', indices_str, ) return await self._generate_content( @@ -7633,7 +7643,7 @@ async def base_async_generator(model, contents, config): # type: ignore[no-unty 'Tools at indices [%s] are not compatible with automatic function ' 'calling (AFC). AFC is disabled. If AFC is intended, please ' 'include python callables in the tool list, and do not include ' - 'function declaration in the tool list.', + 'function declaration and MCP server in the tool list.', indices_str, ) response = await self._generate_content_stream( diff --git a/google/genai/tests/afc/test_find_afc_incompatible_tool_indexes.py b/google/genai/tests/afc/test_find_afc_incompatible_tool_indexes.py index 9785f0f12..c66f9b8aa 100644 --- a/google/genai/tests/afc/test_find_afc_incompatible_tool_indexes.py +++ b/google/genai/tests/afc/test_find_afc_incompatible_tool_indexes.py @@ -83,6 +83,7 @@ def test_empty_tools_list_returns_empty_list(): def test_all_compatible_tools_returns_empty_list_with_empty_fd(): """Verifies that an empty list is returned when all tools are compatible. + A tool is compatible if it's not a `types.Tool` or if its `function_declarations` attribute is empty or None from config. """ @@ -118,6 +119,7 @@ def test_all_compatible_tools_returns_empty_list_with_empty_fd(): def test_all_compatible_tools_returns_empty_list_with_none_fd(): """Verifies that an empty list is returned when all tools are compatible. + A tool is compatible if it's not a `types.Tool` or if its `function_declarations` attribute is empty or None from config. """ @@ -153,6 +155,7 @@ def test_all_compatible_tools_returns_empty_list_with_none_fd(): def test_all_compatible_tools_returns_empty_list(): """Verifies that an empty list is returned when all tools are compatible. + A tool is compatible if it's not a `types.Tool` or if its `function_declarations` attribute is empty or None from config. """ @@ -211,8 +214,7 @@ def test_single_incompatible_tool(): def test_multiple_incompatible_tools(): - """Verifies that all correct indexes are returned for multiple incompatible - tools. """ + """Verifies correct indexes are returned for multiple incompatible tools.""" result = find_afc_incompatible_tool_indexes( config=types.GenerateContentConfig( tools=[ @@ -238,3 +240,29 @@ def test_multiple_incompatible_tools(): ) ) assert result == [2, 5] + +def test_mcp_tool_incompatible(): + """Verifies correct indexes are returned for multiple incompatible tools.""" + result = find_afc_incompatible_tool_indexes( + config=types.GenerateContentConfig( + tools=[ + types.Tool( + google_search_retrieval=types.GoogleSearchRetrieval() + ), + types.Tool(retrieval=types.Retrieval()), + types.Tool( + function_declarations=[ + types.FunctionDeclaration(name='test_function') + ] + ), + types.Tool(code_execution=types.ToolCodeExecution()), + + get_weather_tool, + mcp_to_genai_tool_adapter, + types.Tool( + mcp_servers=[types.McpServer(name='test_mcp_server')] + ), + ] + ) + ) + assert result == [2, 6] diff --git a/google/genai/tests/chats/test_send_message.py b/google/genai/tests/chats/test_send_message.py index 9cfb5acd8..cee6bfe3c 100644 --- a/google/genai/tests/chats/test_send_message.py +++ b/google/genai/tests/chats/test_send_message.py @@ -768,9 +768,8 @@ def test_mcp_tools(client): ) ],}, ) - response = chat.send_message('What is the weather in Boston?'); - response = chat.send_message('What is the weather in San Francisco?'); - + response = chat.send_message('What is the weather in Boston?') + response = chat.send_message('What is the weather in San Francisco?') def test_mcp_tools_stream(client): @@ -842,3 +841,101 @@ async def test_async_mcp_tools_stream(client): 'What is the weather in San Francisco?' ): pass + + +def test_server_side_mcp_tools(client): + with pytest_helper.exception_if_vertex(client, ValueError): + chat = client.chats.create( + model='gemini-2.5-flash', + config={ + 'tools': [ + { + 'mcp_servers': [ + { + 'name': 'weather_server', + 'streamable_http_transport': { + 'url': ( + 'https://gemini-api-demos.uc.r.appspot.com/mcp' + ), + 'headers': { + 'AUTHORIZATION': 'Bearer github_pat_XXXX', + }, + 'timeout': '10s', + }, + }, + ], + }, + ], + }, + ) + response = chat.send_message('What is the weather in Boston on 02/02/2026?') + response = chat.send_message( + 'What is the weather in San Francisco on 02/02/2026?' + ) + + +def test_server_side_mcp_tools_stream(client): + with pytest_helper.exception_if_vertex(client, ValueError): + chat = client.chats.create( + model='gemini-2.5-flash', + config={ + 'tools': [ + { + 'mcp_servers': [ + { + 'name': 'weather_server', + 'streamable_http_transport': { + 'url': ( + 'https://gemini-api-demos.uc.r.appspot.com/mcp' + ), + 'headers': { + 'AUTHORIZATION': 'Bearer github_pat_XXXX', + }, + 'timeout': '10s', + }, + }, + ], + }, + ], + }, + ) + for chunk in chat.send_message_stream( + 'What is the weather in Boston on 02/02/2026?' + ): + pass + for chunk in chat.send_message_stream( + 'What is the weather in San Francisco on 02/02/2026?' + ): + pass + + +@pytest.mark.asyncio +async def test_async_server_side_mcp_tools(client): + with pytest_helper.exception_if_vertex(client, ValueError): + chat = client.aio.chats.create( + model='gemini-2.5-flash', + config={ + 'tools': [ + { + 'mcp_servers': [ + { + 'name': 'weather_server', + 'streamable_http_transport': { + 'url': ( + 'https://gemini-api-demos.uc.r.appspot.com/mcp' + ), + 'headers': { + 'AUTHORIZATION': 'Bearer github_pat_XXXX', + }, + 'timeout': '10s', + }, + }, + ], + }, + ], + }, + ) + await chat.send_message('What is the weather in Boston on 02/02/2026?') + await chat.send_message( + 'What is the weather in San Francisco on 02/02/2026?' + ) diff --git a/google/genai/tests/models/test_generate_content_tools.py b/google/genai/tests/models/test_generate_content_tools.py index bd5ca0c11..142d9f980 100644 --- a/google/genai/tests/models/test_generate_content_tools.py +++ b/google/genai/tests/models/test_generate_content_tools.py @@ -1759,3 +1759,68 @@ async def test_function_declaration_with_callable_async_stream(client): }, ): pass + +def test_server_side_mcp_only(client): + """Test server side mcp, happy path.""" + with pytest_helper.exception_if_vertex(client, ValueError): + response = client.models.generate_content( + model='gemini-2.5-pro', + contents=('What is the weather like in New York (NY) on 02/02/2026?'), + config=types.GenerateContentConfig( + tools=[types.Tool( + mcp_servers=[types.McpServer( + name='get_weather', + streamable_http_transport=types.StreamableHttpTransport( + url='https://gemini-api-demos.uc.r.appspot.com/mcp', + headers={'AUTHORIZATION': 'Bearer github_pat_XXXX'}, + ), + )] + )] + ) + ) + assert response.text + +@pytest.mark.asyncio +async def test_server_side_mcp_only_async(client): + """Test server side mcp, happy path.""" + with pytest_helper.exception_if_vertex(client, ValueError): + response = await client.aio.models.generate_content( + model='gemini-2.5-flash', + contents=( + 'What is the weather like in New York on 02/02/2026?' + ), + config=types.GenerateContentConfig( + tools=[types.Tool( + mcp_servers=[types.McpServer( + name='get_weather', + streamable_http_transport=types.StreamableHttpTransport( + url='https://gemini-api-demos.uc.r.appspot.com/mcp', + headers={'AUTHORIZATION': 'Bearer github_pat_XXXX'}, + ), + )] + + )] + ) + ) + assert response.text + +def test_server_side_mcp_only_stream(client): + """Test server side mcp, happy path.""" + with pytest_helper.exception_if_vertex(client, ValueError): + response = client.models.generate_content_stream( + model='gemini-2.5-pro', + contents=('What is the weather like in New York (NY) on 02/02/2026?'), + config=types.GenerateContentConfig( + tools=[types.Tool( + mcp_servers=[types.McpServer( + name='get_weather', + streamable_http_transport=types.StreamableHttpTransport( + url='https://gemini-api-demos.uc.r.appspot.com/mcp', + headers={'AUTHORIZATION': 'Bearer github_pat_XXXX'}, + ), + )] + )] + ) + ) + for chunk in response: + pass diff --git a/google/genai/tests/types/test_part_type.py b/google/genai/tests/types/test_part_type.py index 8d277a76c..a856b15de 100644 --- a/google/genai/tests/types/test_part_type.py +++ b/google/genai/tests/types/test_part_type.py @@ -247,6 +247,9 @@ def test_none_empty_text(): def test_non_text_part_text(caplog, generate_content_response): + from ... import types as types_module + types_module._response_text_non_text_warning_logged = False + generate_content_response.candidates = [ types.Candidate( content=types.Content( @@ -256,7 +259,6 @@ def test_non_text_part_text(caplog, generate_content_response): ) ), ] - assert generate_content_response.text is None assert any( record.levelname == 'WARNING' diff --git a/google/genai/types.py b/google/genai/types.py index c80a9aad3..b3d023655 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -4223,6 +4223,92 @@ class UrlContextDict(TypedDict, total=False): UrlContextOrDict = Union[UrlContext, UrlContextDict] +class StreamableHttpTransport(_common.BaseModel): + """A transport that can stream HTTP requests and responses. + + Next ID: 6. This data type is not supported in Vertex AI. + """ + + headers: Optional[dict[str, str]] = Field( + default=None, + description="""Optional: Fields for authentication headers, timeouts, etc., if needed.""", + ) + sse_read_timeout: Optional[str] = Field( + default=None, description="""Timeout for SSE read operations.""" + ) + terminate_on_close: Optional[bool] = Field( + default=None, + description="""Whether to close the client session when the transport closes.""", + ) + timeout: Optional[str] = Field( + default=None, description="""HTTP timeout for regular operations.""" + ) + url: Optional[str] = Field( + default=None, + description="""The full URL for the MCPServer endpoint. Example: "https://api.example.com/mcp".""", + ) + + +class StreamableHttpTransportDict(TypedDict, total=False): + """A transport that can stream HTTP requests and responses. + + Next ID: 6. This data type is not supported in Vertex AI. + """ + + headers: Optional[dict[str, str]] + """Optional: Fields for authentication headers, timeouts, etc., if needed.""" + + sse_read_timeout: Optional[str] + """Timeout for SSE read operations.""" + + terminate_on_close: Optional[bool] + """Whether to close the client session when the transport closes.""" + + timeout: Optional[str] + """HTTP timeout for regular operations.""" + + url: Optional[str] + """The full URL for the MCPServer endpoint. Example: "https://api.example.com/mcp".""" + + +StreamableHttpTransportOrDict = Union[ + StreamableHttpTransport, StreamableHttpTransportDict +] + + +class McpServer(_common.BaseModel): + """A MCPServer is a server that can be called by the model to perform actions. + + It is a server that implements the MCP protocol. Next ID: 5. This data type is + not supported in Vertex AI. + """ + + name: Optional[str] = Field( + default=None, description="""The name of the MCPServer.""" + ) + streamable_http_transport: Optional[StreamableHttpTransport] = Field( + default=None, + description="""A transport that can stream HTTP requests and responses.""", + ) + + +class McpServerDict(TypedDict, total=False): + """A MCPServer is a server that can be called by the model to perform actions. + + It is a server that implements the MCP protocol. Next ID: 5. This data type is + not supported in Vertex AI. + """ + + name: Optional[str] + """The name of the MCPServer.""" + + streamable_http_transport: Optional[StreamableHttpTransportDict] + """A transport that can stream HTTP requests and responses.""" + + +McpServerOrDict = Union[McpServer, McpServerDict] + + class Tool(_common.BaseModel): """Tool details of a tool that the model may use to generate a response.""" @@ -4268,6 +4354,10 @@ class Tool(_common.BaseModel): default=None, description="""Optional. Tool to support URL context retrieval.""", ) + mcp_servers: Optional[list[McpServer]] = Field( + default=None, + description="""Optional. MCP Servers to connect to. This field is not supported in Vertex AI.""", + ) class ToolDict(TypedDict, total=False): @@ -4305,6 +4395,9 @@ class ToolDict(TypedDict, total=False): url_context: Optional[UrlContextDict] """Optional. Tool to support URL context retrieval.""" + mcp_servers: Optional[list[McpServerDict]] + """Optional. MCP Servers to connect to. This field is not supported in Vertex AI.""" + ToolOrDict = Union[Tool, ToolDict] if _is_mcp_imported: