scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

test: refactor testing

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-12 16:21:41 -08:00
parent 669862643b
commit dfdb17ae36
2 changed files with 22 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/tests/test_completion.py
View file

@ -1086,51 +1086,6 @@ def test_completion_chat_sagemaker():
pytest.fail(f"Error occurred: {e}")
# test_completion_chat_sagemaker()
def test_completion_chat_sagemaker():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose=True
response = completion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=messages,
max_tokens=100,
temperature=0.7,
stream=True,
)
# Add any assertions here to check the response
complete_response = ""
for chunk in response:
complete_response += chunk.choices[0].delta.content or ""
print(f"complete_response: {complete_response}")
assert len(complete_response) > 0
except Exception as e:
pytest.fail(f"Error occurred: {e}")
import asyncio
@pytest.mark.asyncio
async def test_completion_chat_sagemaker():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose=True
response = await litellm.acompletion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=messages,
max_tokens=100,
temperature=0.7,
stream=True,
)
# Add any assertions here to check the response
complete_response = ""
async for chunk in response:
complete_response += chunk.choices[0].delta.content or ""
print(f"complete_response: {complete_response}")
assert len(complete_response) > 0
except:
pass
asyncio.run(test_completion_chat_sagemaker())
def test_completion_chat_sagemaker_mistral():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]

22
litellm/tests/test_streaming.py
View file

@ -654,6 +654,28 @@ def test_sagemaker_weird_response():
pytest.fail(f"An exception occurred - {str(e)}")
# test_sagemaker_weird_response()
@pytest.mark.asyncio
async def test_sagemaker_streaming_async():
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose=True
response = await litellm.acompletion(
model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
messages=messages,
max_tokens=100,
temperature=0.7,
stream=True,
)
# Add any assertions here to check the response
complete_response = ""
async for chunk in response:
complete_response += chunk.choices[0].delta.content or ""
print(f"complete_response: {complete_response}")
assert len(complete_response) > 0
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# def test_completion_sagemaker_stream():
# try:
# response = completion(