fix replicate test

This commit is contained in:
Ishaan Jaff 2024-05-07 19:48:46 -07:00
parent 4e7b5aa9d7
commit 21d3407b95

View file

@ -2168,9 +2168,9 @@ def test_completion_replicate_vicuna():
def test_replicate_custom_prompt_dict(): def test_replicate_custom_prompt_dict():
litellm.set_verbose = True litellm.set_verbose = True
model_name = "replicate/meta/llama-2-70b-chat" model_name = "replicate/meta/llama-2-7b"
litellm.register_prompt_template( litellm.register_prompt_template(
model="replicate/meta/llama-2-70b-chat", model="replicate/meta/llama-2-7b",
initial_prompt_value="You are a good assistant", # [OPTIONAL] initial_prompt_value="You are a good assistant", # [OPTIONAL]
roles={ roles={
"system": { "system": {
@ -2200,6 +2200,7 @@ def test_replicate_custom_prompt_dict():
repetition_penalty=0.1, repetition_penalty=0.1,
num_retries=3, num_retries=3,
) )
except litellm.APIError as e: except litellm.APIError as e:
pass pass
except litellm.APIConnectionError as e: except litellm.APIConnectionError as e:
@ -3016,6 +3017,7 @@ async def test_acompletion_gemini():
else: else:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# Deepseek tests # Deepseek tests
def test_completion_deepseek(): def test_completion_deepseek():
litellm.set_verbose = True litellm.set_verbose = True