diff --git a/litellm/tests/test_custom_logger.py b/litellm/tests/test_custom_logger.py index c50f2b1ee..26dcdf7d4 100644 --- a/litellm/tests/test_custom_logger.py +++ b/litellm/tests/test_custom_logger.py @@ -78,6 +78,7 @@ class TmpFunction: async def async_test_logging_fn(self, kwargs, completion_obj, start_time, end_time): print(f"ON ASYNC LOGGING") self.async_success = True + print(f'kwargs.get("complete_streaming_response"): {kwargs.get("complete_streaming_response")}') self.complete_streaming_response_in_callback = kwargs.get("complete_streaming_response") @@ -102,7 +103,8 @@ def test_async_chat_openai_stream(): complete_streaming_response = complete_streaming_response.strip("'") response1 = tmp_function.complete_streaming_response_in_callback["choices"][0]["message"]["content"] response2 = complete_streaming_response - assert [ord(c) for c in response1] == [ord(c) for c in response2] + # assert [ord(c) for c in response1] == [ord(c) for c in response2] + assert response1 == response2 assert tmp_function.async_success == True except Exception as e: print(e) diff --git a/litellm/utils.py b/litellm/utils.py index 173abe8ab..a8f469e5c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -830,9 +830,9 @@ class Logging: complete_streaming_response = None if self.stream == True and self.model_call_details.get("litellm_params", {}).get("acompletion", False) == True: # if it's acompletion == True, chunks are built/appended in async_success_handler - self.streaming_chunks.append(result) if result.choices[0].finish_reason is not None: # if it's the last chunk - complete_streaming_response = litellm.stream_chunk_builder(self.streaming_chunks, messages=self.model_call_details.get("messages", None)) + streaming_chunks = self.streaming_chunks + [result] + complete_streaming_response = litellm.stream_chunk_builder(streaming_chunks, messages=self.model_call_details.get("messages", None)) else: # this is a completion() call if self.stream == True: @@ -1053,7 +1053,7 @@ class Logging: if self.stream: if result.choices[0].finish_reason is not None: # if it's the last chunk self.streaming_chunks.append(result) - print_verbose(f"final set of received chunks: {self.streaming_chunks}") + # print_verbose(f"final set of received chunks: {self.streaming_chunks}") try: complete_streaming_response = litellm.stream_chunk_builder(self.streaming_chunks, messages=self.model_call_details.get("messages", None)) except: