fix(vertex_ai.py): add exception mapping for acompletion calls

This commit is contained in:
Krrish Dholakia 2023-12-13 16:35:50 -08:00
parent b6b371b051
commit e678009695
5 changed files with 99 additions and 70 deletions

View file

@ -18,10 +18,12 @@ from openai import (
APIError,
APITimeoutError,
APIConnectionError,
APIResponseValidationError
APIResponseValidationError,
UnprocessableEntityError
)
import httpx
class AuthenticationError(AuthenticationError): # type: ignore
def __init__(self, message, llm_provider, model, response: httpx.Response):
self.status_code = 401
@ -46,6 +48,18 @@ class BadRequestError(BadRequestError): # type: ignore
body=None
) # Call the base class constructor with the parameters it needs
class UnprocessableEntityError(UnprocessableEntityError): # type: ignore
def __init__(self, message, model, llm_provider, response: httpx.Response):
self.status_code = 422
self.message = message
self.model = model
self.llm_provider = llm_provider
super().__init__(
self.message,
response=response,
body=None
) # Call the base class constructor with the parameters it needs
class Timeout(APITimeoutError): # type: ignore
def __init__(self, message, model, llm_provider):
self.status_code = 408