From 768eb89ca1e6b1174d93f5880d9e3ecca02ad02e Mon Sep 17 00:00:00 2001 From: Mason Oh Date: Sun, 1 Feb 2026 15:31:06 +0900 Subject: [PATCH 1/5] feat: add custom client parameters --- src/google/adk/models/google_llm.py | 59 +++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 9015e4b4ee..8c23e42fd4 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -26,8 +26,10 @@ from typing import TYPE_CHECKING from typing import Union +from google.genai import Client from google.genai import types from google.genai.errors import ClientError +from pydantic import Field from typing_extensions import override from ..utils._client_labels_utils import get_client_labels @@ -40,8 +42,6 @@ from .llm_response import LlmResponse if TYPE_CHECKING: - from google.genai import Client - from .llm_request import LlmRequest logger = logging.getLogger('google_adk.' + __name__) @@ -86,6 +86,12 @@ class Gemini(BaseLlm): model: The name of the Gemini model. use_interactions_api: Whether to use the interactions API for model invocation. + custom_api_client: Custom client for standard API calls. If provided, ADK + tracking headers will NOT be automatically added to the constructor, but + will be merged into individual request headers. + custom_live_api_client: Custom client for Live API streaming. If provided, + the api_version must match the backend (v1beta1 for Vertex, v1alpha for + Gemini API). """ model: str = 'gemini-2.5-flash' @@ -127,6 +133,49 @@ class Gemini(BaseLlm): ``` """ + custom_api_client: Optional[Client] = Field( + default=None, exclude=True, frozen=True, repr=False + ) + """Custom API client for generate_content operations. + + Allows injecting a custom Google GenAI Client instance to override the + default api_client. Useful for testing, custom authentication, or using + different configurations. When set, this client is used for all + generate_content_async and interactions API calls. + + Sample: + ```python + from google.genai import Client + + custom_client = Client(api_key="custom_key") + agent = Agent( + model=Gemini(custom_api_client=custom_client) + ) + """ + + custom_live_api_client: Optional[Client] = Field( + default=None, exclude=True, frozen=True, repr=False + ) + """Custom client for Live API (bi-directional streaming) operations. + + Allows injecting a custom Google GenAI Client for ADK Live streaming. When + set, this client is used for all live.connect() calls. The client should be + configured with the appropriate API version for your backend (v1beta1 for + Vertex AI, v1alpha for Gemini API). + + Sample: + ``` + from google.genai import Client, types + + live_client = Client( + http_options=types.HttpOptions(api_version="v1beta1") + ) + agent = Agent( + model=Gemini(custom_live_api_client=live_client) + ) + ``` + """ + @classmethod @override def supported_models(cls) -> list[str]: @@ -298,7 +347,8 @@ def api_client(self) -> Client: Returns: The api client. """ - from google.genai import Client + if self.custom_api_client: + return self.custom_api_client return Client( http_options=types.HttpOptions( @@ -335,7 +385,8 @@ def _live_api_version(self) -> str: @cached_property def _live_api_client(self) -> Client: - from google.genai import Client + if self.custom_live_api_client: + return self.custom_live_api_client return Client( http_options=types.HttpOptions( From 4fa0afc43d727aa940924e883e84fd894e0e1bd0 Mon Sep 17 00:00:00 2001 From: Mason Oh Date: Sun, 1 Feb 2026 15:31:26 +0900 Subject: [PATCH 2/5] test: add test for custom clients --- tests/unittests/models/test_google_llm.py | 73 +++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py index 6837b66b95..394cee7314 100644 --- a/tests/unittests/models/test_google_llm.py +++ b/tests/unittests/models/test_google_llm.py @@ -2139,3 +2139,76 @@ async def __aexit__(self, *args): # Verify the final speech_config is still None assert config_arg.speech_config is None assert isinstance(connection, GeminiLlmConnection) + + +@pytest.mark.asyncio +async def test_custom_api_client_is_used_for_generate_content( + llm_request, generate_content_response +): + """Test that custom_api_client is used when provided.""" + from google.genai import Client + + # Create a mock custom client with proper spec + custom_client = mock.MagicMock(spec=Client) + + # Create a mock coroutine that returns the generate_content_response + async def mock_coro(): + return generate_content_response + + custom_client.aio.models.generate_content.return_value = mock_coro() + + # Create Gemini instance with custom_api_client + gemini_llm = Gemini(model="gemini-1.5-flash", custom_api_client=custom_client) + + # Execute generate_content_async + responses = [ + resp + async for resp in gemini_llm.generate_content_async( + llm_request, stream=False + ) + ] + + # Verify that the custom client was used + assert len(responses) == 1 + assert isinstance(responses[0], LlmResponse) + custom_client.aio.models.generate_content.assert_called_once() + + # Verify that api_client property returns the custom client + assert gemini_llm.api_client is custom_client + + +@pytest.mark.asyncio +async def test_custom_live_api_client_is_used_for_connect(llm_request): + """Test that custom_live_api_client is used when provided.""" + from google.genai import Client + + # Create a mock custom live client with proper spec + custom_live_client = mock.MagicMock(spec=Client) + mock_live_session = mock.AsyncMock() + + class MockLiveConnect: + + async def __aenter__(self): + return mock_live_session + + async def __aexit__(self, *args): + pass + + custom_live_client.aio.live.connect.return_value = MockLiveConnect() + + # Setup live connect config + llm_request.live_connect_config = types.LiveConnectConfig() + + # Create Gemini instance with custom_live_api_client + gemini_llm = Gemini( + model="gemini-1.5-flash", custom_live_api_client=custom_live_client + ) + + # Execute connect + async with gemini_llm.connect(llm_request) as connection: + # Verify that the custom live client was used + custom_live_client.aio.live.connect.assert_called_once() + assert isinstance(connection, GeminiLlmConnection) + + # Verify that _live_api_client property returns the custom client + assert gemini_llm._live_api_client is custom_live_client From e9f5f2e6b921b21cbb852c6ce96910ddf854ae5f Mon Sep 17 00:00:00 2001 From: Mason Oh Date: Sun, 1 Feb 2026 15:37:23 +0900 Subject: [PATCH 3/5] chore: simplify comments --- src/google/adk/models/google_llm.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 8c23e42fd4..b92c67815f 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -86,12 +86,8 @@ class Gemini(BaseLlm): model: The name of the Gemini model. use_interactions_api: Whether to use the interactions API for model invocation. - custom_api_client: Custom client for standard API calls. If provided, ADK - tracking headers will NOT be automatically added to the constructor, but - will be merged into individual request headers. - custom_live_api_client: Custom client for Live API streaming. If provided, - the api_version must match the backend (v1beta1 for Vertex, v1alpha for - Gemini API). + custom_api_client: Custom client for standard API calls. + custom_live_api_client: Custom client for Live API streaming. """ model: str = 'gemini-2.5-flash' From a871f179c240de54f736032a970b0d0cb35ae538 Mon Sep 17 00:00:00 2001 From: Mason Oh Date: Sun, 1 Feb 2026 15:37:33 +0900 Subject: [PATCH 4/5] chore: rm unused import --- src/google/adk/models/google_llm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index b92c67815f..23b773a6eb 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -19,7 +19,6 @@ import copy from functools import cached_property import logging -from typing import Any from typing import AsyncGenerator from typing import cast from typing import Optional From 35b34493baf4d0f594cad447f1d8b8cb01118f0a Mon Sep 17 00:00:00 2001 From: Mason Oh Date: Sun, 1 Feb 2026 16:04:07 +0900 Subject: [PATCH 5/5] chore: add python specifier --- src/google/adk/models/google_llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py index 23b773a6eb..6eedb0f3ca 100644 --- a/src/google/adk/models/google_llm.py +++ b/src/google/adk/models/google_llm.py @@ -159,7 +159,7 @@ class Gemini(BaseLlm): Vertex AI, v1alpha for Gemini API). Sample: - ``` + ```python from google.genai import Client, types live_client = Client(