From 831745e71080b75ca972cde454a7ccbe5ca6749c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 28 Jun 2024 22:19:44 -0700 Subject: [PATCH] test(test_streaming.py): try-except replicate api instability --- litellm/tests/test_streaming.py | 2 ++ litellm/utils.py | 9 ++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index eec6929f2e..fa9e49f87c 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -1265,6 +1265,8 @@ async def test_completion_replicate_llama3_streaming(sync_mode): raise Exception("finish reason not set") if complete_response.strip() == "": raise Exception("Empty response received") + except litellm.UnprocessableEntityError as e: + pass except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index 70d0770632..fef4976ca6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6160,7 +6160,6 @@ def exception_type( ) elif ( original_exception.status_code == 400 - or original_exception.status_code == 422 or original_exception.status_code == 413 ): exception_mapping_worked = True @@ -6170,6 +6169,14 @@ def exception_type( llm_provider="replicate", response=original_exception.response, ) + elif original_exception.status_code == 422: + exception_mapping_worked = True + raise UnprocessableEntityError( + message=f"ReplicateException - {original_exception.message}", + model=model, + llm_provider="replicate", + response=original_exception.response, + ) elif original_exception.status_code == 408: exception_mapping_worked = True raise Timeout(