fix(gemini): address CI failures — MyPy, provider docs, test assertions

- Add google_code_assist to provider_endpoints_support.json (code-quality CI)
- Fix MyPy signature mismatches in GoogleCodeAssistConfig.transform_request
  and transform_response (add missing headers/api_key/json_mode params)
- Fix map_openai_params 4th arg: pass drop_params=False instead of messages list
- Add type: ignore[assignment/index] on vertex_llm_base auth_header dict assignments
- Type load_payload as Dict[str, Any] in chat.py to allow mixed-value dict
- Update test match= strings from "Missing Gemini API key" to "Missing gemini_api_key"
  to match updated error message in _check_custom_proxy and _get_token_and_url

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
balazss 2026-05-04 14:11:30 -07:00
parent 91725bf46c
commit b0d5e0c2cb
6 changed files with 37 additions and 11 deletions

View file

@ -1,6 +1,6 @@
import asyncio
import httpx
from typing import Any, Optional
from typing import Any, Dict, Optional
import litellm
from litellm._logging import verbose_logger
@ -175,7 +175,7 @@ class GoogleCodeAssistChat:
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload = {
load_payload: Dict[str, Any] = {
"metadata": {
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",
@ -207,7 +207,7 @@ class GoogleCodeAssistChat:
"Content-Type": "application/json",
"User-Agent": "GeminiCLI/litellm",
}
load_payload = {
load_payload: Dict[str, Any] = {
"metadata": {
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",

View file

@ -55,6 +55,7 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
messages: list,
optional_params: dict,
litellm_params: dict,
headers: dict = {},
) -> dict:
"""
Transforms standard LiteLLM request to Code Assist API format.
@ -84,7 +85,7 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
generation_config = {}
# Handle parameter mapping
base_params = self.map_openai_params(
{}, optional_params.copy(), model_name, messages
{}, optional_params.copy(), model_name, False
)
for key in ["temperature", "topP", "topK", "maxOutputTokens", "stopSequences"]:
@ -158,6 +159,8 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
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.
@ -168,7 +171,7 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
return super().transform_response(
model=model,
raw_response=ParsedJSONResponseAdapter(gemini_response),
raw_response=ParsedJSONResponseAdapter(gemini_response), # type: ignore[arg-type]
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,

View file

@ -426,9 +426,9 @@ class VertexBase:
)
if gemini_oauth_token:
auth_header = {"Authorization": f"Bearer {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[
auth_header["x-goog-user-project"] = gemini_auth_data[ # type: ignore[index]
"project_id"
]
elif gemini_api_key is not None:
@ -505,9 +505,9 @@ class VertexBase:
stream=stream,
)
if gemini_oauth_token:
auth_header = {"Authorization": f"Bearer {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"]
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`."

View file

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

View file

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

View file

@ -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,