From 61d67cfa43cdcc1f884fe794c66ee2f08769e234 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 6 Jan 2025 15:41:39 -0800
Subject: [PATCH] (perf) - fixes for aiohttp handler to hit 1K RPS (#7590)

* fix getting aiohttp sesson

* fix _get_async_client_session
---
 litellm/llms/custom_httpx/aiohttp_handler.py | 39 ++++++++++++--------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py
index d579391518..f88b48d018 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 ClientSession
 
 import litellm
 import litellm.litellm_core_utils
@@ -13,7 +14,6 @@ from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
     _get_httpx_client,
-    get_async_httpx_client,
 )
 from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager
 
@@ -32,9 +32,21 @@ class BaseLLMAIOHTTPHandler:
     def __init__(self):
         self.client_session: Optional[aiohttp.ClientSession] = None
 
+    def _get_async_client_session(
+        self, dynamic_client_session: Optional[ClientSession] = None
+    ) -> ClientSession:
+        if dynamic_client_session:
+            return dynamic_client_session
+        elif self.client_session:
+            return self.client_session
+        else:
+            # init client session, and then return new session
+            self.client_session = aiohttp.ClientSession()
+            return self.client_session
+
     async def _make_common_async_call(
         self,
-        async_httpx_client: AsyncHTTPHandler,
+        async_client_session: Optional[ClientSession],
         provider_config: BaseConfig,
         api_base: str,
         headers: dict,
@@ -49,12 +61,13 @@ class BaseLLMAIOHTTPHandler:
         )
 
         response: Optional[aiohttp.ClientResponse] = None
-        if self.client_session is None:
-            self.client_session = aiohttp.ClientSession()
+        async_client_session = self._get_async_client_session(
+            dynamic_client_session=async_client_session
+        )
 
         for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
             try:
-                response = await self.client_session.post(
+                response = await async_client_session.post(
                     url=api_base,
                     headers=headers,
                     json=data,
@@ -146,17 +159,11 @@ class BaseLLMAIOHTTPHandler:
         litellm_params: dict,
         encoding: Any,
         api_key: Optional[str] = None,
-        client: Optional[AsyncHTTPHandler] = None,
+        client: Optional[ClientSession] = None,
     ):
-        if client is None:
-            async_httpx_client = get_async_httpx_client(
-                llm_provider=litellm.LlmProviders(custom_llm_provider)
-            )
-        else:
-            async_httpx_client = client
 
         _response = await self._make_common_async_call(
-            async_httpx_client=async_httpx_client,
+            async_client_session=client,
             provider_config=provider_config,
             api_base=api_base,
             headers=headers,
@@ -186,7 +193,7 @@ class BaseLLMAIOHTTPHandler:
         fake_stream: bool = False,
         api_key: Optional[str] = None,
         headers: Optional[dict] = {},
-        client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
+        client: Optional[Union[HTTPHandler, AsyncHTTPHandler, ClientSession]] = None,
     ):
         provider_config = ProviderConfigManager.get_provider_chat_config(
             model=model, provider=litellm.LlmProviders(custom_llm_provider)
@@ -245,7 +252,7 @@ class BaseLLMAIOHTTPHandler:
                 encoding=encoding,
                 client=(
                     client
-                    if client is not None and isinstance(client, AsyncHTTPHandler)
+                    if client is not None and isinstance(client, ClientSession)
                     else None
                 ),
             )