diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index bf7b112ca4..6bebd64510 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/llms/ai21.py b/litellm/llms/ai21.py index 6a22b99e89..882e217e6b 100644 --- a/litellm/llms/ai21.py +++ b/litellm/llms/ai21.py @@ -102,7 +102,7 @@ class AI21LLM: try: model_response["choices"][0]["message"]["content"] = completion_response["completions"][0]["data"]["text"] except: - raise ValueError(f"Unable to parse response. Original response: {response.text}") + raise AI21Error(message=json.dumps(completion_response), status_code=response.status_code) ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. prompt_tokens = len( diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py index 396cdc0790..263d4e0d2a 100644 --- a/litellm/tests/test_exceptions.py +++ b/litellm/tests/test_exceptions.py @@ -33,8 +33,8 @@ litellm.failure_callback = ["sentry"] # Approach: Run each model through the test -> assert if the correct error (always the same one) is triggered # models = ["gpt-3.5-turbo", "chatgpt-test", "claude-instant-1", "command-nightly"] -test_model = "togethercomputer/CodeLlama-34b-Python" -models = ["togethercomputer/CodeLlama-34b-Python"] +test_model = "j2-light" +models = ["j2-light"] def logging_fn(model_call_dict): @@ -98,6 +98,9 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th elif model == "command-nightly": temporary_key = os.environ["COHERE_API_KEY"] os.environ["COHERE_API_KEY"] = "bad-key" + elif "j2" in model: + temporary_key = os.environ["AI21_API_KEY"] + os.environ["AI21_API_KEY"] = "bad-key" elif "togethercomputer" in model: temporary_key = os.environ["TOGETHERAI_API_KEY"] os.environ["TOGETHERAI_API_KEY"] = "84060c79880fc49df126d3e87b53f8a463ff6e1c6d27fe64207cde25cdfcd1f24a" @@ -138,6 +141,8 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th == "replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1" ): os.environ["REPLICATE_API_KEY"] = temporary_key + elif "j2" in model: + os.environ["AI21_API_KEY"] = temporary_key elif ("togethercomputer" in model): os.environ["TOGETHERAI_API_KEY"] = temporary_key return diff --git a/litellm/utils.py b/litellm/utils.py index 994b4916c9..18c32d58b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1446,6 +1446,37 @@ def exception_type(model, original_exception, custom_llm_provider): message=f"HuggingfaceException - {original_exception.message}", llm_provider="huggingface", ) + elif custom_llm_provider == "ai21": + print(f"e: {original_exception}") + if hasattr(original_exception, "message"): + if "Prompt has too many tokens" in original_exception.message: + exception_mapping_worked = True + raise ContextWindowExceededError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21" + ) + if 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"AI21Exception - {original_exception.message}", + llm_provider="ai21", + ) + if original_exception.status_code == 422 or "Prompt has too many tokens" in original_exception.message: + exception_mapping_worked = True + raise InvalidRequestError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + ) + elif original_exception.status_code == 429: + exception_mapping_worked = True + raise RateLimitError( + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + ) elif custom_llm_provider == "together_ai": error_response = json.loads(error_str) if "error" in error_response and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"]: