Merge pull request #18806 from BerriAI/litellm_vertex_ai_api_key_support

[FEAT]: Add support for Vertex AI API keys

docs/my-website/docs/providers/vertex.md

@@ -35,6 +35,8 @@ 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:
@@ -47,7 +49,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
+  vertex_credentials=vertex_credentials_json # Can remove this is added VERTEXAI_API_KEY in env
 )
 ```
 
@@ -1329,15 +1331,41 @@ 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: 
 
-### **Dynamic Params**
+The simplest way to authenticate with Vertex AI. You can set:
+- `api_key` (str) - Your Vertex AI API key
 
-You can set:
+**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:
 - `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
@@ -1392,7 +1420,16 @@ model_list:
 
 ### **Environment Variables**
 
-You can set:
+#### 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:
+
 - `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

litellm/llms/gemini/common_utils.py

@@ -150,6 +150,15 @@ 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(

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -388,6 +388,10 @@ 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
@@ -400,7 +404,7 @@ class VertexBase:
                 stream=stream,
                 gemini_api_key=gemini_api_key,
             )
-            auth_header = None  # this field is not used for gemin
+            auth_header = None  # this field is not used for gemini
         else:
             vertex_location = self.get_vertex_region(
                 vertex_region=vertex_location,
@@ -409,14 +413,32 @@ class VertexBase:
 
             ### SET RUNTIME ENDPOINT ###
             version = "v1beta1" if should_use_v1beta1_features is True else "v1"
-            url, endpoint = _get_vertex_url(
-                mode=mode,
-                model=model,
-                stream=stream,
-                vertex_project=vertex_project,
-                vertex_location=vertex_location,
-                vertex_api_version=version,
-            )
+            
+            # 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,
+                )
 
         return self._check_custom_proxy(
             api_base=api_base,

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
+from .llms.gemini.common_utils import get_api_key_from_env, get_vertex_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,6 +3230,12 @@ 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 {})
@@ -3271,7 +3277,7 @@ def completion(  # type: ignore # noqa: PLR0915
                     vertex_location=vertex_ai_location,
                     vertex_project=vertex_ai_project,
                     vertex_credentials=vertex_credentials,
-                    gemini_api_key=None,
+                    gemini_api_key=vertex_api_key,  # Support for Vertex AI API Key
                     logging_obj=logging,
                     acompletion=acompletion,
                     timeout=timeout,

tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py

@@ -13,6 +13,7 @@ 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):
@@ -1048,3 +1049,139 @@ 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