diff --git a/agentplatform/_genai/_evals_common.py b/agentplatform/_genai/_evals_common.py index 85c3ca6bd9..81541bbe95 100644 --- a/agentplatform/_genai/_evals_common.py +++ b/agentplatform/_genai/_evals_common.py @@ -154,6 +154,8 @@ def _get_runtime_instance( client = agentplatform.Client( project=api_client.project, location=api_client.location, + credentials=api_client._credentials, + http_options=api_client._http_options, ) _thread_local_data.runtime_instances[agent_name] = client.runtimes.get( name=agent_name diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py index 492d9cb429..909e24bc86 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -234,6 +234,41 @@ def test_get_api_client_with_none_location( mock_agentplatform_client.assert_not_called() +class TestGetRuntimeInstance: + _RUNTIME_NAME = "projects/test-project/locations/us-central1/reasoningEngines/123" + + def setup_method(self): + if hasattr(_evals_common._thread_local_data, "runtime_instances"): + del _evals_common._thread_local_data.runtime_instances + + def teardown_method(self): + if hasattr(_evals_common._thread_local_data, "runtime_instances"): + del _evals_common._thread_local_data.runtime_instances + + @mock.patch.object(_evals_common.agentplatform, "Client") + def test_get_runtime_instance_uses_client_credentials_and_http_options( + self, mock_agentplatform_client, mock_api_client_fixture + ): + mock_api_client_fixture._http_options = genai_types.HttpOptions( + base_url="https://custom-endpoint.example.com/" + ) + runtime = _evals_common._get_runtime_instance( + self._RUNTIME_NAME, mock_api_client_fixture + ) + mock_agentplatform_client.assert_called_once_with( + project=mock_api_client_fixture.project, + location=mock_api_client_fixture.location, + credentials=mock_api_client_fixture._credentials, + http_options=mock_api_client_fixture._http_options, + ) + mock_agentplatform_client.return_value.runtimes.get.assert_called_once_with( + name=self._RUNTIME_NAME + ) + assert runtime is ( + mock_agentplatform_client.return_value.runtimes.get.return_value + ) + + class TestNormalizeInferenceModelName: _FQ_PREFIX = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/" diff --git a/vertexai/_genai/_evals_common.py b/vertexai/_genai/_evals_common.py index ca45012b46..4118d35f09 100644 --- a/vertexai/_genai/_evals_common.py +++ b/vertexai/_genai/_evals_common.py @@ -124,6 +124,8 @@ def _get_agent_engine_instance( client = vertexai.Client( project=api_client.project, location=api_client.location, + credentials=api_client._credentials, + http_options=api_client._http_options, ) _thread_local_data.agent_engine_instances[agent_name] = ( client.agent_engines.get(name=agent_name)