tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 15:07:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix/gemini api key environment variable support (#12507)

Browse Source 
* 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
This commit is contained in:
Siddharth Sahu
2025-07-30 04:26:01 +05:30
committed by GitHub
parent 1899549e7b
commit 39d59f1900
8 changed files with 166 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/gemini/common_utils.py
+6 -2
View File
@@ -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")
litellm/llms/gemini/google_genai/transformation.py
+2 -3
View File
@@ -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
)
litellm/llms/gemini/realtime/transformation.py
+49 -40
View File
@@ -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():
litellm/main.py
+3 -2
View File
@@ -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")
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+1 -1
View File
@@ -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)
litellm/utils.py
+3 -2
View File
@@ -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
tests/test_litellm/llms/gemini/test_gemini_client_setup.py
+101
View File
@@ -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}"
)
tests/test_litellm/proxy/test_spend_log_cleanup.py
+1 -1
View File
@@ -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