diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 2e1e1a44594..9637eeb8acb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -368,18 +368,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 = self._response_for_success_logging() end_time: Final = datetime.now() if is_async: @@ -409,6 +398,32 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) + def _response_for_success_logging(self) -> ResponsesAPIStreamingResponse | ResponsesAPIResponse | None: + """A copy of the completed response, so the handlers' usage transforms (Responses API + input_tokens/output_tokens to chat completion prompt_tokens/completion_tokens) don't touch + what the caller gets back. model_dump + model_validate instead of deepcopy avoids pickle + errors with Pydantic ValidatorIterator on tool_choice with allowed_tools (fixes #17192). + + The handlers only unwrap the completion event in their assembled-stream branch, which a + non-streaming caller draining this iterator never reaches, so unwrap it here instead.""" + completed: Final = self.completed_response + copied: Final = self._model_copy_for_logging(completed) + unwrapped: Final = getattr(copied, "response", None) + if getattr(self.logging_obj, "stream", None) is not True and isinstance(unwrapped, ResponsesAPIResponse): + return unwrapped + return copied + + @staticmethod + def _model_copy_for_logging( + completed: ResponsesAPIStreamingResponse | None, + ) -> ResponsesAPIStreamingResponse | None: + if completed is None or not hasattr(completed, "model_dump"): + return completed + try: + return type(completed).model_validate(completed.model_dump()) + except Exception: + return completed + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index da0f608fdb5..5efe4341a80 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -192,7 +192,7 @@ "limit": 0 }, "S110": { - "limit": 218 + "limit": 217 }, "S112": { "limit": 22 diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 5b0f40fdf27..91194e420c3 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -235,3 +235,24 @@ def test_sync_transport_error_before_completed_event_raises(): with pytest.raises(httpx.ReadError): for _ in iterator: pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_completed_event_is_unwrapped_only_for_non_streaming_callers(stream: bool): + """A non-streaming caller can still drain this iterator (chat -> responses bridge + against a provider that always answers with SSE). The success handlers only unwrap + the completion event when logging in stream mode, so a wrapped event there yields no + standard_logging_object and the SpendLogs row is dropped (#36426).""" + logging_obj = _logging_obj_stub() + logging_obj.stream = stream + + iterator = _make_iterator(sse_events=_COMPLETE_STREAM_EVENTS, logging_obj=logging_obj) + async for _ in iterator: + pass + + logged = logging_obj.dispatch_success_handlers.call_args.args[0] + if stream: + assert logged.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + else: + assert isinstance(logged, ResponsesAPIResponse) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index c990ae52ff2..4016b04a07c 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16744 + "limit": 16742 }, "LIT011": { "limit": 5596