scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

feat use identity_pool for vertex

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-24 10:51:46 -07:00
parent a7bcc0e88e
commit 3228125215
1 changed files with 11 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/llms/vertex_httpx.py
View file

@ -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