diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 60eef9604e..73177fdd48 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -42,6 +42,7 @@ guardrails: litellm_settings: callbacks: ["datadog"] + include_cost_in_streaming_usage: true datadog_params: turn_off_message_logging: true datadog_llm_observability_params: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 64ea93028f..93abdee778 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -49,6 +49,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.litellm_metadata: Optional[dict] = litellm_metadata or {} self.collected_chat_completion_chunks: List[ModelResponseStream] = [] self.finished: bool = False + self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj async def __anext__( self, @@ -167,8 +168,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _emit_response_completed_event(self) -> Optional[ResponseCompletedEvent]: litellm_model_response: Optional[ Union[ModelResponse, TextCompletionResponse] - ] = stream_chunk_builder(chunks=self.collected_chat_completion_chunks) + ] = stream_chunk_builder(chunks=self.collected_chat_completion_chunks, logging_obj=self.litellm_logging_obj) if litellm_model_response and isinstance(litellm_model_response, ModelResponse): + # Add cost to usage object if include_cost_in_streaming_usage is True + if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None: + usage = getattr(litellm_model_response, "usage", None) + if usage is not None: + setattr( + usage, "cost", self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response) + ) + # Transform the response responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( request_input=self.request_input, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 82d3980b37..a43e02a0f3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -851,8 +851,15 @@ class LiteLLMCompletionResponsesConfig: output_tokens=0, total_tokens=0, ) - return ResponseAPIUsage( + + response_usage = ResponseAPIUsage( input_tokens=usage.prompt_tokens, output_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, ) + + # Preserve cost field if it exists (for streaming usage with cost calculation) + if hasattr(usage, "cost") and usage.cost is not None: + setattr(response_usage, "cost", usage.cost) + + return response_usage diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index e9e41789f0..eda3e6921d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,6 +5,7 @@ from typing import Any, Dict, Optional import httpx +import litellm from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -13,6 +14,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( OutputTextDeltaEvent, + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -95,6 +97,20 @@ class BaseResponsesAPIStreamingIterator: == ResponsesAPIStreamEvents.RESPONSE_COMPLETED ): self.completed_response = openai_responses_api_chunk + # Add cost to usage object if include_cost_in_streaming_usage is True + if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: + response_obj: Optional[ResponsesAPIResponse] = getattr(openai_responses_api_chunk, "response", None) + if response_obj: + usage_obj: Optional[ResponseAPIUsage] = getattr(response_obj, "usage", None) + if usage_obj is not None: + try: + cost: Optional[float] = self.logging_obj._response_cost_calculator(result=response_obj) + if cost is not None: + setattr(usage_obj, "cost", cost) + except Exception: + # If cost calculation fails, continue without cost + pass + self._handle_logging_completed_response() return openai_responses_api_chunk diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 434035b809..9f4ae03b39 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1033,6 +1033,9 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): total_tokens: int """The total number of tokens used.""" + cost: Optional[float] = None + """The cost of the request.""" + model_config = {"extra": "allow"} diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 855aeff246..afb30f15f0 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -146,6 +146,8 @@ class BaseResponsesAPITest(ABC): @pytest.mark.flaky(retries=3, delay=2) async def test_basic_openai_responses_api_streaming(self, sync_mode): litellm._turn_on_debug() + # Enable cost calculation for streaming usage + litellm.include_cost_in_streaming_usage = True base_completion_call_args = self.get_base_completion_call_args() collected_content_string = "" response_completed_event = None @@ -208,6 +210,14 @@ class BaseResponsesAPITest(ABC): + response_completed_event.response.usage.output_tokens ) + # assert the response completed event includes cost when include_cost_in_streaming_usage is True + assert hasattr(response_completed_event.response.usage, "cost"), "Cost should be included in streaming responses API usage object" + assert response_completed_event.response.usage.cost > 0, "Cost should be greater than 0" + print(f"Cost found in streaming response: {response_completed_event.response.usage.cost}") + + # Reset the setting + litellm.include_cost_in_streaming_usage = False + @pytest.mark.parametrize("sync_mode", [False, True]) @pytest.mark.asyncio async def test_basic_openai_responses_delete_endpoint(self, sync_mode):