(test) embedding stricter testing

This commit is contained in:
ishaan-jaff 2023-11-22 13:50:44 -08:00
parent e8ff4d5eca
commit 40dd38508f

View file

@ -1,6 +1,9 @@
import sys, os
import traceback
import pytest
from dotenv import load_dotenv
load_dotenv()
sys.path.insert(
0, os.path.abspath("../..")
@ -15,12 +18,27 @@ def test_openai_embedding():
response = embedding(
model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"]
)
print(response)
# Add any assertions here to check the response
# print(f"response: {str(response)}")
litellm_response = dict(response)
litellm_response_keys = set(litellm_response.keys())
print(litellm_response_keys)
print("LiteLLM Response\n")
print(litellm_response)
# same request with OpenAI 1.0+
import openai
client = openai.OpenAI(api_key=os.environ['OPENAI_API_KEY'])
response = client.embeddings.create(
model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"]
)
response = dict(response)
openai_response_keys = set(response.keys())
assert litellm_response_keys == openai_response_keys # ENSURE the Keys in litellm response is exactly what the openai package returns
assert len(litellm_response["data"]) == 2 # expect two embedding responses from litellm_response since input had two
print(openai_response_keys)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_openai_embedding()
test_openai_embedding()
def test_openai_azure_embedding_simple():
try: