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

2026-08-02 18:21:56 +00:00 · 2024-02-16 16:29:26 -08:00
parent f57483ea70
commit 9cbdbda8e4
1 changed files with 2 additions and 1 deletions
@@ -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)