phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) select_data_generator

Browse source
This commit is contained in:
ishaan-jaff 2024-01-23 12:13:34 -08:00
parent a61dbc1613
commit 44e213e842
1 changed files with 16 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/proxy_server.py
View file

@ -1423,13 +1423,22 @@ async def async_data_generator(response, user_api_key_dict):
def select_data_generator(response, user_api_key_dict):
# since boto3 - sagemaker does not support async calls
if response.custom_llm_provider == "sagemaker":
return data_generator(
response=response,
)
else:
# default to async_data_generator
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 == "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
)