From 0e69a818425caa76a7dfa67d9d78fd074dc8391a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Jun 2024 20:55:40 -0700 Subject: [PATCH 1/5] cache anthropic httpx client --- litellm/llms/anthropic.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index 8e469a8f48c..5866bf2c6f7 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -161,6 +161,22 @@ def validate_environment(api_key, user_headers): return headers +def _get_async_httpx_client() -> AsyncHTTPHandler: + """ + Retrieves the async HTTP client from the cache + If not present, creates a new client + + Caches the new client and returns it. + """ + _cache_key_name = "anthropic_async_httpx_client" + if _cache_key_name in litellm.in_memory_llm_clients_cache: + return litellm.in_memory_llm_clients_cache[_cache_key_name] + + _new_client = AsyncHTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) + litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client + return _new_client + + async def make_call( client: Optional[AsyncHTTPHandler], api_base: str, @@ -171,7 +187,7 @@ async def make_call( logging_obj, ): if client is None: - client = AsyncHTTPHandler() # Create a new client if none provided + client = _get_async_httpx_client() # Create a new client if none provided response = await client.post(api_base, headers=headers, data=data, stream=True) @@ -463,9 +479,7 @@ class AnthropicChatCompletion(BaseLLM): logger_fn=None, headers={}, ) -> Union[ModelResponse, CustomStreamWrapper]: - async_handler = AsyncHTTPHandler( - timeout=httpx.Timeout(timeout=600.0, connect=5.0) - ) + async_handler = _get_async_httpx_client() response = await async_handler.post(api_base, headers=headers, json=data) if stream and _is_function_call: return self.process_streaming_response( From ec095a814dbcdf6a69e138de802f3fbb0fb4a658 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Jun 2024 21:12:32 -0700 Subject: [PATCH 2/5] fix async client --- .../litellm_core_utils/get_httpx_clients.py | 62 ++++++++++++++ litellm/llms/anthropic.py | 17 +--- litellm/llms/bedrock_httpx.py | 81 +++++++++++++------ 3 files changed, 118 insertions(+), 42 deletions(-) create mode 100644 litellm/litellm_core_utils/get_httpx_clients.py diff --git a/litellm/litellm_core_utils/get_httpx_clients.py b/litellm/litellm_core_utils/get_httpx_clients.py new file mode 100644 index 00000000000..afb48af57ad --- /dev/null +++ b/litellm/litellm_core_utils/get_httpx_clients.py @@ -0,0 +1,62 @@ +from typing import Optional +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +import httpx + + +def _get_async_httpx_client(params: Optional[dict] = None) -> AsyncHTTPHandler: + """ + Retrieves the async HTTP client from the cache + If not present, creates a new client + + Caches the new client and returns it. + """ + _params_key_name = "" + if params is not None: + for key, value in params.items(): + try: + _params_key_name += f"{key}_{value}" + except Exception: + pass + + _cache_key_name = "async_httpx_client" + _params_key_name + if _cache_key_name in litellm.in_memory_llm_clients_cache: + return litellm.in_memory_llm_clients_cache[_cache_key_name] + + if params is not None: + _new_client = AsyncHTTPHandler(**params) + else: + _new_client = AsyncHTTPHandler( + timeout=httpx.Timeout(timeout=600.0, connect=5.0) + ) + litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client + return _new_client + + +def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: + """ + Retrieves the HTTP client from the cache + If not present, creates a new client + + Caches the new client and returns it. + """ + _params_key_name = "" + if params is not None: + for key, value in params.items(): + try: + _params_key_name += f"{key}_{value}" + except Exception: + pass + + _cache_key_name = "httpx_client" + _params_key_name + if _cache_key_name in litellm.in_memory_llm_clients_cache: + return litellm.in_memory_llm_clients_cache[_cache_key_name] + + if params is not None: + _new_client = HTTPHandler(**params) + else: + _new_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) + + litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client + return _new_client diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index 5866bf2c6f7..c01a704600d 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -9,6 +9,7 @@ from litellm.utils import ModelResponse, Usage, map_finish_reason, CustomStreamW import litellm from .prompt_templates.factory import prompt_factory, custom_prompt from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.litellm_core_utils.get_httpx_clients import _get_async_httpx_client from .base import BaseLLM import httpx # type: ignore from litellm.types.llms.anthropic import AnthropicMessagesToolChoice @@ -161,22 +162,6 @@ def validate_environment(api_key, user_headers): return headers -def _get_async_httpx_client() -> AsyncHTTPHandler: - """ - Retrieves the async HTTP client from the cache - If not present, creates a new client - - Caches the new client and returns it. - """ - _cache_key_name = "anthropic_async_httpx_client" - if _cache_key_name in litellm.in_memory_llm_clients_cache: - return litellm.in_memory_llm_clients_cache[_cache_key_name] - - _new_client = AsyncHTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) - litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client - return _new_client - - async def make_call( client: Optional[AsyncHTTPHandler], api_base: str, diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 84b61d4cbd1..123adedefaa 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -42,6 +42,10 @@ from .prompt_templates.factory import ( _bedrock_tools_pt, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.litellm_core_utils.get_httpx_clients import ( + _get_async_httpx_client, + _get_httpx_client, +) from .base import BaseLLM import httpx # type: ignore from .bedrock import BedrockError, convert_messages_to_prompt, ModelResponseIterator @@ -57,6 +61,7 @@ from litellm.caching import DualCache iam_cache = DualCache() + class AmazonCohereChatConfig: """ Reference - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html @@ -167,7 +172,7 @@ async def make_call( logging_obj, ): if client is None: - client = AsyncHTTPHandler() # Create a new client if none provided + client = _get_async_httpx_client() # Create a new client if none provided response = await client.post(api_base, headers=headers, data=data, stream=True) @@ -198,7 +203,7 @@ def make_sync_call( logging_obj, ): if client is None: - client = HTTPHandler() # Create a new client if none provided + client = _get_httpx_client() # Create a new client if none provided response = client.post(api_base, headers=headers, data=data, stream=True) @@ -327,13 +332,19 @@ class BedrockLLM(BaseLLM): ) = params_to_check ### CHECK STS ### - if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: - iam_creds_cache_key = json.dumps({ - "aws_web_identity_token": aws_web_identity_token, - "aws_role_name": aws_role_name, - "aws_session_name": aws_session_name, - "aws_region_name": aws_region_name, - }) + if ( + aws_web_identity_token is not None + and aws_role_name is not None + and aws_session_name is not None + ): + iam_creds_cache_key = json.dumps( + { + "aws_web_identity_token": aws_web_identity_token, + "aws_role_name": aws_role_name, + "aws_session_name": aws_session_name, + "aws_region_name": aws_region_name, + } + ) iam_creds_dict = iam_cache.get_cache(iam_creds_cache_key) if iam_creds_dict is None: @@ -348,7 +359,7 @@ class BedrockLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com", ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html @@ -362,12 +373,18 @@ class BedrockLLM(BaseLLM): iam_creds_dict = { "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], - "aws_secret_access_key": sts_response["Credentials"]["SecretAccessKey"], + "aws_secret_access_key": sts_response["Credentials"][ + "SecretAccessKey" + ], "aws_session_token": sts_response["Credentials"]["SessionToken"], "region_name": aws_region_name, } - iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) + iam_cache.set_cache( + key=iam_creds_cache_key, + value=json.dumps(iam_creds_dict), + ttl=3600 - 60, + ) session = boto3.Session(**iam_creds_dict) @@ -976,7 +993,7 @@ class BedrockLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - self.client = HTTPHandler(**_params) # type: ignore + self.client = _get_httpx_client(**_params) # type: ignore else: self.client = client if (stream is not None and stream == True) and provider != "ai21": @@ -1058,7 +1075,7 @@ class BedrockLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = AsyncHTTPHandler(**_params) # type: ignore + client = _get_async_httpx_client(_params) # type: ignore else: client = client # type: ignore @@ -1433,13 +1450,19 @@ class BedrockConverseLLM(BaseLLM): ) = params_to_check ### CHECK STS ### - if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: - iam_creds_cache_key = json.dumps({ - "aws_web_identity_token": aws_web_identity_token, - "aws_role_name": aws_role_name, - "aws_session_name": aws_session_name, - "aws_region_name": aws_region_name, - }) + if ( + aws_web_identity_token is not None + and aws_role_name is not None + and aws_session_name is not None + ): + iam_creds_cache_key = json.dumps( + { + "aws_web_identity_token": aws_web_identity_token, + "aws_role_name": aws_role_name, + "aws_session_name": aws_session_name, + "aws_region_name": aws_region_name, + } + ) iam_creds_dict = iam_cache.get_cache(iam_creds_cache_key) if iam_creds_dict is None: @@ -1454,7 +1477,7 @@ class BedrockConverseLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com", ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html @@ -1468,12 +1491,18 @@ class BedrockConverseLLM(BaseLLM): iam_creds_dict = { "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], - "aws_secret_access_key": sts_response["Credentials"]["SecretAccessKey"], + "aws_secret_access_key": sts_response["Credentials"][ + "SecretAccessKey" + ], "aws_session_token": sts_response["Credentials"]["SessionToken"], "region_name": aws_region_name, } - iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) + iam_cache.set_cache( + key=iam_creds_cache_key, + value=json.dumps(iam_creds_dict), + ttl=3600 - 60, + ) session = boto3.Session(**iam_creds_dict) @@ -1575,7 +1604,7 @@ class BedrockConverseLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = AsyncHTTPHandler(**_params) # type: ignore + client = _get_async_httpx_client(**_params) # type: ignore else: client = client # type: ignore @@ -1847,7 +1876,7 @@ class BedrockConverseLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = HTTPHandler(**_params) # type: ignore + client = _get_httpx_client(**_params) # type: ignore else: client = client try: From 5ea6fbbe1eb5393d3857e7b78c42fe32c2d40365 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Jun 2024 21:23:13 -0700 Subject: [PATCH 3/5] fix cached httpx client --- litellm/llms/bedrock_httpx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 123adedefaa..6752204088c 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -993,7 +993,7 @@ class BedrockLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - self.client = _get_httpx_client(**_params) # type: ignore + self.client = _get_httpx_client(_params) # type: ignore else: self.client = client if (stream is not None and stream == True) and provider != "ai21": @@ -1604,7 +1604,7 @@ class BedrockConverseLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = _get_async_httpx_client(**_params) # type: ignore + client = _get_async_httpx_client(_params) # type: ignore else: client = client # type: ignore @@ -1876,7 +1876,7 @@ class BedrockConverseLLM(BaseLLM): if isinstance(timeout, float) or isinstance(timeout, int): timeout = httpx.Timeout(timeout) _params["timeout"] = timeout - client = _get_httpx_client(**_params) # type: ignore + client = _get_httpx_client(_params) # type: ignore else: client = client try: From 38995def54162be0950d7ebe891ccc8fa0feb2ce Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Jun 2024 21:30:42 -0700 Subject: [PATCH 4/5] refactor to use _get_async_httpx_client --- .../litellm_core_utils/get_httpx_clients.py | 62 ------------------- litellm/llms/anthropic.py | 7 ++- litellm/llms/bedrock_httpx.py | 5 +- litellm/llms/custom_httpx/http_handler.py | 57 +++++++++++++++++ 4 files changed, 65 insertions(+), 66 deletions(-) delete mode 100644 litellm/litellm_core_utils/get_httpx_clients.py diff --git a/litellm/litellm_core_utils/get_httpx_clients.py b/litellm/litellm_core_utils/get_httpx_clients.py deleted file mode 100644 index afb48af57ad..00000000000 --- a/litellm/litellm_core_utils/get_httpx_clients.py +++ /dev/null @@ -1,62 +0,0 @@ -from typing import Optional -import litellm -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -import httpx - - -def _get_async_httpx_client(params: Optional[dict] = None) -> AsyncHTTPHandler: - """ - Retrieves the async HTTP client from the cache - If not present, creates a new client - - Caches the new client and returns it. - """ - _params_key_name = "" - if params is not None: - for key, value in params.items(): - try: - _params_key_name += f"{key}_{value}" - except Exception: - pass - - _cache_key_name = "async_httpx_client" + _params_key_name - if _cache_key_name in litellm.in_memory_llm_clients_cache: - return litellm.in_memory_llm_clients_cache[_cache_key_name] - - if params is not None: - _new_client = AsyncHTTPHandler(**params) - else: - _new_client = AsyncHTTPHandler( - timeout=httpx.Timeout(timeout=600.0, connect=5.0) - ) - litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client - return _new_client - - -def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: - """ - Retrieves the HTTP client from the cache - If not present, creates a new client - - Caches the new client and returns it. - """ - _params_key_name = "" - if params is not None: - for key, value in params.items(): - try: - _params_key_name += f"{key}_{value}" - except Exception: - pass - - _cache_key_name = "httpx_client" + _params_key_name - if _cache_key_name in litellm.in_memory_llm_clients_cache: - return litellm.in_memory_llm_clients_cache[_cache_key_name] - - if params is not None: - _new_client = HTTPHandler(**params) - else: - _new_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) - - litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client - return _new_client diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index c01a704600d..236f7cd4f8b 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -8,8 +8,11 @@ from typing import Callable, Optional, List, Union from litellm.utils import ModelResponse, Usage, map_finish_reason, CustomStreamWrapper import litellm from .prompt_templates.factory import prompt_factory, custom_prompt -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.litellm_core_utils.get_httpx_clients import _get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + _get_async_httpx_client, + _get_httpx_client, +) from .base import BaseLLM import httpx # type: ignore from litellm.types.llms.anthropic import AnthropicMessagesToolChoice diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 6752204088c..7c7210f84cf 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -41,8 +41,9 @@ from .prompt_templates.factory import ( _bedrock_converse_messages_pt, _bedrock_tools_pt, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.litellm_core_utils.get_httpx_clients import ( +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, _get_async_httpx_client, _get_httpx_client, ) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index f0a5163f39d..a3c5865fa37 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -219,3 +219,60 @@ class HTTPHandler: self.close() except Exception: pass + + +def _get_async_httpx_client(params: Optional[dict] = None) -> AsyncHTTPHandler: + """ + Retrieves the async HTTP client from the cache + If not present, creates a new client + + Caches the new client and returns it. + """ + _params_key_name = "" + if params is not None: + for key, value in params.items(): + try: + _params_key_name += f"{key}_{value}" + except Exception: + pass + + _cache_key_name = "async_httpx_client" + _params_key_name + if _cache_key_name in litellm.in_memory_llm_clients_cache: + return litellm.in_memory_llm_clients_cache[_cache_key_name] + + if params is not None: + _new_client = AsyncHTTPHandler(**params) + else: + _new_client = AsyncHTTPHandler( + timeout=httpx.Timeout(timeout=600.0, connect=5.0) + ) + litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client + return _new_client + + +def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: + """ + Retrieves the HTTP client from the cache + If not present, creates a new client + + Caches the new client and returns it. + """ + _params_key_name = "" + if params is not None: + for key, value in params.items(): + try: + _params_key_name += f"{key}_{value}" + except Exception: + pass + + _cache_key_name = "httpx_client" + _params_key_name + if _cache_key_name in litellm.in_memory_llm_clients_cache: + return litellm.in_memory_llm_clients_cache[_cache_key_name] + + if params is not None: + _new_client = HTTPHandler(**params) + else: + _new_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) + + litellm.in_memory_llm_clients_cache[_cache_key_name] = _new_client + return _new_client From 2c499fbd645863289a09219d0cdcc2feb738feac Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Jun 2024 21:37:35 -0700 Subject: [PATCH 5/5] ci/cd run again --- litellm/tests/test_completion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 1e9014ef4a6..91144684660 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -16,7 +16,7 @@ from litellm.llms.prompt_templates.factory import anthropic_messages_pt from unittest.mock import patch, MagicMock from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler -# litellm.num_retries = 3 +# litellm.num_retries =3 litellm.cache = None litellm.success_callback = [] user_message = "Write a short poem about the sky"