mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
a58cc3d882
commit
9b75e24d8c
@@ -451,6 +451,11 @@ class LlmAgent(BaseAgent):
|
||||
[part.text if part.text else '' for part in event.content.parts]
|
||||
)
|
||||
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(
|
||||
exclude_none=True
|
||||
)
|
||||
|
||||
@@ -208,3 +208,33 @@ class TestLlmAgentOutputSave:
|
||||
"agent1",
|
||||
"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
|
||||
|
||||
Reference in New Issue
Block a user