mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-04-28 01:21:43 +00:00
fix(agent): extract usageMetadata from Gemini streaming finish chunks
Gemini streaming sessions reported 0 tokens because translate_stream_event did not extract usageMetadata from the finish event. Mirrors the non-streaming translate_gemini_response path. Supersedes #15264.
This commit is contained in:
parent
00c3d848d8
commit
3fdde2ccd4
2 changed files with 92 additions and 1 deletions
77
tests/agent/test_gemini_streaming_usage.py
Normal file
77
tests/agent/test_gemini_streaming_usage.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
"""Tests for Gemini streaming usageMetadata extraction.
|
||||
|
||||
The streaming path must extract token counts from usageMetadata on
|
||||
the finish chunk, matching the non-streaming translate_gemini_response
|
||||
behavior.
|
||||
"""
|
||||
import unittest
|
||||
|
||||
from agent.gemini_native_adapter import translate_stream_event
|
||||
|
||||
|
||||
class TestGeminiStreamingUsageMetadata(unittest.TestCase):
|
||||
"""Verify usageMetadata is attached to streaming finish chunks."""
|
||||
|
||||
def _finish_event(self, usage_meta=None):
|
||||
"""Create a Gemini streaming event with finishReason and optional usageMetadata."""
|
||||
event = {
|
||||
"candidates": [{
|
||||
"content": {"parts": [], "role": "model"},
|
||||
"finishReason": "STOP",
|
||||
}],
|
||||
}
|
||||
if usage_meta is not None:
|
||||
event["usageMetadata"] = usage_meta
|
||||
return event
|
||||
|
||||
def _get_usage(self, chunks):
|
||||
"""Extract usage from a list of chunks, if present."""
|
||||
for c in chunks:
|
||||
usage = getattr(c, "usage", None)
|
||||
if usage is not None:
|
||||
return usage
|
||||
return None
|
||||
|
||||
def test_usage_attached_on_finish(self):
|
||||
"""Finish chunk should carry usage from usageMetadata."""
|
||||
event = self._finish_event({
|
||||
"promptTokenCount": 100,
|
||||
"candidatesTokenCount": 50,
|
||||
"totalTokenCount": 150,
|
||||
"cachedContentTokenCount": 20,
|
||||
})
|
||||
chunks = translate_stream_event(event, model="gemini-2.5-pro", tool_call_indices={})
|
||||
usage = self._get_usage(chunks)
|
||||
self.assertIsNotNone(usage, "Finish chunk should have usage")
|
||||
self.assertEqual(usage.prompt_tokens, 100)
|
||||
self.assertEqual(usage.completion_tokens, 50)
|
||||
self.assertEqual(usage.total_tokens, 150)
|
||||
self.assertEqual(usage.prompt_tokens_details.cached_tokens, 20)
|
||||
|
||||
def test_no_usage_when_metadata_absent(self):
|
||||
"""When usageMetadata is missing, finish chunk should not have usage."""
|
||||
chunks = translate_stream_event(
|
||||
self._finish_event(), model="gemini-2.5-pro", tool_call_indices={}
|
||||
)
|
||||
self.assertIsNone(self._get_usage(chunks))
|
||||
|
||||
def test_partial_metadata_defaults_to_zero(self):
|
||||
"""Missing fields in usageMetadata should default to 0."""
|
||||
event = self._finish_event({
|
||||
"promptTokenCount": 200,
|
||||
"totalTokenCount": 200,
|
||||
# candidatesTokenCount and cachedContentTokenCount missing
|
||||
})
|
||||
chunks = translate_stream_event(event, model="gemini-2.5-pro", tool_call_indices={})
|
||||
usage = self._get_usage(chunks)
|
||||
self.assertIsNotNone(usage)
|
||||
self.assertEqual(usage.prompt_tokens, 200)
|
||||
self.assertEqual(usage.completion_tokens, 0)
|
||||
self.assertEqual(usage.prompt_tokens_details.cached_tokens, 0)
|
||||
|
||||
def test_empty_metadata_no_usage(self):
|
||||
"""Empty usageMetadata dict should not create usage."""
|
||||
chunks = translate_stream_event(
|
||||
self._finish_event({}), model="gemini-2.5-pro", tool_call_indices={}
|
||||
)
|
||||
self.assertIsNone(self._get_usage(chunks))
|
||||
Loading…
Add table
Add a link
Reference in a new issue