diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py
index d9bcc89478..de59c68ba9 100644
--- a/litellm/tests/test_custom_callback_input.py
+++ b/litellm/tests/test_custom_callback_input.py
@@ -4,7 +4,7 @@ import sys, os, time, inspect, asyncio, traceback
 from datetime import datetime
 import pytest
 sys.path.insert(0, os.path.abspath('../..'))
-from typing import Optional, Literal, List
+from typing import Optional, Literal, List, Union
 from litellm import completion, embedding
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
@@ -46,9 +46,9 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list) 
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
         except Exception as e: 
             print(f"Assertion Error: {traceback.format_exc()}")
             self.errors.append(traceback.format_exc())
@@ -67,13 +67,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert isinstance(kwargs['input'], (list, dict, str))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs['original_response']) or inspect.isasyncgen(kwargs['original_response'])
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")
@@ -93,13 +93,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list) and isinstance(kwargs['messages'][0], dict)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs['original_response']) or inspect.iscoroutine(kwargs['original_response'])
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")
@@ -119,13 +119,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list) and isinstance(kwargs['messages'][0], dict)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper))
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except:
             print(f"Assertion Error: {traceback.format_exc()}")
@@ -145,13 +145,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list) and isinstance(kwargs['messages'][0], dict)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")
@@ -169,9 +169,9 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list) and isinstance(kwargs['messages'][0], dict)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
         except Exception as e: 
             print(f"Assertion Error: {traceback.format_exc()}")
             self.errors.append(traceback.format_exc())
@@ -190,13 +190,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert isinstance(kwargs['input'], (list, dict, str))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs['original_response']) or inspect.iscoroutine(kwargs['original_response'])
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")
@@ -216,13 +216,13 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['messages'], list)
             assert isinstance(kwargs['optional_params'], dict)
             assert isinstance(kwargs['litellm_params'], dict) 
-            assert isinstance(kwargs['start_time'], Optional[datetime])
+            assert isinstance(kwargs['start_time'], Union[datetime, None])
             assert isinstance(kwargs['stream'], bool)
-            assert isinstance(kwargs['user'], Optional[str])
+            assert isinstance(kwargs['user'], Union[str, None])
             assert isinstance(kwargs['input'], (list, str, dict))
-            assert isinstance(kwargs['api_key'], Optional[str])
+            assert isinstance(kwargs['api_key'], Union[str, None])
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs['original_response']) or kwargs['original_response'] == None
-            assert isinstance(kwargs['additional_args'], Optional[dict])
+            assert isinstance(kwargs['additional_args'], Union[dict, None])
             assert isinstance(kwargs['log_event_type'], str) 
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")