diff --git a/litellm/integrations/gcs_bucket_base.py b/litellm/integrations/gcs_bucket_base.py index 15f5cbf68..ba7e77f22 100644 --- a/litellm/integrations/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket_base.py @@ -38,6 +38,7 @@ class GCSBucketBase(CustomLogger): await vertex_chat_completion._ensure_access_token_async( credentials=self.path_service_account_json, project_id=None, + custom_llm_provider="vertex_ai", ) ) @@ -66,6 +67,7 @@ class GCSBucketBase(CustomLogger): _auth_header, vertex_project = vertex_chat_completion._ensure_access_token( credentials=self.path_service_account_json, project_id=None, + custom_llm_provider="vertex_ai", ) auth_header, _ = vertex_chat_completion._get_token_and_url( diff --git a/litellm/llms/fine_tuning_apis/vertex_ai.py b/litellm/llms/fine_tuning_apis/vertex_ai.py index e24fd3894..b54519412 100644 --- a/litellm/llms/fine_tuning_apis/vertex_ai.py +++ b/litellm/llms/fine_tuning_apis/vertex_ai.py @@ -184,10 +184,10 @@ class VertexFineTuningAPI(VertexLLM): verbose_logger.debug( "creating fine tuning job, args= %s", create_fine_tuning_job_data ) - _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, + custom_llm_provider="vertex_ai_beta", ) auth_header, _ = self._get_token_and_url( @@ -257,10 +257,10 @@ class VertexFineTuningAPI(VertexLLM): vertex_credentials: str, request_route: str, ): - _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, + custom_llm_provider="vertex_ai_beta", ) auth_header, _ = self._get_token_and_url( model="", 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 8d69725dd..231c39079 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 @@ -1079,7 +1079,9 @@ class VertexLLM(VertexBase): ) _auth_header, vertex_project = await self._ensure_access_token_async( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider=custom_llm_provider, ) auth_header, api_base = self._get_token_and_url( @@ -1157,7 +1159,9 @@ class VertexLLM(VertexBase): ) _auth_header, vertex_project = await self._ensure_access_token_async( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider=custom_llm_provider, ) auth_header, api_base = self._get_token_and_url( @@ -1310,7 +1314,9 @@ class VertexLLM(VertexBase): ) _auth_header, vertex_project = self._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider=custom_llm_provider, ) auth_header, url = self._get_token_and_url( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini_embeddings/batch_embed_content_handler.py index d5d3cf1ec..1904ff7e4 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini_embeddings/batch_embed_content_handler.py @@ -46,6 +46,7 @@ class GoogleBatchEmbeddings(VertexLLM): _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, + custom_llm_provider=custom_llm_provider, ) auth_header, url = self._get_token_and_url( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py index 440d0841a..e1969199b 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/image_generation_handler.py @@ -82,7 +82,9 @@ class VertexImageGeneration(VertexLLM): url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:predict" auth_header, _ = self._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai", ) optional_params = optional_params or { "sampleCount": 1 @@ -180,7 +182,9 @@ class VertexImageGeneration(VertexLLM): "https://us-central1-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/us-central1/publishers/google/models/imagegeneration:predict" """ auth_header, _ = self._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai", ) optional_params = optional_params or { "sampleCount": 1 diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/multimodal_embeddings/embedding_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/multimodal_embeddings/embedding_handler.py index aa8c2123a..180939556 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/multimodal_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/multimodal_embeddings/embedding_handler.py @@ -47,6 +47,7 @@ class VertexMultimodalEmbedding(VertexLLM): _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, + custom_llm_provider=custom_llm_provider, ) auth_header, url = self._get_token_and_url( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/text_to_speech/text_to_speech_handler.py b/litellm/llms/vertex_ai_and_google_ai_studio/text_to_speech/text_to_speech_handler.py index 8818d13bc..58c37d24c 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/text_to_speech/text_to_speech_handler.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/text_to_speech/text_to_speech_handler.py @@ -65,10 +65,10 @@ class VertexTextToSpeechAPI(VertexLLM): import base64 ####### Authenticate with Vertex AI ######## - _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, + custom_llm_provider="vertex_ai_beta", ) auth_header, _ = self._get_token_and_url( diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py index 025b27240..13d8edeb9 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_anthropic.py @@ -292,7 +292,9 @@ def completion( vertex_httpx_logic = VertexLLM() access_token, project_id = vertex_httpx_logic._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai", ) anthropic_chat_completions = AnthropicChatCompletion() diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py index c30fa900f..e26a95e8b 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_ai_partner_models/main.py @@ -105,7 +105,9 @@ class VertexAIPartnerModels(BaseLLM): vertex_httpx_logic = VertexLLM() access_token, project_id = vertex_httpx_logic._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai", ) openai_like_chat_completions = DatabricksChatCompletion() diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_llm_base.py b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_llm_base.py index 3c817e9ea..9a828d710 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/vertex_llm_base.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/vertex_llm_base.py @@ -109,11 +109,18 @@ class VertexBase(BaseLLM): credentials.refresh(Request()) def _ensure_access_token( - self, credentials: Optional[str], project_id: Optional[str] + self, + credentials: Optional[str], + project_id: Optional[str], + custom_llm_provider: Literal[ + "vertex_ai", "vertex_ai_beta", "gemini" + ], # if it's vertex_ai or gemini (google ai studio) ) -> Tuple[str, str]: """ Returns auth token and project id """ + if custom_llm_provider == "gemini": + return "", "" if self.access_token is not None: if project_id is not None: return self.access_token, project_id @@ -222,11 +229,18 @@ class VertexBase(BaseLLM): return auth_header, url async def _ensure_access_token_async( - self, credentials: Optional[str], project_id: Optional[str] + self, + credentials: Optional[str], + project_id: Optional[str], + custom_llm_provider: Literal[ + "vertex_ai", "vertex_ai_beta", "gemini" + ], # if it's vertex_ai or gemini (google ai studio) ) -> Tuple[str, str]: """ Async version of _ensure_access_token """ + if custom_llm_provider == "gemini": + return "", "" if self.access_token is not None: if project_id is not None: return self.access_token, project_id diff --git a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py index 68d4146f2..86ffda67e 100644 --- a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py @@ -152,7 +152,9 @@ async def vertex_proxy_route( _auth_header, vertex_project = ( await vertex_fine_tuning_apis_instance._ensure_access_token_async( - credentials=vertex_credentials, project_id=vertex_project + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai_beta", ) )