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

Merge pull request #5354 from BerriAI/vertex_allow_auth

Browse source 
[Feat-Vertex] Support using  workload identity federation
This commit is contained in:
Ishaan Jaff 2024-08-24 16:27:29 -07:00 • committed by GitHub
commit 35740da03d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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