scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

Merge 25ee7b2aeb into 62341e96ae

Browse source
This commit is contained in:
MUKKTINAADH 2026-08-27 20:56:08 +03:00 • committed by GitHub
commit 5efbf2cd02
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 13 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/interactions/litellm_responses_transformation/streaming_iterator.py
View file

@ -135,7 +135,7 @@ class LiteLLMResponsesInteractionsStreamingIterator:
index=0,
)
def _build_completion_event(self, response_id: str) -> InteractionsAPIStreamingResponse:
def _build_completion_event(self, response_id: str, usage: dict | None = None) -> InteractionsAPIStreamingResponse:
if self._use_legacy:
return InteractionsAPIStreamingResponse(
event_type="interaction.complete",
@ -144,6 +144,7 @@ class LiteLLMResponsesInteractionsStreamingIterator:
status="completed",
model=self.model,
outputs=[{"type": "text", "text": self.collected_text}],
usage=usage,
)
return InteractionsAPIStreamingResponse(
event_type="interaction.completed",
@ -157,6 +158,7 @@ class LiteLLMResponsesInteractionsStreamingIterator:
"content": [{"type": "text", "text": self.collected_text}],
}
],
usage=usage,
)
# ------------------------------------------------------------------
@ -218,11 +220,14 @@ class LiteLLMResponsesInteractionsStreamingIterator:
self.finished = True
response: Final = responses_chunk.response
response_id = self._interaction_id or getattr(response, "id", None) or f"interaction_{id(self)}"
# Extract usage from the response if available and convert to dict
usage_obj = getattr(response, "usage", None)
usage = usage_obj.model_dump() if usage_obj is not None else None
terminal: Final[list[InteractionsAPIStreamingResponse]] = []
if self.sent_content_start:
terminal.append(self._build_content_stop_event(response_id))
terminal.append(self._build_completion_event(response_id))
terminal.append(self._build_completion_event(response_id, usage))
self._sent_completion_event = True
return terminal

7
tests/test_litellm/interactions/test_gemini_interactions_transformation.py
View file

@ -444,7 +444,12 @@ class TestStreamingIterator:
text_event = self._make_text_delta("hi")
completed = MagicMock(spec=ResponseCompletedEvent)
completed.response = MagicMock(id="resp_999")
# Create a proper mock response with usage attribute that has model_dump method
mock_response = MagicMock(id="resp_999")
mock_usage = MagicMock()
mock_usage.model_dump.return_value = {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}
mock_response.usage = mock_usage
completed.response = mock_response
sync_iter = MagicMock()
sync_iter.__iter__ = lambda self: self