mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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:
parent
91725bf46c
commit
b0d5e0c2cb
6 changed files with 37 additions and 11 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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`."
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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