diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d6d3e2576f6..8d766cf1cd0 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -450,11 +450,7 @@ class BaseResponsesAPIStreamingIterator: _terminal_chunk: Final = ( openai_responses_api_chunk if _billed_response is None or _billed_response is _response_obj - else ( - openai_responses_api_chunk.model_copy(update={"response": _billed_response}) - if issubclass(type(openai_responses_api_chunk), BaseModel) # pyright: ignore[reportUnnecessaryIsInstance] # test stubs use spec'd Mocks whose __class__ reports BaseModel but whose model_copy returns a Mock - else _replace_response(openai_responses_api_chunk, _billed_response) - ) + else openai_responses_api_chunk.model_copy(update={"response": _billed_response}) ) self.completed_response = _terminal_chunk _stamp_responses_usage_cost(_billed_response, self.logging_obj) @@ -1359,13 +1355,6 @@ def _billed_terminal_response( ) -def _replace_response( - event: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse -) -> ResponsesAPIStreamingResponse: - setattr(event, "response", response) - return event - - def _dump_json_safely(response: BaseModel) -> str | None: try: return response.model_dump_json() diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index b33cc4e93c2..dbf54ec3b9b 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -47,10 +47,10 @@ def _mock_config() -> Mock: def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") if evt_type == "response.completed": - completed = Mock(spec=ResponseCompletedEvent) - completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED - completed.response = mock_responses_api_response - return completed + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=mock_responses_api_response, + ) stub = Mock() stub.type = evt_type return stub @@ -348,10 +348,10 @@ def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") if evt_type == "response.completed": - completed = Mock(spec=ResponseCompletedEvent) - completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED - completed.response = response - return completed + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) stub = Mock() stub.type = evt_type if "delta" in parsed_chunk: