mirror of
https://github.com/strands-agents/harness-sdk.git
synced 2026-10-02 02:44:48 +08:00
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:
co-authored by
Charan Rathore
parent
44bd73a241
commit
4cbc6a78ce
@@ -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},
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user