From 39d59f1900d680e771cdf2d216d2ed4386eac8d4 Mon Sep 17 00:00:00 2001 From: Siddharth Sahu <112792547+sahusiddharth@users.noreply.github.com> Date: Wed, 30 Jul 2025 04:26:01 +0530 Subject: [PATCH] Fix/gemini api key environment variable support (#12507) * Fix: Add support for GOOGLE_API_KEY environment variables for Gemini API authentication * added test cases * incoperated feedback to make it more maintainable * fix failed linting CI --- litellm/llms/gemini/common_utils.py | 8 +- .../gemini/google_genai/transformation.py | 5 +- .../llms/gemini/realtime/transformation.py | 89 ++++++++------- litellm/main.py | 5 +- .../llm_passthrough_endpoints.py | 2 +- litellm/utils.py | 5 +- .../llms/gemini/test_gemini_client_setup.py | 101 ++++++++++++++++++ .../proxy/test_spend_log_cleanup.py | 2 +- 8 files changed, 166 insertions(+), 51 deletions(-) create mode 100644 tests/test_litellm/llms/gemini/test_gemini_client_setup.py diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 31b57434e1..53de6711ca 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -44,7 +44,7 @@ class GeminiModelInfo(BaseLLMModelInfo): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return api_key or (get_secret_str("GEMINI_API_KEY")) + return api_key or (get_secret_str("GOOGLE_API_KEY")) or (get_secret_str("GEMINI_API_KEY")) @staticmethod def get_base_model(model: str) -> Optional[str]: @@ -66,7 +66,7 @@ class GeminiModelInfo(BaseLLMModelInfo): endpoint = f"/{self.api_version}/models" if api_base is None or api_key is None: raise ValueError( - "GEMINI_API_BASE or GEMINI_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint." + "GEMINI_API_BASE or GEMINI_API_KEY/GOOGLE_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint." ) response = litellm.module_level_client.get( @@ -133,3 +133,7 @@ def encode_unserializable_types( else: processed_data[key] = value return processed_data + + +def get_api_key_from_env() -> Optional[str]: + return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 0de6ef9043..5a07dd13a9 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -11,7 +11,6 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM -from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams if TYPE_CHECKING: @@ -24,7 +23,7 @@ else: GenerateContentConfigDict = Any GenerateContentContentListUnionDict = Any GenerateContentResponse = Any - +from ..common_utils import get_api_key_from_env class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ @@ -131,7 +130,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): return ( litellm_params.pop("api_key", None) or litellm_params.pop("gemini_api_key", None) - or get_secret_str("GEMINI_API_KEY") + or get_api_key_from_env() or litellm.api_key ) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 980723eb3f..f32a404c9e 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -3,7 +3,6 @@ This file contains the transformation logic for the Gemini realtime API. """ import json -import os import uuid from typing import Any, Dict, List, Optional, Union, cast @@ -55,7 +54,7 @@ from litellm.types.realtime import ( ) from litellm.utils import get_empty_usage -from ..common_utils import encode_unserializable_types +from ..common_utils import encode_unserializable_types, get_api_key_from_env MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Dict[str, OpenAIRealtimeEventTypes] = { "setupComplete": OpenAIRealtimeEventTypes.SESSION_CREATED, @@ -81,7 +80,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if api_base is None: api_base = "wss://generativelanguage.googleapis.com" if api_key is None: - api_key = os.environ.get("GEMINI_API_KEY") + api_key = get_api_key_from_env() if api_key is None: raise ValueError("api_key is required for Gemini API calls") api_base = api_base.replace("https://", "wss://") @@ -188,9 +187,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): vertex_gemini_config = VertexGeminiConfig() vertex_gemini_config._map_function(value) - optional_params["generationConfig"][ - "tools" - ] = vertex_gemini_config._map_function(value) + optional_params["generationConfig"]["tools"] = ( + vertex_gemini_config._map_function(value) + ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} elif key == "turn_detection": @@ -201,10 +200,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if ( len(transformed_audio_activity_config) > 0 ): # if the config is not empty, add it to the optional params - optional_params[ - "realtimeInputConfig" - ] = BidiGenerateContentRealtimeInputConfig( - automaticActivityDetection=transformed_audio_activity_config + optional_params["realtimeInputConfig"] = ( + BidiGenerateContentRealtimeInputConfig( + automaticActivityDetection=transformed_audio_activity_config + ) ) if len(optional_params["generationConfig"]) == 0: optional_params.pop("generationConfig") @@ -405,15 +404,17 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): output_index=0, event_id="event_{}".format(uuid.uuid4()), item_id=output_item_id, - part={ - "type": "text", - "text": "", - } - if delta_type == "text" - else { - "type": "audio", - "transcript": "", - }, + part=( + { + "type": "text", + "text": "", + } + if delta_type == "text" + else { + "type": "audio", + "transcript": "", + } + ), response_id=response_id, ) response_items.append(response_content_part_added) @@ -440,9 +441,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) return OpenAIRealtimeResponseDelta( - type="response.text.delta" - if delta_type == "text" - else "response.audio.delta", + type=( + "response.text.delta" + if delta_type == "text" + else "response.audio.delta" + ), content_index=0, event_id="event_{}".format(uuid.uuid4()), item_id=output_item_id, @@ -513,12 +516,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): event_id="event_{}".format(uuid.uuid4()), item_id=current_output_item_id, output_index=0, - part={"type": "text", "text": delta_done_event_text} - if delta_done_event_text and delta_type == "text" - else { - "type": "audio", - "transcript": "", # gemini doesn't return transcript for audio - }, + part=( + {"type": "text", "text": delta_done_event_text} + if delta_done_event_text and delta_type == "text" + else { + "type": "audio", + "transcript": "", # gemini doesn't return transcript for audio + } + ), response_id=current_response_id, ) returned_items.append(response_content_part_done) @@ -535,12 +540,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "status": "completed", "role": "assistant", "content": [ - {"type": "text", "text": delta_done_event_text} - if delta_done_event_text and delta_type == "text" - else { - "type": "audio", - "transcript": "", - } + ( + {"type": "text", "text": delta_done_event_text} + if delta_done_event_text and delta_type == "text" + else { + "type": "audio", + "transcript": "", + } + ) ], }, ) @@ -674,9 +681,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): object="realtime.response", id=current_response_id, status="completed", - output=[output_item["item"] for output_item in output_items] - if output_items - else [], + output=( + [output_item["item"] for output_item in output_items] + if output_items + else [] + ), conversation_id=current_conversation_id, modalities=_modalities, usage=responses_api_usage.model_dump(), @@ -828,9 +837,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "session_configuration_request" ] current_item_chunks = realtime_response_transform_input["current_item_chunks"] - current_delta_type: Optional[ - ALL_DELTA_TYPES - ] = realtime_response_transform_input["current_delta_type"] + current_delta_type: Optional[ALL_DELTA_TYPES] = ( + realtime_response_transform_input["current_delta_type"] + ) returned_message: List[OpenAIRealtimeEvents] = [] for key, value in json_message.items(): diff --git a/litellm/main.py b/litellm/main.py index fb6204c5bf..12d9825306 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -148,6 +148,7 @@ 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.groq.chat.handler import GroqChatCompletion +from .llms.gemini.common_utils import get_api_key_from_env from .llms.huggingface.embedding.handler import HuggingFaceEmbedding from .llms.nlp_cloud.chat.handler import completion as nlp_cloud_chat_completion from .llms.ollama.completion import handler as ollama @@ -2595,7 +2596,7 @@ def completion( # type: ignore # noqa: PLR0915 gemini_api_key = ( api_key - or get_secret("GEMINI_API_KEY") + or get_api_key_from_env() or get_secret("PALM_API_KEY") # older palm api key should also work or litellm.api_key ) @@ -3999,7 +4000,7 @@ def embedding( # noqa: PLR0915 ) elif custom_llm_provider == "gemini": gemini_api_key = ( - api_key or get_secret_str("GEMINI_API_KEY") or litellm.api_key + api_key or get_api_key_from_env() or litellm.api_key ) api_base = api_base or litellm.api_base or get_secret_str("GEMINI_API_BASE") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f85ecb47a2..eb06d7c8f4 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -190,7 +190,7 @@ async def gemini_proxy_route( ) if gemini_api_key is None: raise Exception( - "Required 'GEMINI_API_KEY' in environment to make pass-through calls to Google AI Studio." + "Required 'GEMINI_API_KEY'/'GOOGLE_API_KEY' in environment to make pass-through calls to Google AI Studio." ) # Merge query parameters, giving precedence to those in updated_url merged_params = dict(request.query_params) diff --git a/litellm/utils.py b/litellm/utils.py index b5ddce479f..b51be7af1c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5279,10 +5279,11 @@ def validate_environment( # noqa: PLR0915 else: missing_keys.append("FEATHERLESS_AI_API_KEY") elif custom_llm_provider == "gemini": - if "GEMINI_API_KEY" in os.environ: + if ("GOOGLE_API_KEY" in os.environ) or ("GEMINI_API_KEY" in os.environ): keys_in_environment = True else: - missing_keys.append("GEMINI_API_KEY") + missing_keys.append("GOOGLE_API_KEY") + missing_keys.append("GEMINI_API_KEY") elif custom_llm_provider == "groq": if "GROQ_API_KEY" in os.environ: keys_in_environment = True diff --git a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py new file mode 100644 index 0000000000..51c6fedf5b --- /dev/null +++ b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py @@ -0,0 +1,101 @@ +import pytest +import litellm +import os +from unittest.mock import patch, Mock +from litellm import completion + + +@pytest.fixture(autouse=True) +def mock_gemini_api_key(monkeypatch): + monkeypatch.setenv("GOOGLE_API_KEY", "fake-gemini-key-for-testing") + + +def test_gemini_completion(): + response = completion( + model="gemini/gemini-2.0-flash-exp-image-generation", + messages=[{"role": "user", "content": "Test message"}], + mock_response="Test Message", + ) + assert response.choices[0].message.content is not None + + +def test_gemini_completion_no_api_key(): + """Test Gemini completion fails gracefully when no API key is provided.""" + with patch.dict(os.environ, {}, clear=True): + # Remove all API keys + for key in ["GOOGLE_API_KEY", "GEMINI_API_KEY"]: + if key in os.environ: + del os.environ[key] + + # Test without mock_response to ensure actual API key validation + with pytest.raises(Exception) as exc_info: + completion( + model="gemini/gemini-1.5-flash", + messages=[{"role": "user", "content": "Test message"}], + ) + + # Check that the exception message contains API key related text + error_message = str(exc_info.value).lower() + assert any( + keyword in error_message + for keyword in [ + "api key", + "authentication", + "unauthorized", + "invalid", + "missing", + "credential", + ] + ) + + +def test_gemini_completion_no_api_key_with_mock(): + """Alternative test that properly mocks the API key validation.""" + with patch.dict(os.environ, {}, clear=True): + # Remove all API keys + for key in ["GOOGLE_API_KEY", "GEMINI_API_KEY"]: + if key in os.environ: + del os.environ[key] + + with patch("litellm.get_secret") as mock_get_secret: + mock_get_secret.return_value = None + + with pytest.raises(Exception) as exc_info: + completion( + model="gemini/gemini-1.5-flash", + messages=[{"role": "user", "content": "Test message"}], + ) + + error_message = str(exc_info.value).lower() + assert any( + keyword in error_message + for keyword in [ + "api key", + "authentication", + "unauthorized", + "invalid", + "missing", + "credential", + ] + ) + + +@pytest.mark.parametrize("api_key_env", ["GOOGLE_API_KEY", "GEMINI_API_KEY"]) +def test_gemini_completion_both_env_vars(monkeypatch, api_key_env): + """Test Gemini completion works with both environment variable names.""" + # Clear all API keys first + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + + # Set the specific API key being tested + monkeypatch.setenv(api_key_env, f"fake-{api_key_env.lower()}-for-testing") + + response = completion( + model="gemini/gemini-1.5-flash", + messages=[{"role": "user", "content": f"Test with {api_key_env}"}], + mock_response=f"Mocked response using {api_key_env}", + ) + assert ( + response["choices"][0]["message"]["content"] + == f"Mocked response using {api_key_env}" + ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 709ebd662d..6aa18c560c 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -2,7 +2,7 @@ Test cases for spend log cleanup functionality """ -from datetime import UTC, datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest