(test) fix embedding tests

This commit is contained in:
ishaan-jaff 2023-11-24 17:04:59 -08:00
parent 23466107a7
commit ad1672b901

View file

@ -20,7 +20,6 @@ def test_openai_embedding():
model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"] model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"]
) )
litellm_response = dict(response) litellm_response = dict(response)
litellm_response.pop("_response_ms")
litellm_response_keys = set(litellm_response.keys()) litellm_response_keys = set(litellm_response.keys())
print(litellm_response_keys) print(litellm_response_keys)
print("LiteLLM Response\n") print("LiteLLM Response\n")
@ -40,7 +39,7 @@ def test_openai_embedding():
print(openai_response_keys) print(openai_response_keys)
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_openai_embedding() test_openai_embedding()
def test_openai_azure_embedding_simple(): def test_openai_azure_embedding_simple():
try: try:
@ -50,12 +49,12 @@ def test_openai_azure_embedding_simple():
) )
print(response) print(response)
response_keys = dict(response).keys() response_keys = dict(response).keys()
assert set(["usage", "model", "object", "data", "_response_ms"]) == set(response_keys) #assert litellm response has expected keys from OpenAI embedding response assert set(["usage", "model", "object", "data"]) == set(response_keys) #assert litellm response has expected keys from OpenAI embedding response
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_openai_azure_embedding_simple() test_openai_azure_embedding_simple()
def test_openai_azure_embedding_timeouts(): def test_openai_azure_embedding_timeouts():
@ -87,7 +86,7 @@ def test_openai_embedding_timeouts():
pass pass
except Exception as e: except Exception as e:
pytest.fail(f"Expected timeout error, did not get the correct error. Instead got {e}") pytest.fail(f"Expected timeout error, did not get the correct error. Instead got {e}")
test_openai_embedding_timeouts() # test_openai_embedding_timeouts()
def test_openai_azure_embedding(): def test_openai_azure_embedding():
try: try:
@ -184,5 +183,4 @@ def test_aembedding():
pass pass
asyncio.run(embedding_call()) asyncio.run(embedding_call())
# test_aembedding() test_aembedding()