Merge pull request #4178 from BerriAI/litellm_http_handler

Retry on connection disconnect
This commit is contained in:
Krish Dholakia 2024-06-13 21:14:56 -07:00 committed by GitHub
commit d499477f54
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 65 additions and 18 deletions

View File

@ -26,7 +26,7 @@ class AuthenticationError(openai.AuthenticationError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 401
self.message = message
self.message = "litellm.AuthenticationError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -72,7 +72,7 @@ class NotFoundError(openai.NotFoundError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 404
self.message = message
self.message = "litellm.NotFoundError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -117,7 +117,7 @@ class BadRequestError(openai.BadRequestError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 400
self.message = message
self.message = "litellm.BadRequestError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -162,7 +162,7 @@ class UnprocessableEntityError(openai.UnprocessableEntityError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 422
self.message = message
self.message = "litellm.UnprocessableEntityError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -204,7 +204,7 @@ class Timeout(openai.APITimeoutError): # type: ignore
request=request
) # Call the base class constructor with the parameters it needs
self.status_code = 408
self.message = message
self.message = "litellm.Timeout: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -241,7 +241,7 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type:ignore
num_retries: Optional[int] = None,
):
self.status_code = 403
self.message = message
self.message = "litellm.PermissionDeniedError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -280,7 +280,7 @@ class RateLimitError(openai.RateLimitError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 429
self.message = message
self.message = "litellm.RateLimitError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -328,7 +328,7 @@ class ContextWindowExceededError(BadRequestError): # type: ignore
litellm_debug_info: Optional[str] = None,
):
self.status_code = 400
self.message = message
self.message = "litellm.ContextWindowExceededError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -368,7 +368,7 @@ class RejectedRequestError(BadRequestError): # type: ignore
litellm_debug_info: Optional[str] = None,
):
self.status_code = 400
self.message = message
self.message = "litellm.RejectedRequestError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -411,7 +411,7 @@ class ContentPolicyViolationError(BadRequestError): # type: ignore
litellm_debug_info: Optional[str] = None,
):
self.status_code = 400
self.message = message
self.message = "litellm.ContentPolicyViolationError: {}".format(message)
self.model = model
self.llm_provider = llm_provider
self.litellm_debug_info = litellm_debug_info
@ -452,7 +452,7 @@ class ServiceUnavailableError(openai.APIStatusError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 503
self.message = message
self.message = "litellm.ServiceUnavailableError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -501,7 +501,7 @@ class InternalServerError(openai.InternalServerError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = 500
self.message = message
self.message = "litellm.InternalServerError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -552,7 +552,7 @@ class APIError(openai.APIError): # type: ignore
num_retries: Optional[int] = None,
):
self.status_code = status_code
self.message = message
self.message = "litellm.APIError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.litellm_debug_info = litellm_debug_info
@ -589,7 +589,7 @@ class APIConnectionError(openai.APIConnectionError): # type: ignore
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
):
self.message = message
self.message = "litellm.APIConnectionError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
self.status_code = 500
@ -626,7 +626,7 @@ class APIResponseValidationError(openai.APIResponseValidationError): # type: ig
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
):
self.message = message
self.message = "litellm.APIResponseValidationError: {}".format(message)
self.llm_provider = llm_provider
self.model = model
request = httpx.Request(method="POST", url="https://api.openai.com/v1")

View File

@ -12,6 +12,15 @@ class AsyncHTTPHandler:
timeout: Optional[Union[float, httpx.Timeout]] = None,
concurrent_limit=1000,
):
self.timeout = timeout
self.client = self.create_client(
timeout=timeout, concurrent_limit=concurrent_limit
)
def create_client(
self, timeout: Optional[Union[float, httpx.Timeout]], concurrent_limit: int
) -> httpx.AsyncClient:
async_proxy_mounts = None
# Check if the HTTP_PROXY and HTTPS_PROXY environment variables are set and use them accordingly.
http_proxy = os.getenv("HTTP_PROXY", None)
@ -39,7 +48,8 @@ class AsyncHTTPHandler:
if timeout is None:
timeout = _DEFAULT_TIMEOUT
# Create a client with a connection pool
self.client = httpx.AsyncClient(
return httpx.AsyncClient(
timeout=timeout,
limits=httpx.Limits(
max_connections=concurrent_limit,
@ -83,11 +93,48 @@ class AsyncHTTPHandler:
response = await self.client.send(req, stream=stream)
response.raise_for_status()
return response
except httpx.RemoteProtocolError:
# Retry the request with a new session if there is a connection error
new_client = self.create_client(timeout=self.timeout, concurrent_limit=1)
try:
return await self.single_connection_post_request(
url=url,
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
)
finally:
await new_client.aclose()
except httpx.HTTPStatusError as e:
raise e
except Exception as e:
raise e
async def single_connection_post_request(
self,
url: str,
client: httpx.AsyncClient,
data: Optional[Union[dict, str]] = None, # type: ignore
json: Optional[dict] = None,
params: Optional[dict] = None,
headers: Optional[dict] = None,
stream: bool = False,
):
"""
Making POST request for a single connection client.
Used for retrying connection client errors.
"""
req = client.build_request(
"POST", url, data=data, json=json, params=params, headers=headers # type: ignore
)
response = await client.send(req, stream=stream)
response.raise_for_status()
return response
def __del__(self) -> None:
try:
asyncio.get_running_loop().create_task(self.close())

View File

@ -11,7 +11,7 @@ model_list:
- model_name: fake-openai-endpoint
litellm_params:
model: predibase/llama-3-8b-instruct
api_base: "http://0.0.0.0:8081"
api_base: "http://0.0.0.0:8000"
api_key: os.environ/PREDIBASE_API_KEY
tenant_id: os.environ/PREDIBASE_TENANT_ID
max_retries: 0

View File

@ -81,7 +81,7 @@ def test_async_fallbacks(caplog):
# Define the expected log messages
# - error request, falling back notice, success notice
expected_logs = [
"litellm.acompletion(model=gpt-3.5-turbo)\x1b[31m Exception AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: bad-key. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}\x1b[0m",
"litellm.acompletion(model=gpt-3.5-turbo)\x1b[31m Exception litellm.AuthenticationError: AuthenticationError: OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: bad-key. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}\x1b[0m",
"Falling back to model_group = azure/gpt-3.5-turbo",
"litellm.acompletion(model=azure/chatgpt-v-2)\x1b[32m 200 OK\x1b[0m",
"Successful fallback b/w models.",