From 7bddbdd56e38464164d63d6a356fc392bb247687 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 10 Sep 2024 13:49:28 -0700 Subject: [PATCH 1/2] fix vertex only refresh auth when required --- .../vertex_and_google_ai_studio_gemini.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py index 000512d94..4fe194526 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py @@ -9,7 +9,17 @@ import types import uuid from enum import Enum from functools import partial -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + Union, +) import httpx # type: ignore import requests # type: ignore @@ -68,6 +78,11 @@ from .transformation import ( sync_transform_request_body, ) +if TYPE_CHECKING: + from google.auth.credentials import Credentials as GoogleCredentialsObject +else: + GoogleCredentialsObject = Any + class VertexAIConfig: """ @@ -811,7 +826,7 @@ class VertexLLM(BaseLLM): super().__init__() self.access_token: Optional[str] = None self.refresh_token: Optional[str] = None - self._credentials: Optional[Any] = None + self._credentials: Optional[GoogleCredentialsObject] = None self.project_id: Optional[str] = None self.async_handler: Optional[AsyncHTTPHandler] = None @@ -1139,7 +1154,8 @@ class VertexLLM(BaseLLM): if not self.project_id: self.project_id = project_id or cred_project_id else: - self.refresh_auth(self._credentials) + if self._credentials.expired or not self._credentials.token: + self.refresh_auth(self._credentials) if not self.project_id: self.project_id = self._credentials.project_id From 26ae86e59b593806c979b48a1c0debdb956d1178 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 10 Sep 2024 14:02:15 -0700 Subject: [PATCH 2/2] fix types for vertex project id --- .../gemini/vertex_and_google_ai_studio_gemini.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py index 4fe194526..02907789a 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/vertex_and_google_ai_studio_gemini.py @@ -1158,7 +1158,7 @@ class VertexLLM(BaseLLM): self.refresh_auth(self._credentials) if not self.project_id: - self.project_id = self._credentials.project_id + self.project_id = self._credentials.quota_project_id if not self.project_id: raise ValueError("Could not resolve project_id")