mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
Merge b0d5e0c2cb into 934ecdca78
This commit is contained in:
commit
db3991ff33
30 changed files with 2070 additions and 96 deletions
84
docs/my-website/docs/providers/google_code_assist.md
Normal file
84
docs/my-website/docs/providers/google_code_assist.md
Normal 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.
|
||||
|
|
@ -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] = (
|
||||
|
|
|
|||
291
litellm/llms/gemini/authenticator.py
Normal file
291
litellm/llms/gemini/authenticator.py
Normal 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
|
||||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
57
litellm/llms/gemini/fallback_handler.py
Normal file
57
litellm/llms/gemini/fallback_handler.py
Normal 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)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
258
litellm/llms/google_code_assist/chat.py
Normal file
258
litellm/llms/google_code_assist/chat.py
Normal 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
|
||||
182
litellm/llms/google_code_assist/transformation.py
Normal file
182
litellm/llms/google_code_assist/transformation.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
128
litellm/main.py
128
litellm/main.py
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
19
litellm/proxy/client/cli/commands/gemini.py
Normal file
19
litellm/proxy/client/cli/commands/gemini.py
Normal 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}")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
75
tests/test_litellm/llms/gemini/test_fallback_handler.py
Normal file
75
tests/test_litellm/llms/gemini/test_fallback_handler.py
Normal 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()
|
||||
68
tests/test_litellm/llms/gemini/test_gemini_authenticator.py
Normal file
68
tests/test_litellm/llms/gemini/test_gemini_authenticator.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
0
tests/test_litellm/llms/google_code_assist/__init__.py
Normal file
0
tests/test_litellm/llms/google_code_assist/__init__.py
Normal 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()
|
||||
|
|
@ -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"
|
||||
|
|
@ -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"}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue