improvements to exception mapping

This commit is contained in:
Krrish Dholakia 2023-09-11 18:33:54 -07:00
parent cd6b5b229f
commit f946f61b4c
5 changed files with 73 additions and 6 deletions

View file

@ -43,13 +43,14 @@ Base case - we return the original exception.
| | ContextWindowExceededError | AuthenticationError | InvalidRequestError | RateLimitError | ServiceUnavailableError | | | ContextWindowExceededError | AuthenticationError | InvalidRequestError | RateLimitError | ServiceUnavailableError |
|---------------|----------------------------|---------------------|---------------------|---------------|-------------------------| |---------------|----------------------------|---------------------|---------------------|---------------|-------------------------|
| Anthropic | ✅ | ✅ | ✅ | ✅ | | | Anthropic | ✅ | ✅ | ✅ | ✅ | |
| OpenAI | ✅ | ✅ |✅ |✅ |✅ | | OpenAI | ✅ | ✅ |✅ |✅ |✅|
| Replicate | ✅ | ✅ | ✅ | ✅ | ✅ | | Replicate | ✅ | ✅ | ✅ | ✅ | ✅ |
| Cohere | ✅ | ✅ | | ✅ | | | Cohere | ✅ | ✅ | | ✅ | |
| Huggingface | ✅ | ✅ | ✅ | ✅ | | | Huggingface | ✅ | ✅ | ✅ | ✅ | |
| Openrouter | ✅ | ✅ | | ✅ | | | Openrouter | ✅ | ✅ | | ✅ | |
| AI21 | ✅ | ✅ | ✅ | ✅ | | | AI21 | ✅ | ✅ | ✅ | ✅ | |
| TogetherAI | ✅ | ✅ | ✅ | ✅ | | | TogetherAI | ✅ | ✅ | ✅ | ✅ | |
| AlephAlpha | ✅ | ✅ | ✅ | ✅ | ✅ |
> For a deeper understanding of these exceptions, you can check out [this](https://github.com/BerriAI/litellm/blob/d7e58d13bf9ba9edbab2ab2f096f3de7547f35fa/litellm/utils.py#L1544) implementation for additional insights. > For a deeper understanding of these exceptions, you can check out [this](https://github.com/BerriAI/litellm/blob/d7e58d13bf9ba9edbab2ab2f096f3de7547f35fa/litellm/utils.py#L1544) implementation for additional insights.

View file

@ -12,7 +12,7 @@ from litellm import (
embedding, embedding,
completion, completion,
# AuthenticationError, # AuthenticationError,
# InvalidRequestError, InvalidRequestError,
ContextWindowExceededError, ContextWindowExceededError,
# RateLimitError, # RateLimitError,
# ServiceUnavailableError, # ServiceUnavailableError,
@ -34,8 +34,7 @@ litellm.vertex_location = "us-central1"
# 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 = ["gpt-3.5-turbo"] models = ["command-nightly"]
test_model = "claude-instant-1"
# Test 1: Context Window Errors # Test 1: Context Window Errors
@pytest.mark.parametrize("model", models) @pytest.mark.parametrize("model", models)
@ -73,6 +72,9 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
elif model in litellm.openrouter_models: elif model in litellm.openrouter_models:
temporary_key = os.environ["OPENROUTER_API_KEY"] temporary_key = os.environ["OPENROUTER_API_KEY"]
os.environ["OPENROUTER_API_KEY"] = "bad-key" os.environ["OPENROUTER_API_KEY"] = "bad-key"
elif model in litellm.aleph_alpha_models:
temporary_key = os.environ["ALEPH_ALPHA_API_KEY"]
os.environ["ALEPH_ALPHA_API_KEY"] = "bad-key"
elif ( elif (
model model
== "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1" == "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1"
@ -115,8 +117,18 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th
os.environ["AI21_API_KEY"] = temporary_key os.environ["AI21_API_KEY"] = temporary_key
elif ("togethercomputer" in model): elif ("togethercomputer" in model):
os.environ["TOGETHERAI_API_KEY"] = temporary_key os.environ["TOGETHERAI_API_KEY"] = temporary_key
elif model in litellm.aleph_alpha_models:
os.environ["ALEPH_ALPHA_API_KEY"] = temporary_key
return return
# Test 3: Invalid Request Error
@pytest.mark.parametrize("model", models)
def test_invalid_request_error(model):
messages = [{"content": "hey, how's it going?", "role": "user"}]
with pytest.raises(InvalidRequestError):
completion(model=model, messages=messages, max_tokens="hello world")
# Test 3: Rate Limit Errors # Test 3: Rate Limit Errors
# def test_model_call(model): # def test_model_call(model):
# try: # try:

View file

@ -1786,6 +1786,20 @@ def exception_type(model, original_exception, custom_llm_provider):
llm_provider="cohere", llm_provider="cohere",
model=model model=model
) )
elif "invalid type:" in error_str:
exception_mapping_worked = True
raise InvalidRequestError(
message=f"CohereException - {original_exception.message}",
llm_provider="cohere",
model=model
)
elif "Unexpected server error" in error_str:
exception_mapping_worked = True
raise ServiceUnavailableError(
message=f"CohereException - {original_exception.message}",
llm_provider="cohere",
model=model
)
else: else:
if hasattr(original_exception, "status_code"): if hasattr(original_exception, "status_code"):
exception_mapping_worked = True exception_mapping_worked = True
@ -1938,6 +1952,46 @@ def exception_type(model, original_exception, custom_llm_provider):
llm_provider="together_ai", llm_provider="together_ai",
model=model model=model
) )
elif model in litellm.aleph_alpha_models:
if "This is longer than the model's maximum context length" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"AlephAlphaException - {original_exception.message}",
llm_provider="aleph_alpha",
model=model
)
elif hasattr(original_exception, "status_code"):
print(f"status code: {original_exception.status_code}")
if original_exception.status_code == 401:
exception_mapping_worked = True
raise AuthenticationError(
message=f"AlephAlphaException - {original_exception.message}",
llm_provider="aleph_alpha",
model=model
)
elif original_exception.status_code == 400:
exception_mapping_worked = True
raise InvalidRequestError(
message=f"AlephAlphaException - {original_exception.message}",
llm_provider="aleph_alpha",
model=model
)
elif original_exception.status_code == 429:
exception_mapping_worked = True
raise RateLimitError(
message=f"AlephAlphaException - {original_exception.message}",
llm_provider="aleph_alpha",
model=model
)
elif original_exception.status_code == 500:
exception_mapping_worked = True
raise ServiceUnavailableError(
message=f"AlephAlphaException - {original_exception.message}",
llm_provider="aleph_alpha",
model=model
)
raise original_exception
raise original_exception
elif custom_llm_provider == "vllm": elif custom_llm_provider == "vllm":
if hasattr(original_exception, "status_code"): if hasattr(original_exception, "status_code"):
if original_exception.status_code == 0: if original_exception.status_code == 0: