phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_custom_callback_input.py): fix image gen callback test

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-16 16:29:26 -08:00
parent f57483ea70
commit 9cbdbda8e4
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/tests/test_custom_callback_input.py
View file

@ -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)