From dc5dbfa2e475f2461177e80f7fa28c96a2bc3592 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Thu, 6 Aug 2026 14:49:15 -0700 Subject: [PATCH] feat: honor model-declared capabilities when pairing an output schema with tools The basic and output-schema request processors now read `model.capabilities.output_schema_and_tools` instead of inferring support from the model name. A `BaseLlm` subclass that declares the capability is honored, which previously it was not: support was derived from the model id and backend variant regardless of what the model reported. Built-in models are unaffected. `Gemini` and `LiteLlm` already self-report, and any model that does not falls through to the deprecated name-based fallback on `BaseLlm`, which reproduces the previous answer and warns. `utils/output_schema_utils.can_use_output_schema_with_tools()` is marked deprecated. Its body is unchanged and it keeps working; it cannot honor capabilities declared by a subclass, so callers should read the model instead. Co-authored-by: Xuan Yang PiperOrigin-RevId: 960515327 --- .../llm_flows/_output_schema_processor.py | 3 +- src/google/adk/flows/llm_flows/basic.py | 5 +-- src/google/adk/utils/output_schema_utils.py | 6 +++ .../flows/llm_flows/test_basic_processor.py | 42 ++++++------------- .../llm_flows/test_output_schema_processor.py | 30 +++++-------- tests/unittests/testing_utils.py | 23 ++++++++++ 6 files changed, 54 insertions(+), 55 deletions(-) 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 47876c29..85bc5b9a 100644 --- a/src/google/adk/flows/llm_flows/_output_schema_processor.py +++ b/src/google/adk/flows/llm_flows/_output_schema_processor.py @@ -25,7 +25,6 @@ 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 from ._invocation_utils import as_llm_agent from ._invocation_utils import require_agent_name @@ -46,7 +45,7 @@ class _OutputSchemaRequestProcessor(BaseLlmRequestProcessor): if ( not agent.output_schema or not agent.tools - or can_use_output_schema_with_tools(agent.canonical_model) + or agent.canonical_model.capabilities.output_schema_and_tools or getattr(agent, 'mode', None) == 'task' ): return diff --git a/src/google/adk/flows/llm_flows/basic.py b/src/google/adk/flows/llm_flows/basic.py index 3c38bbc9..85797cc6 100644 --- a/src/google/adk/flows/llm_flows/basic.py +++ b/src/google/adk/flows/llm_flows/basic.py @@ -25,7 +25,6 @@ from ...agents.invocation_context import InvocationContext from ...events.event import Event from ...models.llm_request import LlmRequest from ...utils import model_name_utils -from ...utils.output_schema_utils import can_use_output_schema_with_tools from ._base_llm_processor import BaseLlmRequestProcessor from ._invocation_utils import as_llm_agent from ._invocation_utils import require_run_config @@ -71,7 +70,7 @@ def _build_basic_request( agent = as_llm_agent(invocation_context) run_config = require_run_config(invocation_context) model = agent.canonical_model - llm_request.model = model if isinstance(model, str) else model.model + llm_request.model = model.model # Preserved across the agent-config overwrite below, then merged back. run_config_http_options = llm_request.config.http_options @@ -105,7 +104,7 @@ def _build_basic_request( # the basic flow. Structured output for tasks is collected via the # finish_task tool schema instead. if getattr(agent, 'mode', None) != 'task' and agent.output_schema: - if not agent.tools or can_use_output_schema_with_tools(model): + if not agent.tools or model.capabilities.output_schema_and_tools: llm_request.set_output_schema(agent.output_schema) llm_request.live_connect_config.response_modalities = ( diff --git a/src/google/adk/utils/output_schema_utils.py b/src/google/adk/utils/output_schema_utils.py index 1a2a4d5c..0647501d 100644 --- a/src/google/adk/utils/output_schema_utils.py +++ b/src/google/adk/utils/output_schema_utils.py @@ -22,10 +22,16 @@ from __future__ import annotations from typing import Union +from typing_extensions import deprecated + from ..models._capabilities import gemini_output_schema_and_tools from ..models.base_llm import BaseLlm +@deprecated( + 'Use model.capabilities.output_schema_and_tools instead. This function' + ' does not honor capabilities declared by a BaseLlm subclass.' +) def can_use_output_schema_with_tools(model: Union[str, BaseLlm]) -> bool: """Returns True if output schema with tools is supported.""" # LiteLLM handles tools + response_format compatibility per-provider: diff --git a/tests/unittests/flows/llm_flows/test_basic_processor.py b/tests/unittests/flows/llm_flows/test_basic_processor.py index 35923d72..a2e122f7 100644 --- a/tests/unittests/flows/llm_flows/test_basic_processor.py +++ b/tests/unittests/flows/llm_flows/test_basic_processor.py @@ -14,8 +14,6 @@ """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 @@ -28,6 +26,8 @@ from pydantic import BaseModel from pydantic import Field import pytest +from ... import testing_utils + class OutputSchema(BaseModel): """Test schema for output.""" @@ -83,11 +83,13 @@ class TestBasicLlmRequestProcessor: assert llm_request.config.response_mime_type == 'application/json' @pytest.mark.asyncio - async def test_skips_output_schema_when_tools_present(self, mocker): - """Test that processor skips output_schema when agent has tools.""" + async def test_skips_output_schema_when_model_denies_it(self): + """Test that processor skips output_schema when the model cannot pair it.""" agent = LlmAgent( name='test_agent', - model='gemini-2.5-flash', + model=testing_utils.ModelWithCapabilities( + output_schema_and_tools=False + ), output_schema=OutputSchema, tools=[FunctionTool(func=dummy_tool)], # Has tools ) @@ -96,31 +98,21 @@ 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): events.append(event) - # Should NOT have set response_schema since agent has tools + # Should NOT have set response_schema since the model does not support it 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.canonical_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.""" + async def test_sets_output_schema_when_model_declares_it(self): + """Test that processor sets output_schema when the model declares support.""" agent = LlmAgent( name='test_agent', - model='gemini-2.5-flash', + model=testing_utils.ModelWithCapabilities(output_schema_and_tools=True), output_schema=OutputSchema, tools=[FunctionTool(func=dummy_tool)], # Has tools ) @@ -129,25 +121,15 @@ 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=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 + # Should have set response_schema since the model declares support 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.canonical_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 c22fd488..9ae5478b 100644 --- a/tests/unittests/flows/llm_flows/test_output_schema_processor.py +++ b/tests/unittests/flows/llm_flows/test_output_schema_processor.py @@ -14,8 +14,6 @@ """Tests for output schema processor functionality.""" -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 @@ -33,6 +31,8 @@ from pydantic import BaseModel from pydantic import Field import pytest +from ... import testing_utils + class PersonSchema(BaseModel): """Test schema for structured output.""" @@ -151,21 +151,21 @@ async def test_basic_processor_sets_output_schema_without_tools(): @pytest.mark.asyncio @pytest.mark.parametrize( - 'output_schema_with_tools_allowed', + 'output_schema_and_tools', [ False, True, ], ) -async def test_output_schema_request_processor( - output_schema_with_tools_allowed, mocker -): +async def test_output_schema_request_processor(output_schema_and_tools): """Test that output schema processor adds set_model_response tool.""" from google.adk.flows.llm_flows._output_schema_processor import _OutputSchemaRequestProcessor agent = LlmAgent( name='test_agent', - model='gemini-2.5-flash', + model=testing_utils.ModelWithCapabilities( + output_schema_and_tools=output_schema_and_tools + ), output_schema=PersonSchema, tools=[FunctionTool(func=dummy_tool)], ) @@ -175,19 +175,14 @@ 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) - if not output_schema_with_tools_allowed: - # Should have added set_model_response tool if output schema with tools is - # allowed + if not output_schema_and_tools: + # The model cannot pair an output schema with tools, so the prompt-based + # workaround is installed instead. 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 @@ -196,11 +191,6 @@ async def test_output_schema_request_processor( assert not llm_request.tools_dict assert not 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.canonical_model - ) - @pytest.mark.asyncio async def test_set_model_response_tool(): diff --git a/tests/unittests/testing_utils.py b/tests/unittests/testing_utils.py index adaa9acb..84e2bfa3 100644 --- a/tests/unittests/testing_utils.py +++ b/tests/unittests/testing_utils.py @@ -29,6 +29,7 @@ from google.adk.apps.app import App from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService from google.adk.events.event import Event from google.adk.memory.in_memory_memory_service import InMemoryMemoryService +from google.adk.models import LlmCapabilities from google.adk.models.base_llm import BaseLlm from google.adk.models.base_llm_connection import BaseLlmConnection from google.adk.models.llm_request import LlmRequest @@ -332,6 +333,28 @@ class InMemoryRunner: return collected_responses +class ModelWithCapabilities(BaseLlm): + """A model that self-reports fixed capabilities. + + For exercising flows that branch on ``BaseLlm.capabilities``, without + depending on which model ids happen to satisfy ADK's detection today. + """ + + model: str = 'mock' + output_schema_and_tools: bool = False + + @property + @override + def capabilities(self) -> LlmCapabilities: + return LlmCapabilities(output_schema_and_tools=self.output_schema_and_tools) + + @override + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + yield LlmResponse() + + class MockModel(BaseLlm): model: str = 'mock'