diff --git a/litellm/__init__.py b/litellm/__init__.py index c3a4e20bbde..a90b77fb57b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -377,8 +377,7 @@ from .exceptions import ( ServiceUnavailableError, OpenAIError, ContextWindowExceededError, - BudgetExceededError - + BudgetExceededError, ) from .budget_manager import BudgetManager from .proxy.proxy_cli import run_server diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index eec5bb6f4b1..340a3b0ae58 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -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'] diff --git a/litellm/timeout.py b/litellm/timeout.py index 9a8cba2410a..babcc7bea00 100644 --- a/litellm/timeout.py +++ b/litellm/timeout.py @@ -17,7 +17,7 @@ from concurrent import futures from inspect import iscoroutinefunction from functools import wraps 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):