From 32281252158f2fcfbfb8e9b351589d7774ca8481 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 24 Aug 2024 10:51:46 -0700 Subject: [PATCH] feat use identity_pool for vertex --- litellm/llms/vertex_httpx.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index a88925330aa..f40f6e90c5f 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -832,6 +832,7 @@ class VertexLLM(BaseLLM): self, credentials: Optional[str], project_id: Optional[str] ) -> Tuple[Any, str]: import google.auth as google_auth + from google.auth import identity_pool from google.auth.credentials import Credentials # type: ignore[import-untyped] from google.auth.transport.requests import ( Request, # type: ignore[import-untyped] @@ -855,10 +856,16 @@ class VertexLLM(BaseLLM): else: json_obj = json.loads(credentials) - creds = google.oauth2.service_account.Credentials.from_service_account_info( - json_obj, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) + # Check if the JSON object contains Workload Identity Federation configuration + if "type" in json_obj and json_obj["type"] == "external_account": + creds = identity_pool.Credentials.from_info(json_obj) + else: + creds = ( + google.oauth2.service_account.Credentials.from_service_account_info( + json_obj, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) + ) if project_id is None: project_id = creds.project_id