diff --git a/litellm/llms/google_code_assist/chat.py b/litellm/llms/google_code_assist/chat.py index beea6ca997e..07be7bb7ad9 100644 --- a/litellm/llms/google_code_assist/chat.py +++ b/litellm/llms/google_code_assist/chat.py @@ -7,6 +7,10 @@ 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: """ @@ -65,7 +69,7 @@ class GoogleCodeAssistChat: ) # 4. Call Completion API - url = "https://cloudcode-pa.googleapis.com/v1internal:generateContent" + url = _CODE_ASSIST_GENERATE_URL headers = self._get_headers(token) response = client.post( @@ -133,7 +137,7 @@ class GoogleCodeAssistChat: data = self.config.transform_request( model, messages, optional_params, litellm_params ) - url = "https://cloudcode-pa.googleapis.com/v1internal:generateContent" + url = _CODE_ASSIST_GENERATE_URL headers = self._get_headers(token) response = await async_handler.post(url=url, headers=headers, json=data) @@ -165,7 +169,7 @@ class GoogleCodeAssistChat: 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_url = _CODE_ASSIST_LOAD_URL load_headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", @@ -197,7 +201,7 @@ class GoogleCodeAssistChat: 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_url = _CODE_ASSIST_LOAD_URL load_headers = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", diff --git a/litellm/llms/google_code_assist/transformation.py b/litellm/llms/google_code_assist/transformation.py index 273b53d7ba4..88e570cb591 100644 --- a/litellm/llms/google_code_assist/transformation.py +++ b/litellm/llms/google_code_assist/transformation.py @@ -4,6 +4,7 @@ import copy import httpx from typing import Any, List, Optional +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 @@ -127,6 +128,20 @@ class GoogleCodeAssistConfig(VertexGeminiConfig): 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"] + } + + if ( + "thinkingConfig" in optional_params + and "thinkingConfig" not in generation_config + ): + verbose_logger.warning( + "google_code_assist: `thinkingConfig` was provided but not mapped into generationConfig." + ) vertex_request = { "contents": contents, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 3f1bccaccfc..ca127a4391e 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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 @@ -2504,6 +2505,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) @@ -2511,6 +2517,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, @@ -2606,6 +2613,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) @@ -2613,6 +2625,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, @@ -2801,6 +2814,7 @@ class VertexLLM(VertexBase): auth_header, url = self._get_token_and_url( model=model, gemini_api_key=gemini_api_key, + gemini_auth_data=None, auth_header=_auth_header, vertex_project=vertex_project, vertex_location=vertex_location, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 2ffe40d6185..e6f9eb9908a 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -368,6 +368,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 @@ -395,8 +396,7 @@ class VertexBase: ) url = "{}/models/{}:{}".format(api_base, model, endpoint) - gemini_auth_data = None - if gemini_api_key is None: + if gemini_auth_data is None and gemini_api_key is None: from litellm.llms.gemini.common_utils import ( get_gemini_oauth_token, ) @@ -460,6 +460,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. @@ -471,8 +472,7 @@ class VertexBase: """ version: Optional[Literal["v1beta1", "v1"]] = None if custom_llm_provider == "gemini": - gemini_auth_data = None - if gemini_api_key is None: + 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() @@ -491,6 +491,10 @@ class VertexBase: 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 None: + raise ValueError( + "Missing gemini_api_key. Please set `GEMINI_API_KEY` or `GEMINI_OAUTH_TOKEN`." + ) else: auth_header = ( None # this field is not used for gemini when using api key @@ -517,6 +521,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, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 79db733e4aa..d2d1948ebf1 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1074,6 +1074,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(): """