From 34ea2edfd2fa4ae2dcdd2f1d55181c59e7faffd4 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 7 Nov 2025 14:02:07 -0800 Subject: [PATCH] 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 --- src/google/adk/models/apigee_llm.py | 28 ++++++++++++++- tests/unittests/models/test_apigee_llm.py | 43 +++++++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/src/google/adk/models/apigee_llm.py b/src/google/adk/models/apigee_llm.py index 0229f976..0d533a67 100644 --- a/src/google/adk/models/apigee_llm.py +++ b/src/google/adk/models/apigee_llm.py @@ -37,6 +37,8 @@ logger = logging.getLogger('google_adk.' + __name__) _APIGEE_PROXY_URL_ENV_VARIABLE_NAME = 'APIGEE_PROXY_URL' _GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME = 'GOOGLE_GENAI_USE_VERTEXAI' +_PROJECT_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_PROJECT' +_LOCATION_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_LOCATION' class ApigeeLlm(Gemini): @@ -90,6 +92,24 @@ class ApigeeLlm(Gemini): raise ValueError(f'Invalid model string: {model}') self._isvertexai = _identify_vertexai(model) + + # Set the project and location for Vertex AI. + if self._isvertexai: + self._project = os.environ.get(_PROJECT_ENV_VARIABLE_NAME) + self._location = os.environ.get(_LOCATION_ENV_VARIABLE_NAME) + + if not self._project: + raise ValueError( + f'The {_PROJECT_ENV_VARIABLE_NAME} environment variable must be' + ' set.' + ) + + if not self._location: + raise ValueError( + f'The {_LOCATION_ENV_VARIABLE_NAME} environment variable must be' + ' set.' + ) + self._api_version = _identify_api_version(model) self._proxy_url = proxy_url or os.environ.get( _APIGEE_PROXY_URL_ENV_VARIABLE_NAME @@ -128,9 +148,15 @@ class ApigeeLlm(Gemini): **kwargs_for_http_options, ) + kwargs_for_client = {} + kwargs_for_client['vertexai'] = self._isvertexai + if self._isvertexai: + kwargs_for_client['project'] = self._project + kwargs_for_client['location'] = self._location + return Client( - vertexai=self._isvertexai, http_options=http_options, + **kwargs_for_client, ) @override diff --git a/tests/unittests/models/test_apigee_llm.py b/tests/unittests/models/test_apigee_llm.py index eeafd5be..73005001 100644 --- a/tests/unittests/models/test_apigee_llm.py +++ b/tests/unittests/models/test_apigee_llm.py @@ -286,6 +286,42 @@ async def test_proxy_url_from_env_variable(mock_client_constructor): assert kwargs['http_options'].base_url == 'https://env.proxy.url' +@pytest.mark.parametrize( + ('model_string', 'env_vars'), + [ + ( + 'apigee/vertex_ai/gemini-2.5-flash', + {'GOOGLE_CLOUD_LOCATION': 'test-location'}, + ), + ( + 'apigee/vertex_ai/gemini-2.5-flash', + {'GOOGLE_CLOUD_PROJECT': 'test-project'}, + ), + ( + 'apigee/gemini-2.5-flash', + { + 'GOOGLE_GENAI_USE_VERTEXAI': 'true', + 'GOOGLE_CLOUD_LOCATION': 'test-location', + }, + ), + ( + 'apigee/gemini-2.5-flash', + { + 'GOOGLE_GENAI_USE_VERTEXAI': 'true', + 'GOOGLE_CLOUD_PROJECT': 'test-project', + }, + ), + ], +) +def test_vertex_model_missing_project_or_location_raises_error( + model_string, env_vars +): + """Tests that ValueError is raised for Vertex models if project or location is missing.""" + with mock.patch.dict(os.environ, env_vars, clear=True): + with pytest.raises(ValueError, match='environment variable must be set'): + ApigeeLlm(model=model_string, proxy_url=PROXY_URL) + + @pytest.mark.asyncio @pytest.mark.parametrize( ( @@ -359,6 +395,10 @@ async def test_model_string_parsing_and_client_initialization( if use_vertexai_env is not None: env_vars['GOOGLE_GENAI_USE_VERTEXAI'] = use_vertexai_env + if expected_is_vertexai: + env_vars['GOOGLE_CLOUD_PROJECT'] = 'test-project' + env_vars['GOOGLE_CLOUD_LOCATION'] = 'test-location' + # The ApigeeLlm is initialized in the 'with' block to make sure that the mock # of the environment variable is active. with mock.patch.dict(os.environ, env_vars, clear=True): @@ -382,6 +422,9 @@ async def test_model_string_parsing_and_client_initialization( mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args assert kwargs['vertexai'] == expected_is_vertexai + if expected_is_vertexai: + assert kwargs['project'] == 'test-project' + assert kwargs['location'] == 'test-location' http_options = kwargs['http_options'] assert http_options.api_version == expected_api_version