From 6234982fd92c78e63c7e815b3bcf603f0e917fa8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 6 Apr 2024 18:48:33 -0700 Subject: [PATCH] test(test_custom_callback_input.py): fix assert --- litellm/tests/test_custom_callback_input.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 7175f5bc4..0fd2cfabb 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -167,7 +167,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["litellm_params"]["api_base"], str) - assert isinstance(kwargs["cache_hit"], Optional[bool]) + assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) @@ -269,7 +269,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["api_base"], str) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) - assert isinstance(kwargs["cache_hit"], Optional[bool]) + assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) assert isinstance(kwargs["user"], (str, type(None))) assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None)))