diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 04e64a70d..865ed4040 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -48,7 +48,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) except Exception as e: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -69,11 +69,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert isinstance(kwargs['input'], (list, dict, str)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(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'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}") @@ -95,11 +95,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(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'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}") @@ -121,11 +121,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(None))) assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) - assert isinstance(kwargs['additional_args'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}") @@ -147,11 +147,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(None))) assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None - assert isinstance(kwargs['additional_args'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}") @@ -171,7 +171,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) except Exception as e: print(f"Assertion Error: {traceback.format_exc()}") self.errors.append(traceback.format_exc()) @@ -192,11 +192,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert isinstance(kwargs['input'], (list, dict, str)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(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'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}") @@ -218,11 +218,11 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse assert isinstance(kwargs['litellm_params'], dict) assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) - assert isinstance(kwargs['user'], (str, None)) + assert isinstance(kwargs['user'], (str, type(None))) assert isinstance(kwargs['input'], (list, str, dict)) - assert isinstance(kwargs['api_key'], (str, None)) + assert isinstance(kwargs['api_key'], (str, type(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'], (dict, None)) + assert isinstance(kwargs['additional_args'], (dict, type(None))) assert isinstance(kwargs['log_event_type'], str) except: print(f"Assertion Error: {traceback.format_exc()}")