From 39a1b4c3b5b67239cd3a7dcdf6163bd4ea3bae1b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 22 Jan 2024 12:09:43 -0800 Subject: [PATCH] fix(main.py): support custom pricing for embedding calls --- litellm/tests/test_embedding.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index d1f0ee699..630b41d72 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -10,7 +10,7 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm -from litellm import embedding, completion +from litellm import embedding, completion, completion_cost litellm.set_verbose = False @@ -341,8 +341,30 @@ def test_sagemaker_embeddings(): response = litellm.embedding( model="sagemaker/berri-benchmarking-gpt-j-6b-fp16", input=["good morning from litellm", "this is another item"], + input_cost_per_second=0.000420, ) print(f"response: {response}") + cost = completion_cost(completion_response=response) + assert ( + cost > 0.0 and cost < 1.0 + ) # should never be > $1 for a single embedding call + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +@pytest.mark.asyncio +async def test_sagemaker_aembeddings(): + try: + response = await litellm.aembedding( + model="sagemaker/berri-benchmarking-gpt-j-6b-fp16", + input=["good morning from litellm", "this is another item"], + input_cost_per_second=0.000420, + ) + print(f"response: {response}") + cost = completion_cost(completion_response=response) + assert ( + cost > 0.0 and cost < 1.0 + ) # should never be > $1 for a single embedding call except Exception as e: pytest.fail(f"Error occurred: {e}")