test(test_custom_callback_input.py): fix assert

This commit is contained in:
Krrish Dholakia 2024-01-22 20:39:14 -08:00
parent 579dfc3013
commit fb958e5244

View file

@ -152,11 +152,11 @@ class CompletionCustomHandler(
## RESPONSE OBJECT ## RESPONSE OBJECT
assert isinstance( assert isinstance(
response_obj, response_obj,
Union[ (
litellm.ModelResponse, litellm.ModelResponse,
litellm.EmbeddingResponse, litellm.EmbeddingResponse,
litellm.ImageResponse, litellm.ImageResponse,
], ),
) )
## KWARGS ## KWARGS
assert isinstance(kwargs["model"], str) assert isinstance(kwargs["model"], str)