From 86718b75b4aeac22d144620edb485249257678a2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 Jan 2025 14:01:59 -0800 Subject: [PATCH 1/4] ClientSession --- .gitignore | 2 ++ litellm/llms/custom_httpx/aiohttp_handler.py | 8 +++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 4a92cbf31c7..be045dd2651 100644 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,5 @@ litellm/proxy/google-cloud-sdk/* tests/llm_translation/log.txt venv/ tests/local_testing/log.txt + +.codegpt \ No newline at end of file diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index ab6872e0437..be93b68d627 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -29,6 +29,9 @@ class BaseLLMAIOHTTPHandler: + def __init__(self): + self.client_session: Optional[aiohttp.ClientSession] = None + async def _make_common_async_call( self, async_httpx_client: AsyncHTTPHandler, @@ -52,7 +55,10 @@ async def _make_common_async_call( ) ) - async with aiohttp.ClientSession(timeout=timeout_obj) as session: + if self.client_session is None: + self.client_session = aiohttp.ClientSession() + + async with self.client_session as session: for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: response = await session.post( From 56c38b07bce6b3908d6a005f708fefea721c2f64 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 Jan 2025 14:09:59 -0800 Subject: [PATCH 2/4] re use client_session --- litellm/llms/custom_httpx/aiohttp_handler.py | 29 ++++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index be93b68d627..9843d710e4f 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -58,21 +58,20 @@ async def _make_common_async_call( if self.client_session is None: self.client_session = aiohttp.ClientSession() - async with self.client_session as session: - for i in range(max(max_retry_on_unprocessable_entity_error, 1)): - try: - response = await session.post( - url=api_base, - headers=headers, - json=data, - ) - if not response.ok: - response.raise_for_status() - except aiohttp.ClientResponseError as e: - raise self._handle_error(e=e, provider_config=provider_config) - except Exception as e: - raise self._handle_error(e=e, provider_config=provider_config) - break + for i in range(max(max_retry_on_unprocessable_entity_error, 1)): + try: + response = await self.client_session.post( + url=api_base, + headers=headers, + json=data, + ) + if not response.ok: + response.raise_for_status() + except aiohttp.ClientResponseError as e: + raise self._handle_error(e=e, provider_config=provider_config) + except Exception as e: + raise self._handle_error(e=e, provider_config=provider_config) + break if response is None: raise provider_config.get_error_class( From 72acfafb467801867f79d0338d196eabbdfb298d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 Jan 2025 14:27:54 -0800 Subject: [PATCH 3/4] _init_client_session --- litellm/llms/custom_httpx/aiohttp_handler.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 9843d710e4f..92d766e2f4f 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,8 +1,9 @@ import json from typing import TYPE_CHECKING, Any, Optional, Tuple, Union -import aiohttp # Add this import +import aiohttp import httpx # type: ignore +from aiohttp import TCPConnector import litellm import litellm.litellm_core_utils @@ -25,6 +26,7 @@ LiteLLMLoggingObj = Any DEFAULT_TIMEOUT = 600 +DEFAULT_AIOHTTP_CONNECTION_LIMIT = 10000 class BaseLLMAIOHTTPHandler: @@ -56,7 +58,7 @@ async def _make_common_async_call( ) if self.client_session is None: - self.client_session = aiohttp.ClientSession() + self.client_session = self._init_client_session() for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: @@ -136,6 +138,12 @@ def _make_common_sync_call( return response + def _init_client_session(self): + self.client_session = aiohttp.ClientSession( + connector=TCPConnector(limit=10000), + ) + return self.client_session + async def async_completion( self, custom_llm_provider: str, From 8f8de796e79bb2d13511339cb58ad53ab9bd1e69 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 Jan 2025 14:51:27 -0800 Subject: [PATCH 4/4] fix aiohttp --- litellm/llms/custom_httpx/aiohttp_handler.py | 32 ++------------------ litellm/proxy/proxy_config.yaml | 4 +-- 2 files changed, 5 insertions(+), 31 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 92d766e2f4f..f87355097df 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,9 +1,8 @@ import json from typing import TYPE_CHECKING, Any, Optional, Tuple, Union -import aiohttp +import aiohttp # Add this import import httpx # type: ignore -from aiohttp import TCPConnector import litellm import litellm.litellm_core_utils @@ -26,7 +25,6 @@ LiteLLMLoggingObj = Any DEFAULT_TIMEOUT = 600 -DEFAULT_AIOHTTP_CONNECTION_LIMIT = 10000 class BaseLLMAIOHTTPHandler: @@ -58,7 +56,7 @@ async def _make_common_async_call( ) if self.client_session is None: - self.client_session = self._init_client_session() + self.client_session = aiohttp.ClientSession() for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: @@ -138,12 +136,6 @@ def _make_common_sync_call( return response - def _init_client_session(self): - self.client_session = aiohttp.ClientSession( - connector=TCPConnector(limit=10000), - ) - return self.client_session - async def async_completion( self, custom_llm_provider: str, @@ -181,25 +173,7 @@ async def async_completion( ) _json_response = await _response.json() - # cast to httpx.Response - # Todo - use this until we migrate fully to aiohttp - response = httpx.Response( - status_code=_response.status, - headers=_response.headers, - json=_json_response, - ) - return provider_config.transform_response( - model=model, - raw_response=response, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - request_data=data, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - ) + return _json_response def completion( self, diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index e80b764cf96..787347e84d7 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,8 +1,8 @@ model_list: - model_name: "fake-openai-endpoint" litellm_params: - model: openai/any - api_base: https://exampleopenaiendpoint-production.up.railway.app + model: aiohttp_openai/any + api_base: https://example-openai-endpoint.onrender.com/chat/completions api_key: "ishaan" general_settings: