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

feat(gemini): add oauth login and google code assist provider support

Browse source
This commit is contained in:
balazss 2026-03-17 20:05:28 -07:00
parent cfeafbe388
commit f7374ad800
18 changed files with 1324 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
docs/my-website/docs/providers/gemini.md
View file

@ -34,6 +34,26 @@ import os
os.environ["GEMINI_API_KEY"] = "your-api-key"
```
## OAuth Authentication (Optional)
You can also authenticate Gemini requests with OAuth credentials.
```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
```
LiteLLM then reads OAuth credentials from:
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`
For Google Code Assist-specific usage, see [Google Code Assist](./google_code_assist.md).
## Sample Usage
```python
from litellm import completion
@ -2459,4 +2479,3 @@ LiteLLM automatically calculates costs using `output_cost_per_image_token` from
```
For more details, see [Google's Gemini pricing documentation](https://ai.google.dev/gemini-api/docs/pricing).

84
docs/my-website/docs/providers/google_code_assist.md Normal file
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.

1
docs/my-website/sidebars.js
View file

@ -796,6 +796,7 @@ const sidebars = {
label: "Google AI Studio",
items: [
"providers/gemini",
"providers/google_code_assist",
"providers/gemini/videos",
"providers/google_ai_studio/files",
"providers/google_ai_studio/image_gen",

260
litellm/llms/gemini/authenticator.py Normal file
View file

@ -0,0 +1,260 @@
import json
import os
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
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 creds.get("access_token")
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 creds.get("access_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 = os.open(
self.oauth_creds_file,
os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
0o600,
)
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(creds, f)
try:
os.chmod(self.oauth_creds_file, 0o600)
except OSError:
pass
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
server.handle_request()
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

116
litellm/llms/gemini/common_utils.py
View file

@ -5,6 +5,7 @@ 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
@ -154,6 +155,121 @@ 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.
"""
return "ACCESS_TOKEN_SCOPE_INSUFFICIENT" in str(error)
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 = creds_data["token"].get("accessToken")
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."""

47
litellm/llms/gemini/fallback_handler.py Normal file
View file

@ -0,0 +1,47 @@
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 GoogleCodeAssistChat
async def run_gemini_acompletion_with_code_assist_fallback(
primary_call: Awaitable[Any],
fallback_kwargs: Dict[str, Any],
) -> 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 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 GoogleCodeAssistChat().acompletion(**fallback_kwargs)
def run_gemini_completion_with_code_assist_fallback(
primary_call: Callable[[], Any],
fallback_kwargs: Dict[str, Any],
) -> 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 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 GoogleCodeAssistChat().completion(**fallback_kwargs)

226
litellm/llms/google_code_assist/chat.py Normal file
View file

@ -0,0 +1,226 @@
import httpx
from typing import Any, 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
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")
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 = "https://cloudcode-pa.googleapis.com/v1internal:generateContent"
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 = 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")
async_handler = AsyncHTTPHandler()
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 = "https://cloudcode-pa.googleapis.com/v1internal:generateContent"
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,
)
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 = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload = {
"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 = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
load_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload = {
"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))

195
litellm/llms/google_code_assist/transformation.py Normal file
View file

@ -0,0 +1,195 @@
import json
import uuid
import copy
import httpx
from typing import Any, List, Optional
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 __init__(
self,
temperature: Optional[float] = None,
max_output_tokens: Optional[int] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
stop_sequences: Optional[list] = None,
) -> None:
super().__init__(
temperature=temperature,
max_output_tokens=max_output_tokens,
top_p=top_p,
top_k=top_k,
stop_sequences=stop_sequences,
)
def get_supported_openai_params(self, model: str) -> List[str]:
return super().get_supported_openai_params(model)
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
messages: list,
) -> dict:
return super().map_openai_params(
non_default_params, optional_params, model, messages
)
def transform_request(
self,
model: str,
messages: list,
optional_params: dict,
litellm_params: 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)
# 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
)
# 2. Build vertex-style nested request
generation_config = {}
# Handle parameter mapping
base_params = self.map_openai_params(
{}, optional_params.copy(), model, messages
)
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")
}
elif "thinkingConfig" in optional_params:
generation_config["thinkingConfig"] = optional_params["thinkingConfig"]
elif "include_thoughts" in optional_params:
generation_config["thinkingConfig"] = {
"includeThoughts": optional_params["include_thoughts"]
}
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()}"[:13]
model_name = model.split("/")[-1]
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,
) -> 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),
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
)

33
litellm/llms/vertex_ai/common_utils.py
View file

@ -338,6 +338,7 @@ def _get_gemini_url(
model: str,
stream: Optional[bool],
gemini_api_key: Optional[str],
gemini_oauth_token: Optional[str] = None,
) -> Tuple[str, str]:
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
@ -348,32 +349,17 @@ def _get_gemini_url(
"v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta"
)
endpoint = "generateContent"
if mode == "chat":
endpoint = "generateContent"
if stream is True:
endpoint = "streamGenerateContent"
url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}&alt=sse".format(
api_version, _gemini_model_name, endpoint, gemini_api_key
)
else:
url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}".format(
api_version, _gemini_model_name, endpoint, gemini_api_key
)
elif mode == "embedding":
endpoint = "embedContent"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
_gemini_model_name, endpoint, gemini_api_key
)
elif mode == "batch_embedding":
endpoint = "batchEmbedContents"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
_gemini_model_name, endpoint, gemini_api_key
)
elif mode == "count_tokens":
endpoint = "countTokens"
url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format(
_gemini_model_name, endpoint, gemini_api_key
)
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"
@ -381,6 +367,21 @@ 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 gemini_api_key and not gemini_oauth_token:
params.append(f"key={gemini_api_key}")
if stream:
params.append("alt=sse")
if params:
url = f"{base_url}?{'&'.join(params)}"
else:
url = base_url
return url, endpoint

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

@ -394,12 +394,30 @@ class VertexBase:
"Model parameter is required for Gemini custom API base URLs"
)
url = "{}/models/{}:{}".format(api_base, model, endpoint)
gemini_auth_data = None
if gemini_api_key is None:
raise ValueError(
"Missing gemini_api_key, please set `GEMINI_API_KEY`"
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}"}
if gemini_auth_data and gemini_auth_data.get("project_id"):
auth_header["x-goog-user-project"] = gemini_auth_data[
"project_id"
]
elif gemini_api_key is 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:
@ -453,13 +471,30 @@ class VertexBase:
"""
version: Optional[Literal["v1beta1", "v1"]] = None
if custom_llm_provider == "gemini":
gemini_auth_data = None
if 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,
gemini_api_key=gemini_api_key,
gemini_oauth_token=gemini_oauth_token,
)
auth_header = None # this field is not used for gemin
if gemini_oauth_token:
auth_header = {"Authorization": f"Bearer {gemini_oauth_token}"}
if gemini_auth_data and gemini_auth_data.get("project_id"):
auth_header["x-goog-user-project"] = gemini_auth_data["project_id"]
else:
auth_header = (
None # this field is not used for gemini when using api key
)
else:
vertex_location = self.get_vertex_region(
vertex_region=vertex_location,

129
litellm/main.py
View file

@ -101,6 +101,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 GoogleCodeAssistChat
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,
@ -197,7 +202,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
@ -3436,27 +3443,74 @@ 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=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,
"litellm_params": litellm_params, # type: ignore
"logging_obj": logging,
"logger_fn": logger_fn,
},
)
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,
"litellm_params": litellm_params, # type: ignore
"logging_obj": logging,
"logger_fn": logger_fn,
},
)
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (
@ -3632,6 +3686,33 @@ def completion( # type: ignore # noqa: PLR0915
)
return response
response = model_response
elif custom_llm_provider == "google_code_assist":
google_code_assist_chat = GoogleCodeAssistChat()
if acompletion is True:
response = google_code_assist_chat.acompletion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=optional_params,
litellm_params=litellm_params, # type: ignore
logging_obj=logging,
logger_fn=logger_fn,
)
else:
model_response = google_code_assist_chat.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=optional_params,
litellm_params=litellm_params, # type: ignore
logging_obj=logging,
logger_fn=logger_fn,
)
response = model_response
elif custom_llm_provider == "predibase":
tenant_id = (
optional_params.pop("tenant_id", None)
@ -7533,9 +7614,7 @@ def stream_chunk_builder( # noqa: PLR0915
# the final chunk.
all_annotations: list = []
for ac in annotation_chunks:
all_annotations.extend(
ac["choices"][0]["delta"]["annotations"]
)
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations
audio_chunks = [

12
litellm/proxy/client/cli/README.md
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

19
litellm/proxy/client/cli/commands/gemini.py Normal file
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():
"""Login to Gemini using OAuth Device Flow (Loopback)"""
try:
auth = GeminiAuthenticator()
auth.get_token()
click.echo("✅ Successfully authenticated with Gemini.")
except Exception as e:
raise click.ClickException(f"Gemini authentication failed: {e}")

2
litellm/proxy/client/cli/main.py
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
@ -105,6 +106,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

1
litellm/types/utils.py
View file

@ -3151,6 +3151,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"

13
litellm/utils.py
View file

@ -7981,6 +7981,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),
@ -8100,6 +8104,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

73
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand.
[[package]]
name = "a2a-sdk"
@ -7,11 +7,11 @@ description = "A2A Python SDK"
optional = false
python-versions = ">=3.10"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.10\""
files = [
{file = "a2a_sdk-0.3.22-py3-none-any.whl", hash = "sha256:b98701135bb90b0ff85d35f31533b6b7a299bf810658c1c65f3814a6c15ea385"},
{file = "a2a_sdk-0.3.22.tar.gz", hash = "sha256:77a5694bfc4f26679c11b70c7f1062522206d430b34bc1215cfbb1eba67b7e7d"},
]
markers = {main = "python_version >= \"3.10\" and extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
google-api-core = ">=1.26.0"
@ -385,6 +385,7 @@ files = [
{file = "azure_core-1.36.0-py3-none-any.whl", hash = "sha256:fee9923a3a753e94a259563429f3644aaf05c486d45b1215d098115102d91d3b"},
{file = "azure_core-1.36.0.tar.gz", hash = "sha256:22e5605e6d0bf1d229726af56d9e92bc37b6e726b141a18be0b4d424131741b7"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
requests = ">=2.21.0"
@ -405,6 +406,7 @@ files = [
{file = "azure_identity-1.25.1-py3-none-any.whl", hash = "sha256:e9edd720af03dff020223cd269fa3a61e8f345ea75443858273bcb44844ab651"},
{file = "azure_identity-1.25.1.tar.gz", hash = "sha256:87ca8328883de6036443e1c37b40e8dc8fb74898240f61071e09d2e369361456"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
azure-core = ">=1.31.0"
@ -598,7 +600,7 @@ files = [
{file = "cachetools-6.2.2-py3-none-any.whl", hash = "sha256:6c09c98183bf58560c97b2abfcedcbaf6a896a490f534b031b661d3723b45ace"},
{file = "cachetools-6.2.2.tar.gz", hash = "sha256:8e6d266b25e539df852251cfd6f990b4bc3a141db73b939058d809ebd2590fc6"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "certifi"
@ -705,7 +707,7 @@ files = [
{file = "cffi-2.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:b882b3df248017dba09d6b16defe9b5c407fe32fc7c65a9c69798e6175601be9"},
{file = "cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529"},
]
markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"proxy\"", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
[package.dependencies]
pycparser = {version = "*", markers = "implementation_name != \"PyPy\""}
@ -1055,6 +1057,7 @@ files = [
{file = "cryptography-43.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff"},
{file = "cryptography-43.0.3.tar.gz", hash = "sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805"},
]
markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\") or extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""}
@ -1837,11 +1840,11 @@ description = "Google API client core library"
optional = false
python-versions = ">=3.7"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.14\""
files = [
{file = "google_api_core-2.25.2-py3-none-any.whl", hash = "sha256:e9a8f62d363dc8424a8497f4c2a47d6bcda6c16514c935629c257ab5d10210e7"},
{file = "google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300"},
]
markers = {main = "python_version >= \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.14\""}
[package.dependencies]
google-auth = ">=2.14.1,<3.0.0"
@ -1869,7 +1872,7 @@ files = [
{file = "google_api_core-2.28.1-py3-none-any.whl", hash = "sha256:4021b0f8ceb77a6fb4de6fde4502cecab45062e66ff4f2895169e0b35bc9466c"},
{file = "google_api_core-2.28.1.tar.gz", hash = "sha256:2b405df02d68e68ce0fbc138559e6036559e685159d148ae5861013dc201baf8"},
]
markers = {main = "(python_version >= \"3.10\" or extra == \"google\" or extra == \"extra-proxy\") and python_version < \"3.14\"", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
markers = {main = "python_version < \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
[package.dependencies]
google-auth = ">=2.14.1,<3.0.0"
@ -1906,7 +1909,7 @@ files = [
{file = "google_auth-2.43.0-py2.py3-none-any.whl", hash = "sha256:af628ba6fa493f75c7e9dbe9373d148ca9f4399b5ea29976519e0a3848eddd16"},
{file = "google_auth-2.43.0.tar.gz", hash = "sha256:88228eee5fc21b62a1b5fe773ca15e67778cb07dc8363adcb4a8827b52d81483"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
cachetools = ">=2.0.0,<7.0"
@ -2078,11 +2081,11 @@ files = [
]
[package.dependencies]
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev"
grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev"
proto-plus = ">=1.22.3,<2.0.0dev"
protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev"
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0.dev0", extras = ["grpc"]}
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0.dev0"
grpc-google-iam-v1 = ">=0.12.4,<1.0.0.dev0"
proto-plus = ">=1.22.3,<2.0.0.dev0"
protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0"
[[package]]
name = "google-cloud-resource-manager"
@ -2264,7 +2267,7 @@ files = [
{file = "googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038"},
{file = "googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\""}
[package.dependencies]
grpcio = {version = ">=1.44.0,<2.0.0", optional = true, markers = "extra == \"grpc\""}
@ -2673,11 +2676,11 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX."
optional = false
python-versions = ">=3.9"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.10\""
files = [
{file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"},
{file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"},
]
markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "huey"
@ -3042,7 +3045,7 @@ files = [
[package.dependencies]
attrs = ">=22.2.0"
jsonschema-specifications = ">=2023.03.6"
jsonschema-specifications = ">=2023.3.6"
referencing = ">=0.28.4"
rpds-py = ">=0.7.1"
@ -3219,15 +3222,15 @@ files = [
[[package]]
name = "litellm-proxy-extras"
version = "0.4.56"
version = "0.4.57"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
optional = true
python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "litellm_proxy_extras-0.4.56-py3-none-any.whl", hash = "sha256:52dbe3b5358c790e77e12f1ec5ef8e7508b383c2aaf41299750b6fb400908ee7"},
{file = "litellm_proxy_extras-0.4.56.tar.gz", hash = "sha256:63ad59baa0defccc5c929cfd933ee7e32a6614b0fc5fa0fc45a12d7608e33f08"},
{file = "litellm_proxy_extras-0.4.57-py3-none-any.whl", hash = "sha256:04538223cd80318a72d70c6e10f701598e58c763368296a6503c674c92fbdb62"},
{file = "litellm_proxy_extras-0.4.57.tar.gz", hash = "sha256:ef9b95dc42237614216833bd5d46ebf9dea1caa5ea14ea1a66d7f7842b224ec2"},
]
[[package]]
@ -3713,6 +3716,7 @@ files = [
{file = "msal-1.34.0-py3-none-any.whl", hash = "sha256:f669b1644e4950115da7a176441b0e13ec2975c29528d8b9e81316023676d6e1"},
{file = "msal-1.34.0.tar.gz", hash = "sha256:76ba83b716ea5a6d75b0279c0ac353a0e05b820ca1f6682c0eb7f45190c43c2f"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
cryptography = ">=2.5,<49"
@ -3733,6 +3737,7 @@ files = [
{file = "msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca"},
{file = "msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
msal = ">=1.29,<2"
@ -3983,6 +3988,7 @@ files = [
{file = "nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9"},
{file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"},
]
markers = {main = "extra == \"extra-proxy\""}
[[package]]
name = "numpy"
@ -4105,7 +4111,7 @@ files = [
{file = "opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950"},
{file = "opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
importlib-metadata = ">=6.0,<8.8.0"
@ -4220,7 +4226,7 @@ files = [
{file = "opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c"},
{file = "opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
opentelemetry-api = "1.39.1"
@ -4238,7 +4244,7 @@ files = [
{file = "opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb"},
{file = "opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
opentelemetry-api = "1.39.1"
@ -4722,6 +4728,7 @@ files = [
{file = "prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a"},
{file = "prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693"},
]
markers = {main = "extra == \"extra-proxy\""}
[package.dependencies]
click = ">=7.1.2"
@ -4895,7 +4902,7 @@ files = [
{file = "proto_plus-1.26.1-py3-none-any.whl", hash = "sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66"},
{file = "proto_plus-1.26.1.tar.gz", hash = "sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
protobuf = ">=3.19.0,<7.0.0"
@ -4923,7 +4930,7 @@ files = [
{file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"},
{file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\""}
[[package]]
name = "psutil"
@ -5083,7 +5090,7 @@ files = [
{file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"},
{file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "pyasn1-modules"
@ -5096,7 +5103,7 @@ files = [
{file = "pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a"},
{file = "pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
pyasn1 = ">=0.6.1,<0.7.0"
@ -5124,7 +5131,7 @@ files = [
{file = "pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934"},
{file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"},
]
markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
[[package]]
name = "pydantic"
@ -5347,6 +5354,7 @@ files = [
{file = "pyjwt-2.12.1-py3-none-any.whl", hash = "sha256:28ca37c070cad8ba8cd9790cd940535d40274d22f80ab87f3ac6a713e6e8454c"},
{file = "pyjwt-2.12.1.tar.gz", hash = "sha256:c74a7a2adf861c04d002db713dd85f84beb242228e671280bf709d765b03672b"},
]
markers = {main = "(python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")"}
[package.dependencies]
cryptography = {version = ">=3.4.0", optional = true, markers = "extra == \"crypto\""}
@ -6290,7 +6298,7 @@ files = [
{file = "rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762"},
{file = "rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
pyasn1 = ">=0.1.3"
@ -6336,10 +6344,10 @@ files = [
]
[package.dependencies]
botocore = ">=1.37.4,<2.0a.0"
botocore = ">=1.37.4,<2.0a0"
[package.extras]
crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"]
crt = ["botocore[crt] (>=1.37.4,<2.0a0)"]
[[package]]
name = "scikit-learn"
@ -6492,9 +6500,9 @@ tornado = ">=6.4.2,<7"
urllib3 = ">=1.26,<3"
[package.extras]
all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.00)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.0)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
bedrock = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)"]
cohere = ["cohere (>=5.9.4,<6.00)"]
cohere = ["cohere (>=5.9.4,<6.0)"]
dev = ["dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "ipykernel (>=6.25.0,<7)", "mypy (>=1.7.1,<2)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
docs = ["pydoc-markdown (>=4.8.2) ; python_version < \"3.12\""]
fastembed = ["fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\""]
@ -7222,6 +7230,7 @@ files = [
{file = "tomlkit-0.13.3-py3-none-any.whl", hash = "sha256:c89c649d79ee40629a9fda55f8ace8c6a1b42deb912b2a8fd8d942ddadb606b0"},
{file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"},
]
markers = {main = "extra == \"extra-proxy\""}
[[package]]
name = "tornado"
@ -7994,4 +8003,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "1f3bbf967451633fb6290ba88980bdf4fbf83420024b14e862d1da717d903684"
content-hash = "22cdc8096e0c8296827734393f5ab6e66088f397b5295caa1d277466d1fde1e8"

127
tests/test_litellm/llms/google_code_assist/test_google_code_assist.py Normal file
View file

@ -0,0 +1,127 @@
import pytest
from unittest.mock import MagicMock, patch
import httpx
import json
from litellm.llms.google_code_assist.chat import GoogleCodeAssistChat
from litellm.types.utils import ModelResponse
class TestGoogleCodeAssist:
@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
@pytest.mark.asyncio
@patch("litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post")
@patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token")
async def test_acompletion_basic(self, mock_get_token, 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"