test: fix test

2024-03-30 20:22:48 -07:00 · 2024-03-30 20:22:48 -07:00 · aebb0e489c
commit aebb0e489c
parent 17cabf013c
1 changed files with 6 additions and 1 deletions
--- a/litellm/tests/test_custom_callback_input.py
+++ b/litellm/tests/test_custom_callback_input.py
@ -253,7 +253,12 @@ class CompletionCustomHandler(
            assert isinstance(end_time, datetime)
            ## RESPONSE OBJECT
            assert isinstance(
-                response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)
+                response_obj,
+                (
+                    litellm.ModelResponse,
+                    litellm.EmbeddingResponse,
+                    litellm.TextCompletionResponse,
+                ),
            )
            ## KWARGS
            assert isinstance(kwargs["model"], str)