test: fix test

2026-07-14 03:04:12 +00:00 · 2024-03-30 20:22:48 -07:00
parent 17cabf013c
commit aebb0e489c
1 changed files with 6 additions and 1 deletions
@@ -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)