diff --git a/litellm/tests/test_batch_completions.py b/litellm/tests/test_batch_completions.py index 8ada4dcf3c..53e90fd2ac 100644 --- a/litellm/tests/test_batch_completions.py +++ b/litellm/tests/test_batch_completions.py @@ -12,17 +12,24 @@ import litellm from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses # litellm.set_verbose=True -# def test_batch_completions(): -# messages = [[{"role": "user", "content": "Hey, how's it going"}] for _ in range(5)] -# model = "gpt-3.5-turbo" -# try: -# result = batch_completion(model=model, messages=messages) -# print(result) -# print(len(result)) -# except Timeout as e: -# pass -# except Exception as e: -# pytest.fail(f"An error occurred: {e}") +def test_batch_completions(): + messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)] + model = "gpt-3.5-turbo" + try: + result = batch_completion( + model=model, + messages=messages, + max_tokens=10, + temperature=0.2 + ) + 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(): try: @@ -58,27 +65,28 @@ def test_batch_completion_models_all_responses(): -import asyncio -##### USAGE ################ +def test_rate_limit_handler(): + import asyncio + ##### USAGE ################ -jobs = [ - {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]}, - {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]}, -] + jobs = [ + {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]}, + {"model": "gpt-3.5-turbo-16k", "messages": [{"content": "Please provide a summary of the latest scientific discoveries.", "role": "user"}]}, + ] -from litellm import RateLimitManager + from litellm import RateLimitManager -handler = RateLimitManager( - max_requests_per_minute = 60, - max_tokens_per_minute = 20000 -) - -try: - asyncio.run( - handler.batch_completion( - jobs = jobs, - api_key=os.environ['OPENAI_API_KEY'], + handler = RateLimitManager( + max_requests_per_minute = 60, + max_tokens_per_minute = 20000 ) -) -except Exception as e: - print(e) \ No newline at end of file + + try: + asyncio.run( + handler.batch_completion( + jobs = jobs, + api_key=os.environ['OPENAI_API_KEY'], + ) + ) + except Exception as e: + print(e)