From 6a94af24bf3367c05a5d405b7e7b79810a1fac4e Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Mon, 3 Nov 2025 17:11:29 -0800 Subject: [PATCH] chore: Disable SetModelResponseTool workaround for Vertex AI Gemini 2+ models Gemini models now [support Function calling being used together with structured output on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling#structured-output-bp). Co-authored-by: Xuan Yang PiperOrigin-RevId: 827709903 --- .../llm_flows/_output_schema_processor.py | 10 +++- src/google/adk/flows/llm_flows/basic.py | 6 ++- src/google/adk/utils/output_schema_utils.py | 38 ++++++++++++++ .../flows/llm_flows/test_basic_processor.py | 43 +++++++++++++++- .../llm_flows/test_output_schema_processor.py | 35 ++++++++++--- .../utils/test_output_schema_utils.py | 50 +++++++++++++++++++ 6 files changed, 170 insertions(+), 12 deletions(-) create mode 100644 src/google/adk/utils/output_schema_utils.py create mode 100644 tests/unittests/utils/test_output_schema_utils.py diff --git a/src/google/adk/flows/llm_flows/_output_schema_processor.py b/src/google/adk/flows/llm_flows/_output_schema_processor.py index 464e0a71..2298c044 100644 --- a/src/google/adk/flows/llm_flows/_output_schema_processor.py +++ b/src/google/adk/flows/llm_flows/_output_schema_processor.py @@ -25,6 +25,7 @@ from ...agents.invocation_context import InvocationContext from ...events.event import Event from ...models.llm_request import LlmRequest from ...tools.set_model_response_tool import SetModelResponseTool +from ...utils.output_schema_utils import can_use_output_schema_with_tools from ._base_llm_processor import BaseLlmRequestProcessor @@ -39,8 +40,13 @@ class _OutputSchemaRequestProcessor(BaseLlmRequestProcessor): agent = invocation_context.agent - # Check if we need the processor: output_schema + tools - if not agent.output_schema or not agent.tools: + # Check if we need the processor: output_schema + tools + cannot use output + # schema with tools + if ( + not agent.output_schema + or not agent.tools + or can_use_output_schema_with_tools(agent.model) + ): return # Add the set_model_response tool to handle structured output diff --git a/src/google/adk/flows/llm_flows/basic.py b/src/google/adk/flows/llm_flows/basic.py index 24cc7fd6..1468a7ca 100644 --- a/src/google/adk/flows/llm_flows/basic.py +++ b/src/google/adk/flows/llm_flows/basic.py @@ -25,6 +25,7 @@ from typing_extensions import override from ...agents.invocation_context import InvocationContext from ...events.event import Event from ...models.llm_request import LlmRequest +from ...utils.output_schema_utils import can_use_output_schema_with_tools from ._base_llm_processor import BaseLlmRequestProcessor @@ -52,8 +53,9 @@ class _BasicLlmRequestProcessor(BaseLlmRequestProcessor): # support output_schema and tools together. we have a workaround to support # both output_schema and tools at the same time. see # _output_schema_processor.py for details - if agent.output_schema and not agent.tools: - llm_request.set_output_schema(agent.output_schema) + if agent.output_schema: + if not agent.tools or can_use_output_schema_with_tools(agent.model): + llm_request.set_output_schema(agent.output_schema) llm_request.live_connect_config.response_modalities = ( invocation_context.run_config.response_modalities diff --git a/src/google/adk/utils/output_schema_utils.py b/src/google/adk/utils/output_schema_utils.py new file mode 100644 index 00000000..ae14686e --- /dev/null +++ b/src/google/adk/utils/output_schema_utils.py @@ -0,0 +1,38 @@ +# 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. + +"""Utilities for Output Schema. + +This module is for ADK internal use only. +Please do not rely on the implementation details. +""" + +from __future__ import annotations + +from typing import Union + +from ..models.base_llm import BaseLlm +from .model_name_utils import is_gemini_2_or_above +from .variant_utils import get_google_llm_variant +from .variant_utils import GoogleLLMVariant + + +def can_use_output_schema_with_tools(model: Union[str, BaseLlm]): + """Returns True if output schema with tools is supported.""" + model_string = model if isinstance(model, str) else model.model + + return ( + get_google_llm_variant() == GoogleLLMVariant.VERTEX_AI + and is_gemini_2_or_above(model_string) + ) diff --git a/tests/unittests/flows/llm_flows/test_basic_processor.py b/tests/unittests/flows/llm_flows/test_basic_processor.py index 770f3589..e0be7781 100644 --- a/tests/unittests/flows/llm_flows/test_basic_processor.py +++ b/tests/unittests/flows/llm_flows/test_basic_processor.py @@ -14,6 +14,8 @@ """Tests for basic LLM request processor.""" +from unittest import mock + from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.llm_agent import LlmAgent from google.adk.agents.run_config import RunConfig @@ -80,7 +82,7 @@ class TestBasicLlmRequestProcessor: assert llm_request.config.response_mime_type == 'application/json' @pytest.mark.asyncio - async def test_skips_output_schema_when_tools_present(self): + async def test_skips_output_schema_when_tools_present(self, mocker): """Test that processor skips output_schema when agent has tools.""" agent = LlmAgent( name='test_agent', @@ -93,6 +95,11 @@ class TestBasicLlmRequestProcessor: llm_request = LlmRequest() processor = _BasicLlmRequestProcessor() + can_use_output_schema_with_tools = mocker.patch( + 'google.adk.flows.llm_flows.basic.can_use_output_schema_with_tools', + mock.MagicMock(return_value=False), + ) + # Process the request events = [] async for event in processor.run_async(invocation_context, llm_request): @@ -102,6 +109,40 @@ class TestBasicLlmRequestProcessor: assert llm_request.config.response_schema is None assert llm_request.config.response_mime_type != 'application/json' + # Should have checked if output schema can be used with tools + can_use_output_schema_with_tools.assert_called_once_with(agent.model) + + @pytest.mark.asyncio + async def test_sets_output_schema_when_tools_present(self, mocker): + """Test that processor skips output_schema when agent has tools.""" + agent = LlmAgent( + name='test_agent', + model='gemini-2.5-flash', + output_schema=OutputSchema, + tools=[FunctionTool(func=dummy_tool)], # Has tools + ) + + invocation_context = await _create_invocation_context(agent) + llm_request = LlmRequest() + processor = _BasicLlmRequestProcessor() + + can_use_output_schema_with_tools = mocker.patch( + 'google.adk.flows.llm_flows.basic.can_use_output_schema_with_tools', + mock.MagicMock(return_value=True), + ) + + # Process the request + events = [] + async for event in processor.run_async(invocation_context, llm_request): + events.append(event) + + # Should have set response_schema since output schema can be used with tools + assert llm_request.config.response_schema == OutputSchema + assert llm_request.config.response_mime_type == 'application/json' + + # Should have checked if output schema can be used with tools + can_use_output_schema_with_tools.assert_called_once_with(agent.model) + @pytest.mark.asyncio async def test_no_output_schema_no_tools(self): """Test that processor works normally when agent has no output_schema or tools.""" diff --git a/tests/unittests/flows/llm_flows/test_output_schema_processor.py b/tests/unittests/flows/llm_flows/test_output_schema_processor.py index 4c43407c..f7ad8eb3 100644 --- a/tests/unittests/flows/llm_flows/test_output_schema_processor.py +++ b/tests/unittests/flows/llm_flows/test_output_schema_processor.py @@ -14,14 +14,13 @@ """Tests for output schema processor functionality.""" -import json +from unittest import mock from google.adk.agents.invocation_context import InvocationContext from google.adk.agents.llm_agent import LlmAgent from google.adk.agents.run_config import RunConfig from google.adk.flows.llm_flows.single_flow import SingleFlow from google.adk.models.llm_request import LlmRequest -from google.adk.models.llm_response import LlmResponse from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.function_tool import FunctionTool from pydantic import BaseModel @@ -145,7 +144,16 @@ async def test_basic_processor_sets_output_schema_without_tools(): @pytest.mark.asyncio -async def test_output_schema_request_processor(): +@pytest.mark.parametrize( + 'output_schema_with_tools_allowed', + [ + False, + True, + ], +) +async def test_output_schema_request_processor( + output_schema_with_tools_allowed, mocker +): """Test that output schema processor adds set_model_response tool.""" from google.adk.flows.llm_flows._output_schema_processor import _OutputSchemaRequestProcessor @@ -161,16 +169,29 @@ async def test_output_schema_request_processor(): llm_request = LlmRequest() processor = _OutputSchemaRequestProcessor() + can_use_output_schema_with_tools = mocker.patch( + 'google.adk.flows.llm_flows._output_schema_processor.can_use_output_schema_with_tools', + mock.MagicMock(return_value=output_schema_with_tools_allowed), + ) + # Process the request events = [] async for event in processor.run_async(invocation_context, llm_request): events.append(event) - # Should have added set_model_response tool - assert 'set_model_response' in llm_request.tools_dict + if not output_schema_with_tools_allowed: + # Should have added set_model_response tool if output schema with tools is + # allowed + assert 'set_model_response' in llm_request.tools_dict + # Should have added instruction about using set_model_response + assert 'set_model_response' in llm_request.config.system_instruction + else: + # Should skip modifying LlmRequest + assert not llm_request.tools_dict + assert not llm_request.config.system_instruction - # Should have added instruction about using set_model_response - assert 'set_model_response' in llm_request.config.system_instruction + # Should have checked if output schema can be used with tools + can_use_output_schema_with_tools.assert_called_once_with(agent.model) @pytest.mark.asyncio diff --git a/tests/unittests/utils/test_output_schema_utils.py b/tests/unittests/utils/test_output_schema_utils.py new file mode 100644 index 00000000..ca7f88d9 --- /dev/null +++ b/tests/unittests/utils/test_output_schema_utils.py @@ -0,0 +1,50 @@ +# 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 google.adk.models.anthropic_llm import Claude +from google.adk.models.google_llm import Gemini +from google.adk.utils.output_schema_utils import can_use_output_schema_with_tools +import pytest + + +@pytest.mark.parametrize( + "model, env_value, expected", + [ + ("gemini-2.5-pro", "1", True), + ("gemini-2.5-pro", "0", False), + ("gemini-2.5-pro", None, False), + (Gemini(model="gemini-2.5-pro"), "1", True), + (Gemini(model="gemini-2.5-pro"), "0", False), + (Gemini(model="gemini-2.5-pro"), None, False), + ("gemini-2.0-flash", "1", True), + ("gemini-2.0-flash", "0", False), + ("gemini-2.0-flash", None, False), + ("gemini-1.5-pro", "1", False), + ("gemini-1.5-pro", "0", False), + ("gemini-1.5-pro", None, False), + (Claude(model="claude-3.7-sonnet"), "1", False), + (Claude(model="claude-3.7-sonnet"), "0", False), + (Claude(model="claude-3.7-sonnet"), None, False), + ], +) +def test_can_use_output_schema_with_tools( + monkeypatch, model, env_value, expected +): + """Test can_use_output_schema_with_tools.""" + if env_value is not None: + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", env_value) + else: + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + assert can_use_output_schema_with_tools(model) == expected