(feat) improve bedrock, sagemaker exception mapping

This commit is contained in:
ishaan-jaff 2024-01-15 21:22:22 -08:00
parent 893a2c3c7b
commit 00ac18e8b7
2 changed files with 18 additions and 2 deletions

View file

@ -35,7 +35,10 @@ litellm.num_retries = 0
# Approach: Run each model through the test -> assert if the correct error (always the same one) is triggered # Approach: Run each model through the test -> assert if the correct error (always the same one) is triggered
models = ["command-nightly"] models = [
"sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
"bedrock/anthropic.claude-instant-v1",
]
# Test 1: Context Window Errors # Test 1: Context Window Errors
@ -45,7 +48,8 @@ def test_context_window(model):
sample_text = "Say error 50 times" * 1000000 sample_text = "Say error 50 times" * 1000000
messages = [{"content": sample_text, "role": "user"}] messages = [{"content": sample_text, "role": "user"}]
try: try:
litellm.set_verbose = True litellm.set_verbose = False
print("Testing model=", model)
response = completion(model=model, messages=messages) response = completion(model=model, messages=messages)
print(f"response: {response}") print(f"response: {response}")
print("FAILED!") print("FAILED!")

View file

@ -5895,6 +5895,7 @@ def exception_type(
"too many tokens" in error_str "too many tokens" in error_str
or "expected maxLength:" in error_str or "expected maxLength:" in error_str
or "Input is too long" in error_str or "Input is too long" in error_str
or "prompt: length: 1.." in error_str
or "Too many input tokens" in error_str or "Too many input tokens" in error_str
): ):
exception_mapping_worked = True exception_mapping_worked = True
@ -5988,6 +5989,17 @@ def exception_type(
llm_provider="sagemaker", llm_provider="sagemaker",
response=original_exception.response, response=original_exception.response,
) )
elif (
"`inputs` tokens + `max_new_tokens` must be <=" in error_str
or "instance type with more CPU capacity or memory" in error_str
):
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"SagemakerException - {error_str}",
model=model,
llm_provider="sagemaker",
response=original_exception.response,
)
elif custom_llm_provider == "vertex_ai": elif custom_llm_provider == "vertex_ai":
if ( if (
"Vertex AI API has not been used in project" in error_str "Vertex AI API has not been used in project" in error_str