Skip to content

Commit 9b75e24

Browse files
hangfeicopybara-github
authored andcommitted
fix: Avoid pydantic.ValidationError when the model stream returns empty final chunk
Bug: When a model emits a stream of tokens, it sometimes emits a final chunk of whitespace or no content. The agent was trying to parse that content into JSON, causing a validation error. Fix: If a model is expected to return JSON and the last streamed token is empty/whitespace, the agent will no longer try to parse it, and exit gracefully. New unit tests confirm the scenario and the fix. PiperOrigin-RevId: 777609415
1 parent a58cc3d commit 9b75e24

File tree

2 files changed

+35
-0
lines changed

2 files changed

+35
-0
lines changed

src/google/adk/agents/llm_agent.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,11 @@ def __maybe_save_output_to_state(self, event: Event):
451451
[part.text if part.text else '' for part in event.content.parts]
452452
)
453453
if self.output_schema:
454+
# If the result from the final chunk is just whitespace or empty,
455+
# it means this is an empty final chunk of a stream.
456+
# Do not attempt to parse it as JSON.
457+
if not result.strip():
458+
return
454459
result = self.output_schema.model_validate_json(result).model_dump(
455460
exclude_none=True
456461
)

tests/unittests/agents/test_llm_agent_output_save.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,3 +208,33 @@ def test_maybe_save_output_to_state_logging(self, mock_logger):
208208
"agent1",
209209
"agent2",
210210
)
211+
212+
@pytest.mark.parametrize("empty_content", ["", " ", "\n"])
213+
def test_maybe_save_output_to_state_handles_empty_final_chunk_with_schema(
214+
self, empty_content
215+
):
216+
"""Tests that the agent correctly handles an empty final streaming chunk
217+
218+
when an output_schema is specified, preventing a crash.
219+
"""
220+
# ARRANGE: Create an agent that expects a JSON output matching a schema.
221+
agent = LlmAgent(
222+
name="test_agent", output_key="result", output_schema=MockOutputSchema
223+
)
224+
225+
# ARRANGE: Create a final event with empty or whitespace-only content.
226+
# This simulates the final, empty chunk from a model's streaming response.
227+
event = create_test_event(
228+
author="test_agent", content_text=empty_content, is_final=True
229+
)
230+
231+
# ACT: Call the method. The test's primary goal is to ensure this
232+
# does NOT raise a pydantic.ValidationError, which it would have before the fix.
233+
try:
234+
agent._LlmAgent__maybe_save_output_to_state(event)
235+
except Exception as e:
236+
pytest.fail(f"The method unexpectedly raised an exception: {e}")
237+
238+
# ASSERT: Because the method should return early, the state_delta
239+
# should remain empty.
240+
assert len(event.actions.state_delta) == 0

0 commit comments

Comments
 (0)