From 0e26ef858b4dce31729e027bc03c45da988a2a8f Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 15 Jan 2024 21:43:16 -0800 Subject: [PATCH] (feat) sagemaker - map status code and message --- litellm/llms/sagemaker.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 1d341e7e92d..7b50b05af3f 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -184,7 +184,15 @@ def completion( CustomAttributes="accept_eula=true", ) except Exception as e: - raise SagemakerError(status_code=500, message=f"{str(e)}") + status_code = ( + getattr(e, "response", {}) + .get("ResponseMetadata", {}) + .get("HTTPStatusCode", 500) + ) + error_message = ( + getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) + ) + raise SagemakerError(status_code=status_code, message=error_message) response = response["Body"].read().decode("utf8") ## LOGGING @@ -358,7 +366,15 @@ def embedding( CustomAttributes="accept_eula=true", ) except Exception as e: - raise SagemakerError(status_code=500, message=f"{str(e)}") + status_code = ( + getattr(e, "response", {}) + .get("ResponseMetadata", {}) + .get("HTTPStatusCode", 500) + ) + error_message = ( + getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) + ) + raise SagemakerError(status_code=status_code, message=error_message) response = json.loads(response["Body"].read().decode("utf8")) ## LOGGING @@ -368,7 +384,7 @@ def embedding( original_response=response, additional_args={"complete_input_dict": data}, ) - + print_verbose(f"raw model_response: {response}") if "embedding" not in response: raise SagemakerError(status_code=500, message="embedding not found in response")