|
25 | 25 | from google.genai.types import Content |
26 | 26 | from google.genai.types import FunctionCall |
27 | 27 | from google.genai.types import FunctionResponse |
| 28 | +from google.genai.types import GenerateContentResponseUsageMetadata |
28 | 29 | from google.genai.types import Part |
29 | 30 | import pytest |
30 | 31 |
|
@@ -57,7 +58,9 @@ async def test_maybe_compact_events_success(self): |
57 | 58 | expected_prompt = self.compactor._DEFAULT_PROMPT_TEMPLATE.format( |
58 | 59 | conversation_history=expected_conversation_history |
59 | 60 | ) |
60 | | - mock_llm_response = Mock(content=Content(parts=[Part(text='Summary')])) |
| 61 | + mock_llm_response = Mock( |
| 62 | + content=Content(parts=[Part(text='Summary')]), usage_metadata=None |
| 63 | + ) |
61 | 64 |
|
62 | 65 | async def async_gen(): |
63 | 66 | yield mock_llm_response |
@@ -90,11 +93,39 @@ async def async_gen(): |
90 | 93 | self.assertEqual(llm_request.contents[0].parts[0].text, expected_prompt) |
91 | 94 | self.assertFalse(kwargs['stream']) |
92 | 95 |
|
| 96 | + async def test_maybe_compact_events_includes_usage_metadata(self): |
| 97 | + events = [ |
| 98 | + self._create_event(1.0, 'Hello', 'user'), |
| 99 | + self._create_event(2.0, 'Hi there!', 'model'), |
| 100 | + ] |
| 101 | + usage_metadata = GenerateContentResponseUsageMetadata( |
| 102 | + prompt_token_count=10, |
| 103 | + candidates_token_count=5, |
| 104 | + total_token_count=15, |
| 105 | + ) |
| 106 | + mock_llm_response = Mock( |
| 107 | + content=Content(parts=[Part(text='Summary')]), |
| 108 | + usage_metadata=usage_metadata, |
| 109 | + ) |
| 110 | + |
| 111 | + async def async_gen(): |
| 112 | + yield mock_llm_response |
| 113 | + |
| 114 | + self.mock_llm.generate_content_async.return_value = async_gen() |
| 115 | + |
| 116 | + compacted_event = await self.compactor.maybe_summarize_events(events=events) |
| 117 | + |
| 118 | + self.assertIsNotNone(compacted_event) |
| 119 | + self.assertIsNotNone(compacted_event.usage_metadata) |
| 120 | + self.assertEqual(compacted_event.usage_metadata.prompt_token_count, 10) |
| 121 | + self.assertEqual(compacted_event.usage_metadata.candidates_token_count, 5) |
| 122 | + self.assertEqual(compacted_event.usage_metadata.total_token_count, 15) |
| 123 | + |
93 | 124 | async def test_maybe_compact_events_empty_llm_response(self): |
94 | 125 | events = [ |
95 | 126 | self._create_event(1.0, 'Hello', 'user'), |
96 | 127 | ] |
97 | | - mock_llm_response = Mock(content=None) |
| 128 | + mock_llm_response = Mock(content=None, usage_metadata=None) |
98 | 129 |
|
99 | 130 | async def async_gen(): |
100 | 131 | yield mock_llm_response |
|
0 commit comments