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]
)
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
)