Merge pull request #4271 from BerriAI/litellm_vertex_httpx_fix

fix(vertex_httpx.py): Correctly handle Vertex content policy violation error
This commit is contained in:
Krish Dholakia 2024-06-18 20:19:46 -07:00 committed by GitHub
commit d96ffe8075
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 309 additions and 37 deletions

View file

@ -6240,7 +6240,11 @@ def exception_type(
llm_provider="sagemaker",
response=original_exception.response,
)
elif custom_llm_provider == "vertex_ai":
elif (
custom_llm_provider == "vertex_ai"
or custom_llm_provider == "vertex_ai_beta"
or custom_llm_provider == "gemini"
):
if (
"Vertex AI API has not been used in project" in error_str
or "Unable to find your project" in error_str
@ -6259,6 +6263,13 @@ def exception_type(
),
litellm_debug_info=extra_information,
)
if "400 Request payload size exceeds" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"VertexException - {error_str}",
model=model,
llm_provider=custom_llm_provider,
)
elif (
"None Unknown Error." in error_str
or "Content has no parts." in error_str
@ -6292,13 +6303,13 @@ def exception_type(
)
elif "The response was blocked." in error_str:
exception_mapping_worked = True
raise UnprocessableEntityError(
message=f"VertexAIException UnprocessableEntityError - {error_str}",
raise ContentPolicyViolationError(
message=f"VertexAIException ContentPolicyViolationError - {error_str}",
model=model,
llm_provider="vertex_ai",
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=422,
status_code=400,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
@ -6350,6 +6361,27 @@ def exception_type(
),
),
)
if original_exception.status_code == 401:
exception_mapping_worked = True
raise AuthenticationError(
message=f"VertexAIException - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
if original_exception.status_code == 404:
exception_mapping_worked = True
raise NotFoundError(
message=f"VertexAIException - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
if original_exception.status_code == 408:
exception_mapping_worked = True
raise Timeout(
message=f"VertexAIException - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
if original_exception.status_code == 429:
exception_mapping_worked = True
@ -6379,6 +6411,13 @@ def exception_type(
request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
if original_exception.status_code == 503:
exception_mapping_worked = True
raise ServiceUnavailableError(
message=f"VertexAIException - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
elif custom_llm_provider == "palm" or custom_llm_provider == "gemini":
if "503 Getting metadata" in error_str:
# auth errors look like this