diff --git a/litellm/__init__.py b/litellm/__init__.py index 0f48fb4cf4..bc8fb9c135 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -269,6 +269,7 @@ blocked_user_list: Optional[Union[str, List]] = None banned_keywords_list: Optional[Union[str, List]] = None llm_guard_mode: Literal["all", "key-specific", "request-specific"] = "all" guardrail_name_config_map: Dict[str, GuardrailItem] = {} +include_cost_in_streaming_usage: bool = False ### PROMPTS ### from litellm.types.prompts.init_prompts import PromptSpec diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2e9e6770a1..3721851a38 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1584,7 +1584,9 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, messages=self.messages + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, ) response = self.model_response_creator() @@ -1768,7 +1770,9 @@ class CustomStreamWrapper: if self.sent_last_chunk is True: # log the final chunk with accurate streaming values complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, messages=self.messages + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, ) response = self.model_response_creator() if complete_streaming_response is not None: diff --git a/litellm/main.py b/litellm/main.py index 108abbde84..7762e178f6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5704,7 +5704,11 @@ def stream_chunk_builder_text_completion( def stream_chunk_builder( # noqa: PLR0915 - chunks: list, messages: Optional[list] = None, start_time=None, end_time=None + chunks: list, + messages: Optional[list] = None, + start_time=None, + end_time=None, + logging_obj: Optional[Logging] = None, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: try: if chunks is None: @@ -5829,6 +5833,12 @@ def stream_chunk_builder( # noqa: PLR0915 setattr(response, "usage", usage) + # Add cost to usage object if include_cost_in_streaming_usage is True + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + setattr( + usage, "cost", logging_obj._response_cost_calculator(result=response) + ) + return response except Exception as e: verbose_logger.exception( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f8df5781de..6606b0e0b8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -887,6 +887,7 @@ class Usage(CompletionUsage): ) # hidden param for prompt caching. Might change, once openai introduces their equivalent. server_tool_use: Optional[ServerToolUse] = None + cost: Optional[float] = None completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None """Breakdown of tokens used in a completion.""" @@ -907,6 +908,7 @@ class Usage(CompletionUsage): Union[CompletionTokensDetailsWrapper, dict] ] = None, server_tool_use: Optional[ServerToolUse] = None, + cost: Optional[float] = None, **params, ): # handle reasoning_tokens @@ -978,6 +980,11 @@ class Usage(CompletionUsage): else: # maintain openai compatibility in usage object if possible del self.server_tool_use + if cost is not None: + self.cost = cost + else: + del self.cost + ## ANTHROPIC MAPPING ## if "cache_creation_input_tokens" in params and isinstance( params["cache_creation_input_tokens"], int diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 823b635028..445dad57a5 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -3948,3 +3948,45 @@ def test_is_delta_empty(): audio=None, ) ) + + +def test_streaming_with_cost_calculation(): + from litellm.types.utils import Usage + from typing import Optional + + litellm.include_cost_in_streaming_usage = True + + ## Test 1: check if usage object can handle 'cost' field + usage_object = Usage( + prompt_tokens=100, + completion_tokens=100, + total_tokens=200, + cost=1.0, + ) + assert usage_object.cost is not None + + print(f"usage_object: {usage_object}") + + ## Test 2: check if usage object has 'cost' field when streaming + + response = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What is the capital of France?"}], + stream=True, + stream_options={"include_usage": True}, + ) + + usage_object: Optional[Usage] = None + for chunk in response: + _usage_obj = getattr(chunk, "usage", None) + if _usage_obj is not None: + usage_object = _usage_obj + break + + assert usage_object is not None + assert usage_object.total_tokens is not None + assert usage_object.total_tokens > 0 + assert usage_object.prompt_tokens is not None + assert usage_object.prompt_tokens > 0 + assert usage_object.cost is not None + assert usage_object.cost > 0