This commit is contained in:
balazss 2026-05-04 21:11:39 +00:00 • committed by GitHub
commit db3991ff33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
30 changed files with 2070 additions and 96 deletions

View file

@ -0,0 +1,84 @@
# Google Code Assist
Use Google Code Assist models through LiteLLM using Gemini OAuth credentials.
| Property | Details |
|-------|-------|
| Description | Google Code Assist API via Google Cloud Code Assist backend |
| Provider Route on LiteLLM | `google_code_assist/` |
| Supported Endpoints | `/chat/completions` |
| Base Endpoint | `https://cloudcode-pa.googleapis.com` |
## Model Format
Use:
`google_code_assist/<gemini-model>`
Example:
`google_code_assist/gemini-2.5-pro`
## Authentication
Google Code Assist requires OAuth credentials (not API key auth).
LiteLLM checks credentials in this order:
1. `GEMINI_OAUTH_TOKEN`
2. `GEMINI_CREDENTIALS_PATH`
3. `~/.config/litellm/gemini_oauth/oauth_creds.json`
4. `~/.gemini/oauth_creds.json`
5. `~/.config/gcloud/application_default_credentials.json`
### Login via CLI
```bash
export GEMINI_OAUTH_CLIENT_ID="<your-google-oauth-client-id>"
export GEMINI_OAUTH_CLIENT_SECRET="<your-google-oauth-client-secret>"
litellm-proxy gemini login
```
This opens a browser loopback OAuth flow and stores credentials locally.
## Usage - LiteLLM Python SDK
```python
from litellm import completion
response = completion(
model="google_code_assist/gemini-2.5-pro",
messages=[{"role": "user", "content": "Write a Python function to parse CSV safely."}],
)
print(response.choices[0].message.content)
```
## Usage - LiteLLM Proxy
```yaml
model_list:
- model_name: google-code-assist
litellm_params:
model: google_code_assist/gemini-2.5-pro
```
```bash
litellm --config /path/to/config.yaml
```
```bash
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer <YOUR_LITELLM_KEY>" \
-d '{
"model": "google-code-assist",
"messages": [{"role": "user", "content": "Refactor this code for readability."}]
}'
```
## Optional Params
- `google_code_assist_project`: Optional project override sent to the Code Assist backend request.
LiteLLM also performs the required `loadCodeAssist` handshake before `generateContent` automatically.

View file

@ -456,6 +456,10 @@ default_fallbacks: Optional[List] = None
fallbacks: Optional[List] = None
context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
# Backward-compatible default: do not silently reroute gemini/* calls unless explicitly enabled.
auto_fallback_to_google_code_assist: bool = os.getenv(
"LITELLM_AUTO_FALLBACK_TO_GOOGLE_CODE_ASSIST", "false"
).lower() in ("1", "true", "yes", "on")
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
num_retries_per_request: Optional[int] = (

View file

@ -0,0 +1,291 @@
import json
import os
import tempfile
import time
import webbrowser
import http.server
import html
import urllib.parse
import secrets
import hashlib
import base64
from typing import Any, Dict
import httpx
from litellm._logging import verbose_logger
from litellm.secret_managers.main import get_secret_str
GEMINI_SCOPES = [
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
"https://www.googleapis.com/auth/userinfo.profile",
"https://www.googleapis.com/auth/generative-language",
]
# Google OAuth URLs
GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
class GeminiAuthenticator:
def __init__(self) -> None:
"""Initialize the Gemini authenticator with configurable token paths."""
# Token storage paths
self.token_dir = os.getenv(
"GEMINI_OAUTH_TOKEN_DIR",
os.path.expanduser("~/.config/litellm/gemini_oauth"),
)
self.oauth_creds_file = os.path.join(
self.token_dir,
os.getenv("GEMINI_OAUTH_CREDS_FILE", "oauth_creds.json"),
)
self._ensure_token_dir()
@staticmethod
def _get_oauth_client_credentials() -> tuple[str, str]:
client_id = get_secret_str("GEMINI_OAUTH_CLIENT_ID")
client_secret = get_secret_str("GEMINI_OAUTH_CLIENT_SECRET")
if not client_id or not client_secret:
raise ValueError(
"Missing Gemini OAuth client credentials. "
"Set GEMINI_OAUTH_CLIENT_ID and GEMINI_OAUTH_CLIENT_SECRET."
)
return client_id, client_secret
@staticmethod
def _get_access_token_or_raise(creds: Dict[str, Any], source: str) -> str:
token = creds.get("access_token")
if not token or not isinstance(token, str):
raise Exception(
f"OAuth login failed: missing access_token in {source} response."
)
return token
def get_token(self) -> str:
"""
Get the OAuth token, refreshing if necessary.
Returns:
str: The Gemini access token.
Raises:
Exception: If unable to obtain or refresh an access token.
"""
try:
if os.path.exists(self.oauth_creds_file):
with open(self.oauth_creds_file, "r") as f:
creds = json.load(f)
# Check if access_token exists and is valid (rough check)
if (
creds.get("access_token")
and creds.get("expires_at", 0) > time.time() + 60
):
return creds.get("access_token")
# If expired but has refresh_token, try to refresh
if creds.get("refresh_token"):
verbose_logger.debug(
"Gemini access token expired, refreshing..."
)
return self._refresh_token(creds.get("refresh_token"))
except Exception as e:
verbose_logger.warning(f"Error reading Gemini OAuth credentials: {e}")
# If we get here, we need to log in
verbose_logger.info("Starting Gemini OAuth login flow...")
creds = self._login()
self._write_oauth_creds(creds)
return self._get_access_token_or_raise(creds, "authorization_code")
def _refresh_token(self, refresh_token: str) -> str:
"""Refresh the access token using the refresh token."""
client_id, client_secret = self._get_oauth_client_credentials()
data = {
"client_id": client_id,
"client_secret": client_secret,
"refresh_token": refresh_token,
"grant_type": "refresh_token",
}
resp = httpx.post(GOOGLE_TOKEN_URL, data=data)
resp.raise_for_status()
new_creds = resp.json()
# Load existing creds to merge if possible
creds = {}
if os.path.exists(self.oauth_creds_file):
try:
with open(self.oauth_creds_file, "r") as f:
creds = json.load(f)
except Exception:
pass
creds.update(new_creds)
if "expires_in" in new_creds:
creds["expires_at"] = time.time() + new_creds["expires_in"]
self._write_oauth_creds(creds)
return self._get_access_token_or_raise(creds, "refresh_token")
def _ensure_token_dir(self) -> None:
"""Ensure the token directory exists."""
if not os.path.exists(self.token_dir):
os.makedirs(self.token_dir, mode=0o700, exist_ok=True)
else:
try:
os.chmod(self.token_dir, 0o700)
except OSError:
pass
def _write_oauth_creds(self, creds: Dict[str, Any]) -> None:
"""
Write oauth credentials with user-only permissions.
"""
fd, tmp_path = tempfile.mkstemp(dir=self.token_dir, prefix=".tmp_creds_")
try:
os.chmod(tmp_path, 0o600)
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(creds, f)
os.replace(tmp_path, self.oauth_creds_file)
try:
os.chmod(self.oauth_creds_file, 0o600)
except OSError:
pass
except Exception:
try:
os.unlink(tmp_path)
except OSError:
pass
raise
def _login(self) -> Dict[str, Any]:
"""Perform loopback flow login."""
client_id, client_secret = self._get_oauth_client_credentials()
# PKCE
code_verifier = secrets.token_urlsafe(64)
code_challenge = (
base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest())
.decode()
.replace("=", "")
)
state = secrets.token_urlsafe(32)
# Local web server for callback
auth_code = None
error = None
class CallbackHandler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
nonlocal auth_code, error
parsed_url = urllib.parse.urlparse(self.path)
if parsed_url.path != "/oauth2callback":
self.send_response(404)
self.end_headers()
return
query = parsed_url.query
params = urllib.parse.parse_qs(query)
if params.get("state", [None])[0] != state:
error = "state_mismatch"
self.send_response(400)
self.send_header("Content-type", "text/html")
self.end_headers()
self.wfile.write(
b"<html><body><h1>Authentication failed</h1><p>State mismatch. Possible CSRF attack.</p></body></html>"
)
elif "code" in params:
auth_code = params["code"][0]
self.send_response(200)
self.send_header("Content-type", "text/html")
self.end_headers()
self.wfile.write(
b"<html><body><h1>Authentication successful!</h1><p>You can close this tab and return to the terminal.</p></body></html>"
)
elif "error" in params:
error = params["error"][0]
self.send_response(200)
self.send_header("Content-type", "text/html")
self.end_headers()
self.wfile.write(
(
"<html><body><h1>Authentication failed</h1><p>"
f"{html.escape(error)}"
"</p></body></html>"
).encode()
)
else:
self.send_response(404)
self.end_headers()
def log_message(self, format, *args):
# Suppress logging to avoid noise in the terminal
pass
server = http.server.HTTPServer(("127.0.0.1", 0), CallbackHandler)
port = server.server_port
redirect_uri = f"http://127.0.0.1:{port}/oauth2callback"
params = {
"client_id": client_id,
"redirect_uri": redirect_uri,
"response_type": "code",
"scope": " ".join(GEMINI_SCOPES),
"state": state,
"code_challenge": code_challenge,
"code_challenge_method": "S256",
"access_type": "offline",
"prompt": "consent",
}
auth_url = f"{GOOGLE_AUTH_URL}?{urllib.parse.urlencode(params)}"
print( # noqa: T201
f"Please visit the following URL to authenticate:\n\n{auth_url}\n"
)
webbrowser.open(auth_url)
# Wait for callback; browsers may hit non-callback paths first (e.g. /favicon.ico).
# Use a bounded loop so headless/CI environments fail fast instead of hanging forever.
loopback_timeout_s = float(
os.getenv("GEMINI_OAUTH_LOOPBACK_TIMEOUT_SECONDS", "120")
)
loopback_timeout_s = max(loopback_timeout_s, 1.0)
deadline = time.monotonic() + loopback_timeout_s
server.timeout = 1.0
try:
# Note: deadline is checked after handle_request() returns, so the loop can overshoot by up to server.timeout (1s) — acceptable for a one-shot CLI login.
while auth_code is None and error is None:
server.handle_request()
if time.monotonic() >= deadline:
raise Exception(
"OAuth login timed out. Re-run `litellm-proxy gemini login` "
f"and complete the browser flow within {int(loopback_timeout_s)} seconds."
)
finally:
server.server_close()
if error:
raise Exception(f"OAuth login failed: {error}")
if not auth_code:
raise Exception("OAuth login failed: No code received")
# Exchange code for token
data = {
"client_id": client_id,
"client_secret": client_secret,
"code": auth_code,
"code_verifier": code_verifier,
"grant_type": "authorization_code",
"redirect_uri": redirect_uri,
}
resp = httpx.post(GOOGLE_TOKEN_URL, data=data)
resp.raise_for_status()
creds = resp.json()
if "expires_in" in creds:
creds["expires_at"] = time.time() + creds["expires_in"]
return creds

View file

@ -1,10 +1,12 @@
import base64
import datetime
import json
from typing import Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@ -155,6 +157,196 @@ def get_api_key_from_env() -> Optional[str]:
return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY")
def should_fallback_to_google_code_assist(error: Exception) -> bool:
"""
Returns True if the error indicates missing OAuth scope for Gemini calls.
"""
def _iter_exception_chain(exc: BaseException):
seen = set()
current: Optional[BaseException] = exc
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
current = current.__cause__ or current.__context__
def _contains_scope_code(value: Any) -> bool:
if isinstance(value, str):
return "ACCESS_TOKEN_SCOPE_INSUFFICIENT" in value
if isinstance(value, dict):
for v in value.values():
if _contains_scope_code(v):
return True
return False
if isinstance(value, list):
for item in value:
if _contains_scope_code(item):
return True
return False
return False
def _extract_json_payload_from_response(response: Any) -> Optional[dict]:
if response is None:
return None
try:
payload = response.json()
if isinstance(payload, dict):
return payload
except Exception:
pass
try:
text = getattr(response, "text", None)
if isinstance(text, str):
payload = json.loads(text)
if isinstance(payload, dict):
return payload
except Exception:
pass
return None
for exc in _iter_exception_chain(error):
if isinstance(exc, httpx.HTTPStatusError):
response = getattr(exc, "response", None)
if response is None or getattr(response, "status_code", None) != 403:
continue
payload = _extract_json_payload_from_response(response)
if payload and _contains_scope_code(payload):
return True
continue
status_code = getattr(exc, "status_code", None)
if status_code != 403:
continue
body = getattr(exc, "body", None)
if body and _contains_scope_code(body):
return True
response = getattr(exc, "response", None)
payload = _extract_json_payload_from_response(response)
if payload and _contains_scope_code(payload):
return True
message = getattr(exc, "message", None)
if message and _contains_scope_code(message):
return True
return False
def get_gemini_oauth_token() -> Optional[dict]: # noqa: PLR0915
"""
Returns the Gemini OAuth token and metadata.
Check:
1. GEMINI_OAUTH_TOKEN env var
2. GEMINI_CREDENTIALS_PATH env var
3. ~/.config/litellm/gemini_oauth/oauth_creds.json
4. ~/.gemini/oauth_creds.json (gemini-cli default)
5. ~/.config/gcloud/application_default_credentials.json (Standard ADC)
"""
import json
import time
from pathlib import Path
def _is_expired(creds_data: dict) -> bool:
expires_at = creds_data.get("expires_at")
if isinstance(expires_at, (int, float)):
return expires_at <= (time.time() + 60)
expiry = creds_data.get("expiry") or creds_data.get("expires_on")
if isinstance(expiry, str):
try:
expires_dt = datetime.datetime.fromisoformat(
expiry.replace("Z", "+00:00")
)
return expires_dt.timestamp() <= (time.time() + 60)
except Exception:
return False
return False
# 1. Check GEMINI_OAUTH_TOKEN
token = get_secret_str("GEMINI_OAUTH_TOKEN")
if token:
return {"token": token}
# 2. Check GEMINI_CREDENTIALS_PATH or default paths
paths_to_check = []
env_path = get_secret_str("GEMINI_CREDENTIALS_PATH")
if env_path:
paths_to_check.append(Path(env_path).expanduser())
paths_to_check.extend(
[
Path("~/.config/litellm/gemini_oauth/oauth_creds.json").expanduser(),
Path("~/.gemini/oauth_creds.json").expanduser(),
Path("~/.config/gcloud/application_default_credentials.json").expanduser(),
]
)
for creds_path in paths_to_check:
if creds_path.exists():
try:
# Try to use google-auth if available
try:
from google.oauth2 import credentials
from google.auth.transport.requests import Request
creds = credentials.Credentials.from_authorized_user_file(
str(creds_path)
)
if not creds.valid:
verbose_logger.debug(
f"Refreshing Gemini token for {creds_path}"
)
creds.refresh(Request())
if creds.token:
result = {"token": creds.token}
# Extract project ID if present
with open(creds_path, "r") as f:
creds_data = json.load(f)
project_id = creds_data.get(
"quota_project_id"
) or creds_data.get("project_id")
if project_id:
result["project_id"] = project_id
return result
except Exception:
# Fallback to manual reading
with open(creds_path, "r") as f:
creds_data = json.load(f)
if _is_expired(creds_data):
verbose_logger.warning(
"Gemini OAuth token appears expired in %s. "
"Run `litellm-proxy gemini login` to refresh credentials.",
creds_path,
)
continue
token = creds_data.get("access_token")
if not token and "token" in creds_data:
token_field = creds_data["token"]
if isinstance(token_field, dict):
token = token_field.get("accessToken")
elif isinstance(token_field, str):
token = token_field
if token:
result = {"token": token}
project_id = creds_data.get(
"quota_project_id"
) or creds_data.get("project_id")
if project_id:
result["project_id"] = project_id
return result
except Exception as e:
verbose_logger.error(
f"Error reading Gemini credentials from {creds_path}: {e}"
)
return None
class GoogleAIStudioTokenCounter(BaseTokenCounter):
"""Token counter implementation for Google AI Studio provider."""

View file

@ -0,0 +1,57 @@
from typing import Any, Awaitable, Callable, Dict
from litellm._logging import verbose_logger
from litellm.llms.gemini.common_utils import should_fallback_to_google_code_assist
from litellm.llms.google_code_assist.chat import get_google_code_assist_chat
_google_code_assist_chat = get_google_code_assist_chat()
async def run_gemini_acompletion_with_code_assist_fallback(
primary_call: Callable[[], Awaitable[Any]],
fallback_kwargs: Dict[str, Any],
auto_fallback_to_google_code_assist: bool = False,
) -> Any:
"""
Execute Gemini async completion and fallback to Google Code Assist when
OAuth scope is insufficient.
"""
try:
return await primary_call()
except Exception as e:
if not auto_fallback_to_google_code_assist:
raise e
if not should_fallback_to_google_code_assist(e):
raise e
verbose_logger.warning(
"Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. "
"Falling back to google_code_assist."
)
return await _google_code_assist_chat.acompletion(**fallback_kwargs)
def run_gemini_completion_with_code_assist_fallback(
primary_call: Callable[[], Any],
fallback_kwargs: Dict[str, Any],
auto_fallback_to_google_code_assist: bool = False,
) -> Any:
"""
Execute Gemini sync completion and fallback to Google Code Assist when
OAuth scope is insufficient.
"""
try:
return primary_call()
except Exception as e:
if not auto_fallback_to_google_code_assist:
raise e
if not should_fallback_to_google_code_assist(e):
raise e
verbose_logger.warning(
"Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. "
"Falling back to google_code_assist."
)
return _google_code_assist_chat.completion(**fallback_kwargs)

View file

@ -2,6 +2,7 @@
Transformation for Calling Google models in their native format.
"""
import asyncio
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
import httpx
@ -11,6 +12,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.base_llm.google_genai.transformation import (
BaseGoogleGenAIGenerateContentConfig,
)
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.types.router import GenericLiteLLMParams
@ -200,6 +202,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
stream: bool,
api_base: Optional[str],
litellm_params: dict,
gemini_auth_data: Optional[dict] = None,
) -> Tuple[dict, str]:
"""
Build final headers and API URL from auth components.
@ -217,6 +220,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
custom_llm_provider=self.custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=True,
gemini_auth_data=gemini_auth_data,
)
headers = self.validate_environment(
@ -291,6 +295,12 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
custom_llm_provider=self.custom_llm_provider,
)
gemini_auth_data: Optional[dict] = None
if self.custom_llm_provider == "gemini":
gemini_api_key = self._get_google_ai_studio_api_key(dict(litellm_params))
if gemini_api_key is None:
gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token)
return self._build_final_headers_and_url(
model=model,
auth_header=_auth_header,
@ -300,6 +310,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
stream=stream,
api_base=api_base,
litellm_params=litellm_params,
gemini_auth_data=gemini_auth_data,
)
def transform_generate_content_request(

View file

@ -0,0 +1,258 @@
import asyncio
import httpx
from typing import Any, Dict, Optional
import litellm
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:
"""
Handler for Google Code Assist API.
Provides both synchronous and asynchronous completion methods.
"""
def __init__(self) -> None:
self.config = GoogleCodeAssistConfig()
def completion(
self,
model: str,
messages: list,
model_response: litellm.utils.ModelResponse,
print_verbose: Any,
logging_obj: Any,
optional_params: dict,
litellm_params: dict,
logger_fn=None,
) -> litellm.utils.ModelResponse:
"""
Synchronous completion for Google Code Assist.
"""
try:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
# 1. Get OAuth data
gemini_auth_data = get_gemini_oauth_token()
if not gemini_auth_data:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token. Run 'litellm-proxy gemini login' or set GEMINI_OAUTH_TOKEN.",
)
token = gemini_auth_data.get("token")
initial_project_id = gemini_auth_data.get("project_id")
if not token:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.",
)
client = _get_httpx_client()
# 2. MANDATORY HANDSHAKE: loadCodeAssist
final_project_id = self._handle_handshake(client, token, initial_project_id)
litellm_params["google_code_assist_project"] = final_project_id
# 3. Transform request
data = self.config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
)
# 4. Call Completion API
url = _CODE_ASSIST_GENERATE_URL
headers = self._get_headers(token)
response = client.post(
url=url,
headers=headers,
json=data,
)
response.raise_for_status()
# 5. Transform response
return self.config.transform_response(
model=model,
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
)
except Exception as e:
raise self._handle_error(e)
async def acompletion(
self,
model: str,
messages: list,
model_response: litellm.utils.ModelResponse,
print_verbose: Any,
logging_obj: Any,
optional_params: dict,
litellm_params: dict,
logger_fn=None,
) -> litellm.utils.ModelResponse:
"""
Asynchronous completion for Google Code Assist.
"""
try:
from litellm.llms.gemini.common_utils import get_gemini_oauth_token
gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token)
if not gemini_auth_data:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token. Run 'litellm-proxy gemini login' or set GEMINI_OAUTH_TOKEN.",
)
token = gemini_auth_data.get("token")
initial_project_id = gemini_auth_data.get("project_id")
if not token:
raise GoogleCodeAssistError(
status_code=401,
message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.",
)
async_handler = AsyncHTTPHandler()
try:
final_project_id = await self._ahandle_handshake(
async_handler, token, initial_project_id
)
litellm_params["google_code_assist_project"] = final_project_id
data = self.config.transform_request(
model, messages, optional_params, litellm_params
)
url = _CODE_ASSIST_GENERATE_URL
headers = self._get_headers(token)
response = await async_handler.post(url=url, headers=headers, json=data)
response.raise_for_status()
return self.config.transform_response(
model=model,
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
request_data=data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=None,
)
finally:
try:
await async_handler.close()
except Exception as close_error:
verbose_logger.debug(
f"Failed to close Google Code Assist async HTTP handler: {close_error}"
)
except Exception as e:
raise self._handle_error(e)
def _handle_handshake(
self, client, token: str, initial_project_id: Optional[str]
) -> Optional[str]:
"""Performs the loadCodeAssist handshake to establish session context."""
load_url = _CODE_ASSIST_LOAD_URL
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload: Dict[str, Any] = {
"metadata": {
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
}
}
if initial_project_id:
load_payload["cloudaicompanionProject"] = initial_project_id
load_payload["metadata"]["duetProject"] = initial_project_id
try:
load_resp = client.post(load_url, headers=load_headers, json=load_payload)
load_resp.raise_for_status()
return load_resp.json().get("cloudaicompanionProject") or initial_project_id
except Exception as e:
verbose_logger.debug(f"Gemini Code Assist handshake failed: {e}")
return initial_project_id
async def _ahandle_handshake(
self,
async_handler: AsyncHTTPHandler,
token: str,
initial_project_id: Optional[str],
) -> Optional[str]:
"""Async version of loadCodeAssist handshake for non-blocking async calls."""
load_url = _CODE_ASSIST_LOAD_URL
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload: Dict[str, Any] = {
"metadata": {
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
}
}
if initial_project_id:
load_payload["cloudaicompanionProject"] = initial_project_id
load_payload["metadata"]["duetProject"] = initial_project_id
try:
load_resp = await async_handler.post(
url=load_url, headers=load_headers, json=load_payload
)
load_resp.raise_for_status()
return load_resp.json().get("cloudaicompanionProject") or initial_project_id
except Exception as e:
verbose_logger.debug(f"Gemini Code Assist async handshake failed: {e}")
return initial_project_id
def _get_headers(self, token: str) -> dict:
"""Returns standard headers for Code Assist API calls."""
return {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
"X-Goog-Api-Client": "litellm-google-code-assist",
}
def _handle_error(self, e: Exception) -> Exception:
"""Centralized error mapping for Code Assist."""
if isinstance(e, httpx.HTTPStatusError):
return GoogleCodeAssistError(
status_code=e.response.status_code, message=e.response.text
)
if isinstance(e, GoogleCodeAssistError):
return e
return GoogleCodeAssistError(status_code=500, message=str(e))
_shared_google_code_assist_chat: Optional[GoogleCodeAssistChat] = None
def get_google_code_assist_chat() -> GoogleCodeAssistChat:
global _shared_google_code_assist_chat
if _shared_google_code_assist_chat is None:
_shared_google_code_assist_chat = GoogleCodeAssistChat()
return _shared_google_code_assist_chat

View file

@ -0,0 +1,182 @@
import json
import uuid
import copy
import httpx
from typing import Any
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
class GoogleCodeAssistError(BaseLLMException):
"""
Exception raised for errors in the Google Code Assist API.
"""
def __init__(self, status_code, message):
super().__init__(status_code=status_code, message=message)
class ParsedJSONResponseAdapter:
"""
Adapter that provides the subset of httpx.Response surface used by
VertexGeminiConfig.transform_response.
"""
def __init__(self, json_data: dict):
self._json = json_data
self.status_code = 200
self.text = json.dumps(json_data)
self.headers = httpx.Headers({"content-type": "application/json"})
def json(self):
return self._json
class GoogleCodeAssistConfig(VertexGeminiConfig):
"""
Reference: https://cloud.google.com/gemini/docs/api/reference/rest/v1internal/projects.locations.codeAssist/generateContent
The class `GoogleCodeAssistConfig` provides configuration for the Google Code Assist API.
It inherits from `VertexGeminiConfig` to provide consistent parameter mapping.
- `temperature` (float): This controls the degree of randomness in token selection.
- `max_output_tokens` (integer): This sets the limitation for the maximum amount of token in the text output.
- `top_p` (float): The tokens are selected from the most probable to the least probable until the sum of their probabilities equals the `top_p` value.
- `top_k` (integer): The value of `top_k` determines how many of the most probable tokens are considered in the selection.
- `stop_sequences` (List[str]): The set of character sequences that will stop output generation.
"""
def transform_request(
self,
model: str,
messages: list,
optional_params: dict,
litellm_params: dict,
headers: dict = {},
) -> dict:
"""
Transforms standard LiteLLM request to Code Assist API format.
Matches gemini-cli 'toGenerateContentRequest'.
"""
# 1. Map messages to Gemini format
from ..vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
_transform_system_message,
)
# Create a copy to avoid mutating the original list
messages_copy = copy.deepcopy(messages)
model_name = model.split("/")[-1]
# Separate system instruction
system_instruction, filtered_messages = _transform_system_message(
supports_system_message=True, messages=messages_copy
)
# Convert the rest of messages
contents = _gemini_convert_messages_with_history(
messages=filtered_messages, model=model_name
)
# 2. Build vertex-style nested request
generation_config = {}
# Handle parameter mapping
base_params = self.map_openai_params(
{}, optional_params.copy(), model_name, False
)
for key in ["temperature", "topP", "topK", "maxOutputTokens", "stopSequences"]:
if key in base_params:
generation_config[key] = base_params.pop(key)
# Support thinkingConfig
if "thinkingConfig" in base_params:
generation_config["thinkingConfig"] = base_params.pop("thinkingConfig")
elif "include_thoughts" in base_params:
generation_config["thinkingConfig"] = {
"includeThoughts": base_params.pop("include_thoughts")
}
# Forward any remaining mapped params into generation_config so they are not silently dropped.
for key, value in base_params.items():
if key not in generation_config:
verbose_logger.debug(
"google_code_assist: forwarding mapped param '%s' into generationConfig",
key,
)
generation_config[key] = value
# Log unmapped/unrecognized caller keys so users have visibility (do not auto-forward —
# they may not be valid Gemini fields).
for key in optional_params:
if key not in base_params and key not in generation_config:
verbose_logger.debug(
"google_code_assist: dropping unrecognized optional_param '%s' "
"(not produced by VertexGeminiConfig.map_openai_params)",
key,
)
vertex_request = {
"contents": contents,
"session_id": litellm_params.get("session_id", str(uuid.uuid4())),
}
if system_instruction:
vertex_request["systemInstruction"] = {
"role": "system",
"parts": system_instruction["parts"],
}
if generation_config:
vertex_request["generationConfig"] = generation_config
# 3. Wrap in Code Assist envelope (matches verified gemini-cli structure)
user_prompt_id = f"litellm-{uuid.uuid4()}"
ca_request = {
"model": model_name,
"user_prompt_id": user_prompt_id,
"request": vertex_request,
}
# Add project ID if available
if litellm_params.get("google_code_assist_project"):
ca_request["project"] = litellm_params["google_code_assist_project"]
return ca_request
def transform_response(
self,
model: str,
raw_response: Any,
model_response: ModelResponse,
logging_obj: Any,
request_data: dict,
messages: list,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Any = None,
json_mode: Any = None,
) -> ModelResponse:
"""
Transforms Code Assist API response to standard LiteLLM format.
"""
data = raw_response.json()
# Code Assist wraps the response in a "response" key
gemini_response = data.get("response", data)
return super().transform_response(
model=model,
raw_response=ParsedJSONResponseAdapter(gemini_response), # type: ignore[arg-type]
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
)

View file

@ -413,28 +413,12 @@ def _get_gemini_url(
endpoint = "generateContent"
if stream is True:
endpoint = "streamGenerateContent"
url = "https://generativelanguage.googleapis.com/{}/{}:{}?alt=sse".format(
api_version, _gemini_model_name, endpoint
)
else:
url = "https://generativelanguage.googleapis.com/{}/{}:{}".format(
api_version, _gemini_model_name, endpoint
)
elif mode == "embedding":
endpoint = "embedContent"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
_gemini_model_name, endpoint
)
elif mode == "batch_embedding":
endpoint = "batchEmbedContents"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
_gemini_model_name, endpoint
)
elif mode == "count_tokens":
endpoint = "countTokens"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format(
_gemini_model_name, endpoint
)
elif mode == "image_generation":
raise ValueError(
"LiteLLM's `gemini/` route does not support image generation yet. Let us know if you need this feature by opening an issue at https://github.com/BerriAI/litellm/issues"
@ -442,6 +426,19 @@ def _get_gemini_url(
else:
raise ValueError(f"Unsupported mode: {mode}")
base_url = "https://generativelanguage.googleapis.com/{}/{}:{}".format(
api_version if mode == "chat" else "v1beta", _gemini_model_name, endpoint
)
params = []
if mode == "chat" and stream is True:
params.append("alt=sse")
if params:
url = f"{base_url}?{'&'.join(params)}"
else:
url = base_url
return url, endpoint

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
@ -2721,6 +2722,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)
@ -2728,6 +2734,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,
@ -2823,6 +2830,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)
@ -2830,6 +2842,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,

View file

@ -2,6 +2,7 @@
Google AI Studio /batchEmbedContents Embeddings Endpoint
"""
import asyncio
import json
from typing import Any, Dict, Literal, Optional, Union
@ -130,6 +131,36 @@ class GoogleBatchEmbeddings(VertexLLM):
client=None,
extra_headers: Optional[dict] = None,
) -> EmbeddingResponse:
optional_params = optional_params or {}
is_multimodal = _is_multimodal_input(input)
use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
mode: Literal["embedding", "batch_embedding"]
if use_embed_content:
mode = "embedding"
else:
mode = "batch_embedding"
if aembedding is True:
return self._async_batch_embeddings_with_auth_resolution( # type: ignore
model=model,
input=input,
model_response=model_response,
custom_llm_provider=custom_llm_provider,
optional_params=optional_params,
logging_obj=logging_obj,
api_key=api_key,
api_base=api_base,
use_embed_content=use_embed_content,
mode=mode,
timeout=timeout,
client=client,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_credentials=vertex_credentials,
extra_headers=extra_headers,
)
_auth_header, vertex_project = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
@ -149,16 +180,6 @@ class GoogleBatchEmbeddings(VertexLLM):
else:
sync_handler = client # type: ignore
optional_params = optional_params or {}
is_multimodal = _is_multimodal_input(input)
use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
mode: Literal["embedding", "batch_embedding"]
if use_embed_content:
mode = "embedding"
else:
mode = "batch_embedding"
auth_header, url = self._get_token_and_url(
model=model,
auth_header=_auth_header,
@ -184,22 +205,6 @@ class GoogleBatchEmbeddings(VertexLLM):
if extra_headers is not None:
headers.update(extra_headers)
if aembedding is True:
return self.async_batch_embeddings( # type: ignore
model=model,
api_base=api_base,
url=url,
data=None,
model_response=model_response,
timeout=timeout,
headers=headers,
input=input,
use_embed_content=use_embed_content,
api_key=api_key,
optional_params=optional_params,
logging_obj=logging_obj,
)
### TRANSFORMATION (sync path) ###
request_data: Any
if use_embed_content:
@ -257,6 +262,78 @@ class GoogleBatchEmbeddings(VertexLLM):
input=input,
)
async def _async_batch_embeddings_with_auth_resolution(
self,
model: str,
input: EmbeddingInput,
model_response: EmbeddingResponse,
custom_llm_provider: Literal["gemini", "vertex_ai"],
optional_params: dict,
logging_obj: Any,
api_key: Optional[str],
api_base: Optional[str],
use_embed_content: bool,
mode: Literal["embedding", "batch_embedding"],
timeout: Optional[Union[float, httpx.Timeout]],
client: Optional[AsyncHTTPHandler],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_credentials: Optional[Any],
extra_headers: Optional[dict],
) -> EmbeddingResponse:
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
gemini_auth_data = None
if custom_llm_provider == "gemini" and 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)
auth_header, url = self._get_token_and_url(
model=model,
auth_header=_auth_header,
gemini_api_key=api_key,
gemini_auth_data=gemini_auth_data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_credentials=vertex_credentials,
stream=None,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=False,
mode=mode,
)
headers = {
"Content-Type": "application/json; charset=utf-8",
}
if auth_header is not None:
if isinstance(auth_header, dict):
headers.update(auth_header)
else:
headers["Authorization"] = f"Bearer {auth_header}"
if extra_headers is not None:
headers.update(extra_headers)
return await self.async_batch_embeddings(
model=model,
api_base=api_base,
url=url,
data=None,
model_response=model_response,
timeout=timeout,
headers=headers,
client=client,
input=input,
use_embed_content=use_embed_content,
api_key=api_key,
optional_params=optional_params,
logging_obj=logging_obj,
)
async def async_batch_embeddings(
self,
model: str,

View file

@ -1,3 +1,4 @@
import asyncio
import json
from typing import Literal, Optional, Union
@ -14,6 +15,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexAIError,
VertexLLM,
)
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
from litellm.types.utils import EmbeddingResponse
from .transformation import VertexAIMultimodalEmbeddingConfig
@ -50,6 +52,25 @@ class VertexMultimodalEmbedding(VertexLLM):
timeout=300,
client=None,
) -> EmbeddingResponse:
if aembedding is True:
return self._async_multimodal_embedding_with_auth_resolution( # type: ignore
model=model,
input=input,
model_response=model_response,
custom_llm_provider=custom_llm_provider,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
api_key=api_key,
api_base=api_base,
headers=headers,
timeout=timeout,
client=client,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_credentials=vertex_credentials,
)
_auth_header, vertex_project = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
@ -108,21 +129,6 @@ class VertexMultimodalEmbedding(VertexLLM):
},
)
if aembedding is True:
return self.async_multimodal_embedding( # type: ignore
model=model,
api_base=url,
data=request_data,
timeout=timeout,
headers=headers,
client=client,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
api_key=api_key,
)
response = sync_handler.post(
url=url,
headers=headers,
@ -140,6 +146,87 @@ class VertexMultimodalEmbedding(VertexLLM):
litellm_params=litellm_params,
)
async def _async_multimodal_embedding_with_auth_resolution(
self,
model: str,
input: Union[list, str],
model_response: EmbeddingResponse,
custom_llm_provider: Literal["gemini", "vertex_ai"],
optional_params: dict,
litellm_params: dict,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str],
api_base: Optional[str],
headers: dict,
timeout: Optional[Union[float, httpx.Timeout]],
client: Optional[AsyncHTTPHandler],
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
) -> EmbeddingResponse:
_auth_header, vertex_project = await self._ensure_access_token_async(
credentials=vertex_credentials,
project_id=vertex_project,
custom_llm_provider=custom_llm_provider,
)
gemini_auth_data = None
if custom_llm_provider == "gemini" and 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)
auth_header, url = self._get_token_and_url(
model=model,
auth_header=_auth_header,
gemini_api_key=api_key,
gemini_auth_data=gemini_auth_data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_credentials=vertex_credentials,
stream=None,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
should_use_v1beta1_features=False,
mode="embedding",
)
request_data = vertex_multimodal_embedding_handler.transform_embedding_request(
model, input, optional_params, headers
)
request_headers = vertex_multimodal_embedding_handler.validate_environment(
headers=headers,
model=model,
messages=[],
optional_params=optional_params,
api_key=auth_header,
api_base=api_base,
litellm_params=litellm_params,
)
logging_obj.pre_call(
input=input,
api_key="",
additional_args={
"complete_input_dict": request_data,
"api_base": url,
"headers": request_headers,
},
)
return await self.async_multimodal_embedding(
model=model,
api_base=url,
data=request_data,
timeout=timeout,
headers=request_headers,
client=client,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
api_key=api_key,
)
async def async_multimodal_embedding(
self,
model: str,

View file

@ -1,3 +1,4 @@
import asyncio
from typing import Dict, Literal, Optional, Union
import httpx
@ -174,12 +175,18 @@ class VertexEmbedding(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)
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,

View file

@ -387,6 +387,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
@ -413,12 +414,29 @@ class VertexBase:
"Model parameter is required for Gemini custom API base URLs"
)
url = "{}/models/{}:{}".format(api_base, model, endpoint)
if gemini_api_key is None:
raise ValueError(
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
if gemini_auth_data is None and gemini_api_key is None:
from litellm.llms.gemini.common_utils import (
get_gemini_oauth_token,
)
if gemini_api_key is not None:
gemini_auth_data = get_gemini_oauth_token()
gemini_oauth_token = (
gemini_auth_data.get("token") if gemini_auth_data else None
)
if gemini_oauth_token:
auth_header = {"Authorization": f"Bearer {gemini_oauth_token}"} # type: ignore[assignment]
if gemini_auth_data and gemini_auth_data.get("project_id"):
auth_header["x-goog-user-project"] = gemini_auth_data[ # type: ignore[index]
"project_id"
]
elif gemini_api_key is not None:
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
else:
raise ValueError(
"Missing gemini_api_key. Please set `GEMINI_API_KEY` or `GEMINI_OAUTH_TOKEN`."
)
else:
# For Vertex AI
if use_psc_endpoint_format:
@ -461,6 +479,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.
@ -472,16 +491,29 @@ class VertexBase:
"""
version: Optional[Literal["v1beta1", "v1"]] = None
if custom_llm_provider == "gemini":
if not gemini_api_key:
raise ValueError(
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
)
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()
gemini_oauth_token = (
gemini_auth_data.get("token") if gemini_auth_data else None
)
url, endpoint = _get_gemini_url(
mode=mode,
model=model,
stream=stream,
)
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
if gemini_oauth_token:
auth_header = {"Authorization": f"Bearer {gemini_oauth_token}"} # type: ignore[assignment]
if gemini_auth_data and gemini_auth_data.get("project_id"):
auth_header["x-goog-user-project"] = gemini_auth_data["project_id"] # type: ignore[index]
elif gemini_api_key is None:
raise ValueError(
"Missing gemini_api_key. Please set `GEMINI_API_KEY` or `GEMINI_OAUTH_TOKEN`."
)
else:
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
else:
vertex_location = self.get_vertex_region(
vertex_region=vertex_location,
@ -504,6 +536,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,

View file

@ -103,6 +103,11 @@ from litellm.llms.cohere.common_utils import CohereModelInfo
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.llms.google_code_assist.chat import get_google_code_assist_chat
from litellm.llms.gemini.fallback_handler import (
run_gemini_acompletion_with_code_assist_fallback,
run_gemini_completion_with_code_assist_fallback,
)
from litellm.llms.vertex_ai.common_utils import (
VertexAIModelRoute,
get_vertex_ai_model_route,
@ -158,6 +163,8 @@ from litellm.utils import (
validate_openai_optional_params,
)
_google_code_assist_chat = get_google_code_assist_chat()
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
from .litellm_core_utils.core_helpers import safe_deep_copy
@ -199,7 +206,9 @@ 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,
)
from .llms.groq.chat.handler import GroqChatCompletion
from .llms.heroku.chat.transformation import HerokuChatConfig
from .llms.huggingface.embedding.handler import HuggingFaceEmbedding
@ -3448,27 +3457,76 @@ def completion( # type: ignore # noqa: PLR0915
api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE")
new_params = safe_deep_copy(optional_params or {})
response = vertex_chat_completion.completion( # type: ignore
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=new_params,
litellm_params=litellm_params, # type: ignore
logger_fn=logger_fn,
encoding=_get_encoding(),
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
logging_obj=logging,
acompletion=acompletion,
timeout=timeout,
custom_llm_provider=custom_llm_provider, # type: ignore
client=client,
api_base=api_base,
extra_headers=headers,
)
if acompletion is True:
response = run_gemini_acompletion_with_code_assist_fallback(
primary_call=lambda: vertex_chat_completion.completion( # type: ignore
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=new_params,
litellm_params=litellm_params, # type: ignore
logger_fn=logger_fn,
encoding=_get_encoding(),
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
logging_obj=logging,
acompletion=True,
timeout=timeout,
custom_llm_provider=custom_llm_provider, # type: ignore
client=client,
api_base=api_base,
extra_headers=headers,
),
fallback_kwargs={
"model": model,
"messages": messages,
"model_response": model_response,
"print_verbose": print_verbose,
"optional_params": optional_params or {},
"litellm_params": litellm_params, # type: ignore
"logging_obj": logging,
"logger_fn": logger_fn,
},
auto_fallback_to_google_code_assist=litellm.auto_fallback_to_google_code_assist,
)
else:
response = run_gemini_completion_with_code_assist_fallback(
primary_call=lambda: vertex_chat_completion.completion( # type: ignore
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=new_params,
litellm_params=litellm_params, # type: ignore
logger_fn=logger_fn,
encoding=_get_encoding(),
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
logging_obj=logging,
acompletion=False,
timeout=timeout,
custom_llm_provider=custom_llm_provider, # type: ignore
client=client,
api_base=api_base,
extra_headers=headers,
),
fallback_kwargs={
"model": model,
"messages": messages,
"model_response": model_response,
"print_verbose": print_verbose,
"optional_params": optional_params or {},
"litellm_params": litellm_params, # type: ignore
"logging_obj": logging,
"logger_fn": logger_fn,
},
auto_fallback_to_google_code_assist=litellm.auto_fallback_to_google_code_assist,
)
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (
@ -3644,6 +3702,32 @@ def completion( # type: ignore # noqa: PLR0915
)
return response
response = model_response
elif custom_llm_provider == "google_code_assist":
_ca_optional_params = optional_params or {}
if acompletion is True:
response = _google_code_assist_chat.acompletion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=_ca_optional_params,
litellm_params=litellm_params, # type: ignore
logging_obj=logging,
logger_fn=logger_fn,
)
else:
response = _google_code_assist_chat.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=_ca_optional_params,
litellm_params=litellm_params, # type: ignore
logging_obj=logging,
logger_fn=logger_fn,
)
elif custom_llm_provider == "predibase":
tenant_id = (
optional_params.pop("tenant_id", None)

View file

@ -29,6 +29,18 @@ litellm-proxy --version
litellm-proxy -v
```
## Gemini OAuth Login
Authenticate locally for Gemini OAuth-based providers (for example `google_code_assist/*`):
```bash
export GEMINI_OAUTH_CLIENT_ID="<your-google-oauth-client-id>"
export GEMINI_OAUTH_CLIENT_SECRET="<your-google-oauth-client-secret>"
litellm-proxy gemini login
```
This starts a browser-based loopback OAuth flow and stores credentials for reuse.
## Commands
### Models Management

View file

@ -0,0 +1,19 @@
import click
from litellm.llms.gemini.authenticator import GeminiAuthenticator
@click.group(name="gemini")
def gemini():
"""Gemini-specific management commands"""
pass
@gemini.command(name="login")
def login():
"""Authenticate with Gemini using the OAuth Authorization Code flow (loopback redirect)"""
try:
auth = GeminiAuthenticator()
auth.get_token()
click.echo("✅ Successfully authenticated with Gemini.")
except Exception as e:
raise click.ClickException(f"Gemini authentication failed: {e}")

View file

@ -12,6 +12,7 @@ from .commands.chat import chat
from .commands.credentials import credentials
from .commands.http import http
from .commands.keys import keys
from .commands.gemini import gemini
# local imports
from .commands.models import models
@ -107,6 +108,7 @@ cli.add_command(chat)
cli.add_command(http)
# Add the keys command group
cli.add_command(keys)
cli.add_command(gemini)
# Add the teams command group
cli.add_command(teams)
# Add the users command group

View file

@ -3222,6 +3222,7 @@ class LlmProviders(str, Enum):
VERTEX_AI = "vertex_ai"
VERTEX_AI_BETA = "vertex_ai_beta"
GEMINI = "gemini"
GOOGLE_CODE_ASSIST = "google_code_assist"
AI21 = "ai21"
BASETEN = "baseten"
BLACK_FOREST_LABS = "black_forest_labs"

View file

@ -8119,6 +8119,10 @@ class ProviderConfigManager:
LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False),
LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False),
LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False),
LlmProviders.GOOGLE_CODE_ASSIST: (
lambda: ProviderConfigManager._get_google_code_assist_config(),
False,
),
LlmProviders.AI21: (lambda: litellm.AI21ChatConfig(), False),
LlmProviders.AI21_CHAT: (lambda: litellm.AI21ChatConfig(), False),
LlmProviders.AZURE_TEXT: (lambda: litellm.AzureOpenAITextConfig(), False),
@ -8238,6 +8242,15 @@ class ProviderConfigManager:
return LangGraphConfig()
@staticmethod
def _get_google_code_assist_config() -> BaseConfig:
"""Get Google Code Assist config."""
from litellm.llms.google_code_assist.transformation import (
GoogleCodeAssistConfig,
)
return GoogleCodeAssistConfig()
@staticmethod
def get_provider_chat_config( # noqa: PLR0915
model: str, provider: LlmProviders

View file

@ -1105,6 +1105,29 @@
"realtime": true
}
},
"google_code_assist": {
"display_name": "Google Code Assist (`google_code_assist`)",
"url": "https://docs.litellm.ai/docs/providers/google_code_assist",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"interactions": false,
"a2a": false,
"vector_stores_search": false,
"count_tokens": false,
"rag_ingest": false,
"realtime": false,
"generateContent": false
}
},
"gemini": {
"display_name": "Google AI Studio - Gemini (`gemini`)",
"url": "https://docs.litellm.ai/docs/providers/gemini",

View file

@ -419,7 +419,7 @@ async def test_gemini_custom_api_base_proxy_integration():
print(f"✅ Custom API base streaming URL test passed: {result_url_streaming}")
# Test case 3: Error handling - missing API key
with pytest.raises(ValueError, match="Missing Gemini API key"):
with pytest.raises(ValueError, match="Missing gemini_api_key"):
vertex_base._check_custom_proxy(
api_base=custom_api_base,
custom_llm_provider="gemini",

View file

@ -0,0 +1,75 @@
from unittest.mock import AsyncMock, patch
import pytest
from litellm.llms.gemini.fallback_handler import (
run_gemini_acompletion_with_code_assist_fallback,
run_gemini_completion_with_code_assist_fallback,
)
def test_run_gemini_completion_with_code_assist_fallback_disabled():
def _raise_scope_error():
raise Exception("ACCESS_TOKEN_SCOPE_INSUFFICIENT")
with patch(
"litellm.llms.gemini.fallback_handler._google_code_assist_chat.completion"
) as mock_completion:
with pytest.raises(Exception, match="ACCESS_TOKEN_SCOPE_INSUFFICIENT"):
run_gemini_completion_with_code_assist_fallback(
primary_call=_raise_scope_error,
fallback_kwargs={},
auto_fallback_to_google_code_assist=False,
)
mock_completion.assert_not_called()
def test_run_gemini_completion_with_code_assist_fallback_enabled_but_not_match():
def _raise_other_error():
raise Exception("SOME_OTHER_ERROR")
with (
patch(
"litellm.llms.gemini.fallback_handler.should_fallback_to_google_code_assist",
return_value=False,
) as mock_should_fallback,
patch(
"litellm.llms.gemini.fallback_handler._google_code_assist_chat.completion"
) as mock_completion,
):
with pytest.raises(Exception, match="SOME_OTHER_ERROR"):
run_gemini_completion_with_code_assist_fallback(
primary_call=_raise_other_error,
fallback_kwargs={},
auto_fallback_to_google_code_assist=True,
)
mock_should_fallback.assert_called_once()
mock_completion.assert_not_called()
@pytest.mark.asyncio
async def test_run_gemini_acompletion_with_code_assist_fallback_enabled():
async def _raise_scope_error():
raise Exception("ACCESS_TOKEN_SCOPE_INSUFFICIENT")
with (
patch(
"litellm.llms.gemini.fallback_handler.should_fallback_to_google_code_assist",
return_value=True,
),
patch(
"litellm.llms.gemini.fallback_handler._google_code_assist_chat.acompletion",
new_callable=AsyncMock,
) as mock_acompletion,
):
mock_acompletion.return_value = "fallback-ok"
result = await run_gemini_acompletion_with_code_assist_fallback(
primary_call=_raise_scope_error,
fallback_kwargs={},
auto_fallback_to_google_code_assist=True,
)
assert result == "fallback-ok"
mock_acompletion.assert_awaited_once()

View file

@ -0,0 +1,68 @@
from unittest.mock import MagicMock, patch
import pytest
from litellm.llms.gemini.authenticator import GeminiAuthenticator
def test_get_token_raises_if_login_response_missing_access_token(tmp_path):
auth = GeminiAuthenticator()
auth.oauth_creds_file = str(tmp_path / "oauth_creds.json")
with (
patch.object(auth, "_login", return_value={"refresh_token": "abc"}),
patch.object(auth, "_write_oauth_creds"),
):
with pytest.raises(Exception, match="missing access_token"):
auth.get_token()
def test_refresh_token_raises_if_response_missing_access_token(tmp_path):
auth = GeminiAuthenticator()
auth.oauth_creds_file = str(tmp_path / "oauth_creds.json")
mock_resp = MagicMock()
mock_resp.raise_for_status.return_value = None
mock_resp.json.return_value = {"expires_in": 3600}
with (
patch.object(
GeminiAuthenticator,
"_get_oauth_client_credentials",
return_value=("client-id", "client-secret"),
),
patch("litellm.llms.gemini.authenticator.httpx.post", return_value=mock_resp),
patch.object(auth, "_write_oauth_creds"),
):
with pytest.raises(Exception, match="missing access_token"):
auth._refresh_token("refresh-token")
def test_login_times_out_without_oauth_callback(monkeypatch):
auth = GeminiAuthenticator()
monkeypatch.setenv("GEMINI_OAUTH_LOOPBACK_TIMEOUT_SECONDS", "1")
class _NoRequestServer:
def __init__(self, *args, **kwargs):
self.server_port = 12345
self.timeout = None
def handle_request(self):
return None
def server_close(self):
return None
with (
patch.object(
GeminiAuthenticator,
"_get_oauth_client_credentials",
return_value=("client-id", "client-secret"),
),
patch(
"litellm.llms.gemini.authenticator.http.server.HTTPServer", _NoRequestServer
),
patch("litellm.llms.gemini.authenticator.webbrowser.open", return_value=True),
):
with pytest.raises(Exception, match="OAuth login timed out"):
auth._login()

View file

@ -1,8 +1,13 @@
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter
from litellm.llms.gemini.common_utils import (
GeminiModelInfo,
GoogleAIStudioTokenCounter,
should_fallback_to_google_code_assist,
)
class TestGeminiModelInfo:
@ -221,3 +226,48 @@ class TestGoogleAIStudioTokenCounter:
# Verify the contents are unchanged
assert cleaned_contents == contents_without_function_response
class TestGeminiFallbackDetection:
def test_should_fallback_to_google_code_assist_for_structured_403_scope_error(self):
request = httpx.Request(
"POST", "https://generativelanguage.googleapis.com/test"
)
response = httpx.Response(
status_code=403,
json={
"error": {
"code": 403,
"message": "Request had insufficient authentication scopes.",
"status": "PERMISSION_DENIED",
"details": [{"reason": "ACCESS_TOKEN_SCOPE_INSUFFICIENT"}],
}
},
request=request,
)
err = httpx.HTTPStatusError(
"403 Client Error: Forbidden for url", request=request, response=response
)
assert should_fallback_to_google_code_assist(err) is True
def test_should_not_fallback_to_google_code_assist_for_non_403_scope_string(self):
request = httpx.Request(
"POST", "https://generativelanguage.googleapis.com/test"
)
response = httpx.Response(
status_code=500,
json={
"error": {
"message": "ACCESS_TOKEN_SCOPE_INSUFFICIENT",
}
},
request=request,
)
err = httpx.HTTPStatusError(
"500 Server Error: Internal Server Error for url",
request=request,
response=response,
)
assert should_fallback_to_google_code_assist(err) is False

View file

@ -0,0 +1,216 @@
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import json
from litellm.llms.google_code_assist.chat import (
GoogleCodeAssistChat,
get_google_code_assist_chat,
)
from litellm.types.utils import ModelResponse
class TestGoogleCodeAssist:
def test_get_google_code_assist_chat_returns_singleton(self):
assert get_google_code_assist_chat() is get_google_code_assist_chat()
@patch("litellm.llms.google_code_assist.chat._get_httpx_client")
@patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token")
def test_completion_basic(self, mock_get_token, mock_get_client):
"""
Test basic completion with mocked handshake and API call.
"""
mock_get_token.return_value = {
"token": "test-token",
"project_id": "test-project",
}
mock_client = MagicMock()
mock_get_client.return_value = mock_client
# Mock handshake response
handshake_data = {"cloudaicompanionProject": "final-project"}
mock_handshake_resp = httpx.Response(
status_code=200,
content=json.dumps(handshake_data).encode(),
request=httpx.Request("POST", "https://handshake"),
)
# Mock completion response
completion_data = {
"response": {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello world"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 2,
"totalTokenCount": 7,
},
}
}
mock_completion_resp = httpx.Response(
status_code=200,
content=json.dumps(completion_data).encode(),
request=httpx.Request("POST", "https://completion"),
)
mock_client.post.side_effect = [mock_handshake_resp, mock_completion_resp]
handler = GoogleCodeAssistChat()
response = handler.completion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
assert response.choices[0].message.content == "Hello world"
assert response.usage.total_tokens == 7
@patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token")
def test_completion_raises_when_token_missing(self, mock_get_token):
mock_get_token.return_value = {"token": None}
handler = GoogleCodeAssistChat()
with pytest.raises(Exception, match="Missing Gemini OAuth token value"):
handler.completion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
@pytest.mark.asyncio
@patch("litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post")
@patch(
"litellm.llms.google_code_assist.chat.AsyncHTTPHandler.close",
new_callable=AsyncMock,
)
@patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token")
async def test_acompletion_basic(
self, mock_get_token, mock_async_close, mock_async_post
):
"""
Test async completion.
"""
mock_get_token.return_value = {"token": "test-token"}
# Handshake (async)
handshake_data = {"cloudaicompanionProject": "final-project"}
mock_handshake_resp = httpx.Response(
status_code=200,
content=json.dumps(handshake_data).encode(),
request=httpx.Request("POST", "https://handshake"),
)
# Completion (async)
completion_data = {
"response": {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Async success"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 2,
"totalTokenCount": 7,
},
}
}
mock_completion_resp = httpx.Response(
status_code=200,
content=json.dumps(completion_data).encode(),
request=httpx.Request("POST", "https://completion"),
)
mock_async_post.side_effect = [mock_handshake_resp, mock_completion_resp]
handler = GoogleCodeAssistChat()
response = await handler.acompletion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
assert response.choices[0].message.content == "Async success"
mock_async_close.assert_awaited_once()
@pytest.mark.asyncio
@patch(
"litellm.llms.google_code_assist.chat.asyncio.to_thread", new_callable=AsyncMock
)
async def test_acompletion_uses_thread_for_oauth_lookup(self, mock_to_thread):
mock_to_thread.return_value = {"token": "test-token"}
with (
patch(
"litellm.llms.google_code_assist.chat.GoogleCodeAssistChat._ahandle_handshake",
new_callable=AsyncMock,
return_value="final-project",
),
patch(
"litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_async_post,
patch(
"litellm.llms.google_code_assist.chat.AsyncHTTPHandler.close",
new_callable=AsyncMock,
),
):
completion_data = {
"response": {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Async success"}],
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 2,
"totalTokenCount": 7,
},
}
}
mock_completion_resp = httpx.Response(
status_code=200,
content=json.dumps(completion_data).encode(),
request=httpx.Request("POST", "https://completion"),
)
mock_async_post.return_value = mock_completion_resp
handler = GoogleCodeAssistChat()
await handler.acompletion(
model="google_code_assist/gemini-1.5-flash",
messages=[{"role": "user", "content": "hi"}],
model_response=ModelResponse(),
print_verbose=False,
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
mock_to_thread.assert_awaited_once()

View file

@ -0,0 +1,42 @@
from unittest.mock import patch
from litellm.llms.google_code_assist.transformation import GoogleCodeAssistConfig
def test_transform_request_uses_base_model_name_for_internal_gemini_helpers():
class _CaptureConfig(GoogleCodeAssistConfig):
def __init__(self):
super().__init__()
self.model_seen_by_map = None
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
messages: list,
) -> dict:
self.model_seen_by_map = model
return {"thinkingConfig": {"includeThoughts": True}}
config = _CaptureConfig()
with (
patch(
"litellm.llms.vertex_ai.gemini.transformation._transform_system_message",
return_value=(None, [{"role": "user", "content": "hi"}]),
),
patch(
"litellm.llms.vertex_ai.gemini.transformation._gemini_convert_messages_with_history",
return_value=[{"role": "user", "parts": [{"text": "hi"}]}],
) as mock_convert,
):
result = config.transform_request(
model="google_code_assist/gemini-2.5-pro",
messages=[{"role": "user", "content": "hi"}],
optional_params={"thinkingConfig": {"includeThoughts": False}},
litellm_params={},
)
assert config.model_seen_by_map == "gemini-2.5-pro"
assert mock_convert.call_args.kwargs["model"] == "gemini-2.5-pro"
assert result["model"] == "gemini-2.5-pro"

View file

@ -11,6 +11,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm.llms.vertex_ai.common_utils import (
_get_gemini_url,
_get_vertex_url,
convert_anyof_null_to_nullable,
get_vertex_location_from_url,
@ -603,6 +604,33 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
assert url == expected_url
def test_get_gemini_url_stream_query_param_only_for_chat_mode():
chat_url, _ = _get_gemini_url(
mode="chat",
model="gemini-1.5-flash",
stream=True,
)
embedding_url, _ = _get_gemini_url(
mode="embedding",
model="gemini-1.5-flash",
stream=True,
)
assert "alt=sse" in chat_url
assert "alt=sse" not in embedding_url
def test_get_gemini_url_requires_literal_true_for_streaming_endpoint_and_alt_sse():
url, endpoint = _get_gemini_url(
mode="chat",
model="gemini-1.5-flash",
stream=1, # type: ignore[arg-type] # truthy non-bool should not be treated as streaming=True
)
assert endpoint == "generateContent"
assert "alt=sse" not in url
@pytest.mark.parametrize(
"model_cost_entry, vertex_region, expected_region",
[
@ -1096,6 +1124,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():
"""
@ -1457,11 +1537,7 @@ def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata()
def test_vertex_request_labels_from_litellm_params_accepts_litellm_metadata():
lp = {
"litellm_metadata": {
"requester_metadata": {"team": "platform", "count": 3}
}
}
lp = {"litellm_metadata": {"requester_metadata": {"team": "platform", "count": 3}}}
assert vertex_request_labels_from_litellm_params(lp) == {"team": "platform"}

View file

@ -821,7 +821,7 @@ class TestVertexBase:
if custom_llm_provider == "gemini" and api_base and gemini_api_key is None:
# Test case 5: Should raise ValueError for Gemini without API key
with pytest.raises(ValueError, match="Missing Gemini API key"):
with pytest.raises(ValueError, match="Missing gemini_api_key"):
vertex_base._check_custom_proxy(
api_base=api_base,
custom_llm_provider=custom_llm_provider,