From 5246cb40c9b8fc5b22e54797c6951bf1209a88d1 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 23 May 2026 03:22:22 +0800 Subject: [PATCH 1/2] fix: preserve live thinking config --- src/google/adk/models/google_llm.py | 4 +++ tests/unittests/models/test_google_llm.py | 32 +++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index d5923ffd25..0646a26912 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -453,6 +453,10 @@ async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection: ' backend. Please use Vertex AI backend.' ) llm_request.live_connect_config.tools = llm_request.config.tools + if llm_request.config.thinking_config is not None: + llm_request.live_connect_config.thinking_config = ( + llm_request.config.thinking_config + ) logger.debug('Connecting to live with llm_request:%s', llm_request) logger.debug('Live connect config: %s', llm_request.live_connect_config) async with self._live_api_client.aio.live.connect( diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py index 7f7ce39895..87189efec4 100644 --- a/tests/unittests/models/test_google_llm.py +++ b/tests/unittests/models/test_google_llm.py @@ -852,6 +852,38 @@ async def __aexit__(self, *args): ) +@pytest.mark.asyncio +async def test_connect_copies_thinking_config_to_live_config( + gemini_llm, llm_request +): + """Test that live connections preserve thinking_config from generate config.""" + thinking_config = types.ThinkingConfig( + thinking_budget=10, + include_thoughts=True, + ) + llm_request.config.thinking_config = thinking_config + llm_request.live_connect_config = types.LiveConnectConfig() + + mock_live_session = mock.AsyncMock() + + with mock.patch.object(gemini_llm, "_live_api_client") as mock_live_client: + + class MockLiveConnect: + + async def __aenter__(self): + return mock_live_session + + async def __aexit__(self, *args): + pass + + mock_live_client.aio.live.connect.return_value = MockLiveConnect() + + async with gemini_llm.connect(llm_request): + mock_live_client.aio.live.connect.assert_called_once() + config_arg = mock_live_client.aio.live.connect.call_args.kwargs["config"] + assert config_arg.thinking_config == thinking_config + + @pytest.mark.parametrize( ( "api_backend, " From f8772a0d24eb92bd2ade0e4c67d58f276dea9ef7 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 28 May 2026 08:07:31 +0800 Subject: [PATCH 2/2] test: cover mTLS google llm base URLs --- tests/unittests/models/test_google_llm.py | 53 ++++++++++++++++++----- 1 file changed, 43 insertions(+), 10 deletions(-) diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py index 87189efec4..1bb29b0bb0 100644 --- a/tests/unittests/models/test_google_llm.py +++ b/tests/unittests/models/test_google_llm.py @@ -40,6 +40,17 @@ from google.genai.types import Part import pytest +_GENERATIVE_LANGUAGE_API_BASE_URL = ( + "https://generativelanguage.googleapis.com/v1alpha" +) +_GENERATIVE_LANGUAGE_API_ROOT = "https://generativelanguage.googleapis.com/" +_GENERATIVE_LANGUAGE_MTLS_API_BASE_URL = ( + "https://generativelanguage.mtls.googleapis.com/v1alpha" +) +_GENERATIVE_LANGUAGE_MTLS_API_ROOT = ( + "https://generativelanguage.mtls.googleapis.com/" +) + class MockAsyncIterator: """Mock for async iterator.""" @@ -252,17 +263,27 @@ def test_client_version_header_with_agent_engine(monkeypatch): ) -def test_api_client_uses_api_version_from_google_base_url(): +@pytest.mark.parametrize( + ("base_url", "normalized_base_url"), + ( + (_GENERATIVE_LANGUAGE_API_BASE_URL, _GENERATIVE_LANGUAGE_API_ROOT), + ( + _GENERATIVE_LANGUAGE_MTLS_API_BASE_URL, + _GENERATIVE_LANGUAGE_MTLS_API_ROOT, + ), + ), +) +def test_api_client_uses_api_version_from_google_base_url( + base_url, normalized_base_url +): model = Gemini( model="gemini-2.5-flash", - base_url="https://generativelanguage.googleapis.com/v1alpha", + base_url=base_url, ) client = model.api_client - assert client._api_client._http_options.base_url == ( - "https://generativelanguage.googleapis.com/" - ) + assert client._api_client._http_options.base_url == normalized_base_url assert client._api_client._http_options.api_version == "v1alpha" @@ -670,7 +691,7 @@ async def test_generate_content_async_patches_api_version( ): gemini_llm = Gemini( model="gemini-2.5-flash", - base_url="https://generativelanguage.googleapis.com/v1alpha", + base_url=_GENERATIVE_LANGUAGE_API_BASE_URL, ) llm_request.config.http_options = types.HttpOptions( headers={"custom-header": "custom-value"} @@ -718,7 +739,7 @@ def test_live_api_version_vertex_ai(gemini_llm): def test_live_api_version_uses_google_base_url_version(): gemini_llm = Gemini( model="gemini-2.5-flash", - base_url="https://generativelanguage.googleapis.com/v1alpha", + base_url=_GENERATIVE_LANGUAGE_API_BASE_URL, ) assert gemini_llm._live_api_version == "v1alpha" @@ -732,16 +753,28 @@ def test_live_api_version_gemini_api(gemini_llm): assert gemini_llm._live_api_version == "v1alpha" -def test_live_api_client_uses_api_version_from_google_base_url(): +@pytest.mark.parametrize( + ("base_url", "normalized_base_url"), + ( + (_GENERATIVE_LANGUAGE_API_BASE_URL, _GENERATIVE_LANGUAGE_API_ROOT), + ( + _GENERATIVE_LANGUAGE_MTLS_API_BASE_URL, + _GENERATIVE_LANGUAGE_MTLS_API_ROOT, + ), + ), +) +def test_live_api_client_uses_api_version_from_google_base_url( + base_url, normalized_base_url +): gemini_llm = Gemini( model="gemini-2.5-flash", - base_url="https://generativelanguage.googleapis.com/v1alpha", + base_url=base_url, ) client = gemini_llm._live_api_client http_options = client._api_client._http_options - assert http_options.base_url == "https://generativelanguage.googleapis.com/" + assert http_options.base_url == normalized_base_url assert http_options.api_version == "v1alpha"