diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index b42210e2793..39f675099d7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -46,6 +46,7 @@ from litellm.types.utils import ( StreamingChoices, TextChoices, TextCompletionResponse, + UnreportedUsage, Usage, ) from litellm.utils import print_verbose, token_counter @@ -876,6 +877,8 @@ class ChunkProcessor: if isinstance(usage_chunk, dict): return Usage(**usage_chunk) + if isinstance(usage_chunk, UnreportedUsage): + return None if usage_chunk is None or isinstance(usage_chunk, Usage): return usage_chunk return Usage(**usage_chunk.model_dump()) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a627b26b9af..6ec426d6da5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,7 +24,12 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, stored_attribute +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + set_hidden_param, + stored_attribute, +) from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) @@ -43,6 +48,8 @@ from litellm.types.utils import ( ModelResponseStream, PromptTokensDetailsWrapper, StreamingChoices, + TextCompletionResponse, + UnreportedUsage, Usage, ) from litellm.types.utils import GenericStreamingChunk as GChunk @@ -62,6 +69,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) +_CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None _USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) @@ -335,7 +343,8 @@ class CustomStreamWrapper: self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False self.tool_call = False - self.chunks: list = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options + # keep track of the returned chunks - used for calculating the input/output tokens for stream options + self.chunks: list[ModelResponseStream] = [] self._repeated_messages_count = 1 self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: int | None = None @@ -544,7 +553,7 @@ class CustomStreamWrapper: last_content: Final = self.chunks[-1].choices[0].delta.content if ( - last_content is None or not isinstance(last_content, str) or len(last_content) <= 2 + not isinstance(last_content, str) or len(last_content) <= 2 ): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946 self._repeated_messages_count = 1 return @@ -1927,6 +1936,51 @@ class CustomStreamWrapper: response.hidden_params["additional_headers"] = {} response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost + def _builder_failure_response(self, error: Exception) -> ModelResponseStream | None: + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", + str(error), + ) + try: + return self.model_response_creator(chunk={"usage": calculate_total_usage(chunks=self.chunks)}) + except Exception: + return None + + def _build_complete_streaming_response(self) -> _CompleteStreamingResponse: + try: + return litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: + return self._builder_failure_response(e) + + async def _async_build_complete_streaming_response(self) -> _CompleteStreamingResponse: + try: + return await asyncify(litellm.stream_chunk_builder)( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: + return self._builder_failure_response(e) + + def _set_final_chunk_usage(self, chunk: ModelResponseStream) -> None: + reported_usage: Final = _reported_total_usage(chunks=self.chunks) + set_hidden_param( + chunk, + "usage", + UnreportedUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + if reported_usage is None + else reported_usage, + ) + self._last_returned_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior + dict[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ) + def __next__(self) -> "ModelResponseStream": if self._closed: raise StopIteration @@ -2007,12 +2061,7 @@ class CustomStreamWrapper: continue # add usage as hidden param if self.sent_last_chunk is True and self.stream_options is None: - usage = calculate_total_usage(chunks=self.chunks) - response_hidden_params_for_final_chunk = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(response, HIDDEN_PARAMS_ATTR) - ) - response_hidden_params_for_final_chunk["usage"] = usage - self._last_returned_hidden_params = response_hidden_params_for_final_chunk + self._set_final_chunk_usage(response) # Add MCP metadata to final chunk if present response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT @@ -2020,30 +2069,7 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: - try: - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) - except Exception as e: - # stream_chunk_builder can re-raise (as APIError) on large agentic - # streams. The raise originates inside this except-StopIteration block, - # so the sibling `except Exception` below does not catch it; it would - # escape __next__ and drop the request from SpendLogs. Recover - # best-effort usage from the raw chunks so cost is still tracked - verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", - str(e), - ) - try: - complete_streaming_response = self.model_response_creator( - chunk={"usage": calculate_total_usage(chunks=self.chunks)} - ) - except Exception: - complete_streaming_response = None - + complete_streaming_response: Final = self._build_complete_streaming_response() response = self.model_response_creator() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2107,9 +2133,8 @@ class CustomStreamWrapper: # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: self.chunks.append(processed_chunk) - if self.stream_options is None: # add usage as hidden param - usage = calculate_total_usage(chunks=self.chunks) - processed_chunk.hidden_params["usage"] = usage + if self.stream_options is None: + self._set_final_chunk_usage(processed_chunk) ## LOGGING executor.submit( self.run_success_logging_and_cache_storage, @@ -2224,12 +2249,7 @@ class CustomStreamWrapper: # add usage as hidden param if self.sent_last_chunk is True and self.stream_options is None: - usage = calculate_total_usage(chunks=self.chunks) - processed_chunk_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ) - processed_chunk_hidden_params["usage"] = usage - self._last_returned_hidden_params = processed_chunk_hidden_params + self._set_final_chunk_usage(processed_chunk) # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: @@ -2285,29 +2305,7 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: - # log the final chunk with accurate streaming values - try: - complete_streaming_response = await asyncify(litellm.stream_chunk_builder)( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) - except Exception as e: - # see sync __next__: a raise from stream_chunk_builder inside this - # except handler escapes __anext__ and drops the request from SpendLogs. - # Recover best-effort usage from the raw chunks so cost is still tracked - verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", - str(e), - ) - try: - complete_streaming_response = self.model_response_creator( - chunk={"usage": calculate_total_usage(chunks=self.chunks)} - ) - except Exception: - complete_streaming_response = None - + complete_streaming_response: Final = await self._async_build_complete_streaming_response() response: Final = self.model_response_creator() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2384,11 +2382,7 @@ class CustomStreamWrapper: if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: self.chunks.append(processed_chunk) if self.stream_options is None: - usage: Final = calculate_total_usage(chunks=self.chunks) - processed_chunk_hidden_params: Final = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ) - processed_chunk_hidden_params["usage"] = usage + self._set_final_chunk_usage(processed_chunk) # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log @@ -2659,7 +2653,13 @@ def _coerce_token_details( return details_type(**(raw if isinstance(raw, dict) else raw.model_dump())) -def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: +def _reported_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage | None: + if not any("usage" in chunk and chunk["usage"] is not None for chunk in chunks): + return None + return calculate_total_usage(chunks=chunks) + + +def calculate_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( attach_cache_creation_token_details, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 576c4f093d3..2721a0c1cc1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2062,6 +2062,10 @@ class Usage(SafeAttributeModel, CompletionUsage): setattr(self, key, value) +class UnreportedUsage(Usage): + pass + + class StreamingChoices(OpenAIObject): def __init__( self, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 691442e80d0..68aabc3c01e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,4 +1,5 @@ import asyncio +import datetime import json import sys import threading @@ -18,6 +19,7 @@ import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( @@ -26,6 +28,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, @@ -2820,6 +2823,238 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): ), f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("ends_with_finish_chunk", [True, False]) +@pytest.mark.asyncio +async def test_stream_without_provider_usage_falls_back_to_the_token_estimate( + sync_mode: bool, ends_with_finish_chunk: bool +) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + content_chunk: Final = ModelResponseStream( + id="chatcmpl-no-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + ) + finish_chunk: Final = ModelResponseStream( + id="chatcmpl-no-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk] + ), + model=model, + custom_llm_provider="openai", + logging_obj=Logging( + model=model, + messages=messages, + stream=True, + call_type="completion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="no-usage-call", + function_id="no-usage-fn", + ), + stream_options=None, + ) + + yielded: Final = ( + [(chunk, chunk._hidden_params.get("usage")) for chunk in wrapper] + if sync_mode + else [(chunk, chunk._hidden_params.get("usage")) async for chunk in wrapper] + ) + collected: Final = [chunk for chunk, _ in yielded] + + expected: Final = ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + assembled: Final = litellm.stream_chunk_builder(chunks=collected, messages=messages) + assert assembled is not None + assert (assembled.usage.prompt_tokens, assembled.usage.completion_tokens) == expected + last_chunk, usage_when_yielded = yielded[-1] + assert usage_when_yielded is not None + usage_after_drain: Final = last_chunk._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == expected + + rates: Final = litellm.model_cost[model] + spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"] + assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_with_usage_after_the_finish_chunk_bills_only_the_reported_usage(sync_mode: bool) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + reported: Final = (6, 20) + chunks: Final = [ + ModelResponseStream(id="cmpl-late-usage", created=1741037890, model=model, choices=choices, usage=usage) + for choices, usage in ( + ([StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], None), + ([StreamingChoices(index=0, delta=Delta(content=""), finish_reason="length")], None), + ([], Usage(prompt_tokens=reported[0], completion_tokens=reported[1], total_tokens=sum(reported))), + ) + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model=model, + custom_llm_provider="openai", + logging_obj=Logging( + model=model, + messages=messages, + stream=True, + call_type="completion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="late-usage-call", + function_id="late-usage-fn", + ), + stream_options=None, + ) + + collected: Final = [chunk for chunk in wrapper] if sync_mode else [chunk async for chunk in wrapper] + + usage_after_drain: Final = collected[-1]._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == reported + assembled: Final = litellm.stream_chunk_builder(chunks=collected, messages=messages) + assert assembled is not None + rates: Final = litellm.model_cost[model] + spend: Final = reported[0] * rates["input_cost_per_token"] + reported[1] * rates["output_cost_per_token"] + assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend) + + +class _LoggedResponseRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.responses: list[object] = [] + self.logged: Final = threading.Event() + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.responses.append(response_obj) + self.logged.set() + + async def async_log_success_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.responses.append(response_obj) + self.logged.set() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_without_provider_usage_logs_the_estimate_with_the_final_chunk_mcp_metadata( + sync_mode: bool, +) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + tool_calls: Final = [{"name": "lookup_fox", "arguments": "{}"}] + recorder: Final = _LoggedResponseRecorder() + logging_obj: Final = Logging( + model=model, + messages=messages, + stream=True, + call_type="completion" if sync_mode else "acompletion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="mcp-no-usage-call", + function_id="mcp-no-usage-fn", + dynamic_success_callbacks=[recorder], + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "mcp-no-usage-call"} + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream(id="chatcmpl-mcp", created=1741037890, model=model, choices=choices) + for choices in ( + [StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + [StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + ] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + _add_mcp_metadata_to_response(wrapper, openai_tools=None, tool_calls=tool_calls) + + _ = [chunk for chunk in wrapper] if sync_mode else [chunk async for chunk in wrapper] + + assert await asyncio.to_thread(recorder.logged.wait, 5) + assert len(recorder.responses) == 1 + logged: Final = recorder.responses[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.provider_specific_fields == {"mcp_tool_calls": tool_calls} + assert (logged.usage.prompt_tokens, logged.usage.completion_tokens) == ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("ends_with_finish_chunk", [True, False]) +@pytest.mark.parametrize("reported_tokens", [(30, 12), (0, 0)]) +@pytest.mark.asyncio +async def test_stream_with_provider_usage_hands_that_usage_to_the_last_chunk( + logging_obj: Logging, sync_mode: bool, ends_with_finish_chunk: bool, reported_tokens: tuple[int, int] +) -> None: + model: Final = "gpt-4o" + reported: Final = Usage( + prompt_tokens=reported_tokens[0], + completion_tokens=reported_tokens[1], + total_tokens=sum(reported_tokens), + ) + content_chunk: Final = ModelResponseStream( + id="chatcmpl-with-usage", + created=1741037890, + model=model, + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content="Hello there"), + finish_reason=None if ends_with_finish_chunk else "stop", + ) + ], + **({} if ends_with_finish_chunk else {"usage": reported}), + ) + finish_chunk: Final = ModelResponseStream( + id="chatcmpl-with-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + usage=reported, + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + + yielded: Final = ( + [(chunk, chunk._hidden_params.get("usage")) for chunk in wrapper] + if sync_mode + else [(chunk, chunk._hidden_params.get("usage")) async for chunk in wrapper] + ) + + last_chunk, usage_when_yielded = yielded[-1] + assert usage_when_yielded is not None + assert (usage_when_yielded.prompt_tokens, usage_when_yielded.completion_tokens) == reported_tokens + usage_after_drain: Final = last_chunk._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == reported_tokens + + @pytest.mark.asyncio async def test_custom_stream_wrapper_aclose(): """Test that aclose() delegates to the underlying completion_stream's aclose()""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 679e1bfd85a..f837dc3fb5f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails +from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -76,6 +76,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.utils import ModelResponseListIterator from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment @@ -8480,9 +8482,11 @@ class _RecordingSuccessLogger(CustomLogger): def __init__(self): super().__init__() self.success_events = [] + self.logged: Final = asyncio.Event() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) + self.logged.set() class TestStreamingClientDisconnectBilling: @@ -8593,6 +8597,60 @@ class TestStreamingClientDisconnectBilling: assert getattr(partial_response, "service_tier") == "priority" assert partial_response.usage.total_tokens > 0 + @pytest.mark.asyncio + async def test_disconnect_after_finish_chunk_bills_the_estimate_when_the_provider_sent_no_usage(self): + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + recorder: Final = _RecordingSuccessLogger() + logging_obj: Final = LiteLLMLoggingObj( + model=model, + messages=messages, + stream=True, + call_type="acompletion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="no-usage-disconnect-call", + function_id="no-usage-disconnect-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "no-usage-disconnect-call"} + ) + response: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream(id="chatcmpl-no-usage", created=1741037890, model=model, choices=choices) + for choices in ( + [StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + [StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + ] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + stream_iter: Final = response.__aiter__() + first_chunk: Final = await stream_iter.__anext__() + finish_chunk: Final = await stream_iter.__anext__() + assert first_chunk.choices[0].finish_reason is None + assert finish_chunk.choices[0].finish_reason == "stop" + + billed: Final = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": logging_obj}, + response, + ) + await asyncio.wait_for(recorder.logged.wait(), timeout=5) + + assert billed is True + assert len(recorder.success_events) == 1 + usage: Final = recorder.success_events[0]["response_obj"].usage + assert (usage.prompt_tokens, usage.completion_tokens) == ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + @pytest.mark.asyncio async def test_completed_stream_does_not_double_bill_on_late_disconnect(self): recorder = _RecordingSuccessLogger()