feat - predibase exceptions

This commit is contained in:
Ishaan Jaff 2024-05-15 16:52:33 -07:00
parent f138c15859
commit b88ad61490

View file

@ -8155,8 +8155,11 @@ def exception_type(
# Common Extra information needed for all providers # Common Extra information needed for all providers
# We pass num retries, api_base, vertex_deployment etc to the exception here # We pass num retries, api_base, vertex_deployment etc to the exception here
################################################################################ ################################################################################
extra_information = ""
_api_base = litellm.get_api_base(model=model, optional_params=extra_kwargs) try:
_api_base = litellm.get_api_base(
model=model, optional_params=extra_kwargs
)
messages = litellm.get_first_chars_messages(kwargs=completion_kwargs) messages = litellm.get_first_chars_messages(kwargs=completion_kwargs)
_vertex_project = extra_kwargs.get("vertex_project") _vertex_project = extra_kwargs.get("vertex_project")
_vertex_location = extra_kwargs.get("vertex_location") _vertex_location = extra_kwargs.get("vertex_location")
@ -8182,6 +8185,9 @@ def exception_type(
extra_information = _add_key_name_and_team_to_alert( extra_information = _add_key_name_and_team_to_alert(
request_info=extra_information, metadata=_metadata request_info=extra_information, metadata=_metadata
) )
except:
# DO NOT LET this Block raising the original exception
pass
################################################################################ ################################################################################
# End of Common Extra information Needed for all providers # End of Common Extra information Needed for all providers
@ -8546,6 +8552,28 @@ def exception_type(
model=model, model=model,
response=original_exception.response, response=original_exception.response,
) )
elif custom_llm_provider == "predibase":
if "authorization denied for" in error_str:
exception_mapping_worked = True
# Predibase returns the raw API Key in the response - this block ensures it's not returned in the exception
if (
error_str is not None
and isinstance(error_str, str)
and "bearer" in error_str.lower()
):
# only keep the first 10 chars after the occurnence of "bearer"
_bearer_token_start_index = error_str.lower().find("bearer")
error_str = error_str[: _bearer_token_start_index + 14]
error_str += "XXXXXXX" + '"'
raise AuthenticationError(
message=f"PredibaseException: Authentication Error - {error_str}",
llm_provider="predibase",
model=model,
response=original_exception.response,
litellm_debug_info=extra_information,
)
elif custom_llm_provider == "bedrock": elif custom_llm_provider == "bedrock":
if ( if (
"too many tokens" in error_str "too many tokens" in error_str