From 72275ad8cb1a8fe7558726cacda6289a76cdb559 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 25 Jan 2024 15:59:53 -0800
Subject: [PATCH] fix(main.py): fix logging event loop for async logging but
 sync streaming

---
 litellm/main.py  | 12 ++++--------
 litellm/utils.py | 21 +++++++++++++++------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index 6b40354739d..89750ef46b4 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -274,14 +274,10 @@ async def acompletion(
         else:
             # Call the synchronous function using run_in_executor
             response = await loop.run_in_executor(None, func_with_context)  # type: ignore
-        # if kwargs.get("stream", False):  # return an async generator
-        #     return _async_streaming(
-        #         response=response,
-        #         model=model,
-        #         custom_llm_provider=custom_llm_provider,
-        #         args=args,
-        #     )
-        # else:
+        if isinstance(response, CustomStreamWrapper):
+            response.set_logging_event_loop(
+                loop=loop
+            )  # sets the logging event loop if the user does sync streaming (e.g. on proxy for sagemaker calls)
         return response
     except Exception as e:
         custom_llm_provider = custom_llm_provider or "openai"
diff --git a/litellm/utils.py b/litellm/utils.py
index 02ac83d0659..2bc1d34e9e7 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -7116,6 +7116,7 @@ class CustomStreamWrapper:
             "model_id": (_model_info.get("id", None))
         }  # returned as x-litellm-model-id response header in proxy
         self.response_id = None
+        self.logging_loop = None
 
     def __iter__(self):
         return self
@@ -8016,16 +8017,24 @@ class CustomStreamWrapper:
                 original_exception=e,
             )
 
+    def set_logging_event_loop(self, loop):
+        self.logging_loop = loop
+
+    async def your_async_function(self):
+        # Your asynchronous code here
+        return "Your asynchronous code is running"
+
     def run_success_logging_in_thread(self, processed_chunk):
         # Create an event loop for the new thread
         ## ASYNC LOGGING
-        # Run the asynchronous function in the new thread's event loop
-        asyncio.run(
-            self.logging_obj.async_success_handler(
-                processed_chunk,
+        if self.logging_loop is not None:
+            future = asyncio.run_coroutine_threadsafe(
+                self.logging_obj.async_success_handler(processed_chunk),
+                loop=self.logging_loop,
             )
-        )
-
+            result = future.result()
+        else:
+            asyncio.run(self.logging_obj.async_success_handler(processed_chunk))
         ## SYNC LOGGING
         self.logging_obj.success_handler(processed_chunk)