fix: GenAI Client(evals) - Fix run_inference for local ADK agents and remote runtime N+1

agentplatform/_genai/_evals_common.py

@@ -3510,7 +3510,9 @@ def _create_runtime_session(
             raise RuntimeError(
                 "Failed to create session: runtime.api_client is None."
             ) from exc
-        operation = runtime.api_client.sessions.create(
+        from . import sessions
+
+        operation = sessions.Sessions(runtime.api_client._api_client).create(
             name=runtime.api_resource.name,
             user_id=user_id,
             config=types.CreateRuntimeSessionConfig(
@@ -3574,10 +3576,12 @@ def _execute_agent_run_with_retry(
             return {"error": "runtime.api_resource is None."}
         if runtime.api_client is None:
             return {"error": "runtime.api_client is None."}
+        from . import sessions
+
         session_name = f"{runtime.api_resource.name}/sessions/{session_id}"
         base_ts = datetime.datetime(2000, 1, 1, tzinfo=datetime.timezone.utc)
         for i, ag_event in enumerate(history_events):
-            runtime.api_client.sessions.events.append(
+            sessions.Sessions(runtime.api_client._api_client).events.append(
                 name=session_name,
                 author=ag_event.author or "user",
                 invocation_id="history",
@@ -3635,7 +3639,7 @@ async def _execute_local_agent_run_with_retry_async(
     # Lazy-import ADK dependencies to avoid top-level import failures when
     # google-adk is not installed.
     from google.adk.runners import Runner
-    from google.adk.sessions import InMemorySessionService
+    from google.adk.sessions.in_memory_session_service import InMemorySessionService
 
     # Multi-turn agent scraping with user simulation.
     if user_simulator_config or "conversation_plan" in row:

tests/unit/agentplatform/genai/replays/test_run_inference.py

@@ -14,19 +14,11 @@
 #
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
-import pytest
-
+from google.adk.agents import LlmAgent
 from tests.unit.agentplatform.genai.replays import pytest_helper
 from agentplatform._genai import types
 from google.genai import types as genai_types
 
-pytest.importorskip(
-    "google.adk", reason="google-adk not installed, skipping ADK agent tests"
-)
-from google.adk.agents import (  # noqa: E402
-    LlmAgent,
-)  # pylint: disable=g-import-not-at-top,g-bad-import-order
-
 
 def test_inference_with_eval_cases_multi_turn_agent_data(client):
     """Tests run_inference with multi-turn agent_data in eval_cases.
@@ -187,6 +179,7 @@ def test_inference_with_prompt_column_local_agent(client):
     assert response_val is not None
     assert isinstance(response_val, str)
     assert len(response_val) > 0
+    assert '"error"' not in response_val
 
 
 def test_inference_with_completed_and_incomplete_agent_data(client):
@@ -298,6 +291,7 @@ def test_inference_with_completed_and_incomplete_agent_data(client):
     assert row1_response is not None
     assert isinstance(row1_response, str)
     assert len(row1_response) > 0
+    assert '"error"' not in row1_response
 
 
 def test_inference_with_gemini_agent(client):

tests/unit/agentplatform/genai/test_evals.py

@@ -3655,12 +3655,14 @@ def test_run_inference_with_runtime_with_response_column_raises_error(
             "'intermediate_events' or 'response' columns"
         ) in str(excinfo.value)
 
+    @mock.patch("agentplatform._genai.sessions.Sessions")
     @mock.patch.object(_evals_utils, "EvalDatasetLoader")
     @mock.patch.object(_evals_common.agentplatform, "Client")
     def test_run_inference_with_runtime_falls_back_to_managed_sessions_api(
         self,
         mock_agentplatform_client,
         mock_eval_dataset_loader,
+        mock_sessions,
     ):
         """Tests that run_inference falls back to the managed Sessions API
         when the agent engine does not have create_session registered."""
@@ -3694,7 +3696,7 @@ def test_run_inference_with_runtime_falls_back_to_managed_sessions_api(
             "projects/test-project/locations/us-central1"
             "/reasoningEngines/123/sessions/managed-session-1"
         )
-        mock_runtime.api_client.sessions.create.return_value = mock_session_operation
+        mock_sessions.return_value.create.return_value = mock_session_operation
 
         stream_query_return_value = [
             {
@@ -3719,7 +3721,8 @@ def test_run_inference_with_runtime_falls_back_to_managed_sessions_api(
         )
 
         # Verify the managed Sessions API was called as fallback.
-        mock_runtime.api_client.sessions.create.assert_called_once_with(
+        mock_sessions.assert_called_once_with(mock_runtime.api_client._api_client)
+        mock_sessions.return_value.create.assert_called_once_with(
             name="projects/test-project/locations/us-central1/reasoningEngines/123",
             user_id="123",
             config=agentplatform_genai_types.CreateRuntimeSessionConfig(
@@ -3832,7 +3835,9 @@ def run_async_side_effect(*args, **kwargs):
             sys.modules,
             {
                 "google.adk": mock.MagicMock(),
-                "google.adk.sessions": mock_adk_sessions_module,
+                "google.adk.sessions.in_memory_session_service": (
+                    mock_adk_sessions_module
+                ),
                 "google.adk.runners": mock_adk_runners_module,
                 "google.adk.agents": mock.MagicMock(),
             },
@@ -11359,7 +11364,11 @@ def test_local_agent_run_retries_429_and_fails_fast_on_400(self, mock_sleep):
 
         with mock.patch.dict(
             sys.modules,
-            {"google.adk.runners": mock_runners, "google.adk.sessions": mock_sessions},
+            {
+                "google.adk.runners": mock_runners,
+                "google.adk.sessions": mock_sessions,
+                "google.adk.sessions.in_memory_session_service": mock_sessions,
+            },
         ):
             result = _evals_common._execute_local_agent_run_with_retry(
                 row=pd.Series({"prompt": "prompt"}),

vertexai/_genai/_evals_common.py

@@ -2416,7 +2416,7 @@ async def _execute_local_agent_run_with_retry_async(
     # Lazy-import ADK dependencies to avoid top-level import failures when
     # google-adk is not installed.
     from google.adk.runners import Runner
-    from google.adk.sessions import InMemorySessionService
+    from google.adk.sessions.in_memory_session_service import InMemorySessionService
 
     # Multi-turn agent scraping with user simulation.
     if user_simulator_config or "conversation_plan" in row: