scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(responses): log spend for incomplete replayed streams

Browse source 
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Tom Ron 2026-10-02 12:33:44 +03:00
parent 276b569568
commit 1fd2f32853
2 changed files with 38 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/responses/streaming_iterator.py
View file

@ -902,7 +902,10 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
if getattr(evt, "type", None) in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
):
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
@ -916,7 +919,10 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
if getattr(evt, "type", None) in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
):
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt
@ -969,7 +975,10 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
if getattr(evt, "type", None) in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
):
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
@ -983,7 +992,10 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
if getattr(evt, "type", None) in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
):
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt

24
tests/test_litellm/responses/litellm_completion_transformation/test_incomplete_status_transformation.py
View file

@ -1,4 +1,6 @@
from unittest.mock import AsyncMock
import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -8,7 +10,10 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.responses.streaming_iterator import _build_synthetic_response_events
from litellm.responses.streaming_iterator import (
CachedResponsesAPIStreamingIterator,
_build_synthetic_response_events,
)
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.utils import Choices, Message, ModelResponse, Usage
@ -91,3 +96,18 @@ def test_replayed_stream_terminal_event_follows_status(finish_reason, event_type
def test_omitted_temperature_defaults_to_zero():
assert _transform("stop", {}).temperature == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("finish_reason", ["length", "stop"])
async def test_replayed_stream_logs_success_exactly_once(finish_reason):
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
iterator = CachedResponsesAPIStreamingIterator(
response=_transform(finish_reason, {}),
logging_obj=logging_obj,
)
async for _ in iterator:
pass
await asyncio.sleep(0)
assert logging_obj.dispatch_success_handlers.await_count == 1