diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 658a0e8ad5d..40ff88fc557 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti import httpx from openai._streaming import SSEDecoder -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from typing_extensions import TypeIs import litellm @@ -439,18 +439,7 @@ class BaseResponsesAPIStreamingIterator: if self._persist_completed_response_before_logging: self._persist_completed_response_to_cache(is_async=is_async) - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, "model_dump"): - try: - logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump()) - except Exception: - # Fallback to original if serialization fails - pass + logging_response: Final[object] = _logging_copy(self.completed_response) self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() @@ -489,10 +478,10 @@ class BaseResponsesAPIStreamingIterator: def _restore_provider_response_headers(self, logging_response: object) -> None: """Re-apply the provider's response headers to the copy handed to logging callbacks. - ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the ``_hidden_params`` the provider transform set on the nested response are lost. Returns early - when that copy fell back to the original event, so logging-only state never lands on the - object the caller is iterating. + when the event was not a pydantic model and logging got the original, so logging-only state + never lands on the object the caller is iterating. """ if logging_response is self.completed_response: return @@ -1294,6 +1283,26 @@ def _add_text_like_part_events( ) +def _logging_copy(event: object) -> object: + """Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never + reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the + deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow + copies of the event and its nested response still keep the caller's ``usage`` attribute separate.""" + if not isinstance(event, BaseModel): + return event + try: + return type(event).model_validate(event.model_dump()) + except Exception: + return _detached_shallow_copy(event) + + +def _detached_shallow_copy(event: BaseModel) -> BaseModel: + nested: Final[object] = getattr(event, "response", None) + if isinstance(nested, BaseModel): + return event.model_copy(update={"response": nested.model_copy()}) + return event.model_copy() + + def _usage_as_model(usage: object) -> ResponseAPIUsage | None: if isinstance(usage, ResponseAPIUsage): return usage diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 60ce08666b6..5e0e794d93e 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -18,6 +18,7 @@ from litellm.responses.streaming_iterator import ( SyncResponsesAPIStreamingIterator, ) from litellm.types.llms.openai import ( + ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, @@ -329,8 +330,6 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): def _responses_api_response_with_usage() -> ResponsesAPIResponse: - from litellm.types.llms.openai import ResponseAPIUsage - return ResponsesAPIResponse( id="resp_lit6427", created_at=int(datetime(2025, 1, 1).timestamp()), @@ -383,8 +382,6 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage( { "input_tokens": 29, @@ -406,8 +403,6 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): from litellm.responses.streaming_iterator import _stamp_responses_usage_cost - from litellm.types.llms.openai import ResponseAPIUsage - response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj._response_cost_calculator.return_value = 0.000704 @@ -586,5 +581,50 @@ async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): iterator._log_completed_response(is_async=True) - assert logged == [iterator.completed_response] + assert len(logged) == 1 + assert logged[0] is not iterator.completed_response + assert logged[0].response is not iterator.completed_response.response + assert logged[0].response._hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" assert iterator.completed_response.response._hidden_params == {} + + +def _unvalidated_completed_config() -> Mock: + """Config whose completed event carries a Perplexity-style response that fails validation + (``truncation: ""``) and already holds the stamped ``ResponseAPIUsage``.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + response = _unvalidated_response_with_dict_usage( + ResponseAPIUsage(input_tokens=29, output_tokens=373, total_tokens=402, cost={"total_cost": 0.0001}) + ) + return ResponseCompletedEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_validation(): + """LIT-7391: the logging copy cannot round-trip a response that fails validation, and logging + rewrites the assembled response's usage to chat shape in place, so the event handed to logging + must never be the one the caller receives.""" + logging_obj = _logging_obj_stub() + logging_obj.stream = True + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator(headers={}, config=_unvalidated_completed_config(), logging_obj=logging_obj) + events = [event async for event in iterator] + + assert len(logged) == 1 + now = datetime.now() + LiteLLMLoggingObj._get_assembled_streaming_response( + logging_obj, logged[0], start_time=now, end_time=now, is_async=True, streaming_chunks=[] + ) + assert logged[0].response.usage["prompt_tokens"] == 29 + + client_usage = events[-1].response.usage + assert isinstance(client_usage, ResponseAPIUsage) + assert client_usage.input_tokens == 29 + assert client_usage.cost == pytest.approx(0.0001)