test(test_completion.py): num retries 3

This commit is contained in:
Krrish Dholakia 2023-11-04 15:55:56 -07:00
parent 5b3978eff4
commit 5bd6d6ff4c
3 changed files with 6 additions and 4 deletions

View file

@ -9,9 +9,11 @@ sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
from openai.error import Timeout
import litellm
from litellm import embedding, completion, completion_cost
from litellm import RateLimitError
litellm.num_retries = 3
user_message = "Write a short poem about the sky"
messages = [{"content": user_message, "role": "user"}]
@ -405,10 +407,11 @@ def test_completion_openai_with_optional_params():
)
# Add any assertions here to check the response
print(response)
except Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_openai_litellm_key():
try:
litellm.api_key = os.environ['OPENAI_API_KEY']