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
This commit is contained in:
Hangfei Lin
2025-06-30 09:50:04 -07:00
committed by Copybara-Service
parent a58cc3d882
commit 9b75e24d8c
2 changed files with 35 additions and 0 deletions
+5
View File
@@ -451,6 +451,11 @@ class LlmAgent(BaseAgent):
[part.text if part.text else '' for part in event.content.parts] [part.text if part.text else '' for part in event.content.parts]
) )
if self.output_schema: if self.output_schema:
# If the result from the final chunk is just whitespace or empty,
# it means this is an empty final chunk of a stream.
# Do not attempt to parse it as JSON.
if not result.strip():
return
result = self.output_schema.model_validate_json(result).model_dump( result = self.output_schema.model_validate_json(result).model_dump(
exclude_none=True exclude_none=True
) )
@@ -208,3 +208,33 @@ class TestLlmAgentOutputSave:
"agent1", "agent1",
"agent2", "agent2",
) )
@pytest.mark.parametrize("empty_content", ["", " ", "\n"])
def test_maybe_save_output_to_state_handles_empty_final_chunk_with_schema(
self, empty_content
):
"""Tests that the agent correctly handles an empty final streaming chunk
when an output_schema is specified, preventing a crash.
"""
# ARRANGE: Create an agent that expects a JSON output matching a schema.
agent = LlmAgent(
name="test_agent", output_key="result", output_schema=MockOutputSchema
)
# ARRANGE: Create a final event with empty or whitespace-only content.
# This simulates the final, empty chunk from a model's streaming response.
event = create_test_event(
author="test_agent", content_text=empty_content, is_final=True
)
# ACT: Call the method. The test's primary goal is to ensure this
# does NOT raise a pydantic.ValidationError, which it would have before the fix.
try:
agent._LlmAgent__maybe_save_output_to_state(event)
except Exception as e:
pytest.fail(f"The method unexpectedly raised an exception: {e}")
# ASSERT: Because the method should return early, the state_delta
# should remain empty.
assert len(event.actions.state_delta) == 0