From a4b0740c4aa2aaffc406a3bd4baf698b44650e90 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 15 Jul 2024 21:46:51 -0700 Subject: [PATCH] fix(test_custom_callback_input.py): fix unit test for embedding logging --- litellm/tests/test_custom_callback_input.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 8514b237c..eae0412d3 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -198,7 +198,10 @@ class CompletionCustomHandler( assert isinstance(kwargs["user"], (str, type(None))) assert ( isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) + and ( + isinstance(kwargs["input"][0], dict) + or isinstance(kwargs["input"][0], str) + ) ) or isinstance(kwargs["input"], (dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert isinstance(