diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 33ebf535d2..f46608aa57 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/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 diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e53829d332..30c5b4f17c 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/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( diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 826f151df3..a3606ff9de 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/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, diff --git a/litellm/main.py b/litellm/main.py index e8a8b504d9..0264d0a7a7 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 +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, 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 389c844613..80d65991ac 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,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 \ No newline at end of file