fix(vertex_ai.py): fix exception mapping for vertex ai

This commit is contained in:
Krrish Dholakia 2023-11-23 17:35:26 -08:00
parent 7fda40e2be
commit 78d13ea6eb
4 changed files with 107 additions and 87 deletions

View file

@ -1,7 +1,7 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Local OpenAI Proxy Server
# [OLD PROXY 👉 [**NEW** proxy here](./simple_proxy.md)] Local OpenAI Proxy Server
A fast, and lightweight OpenAI-compatible server to call 100+ LLM APIs.

View file

@ -73,7 +73,8 @@ def completion(
try:
import vertexai
except:
raise Exception("vertexai import failed please run `pip install google-cloud-aiplatform`")
raise VertexAIError(status_code=400,message="vertexai import failed please run `pip install google-cloud-aiplatform`")
try:
from vertexai.preview.language_models import ChatModel, CodeChatModel, InputOutputTextPair
from vertexai.language_models import TextGenerationModel, CodeGenerationModel
@ -160,6 +161,8 @@ def completion(
)
model_response.usage = usage
return model_response
except Exception as e:
raise VertexAIError(status_code=500, message=str(e))
def embedding():

View file

@ -64,7 +64,7 @@ def test_context_window_with_fallbacks(model):
# for model in litellm.models_by_provider["bedrock"]:
# test_context_window(model=model)
# test_context_window(model="command-nightly")
# test_context_window(model="chat-bison")
# test_context_window_with_fallbacks(model="command-nightly")
# Test 2: InvalidAuth Errors
@pytest.mark.parametrize("model", models)

View file

@ -3816,6 +3816,23 @@ def exception_type(
llm_provider="vertex_ai",
response=original_exception.response
)
if hasattr(original_exception, "status_code"):
if original_exception.status_code == 400:
exception_mapping_worked = True
raise BadRequestError(
message=f"VertexAIException - {error_str}",
model=model,
llm_provider="vertex_ai",
response=original_exception.response
)
if original_exception.status_code == 500:
exception_mapping_worked = True
raise APIError(
message=f"VertexAIException - {error_str}",
model=model,
llm_provider="vertex_ai",
request=original_exception.request
)
elif custom_llm_provider == "palm":
if "503 Getting metadata" in error_str:
# auth errors look like this