From dd9e79adbd45f583d764dede7c3a3cdf0707eb82 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 17 Mar 2025 17:26:21 -0700
Subject: [PATCH] fix(streaming_handler.py): emit deep copy of completed chunk

---
 litellm/litellm_core_utils/logging_utils.py   |  3 +++
 .../litellm_core_utils/streaming_handler.py   | 23 +++++++++++++------
 tests/local_testing/test_caching.py           |  4 ++--
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py
index 3c934a4276..c7512ea146 100644
--- a/litellm/litellm_core_utils/logging_utils.py
+++ b/litellm/litellm_core_utils/logging_utils.py
@@ -78,6 +78,9 @@ def _assemble_complete_response_from_streaming_chunks(
         Union[ModelResponse, TextCompletionResponse]
     ] = None
 
+    if isinstance(result, ModelResponse):
+        return result
+
     if result.choices[0].finish_reason is not None:  # if it's the last chunk
         streaming_chunks.append(result)
         try:
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 15d94b31a9..8fc63db5eb 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -1472,6 +1472,15 @@ class CustomStreamWrapper:
         """
         self.logging_loop = loop
 
+    def cache_streaming_response(self, processed_chunk, cache_hit: bool):
+        """
+        Caches the streaming response
+        """
+        if not cache_hit and self.logging_obj._llm_caching_handler is not None:
+            self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache(
+                processed_chunk
+            )
+
     def run_success_logging_and_cache_storage(self, processed_chunk, cache_hit: bool):
         """
         Runs success logging in a thread and adds the response to the cache
@@ -1503,12 +1512,6 @@ class CustomStreamWrapper:
         ## SYNC LOGGING
         self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
 
-        ## Sync store in cache
-        if self.logging_obj._llm_caching_handler is not None:
-            self.logging_obj._llm_caching_handler._sync_add_streaming_response_to_cache(
-                processed_chunk
-            )
-
     def finish_reason_handler(self):
         model_response = self.model_response_creator()
         _finish_reason = self.received_finish_reason or self.intermittent_finish_reason
@@ -1603,9 +1606,15 @@ class CustomStreamWrapper:
                         "usage",
                         getattr(complete_streaming_response, "usage"),
                     )
+                    self.cache_streaming_response(
+                        processed_chunk=complete_streaming_response.model_copy(
+                            deep=True
+                        ),
+                        cache_hit=cache_hit,
+                    )
                     executor.submit(
                         self.logging_obj.success_handler,
-                        complete_streaming_response,
+                        complete_streaming_response.model_copy(deep=True),
                         None,
                         None,
                         cache_hit,
diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py
index df2afdc167..7c6e400c05 100644
--- a/tests/local_testing/test_caching.py
+++ b/tests/local_testing/test_caching.py
@@ -2151,7 +2151,7 @@ def test_logging_turn_off_message_logging_streaming():
     mock_obj = Cache(type="local")
     litellm.cache = mock_obj
 
-    with patch.object(mock_obj, "add_cache", new=MagicMock()) as mock_client:
+    with patch.object(mock_obj, "add_cache") as mock_client:
         print(f"mock_obj.add_cache: {mock_obj.add_cache}")
 
         resp = litellm.completion(
@@ -2167,7 +2167,7 @@ def test_logging_turn_off_message_logging_streaming():
         time.sleep(1)
 
         mock_client.assert_called_once()
-
+        print(f"mock_client.call_args: {mock_client.call_args}")
         assert mock_client.call_args.args[0].choices[0].message.content == "hello"