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