fix(gemini): tolerate missing usage metadata on final response (#4724)

Co-authored-by: Charan Rathore <charan-rathore@users.noreply.github.com>
This commit is contained in:
Charan Rathore
2026-09-30 14:04:45 -04:00
committed by GitHub
co-authored by Charan Rathore
parent 44bd73a241
commit 4cbc6a78ce
2 changed files with 46 additions and 1 deletions
+1 -1
View File
@@ -677,7 +677,7 @@ class GeminiModel(Model):
"data": "TOOL_USE" if tool_used else (candidate.finish_reason if candidate else "STOP"),
}
)
if event:
if event and event.usage_metadata is not None:
yield self._format_chunk({"chunk_type": "metadata", "data": event.usage_metadata})
except genai.errors.ClientError as error:
@@ -1753,3 +1753,48 @@ class TestCountTokens:
gemini_client.aio.models.count_tokens.assert_not_called()
assert isinstance(result, int)
assert result >= 0
@pytest.mark.asyncio
@pytest.mark.parametrize("finish_reason,stop_reason", [("STOP", "end_turn"), ("SAFETY", "guardrail_intervened")])
async def test_stream_response_without_usage_metadata(
gemini_client, model, messages, agenerator, alist, finish_reason, stop_reason
):
response = genai.types.GenerateContentResponse.model_validate(
{
"candidates": [
{
"content": {"role": "model", "parts": [{"text": "Hello"}]},
"finishReason": finish_reason,
}
]
}
)
assert response.usage_metadata is None
gemini_client.aio.models.generate_content_stream.return_value = agenerator([response])
chunks = await alist(model.stream(messages))
assert {"contentBlockDelta": {"delta": {"text": "Hello"}}} in chunks
assert chunks[-1] == {"messageStop": {"stopReason": stop_reason}}
assert not any("metadata" in chunk for chunk in chunks)
@pytest.mark.asyncio
async def test_stream_response_empty_usage_metadata_is_preserved(gemini_client, model, messages, agenerator, alist):
response = genai.types.GenerateContentResponse.model_validate(
{
"candidates": [{"content": {"role": "model", "parts": [{"text": "Hello"}]}, "finishReason": "STOP"}],
"usageMetadata": {},
}
)
gemini_client.aio.models.generate_content_stream.return_value = agenerator([response])
chunks = await alist(model.stream(messages))
assert chunks[-1] == {
"metadata": {
"usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
"metrics": {"latencyMs": 0},
}
}