fix linting errors

This commit is contained in:
Ishaan Jaff 2024-09-02 19:39:10 -07:00
parent f5512b1e8d
commit 778cba702e
2 changed files with 5 additions and 2 deletions

View file

@ -16,7 +16,7 @@ class VertexImageGeneration(VertexLLM):
self, self,
json_response: Dict[str, Any], json_response: Dict[str, Any],
model_response: litellm.ImageResponse, model_response: litellm.ImageResponse,
model: str, model: Optional[str] = None,
) -> litellm.ImageResponse: ) -> litellm.ImageResponse:
if "predictions" not in json_response: if "predictions" not in json_response:
raise litellm.InternalServerError( raise litellm.InternalServerError(

View file

@ -1,5 +1,6 @@
import re import re
from datetime import datetime from datetime import datetime
from typing import Union
import httpx import httpx
@ -117,7 +118,9 @@ class PassThroughEndpointLogging:
model = self.extract_model_from_url(url_route) model = self.extract_model_from_url(url_route)
_json_response = httpx_response.json() _json_response = httpx_response.json()
litellm_model_response = litellm.ModelResponse() litellm_model_response: Union[
litellm.ModelResponse, litellm.EmbeddingResponse, litellm.ImageResponse
] = litellm.ModelResponse()
if vertex_image_generation_class.is_image_generation_response( if vertex_image_generation_class.is_image_generation_response(
_json_response _json_response
): ):