From 9cbdbda8e4e349847c1eaebc1a8af1e7e4b85a13 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 16 Feb 2024 16:29:26 -0800 Subject: [PATCH] test(test_custom_callback_input.py): fix image gen callback test --- litellm/tests/test_custom_callback_input.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 062f78a76..579fe6583 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -176,7 +176,8 @@ class CompletionCustomHandler( ) or isinstance(kwargs["input"], (dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + kwargs["original_response"], + (str, litellm.CustomStreamWrapper, BaseModel), ) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str)