Merge pull request #1952 from BerriAI/litellm_aioboto3_sagemaker

Implements aioboto3 for sagemaker
This commit is contained in:
Krish Dholakia 2024-02-14 21:47:22 -08:00 committed by GitHub
commit 122ad77d56
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 203 additions and 75 deletions

View file

@ -1907,24 +1907,7 @@ async def async_data_generator(response, user_api_key_dict):
def select_data_generator(response, user_api_key_dict):
try:
# since boto3 - sagemaker does not support async calls, we should use a sync data_generator
if hasattr(
response, "custom_llm_provider"
) and response.custom_llm_provider in ["sagemaker"]:
return data_generator(
response=response,
)
else:
# default to async_data_generator
return async_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
except:
# worst case - use async_data_generator
return async_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
return async_data_generator(response=response, user_api_key_dict=user_api_key_dict)
def get_litellm_model_info(model: dict = {}):