From 0872a655fdba4b6767d399ce1b92797b26dd835a Mon Sep 17 00:00:00 2001 From: balazss Date: Wed, 18 Mar 2026 00:43:26 -0700 Subject: [PATCH] fix(vertex): align stream checks and share code-assist singleton --- litellm/llms/gemini/fallback_handler.py | 4 ++-- litellm/llms/google_code_assist/chat.py | 10 ++++++++++ litellm/llms/vertex_ai/common_utils.py | 3 +-- litellm/main.py | 4 ++-- .../google_code_assist/test_google_code_assist.py | 8 +++++++- .../llms/vertex_ai/test_vertex_ai_common_utils.py | 13 +++++++++++++ 6 files changed, 35 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini/fallback_handler.py b/litellm/llms/gemini/fallback_handler.py index 08cccf222c7..548cf5d5918 100644 --- a/litellm/llms/gemini/fallback_handler.py +++ b/litellm/llms/gemini/fallback_handler.py @@ -2,9 +2,9 @@ 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 +from litellm.llms.google_code_assist.chat import get_google_code_assist_chat -_google_code_assist_chat = GoogleCodeAssistChat() +_google_code_assist_chat = get_google_code_assist_chat() async def run_gemini_acompletion_with_code_assist_fallback( diff --git a/litellm/llms/google_code_assist/chat.py b/litellm/llms/google_code_assist/chat.py index 07be7bb7ad9..b71f714abdc 100644 --- a/litellm/llms/google_code_assist/chat.py +++ b/litellm/llms/google_code_assist/chat.py @@ -246,3 +246,13 @@ class GoogleCodeAssistChat: if isinstance(e, GoogleCodeAssistError): return e return GoogleCodeAssistError(status_code=500, message=str(e)) + + +_shared_google_code_assist_chat: Optional[GoogleCodeAssistChat] = None + + +def get_google_code_assist_chat() -> GoogleCodeAssistChat: + global _shared_google_code_assist_chat + if _shared_google_code_assist_chat is None: + _shared_google_code_assist_chat = GoogleCodeAssistChat() + return _shared_google_code_assist_chat diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 754227e8873..511c0891e8c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -349,7 +349,6 @@ 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: @@ -374,7 +373,7 @@ def _get_gemini_url( params = [] if gemini_api_key and not gemini_oauth_token: params.append(f"key={gemini_api_key}") - if mode == "chat" and stream: + if mode == "chat" and stream is True: params.append("alt=sse") if params: diff --git a/litellm/main.py b/litellm/main.py index 26eaced03f6..d1971c53c0c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -101,7 +101,7 @@ 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.google_code_assist.chat import get_google_code_assist_chat from litellm.llms.gemini.fallback_handler import ( run_gemini_acompletion_with_code_assist_fallback, run_gemini_completion_with_code_assist_fallback, @@ -161,7 +161,7 @@ from litellm.utils import ( validate_openai_optional_params, ) -_google_code_assist_chat = GoogleCodeAssistChat() +_google_code_assist_chat = get_google_code_assist_chat() from ._logging import verbose_logger from .caching.caching import disable_cache, enable_cache, update_cache diff --git a/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py b/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py index c760e8d6ead..aaafdd5a5ba 100644 --- a/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py +++ b/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py @@ -2,11 +2,17 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch import httpx import json -from litellm.llms.google_code_assist.chat import GoogleCodeAssistChat +from litellm.llms.google_code_assist.chat import ( + GoogleCodeAssistChat, + get_google_code_assist_chat, +) from litellm.types.utils import ModelResponse class TestGoogleCodeAssist: + def test_get_google_code_assist_chat_returns_singleton(self): + assert get_google_code_assist_chat() is get_google_code_assist_chat() + @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): 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 d2d1948ebf1..472b24314e4 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 @@ -581,6 +581,19 @@ def test_get_gemini_url_stream_query_param_only_for_chat_mode(): assert "alt=sse" not in embedding_url +def test_get_gemini_url_requires_literal_true_for_streaming_endpoint_and_alt_sse(): + url, endpoint = _get_gemini_url( + mode="chat", + model="gemini-1.5-flash", + stream=1, # truthy non-bool should not be treated as streaming=True + gemini_api_key="test-key", + gemini_oauth_token=None, + ) + + assert endpoint == "generateContent" + assert "alt=sse" not in url + + @pytest.mark.parametrize( "model_cost_entry, vertex_region, expected_region", [