diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index 8e469a8f48c..236f7cd4f8b 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -8,7 +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.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 @@ -171,7 +175,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 +467,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( diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 84b61d4cbd1..7c7210f84cf 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -41,7 +41,12 @@ from .prompt_templates.factory import ( _bedrock_converse_messages_pt, _bedrock_tools_pt, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _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 +62,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 +173,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 +204,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 +333,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 +360,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 +374,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 +994,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 +1076,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 +1451,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 +1478,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 +1492,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 +1605,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 +1877,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: 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 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"