fix(utils.py): fix sync/async stream logging

litellm/tests/test_custom_logger.py

@@ -1,14 +1,17 @@
 ### What this tests ####
 import sys, os, time, inspect, asyncio, traceback
 import pytest
-sys.path.insert(0, os.path.abspath('../..'))
+
+sys.path.insert(0, os.path.abspath("../.."))
 
 from litellm import completion, embedding
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
 
+
 class MyCustomHandler(CustomLogger):
     complete_streaming_response_in_callback = ""
+
     def __init__(self):
         self.success: bool = False  # type: ignore
         self.failure: bool = False  # type: ignore
@@ -45,7 +48,6 @@ class MyCustomHandler(CustomLogger):
         if kwargs.get("stream") == True:
             self.sync_stream_collected_response = response_obj
 
-
     def log_failure_event(self, kwargs, response_obj, start_time, end_time):
         print(f"On Failure")
         self.failure = True
@@ -72,14 +74,20 @@ class MyCustomHandler(CustomLogger):
 
         self.async_completion_kwargs_fail = kwargs
 
+
 class TmpFunction:
     complete_streaming_response_in_callback = ""
     async_success: bool = False
+
     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")}')
+        print(
-        self.complete_streaming_response_in_callback = kwargs.get("complete_streaming_response")
+            f'kwargs.get("complete_streaming_response"): {kwargs.get("complete_streaming_response")}'
+        )
+        self.complete_streaming_response_in_callback = kwargs.get(
+            "complete_streaming_response"
+        )
 
 
 def test_async_chat_openai_stream():
@@ -88,29 +96,39 @@ def test_async_chat_openai_stream():
         # litellm.set_verbose = True
         litellm.success_callback = [tmp_function.async_test_logging_fn]
         complete_streaming_response = ""
+
         async def call_gpt():
             nonlocal complete_streaming_response
-            response = await litellm.acompletion(model="gpt-3.5-turbo",
+            response = await litellm.acompletion(
-                                messages=[{
+                model="gpt-3.5-turbo",
-                                    "role": "user",
+                messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
-                                    "content": "Hi 👋 - i'm openai"
+                stream=True,
-                                }],
+            )
-                                stream=True)
             async for chunk in response:
-                complete_streaming_response += chunk["choices"][0]["delta"]["content"] or ""
+                complete_streaming_response += (
+                    chunk["choices"][0]["delta"]["content"] or ""
+                )
                 print(complete_streaming_response)
+
         asyncio.run(call_gpt())
         complete_streaming_response = complete_streaming_response.strip("'")
-        response1 = tmp_function.complete_streaming_response_in_callback["choices"][0]["message"]["content"]
+        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]
+        print(f"response1: {response1}")
+        print(f"response2: {response2}")
         assert response1 == response2
         assert tmp_function.async_success == True
     except Exception as e:
         print(e)
         pytest.fail(f"An error occurred - {str(e)}")
+
+
 # test_async_chat_openai_stream()
 
+
 def test_completion_azure_stream_moderation_failure():
     try:
         customHandler = MyCustomHandler()
@@ -152,27 +170,32 @@ def test_async_custom_handler_stream():
             },
         ]
         complete_streaming_response = ""
+
         async def test_1():
             nonlocal complete_streaming_response
             response = await litellm.acompletion(
-                model="azure/chatgpt-v-2", 
+                model="azure/chatgpt-v-2", messages=messages, stream=True
-                messages=messages,
-                stream=True
             )
             async for chunk in response:
-                complete_streaming_response += chunk["choices"][0]["delta"]["content"] or ""
+                complete_streaming_response += (
+                    chunk["choices"][0]["delta"]["content"] or ""
+                )
                 print(complete_streaming_response)
 
         asyncio.run(test_1())
 
         response_in_success_handler = customHandler2.stream_collected_response
-        response_in_success_handler = response_in_success_handler["choices"][0]["message"]["content"]
+        response_in_success_handler = response_in_success_handler["choices"][0][
+            "message"
+        ]["content"]
         print("\n\n")
         print("response_in_success_handler: ", response_in_success_handler)
         print("complete_streaming_response: ", complete_streaming_response)
         assert response_in_success_handler == complete_streaming_response
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
+
+
 # test_async_custom_handler_stream()
 
 
@@ -194,9 +217,7 @@ def test_azure_completion_stream():
         complete_streaming_response = ""
 
         response = litellm.completion(
-            model="azure/chatgpt-v-2", 
+            model="azure/chatgpt-v-2", messages=messages, stream=True
-            messages=messages,
-            stream=True
         )
         for chunk in response:
             complete_streaming_response += chunk["choices"][0]["delta"]["content"] or ""
@@ -204,7 +225,9 @@ def test_azure_completion_stream():
 
         time.sleep(0.5)  # wait 1/2 second before checking callbacks
         response_in_success_handler = customHandler2.sync_stream_collected_response
-        response_in_success_handler = response_in_success_handler["choices"][0]["message"]["content"]
+        response_in_success_handler = response_in_success_handler["choices"][0][
+            "message"
+        ]["content"]
         print("\n\n")
         print("response_in_success_handler: ", response_in_success_handler)
         print("complete_streaming_response: ", complete_streaming_response)
@@ -212,6 +235,7 @@ def test_azure_completion_stream():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
+
 @pytest.mark.asyncio
 async def test_async_custom_handler_completion():
     try:
@@ -222,14 +246,21 @@ async def test_async_custom_handler_completion():
         litellm.callbacks = [customHandler_success]
         response = await litellm.acompletion(
             model="gpt-3.5-turbo",
-                messages=[{
+            messages=[
+                {
                     "role": "user",
                     "content": "hello from litellm test",
-                }]
+                }
+            ],
         )
         await asyncio.sleep(1)
-        assert customHandler_success.async_success == True, "async success is not set to True even after success"
+        assert (
-        assert customHandler_success.async_completion_kwargs.get("model") == "gpt-3.5-turbo"
+            customHandler_success.async_success == True
+        ), "async success is not set to True even after success"
+        assert (
+            customHandler_success.async_completion_kwargs.get("model")
+            == "gpt-3.5-turbo"
+        )
         # failure
         litellm.callbacks = [customHandler_failure]
         messages = [
@@ -249,15 +280,28 @@ async def test_async_custom_handler_completion():
             )
         except:
             pass
-        assert customHandler_failure.async_failure == True, "async failure is not set to True even after failure"        
+        assert (
-        assert customHandler_failure.async_completion_kwargs_fail.get("model") == "gpt-3.5-turbo"
+            customHandler_failure.async_failure == True
-        assert len(str(customHandler_failure.async_completion_kwargs_fail.get("exception"))) > 10 # expect APIError("OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: test. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}"), 'traceback_exception': 'Traceback (most recent call last):\n  File "/Users/ishaanjaffer/Github/litellm/litellm/llms/openai.py", line 269, in acompletion\n    response = await openai_aclient.chat.completions.create(**data)\n  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/openai/resources/chat/completions.py", line 119
+        ), "async failure is not set to True even after failure"
+        assert (
+            customHandler_failure.async_completion_kwargs_fail.get("model")
+            == "gpt-3.5-turbo"
+        )
+        assert (
+            len(
+                str(customHandler_failure.async_completion_kwargs_fail.get("exception"))
+            )
+            > 10
+        )  # expect APIError("OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: test. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}"), 'traceback_exception': 'Traceback (most recent call last):\n  File "/Users/ishaanjaffer/Github/litellm/litellm/llms/openai.py", line 269, in acompletion\n    response = await openai_aclient.chat.completions.create(**data)\n  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/openai/resources/chat/completions.py", line 119
         litellm.callbacks = []
         print("Passed setting async failure")
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
+
+
 # asyncio.run(test_async_custom_handler_completion())
 
+
 @pytest.mark.asyncio
 async def test_async_custom_handler_embedding():
     try:
@@ -267,30 +311,53 @@ async def test_async_custom_handler_embedding():
         assert customHandler_embedding.async_success_embedding == False
         response = await litellm.aembedding(
             model="text-embedding-ada-002",
-                input = ["hello world"],
+            input=["hello world"],
         )
         await asyncio.sleep(1)
-        assert customHandler_embedding.async_success_embedding == True, "async_success_embedding is not set to True even after success"
+        assert (
-        assert customHandler_embedding.async_embedding_kwargs.get("model") == "text-embedding-ada-002"
+            customHandler_embedding.async_success_embedding == True
-        assert customHandler_embedding.async_embedding_response["usage"]["prompt_tokens"] ==2
+        ), "async_success_embedding is not set to True even after success"
+        assert (
+            customHandler_embedding.async_embedding_kwargs.get("model")
+            == "text-embedding-ada-002"
+        )
+        assert (
+            customHandler_embedding.async_embedding_response["usage"]["prompt_tokens"]
+            == 2
+        )
         print("Passed setting async success: Embedding")
         # failure
         assert customHandler_embedding.async_failure_embedding == False
         try:
             response = await litellm.aembedding(
                 model="text-embedding-ada-002",
-                        input = ["hello world"],
+                input=["hello world"],
                 api_key="my-bad-key",
             )
         except:
             pass
-        assert customHandler_embedding.async_failure_embedding == True, "async failure embedding is not set to True even after failure"        
+        assert (
-        assert customHandler_embedding.async_embedding_kwargs_fail.get("model") == "text-embedding-ada-002"
+            customHandler_embedding.async_failure_embedding == True
-        assert len(str(customHandler_embedding.async_embedding_kwargs_fail.get("exception"))) > 10 # exppect APIError("OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: test. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}"), 'traceback_exception': 'Traceback (most recent call last):\n  File "/Users/ishaanjaffer/Github/litellm/litellm/llms/openai.py", line 269, in acompletion\n    response = await openai_aclient.chat.completions.create(**data)\n  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/openai/resources/chat/completions.py", line 119
+        ), "async failure embedding is not set to True even after failure"
+        assert (
+            customHandler_embedding.async_embedding_kwargs_fail.get("model")
+            == "text-embedding-ada-002"
+        )
+        assert (
+            len(
+                str(
+                    customHandler_embedding.async_embedding_kwargs_fail.get("exception")
+                )
+            )
+            > 10
+        )  # exppect APIError("OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: test. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}"), 'traceback_exception': 'Traceback (most recent call last):\n  File "/Users/ishaanjaffer/Github/litellm/litellm/llms/openai.py", line 269, in acompletion\n    response = await openai_aclient.chat.completions.create(**data)\n  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/openai/resources/chat/completions.py", line 119
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
+
+
 # asyncio.run(test_async_custom_handler_embedding())
 
+
 @pytest.mark.asyncio
 async def test_async_custom_handler_embedding_optional_param():
     """
@@ -300,16 +367,19 @@ async def test_async_custom_handler_embedding_optional_param():
     customHandler_optional_params = MyCustomHandler()
     litellm.callbacks = [customHandler_optional_params]
     response = await litellm.aembedding(
-                model="azure/azure-embedding-model", 
+        model="azure/azure-embedding-model", input=["hello world"], user="John"
-                input = ["hello world"],
-                user = "John"
     )
     await asyncio.sleep(1)  # success callback is async
     assert customHandler_optional_params.user == "John"
-    assert customHandler_optional_params.user == customHandler_optional_params.data_sent_to_api["user"]
+    assert (
+        customHandler_optional_params.user
+        == customHandler_optional_params.data_sent_to_api["user"]
+    )
+
 
 # asyncio.run(test_async_custom_handler_embedding_optional_param())
 
+
 @pytest.mark.asyncio
 async def test_async_custom_handler_embedding_optional_param_bedrock():
     """
@@ -323,9 +393,7 @@ async def test_async_custom_handler_embedding_optional_param_bedrock():
     customHandler_optional_params = MyCustomHandler()
     litellm.callbacks = [customHandler_optional_params]
     response = await litellm.aembedding(
-                model="bedrock/amazon.titan-embed-text-v1", 
+        model="bedrock/amazon.titan-embed-text-v1", input=["hello world"], user="John"
-                input = ["hello world"],
-                user = "John"
     )
     await asyncio.sleep(1)  # success callback is async
     assert customHandler_optional_params.user == "John"
@@ -334,16 +402,36 @@ async def test_async_custom_handler_embedding_optional_param_bedrock():
 
 def test_redis_cache_completion_stream():
     from litellm import Cache
+
     # Important Test - This tests if we can add to streaming cache, when custom callbacks are set
     import random
+
     try:
         print("\nrunning test_redis_cache_completion_stream")
         litellm.set_verbose = True
-        random_number = random.randint(1, 100000) # add a random number to ensure it's always adding / reading from cache
+        random_number = random.randint(
-        messages = [{"role": "user", "content": f"write a one sentence poem about: {random_number}"}]
+            1, 100000
-        litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
+        )  # add a random number to ensure it's always adding / reading from cache
+        messages = [
+            {
+                "role": "user",
+                "content": f"write a one sentence poem about: {random_number}",
+            }
+        ]
+        litellm.cache = Cache(
+            type="redis",
+            host=os.environ["REDIS_HOST"],
+            port=os.environ["REDIS_PORT"],
+            password=os.environ["REDIS_PASSWORD"],
+        )
         print("test for caching, streaming + completion")
-        response1 = completion(model="gpt-3.5-turbo", messages=messages, max_tokens=40, temperature=0.2, stream=True)
+        response1 = completion(
+            model="gpt-3.5-turbo",
+            messages=messages,
+            max_tokens=40,
+            temperature=0.2,
+            stream=True,
+        )
         response_1_content = ""
         for chunk in response1:
             print(chunk)
@@ -351,14 +439,22 @@ def test_redis_cache_completion_stream():
         print(response_1_content)
 
         time.sleep(0.1)  # sleep for 0.1 seconds allow set cache to occur
-        response2 = completion(model="gpt-3.5-turbo", messages=messages, max_tokens=40, temperature=0.2, stream=True)
+        response2 = completion(
+            model="gpt-3.5-turbo",
+            messages=messages,
+            max_tokens=40,
+            temperature=0.2,
+            stream=True,
+        )
         response_2_content = ""
         for chunk in response2:
             print(chunk)
             response_2_content += chunk.choices[0].delta.content or ""
         print("\nresponse 1", response_1_content)
         print("\nresponse 2", response_2_content)
-        assert response_1_content == response_2_content, f"Response 1 != Response 2. Same params, Response 1{response_1_content} != Response 2{response_2_content}"
+        assert (
+            response_1_content == response_2_content
+        ), f"Response 1 != Response 2. Same params, Response 1{response_1_content} != Response 2{response_2_content}"
         litellm.success_callback = []
         litellm._async_success_callback = []
         litellm.cache = None
@@ -366,4 +462,6 @@ def test_redis_cache_completion_stream():
         print(e)
         litellm.success_callback = []
         raise e
+
+
 # test_redis_cache_completion_stream()