From c7a41c35d5cdf2fc2d270c6fee6c793729eeee06 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:48:25 +0000 Subject: [PATCH] perf(mock): emit admission-time usage chunk on streaming mock_response (#40637) * perf(mock): emit admission-time usage chunk on streaming mock_response Streaming mock_response chunks carried no usage, so the chunk builder re-tokenized the whole prompt in Python after the stream ended even when budget reservation had already counted it at admission. The mock streaming generators now yield a final usage-only chunk carrying the admission prompt count (same completion count as the non-streaming path). Without an admission count the old tokenizer fallback stays. * fix(mock): type the mock stream generators and keep the usage chunk on the content stream id Review follow-up: the usage-only chunk was built with a fresh id, so CustomStreamWrapper switched response_id for the finish-reason and usage chunks. It now copies the content stream id. The generators also get full parameter and return annotations. --------- Co-authored-by: yassin --- litellm/main.py | 4 +- litellm/utils.py | 35 +++- .../test_streaming_chunk_builder_utils.py | 49 ++++++ tests/test_litellm/test_main.py | 158 ++++++++++++++++++ tests/test_litellm/test_utils.py | 90 ++++++++++ 5 files changed, 329 insertions(+), 7 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 1a4beb787bc..10fb32828f1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -939,7 +939,7 @@ def mock_completion( if kwargs.get("acompletion", False) is True: return CustomStreamWrapper( completion_stream=async_mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", @@ -947,7 +947,7 @@ def mock_completion( ) return CustomStreamWrapper( completion_stream=mock_completion_streaming_obj( - model_response, mock_response=mock_response, model=model, n=n + model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens ), model=model, custom_llm_provider="openai", diff --git a/litellm/utils.py b/litellm/utils.py index aced4c9b312..b93938a480e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -66,6 +66,7 @@ from litellm.constants import ( DEFAULT_EMBEDDING_PARAM_VALUES, DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, @@ -278,7 +279,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args from litellm import utils as litellm_utils @@ -7001,7 +7002,26 @@ class TextCompletionStreamWrapper: raise StopAsyncIteration -def mock_completion_streaming_obj(model_response, mock_response, model, n: int | None = None): +def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream: + return ModelResponseStream( + id=model_response.id, + choices=[], # mutable-ok: ModelResponseStream only treats a list as explicit choices, a tuple gets a default choice + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + ), + ) + + +def mock_completion_streaming_obj( + model_response: ModelResponseStream, + mock_response: str | MockException | ModelResponseStream, + model: str, + n: int | None = None, + prompt_tokens: int | None = None, +) -> Iterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7021,14 +7041,17 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int | _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) async def async_mock_completion_streaming_obj( - model_response, + model_response: ModelResponseStream, mock_response: str | MockException | ModelResponseStream, - model, + model: str, n: int | None = None, -): + prompt_tokens: int | None = None, +) -> AsyncIterator[ModelResponseStream]: if isinstance(mock_response, litellm.MockException): raise mock_response if isinstance(mock_response, ModelResponseStream): @@ -7048,6 +7071,8 @@ async def async_mock_completion_streaming_obj( _all_choices.append(_streaming_choice) model_response.choices = _all_choices yield model_response + if prompt_tokens is not None: + yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens) ########## Reading Config File ############################ diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 626b8a63b20..efe4209c1c9 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1554,3 +1554,52 @@ def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None assert response is not None assert response.choices[0].message.role == "user" assert response.choices[0].message.content == "Hi" + + +def _fail_prompt_token_count() -> int: + raise AssertionError("prompt tokens must come from the usage chunk, not the tokenizer") + + +def test_calculate_usage_reads_prompt_tokens_from_mock_stream_usage_chunk_without_tokenizer_fallback() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), + mock_response="ok", + model="gpt-5.4-mini", + prompt_tokens=51234, + ) + ) + assert chunks[-1].choices == [] + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=_fail_prompt_token_count, + ) + + assert usage.prompt_tokens == 51234 + assert usage.completion_tokens == chunks[-1].usage.completion_tokens + assert usage.total_tokens == 51234 + usage.completion_tokens + + +def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_admission_count() -> None: + from litellm.utils import mock_completion_streaming_obj + + chunks: Final = list( + mock_completion_streaming_obj( + ModelResponseStream(model="gpt-5.4-mini"), mock_response="ok", model="gpt-5.4-mini" + ) + ) + assert all(chunk.choices for chunk in chunks) + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="ok", + count_prompt_tokens=lambda: 77, + ) + + assert usage.prompt_tokens == 77 diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index f8841af9750..a36ca229981 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2421,6 +2421,164 @@ def test_mock_completion_usage_falls_back_to_default_without_admission_count(): assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT +_ADMISSION_INPUT_TOKENS: Final = 51234 +_ADMISSION_METADATA: Final = { + "user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": _ADMISSION_INPUT_TOKENS} +} +_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}] +_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter" + + +def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]: + return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None] + + +def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]: + return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + +@pytest.mark.parametrize("n", (None, 2)) +def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + usage_chunks = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", (None, 2)) +async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback( + n: int | None, +): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + n=n, + stream_options={"include_usage": True}, + litellm_metadata=_ADMISSION_METADATA, + ) + chunks = [chunk async for chunk in response] + + usage_chunks = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens + assert _prompt_token_counter_calls(token_counter) == [] + assert all(chunk.choices for chunk in chunks[:-1]) + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + +def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count(): + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + metadata=_ADMISSION_METADATA, + ) + ) + + assert _client_usage_chunks(chunks) == [] + assert all(len(chunk.choices) == 1 for chunk in chunks) + assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS + assert _prompt_token_counter_calls(token_counter) == [] + + +def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + chunks = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + ) + + usage_chunks = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +@pytest.mark.asyncio +async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer(): + expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES) + with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter: + response = await litellm.acompletion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + ) + chunks = [chunk async for chunk in response] + + usage_chunks = _client_usage_chunks(chunks) + assert len(usage_chunks) == 1 + assert usage_chunks[0].prompt_tokens == expected_prompt_tokens + assert len(_prompt_token_counter_calls(token_counter)) >= 1 + + +def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(): + non_stream = litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + metadata=_ADMISSION_METADATA, + ) + chunks = list( + litellm.completion( + model="openai/gpt-5.4-mini", + messages=_MOCK_STREAM_MESSAGES, + mock_response="ok", + api_key="mock", + stream=True, + stream_options={"include_usage": True}, + metadata=_ADMISSION_METADATA, + ) + ) + + stream_usage: Final = _client_usage_chunks(chunks)[0] + assert (non_stream.usage.prompt_tokens, non_stream.usage.completion_tokens, non_stream.usage.total_tokens) == ( + stream_usage.prompt_tokens, + stream_usage.completion_tokens, + stream_usage.total_tokens, + ) + + def test_mock_completion_stream_with_model_response(): """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" from litellm import completion diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8b186be43e5..2bdde43ff2d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -15,6 +15,7 @@ from jsonschema import validate import litellm from litellm._internal_context import is_internal_call +from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm._logging import ( CorrelationContextFilter, JsonFormatter, @@ -6208,3 +6209,92 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with "api_key": "sk-from-db", } assert _credential_warnings(caplog) == [] + + +_MOCK_STREAM_ID: Final = "chatcmpl-mock-stream" +_ChunkSnapshot = tuple[str, tuple[str | None, ...], Usage | None] + + +def _snapshot(chunk: ModelResponseStream) -> _ChunkSnapshot: + return chunk.id, tuple(choice.delta.content for choice in chunk.choices), getattr(chunk, "usage", None) + + +def _mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import mock_completion_streaming_obj + + return [ + _snapshot(chunk) + for chunk in mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +async def _async_mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]: + from litellm.utils import async_mock_completion_streaming_obj + + return [ + _snapshot(chunk) + async for chunk in async_mock_completion_streaming_obj( + ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"), + mock_response=mock_response, + model="gpt-5.4-mini", + prompt_tokens=prompt_tokens, + ) + ] + + +_CONTENT_SNAPSHOTS: Final = [(_MOCK_STREAM_ID, (content,), None) for content in ("hel", "lo ", "wor", "ld")] + + +def _assert_trailing_usage_chunk(snapshots: list[_ChunkSnapshot], prompt_tokens: int) -> None: + assert snapshots[:-1] == _CONTENT_SNAPSHOTS + chunk_id, choices, usage = snapshots[-1] + assert chunk_id == _MOCK_STREAM_ID + assert choices == () + assert usage is not None + assert usage.prompt_tokens == prompt_tokens + assert usage.completion_tokens == DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + assert usage.total_tokens == prompt_tokens + usage.completion_tokens + + +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +def test_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens(prompt_tokens: int) -> None: + _assert_trailing_usage_chunk(_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prompt_tokens", (51234, 0)) +async def test_async_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens( + prompt_tokens: int, +) -> None: + _assert_trailing_usage_chunk(await _async_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens) + + +def test_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert _mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None: + assert await _async_mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS + + +def test_mock_completion_streaming_obj_passes_prebuilt_stream_chunk_through_without_usage_chunk() -> None: + prebuilt: Final = ModelResponseStream( + model="gpt-5.4-mini", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="prebuilt"))] + ) + + assert _mock_stream_snapshots(prebuilt, 51234) == [(prebuilt.id, ("prebuilt",), None)] + + +@pytest.mark.asyncio +async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_usage_chunk() -> None: + mock_exception: Final = litellm.MockException( + status_code=500, message="boom", llm_provider="openai", model="gpt-5.4-mini" + ) + with pytest.raises(litellm.MockException): + await _async_mock_stream_snapshots(mock_exception, 51234)