From 14d331e9ea08c7bd67b878635f1794dfd94bd7aa Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 6 Dec 2025 16:03:51 -0800 Subject: [PATCH] fix(vertex_llm_base.py): Add missing scopes to identity pool credentials (#17560) * fix(vertex_llm_base.py): Fixes https://github.com/BerriAI/litellm/issues/17377 * test: update tests --- litellm/llms/vertex_ai/vertex_llm_base.py | 50 +++-- .../llms/vertex_ai/test_vertex_llm_base.py | 201 ++++++++++++++---- 2 files changed, 186 insertions(+), 65 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 251d3c0c45..826f151df3 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -90,9 +90,15 @@ class VertexBase: else "" ) if isinstance(environment_id, str) and "aws" in environment_id: - creds = self._credentials_from_identity_pool_with_aws(json_obj) + creds = self._credentials_from_identity_pool_with_aws( + json_obj, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) else: - creds = self._credentials_from_identity_pool(json_obj) + creds = self._credentials_from_identity_pool( + json_obj, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) # Check if the JSON object contains Authorized User configuration (via gcloud auth application-default login) elif "type" in json_obj and json_obj["type"] == "authorized_user": creds = self._credentials_from_authorized_user( @@ -131,15 +137,21 @@ class VertexBase: return creds, project_id # Google Auth Helpers -- extracted for mocking purposes in tests - def _credentials_from_identity_pool(self, json_obj): + def _credentials_from_identity_pool(self, json_obj, scopes): from google.auth import identity_pool - return identity_pool.Credentials.from_info(json_obj) + creds = identity_pool.Credentials.from_info(json_obj) + if scopes and hasattr(creds, "requires_scopes") and creds.requires_scopes: + creds = creds.with_scopes(scopes) + return creds - def _credentials_from_identity_pool_with_aws(self, json_obj): + def _credentials_from_identity_pool_with_aws(self, json_obj, scopes): from google.auth import aws - return aws.Credentials.from_info(json_obj) + creds = aws.Credentials.from_info(json_obj) + if scopes and hasattr(creds, "requires_scopes") and creds.requires_scopes: + creds = creds.with_scopes(scopes) + return creds def _credentials_from_authorized_user(self, json_obj, scopes): import google.oauth2.credentials @@ -300,7 +312,7 @@ class VertexBase: ) -> Tuple[Optional[str], str]: """ for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317 - + Handles custom api_base for: 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint} 2. Vertex AI with standard proxies - constructs {api_base}:{endpoint} @@ -328,7 +340,7 @@ class VertexBase: "Missing gemini_api_key, please set `GEMINI_API_KEY`" ) if gemini_api_key is not None: - auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] + auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] else: # For Vertex AI if use_psc_endpoint_format: @@ -396,9 +408,7 @@ class VertexBase: ) ### SET RUNTIME ENDPOINT ### - version = ( - "v1beta1" if should_use_v1beta1_features is True else "v1" - ) + version = "v1beta1" if should_use_v1beta1_features is True else "v1" url, endpoint = _get_vertex_url( mode=mode, model=model, @@ -675,13 +685,13 @@ class VertexBase: def safe_get_vertex_ai_project(litellm_params: dict) -> Optional[str]: """ Safely get Vertex AI project without mutating the litellm_params dict. - + Unlike get_vertex_ai_project(), this does NOT pop values from the dict, making it safe to call multiple times with the same litellm_params. - + Args: litellm_params: Dictionary containing Vertex AI parameters - + Returns: Vertex AI project ID or None """ @@ -696,13 +706,13 @@ class VertexBase: def safe_get_vertex_ai_credentials(litellm_params: dict) -> Optional[str]: """ Safely get Vertex AI credentials without mutating the litellm_params dict. - + Unlike get_vertex_ai_credentials(), this does NOT pop values from the dict, making it safe to call multiple times with the same litellm_params. - + Args: litellm_params: Dictionary containing Vertex AI parameters - + Returns: Vertex AI credentials or None """ @@ -716,13 +726,13 @@ class VertexBase: def safe_get_vertex_ai_location(litellm_params: dict) -> Optional[str]: """ Safely get Vertex AI location without mutating the litellm_params dict. - + Unlike get_vertex_ai_location(), this does NOT pop values from the dict, making it safe to call multiple times with the same litellm_params. - + Args: litellm_params: Dictionary containing Vertex AI parameters - + Returns: Vertex AI location/region or None """ diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 129534ffc7..389c844613 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -286,6 +286,9 @@ class TestVertexBase: ) assert mock_credentials_from_identity_pool.called + assert mock_credentials_from_identity_pool.call_args[1]["scopes"] == [ + "https://www.googleapis.com/auth/cloud-platform" + ] assert token == "refreshed-token" @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) @@ -299,9 +302,7 @@ class TestVertexBase: "project_id": "test-project", "refresh_token": "fake-refresh-token", "type": "external_account", - "credential_source": { - "environment_id": "aws1" - } + "credential_source": {"environment_id": "aws1"}, } mock_creds = MagicMock() mock_creds.token = "token-1" @@ -309,7 +310,9 @@ class TestVertexBase: mock_creds.project_id = "test-project" with patch.object( - vertex_base, "_credentials_from_identity_pool_with_aws", return_value=mock_creds + vertex_base, + "_credentials_from_identity_pool_with_aws", + return_value=mock_creds, ) as mock_credentials_from_identity_pool_with_aws, patch.object( vertex_base, "refresh_auth" ) as mock_refresh: @@ -333,6 +336,9 @@ class TestVertexBase: ) assert mock_credentials_from_identity_pool_with_aws.called + assert mock_credentials_from_identity_pool_with_aws.call_args[1][ + "scopes" + ] == ["https://www.googleapis.com/auth/cloud-platform"] assert token == "refreshed-token" @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) @@ -447,7 +453,7 @@ class TestVertexBase: assert token == "token-1" assert project == "resolved-project" - # Verify both cache entries exist + # Verify both cache entries exist original_cache_key = (json.dumps(credentials), None) resolved_cache_key = (json.dumps(credentials), "resolved-project") @@ -455,8 +461,12 @@ class TestVertexBase: assert resolved_cache_key in vertex_base._credentials_project_mapping # Both should contain the same tuple - original_entry = vertex_base._credentials_project_mapping[original_cache_key] - resolved_entry = vertex_base._credentials_project_mapping[resolved_cache_key] + original_entry = vertex_base._credentials_project_mapping[ + original_cache_key + ] + resolved_entry = vertex_base._credentials_project_mapping[ + resolved_cache_key + ] assert isinstance(original_entry, tuple) assert isinstance(resolved_entry, tuple) @@ -594,7 +604,9 @@ class TestVertexBase: credentials = {"type": "service_account"} with patch.object( - vertex_base, "load_auth", return_value=(mock_creds, "resolved-from-credentials") + vertex_base, + "load_auth", + return_value=(mock_creds, "resolved-from-credentials"), ) as mock_load_auth: # First call: User provides NO project_id, should resolve from credentials @@ -624,8 +636,12 @@ class TestVertexBase: assert resolved_cache_key in vertex_base._credentials_project_mapping # Both should contain the tuple with resolved project_id - original_entry = vertex_base._credentials_project_mapping[original_cache_key] - resolved_entry = vertex_base._credentials_project_mapping[resolved_cache_key] + original_entry = vertex_base._credentials_project_mapping[ + original_cache_key + ] + resolved_entry = vertex_base._credentials_project_mapping[ + resolved_cache_key + ] assert isinstance(original_entry, tuple) assert isinstance(resolved_entry, tuple) @@ -719,7 +735,7 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", {"x-goog-api-key": "test-api-key"}, - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", ), # Test case 2: Gemini with custom API base and streaming ( @@ -732,7 +748,7 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", {"x-goog-api-key": "test-api-key"}, - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse", ), # Test case 3: Non-Gemini provider with custom API base ( @@ -745,7 +761,7 @@ class TestVertexBase: "https://aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent", "gemini-pro", "Bearer token123", - "https://custom-vertex-api.com:generateContent" + "https://custom-vertex-api.com:generateContent", ), # Test case 4: No API base provided (should return original values) ( @@ -758,7 +774,7 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", "Bearer token123", - "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", ), # Test case 5: Gemini without API key (should raise ValueError) ( @@ -771,26 +787,26 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", None, # This should raise an exception - None + None, ), ], ) def test_check_custom_proxy( - self, - api_base, - custom_llm_provider, - gemini_api_key, - endpoint, - stream, - auth_header, - url, - model, - expected_auth_header, - expected_url + self, + api_base, + custom_llm_provider, + gemini_api_key, + endpoint, + stream, + auth_header, + url, + model, + expected_auth_header, + expected_url, ): """Test the _check_custom_proxy method for handling custom API base URLs""" vertex_base = VertexBase() - + if custom_llm_provider == "gemini" and api_base and gemini_api_key is None: # Test case 5: Should raise ValueError for Gemini without API key with pytest.raises(ValueError, match="Missing gemini_api_key"): @@ -816,21 +832,37 @@ class TestVertexBase: url=url, model=model, ) - - assert result_auth_header == expected_auth_header, f"Expected auth_header {expected_auth_header}, got {result_auth_header}" - assert result_url == expected_url, f"Expected URL {expected_url}, got {result_url}" + + assert ( + result_auth_header == expected_auth_header + ), f"Expected auth_header {expected_auth_header}, got {result_auth_header}" + assert ( + result_url == expected_url + ), f"Expected URL {expected_url}, got {result_url}" def test_check_custom_proxy_gemini_url_construction(self): """Test that Gemini URLs are constructed correctly with custom API base""" vertex_base = VertexBase() - + # Test various Gemini models with custom API base test_cases = [ - ("gemini-2.5-flash-lite", "generateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"), - ("gemini-2.5-pro", "generateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"), - ("gemini-1.5-flash", "streamGenerateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent"), + ( + "gemini-2.5-flash-lite", + "generateContent", + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", + ), + ( + "gemini-2.5-pro", + "generateContent", + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent", + ), + ( + "gemini-1.5-flash", + "streamGenerateContent", + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent", + ), ] - + for model, endpoint, expected_url in test_cases: _, result_url = vertex_base._check_custom_proxy( api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", @@ -842,13 +874,15 @@ class TestVertexBase: url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}", model=model, ) - - assert result_url == expected_url, f"Expected {expected_url}, got {result_url} for model {model}" + + assert ( + result_url == expected_url + ), f"Expected {expected_url}, got {result_url} for model {model}" def test_check_custom_proxy_streaming_parameter(self): """Test that streaming parameter correctly adds ?alt=sse to URLs""" vertex_base = VertexBase() - + # Test with streaming enabled _, result_url_streaming = vertex_base._check_custom_proxy( api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", @@ -860,10 +894,12 @@ class TestVertexBase: url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", model="gemini-2.5-flash-lite", ) - + expected_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" - assert result_url_streaming == expected_streaming_url, f"Expected {expected_streaming_url}, got {result_url_streaming}" - + assert ( + result_url_streaming == expected_streaming_url + ), f"Expected {expected_streaming_url}, got {result_url_streaming}" + # Test with streaming disabled _, result_url_no_streaming = vertex_base._check_custom_proxy( api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", @@ -875,9 +911,11 @@ class TestVertexBase: url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", model="gemini-2.5-flash-lite", ) - + expected_no_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" - assert result_url_no_streaming == expected_no_streaming_url, f"Expected {expected_no_streaming_url}, got {result_url_no_streaming}" + assert ( + result_url_no_streaming == expected_no_streaming_url + ), f"Expected {expected_no_streaming_url}, got {result_url_no_streaming}" @pytest.mark.parametrize( "api_base, custom_llm_provider, gemini_api_key, endpoint, stream, auth_header, url, model, expected_auth_header, expected_url", @@ -893,7 +931,7 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", {"x-goog-api-key": "test-api-key"}, - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", ), # Test case 2: Gemini with custom API base and streaming ( @@ -906,7 +944,7 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", {"x-goog-api-key": "test-api-key"}, - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse", ), ], ) @@ -937,3 +975,76 @@ class TestVertexBase: ) assert result_auth_header == expected_auth_header assert result_url == expected_url + + def test_credentials_from_identity_pool_implementation(self): + """Test the actual implementation of _credentials_from_identity_pool""" + vertex_base = VertexBase() + json_obj = {"type": "external_account", "audience": "test"} + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + + # Mock the credentials object + mock_creds = MagicMock() + mock_creds.requires_scopes = True + mock_creds.with_scopes.return_value = "scoped_creds" + + # Mock the identity_pool module + with patch("google.auth.identity_pool.Credentials") as MockCredentials: + MockCredentials.from_info.return_value = mock_creds + + # Call the method + result = vertex_base._credentials_from_identity_pool(json_obj, scopes) + + # Verify calls + MockCredentials.from_info.assert_called_once_with(json_obj) + mock_creds.with_scopes.assert_called_once_with(scopes) + assert result == "scoped_creds" + + def test_credentials_from_identity_pool_no_scopes_needed(self): + """Test _credentials_from_identity_pool when scopes are not needed""" + vertex_base = VertexBase() + json_obj = {"type": "external_account", "audience": "test"} + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + + # Mock the credentials object + mock_creds = MagicMock() + mock_creds.requires_scopes = False # Scopes not required + + # Mock the identity_pool module + with patch("google.auth.identity_pool.Credentials") as MockCredentials: + MockCredentials.from_info.return_value = mock_creds + + # Call the method + result = vertex_base._credentials_from_identity_pool(json_obj, scopes) + + # Verify calls + MockCredentials.from_info.assert_called_once_with(json_obj) + mock_creds.with_scopes.assert_not_called() + assert result == mock_creds + + def test_credentials_from_identity_pool_with_aws_implementation(self): + """Test the actual implementation of _credentials_from_identity_pool_with_aws""" + vertex_base = VertexBase() + json_obj = { + "type": "external_account", + "credential_source": {"environment_id": "aws1"}, + } + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + + # Mock the credentials object + mock_creds = MagicMock() + mock_creds.requires_scopes = True + mock_creds.with_scopes.return_value = "scoped_creds" + + # Mock the aws module + with patch("google.auth.aws.Credentials") as MockCredentials: + MockCredentials.from_info.return_value = mock_creds + + # Call the method + result = vertex_base._credentials_from_identity_pool_with_aws( + json_obj, scopes + ) + + # Verify calls + MockCredentials.from_info.assert_called_once_with(json_obj) + mock_creds.with_scopes.assert_called_once_with(scopes) + assert result == "scoped_creds"