test(batch_completion): add tests for batch_completion and rate limit handler

This commit is contained in:
ishaan-jaff 2023-10-06 15:36:13 -07:00
parent 5953a992b7
commit 6f47a2d6f6

View file

@ -12,17 +12,24 @@ import litellm
from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses
# litellm.set_verbose=True # litellm.set_verbose=True
# def test_batch_completions(): def test_batch_completions():
# messages = [[{"role": "user", "content": "Hey, how's it going"}] for _ in range(5)] messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)]
# model = "gpt-3.5-turbo" model = "gpt-3.5-turbo"
# try: try:
# result = batch_completion(model=model, messages=messages) result = batch_completion(
# print(result) model=model,
# print(len(result)) messages=messages,
# except Timeout as e: max_tokens=10,
# pass temperature=0.2
# except Exception as e: )
# pytest.fail(f"An error occurred: {e}") print(result)
print(len(result))
assert(len(result)==3)
except Timeout as e:
pass
except Exception as e:
pytest.fail(f"An error occurred: {e}")
# test_batch_completions()
def test_batch_completions_models(): def test_batch_completions_models():
try: try:
@ -58,6 +65,7 @@ def test_batch_completion_models_all_responses():
def test_rate_limit_handler():
import asyncio import asyncio
##### USAGE ################ ##### USAGE ################