Skip to content

Commit 22ca7ee

Browse files
wuliang229copybara-github
authored andcommitted
chore: Defer import of live, Client and _transformers in google.genai
This is targeting the Client import mostly, but also prevents future latency increase if the other modules in genai adds more 3p dependencies. This change will make ADK only import `live`, `Client` and `_transformers` just-in-time, therefore cutting down cold start latency. Co-authored-by: Liang Wu <wuliang@google.com> PiperOrigin-RevId: 831124329
1 parent a550509 commit 22ca7ee

File tree

6 files changed

+26
-12
lines changed

6 files changed

+26
-12
lines changed

src/google/adk/flows/llm_flows/instructions.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
from typing import AsyncGenerator
2020
from typing import TYPE_CHECKING
2121

22-
from google.genai import _transformers
2322
from typing_extensions import override
2423

2524
from ...agents.readonly_context import ReadonlyContext
@@ -85,6 +84,8 @@ async def run_async(
8584

8685
# Handle static_instruction - add via append_instructions
8786
if agent.static_instruction:
87+
from google.genai import _transformers
88+
8889
# Convert ContentUnion to Content using genai transformer
8990
static_content = _transformers.t_content(agent.static_instruction)
9091
llm_request.append_instructions(static_content)

src/google/adk/models/apigee_llm.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,22 @@
1818
from functools import cached_property
1919
import logging
2020
import os
21-
import re
2221
from typing import Optional
2322
from typing import TYPE_CHECKING
2423

2524
from google.adk import version as adk_version
26-
from google.genai import Client
2725
from google.genai import types
2826
from typing_extensions import override
2927

3028
from ..utils.env_utils import is_env_enabled
3129
from .google_llm import Gemini
3230

3331
if TYPE_CHECKING:
32+
from google.genai import Client
33+
3434
from .llm_request import LlmRequest
3535

36+
3637
logger = logging.getLogger('google_adk.' + __name__)
3738

3839
_APIGEE_PROXY_URL_ENV_VARIABLE_NAME = 'APIGEE_PROXY_URL'
@@ -137,6 +138,7 @@ def api_client(self) -> Client:
137138
Returns:
138139
The api client.
139140
"""
141+
from google.genai import Client
140142

141143
kwargs_for_http_options = {}
142144
if self._api_version:

src/google/adk/models/gemini_context_cache_manager.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
import logging
2222
import time
2323
from typing import Optional
24+
from typing import TYPE_CHECKING
2425

25-
from google.genai import Client
2626
from google.genai import types
2727

2828
from ..utils.feature_decorator import experimental
@@ -32,6 +32,9 @@
3232

3333
logger = logging.getLogger("google_adk." + __name__)
3434

35+
if TYPE_CHECKING:
36+
from google.genai import Client
37+
3538

3639
@experimental
3740
class GeminiContextCacheManager:

src/google/adk/models/gemini_llm_connection.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
from typing import AsyncGenerator
1919
from typing import Union
2020

21-
from google.genai import live
2221
from google.genai import types
2322

2423
from ..utils.context_utils import Aclosing
@@ -28,6 +27,10 @@
2827
logger = logging.getLogger('google_adk.' + __name__)
2928

3029
RealtimeInput = Union[types.Blob, types.ActivityStart, types.ActivityEnd]
30+
from typing import TYPE_CHECKING
31+
32+
if TYPE_CHECKING:
33+
from google.genai import live
3134

3235

3336
class GeminiLlmConnection(BaseLlmConnection):

src/google/adk/models/google_llm.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
from typing import TYPE_CHECKING
2828
from typing import Union
2929

30-
from google.genai import Client
3130
from google.genai import types
3231
from typing_extensions import override
3332

@@ -41,6 +40,8 @@
4140
from .llm_response import LlmResponse
4241

4342
if TYPE_CHECKING:
43+
from google.genai import Client
44+
4445
from .llm_request import LlmRequest
4546

4647
logger = logging.getLogger('google_adk.' + __name__)
@@ -200,6 +201,8 @@ def api_client(self) -> Client:
200201
Returns:
201202
The api client.
202203
"""
204+
from google.genai import Client
205+
203206
return Client(
204207
http_options=types.HttpOptions(
205208
headers=self._tracking_headers,
@@ -239,6 +242,8 @@ def _live_api_version(self) -> str:
239242

240243
@cached_property
241244
def _live_api_client(self) -> Client:
245+
from google.genai import Client
246+
242247
return Client(
243248
http_options=types.HttpOptions(
244249
headers=self._tracking_headers, api_version=self._live_api_version

tests/unittests/models/test_apigee_llm.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def llm_request():
4646

4747

4848
@pytest.mark.asyncio
49-
@mock.patch('google.adk.models.apigee_llm.Client')
49+
@mock.patch('google.genai.Client')
5050
async def test_generate_content_async_non_streaming(
5151
mock_client_constructor, llm_request
5252
):
@@ -96,7 +96,7 @@ async def test_generate_content_async_non_streaming(
9696

9797

9898
@pytest.mark.asyncio
99-
@mock.patch('google.adk.models.apigee_llm.Client')
99+
@mock.patch('google.genai.Client')
100100
async def test_generate_content_async_streaming(
101101
mock_client_constructor, llm_request
102102
):
@@ -167,7 +167,7 @@ async def mock_stream_generator():
167167

168168

169169
@pytest.mark.asyncio
170-
@mock.patch('google.adk.models.apigee_llm.Client')
170+
@mock.patch('google.genai.Client')
171171
async def test_generate_content_async_with_custom_headers(
172172
mock_client_constructor, llm_request
173173
):
@@ -207,7 +207,7 @@ async def test_generate_content_async_with_custom_headers(
207207

208208

209209
@pytest.mark.asyncio
210-
@mock.patch('google.adk.models.apigee_llm.Client')
210+
@mock.patch('google.genai.Client')
211211
async def test_vertex_model_path_parsing(mock_client_constructor):
212212
"""Tests that Vertex AI model paths are parsed correctly."""
213213
apigee_llm = ApigeeLlm(model=APIGEE_VERTEX_MODEL_ID, proxy_url=PROXY_URL)
@@ -249,7 +249,7 @@ async def test_vertex_model_path_parsing(mock_client_constructor):
249249

250250

251251
@pytest.mark.asyncio
252-
@mock.patch('google.adk.models.apigee_llm.Client')
252+
@mock.patch('google.genai.Client')
253253
async def test_proxy_url_from_env_variable(mock_client_constructor):
254254
"""Tests that proxy_url is read from environment variable."""
255255
with mock.patch.dict(
@@ -381,7 +381,7 @@ def test_vertex_model_missing_project_or_location_raises_error(
381381
),
382382
],
383383
)
384-
@mock.patch('google.adk.models.apigee_llm.Client')
384+
@mock.patch('google.genai.Client')
385385
async def test_model_string_parsing_and_client_initialization(
386386
mock_client_constructor,
387387
model_string,

0 commit comments

Comments
 (0)