From 6fae3fde22437c8da7d717c9cbabe91fb84bdacf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 6 Dec 2023 22:46:45 -0800 Subject: [PATCH] test(test_embeddings.py): improve testing for titan embeddings --- litellm/tests/test_embedding.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 5361d1927..71e59819f 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -159,6 +159,9 @@ def test_bedrock_embedding_titan(): "lets test a second string for good measure"] ) print(f"response:", response) + assert isinstance(response['data'][0]['embedding'], list), "Expected response to be a list" + print(f"type of first embedding:", type(response['data'][0]['embedding'][0])) + assert all(isinstance(x, float) for x in response['data'][0]['embedding']), "Expected response to be a list of floats" except Exception as e: pytest.fail(f"Error occurred: {e}") # test_bedrock_embedding_titan() @@ -177,7 +180,7 @@ def test_bedrock_embedding_cohere(): except Exception as e: pytest.fail(f"Error occurred: {e}") -test_bedrock_embedding_cohere() +# test_bedrock_embedding_cohere() # comment out hf tests - since hf endpoints are unstable def test_hf_embedding():