From ff8e9aeb5c11e2c3560c6bec8ae97ce3fd5be633 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 10 Jan 2026 13:39:49 -0800 Subject: [PATCH] Revert "Add support for Vertex AI API keys" This reverts commit ad501048f391ee473fe1fa2e6b69095a076f4630. --- docs/my-website/docs/providers/vertex.md | 45 +----- litellm/llms/gemini/common_utils.py | 9 -- litellm/llms/vertex_ai/vertex_llm_base.py | 40 ++--- litellm/main.py | 10 +- .../llms/vertex_ai/test_vertex_llm_base.py | 137 ------------------ 5 files changed, 15 insertions(+), 226 deletions(-) diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index f46608aa57c..33ebf535d29 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -35,8 +35,6 @@ import json # !gcloud auth application-default login - run this to add vertex credentials to your env ## OR ## file_path = 'path/to/vertex_ai_service_account.json' -## OR ## -export VERTEXAI_API_KEY="your-api-key" # Load the JSON file with open(file_path, 'r') as file: @@ -49,7 +47,7 @@ vertex_credentials_json = json.dumps(vertex_credentials) response = completion( model="vertex_ai/gemini-2.5-pro", messages=[{ "content": "Hello, how are you?","role": "user"}], - vertex_credentials=vertex_credentials_json # Can remove this is added VERTEXAI_API_KEY in env + vertex_credentials=vertex_credentials_json ) ``` @@ -1331,41 +1329,15 @@ Here's how to use Vertex AI with the LiteLLM Proxy Server ## Authentication - vertex_project, vertex_location, etc. -LiteLLM supports two authentication methods for Vertex AI: - -1. **API Key Authentication** (Recommended for getting started) -2. **Service Account Credentials** (Recommended for production) - Set your vertex credentials via: - dynamic params OR - env vars -### **Authentication Method 1: -The simplest way to authenticate with Vertex AI. You can set: -- `api_key` (str) - Your Vertex AI API key +### **Dynamic Params** -**Environment Variables:** -```bash -export VERTEXAI_API_KEY="your-api-key" -``` - -**Or pass as parameters:** -```python -from litellm import completion - -response = completion( - model="vertex_ai/gemini-2.0-flash-exp", - messages=[{"role": "user", "content": "Hello!"}], - api_key="your-vertex-api-key", - -) -``` - -### **Authentication Method 2: Service Account Credentials** - -For production environments with fine-grained access control. You can set: +You can set: - `vertex_credentials` (str) - can be a json string or filepath to your vertex ai service account.json - `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.). Some models support the global location, please see [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#supported_models) - `vertex_project` Optional[str] - use if vertex project different from the one in vertex_credentials @@ -1420,16 +1392,7 @@ model_list: ### **Environment Variables** -#### For API Key Authentication: - -- `VERTEXAI_API_KEY` or `VERTEX_API_KEY` - Your Vertex AI API key - -```bash -export VERTEXAI_API_KEY="your-vertex-api-key" -``` - -#### For Service Account Authentication: - +You can set: - `GOOGLE_APPLICATION_CREDENTIALS` - store the filepath for your service_account.json in here (used by vertex sdk directly). - VERTEXAI_LOCATION - place where vertex model is deployed (us-central1, asia-southeast1, etc.) - VERTEXAI_PROJECT - Optional[str] - use if vertex project different from the one in vertex_credentials diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 30c5b4f17c5..e53829d3329 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -150,15 +150,6 @@ def get_api_key_from_env() -> Optional[str]: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") -def get_vertex_api_key_from_env() -> Optional[str]: - """ - Get API key from environment for Vertex AI. - Checks VERTEXAI_API_KEY and VERTEX_API_KEY environment variables. - This allows using Vertex AI with API keys instead of service account credentials. - """ - return get_secret_str("VERTEXAI_API_KEY") or get_secret_str("VERTEX_API_KEY") - - class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" def should_use_token_counting_api( diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index a3606ff9deb..826f151df35 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -388,10 +388,6 @@ class VertexBase: Internal function. Returns the token and url for the call. Handles logic if it's google ai studio vs. vertex ai. - - For Vertex AI: - - If gemini_api_key is provided, use API key authentication (x-goog-api-key header) - - Otherwise, use service account credentials (OAuth2 Bearer token) Returns token, url @@ -404,7 +400,7 @@ class VertexBase: stream=stream, gemini_api_key=gemini_api_key, ) - auth_header = None # this field is not used for gemini + auth_header = None # this field is not used for gemin else: vertex_location = self.get_vertex_region( vertex_region=vertex_location, @@ -413,32 +409,14 @@ class VertexBase: ### SET RUNTIME ENDPOINT ### version = "v1beta1" if should_use_v1beta1_features is True else "v1" - - # Check if using API key authentication for Vertex AI - if gemini_api_key and not vertex_credentials: - # When using API key with Vertex AI, use the Google AI Studio endpoint - # This is because Vertex AI API keys work with generativelanguage.googleapis.com - verbose_logger.debug( - f"Using Vertex AI API key authentication for model: {model} - routing to Google AI Studio endpoint" - ) - url, endpoint = _get_gemini_url( - mode=mode, - model=model, - stream=stream, - gemini_api_key=gemini_api_key, - ) - # API key is already included in the URL by _get_gemini_url - auth_header = None - else: - # Use OAuth2 Bearer token authentication (traditional Vertex AI) - url, endpoint = _get_vertex_url( - mode=mode, - model=model, - stream=stream, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_api_version=version, - ) + url, endpoint = _get_vertex_url( + mode=mode, + model=model, + stream=stream, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=version, + ) return self._check_custom_proxy( api_base=api_base, diff --git a/litellm/main.py b/litellm/main.py index 10e3bcac04b..6905d8f6a87 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -189,7 +189,7 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha, palm -from .llms.gemini.common_utils import get_api_key_from_env, get_vertex_api_key_from_env +from .llms.gemini.common_utils import get_api_key_from_env from .llms.groq.chat.handler import GroqChatCompletion from .llms.heroku.chat.transformation import HerokuChatConfig from .llms.huggingface.embedding.handler import HuggingFaceEmbedding @@ -3230,12 +3230,6 @@ def completion( # type: ignore # noqa: PLR0915 or get_secret("VERTEXAI_CREDENTIALS") ) - vertex_api_key = ( - api_key - or get_vertex_api_key_from_env() - or litellm.api_key - ) - api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") new_params = safe_deep_copy(optional_params or {}) @@ -3277,7 +3271,7 @@ def completion( # type: ignore # noqa: PLR0915 vertex_location=vertex_ai_location, vertex_project=vertex_ai_project, vertex_credentials=vertex_credentials, - gemini_api_key=vertex_api_key, # Support for Vertex AI API Key + gemini_api_key=None, logging_obj=logging, acompletion=acompletion, timeout=timeout, 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 80d65991acb..389c8446135 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 @@ -13,7 +13,6 @@ sys.path.insert( import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.llms.vertex_ai.common_utils import _get_gemini_url def run_sync(coro): @@ -1049,139 +1048,3 @@ class TestVertexBase: MockCredentials.from_info.assert_called_once_with(json_obj) mock_creds.with_scopes.assert_called_once_with(scopes) assert result == "scoped_creds" - - def test_get_token_and_url_with_api_key(self): - """Test that API key authentication routes to Google AI Studio endpoint""" - vertex_base = VertexBase() - - # Test with API key and no credentials - should use Google AI Studio endpoint - auth_header, url = vertex_base._get_token_and_url( - model="gemini-2.0-flash-exp", - auth_header=None, - gemini_api_key="test-api-key-123", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, # No service account credentials - stream=False, - custom_llm_provider="vertex_ai", - api_base=None, - should_use_v1beta1_features=False, - mode="chat", - ) - - # Should route to Google AI Studio endpoint - assert "generativelanguage.googleapis.com" in url - assert "gemini-2.0-flash-exp" in url - assert "key=test-api-key-123" in url - assert auth_header is None # API key is in URL, not header - - def test_get_token_and_url_with_credentials(self): - """Test that service account credentials route to Vertex AI endpoint""" - vertex_base = VertexBase() - - mock_creds = MagicMock() - mock_creds.token = "mock-bearer-token" - mock_creds.expired = False - - with patch.object( - vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") - ): - # Test with credentials - should use Vertex AI endpoint - auth_header, url = vertex_base._get_token_and_url( - model="gemini-2.0-flash-exp", - auth_header="mock-bearer-token", - gemini_api_key=None, - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials={"type": "service_account"}, - stream=False, - custom_llm_provider="vertex_ai", - api_base=None, - should_use_v1beta1_features=False, - mode="chat", - ) - - # Should route to Vertex AI endpoint - assert "aiplatform.googleapis.com" in url - assert "projects/test-project" in url - assert "locations/us-central1" in url - assert auth_header == "mock-bearer-token" - - def test_get_token_and_url_api_key_with_streaming(self): - """Test API key authentication with streaming enabled""" - vertex_base = VertexBase() - - auth_header, url = vertex_base._get_token_and_url( - model="gemini-2.0-flash-exp", - auth_header=None, - gemini_api_key="test-api-key-456", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - stream=True, # Streaming enabled - custom_llm_provider="vertex_ai", - api_base=None, - should_use_v1beta1_features=False, - mode="chat", - ) - - # Should route to Google AI Studio endpoint with streaming - assert "generativelanguage.googleapis.com" in url - assert "streamGenerateContent" in url - assert "key=test-api-key-456" in url - assert "alt=sse" in url - assert auth_header is None - - def test_get_token_and_url_api_key_priority(self): - """Test that credentials take priority over API key when both are provided""" - vertex_base = VertexBase() - - # When both API key and credentials are provided, credentials take priority - mock_creds = MagicMock() - mock_creds.token = "mock-bearer-token" - mock_creds.expired = False - - with patch.object( - vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") - ): - auth_header, url = vertex_base._get_token_and_url( - model="gemini-2.0-flash-exp", - auth_header="mock-bearer-token", - gemini_api_key="test-api-key-789", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials={"type": "service_account"}, # Credentials provided - stream=False, - custom_llm_provider="vertex_ai", - api_base=None, - should_use_v1beta1_features=False, - mode="chat", - ) - - # Should use Vertex AI endpoint with Bearer token (credentials take priority) - assert "aiplatform.googleapis.com" in url - assert auth_header == "mock-bearer-token" - - def test_get_token_and_url_with_embedding_mode(self): - """Test API key authentication with embedding mode""" - vertex_base = VertexBase() - - auth_header, url = vertex_base._get_token_and_url( - model="text-embedding-004", - auth_header=None, - gemini_api_key="test-embedding-key", - vertex_project="test-project", - vertex_location="us-central1", - vertex_credentials=None, - stream=False, - custom_llm_provider="vertex_ai", - api_base=None, - should_use_v1beta1_features=False, - mode="embedding", - ) - - # Should route to Google AI Studio endpoint for embeddings - assert "generativelanguage.googleapis.com" in url - assert "embedContent" in url - assert "key=test-embedding-key" in url - assert auth_header is None \ No newline at end of file