scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(gemini): harden vertex oauth path and dedupe code-assist urls

Browse source
This commit is contained in:
balazss 2026-03-17 23:26:00 -07:00
parent 6714590c9c
commit 2b8ca172eb
5 changed files with 98 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/google_code_assist/chat.py
View file

@ -7,6 +7,10 @@ from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import _get_httpx_client, AsyncHTTPHandler
from .transformation import GoogleCodeAssistConfig, GoogleCodeAssistError
_CODE_ASSIST_BASE_URL = "https://cloudcode-pa.googleapis.com"
_CODE_ASSIST_GENERATE_URL = f"{_CODE_ASSIST_BASE_URL}/v1internal:generateContent"
_CODE_ASSIST_LOAD_URL = f"{_CODE_ASSIST_BASE_URL}/v1internal:loadCodeAssist"
class GoogleCodeAssistChat:
"""
@ -65,7 +69,7 @@ class GoogleCodeAssistChat:
)
# 4. Call Completion API
url = "https://cloudcode-pa.googleapis.com/v1internal:generateContent"
url = _CODE_ASSIST_GENERATE_URL
headers = self._get_headers(token)
response = client.post(
@ -133,7 +137,7 @@ class GoogleCodeAssistChat:
data = self.config.transform_request(
model, messages, optional_params, litellm_params
)
url = "https://cloudcode-pa.googleapis.com/v1internal:generateContent"
url = _CODE_ASSIST_GENERATE_URL
headers = self._get_headers(token)
response = await async_handler.post(url=url, headers=headers, json=data)
@ -165,7 +169,7 @@ class GoogleCodeAssistChat:
self, client, token: str, initial_project_id: Optional[str]
) -> Optional[str]:
"""Performs the loadCodeAssist handshake to establish session context."""
load_url = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
load_url = _CODE_ASSIST_LOAD_URL
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
@ -197,7 +201,7 @@ class GoogleCodeAssistChat:
initial_project_id: Optional[str],
) -> Optional[str]:
"""Async version of loadCodeAssist handshake for non-blocking async calls."""
load_url = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
load_url = _CODE_ASSIST_LOAD_URL
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",

15
litellm/llms/google_code_assist/transformation.py
View file

@ -4,6 +4,7 @@ import copy
import httpx
from typing import Any, List, Optional
from litellm._logging import verbose_logger
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ModelResponse
from ..vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
@ -127,6 +128,20 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
generation_config["thinkingConfig"] = {
"includeThoughts": base_params.pop("include_thoughts")
}
elif "thinkingConfig" in optional_params:
generation_config["thinkingConfig"] = optional_params["thinkingConfig"]
elif "include_thoughts" in optional_params:
generation_config["thinkingConfig"] = {
"includeThoughts": optional_params["include_thoughts"]
}
if (
"thinkingConfig" in optional_params
and "thinkingConfig" not in generation_config
):
verbose_logger.warning(
"google_code_assist: `thinkingConfig` was provided but not mapped into generationConfig."
)
vertex_request = {
"contents": contents,

14
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -2,6 +2,7 @@
## httpx client for vertex ai calls
## Initial implementation - covers gemini + image gen calls
import json
import asyncio
import time
from copy import deepcopy
from functools import partial
@ -2504,6 +2505,11 @@ class VertexLLM(VertexBase):
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
gemini_auth_data = None
if custom_llm_provider == "gemini" and gemini_api_key is None:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token)
# Extract use_psc_endpoint_format from optional_params
use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False)
@ -2511,6 +2517,7 @@ class VertexLLM(VertexBase):
auth_header, api_base = self._get_token_and_url(
model=model,
gemini_api_key=gemini_api_key,
gemini_auth_data=gemini_auth_data,
auth_header=_auth_header,
vertex_project=vertex_project,
vertex_location=vertex_location,
@ -2606,6 +2613,11 @@ class VertexLLM(VertexBase):
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
gemini_auth_data = None
if custom_llm_provider == "gemini" and gemini_api_key is None:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token)
# Extract use_psc_endpoint_format from optional_params
use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False)
@ -2613,6 +2625,7 @@ class VertexLLM(VertexBase):
auth_header, api_base = self._get_token_and_url(
model=model,
gemini_api_key=gemini_api_key,
gemini_auth_data=gemini_auth_data,
auth_header=_auth_header,
vertex_project=vertex_project,
vertex_location=vertex_location,
@ -2801,6 +2814,7 @@ class VertexLLM(VertexBase):
auth_header, url = self._get_token_and_url(
model=model,
gemini_api_key=gemini_api_key,
gemini_auth_data=None,
auth_header=_auth_header,
vertex_project=vertex_project,
vertex_location=vertex_location,

13
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -368,6 +368,7 @@ class VertexBase:
vertex_location: Optional[str] = None,
vertex_api_version: Optional[Literal["v1", "v1beta1"]] = None,
use_psc_endpoint_format: bool = False,
gemini_auth_data: Optional[dict] = None,
) -> Tuple[Optional[str], str]:
"""
for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
@ -395,8 +396,7 @@ class VertexBase:
)
url = "{}/models/{}:{}".format(api_base, model, endpoint)
gemini_auth_data = None
if gemini_api_key is None:
if gemini_auth_data is None and gemini_api_key is None:
from litellm.llms.gemini.common_utils import (
get_gemini_oauth_token,
)
@ -460,6 +460,7 @@ class VertexBase:
should_use_v1beta1_features: Optional[bool] = False,
mode: all_gemini_url_modes = "chat",
use_psc_endpoint_format: bool = False,
gemini_auth_data: Optional[dict] = None,
) -> Tuple[Optional[str], str]:
"""
Internal function. Returns the token and url for the call.
@ -471,8 +472,7 @@ class VertexBase:
"""
version: Optional[Literal["v1beta1", "v1"]] = None
if custom_llm_provider == "gemini":
gemini_auth_data = None
if gemini_api_key is None:
if gemini_auth_data is None and gemini_api_key is None:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
gemini_auth_data = get_gemini_oauth_token()
@ -491,6 +491,10 @@ class VertexBase:
auth_header = {"Authorization": f"Bearer {gemini_oauth_token}"}
if gemini_auth_data and gemini_auth_data.get("project_id"):
auth_header["x-goog-user-project"] = gemini_auth_data["project_id"]
elif gemini_api_key is None:
raise ValueError(
"Missing gemini_api_key. Please set `GEMINI_API_KEY` or `GEMINI_OAUTH_TOKEN`."
)
else:
auth_header = (
None # this field is not used for gemini when using api key
@ -517,6 +521,7 @@ class VertexBase:
auth_header=auth_header,
custom_llm_provider=custom_llm_provider,
gemini_api_key=gemini_api_key,
gemini_auth_data=gemini_auth_data,
endpoint=endpoint,
stream=stream,
url=url,

52
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -1074,6 +1074,58 @@ def test_get_token_url():
pass
def test_get_token_and_url_gemini_raises_if_no_api_key_and_no_oauth_token():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
)
vertex_llm = VertexLLM()
with patch(
"litellm.llms.gemini.common_utils.get_gemini_oauth_token", return_value=None
):
with pytest.raises(ValueError, match="Missing gemini_api_key"):
vertex_llm._get_token_and_url(
auth_header=None,
vertex_project="",
vertex_location="",
vertex_credentials="",
gemini_api_key=None,
custom_llm_provider="gemini",
should_use_v1beta1_features=False,
api_base=None,
model="gemini-2.5-pro",
stream=False,
)
def test_get_token_and_url_gemini_uses_prefetched_auth_data_without_lookup():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
)
vertex_llm = VertexLLM()
with patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token") as mock_get:
auth_header, _ = vertex_llm._get_token_and_url(
auth_header=None,
vertex_project="",
vertex_location="",
vertex_credentials="",
gemini_api_key=None,
gemini_auth_data={"token": "oauth-token", "project_id": "my-project"},
custom_llm_provider="gemini",
should_use_v1beta1_features=False,
api_base=None,
model="gemini-2.5-pro",
stream=False,
)
assert auth_header == {
"Authorization": "Bearer oauth-token",
"x-goog-user-project": "my-project",
}
mock_get.assert_not_called()
@pytest.mark.asyncio
async def test_vertex_ai_token_counter_routes_partner_models():
"""