feat: Add async support for reasoning engine.

PiperOrigin-RevId: 763862748

Author: yeesian

google/cloud/aiplatform/compat/services/__init__.py

@@ -179,6 +179,9 @@
 from google.cloud.aiplatform_v1.services.reasoning_engine_execution_service import (
     client as reasoning_engine_execution_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.reasoning_engine_execution_service import (
+    async_client as reasoning_engine_execution_async_client_v1,
+)
 from google.cloud.aiplatform_v1.services.schedule_service import (
     client as schedule_service_client_v1,
 )

google/cloud/aiplatform/utils/__init__.py

@@ -90,6 +90,7 @@
     prediction_service_async_client_v1,
     reasoning_engine_service_client_v1,
     reasoning_engine_execution_service_client_v1,
+    reasoning_engine_execution_async_client_v1,
     schedule_service_client_v1,
     tensorboard_service_client_v1,
     vizier_service_client_v1,
@@ -1007,6 +1008,17 @@ class AgentEngineExecutionClientWithOverride(ClientWithOverride):
     )
 
 
+class AgentEngineExecutionAsyncClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1
+    _version_map = (
+        (
+            compat.V1,
+            reasoning_engine_execution_async_client_v1.ReasoningEngineExecutionServiceAsyncClient,
+        ),
+    )
+
+
 class VertexRagDataClientWithOverride(ClientWithOverride):
     _is_temporary = True
     _default_version = compat.DEFAULT_VERSION

tests/unit/vertex_langchain/test_agent_engines.py

@@ -63,6 +63,20 @@ def clone(self):
         return self
 
 
+class AsyncQueryEngine:
+    """A sample Agent Engine that implements `async_query`."""
+
+    def set_up(self):
+        pass
+
+    async def async_query(self, unused_arbitrary_string_name: str):
+        """Runs the query asynchronously."""
+        return unused_arbitrary_string_name.upper()
+
+    def clone(self):
+        return self
+
+
 class AsyncStreamQueryEngine:
     """A sample Agent Engine that implements `async_stream_query`."""
 
@@ -104,10 +118,18 @@ def query(self, unused_arbitrary_string_name: str) -> str:
         """Runs the engine."""
         return unused_arbitrary_string_name.upper()
 
+    async def async_query(self, unused_arbitrary_string_name: str) -> str:
+        """Runs the query asynchronously."""
+        return unused_arbitrary_string_name.upper()
+
     # Add a custom method to test the custom method registration.
     def custom_method(self, x: str) -> str:
         return x.upper()
 
+    # Add a custom async method to test the custom async method registration.
+    async def custom_async_method(self, x: str):
+        return x.upper()
+
     def stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]:
         """Runs the stream engine."""
         for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE:
@@ -146,6 +168,10 @@ def register_operations(self) -> Dict[str, List[str]]:
                 _TEST_DEFAULT_METHOD_NAME,
                 _TEST_CUSTOM_METHOD_NAME,
             ],
+            _TEST_ASYNC_API_MODE: [
+                _TEST_DEFAULT_ASYNC_METHOD_NAME,
+                _TEST_CUSTOM_ASYNC_METHOD_NAME,
+            ],
             _TEST_STREAM_API_MODE: [
                 _TEST_DEFAULT_STREAM_METHOD_NAME,
                 _TEST_CUSTOM_STREAM_METHOD_NAME,
@@ -164,14 +190,22 @@ def query(self, unused_arbitrary_string_name: str) -> str:
         """Runs the engine."""
         return unused_arbitrary_string_name.upper()
 
+    async def async_query(self, unused_arbitrary_string_name: str) -> str:
+        """Runs the query asynchronously."""
+        return unused_arbitrary_string_name.upper()
+
     # Add a custom method to test the custom method registration
     def custom_method(self, x: str) -> str:
         return x.upper()
 
-    # Add a custom method that is not registered.ration
+    # Add a custom method that is not registered.
     def custom_method_2(self, x: str) -> str:
         return x.upper()
 
+    # Add a custom async method to test the custom async method registration.
+    async def custom_async_method(self, x: str):
+        return x.upper()
+
     def stream_query(self, unused_arbitrary_string_name: str) -> Iterable[Any]:
         """Runs the stream engine."""
         for chunk in _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE:
@@ -204,6 +238,10 @@ def register_operations(self) -> Dict[str, List[str]]:
                 _TEST_DEFAULT_METHOD_NAME,
                 _TEST_CUSTOM_METHOD_NAME,
             ],
+            _TEST_ASYNC_API_MODE: [
+                _TEST_DEFAULT_ASYNC_METHOD_NAME,
+                _TEST_CUSTOM_ASYNC_METHOD_NAME,
+            ],
             _TEST_STREAM_API_MODE: [
                 _TEST_DEFAULT_STREAM_METHOD_NAME,
                 _TEST_CUSTOM_STREAM_METHOD_NAME,
@@ -291,9 +329,11 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_REQUIREMENTS_FILE = _agent_engines._REQUIREMENTS_FILE
 _TEST_EXTRA_PACKAGES_FILE = _agent_engines._EXTRA_PACKAGES_FILE
 _TEST_STANDARD_API_MODE = _agent_engines._STANDARD_API_MODE
+_TEST_ASYNC_API_MODE = _agent_engines._ASYNC_API_MODE
 _TEST_STREAM_API_MODE = _agent_engines._STREAM_API_MODE
 _TEST_ASYNC_STREAM_API_MODE = _agent_engines._ASYNC_STREAM_API_MODE
 _TEST_DEFAULT_METHOD_NAME = _agent_engines._DEFAULT_METHOD_NAME
+_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines._DEFAULT_ASYNC_METHOD_NAME
 _TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines._DEFAULT_STREAM_METHOD_NAME
 _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = (
     _agent_engines._DEFAULT_ASYNC_STREAM_METHOD_NAME
@@ -304,6 +344,7 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_MODE_KEY_IN_SCHEMA = _agent_engines._MODE_KEY_IN_SCHEMA
 _TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines._METHOD_NAME_KEY_IN_SCHEMA
 _TEST_CUSTOM_METHOD_NAME = "custom_method"
+_TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method"
 _TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method"
 _TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME = "custom_async_stream_method"
 _TEST_CUSTOM_METHOD_DEFAULT_DOCSTRING = """
@@ -320,6 +361,20 @@ def register_operations(self) -> Dict[str, List[str]]:
     Returns:
         dict[str, Any]: The response from serving the user request.
 """
+_TEST_CUSTOM_ASYNC_METHOD_DEFAULT_DOCSTRING = """
+    Runs the Agent Engine to serve the user request.
+
+    This will be based on the `.custom_async_method(...)` of the python object that
+    was passed in when creating the Agent Engine. The method will invoke the
+    `async_query` API client of the python object.
+
+    Args:
+        **kwargs:
+            Optional. The arguments of the `.custom_async_method(...)` method.
+
+    Returns:
+        Coroutine[Any]: The response from serving the user request.
+"""
 _TEST_CUSTOM_STREAM_METHOD_DEFAULT_DOCSTRING = """
     Runs the Agent Engine to serve the user request.
 
@@ -429,6 +484,13 @@ def register_operations(self) -> Dict[str, List[str]]:
 ]
 _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = []
 _TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py"
+_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _utils.to_proto(
+    _utils.generate_schema(
+        AsyncQueryEngine().async_query,
+        schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME,
+    )
+)
+_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE
 _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _utils.to_proto(
     _utils.generate_schema(
         OperationRegistrableEngine().custom_method,
@@ -438,6 +500,15 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[
     _TEST_MODE_KEY_IN_SCHEMA
 ] = _TEST_STANDARD_API_MODE
+_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _utils.to_proto(
+    _utils.generate_schema(
+        OperationRegistrableEngine().custom_async_method,
+        schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
+    )
+)
+_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[
+    _TEST_MODE_KEY_IN_SCHEMA
+] = _TEST_ASYNC_API_MODE
 _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _utils.to_proto(
     _utils.generate_schema(
         StreamQueryEngine().stream_query,
@@ -475,6 +546,8 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_OPERATION_REGISTRABLE_SCHEMAS = [
     _TEST_AGENT_ENGINE_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA,
+    _TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA,
+    _TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA,
     _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA,
     _TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA,
@@ -499,6 +572,7 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[
     _TEST_MODE_KEY_IN_SCHEMA
 ] = _TEST_STANDARD_API_MODE
+_TEST_ASYNC_QUERY_SCHEMAS = [_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA]
 _TEST_STREAM_QUERY_SCHEMAS = [
     _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA,
 ]
@@ -758,6 +832,17 @@ def query() -> str:
         return "RESPONSE"
 
 
+class InvalidCapitalizeEngineWithoutAsyncQuerySelf:
+    """A sample Agent Engine with an invalid async_query method."""
+
+    def set_up(self):
+        pass
+
+    async def async_query() -> str:
+        """Runs the engine."""
+        return "RESPONSE"
+
+
 class InvalidCapitalizeEngineWithoutStreamQuerySelf:
     """A sample Agent Engine with an invalid query_stream_query method."""
 
@@ -1161,6 +1246,23 @@ def test_get_agent_framework(
                     ),
                 ),
             ),
+            (
+                "Update the async query engine",
+                {"agent_engine": AsyncQueryEngine()},
+                types.reasoning_engine_service.UpdateReasoningEngineRequest(
+                    reasoning_engine=_generate_agent_engine_with_class_methods_and_agent_framework(
+                        _TEST_ASYNC_QUERY_SCHEMAS,
+                        _agent_engines._DEFAULT_AGENT_FRAMEWORK,
+                    ),
+                    update_mask=field_mask_pb2.FieldMask(
+                        paths=[
+                            "spec.package_spec.pickle_object_gcs_uri",
+                            "spec.class_methods",
+                            "spec.agent_framework",
+                        ]
+                    ),
+                ),
+            ),
             (
                 "Update the stream query engine",
                 {"agent_engine": StreamQueryEngine()},
@@ -1534,6 +1636,20 @@ def test_query_agent_engine(
                         ),
                         _TEST_STANDARD_API_MODE,
                     ),
+                    (
+                        _utils.generate_schema(
+                            OperationRegistrableEngine().async_query,
+                            schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME,
+                        ),
+                        _TEST_ASYNC_API_MODE,
+                    ),
+                    (
+                        _utils.generate_schema(
+                            OperationRegistrableEngine().custom_async_method,
+                            schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
+                        ),
+                        _TEST_ASYNC_API_MODE,
+                    ),
                     (
                         _utils.generate_schema(
                             OperationRegistrableEngine().stream_query,
@@ -2320,8 +2436,8 @@ def test_create_agent_engine_no_query_method(
             TypeError,
             match=(
                 "agent_engine has none of the following callable methods: "
-                "`query`, `stream_query`, `async_stream_query` or "
-                "`register_operations`."
+                "`query`, `async_query`, `stream_query`, `async_stream_query` "
+                "or `register_operations`."
             ),
         ):
             agent_engines.create(
@@ -2344,8 +2460,8 @@ def test_create_agent_engine_noncallable_query_attribute(
             TypeError,
             match=(
                 "agent_engine has none of the following callable methods: "
-                "`query`, `stream_query`, `async_stream_query` or "
-                "`register_operations`."
+                "`query`, `async_query`, `stream_query`, `async_stream_query` "
+                "or `register_operations`."
             ),
         ):
             agent_engines.create(
@@ -2406,6 +2522,23 @@ def test_create_agent_engine_with_invalid_query_method(
                 requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
             )
 
+    def test_create_agent_engine_with_invalid_async_query_method(
+        self,
+        create_agent_engine_mock,
+        cloud_storage_create_bucket_mock,
+        tarfile_open_mock,
+        cloudpickle_dump_mock,
+        cloudpickle_load_mock,
+        importlib_metadata_version_mock,
+        get_agent_engine_mock,
+    ):
+        with pytest.raises(ValueError, match="Invalid async_query signature"):
+            agent_engines.create(
+                InvalidCapitalizeEngineWithoutAsyncQuerySelf(),
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
+            )
+
     def test_create_agent_engine_with_invalid_stream_query_method(
         self,
         create_agent_engine_mock,
@@ -2574,8 +2707,8 @@ def test_update_agent_engine_no_query_method(
             TypeError,
             match=(
                 "agent_engine has none of the following callable methods: "
-                "`query`, `stream_query`, `async_stream_query` or "
-                "`register_operations`."
+                "`query`, `async_query`, `stream_query`, `async_stream_query` "
+                "or `register_operations`."
             ),
         ):
             test_agent_engine = _generate_agent_engine_to_update()
@@ -2597,8 +2730,8 @@ def test_update_agent_engine_noncallable_query_attribute(
             TypeError,
             match=(
                 "agent_engine has none of the following callable methods: "
-                "`query`, `stream_query`, `async_stream_query` or "
-                "`register_operations`."
+                "`query`, `async_query`, `stream_query`, `async_stream_query` "
+                "or `register_operations`."
             ),
         ):
             test_agent_engine = _generate_agent_engine_to_update()
@@ -2737,7 +2870,7 @@ def test_update_class_methods_spec_with_registered_operation_not_found(self):
                     "register the API methods: "
                     "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
                     "Error: {Unsupported api mode: `UNKNOWN_API_MODE`, "
-                    "Supported modes are: ``, `stream` and `async_stream`.}"
+                    "Supported modes are: ``, `async`, `stream` and `async_stream`.}"
                 ),
             ),
         ],

vertexai/agent_engines/__init__.py

@@ -29,6 +29,7 @@
     ModuleAgent,
     OperationRegistrable,
     Queryable,
+    AsyncQueryable,
     StreamQueryable,
     AsyncStreamQueryable,
 )
@@ -313,6 +314,7 @@ def update(
     "Cloneable",
     "OperationRegistrable",
     "Queryable",
+    "AsyncQueryable",
     "StreamQueryable",
     "AsyncStreamQueryable",
     # Methods