phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_completion.py): num retries 3

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-04 15:55:56 -07:00
parent f7c5595a0d
commit 2310e0c6eb
3 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/__init__.py
View file

@ -377,8 +377,7 @@ from .exceptions import (
ServiceUnavailableError, ServiceUnavailableError,
OpenAIError, OpenAIError,
ContextWindowExceededError, ContextWindowExceededError,
BudgetExceededError BudgetExceededError,
) )
from .budget_manager import BudgetManager from .budget_manager import BudgetManager
from .proxy.proxy_cli import run_server from .proxy.proxy_cli import run_server

5
litellm/tests/test_completion.py
View file

@ -9,9 +9,11 @@ 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
import pytest import pytest
from openai.error import Timeout
import litellm import litellm
from litellm import embedding, completion, completion_cost from litellm import embedding, completion, completion_cost
from litellm import RateLimitError from litellm import RateLimitError
litellm.num_retries = 3
user_message = "Write a short poem about the sky" user_message = "Write a short poem about the sky"
messages = [{"content": user_message, "role": "user"}] 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 # Add any assertions here to check the response
print(response) print(response)
except Timeout as e:
pass
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
def test_completion_openai_litellm_key(): def test_completion_openai_litellm_key():
try: try:
litellm.api_key = os.environ['OPENAI_API_KEY'] litellm.api_key = os.environ['OPENAI_API_KEY']

2
litellm/timeout.py
View file

@ -17,7 +17,7 @@ from concurrent import futures
from inspect import iscoroutinefunction from inspect import iscoroutinefunction
from functools import wraps from functools import wraps
from threading import Thread from threading import Thread
from openai.error import Timeout from litellm.error import Timeout
def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout): def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout):