From 40dd38508f8139115432c358dfc4ef3479c83822 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 22 Nov 2023 13:50:44 -0800 Subject: [PATCH] (test) embedding stricter testing --- litellm/tests/test_embedding.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 546c3dd59..deff80632 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -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: