refactor: simplify litellm exception imports

This commit is contained in:
Paul Gauthier 2024-10-28 14:27:19 -07:00 committed by Paul Gauthier (aider)
parent e2dff0a74b
commit bf63e7045b
2 changed files with 15 additions and 16 deletions

View file

@ -15,7 +15,7 @@ CACHE = None
RETRY_TIMEOUT = 60 RETRY_TIMEOUT = 60
#ai
def retry_exceptions(): def retry_exceptions():
import httpx import httpx
@ -25,19 +25,18 @@ def retry_exceptions():
httpx.RemoteProtocolError, httpx.RemoteProtocolError,
httpx.ReadTimeout, httpx.ReadTimeout,
# litellm # litellm
litellm.exceptions.BadRequestError, litellm.AuthenticationError,
litellm.exceptions.AuthenticationError, litellm.PermissionDeniedError,
litellm.exceptions.PermissionDeniedError, litellm.NotFoundError,
litellm.exceptions.NotFoundError, litellm.UnprocessableEntityError,
litellm.exceptions.UnprocessableEntityError, litellm.RateLimitError,
litellm.exceptions.RateLimitError, litellm.InternalServerError,
litellm.exceptions.InternalServerError, litellm.ContextWindowExceededError,
litellm.exceptions.ContextWindowExceededError, litellm.ContentPolicyViolationError,
litellm.exceptions.ContentPolicyViolationError, litellm.APIConnectionError,
litellm.exceptions.APIConnectionError, litellm.APIError,
litellm.exceptions.APIError, litellm.ServiceUnavailableError,
litellm.exceptions.ServiceUnavailableError, litellm.Timeout,
litellm.exceptions.Timeout,
) )
@ -111,5 +110,5 @@ def simple_send_with_retries(model_name, messages, extra_params=None):
_hash, response = send_completion(**kwargs) _hash, response = send_completion(**kwargs)
return response.choices[0].message.content return response.choices[0].message.content
except (AttributeError, litellm.exceptions.BadRequestError): except (AttributeError, litellm.BadRequestError):
return return

View file

@ -10,7 +10,7 @@ from aider.sendchat import simple_send_with_retries
class PrintCalled(Exception): class PrintCalled(Exception):
pass pass
#ai add a test that simply invokes retry_exceptions() and ensures no error raised!
class TestSendChat(unittest.TestCase): class TestSendChat(unittest.TestCase):
@patch("litellm.completion") @patch("litellm.completion")
@patch("builtins.print") @patch("builtins.print")