scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

Merge f75dd80bc0 into 2c6b280755

Browse source
This commit is contained in:
Akilan Chithra Sathish 2026-10-20 16:58:07 +00:00 • committed by GitHub
commit dc14547ced
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 279 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/litellm_core_utils/streaming_handler.py
View file

@ -507,6 +507,14 @@ class CustomStreamWrapper:
def check_send_stream_usage(self, stream_options: dict | None):
return stream_options is not None and stream_options.get("include_usage", False) is True
def _should_forward_continuous_usage(self) -> bool:
options: Final[Mapping[str, object] | None] = getattr(self, "stream_options", None)
return (
self.send_stream_usage is True
and options is not None
and options.get("continuous_usage_stats", False) is True
)
def check_is_function_call(self, logging_obj) -> bool:
from litellm.litellm_core_utils.prompt_templates.common_utils import (
is_function_call,
@ -1236,6 +1244,8 @@ class CustomStreamWrapper:
return self._handle_special_delta_content(model_response)
else:
if hasattr(model_response, "usage"):
if self._should_forward_continuous_usage() and getattr(model_response, "usage", None) is not None:
return model_response
self.chunks.append(model_response)
return
@ -1570,6 +1580,11 @@ class CustomStreamWrapper:
return _ProviderChunkParsed(response_obj)
def chunk_creator(self, chunk: Any):
finish_reason_seen: Final = (
self.sent_last_chunk is True
or self.received_finish_reason is not None
or self.intermittent_finish_reason is not None
)
if self.custom_llm_provider != "cached_response" and isinstance(
chunk, (ChatCompletionChunk, ModelResponseStream)
):
@ -1666,7 +1681,9 @@ class CustomStreamWrapper:
model_response.choices[0].delta = plain_delta if plain_delta is not None else Delta()
plain_delta_created = plain_delta is not None # rebind-ok: the first chunk builds its delta
else:
if self.send_stream_usage is True:
if self.send_stream_usage is True and (
not self._should_forward_continuous_usage() or not finish_reason_seen
):
model_response.choices = []
return model_response
self._record_usage_only_chunk(model_response=model_response)
@ -2045,7 +2062,7 @@ class CustomStreamWrapper:
obj_dict = response.model_dump()
if "usage" in obj_dict:
if "usage" in obj_dict and not self._should_forward_continuous_usage():
del obj_dict["usage"]
response_hidden_params_for_model = cast( # cast-ok: preserve dynamic mapping behavior
@ -2227,10 +2244,8 @@ class CustomStreamWrapper:
# directly to avoid expensive model_copy() per chunk.
self.chunks.append(processed_chunk.model_copy())
# Strip usage from the outgoing chunk so it's not sent twice
# (once in the chunk, once in _hidden_params).
obj_dict = processed_chunk.model_dump()
if "usage" in obj_dict:
if "usage" in obj_dict and not self._should_forward_continuous_usage():
del obj_dict["usage"]
processed_chunk = self.model_response_creator(
chunk=obj_dict,

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

@ -523,6 +523,265 @@ def test_strip_sse_data_from_chunk():
assert CustomStreamWrapper.strip_sse_data_from_chunk(None) is None
def _continuous_usage_fixture_chunks() -> tuple[ModelResponseStream, ...]:
return (
ModelResponseStream(
id="chatcmpl-continuous-usage",
created=1,
model="stream-usage-test",
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="a"))],
usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11),
),
ModelResponseStream(
id="chatcmpl-continuous-usage",
created=1,
model="stream-usage-test",
choices=[StreamingChoices(index=0, delta=Delta(content="b"), finish_reason="stop")],
usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
),
ModelResponseStream(
id="chatcmpl-continuous-usage",
created=1,
model="stream-usage-test",
choices=[],
usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
),
)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_forwards_running_counts(logging_obj: Logging, sync_mode: bool) -> None:
wrapper: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=_continuous_usage_fixture_chunks()),
model="stream-usage-test",
logging_obj=logging_obj,
custom_llm_provider="hosted_vllm",
stream_options={"include_usage": True, "continuous_usage_stats": True},
)
collected: Final = tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper])
content_chunks: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content)
assert tuple(chunk.choices[0].delta.content for chunk in content_chunks) == ("a", "b")
assert tuple(getattr(chunk, "usage", None) for chunk in content_chunks) == (
Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11),
Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
)
expected_final_usage: Final = Usage(
prompt_tokens=10,
completion_tokens=2,
total_tokens=12,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2),
)
assert collected[-1].usage == expected_final_usage
assert tuple(getattr(chunk, "usage", None) for chunk in collected) == (
Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11),
Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12),
None,
expected_final_usage,
)
def _continuous_usage_stream(
logging_obj: Logging,
chunks: tuple[ModelResponseStream, ...],
stream_options: dict[str, bool] | None,
) -> CustomStreamWrapper:
return CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=chunks),
model="stream-usage-test",
logging_obj=logging_obj,
custom_llm_provider="hosted_vllm",
stream_options=stream_options,
)
async def _continuous_usage_collected(wrapper: CustomStreamWrapper, sync_mode: bool) -> tuple[ModelResponseStream, ...]:
return tuple(wrapper) if sync_mode else tuple([chunk async for chunk in wrapper])
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.parametrize(
("stream_options", "include_final"),
(
(None, False),
({"include_usage": True}, True),
({"include_usage": True, "continuous_usage_stats": False}, True),
({"include_usage": False, "continuous_usage_stats": True}, False),
({"continuous_usage_stats": True}, False),
),
)
@pytest.mark.asyncio
async def test_continuous_usage_default_contract(
logging_obj: Logging, sync_mode: bool, stream_options: dict[str, bool] | None, include_final: bool
) -> None:
wrapper: Final = _continuous_usage_stream(logging_obj, _continuous_usage_fixture_chunks(), stream_options)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content)
assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None)
assert getattr(collected[-1], "usage", None) == (
Usage(
prompt_tokens=10,
completion_tokens=2,
total_tokens=12,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=2),
)
if include_final
else None
)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_preserves_zero_counts(logging_obj: Logging, sync_mode: bool) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
zero: Final = chunks[0].model_copy(
update={
"choices": [StreamingChoices(index=0, delta=Delta(role="assistant", content=""))],
"usage": Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10),
}
)
wrapper: Final = _continuous_usage_stream(
logging_obj, (zero, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert getattr(collected[0], "usage", None) == Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_does_not_invent_intermediate_counts(logging_obj: Logging, sync_mode: bool) -> None:
chunks: Final = tuple(chunk.model_copy(update={"usage": None}) for chunk in _continuous_usage_fixture_chunks()[:2])
wrapper: Final = _continuous_usage_stream(
logging_obj, chunks, {"include_usage": True, "continuous_usage_stats": True}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
content: Final = tuple(chunk for chunk in collected if chunk.choices and chunk.choices[0].delta.content)
assert tuple(getattr(chunk, "usage", None) for chunk in content) == (None, None)
assert collected[-1].usage is not None
assert collected[-1].usage.total_tokens == (
collected[-1].usage.prompt_tokens + collected[-1].usage.completion_tokens
)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_only_chunk_is_visible(logging_obj: Logging, sync_mode: bool) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
progress: Final = chunks[-1].model_copy(
update={"usage": Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)}
)
wrapper: Final = _continuous_usage_stream(
logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert collected[1].choices == [StreamingChoices(index=0, delta=Delta())]
assert getattr(collected[1], "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)
assert len(collected) == 5
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.parametrize("content", (None, ""))
@pytest.mark.parametrize("completion_tokens", (0, 1))
@pytest.mark.parametrize("continuous", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_empty_choice_is_visible(
logging_obj: Logging, sync_mode: bool, content: str | None, completion_tokens: int, continuous: bool
) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
progress: Final = chunks[0].model_copy(
update={
"choices": [StreamingChoices(index=0, delta=Delta(content=content))],
"usage": Usage(prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens),
}
)
wrapper: Final = _continuous_usage_stream(
logging_obj, (chunks[0], progress, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": continuous}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert len(collected) == (5 if continuous else 4)
if continuous:
assert collected[1].choices == [StreamingChoices(index=0, delta=Delta(content=content))]
assert getattr(collected[1], "usage", None) == Usage(
prompt_tokens=10, completion_tokens=completion_tokens, total_tokens=10 + completion_tokens
)
else:
assert collected[1].choices[0].delta.content == "b"
assert getattr(collected[1], "usage", None) is None
assert collected[-1].usage.total_tokens == 12
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_keeps_final_trailer_contract(logging_obj: Logging, sync_mode: bool) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
late: Final = chunks[-1].model_copy(update={"usage": Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13)})
wrapper: Final = _continuous_usage_stream(
logging_obj, (*chunks[:2], late), {"include_usage": True, "continuous_usage_stats": True}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert len(collected) == 4
assert collected[-2].choices[0].finish_reason == "stop"
assert collected[-1].usage == Usage(
prompt_tokens=10,
completion_tokens=3,
total_tokens=13,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=3),
)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_preserves_details_and_choices(logging_obj: Logging, sync_mode: bool) -> None:
chunks: Final = _continuous_usage_fixture_chunks()
detailed: Final = chunks[0].model_copy(
update={
"choices": [
StreamingChoices(index=0, delta=Delta(role="assistant", content="a")),
StreamingChoices(index=1, delta=Delta(content="other")),
],
"usage": Usage(
prompt_tokens=10,
completion_tokens=1,
total_tokens=11,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1),
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4),
),
}
)
default_wrapper: Final = _continuous_usage_stream(logging_obj, (detailed, *chunks[1:]), {"include_usage": True})
default_first: Final = next(default_wrapper) if sync_mode else await anext(default_wrapper)
await default_wrapper.aclose()
wrapper: Final = _continuous_usage_stream(
logging_obj, (detailed, *chunks[1:]), {"include_usage": True, "continuous_usage_stats": True}
)
collected: Final = await _continuous_usage_collected(wrapper, sync_mode)
assert collected[0].choices == [
StreamingChoices(index=0, delta=Delta(role="assistant", content="a")),
StreamingChoices(index=1, delta=Delta(content="other")),
]
assert collected[0].choices == default_first.choices
assert getattr(collected[0], "usage", None) == Usage(
prompt_tokens=10,
completion_tokens=1,
total_tokens=11,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=1),
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=4),
)
@pytest.mark.parametrize("sync_mode", (True, False))
@pytest.mark.asyncio
async def test_continuous_usage_is_available_before_completion(logging_obj: Logging, sync_mode: bool) -> None:
wrapper: Final = _continuous_usage_stream(
logging_obj, _continuous_usage_fixture_chunks(), {"include_usage": True, "continuous_usage_stats": True}
)
first: Final = next(wrapper) if sync_mode else await anext(wrapper)
assert first.choices[0].delta.content == "a"
assert getattr(first, "usage", None) == Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)
await wrapper.aclose()
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_streaming_handler_with_usage(