diff --git a/litellm/llms/google_code_assist/chat.py b/litellm/llms/google_code_assist/chat.py index b71f714abdc..26c30e2b3d1 100644 --- a/litellm/llms/google_code_assist/chat.py +++ b/litellm/llms/google_code_assist/chat.py @@ -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", diff --git a/litellm/llms/google_code_assist/transformation.py b/litellm/llms/google_code_assist/transformation.py index 6de5b608ea0..49c87c10b04 100644 --- a/litellm/llms/google_code_assist/transformation.py +++ b/litellm/llms/google_code_assist/transformation.py @@ -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, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index fcc1073fe62..a3c6637408b 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -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`." diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index ed49c146210..cf1a011032e 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -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", diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py index dbe30037313..8005dd20845 100644 --- a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py +++ b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py @@ -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", diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 88aac07a0c9..dbd4581cf9a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -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,