diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5ec9c79bb2..4ced0b8fba 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -5,6 +5,8 @@ from typing import Optional, Union, Mapping, Any # https://www.python-httpx.org/advanced/timeouts _DEFAULT_TIMEOUT = httpx.Timeout(timeout=5.0, connect=5.0) +_DEFAULT_CONNECTION_RETRIES = 2 # retries if error connecting to api + class AsyncHTTPHandler: def __init__( @@ -12,6 +14,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 +50,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 +95,49 @@ class AsyncHTTPHandler: response = await self.client.send(req, stream=stream) response.raise_for_status() return response + except httpx.ConnectError: + # Retry the request with a new session if there is a connection error + new_client = self.create_client(timeout=self.timeout, concurrent_limit=1) + for _ in range(_DEFAULT_CONNECTION_RETRIES): + try: + return await self.single_connection_post_request( + url=url, + client=new_client, + data=data, + json=json, + params=params, + headers=headers, + stream=stream, + ) + except httpx.ConnectError: + pass 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())