From 56a9d3dfb6c8d3a4befba3d20618ed60d42dec16 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 11 Dec 2023 16:25:27 -0800 Subject: [PATCH] (test) custom input callback test --- litellm/tests/test_custom_callback_input.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index de59c68ba9..72f2aa6011 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -46,7 +46,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) except Exception as e: @@ -67,7 +67,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert isinstance(kwargs['input'], (list, dict, str)) @@ -93,7 +93,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) @@ -119,7 +119,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) @@ -145,7 +145,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert (isinstance(kwargs['input'], list) and isinstance(kwargs['input'][0], dict)) or isinstance(kwargs['input'], (dict, str)) @@ -169,7 +169,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) except Exception as e: @@ -190,7 +190,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert isinstance(kwargs['input'], (list, dict, str)) @@ -216,7 +216,7 @@ 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'], Union[datetime, None]) + assert isinstance(kwargs['start_time'], (datetime, type(None))) assert isinstance(kwargs['stream'], bool) assert isinstance(kwargs['user'], Union[str, None]) assert isinstance(kwargs['input'], (list, str, dict))