From fb958e524475b01aee75ddae82eb8d156640602c Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 22 Jan 2024 20:39:14 -0800
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 532fac5d79..556628d828 100644
--- a/litellm/tests/test_custom_callback_input.py
+++ b/litellm/tests/test_custom_callback_input.py
@@ -152,11 +152,11 @@ class CompletionCustomHandler(
             ## RESPONSE OBJECT
             assert isinstance(
                 response_obj,
-                Union[
+                (
                     litellm.ModelResponse,
                     litellm.EmbeddingResponse,
                     litellm.ImageResponse,
-                ],
+                ),
             )
             ## KWARGS
             assert isinstance(kwargs["model"], str)