Skip to content

Commit 34ea2ed

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: Fixes a bug introduced by the underlying google genai library
We need to set project and location parameters when we instantiate the Google GenAI library Client for VertexAi. This used to work before through environment variables but it doesn't seem to work anymore. So we are updating the `ApigeeLlm` wrapper to read from `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` environment variables and pass to the Client constructor. PiperOrigin-RevId: 829569362
1 parent 9e22cc4 commit 34ea2ed

File tree

2 files changed

+70
-1
lines changed

2 files changed

+70
-1
lines changed

src/google/adk/models/apigee_llm.py

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737

3838
_APIGEE_PROXY_URL_ENV_VARIABLE_NAME = 'APIGEE_PROXY_URL'
3939
_GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME = 'GOOGLE_GENAI_USE_VERTEXAI'
40+
_PROJECT_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_PROJECT'
41+
_LOCATION_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_LOCATION'
4042

4143

4244
class ApigeeLlm(Gemini):
@@ -90,6 +92,24 @@ def __init__(
9092
raise ValueError(f'Invalid model string: {model}')
9193

9294
self._isvertexai = _identify_vertexai(model)
95+
96+
# Set the project and location for Vertex AI.
97+
if self._isvertexai:
98+
self._project = os.environ.get(_PROJECT_ENV_VARIABLE_NAME)
99+
self._location = os.environ.get(_LOCATION_ENV_VARIABLE_NAME)
100+
101+
if not self._project:
102+
raise ValueError(
103+
f'The {_PROJECT_ENV_VARIABLE_NAME} environment variable must be'
104+
' set.'
105+
)
106+
107+
if not self._location:
108+
raise ValueError(
109+
f'The {_LOCATION_ENV_VARIABLE_NAME} environment variable must be'
110+
' set.'
111+
)
112+
93113
self._api_version = _identify_api_version(model)
94114
self._proxy_url = proxy_url or os.environ.get(
95115
_APIGEE_PROXY_URL_ENV_VARIABLE_NAME
@@ -128,9 +148,15 @@ def api_client(self) -> Client:
128148
**kwargs_for_http_options,
129149
)
130150

151+
kwargs_for_client = {}
152+
kwargs_for_client['vertexai'] = self._isvertexai
153+
if self._isvertexai:
154+
kwargs_for_client['project'] = self._project
155+
kwargs_for_client['location'] = self._location
156+
131157
return Client(
132-
vertexai=self._isvertexai,
133158
http_options=http_options,
159+
**kwargs_for_client,
134160
)
135161

136162
@override

tests/unittests/models/test_apigee_llm.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,42 @@ async def test_proxy_url_from_env_variable(mock_client_constructor):
286286
assert kwargs['http_options'].base_url == 'https://env.proxy.url'
287287

288288

289+
@pytest.mark.parametrize(
290+
('model_string', 'env_vars'),
291+
[
292+
(
293+
'apigee/vertex_ai/gemini-2.5-flash',
294+
{'GOOGLE_CLOUD_LOCATION': 'test-location'},
295+
),
296+
(
297+
'apigee/vertex_ai/gemini-2.5-flash',
298+
{'GOOGLE_CLOUD_PROJECT': 'test-project'},
299+
),
300+
(
301+
'apigee/gemini-2.5-flash',
302+
{
303+
'GOOGLE_GENAI_USE_VERTEXAI': 'true',
304+
'GOOGLE_CLOUD_LOCATION': 'test-location',
305+
},
306+
),
307+
(
308+
'apigee/gemini-2.5-flash',
309+
{
310+
'GOOGLE_GENAI_USE_VERTEXAI': 'true',
311+
'GOOGLE_CLOUD_PROJECT': 'test-project',
312+
},
313+
),
314+
],
315+
)
316+
def test_vertex_model_missing_project_or_location_raises_error(
317+
model_string, env_vars
318+
):
319+
"""Tests that ValueError is raised for Vertex models if project or location is missing."""
320+
with mock.patch.dict(os.environ, env_vars, clear=True):
321+
with pytest.raises(ValueError, match='environment variable must be set'):
322+
ApigeeLlm(model=model_string, proxy_url=PROXY_URL)
323+
324+
289325
@pytest.mark.asyncio
290326
@pytest.mark.parametrize(
291327
(
@@ -359,6 +395,10 @@ async def test_model_string_parsing_and_client_initialization(
359395
if use_vertexai_env is not None:
360396
env_vars['GOOGLE_GENAI_USE_VERTEXAI'] = use_vertexai_env
361397

398+
if expected_is_vertexai:
399+
env_vars['GOOGLE_CLOUD_PROJECT'] = 'test-project'
400+
env_vars['GOOGLE_CLOUD_LOCATION'] = 'test-location'
401+
362402
# The ApigeeLlm is initialized in the 'with' block to make sure that the mock
363403
# of the environment variable is active.
364404
with mock.patch.dict(os.environ, env_vars, clear=True):
@@ -382,6 +422,9 @@ async def test_model_string_parsing_and_client_initialization(
382422
mock_client_constructor.assert_called_once()
383423
_, kwargs = mock_client_constructor.call_args
384424
assert kwargs['vertexai'] == expected_is_vertexai
425+
if expected_is_vertexai:
426+
assert kwargs['project'] == 'test-project'
427+
assert kwargs['location'] == 'test-location'
385428
http_options = kwargs['http_options']
386429
assert http_options.api_version == expected_api_version
387430

0 commit comments

Comments
 (0)