feat: add workload identity federation between GCP and AWS

Author: pascallim

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -12,6 +12,7 @@
 from litellm.litellm_core_utils.asyncify import asyncify
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from google.auth import aws, identity_pool
 
 from .common_utils import _get_gemini_url, _get_vertex_url, all_gemini_url_modes
 
@@ -41,7 +42,6 @@ def load_auth(
         self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str]
     ) -> Tuple[Any, str]:
         import google.auth as google_auth
-        from google.auth import identity_pool
         from google.auth.transport.requests import (
             Request,  # type: ignore[import-untyped]
         )
@@ -80,7 +80,15 @@ def load_auth(
 
             # 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)
+                # If environment_id key contains "aws" value it corresponds to an AWS config file
+                if (
+                    "credential_source" in json_obj
+                    and "environment_id" in json_obj["credential_source"]
+                    and "aws" in json_obj["credential_source"]["environment_id"]
+                ):
+                    creds = aws.Credentials.from_info(json_obj)
+                else:
+                    creds = identity_pool.Credentials.from_info(json_obj)
             else:
                 creds = (
                     google.oauth2.service_account.Credentials.from_service_account_info(

tests/litellm/llms/vertex_ai/test_vertex_llm_base.py

@@ -174,3 +174,40 @@ async def test_gemini_credentials(self, is_async):
             )
         assert token == ""
         assert project == ""
+
+
+    @patch("litellm.llms.vertex_ai.vertex_llm_base.aws.Credentials.from_info")
+    @patch("litellm.llms.vertex_ai.vertex_llm_base.identity_pool.Credentials.from_info")
+    def test_load_auth_wif(self, mock_identity_pool, mock_aws):
+        vertex_base = VertexBase()
+        input_project_id = "some_project_id"
+
+        # Test case 1: Using Workload Identity Federation for Microsoft Azure and
+        # OIDC identity providers (default behavior)
+        json_obj_1  = {
+            "type": "external_account",
+        }
+        mock_creds_1 = MagicMock()
+        mock_identity_pool.return_value = mock_creds_1
+
+        creds_1, project_id= vertex_base.load_auth(
+            credentials=json_obj_1, project_id=input_project_id
+        )
+        mock_identity_pool.assert_called_once_with(json_obj_1)
+        assert creds_1 == mock_creds_1
+        assert project_id == input_project_id
+
+        # Test case 2: Using Workload Identity Federation for AWS
+        json_obj_2  = {
+            "type": "external_account",
+            "credential_source": {"environment_id": "aws1"}
+        }
+        mock_creds_2 = MagicMock()
+        mock_aws.return_value = mock_creds_2
+
+        creds_2, project_id= vertex_base.load_auth(
+            credentials=json_obj_2, project_id=input_project_id
+        )
+        mock_aws.assert_called_once_with(json_obj_2)
+        assert creds_2 == mock_creds_2
+        assert project_id == input_project_id