update test_batch_completions

This commit is contained in:
Krrish Dholakia 2023-09-11 08:44:23 -07:00
parent 8fc25c26ae
commit 08132975f4

View file

@ -1,28 +1,25 @@
#### What this tests #### # #### What this tests ####
# This tests calling batch_completions by running 100 messages together # # This tests calling batch_completions by running 100 messages together
import sys, os # import sys, os
import traceback # import traceback
import pytest # import pytest
sys.path.insert( # sys.path.insert(
0, os.path.abspath("../..") # 0, os.path.abspath("../..")
) # Adds the parent directory to the system path # ) # Adds the parent directory to the system path
from openai.error import Timeout # from openai.error import Timeout
import litellm # import litellm
from litellm import batch_completion # from litellm import batch_completion
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": "Hey, how's it going"}] for _ in range(5)]
print(messages[0:5]) # model = "gpt-3.5-turbo"
print(len(messages)) # try:
# model = "vllm/facebook/opt-125m" # result = batch_completion(model=model, messages=messages)
model = "gpt-3.5-turbo" # print(result)
try: # print(len(result))
result = batch_completion(model=model, messages=messages) # except Timeout as e:
print(result) # pass
print(len(result)) # except Exception as e:
except Timeout as e: # pytest.fail(f"An error occurred: {e}")
pass
except Exception as e:
pytest.fail(f"An error occurred: {e}")