mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(streaming): do not add close() to CustomStreamWrapper
`tests/unit/test_router_silent_experiment.py` fails on this branch: the shadow
and fallback streaming tests see zero success callbacks where they expect one.
The cause is a name collision. `Router.stream_with_fallbacks` ends with
finally:
if hasattr(model_response, "close"):
model_response.close()
`CustomStreamWrapper` had no `close`, so that branch never ran. Adding one
turned it on, so the fallback path now invokes the terminal-log flush from its
`finally`, and the shadow stream reports its completion through a path the
success callbacks were never accounting for.
The flush API is renamed to `close_stream()`. Same behaviour, same context
manager, but `hasattr(stream, "close")` stays False exactly as before, so no
existing call path changes. `test_router_close_duck_typing_contract` pins that
so it cannot be reintroduced by rename.
This commit is contained in:
parent
d745030539
commit
5643b80340
3 changed files with 207 additions and 501 deletions
|
|
@ -363,14 +363,21 @@ class CustomStreamWrapper:
|
|||
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
|
||||
return self
|
||||
|
||||
def close(self) -> None:
|
||||
def close_stream(self) -> None:
|
||||
"""Deterministically end a stream the consumer stopped reading early.
|
||||
|
||||
A stream abandoned without reaching its final chunk would otherwise
|
||||
never emit its terminal success log (spans, spend, custom callbacks
|
||||
all miss the call). ``close()`` flushes that log from the chunks
|
||||
all miss the call). ``close_stream()`` flushes that log from the chunks
|
||||
received so far and releases the underlying provider stream; safe to
|
||||
call after the stream has been fully consumed (it is a no-op then).
|
||||
|
||||
Deliberately *not* named ``close``: ``Router.stream_with_fallbacks``
|
||||
duck-types on ``hasattr(model_response, "close")`` and calls it from a
|
||||
``finally``, so adding a ``close`` here silently starts firing the
|
||||
terminal success log from the fallback path, which changes how shadow
|
||||
and fallback streams report completions. Callers that hold a concrete
|
||||
``CustomStreamWrapper`` use this name or the context manager below.
|
||||
"""
|
||||
self._flush_abandoned_stream_logging()
|
||||
stream_to_close: Final = self.completion_stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the underlying stream object is provider-specific
|
||||
|
|
@ -386,7 +393,7 @@ class CustomStreamWrapper:
|
|||
return self
|
||||
|
||||
def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
|
||||
self.close()
|
||||
self.close_stream()
|
||||
return False
|
||||
|
||||
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:
|
||||
|
|
|
|||
0
litellm/litellm_core_utils/test_streaming_handler.py
Normal file
0
litellm/litellm_core_utils/test_streaming_handler.py
Normal file
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue