diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..062abcfd142 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -412,13 +412,7 @@ class BaseResponsesAPIStreamingIterator: # 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 = self._response_for_success_logging() end_time: Final = datetime.now() if is_async: @@ -448,6 +442,29 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) + def _response_for_success_logging( + self, + ) -> ResponsesAPIStreamingResponse | ResponsesAPIResponse | None: + """Build the response object to pass to success handlers. + + ``self.completed_response`` is a ``ResponseCompletedEvent`` wrapper. + The success handlers only unwrap it in their assembled-stream branch, + which a non-streaming caller draining this iterator never reaches. + Unwrap it here so the inner ``ResponsesAPIResponse`` reaches the + handlers and ``standard_logging_object`` is built for cost tracking. + """ + completed = self.completed_response + if completed is None or not hasattr(completed, "model_dump"): + return completed + try: + copied = type(completed).model_validate(completed.model_dump()) + except Exception: + return completed + unwrapped = getattr(copied, "response", None) + if isinstance(unwrapped, ResponsesAPIResponse): + return unwrapped + return copied + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 38407c94fe7..d322dab6d6d 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -305,3 +305,62 @@ def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypat asyncio.run(_short_lived_script()) assert len(writes) == 1 +async def test_completed_response_unwrapped_for_success_logging(): + """The iterator must pass the unwrapped ResponsesAPIResponse (not the + ResponseCompletedEvent wrapper) to dispatch_success_handlers so that + _success_handler_helper_fn recognises it and builds + standard_logging_object for cost tracking. + + Without this fix, non-streaming callers that drain a force-streaming + provider (e.g. the Anthropic /v1/messages -> Responses API bridge) get + ``standard_logging_object not found`` because the wrapper event is not + recognised by _is_recognized_call_type_for_logging in the non-streaming + branch of _success_handler_helper_fn. + """ + logging_obj = _logging_obj_stub() + dispatched_args: list = [] + + async def _dispatch(result, **kwargs): + dispatched_args.append(result) + + logging_obj.dispatch_success_handlers.side_effect = _dispatch + + iterator = _make_iterator( + sse_events=_COMPLETE_STREAM_EVENTS, + logging_obj=logging_obj, + ) + + async for _ in iterator: + pass + + assert len(dispatched_args) == 1, ( + f"Expected dispatch_success_handlers called once; got {len(dispatched_args)}" + ) + logged = dispatched_args[0] + assert isinstance(logged, ResponsesAPIResponse), ( + f"Expected ResponsesAPIResponse unwrapped for logging; got {type(logged).__name__}" + ) + + +def test_sync_completed_response_unwrapped_for_success_logging(): + """Sync counterpart of the unwrapping test.""" + logging_obj = _logging_obj_stub() + success_args: list = [] + + logging_obj.async_success_handler.side_effect = lambda result, **kw: success_args.append(result) + + iterator = _make_sync_iterator( + sse_events=_COMPLETE_STREAM_EVENTS, + logging_obj=logging_obj, + ) + + for _ in iterator: + pass + + assert len(success_args) == 1, ( + f"Expected async_success_handler called once; got {len(success_args)}" + ) + logged = success_args[0] + assert isinstance(logged, ResponsesAPIResponse), ( + f"Expected ResponsesAPIResponse unwrapped for logging; got {type(logged).__name__}" + )