phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_completion.py): fix replicate streaming test

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-21 10:38:05 -08:00
parent 1218121e47
commit 37187e8e88
1 changed files with 9 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

19
litellm/tests/test_completion.py
View file

@ -54,14 +54,11 @@ def test_completion_claude():
print(response.usage) print(response.usage)
print(response.usage.completion_tokens) print(response.usage.completion_tokens)
print(response["usage"]["completion_tokens"]) print(response["usage"]["completion_tokens"])
response = completion(
model="claude-2.1", messages=messages, request_timeout=10,
)
# print("new cost tracking") # print("new cost tracking")
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_completion_claude() # test_completion_claude()
def test_completion_claude2_1(): def test_completion_claude2_1():
try: try:
@ -804,7 +801,6 @@ def test_completion_replicate_vicuna():
# test_completion_replicate_vicuna() # test_completion_replicate_vicuna()
def test_completion_replicate_llama2_stream(): def test_completion_replicate_llama2_stream():
print("TESTING REPLICATE streaming")
litellm.set_verbose=False litellm.set_verbose=False
model_name = "replicate/meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0" model_name = "replicate/meta/llama-2-7b-chat:13c3cdee13ee059ab779f0291d29054dab00a47dad8261375654de5540165fb0"
try: try:
@ -820,15 +816,18 @@ def test_completion_replicate_llama2_stream():
max_tokens=20, max_tokens=20,
num_retries=3 num_retries=3
) )
print(response) print(f"response: {response}")
# Add any assertions here to check the response # Add any assertions here to check the response
complete_response = ""
for i, chunk in enumerate(response): for i, chunk in enumerate(response):
if i == 0: complete_response += chunk.choices[0].delta["content"]
assert len(chunk.choices[0].delta["content"]) > 5 # if i == 0:
print(chunk) # assert len(chunk.choices[0].delta["content"]) > 2
# print(chunk)
assert len(complete_response) > 5
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_replicate_llama2_stream() test_completion_replicate_llama2_stream()
# commenthing this out since we won't be always testing a custom replicate deployment # commenthing this out since we won't be always testing a custom replicate deployment
# def test_completion_replicate_deployments(): # def test_completion_replicate_deployments():