phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
ishaan-jaff 2023-10-06 15:36:13 -07:00
parent 5953a992b7
commit 6f47a2d6f6
1 changed files with 39 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

70
litellm/tests/test_batch_completions.py
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,27 +65,28 @@ def test_batch_completion_models_all_responses():
import asyncio def test_rate_limit_handler():
##### USAGE ################ import asyncio
##### USAGE ################
jobs = [ 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"}]},
{"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( handler = RateLimitManager(
max_requests_per_minute = 60, max_requests_per_minute = 60,
max_tokens_per_minute = 20000 max_tokens_per_minute = 20000
)
try:
asyncio.run(
handler.batch_completion(
jobs = jobs,
api_key=os.environ['OPENAI_API_KEY'],
) )
)
except Exception as e: try:
print(e) asyncio.run(
handler.batch_completion(
jobs = jobs,
api_key=os.environ['OPENAI_API_KEY'],
)
)
except Exception as e:
print(e)