diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 0d840d136d0..3ef9c7bee4d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,7 +41,7 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base -from litellm.litellm_core_utils.hidden_params import stamp_billed_usage_cost +from litellm.litellm_core_utils.hidden_params import USAGE_COST_BEFORE_LITELLM_STAMP, stamp_billed_usage_cost from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -676,10 +676,17 @@ class BaseResponsesAPIStreamingIterator: headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy - # splats into the client's HTTP headers, and copying non-header keys would carry response_cost + # splats into the client's HTTP headers, and copying non-header keys would carry response_cost. + # The provider's own usage.cost travels along, or the stamped bill would be priced as a provider report + stamp: Final[Mapping[str, object]] = ( + {USAGE_COST_BEFORE_LITELLM_STAMP: source[USAGE_COST_BEFORE_LITELLM_STAMP]} + if USAGE_COST_BEFORE_LITELLM_STAMP in source + else EMPTY_MAPPING + ) target.hidden_params = { "additional_headers": {**headers}, "headers": {**raw_headers}, + **stamp, **existing, } diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 6c3e6177503..5716f46c210 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2781,3 +2781,39 @@ async def test_websocket_turns_logging_survives_garbage_collection(terminal_even await _run_websocket_turn(logger, terminal_event) assert await _logging_tasks_left_after_garbage_collection(logger) == 1 + + +def test_logging_copy_keeps_the_provider_usage_cost_behind_a_stamped_bill() -> None: + from litellm.litellm_core_utils.hidden_params import USAGE_COST_BEFORE_LITELLM_STAMP + from litellm.responses.streaming_iterator import _logging_copy + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="aresponses", + start_time=datetime.now(), + litellm_call_id="stamped-copy", + function_id="stamped-copy", + ) + iterator = _make_iterator(sse_events=[], logging_obj=logging_obj) + billed = ResponsesAPIResponse( + id="resp_stamped", + created_at=1, + status="completed", + model="gpt-4o-mini", + object="response", + output=[], + usage={"input_tokens": 100, "output_tokens": 20, "total_tokens": 120, "cost": 0.63}, + ) + billed.hidden_params = {USAGE_COST_BEFORE_LITELLM_STAMP: 0.42, "response_cost": 0.63} + iterator.completed_response = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=billed + ) + logging_response = _logging_copy(iterator.completed_response) + + iterator._restore_provider_response_headers(logging_response) + + copied = logging_response.response.hidden_params + assert copied[USAGE_COST_BEFORE_LITELLM_STAMP] == 0.42 + assert "response_cost" not in copied