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

refactor(bedrock.py): better exception mapping for bedrock + huggingface

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-04 16:12:12 -07:00
parent ab54262d37
commit 7c46e85ed6
2 changed files with 158 additions and 141 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/bedrock.py
View file

@ -278,6 +278,8 @@ def completion(
litellm_params=None, litellm_params=None,
logger_fn=None, logger_fn=None,
): ):
exception_mapping_worked = False
try:
# pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them
aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = optional_params.pop("aws_access_key_id", None) aws_access_key_id = optional_params.pop("aws_access_key_id", None)
@ -423,6 +425,16 @@ def completion(
model_response.usage.prompt_tokens = prompt_tokens model_response.usage.prompt_tokens = prompt_tokens
model_response.usage.total_tokens = prompt_tokens + completion_tokens model_response.usage.total_tokens = prompt_tokens + completion_tokens
return model_response return model_response
except BedrockError as e:
exception_mapping_worked = True
raise e
except Exception as e:
if exception_mapping_worked:
raise e
else:
import traceback
raise BedrockError(status_code=500, message=traceback.format_exc())
def embedding( def embedding(

5
litellm/llms/huggingface_restapi.py
View file

@ -141,6 +141,7 @@ def completion(
litellm_params=None, litellm_params=None,
logger_fn=None, logger_fn=None,
): ):
exception_mapping_worked = False
try: try:
headers = validate_environment(api_key, headers) headers = validate_environment(api_key, headers)
task = get_hf_task_for_model(model) task = get_hf_task_for_model(model)
@ -365,8 +366,12 @@ def completion(
model_response._hidden_params["original_response"] = completion_response model_response._hidden_params["original_response"] = completion_response
return model_response return model_response
except HuggingfaceError as e: except HuggingfaceError as e:
exception_mapping_worked = True
raise e raise e
except Exception as e: except Exception as e:
if exception_mapping_worked:
raise e
else:
import traceback import traceback
raise HuggingfaceError(status_code=500, message=traceback.format_exc()) raise HuggingfaceError(status_code=500, message=traceback.format_exc())