fix: GenAI SDK client - switch to nested modules for agent engines

PiperOrigin-RevId: 797839264

Author: yeesian

tests/unit/vertexai/genai/test_agent_engines.py

@@ -926,16 +926,17 @@ def test_await_operation(self):
                     }
                 ),
             )
-            agent_engine = self.client.agent_engines._await_operation(
+            operation = _agent_engines_utils._await_operation(
                 operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME,
+                get_operation_fn=self.client.agent_engines._get_agent_operation,
             )
             request_mock.assert_called_with(
                 "get",
                 _TEST_AGENT_ENGINE_OPERATION_NAME,
                 {"_url": {"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME}},
                 None,
             )
-            assert isinstance(agent_engine, _genai_types.AgentEngineOperation)
+            assert isinstance(operation, _genai_types.AgentEngineOperation)
 
     def test_register_api_methods(self):
         agent = self.client.agent_engines._register_api_methods(
@@ -1145,7 +1146,7 @@ def test_list_agent_engine(self):
 
     @pytest.mark.usefixtures("caplog")
     @mock.patch.object(_agent_engines_utils, "_prepare")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
         mock_await_operation.return_value = _genai_types.AgentEngineOperation(
             response=_genai_types.ReasoningEngine(
@@ -1170,10 +1171,6 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
                     env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT,
                 ),
             )
-            mock_await_operation.assert_called_once_with(
-                operation_name=None,
-                poll_interval_seconds=10,
-            )
             request_mock.assert_called_with(
                 "post",
                 "reasoningEngines",
@@ -1199,7 +1196,7 @@ def test_create_agent_engine(self, mock_await_operation, mock_prepare, caplog):
             )
 
     @mock.patch.object(agent_engines.AgentEngines, "_create_config")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_create_agent_engine_lightweight(
         self,
         mock_await_operation,
@@ -1225,10 +1222,6 @@ def test_create_agent_engine_lightweight(
                     description=_TEST_AGENT_ENGINE_DESCRIPTION,
                 )
             )
-            mock_await_operation.assert_called_once_with(
-                operation_name=None,
-                poll_interval_seconds=1,
-            )
             request_mock.assert_called_with(
                 "post",
                 "reasoningEngines",
@@ -1240,7 +1233,7 @@ def test_create_agent_engine_lightweight(
             )
 
     @mock.patch.object(agent_engines.AgentEngines, "_create_config")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_create_agent_engine_with_env_vars_dict(
         self,
         mock_await_operation,
@@ -1317,7 +1310,7 @@ def test_create_agent_engine_with_env_vars_dict(
 
     @pytest.mark.usefixtures("caplog")
     @mock.patch.object(_agent_engines_utils, "_prepare")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_update_agent_engine_requirements(
         self, mock_await_operation, mock_prepare, caplog
     ):
@@ -1374,7 +1367,7 @@ def test_update_agent_engine_requirements(
             )
 
     @mock.patch.object(_agent_engines_utils, "_prepare")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_update_agent_engine_extra_packages(
         self, mock_await_operation, mock_prepare
     ):
@@ -1428,7 +1421,7 @@ def test_update_agent_engine_extra_packages(
             )
 
     @mock.patch.object(_agent_engines_utils, "_prepare")
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_update_agent_engine_env_vars(
         self, mock_await_operation, mock_prepare, caplog
     ):
@@ -1484,7 +1477,7 @@ def test_update_agent_engine_env_vars(
                 None,
             )
 
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_update_agent_engine_display_name(self, mock_await_operation):
         mock_await_operation.return_value = _genai_types.AgentEngineOperation(
             response=_genai_types.ReasoningEngine(
@@ -1513,7 +1506,7 @@ def test_update_agent_engine_display_name(self, mock_await_operation):
                 None,
             )
 
-    @mock.patch.object(agent_engines.AgentEngines, "_await_operation")
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_update_agent_engine_description(self, mock_await_operation):
         mock_await_operation.return_value = _genai_types.AgentEngineOperation(
             response=_genai_types.ReasoningEngine(

vertexai/_genai/_agent_engines_utils.py

@@ -23,6 +23,7 @@
 import os
 import sys
 import tarfile
+import time
 import types
 import typing
 from typing import (
@@ -350,6 +351,46 @@ class _RequirementsValidationResult(TypedDict):
     actions: _RequirementsValidationActions
 
 
+AgentEngineOperationUnion = Union[
+    genai_types.AgentEngineOperation,
+    genai_types.AgentEngineMemoryOperation,
+    genai_types.AgentEngineGenerateMemoriesOperation,
+]
+
+
+class GetOperationFunction(Protocol):
+    def __call__(self, *, operation_name: str, **kwargs) -> AgentEngineOperationUnion:
+        pass
+
+
+def _await_operation(
+    *,
+    operation_name: str,
+    get_operation_fn: GetOperationFunction,
+    poll_interval_seconds: int = 10,
+) -> Any:
+    """Waits for the operation for creating an agent engine to complete.
+
+    Args:
+        operation_name (str):
+            Required. The name of the operation for creating the Agent Engine.
+        poll_interval_seconds (int):
+            The number of seconds to wait between each poll.
+        get_operation_fn (Callable[[str], Any]):
+            Optional. The function to use for getting the operation. If not
+            provided, `self._get_agent_operation` will be used.
+
+    Returns:
+        The operation that has completed (i.e. `operation.done==True`).
+    """
+    operation = get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        time.sleep(poll_interval_seconds)
+        operation = get_operation_fn(operation_name=operation.name)
+
+    return operation
+
+
 def _compare_requirements(
     *,
     requirements: Mapping[str, str],
@@ -1286,18 +1327,6 @@ def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
     return _method
 
 
-AgentEngineOperationUnion = Union[
-    genai_types.AgentEngineOperation,
-    genai_types.AgentEngineMemoryOperation,
-    genai_types.AgentEngineGenerateMemoriesOperation,
-]
-
-
-class GetOperationFunction(Protocol):
-    def __call__(self, *, operation_name: str, **kwargs) -> AgentEngineOperationUnion:
-        ...
-
-
 def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
     """Wraps an Agent Engine method, creating a callable for `query` API.