phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_batch_completions.py): handle timeouts

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-04 13:34:55 -07:00
parent 3477604c90
commit d59c9d96bb
1 changed files with 18 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/tests/test_batch_completions.py
View file

@ -14,18 +14,20 @@ from litellm import batch_completion, batch_completion_models, completion, batch
def test_batch_completions(): def test_batch_completions():
messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)] messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)]
model = "gpt-3.5-turbo" model = "j2-mid"
try: try:
result = batch_completion( result = batch_completion(
model=model, model=model,
messages=messages, messages=messages,
max_tokens=10, max_tokens=10,
temperature=0.2 temperature=0.2,
request_timeout=1
) )
print(result) print(result)
print(len(result)) print(len(result))
assert(len(result)==3) assert(len(result)==3)
except Timeout as e: except Timeout as e:
print(f"IN TIMEOUT")
pass pass
except Exception as e: except Exception as e:
pytest.fail(f"An error occurred: {e}") pytest.fail(f"An error occurred: {e}")
@ -38,18 +40,25 @@ def test_batch_completions_models():
messages=[{"role": "user", "content": "Hey, how's it going"}] messages=[{"role": "user", "content": "Hey, how's it going"}]
) )
print(result) print(result)
except Timeout as e:
pass
except Exception as e: except Exception as e:
pytest.fail(f"An error occurred: {e}") pytest.fail(f"An error occurred: {e}")
# test_batch_completions_models() # test_batch_completions_models()
def test_batch_completion_models_all_responses(): def test_batch_completion_models_all_responses():
responses = batch_completion_models_all_responses( try:
models=["j2-light", "claude-instant-1.2", "command-nightly"], responses = batch_completion_models_all_responses(
messages=[{"role": "user", "content": "write a poem"}], models=["j2-light", "claude-instant-1.2", "command-nightly"],
max_tokens=500 messages=[{"role": "user", "content": "write a poem"}],
) max_tokens=500
print(responses) )
assert(len(responses) == 3) print(responses)
assert(len(responses) == 3)
except Timeout as e:
pass
except Exception as e:
pytest.fail(f"An error occurred: {e}")
# test_batch_completion_models_all_responses() # test_batch_completion_models_all_responses()
# def test_batch_completions(): # def test_batch_completions():