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]
|
[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
|
||||||
|
|||||||
Reference in New Issue
Block a user