diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index 72f82bbdf41..b9870544530 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -890,9 +890,14 @@ def _prepare_config( kwargs["response_schema"] = response_schema elif (schema := self._extract_response_schema(response_format)) is not None: kwargs["response_schema"] = schema - if tools := self._prepare_tools(options): + tools = self._prepare_tools(options) + if tools: kwargs["tools"] = tools - if tool_config := self._prepare_tool_config(options.get("tool_choice")): + tool_config = self._prepare_tool_config(options.get("tool_choice")) + if tools and not self._vertexai and self._requires_server_side_tool_invocations(tools): + tool_config = tool_config or types.ToolConfig() + tool_config.include_server_side_tool_invocations = True + if tool_config: kwargs["tool_config"] = tool_config if thinking_config := options.get("thinking_config"): thinking_config_kwargs = {k: v for k, v in thinking_config.items() if v is not None} @@ -975,7 +980,7 @@ def _prepare_tools(self, options: Mapping[str, Any]) -> list[types.Tool] | None: types.FunctionDeclaration( name=tool.name, description=tool.description or "", - parameters=tool.parameters(), # type: ignore[arg-type] + parameters_json_schema=tool.parameters(), ) for tool in tools_option if isinstance(tool, FunctionTool) @@ -988,6 +993,20 @@ def _prepare_tools(self, options: Mapping[str, Any]) -> list[types.Tool] | None: return result or None + @staticmethod + def _requires_server_side_tool_invocations(tools: Sequence[types.Tool]) -> bool: + """Return whether Gemini Developer API requires server-side tool invocation reporting.""" + has_function_declarations = any(tool.function_declarations for tool in tools) + return has_function_declarations and any(RawGeminiChatClient._has_server_side_tool(tool) for tool in tools) + + @staticmethod + def _has_server_side_tool(tool: types.Tool) -> bool: + """Return whether the Gemini tool declares a native server-side tool.""" + return any( + field_name != "function_declarations" and getattr(tool, field_name, None) is not None + for field_name in type(tool).model_fields + ) + def _prepare_tool_config(self, tool_choice: Any) -> types.ToolConfig | None: """Build a Gemini ``ToolConfig`` from the framework ``tool_choice`` value. diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index f2af3cfd7ae..dceb7b72aac 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -14,6 +14,7 @@ from agent_framework import Agent, Content, FunctionTool, Message from google.genai import types from pydantic import BaseModel +from typing_extensions import NotRequired, TypedDict from agent_framework_gemini import GeminiChatClient, GeminiChatOptions, RawGeminiChatClient, ThinkingConfig @@ -40,6 +41,12 @@ def _has_gemini_integration_credentials() -> bool: _TEST_MODEL = os.getenv("GOOGLE_MODEL") or os.getenv("GEMINI_MODEL", "gemini-2.5-flash-lite") + +class _ToolListItem(TypedDict): + title: str + description: NotRequired[str] + + # stub helpers @@ -1733,6 +1740,31 @@ def get_weather(city: str) -> str: assert function_declaration.name == "get_weather" +async def test_function_tool_json_schema_forwarded_to_parameters_json_schema() -> None: + """Forwards full JSON Schema from FunctionTool instead of coercing it into Gemini's Schema subset.""" + + def add_items(items: list[_ToolListItem]) -> str: + """Add items.""" + return "ok" + + tool = FunctionTool(name="add_items", func=add_items) + schema = tool.parameters() + assert "$defs" in schema + + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="Done")])) + + await client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Add items")])], + options={"tools": [tool]}, + ) + + config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"] + function_declaration = _first_function_declaration(config) + assert function_declaration.parameters is None + assert function_declaration.parameters_json_schema == schema + + async def test_callable_tool_resolved_via_validate_options() -> None: """Raw callables passed as tools must be normalized by _validate_options into FunctionTools and reach the Gemini config as function declarations. @@ -1756,6 +1788,63 @@ def get_weather(city: str) -> str: assert function_declaration.name == "get_weather" +async def test_developer_api_mixed_native_and_function_tools_enable_server_side_invocations() -> None: + """Enables Gemini Developer API server-side tool invocations when built-ins and function tools are mixed.""" + tool = _make_dummy_tool() + search_tool = GeminiChatClient.get_web_search_tool() + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="Done")])) + + await client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Search and call a tool")])], + options={"tools": [search_tool, tool]}, + ) + + config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"] + assert config.tool_config is not None + assert config.tool_config.include_server_side_tool_invocations is True + + +async def test_developer_api_mixed_native_and_function_tools_preserve_tool_choice() -> None: + """Preserves function calling mode while enabling Developer API server-side built-in tool invocations.""" + tool = _make_dummy_tool() + search_tool = GeminiChatClient.get_web_search_tool() + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="Done")])) + + await client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Search and call a tool")])], + options={"tools": [search_tool, tool], "tool_choice": "auto"}, + ) + + config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"] + function_calling_config = _function_calling_config(config) + assert function_calling_config.mode == "AUTO" + assert config.tool_config is not None + assert config.tool_config.include_server_side_tool_invocations is True + + +async def test_vertex_ai_mixed_native_and_function_tools_do_not_enable_server_side_invocations() -> None: + """Does not set the Developer API-only server-side invocation flag for Vertex AI.""" + tool = _make_dummy_tool() + search_tool = GeminiChatClient.get_web_search_tool() + mock = MagicMock() + mock._api_client.vertexai = True + client = GeminiChatClient(client=mock, model="gemini-2.5-flash") + mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="Done")])) + + await client.get_response( + messages=[Message(role="user", contents=[Content.from_text("Search and call a tool")])], + options={"tools": [search_tool, tool], "tool_choice": "auto"}, + ) + + config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"] + function_calling_config = _function_calling_config(config) + assert function_calling_config.mode == "AUTO" + assert config.tool_config is not None + assert config.tool_config.include_server_side_tool_invocations is None + + # _coerce_to_dict