scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(streaming): log the streamed part's usage when a stream is closed early

Browse source 
close(), aclose() and both context managers on CustomStreamWrapper now log one partial success built from the chunks already streamed, so an SDK app that closes a stream when its client hangs up keeps that usage and cost. Nothing extra logs when the stream already ended, failed, was never read, when streaming logging is disabled, or when the proxy already logged a failure for it.
This commit is contained in:
mateo-berri 2026-11-04 19:58:08 -08:00
parent 5c60824024
commit ba550f3970
3 changed files with 552 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2803,6 +2803,10 @@ class Logging(LiteLLMLoggingBaseClass):
if self.cost_breakdown is not None:
self.cost_breakdown = _with_failed_attempts_cost(self.cost_breakdown, own_cost, failed_attempts_cost)
@property
def stream_failure_claimed(self) -> bool:
return self._stream_failure_claimed
def claim_stream_failure_report(self) -> bool:
if self._stream_failure_claimed:
return False

100
litellm/litellm_core_utils/streaming_handler.py
View file

@ -341,6 +341,21 @@ class _VertexChunkLike(Protocol):
candidates: Sequence[_VertexCandidateLike]
class _DispatchesSuccessHandlers(Protocol):
async def dispatch_success_handlers(
self,
result: object = None,
start_time: object = None,
end_time: object = None,
cache_hit: object = None,
prefer_async_handlers: bool = False,
) -> None: ...
def _success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _DispatchesSuccessHandlers:
return logging_obj
class _ParsedChunkHiddenParams(LiteLLMBaseModel):
provider_specific_fields: Mapping[str, object] | None = None
upstream_provider: str | None = None
@ -429,6 +444,8 @@ class CustomStreamWrapper:
self.sent_last_chunk = False
self._openai_handler_receives_chat_chunks = False
self._failure_logging: asyncio.Task[None] | None = None
self._outcome_logged = False
self._iterated_provider_stream = False
self._closed = False
self._stream_created_time: float = time.time()
@ -613,6 +630,54 @@ class CustomStreamWrapper:
self._restore_consumer_correlation_context(guarded=True)
def close(self) -> None:
self._log_abandoned_stream()
self._release()
async def aclose(self) -> None:
with anyio.CancelScope(shield=True):
await self._alog_abandoned_stream()
await self._arelease()
def _claim_abandoned_stream_log(self) -> bool:
if (
self._outcome_logged
or not self._iterated_provider_stream
or litellm.disable_streaming_logging is True
or self.logging_obj.stream_failure_claimed
):
return False
self._outcome_logged = True
return True
def _log_abandoned_stream(self) -> None:
if self._claim_abandoned_stream_log():
executor.submit(self._log_streamed_part)
def _log_streamed_part(self) -> None:
partial: Final = self.partial_response()
if partial is not None:
self.logging_obj.success_handler(partial, None, None, self._is_cache_hit())
async def _alog_abandoned_stream(self) -> None:
if not self._claim_abandoned_stream_log():
return
try:
partial: Final = await asyncify(self.partial_response)()
if partial is not None:
await _success_dispatcher(self.logging_obj).dispatch_success_handlers(
partial,
cache_hit=self._is_cache_hit(),
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
except Exception as e: # noqa: BLE001 # best-effort usage logging; closing must not raise into the caller
verbose_logger.debug("CustomStreamWrapper.aclose: could not log the abandoned stream's usage: %s", e)
def _is_cache_hit(self) -> bool:
return self.custom_llm_provider == "cached_response"
def _release(self) -> None:
stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped
self.completion_stream = None
self._closed = True
@ -626,7 +691,7 @@ class CustomStreamWrapper:
verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e)
self._restore_consumer_correlation_context()
async def aclose(self):
async def _arelease(self) -> None:
# Restore the consumer's outer context only after the underlying
# provider stream's own close (and its diagnostic logging below, if
# closing fails) completes - not before - so those log lines still
@ -2313,9 +2378,8 @@ class CustomStreamWrapper:
def __next__(self) -> "ModelResponseStream":
if self._closed:
raise StopIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
self._iterated_provider_stream = True
cache_hit: Final = self._is_cache_hit()
self._check_max_streaming_duration()
try:
if self.completion_stream is None:
@ -2401,12 +2465,14 @@ class CustomStreamWrapper:
error: Final = self._stream_eof_error()
eof_traceback: Final = traceback.format_exc()
self._record_partial_usage_for_failure()
self._outcome_logged = True
threading.Thread(target=self.logging_obj.failure_handler, args=(error, eof_traceback)).start()
self.close()
self._release()
self._handle_stream_fallback_error(error)
if self.sent_last_chunk is True:
complete_streaming_response: Final = self._build_complete_streaming_response()
response = self._new_usage_chunk()
self._outcome_logged = True
if complete_streaming_response is not None:
self._propagate_usage_cost_to_hidden_params(
complete_streaming_response, self.custom_llm_provider, self.chunks
@ -2463,7 +2529,7 @@ class CustomStreamWrapper:
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
self.close()
self._release()
raise # Re-raise StopIteration
else:
self.sent_last_chunk = True
@ -2496,8 +2562,9 @@ class CustomStreamWrapper:
except Exception as e:
traceback_exception: Final = traceback.format_exc()
# LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
self._outcome_logged = True
threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
self.close()
self._release()
self._handle_stream_fallback_error(e)
def fetch_sync_stream(self):
@ -2519,9 +2586,8 @@ class CustomStreamWrapper:
async def __anext__(self) -> "ModelResponseStream":
if self._closed:
raise StopAsyncIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
self._iterated_provider_stream = True
cache_hit: Final = self._is_cache_hit()
try:
# Inside the try (not before it) so a raised litellm.Timeout flows
# through the same except Exception -> _handle_stream_fallback_error
@ -2608,7 +2674,7 @@ class CustomStreamWrapper:
return usage_free_chunk
except (StopAsyncIteration, StopIteration):
if self._should_fail_on_clean_eof():
await self.aclose()
await self._arelease()
self._log_stream_failure_and_raise(self._stream_eof_error())
return await self._finalize_completed_stream(cache_hit=cache_hit)
except httpx.TimeoutException as e: # if httpx read timeout error occues
@ -2616,15 +2682,15 @@ class CustomStreamWrapper:
## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT
traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}"
self._dispatch_stream_failure(e, traceback_exception)
await self.aclose()
await self._arelease()
self._handle_stream_fallback_error(e)
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
if self.received_finish_reason is None:
await self.aclose()
await self._arelease()
self._log_stream_failure_and_raise(e)
return await self._finalize_completed_stream(cache_hit=cache_hit)
except Exception as e:
await self.aclose()
await self._arelease()
self._log_stream_failure_and_raise(e)
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
@ -2668,6 +2734,7 @@ class CustomStreamWrapper:
self.sent_stream_usage = True
return response
self._outcome_logged = True
_deferred_cb: Final = getattr(
self.logging_obj,
"_on_deferred_stream_complete",
@ -2690,7 +2757,7 @@ class CustomStreamWrapper:
# when consumers use ``async for`` on sync-SDK streams. Legacy string
# callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
self.logging_obj.dispatch_success_handlers(
_success_dispatcher(self.logging_obj).dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
@ -2699,7 +2766,7 @@ class CustomStreamWrapper:
)
)
await self.aclose()
await self._arelease()
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
@ -2722,6 +2789,7 @@ class CustomStreamWrapper:
self._handle_stream_fallback_error(e)
def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None:
self._outcome_logged = True
self._record_partial_usage_for_failure()
self._failure_logging = asyncio.create_task(
self.logging_obj.dispatch_attempt_failure_handlers(e, traceback_exception)

465
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -13,14 +13,16 @@ from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from types import CodeType, FrameType, MappingProxyType, TracebackType
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, TypeAlias, TypedDict, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import pytest
import respx
from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
from openai.types.chat import ChatCompletionChunk
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly
import litellm
from litellm import verbose_logger
from litellm._logging import session_id_var, trace_id_var
@ -47,6 +49,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
from litellm.types.utils import (
ChoiceLogprobs,
Choices,
CompletionTokensDetailsWrapper,
Delta,
ModelResponse,
@ -9571,3 +9574,463 @@ def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details(
assert usage.prompt_tokens_details.cached_tokens == 80
assert usage.completion_tokens_details is not None
assert usage.completion_tokens_details.reasoning_tokens == 5
_CUT_STREAM_URL: Final = "https://api.openai.com/v1/chat/completions"
_CUT_STREAM_MODEL: Final = "gpt-5.6"
_CUT_STREAM_PROMPT: Final = "Write a story about a lighthouse keeper."
_CUT_STREAM_DELTAS: Final = ("Once", " upon", " a", " time", " a", " keeper", " lit", " the", " lamp", ".")
_CUT_STREAM_PROMPT_TOKENS: Final = 431
_CUT_STREAM_COMPLETION_TOKENS: Final = 57
_CloseWith: TypeAlias = Literal["aclose", "async_with", "close", "with"]
def _cut_stream_event(fields: Mapping[str, object]) -> str:
chunk: Final = {
"id": "chatcmpl-cut-stream",
"object": "chat.completion.chunk",
"created": 1700000000,
"model": _CUT_STREAM_MODEL,
**fields,
}
return f"data: {json.dumps(chunk)}\n\n"
def _cut_stream_content(position: int, text: str) -> str:
delta: Final = {"role": "assistant", "content": text} if position == 0 else {"content": text}
return _cut_stream_event({"choices": [{"index": 0, "delta": delta, "finish_reason": None}]})
def _cut_stream_content_sse(deltas: Sequence[str]) -> str:
return "".join(_cut_stream_content(position, text) for position, text in enumerate(deltas))
def _cut_stream_sse(deltas: Sequence[str]) -> str:
return (
_cut_stream_content_sse(deltas)
+ _cut_stream_event({"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
+ _cut_stream_event(
{
"choices": [],
"usage": {
"prompt_tokens": _CUT_STREAM_PROMPT_TOKENS,
"completion_tokens": _CUT_STREAM_COMPLETION_TOKENS,
"total_tokens": _CUT_STREAM_PROMPT_TOKENS + _CUT_STREAM_COMPLETION_TOKENS,
},
}
)
+ "data: [DONE]\n\n"
)
_CUT_STREAM_SSE: Final = _cut_stream_sse(_CUT_STREAM_DELTAS)
class _ConnectionResetMidStream(httpx.SyncByteStream, httpx.AsyncByteStream):
def __iter__(self) -> Iterator[bytes]:
yield _cut_stream_content_sse(_CUT_STREAM_DELTAS).encode()
raise httpx.ReadError("connection reset by peer")
async def __aiter__(self) -> AsyncIterator[bytes]:
yield _cut_stream_content_sse(_CUT_STREAM_DELTAS).encode()
raise httpx.ReadError("connection reset by peer")
class _LoggedStreamOutcome(TypedDict):
model: ReadOnly[str]
response_cost: ReadOnly[float]
prompt_tokens: ReadOnly[int]
completion_tokens: ReadOnly[int]
@dataclass(frozen=True, slots=True)
class _LoggedStreamSuccess:
outcome: _LoggedStreamOutcome
text: str
def _logged_text(response_obj: object) -> str:
choices: Final = response_obj.choices if isinstance(response_obj, ModelResponse) else []
first: Final = choices[0] if choices else None
return (first.message.content or "") if isinstance(first, Choices) else ""
class _StreamOutcomeRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.successes: tuple[_LoggedStreamSuccess, ...] = ()
self.failures: tuple[str, ...] = ()
self._recording: Final = threading.Lock()
self.logged: Final = threading.Event()
def _record_success(self, kwargs: Mapping[str, object], response_obj: object) -> None:
success: Final = _LoggedStreamSuccess(
outcome=TypeAdapter(_LoggedStreamOutcome).validate_python(kwargs["standard_logging_object"]),
text=_logged_text(response_obj),
)
with self._recording:
self.successes = (*self.successes, success)
self.logged.set()
def _record_failure(self, mode: str) -> None:
with self._recording:
self.failures = (*self.failures, mode)
self.logged.set()
def log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._record_success(kwargs, response_obj)
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._record_success(kwargs, response_obj)
def log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._record_failure("sync")
async def async_log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._record_failure("async")
async def settled(self) -> None:
assert await asyncio.to_thread(self.logged.wait, 5), "no success or failure was logged"
await asyncio.sleep(0.5)
@pytest.fixture
def stream_outcomes(monkeypatch: pytest.MonkeyPatch) -> _StreamOutcomeRecorder:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
recorder: Final = _StreamOutcomeRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
return recorder
def _serve_openai_stream(respx_mock: respx.MockRouter, stream: str | _ConnectionResetMidStream) -> None:
headers: Final = {"content-type": "text/event-stream"}
response: Final = (
httpx.Response(200, text=stream, headers=headers)
if isinstance(stream, str)
else httpx.Response(200, stream=stream, headers=headers)
)
respx_mock.post(_CUT_STREAM_URL).mock(return_value=response)
def _open_sync_stream() -> CustomStreamWrapper:
stream: Final = litellm.completion(
model=_CUT_STREAM_MODEL,
messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}],
stream=True,
stream_options={"include_usage": True},
api_key="sk-unit-test",
)
assert isinstance(stream, CustomStreamWrapper)
return stream
async def _open_async_stream() -> CustomStreamWrapper:
stream: Final = await litellm.acompletion(
model=_CUT_STREAM_MODEL,
messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}],
stream=True,
stream_options={"include_usage": True},
api_key="sk-unit-test",
)
assert isinstance(stream, CustomStreamWrapper)
return stream
def _text_of(chunks: Sequence[ModelResponseStream]) -> str:
return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
async def _read_async(stream: CustomStreamWrapper, chunks: int) -> str:
return _text_of([await stream.__anext__() for _ in range(chunks)])
async def _read_two_chunks_then_close(close_with: _CloseWith) -> str:
if close_with == "close":
sync_stream: Final = _open_sync_stream()
sync_seen: Final = _text_of(tuple(itertools.islice(sync_stream, 2)))
sync_stream.close()
return sync_seen
if close_with == "with":
with _open_sync_stream() as managed_sync_stream:
return _text_of(tuple(itertools.islice(managed_sync_stream, 2)))
if close_with == "async_with":
async with await _open_async_stream() as managed_stream:
return await _read_async(managed_stream, 2)
stream: Final = await _open_async_stream()
seen: Final = await _read_async(stream, 2)
await stream.aclose()
return seen
@pytest.mark.asyncio
@pytest.mark.parametrize("close_with", ["aclose", "async_with", "close", "with"])
async def test_stream_closed_by_its_consumer_before_the_end_logs_the_streamed_part_once(
close_with: _CloseWith, stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
seen: Final = await _read_two_chunks_then_close(close_with)
await stream_outcomes.settled()
assert seen == "".join(_CUT_STREAM_DELTAS[:2])
assert stream_outcomes.failures == ()
assert len(stream_outcomes.successes) == 1
logged: Final = stream_outcomes.successes[0]
assert logged.text == seen
assert logged.outcome["prompt_tokens"] == litellm.token_counter(
model=_CUT_STREAM_MODEL, messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}]
)
assert logged.outcome["completion_tokens"] == litellm.token_counter(
model=_CUT_STREAM_MODEL, text=seen, count_response_tokens=True
)
assert logged.outcome["response_cost"] > 0
@pytest.mark.asyncio
async def test_async_stream_closed_right_after_its_usage_chunk_logs_the_provider_usage(
stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
stream: Final = await _open_async_stream()
received: Final = []
async for chunk in stream:
received.append(chunk)
if isinstance(getattr(chunk, "usage", None), Usage):
break
await stream.aclose()
await stream_outcomes.settled()
assert _text_of(received) == "".join(_CUT_STREAM_DELTAS)
assert len(stream_outcomes.successes) == 1
logged: Final = stream_outcomes.successes[0]
assert logged.text == "".join(_CUT_STREAM_DELTAS)
assert (logged.outcome["prompt_tokens"], logged.outcome["completion_tokens"]) == (
_CUT_STREAM_PROMPT_TOKENS,
_CUT_STREAM_COMPLETION_TOKENS,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_stream_read_to_its_end_then_closed_logs_the_provider_usage_once(
sync_mode: bool, stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
if sync_mode:
sync_stream: Final = _open_sync_stream()
sync_seen: Final = _text_of(tuple(sync_stream))
sync_stream.close()
else:
stream: Final = await _open_async_stream()
seen: Final = _text_of([chunk async for chunk in stream])
await stream.aclose()
await stream_outcomes.settled()
assert (sync_seen if sync_mode else seen) == "".join(_CUT_STREAM_DELTAS)
assert len(stream_outcomes.successes) == 1
logged: Final = stream_outcomes.successes[0]
assert (logged.outcome["prompt_tokens"], logged.outcome["completion_tokens"]) == (
_CUT_STREAM_PROMPT_TOKENS,
_CUT_STREAM_COMPLETION_TOKENS,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.parametrize(
"broken_stream",
[
pytest.param(_ConnectionResetMidStream(), id="connection_reset"),
pytest.param(_cut_stream_content_sse(_CUT_STREAM_DELTAS), id="body_ended_before_finish_reason"),
],
)
async def test_stream_that_broke_mid_flight_logs_no_success_when_closed(
sync_mode: bool,
broken_stream: str | _ConnectionResetMidStream,
stream_outcomes: _StreamOutcomeRecorder,
respx_mock: respx.MockRouter,
) -> None:
_serve_openai_stream(respx_mock, broken_stream)
if sync_mode:
sync_stream: Final = _open_sync_stream()
with pytest.raises(MidStreamFallbackError):
tuple(sync_stream)
sync_stream.close()
else:
stream: Final = await _open_async_stream()
with pytest.raises(MidStreamFallbackError):
[chunk async for chunk in stream]
await stream.aclose()
await stream_outcomes.settled()
assert stream_outcomes.successes == ()
assert len(stream_outcomes.failures) >= 1
@pytest.mark.asyncio
async def test_stream_whose_failure_was_already_reported_logs_no_success_when_closed(
stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
stream: Final = await _open_async_stream()
await _read_async(stream, 2)
assert stream.logging_obj.claim_stream_failure_report() is True
await stream.aclose()
await asyncio.sleep(0.5)
assert stream_outcomes.successes == ()
@pytest.mark.asyncio
async def test_router_stream_closed_by_its_consumer_before_the_end_logs_the_streamed_part_once(
stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
router: Final = litellm.Router(
model_list=[
{
"model_name": "story-writer",
"litellm_params": {"model": f"openai/{_CUT_STREAM_MODEL}", "api_key": "sk-unit-test"},
}
]
)
stream: Final = await router.acompletion(
model="story-writer",
messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}],
stream=True,
stream_options={"include_usage": True},
)
assert isinstance(stream, CustomStreamWrapper)
seen: Final = await _read_async(stream, 2)
await stream.aclose()
await stream_outcomes.settled()
assert seen == "".join(_CUT_STREAM_DELTAS[:2])
assert len(stream_outcomes.successes) == 1
logged: Final = stream_outcomes.successes[0]
assert logged.text == seen
assert logged.outcome["response_cost"] > 0
@pytest.mark.asyncio
async def test_router_stream_closed_after_a_mid_stream_fallback_logs_only_the_fallback_part(
stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter
) -> None:
backup_api_base: Final = "https://backup.openai.test/v1"
backup_deltas: Final = (" Then", " the", " storm", " came", ".")
_serve_openai_stream(respx_mock, _ConnectionResetMidStream())
respx_mock.post(f"{backup_api_base}/chat/completions").mock(
return_value=httpx.Response(
200, text=_cut_stream_sse(backup_deltas), headers={"content-type": "text/event-stream"}
)
)
router: Final = litellm.Router(
model_list=[
{
"model_name": "story-writer",
"litellm_params": {"model": f"openai/{_CUT_STREAM_MODEL}", "api_key": "sk-unit-test"},
},
{
"model_name": "story-writer-backup",
"litellm_params": {
"model": f"openai/{_CUT_STREAM_MODEL}",
"api_key": "sk-unit-test",
"api_base": backup_api_base,
},
"model_info": {"supports_assistant_prefill": True},
},
],
fallbacks=[{"story-writer": ["story-writer-backup"]}],
enable_mid_stream_fallback_continuation=True,
)
stream: Final = await router.acompletion(
model="story-writer",
messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}],
stream=True,
stream_options={"include_usage": True},
)
assert isinstance(stream, CustomStreamWrapper)
primary_part: Final = await _read_async(stream, len(_CUT_STREAM_DELTAS))
fallback_part: Final = await _read_async(stream, 2)
await stream.aclose()
await stream_outcomes.settled()
assert primary_part == "".join(_CUT_STREAM_DELTAS)
assert fallback_part == "".join(backup_deltas[:2])
assert [logged.text for logged in stream_outcomes.successes] == [fallback_part]
@pytest.mark.asyncio
async def test_stream_closed_early_logs_nothing_while_streaming_logging_is_disabled(
stream_outcomes: _StreamOutcomeRecorder, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_streaming_logging", True)
_serve_openai_stream(respx_mock, _CUT_STREAM_SSE)
stream: Final = await _open_async_stream()
await _read_async(stream, 2)
await stream.aclose()
await asyncio.sleep(0.5)
assert stream_outcomes.successes == ()
class _SuccessLoggingDown(Logging):
async def dispatch_success_handlers(self, *args: object, **kwargs: object) -> None:
raise RuntimeError("success logging is down")
class _ProviderChunks:
def __init__(self, chunks: Sequence[ModelResponseStream]) -> None:
self._chunks: Final = iter(chunks)
self.closed = False
def __aiter__(self) -> "_ProviderChunks":
return self
async def __anext__(self) -> ModelResponseStream:
chunk: Final = next(self._chunks, None)
if chunk is None:
raise StopAsyncIteration
return chunk
async def aclose(self) -> None:
self.closed = True
@pytest.mark.asyncio
async def test_async_stream_closed_early_still_closes_the_provider_stream_when_logging_its_part_fails() -> None:
provider_stream: Final = _ProviderChunks(tuple(_stream_chunk({"content": text}) for text in _CUT_STREAM_DELTAS))
stream: Final = CustomStreamWrapper(
completion_stream=provider_stream,
model=_CUT_STREAM_MODEL,
logging_obj=_SuccessLoggingDown(
model=_CUT_STREAM_MODEL,
messages=[{"role": "user", "content": _CUT_STREAM_PROMPT}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="cut-stream-logging-down",
function_id="cut-stream-logging-down",
),
)
seen: Final = await _read_async(stream, 2)
await stream.aclose()
assert seen == "".join(_CUT_STREAM_DELTAS[:2])
assert provider_stream.closed