Skip to content

Commit 67d087a

Browse files
committed
fix(agent_tool): only apply ReAct wrapper when output_schema is not set
- Fix: only apply the ReAct wrapper in agent_tool.py when output_schema is not set on the inner agent, preventing breaking of single-shot structured output mode - Add regression test test_run_async_with_input_and_output_schema_passes_raw_json documenting that raw JSON is passed when both input_schema and output_schema are set - Apply pre-commit formatting fixes (isort + pyink)
1 parent fd0073a commit 67d087a

3 files changed

Lines changed: 111 additions & 38 deletions

File tree

src/google/adk/tools/agent_tool.py

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -218,19 +218,29 @@ async def run_async(
218218
if input_schema:
219219
input_value = input_schema.model_validate(args)
220220
json_payload = input_value.model_dump_json(exclude_none=True)
221-
content = types.Content(
222-
role='user',
223-
parts=[
224-
types.Part.from_text(
225-
text=(
226-
'Process the following structured request. Use your'
227-
' available tools as needed to gather information or'
228-
' perform actions before producing the final'
229-
' response.\n\nRequest:\n' + json_payload
230-
)
231-
)
232-
],
233-
)
221+
output_schema = _get_output_schema(self.agent)
222+
if output_schema:
223+
# Single-shot structured output mode: pass raw JSON, no ReAct wrapper.
224+
content = types.Content(
225+
role='user',
226+
parts=[types.Part.from_text(text=json_payload)],
227+
)
228+
else:
229+
# Tool-calling mode: wrap with ReAct-style prompt.
230+
content = types.Content(
231+
role='user',
232+
parts=[
233+
types.Part.from_text(
234+
text=(
235+
'Process the following structured request. Use your'
236+
' available tools as needed to gather information or'
237+
' perform actions before producing the final'
238+
' response.\n\nRequest:\n'
239+
+ json_payload
240+
)
241+
)
242+
],
243+
)
234244
else:
235245
content = types.Content(
236246
role='user',

tests/unittests/models/test_litellm.py

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4999,9 +4999,7 @@ async def test_get_completion_inputs_tool_choice_none_without_tool_config():
49994999
"""tool_choice must be None when no tool_config is present."""
50005000
llm_request = LlmRequest(
50015001
contents=[
5002-
types.Content(
5003-
role="user", parts=[types.Part.from_text(text="Hello")]
5004-
)
5002+
types.Content(role="user", parts=[types.Part.from_text(text="Hello")])
50055003
],
50065004
)
50075005

@@ -5017,9 +5015,7 @@ async def test_get_completion_inputs_tool_choice_required_for_any_mode():
50175015
"""tool_choice must be 'required' when mode=ANY."""
50185016
llm_request = LlmRequest(
50195017
contents=[
5020-
types.Content(
5021-
role="user", parts=[types.Part.from_text(text="Hello")]
5022-
)
5018+
types.Content(role="user", parts=[types.Part.from_text(text="Hello")])
50235019
],
50245020
config=types.GenerateContentConfig(
50255021
tool_config=types.ToolConfig(
@@ -5042,9 +5038,7 @@ async def test_get_completion_inputs_tool_choice_none_for_none_mode():
50425038
"""tool_choice must be 'none' when mode=NONE."""
50435039
llm_request = LlmRequest(
50445040
contents=[
5045-
types.Content(
5046-
role="user", parts=[types.Part.from_text(text="Hello")]
5047-
)
5041+
types.Content(role="user", parts=[types.Part.from_text(text="Hello")])
50485042
],
50495043
config=types.GenerateContentConfig(
50505044
tool_config=types.ToolConfig(
@@ -5067,9 +5061,7 @@ async def test_get_completion_inputs_tool_choice_none_for_auto_mode():
50675061
"""tool_choice must be None (provider default) when mode=AUTO."""
50685062
llm_request = LlmRequest(
50695063
contents=[
5070-
types.Content(
5071-
role="user", parts=[types.Part.from_text(text="Hello")]
5072-
)
5064+
types.Content(role="user", parts=[types.Part.from_text(text="Hello")])
50735065
],
50745066
config=types.GenerateContentConfig(
50755067
tool_config=types.ToolConfig(
@@ -5159,9 +5151,7 @@ async def test_generate_content_async_omits_tool_choice_for_auto_mode(
51595151

51605152
llm_request = LlmRequest(
51615153
contents=[
5162-
types.Content(
5163-
role="user", parts=[types.Part.from_text(text="Hi")]
5164-
)
5154+
types.Content(role="user", parts=[types.Part.from_text(text="Hi")])
51655155
],
51665156
config=types.GenerateContentConfig(
51675157
tool_config=types.ToolConfig(
@@ -5190,9 +5180,7 @@ async def test_generate_content_async_omits_tool_choice_without_tool_config(
51905180

51915181
llm_request = LlmRequest(
51925182
contents=[
5193-
types.Content(
5194-
role="user", parts=[types.Part.from_text(text="Hi")]
5195-
)
5183+
types.Content(role="user", parts=[types.Part.from_text(text="Hi")])
51965184
],
51975185
)
51985186

tests/unittests/tools/test_agent_tool.py

Lines changed: 82 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1441,15 +1441,17 @@ def test_empty_sequential_agent_falls_back_to_request(self):
14411441
async def _run_agent_tool_and_capture_content(
14421442
args: dict,
14431443
input_schema=None,
1444+
output_schema=None,
14441445
) -> types.Content:
14451446
"""Drives AgentTool and captures the Content passed to the inner agent.
14461447
14471448
This uses a stub Runner (same pattern as test_agent_tool_inherits_parent_app_name)
14481449
to intercept the new_message without executing the actual agent pipeline.
14491450
"""
1451+
from unittest.mock import patch
1452+
14501453
from google.adk.agents.llm_agent import LlmAgent
14511454
from google.adk.plugins.plugin_manager import PluginManager
1452-
from unittest.mock import patch
14531455
import google.adk.runners as _runners_module
14541456

14551457
if input_schema is not None:
@@ -1458,6 +1460,7 @@ async def _run_agent_tool_and_capture_content(
14581460
description='captures input',
14591461
model=testing_utils.MockModel.create(responses=['done']),
14601462
input_schema=input_schema,
1463+
output_schema=output_schema,
14611464
)
14621465
else:
14631466
inner = Agent(name='inner_agent', model='test-model')
@@ -1470,16 +1473,33 @@ async def _empty_async_generator():
14701473

14711474
class _StubRunner:
14721475

1473-
def __init__(self, *, app_name, agent, artifact_service,
1474-
session_service, memory_service, credential_service, plugins):
1476+
def __init__(
1477+
self,
1478+
*,
1479+
app_name,
1480+
agent,
1481+
artifact_service,
1482+
session_service,
1483+
memory_service,
1484+
credential_service,
1485+
plugins,
1486+
):
14751487
del artifact_service, memory_service, credential_service
14761488
self.agent = agent
14771489
self.session_service = session_service
14781490
self.plugin_manager = PluginManager(plugins=plugins)
14791491
self.app_name = app_name
14801492

1481-
def run_async(self, *, user_id, session_id, invocation_id=None,
1482-
new_message=None, state_delta=None, run_config=None):
1493+
def run_async(
1494+
self,
1495+
*,
1496+
user_id,
1497+
session_id,
1498+
invocation_id=None,
1499+
new_message=None,
1500+
state_delta=None,
1501+
run_config=None,
1502+
):
14831503
new_message_holder.append(new_message)
14841504
return _empty_async_generator()
14851505

@@ -1538,6 +1558,7 @@ class MyInput(BaseModel):
15381558
assert 'Request:\n' in text
15391559
json_part = text.split('Request:\n', 1)[1]
15401560
import json as _json
1561+
15411562
payload = _json.loads(json_part)
15421563
assert payload['custom_input'] == 'test_value'
15431564
# The full text must NOT be just the raw JSON blob
@@ -1559,6 +1580,60 @@ class MyInput(BaseModel):
15591580
assert content is not None
15601581
text = content.parts[0].text
15611582
# A bare JSON blob would start with '{'; the wrapped version must not
1562-
assert not text.startswith('{'), (
1563-
'Content text is raw JSON instead of a natural-language instruction'
1583+
assert not text.startswith(
1584+
'{'
1585+
), 'Content text is raw JSON instead of a natural-language instruction'
1586+
1587+
1588+
@mark.asyncio
1589+
async def test_run_async_with_input_and_output_schema_passes_raw_json():
1590+
"""With both input_schema AND output_schema, the raw JSON payload is passed
1591+
directly to the inner runner WITHOUT the ReAct wrapper prefix.
1592+
1593+
The wrapper ('Process the following structured request...') is only added
1594+
when input_schema is set and output_schema is NOT set (tool-calling mode).
1595+
When output_schema is also present the agent operates in single-shot
1596+
structured-output mode, so the runner receives the bare JSON string that the
1597+
inner agent can parse deterministically — adding the prose prefix would
1598+
corrupt the structured input.
1599+
"""
1600+
import json as _json
1601+
1602+
class MyInput(BaseModel):
1603+
query: str
1604+
limit: int
1605+
1606+
class MyOutput(BaseModel):
1607+
result: str
1608+
1609+
content = await _run_agent_tool_and_capture_content(
1610+
args={'query': 'hello', 'limit': 5},
1611+
input_schema=MyInput,
1612+
output_schema=MyOutput,
15641613
)
1614+
1615+
assert content is not None
1616+
assert len(content.parts) == 1
1617+
text = content.parts[0].text
1618+
1619+
# output_schema mode is single-shot; wrapper must not be applied
1620+
assert not text.startswith('Process'), (
1621+
'output_schema mode is single-shot; wrapper must not be applied,'
1622+
f' but text starts with: {text[:60]!r}'
1623+
)
1624+
1625+
# The payload must be valid JSON
1626+
try:
1627+
payload = _json.loads(text)
1628+
except _json.JSONDecodeError as exc:
1629+
raise AssertionError(
1630+
f'Content text is not valid JSON in output_schema mode: {text!r}'
1631+
) from exc
1632+
1633+
# The JSON must match the input args
1634+
assert (
1635+
payload['query'] == 'hello'
1636+
), f"Expected query='hello', got {payload.get('query')!r}"
1637+
assert (
1638+
payload['limit'] == 5
1639+
), f"Expected limit=5, got {payload.get('limit')!r}"

0 commit comments

Comments
 (0)