phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_embeddings.py): improve testing for titan embeddings

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-06 22:46:45 -08:00
parent c1e95740b0
commit 6fae3fde22
1 changed files with 4 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/tests/test_embedding.py
View file

@ -159,6 +159,9 @@ def test_bedrock_embedding_titan():
"lets test a second string for good measure"] "lets test a second string for good measure"]
) )
print(f"response:", response) 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: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_bedrock_embedding_titan() # test_bedrock_embedding_titan()
@ -177,7 +180,7 @@ def test_bedrock_embedding_cohere():
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_bedrock_embedding_cohere() # test_bedrock_embedding_cohere()
# comment out hf tests - since hf endpoints are unstable # comment out hf tests - since hf endpoints are unstable
def test_hf_embedding(): def test_hf_embedding():