fix(timeout.py): fix timeout issue

This commit is contained in:
Krrish Dholakia 2023-11-13 11:07:09 -08:00
parent 8ae511b178
commit b6eea3cf8e
3 changed files with 5 additions and 4 deletions

View file

@ -46,11 +46,12 @@ class BadRequestError(BadRequestError): # type: ignore
) # Call the base class constructor with the parameters it needs ) # Call the base class constructor with the parameters it needs
class Timeout(APITimeoutError): # type: ignore class Timeout(APITimeoutError): # type: ignore
def __init__(self, message, model, llm_provider, request: httpx.Request): def __init__(self, message, model, llm_provider):
self.status_code = 408 self.status_code = 408
self.message = message self.message = message
self.model = model self.model = model
self.llm_provider = llm_provider self.llm_provider = llm_provider
request = httpx.Request(method="POST", url="https://api.openai.com/v1")
super().__init__( super().__init__(
request=request request=request
) # Call the base class constructor with the parameters it needs ) # Call the base class constructor with the parameters it needs

View file

@ -7,7 +7,7 @@ import pytest
sys.path.insert( 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
from openai import Timeout from openai import APITimeoutError as Timeout
import litellm import litellm
from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses from litellm import batch_completion, batch_completion_models, completion, batch_completion_models_all_responses
# litellm.set_verbose=True # litellm.set_verbose=True
@ -31,7 +31,7 @@ def test_batch_completions():
pass pass
except Exception as e: except Exception as e:
pytest.fail(f"An error occurred: {e}") pytest.fail(f"An error occurred: {e}")
# test_batch_completions() test_batch_completions()
def test_batch_completions_models(): def test_batch_completions_models():
try: try:

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 import Timeout from litellm.exceptions import Timeout
def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout): def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout):