(test) bedrock - embedding with strings

This commit is contained in:
ishaan-jaff 2024-01-11 23:04:41 +05:30
parent a9d812eb8d
commit 276d11946e

View file

@ -186,10 +186,12 @@ def test_cohere_embedding3():
def test_bedrock_embedding_titan(): def test_bedrock_embedding_titan():
try: try:
litellm.set_verbose = True # this tests if we support str input for bedrock embedding
litellm.set_verbose = False
# DO NOT MAKE THE INPUT A LIST in this test
response = embedding( response = embedding(
model="amazon.titan-embed-text-v1", model="amazon.titan-embed-text-v1",
input="good morning from litellm, attempting to embed data", input="good morning from litellm, attempting to embed data", # input should always be a string in this test
) )
print(f"response:", response) print(f"response:", response)
assert isinstance( assert isinstance(
@ -199,6 +201,8 @@ def test_bedrock_embedding_titan():
assert all( assert all(
isinstance(x, float) for x in response["data"][0]["embedding"] isinstance(x, float) for x in response["data"][0]["embedding"]
), "Expected response to be a list of floats" ), "Expected response to be a list of floats"
print("Response Usage", response.usage)
assert response.usage.prompt_tokens == 11
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
@ -277,7 +281,7 @@ def test_aembedding():
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_aembedding() # test_aembedding()
def test_aembedding_azure(): def test_aembedding_azure():