chore: Dynamic set query method

shawn-yang-google · copybara-github · commit 653ba8815845 · 2024-11-21T14:55:02.000-08:00
PiperOrigin-RevId: 698923437
diff --git a/tests/unit/vertex_langchain/test_reasoning_engines.py b/tests/unit/vertex_langchain/test_reasoning_engines.py
@@ -146,6 +146,7 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_STANDARD_API_MODE = _reasoning_engines._STANDARD_API_MODE
 _TEST_MODE_KEY_IN_SCHEMA = _reasoning_engines._MODE_KEY_IN_SCHEMA
 _TEST_DEFAULT_METHOD_NAME = _reasoning_engines._DEFAULT_METHOD_NAME
+_TEST_DEFAULT_METHOD_DOCSTRING = _reasoning_engines._DEFAULT_METHOD_DOCSTRING
 _TEST_CUSTOM_METHOD_NAME = "custom_method"
 _TEST_QUERY_PROMPT = "Find the first fibonacci number greater than 999"
 _TEST_REASONING_ENGINE_GCS_URI = "{}/{}/{}".format(
@@ -413,13 +414,6 @@ def query_reasoning_engine_mock():
         yield query_reasoning_engine_mock
 
 
-@pytest.fixture(scope="module")
-def to_dict_mock():
-    with mock.patch.object(_utils, "to_dict") as to_dict_mock:
-        to_dict_mock.return_value = {}
-        yield to_dict_mock
-
-
 # Function scope is required for the pytest parameterized tests.
 @pytest.fixture(scope="function")
 def types_reasoning_engine_mock():
@@ -853,23 +847,49 @@ def test_delete_after_get_reasoning_engine(
             name=test_reasoning_engine.resource_name,
         )
 
+    def test_query_after_create_reasoning_engine(
+        self,
+        get_reasoning_engine_mock,
+        query_reasoning_engine_mock,
+        get_gca_resource_mock,
+    ):
+        test_reasoning_engine = reasoning_engines.ReasoningEngine.create(
+            self.test_app,
+            display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME,
+            requirements=_TEST_REASONING_ENGINE_REQUIREMENTS,
+            extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH],
+        )
+        get_reasoning_engine_mock.assert_called_with(
+            name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
+            retry=_TEST_RETRY,
+        )
+        with mock.patch.object(_utils, "to_dict") as to_dict_mock:
+            to_dict_mock.return_value = {}
+            test_reasoning_engine.query(query=_TEST_QUERY_PROMPT)
+            assert test_reasoning_engine.query.__doc__ == _TEST_DEFAULT_METHOD_DOCSTRING
+            query_reasoning_engine_mock.assert_called_with(
+                request=_TEST_REASONING_ENGINE_QUERY_REQUEST
+            )
+            to_dict_mock.assert_called_once()
+
     def test_query_reasoning_engine(
         self,
         get_reasoning_engine_mock,
         query_reasoning_engine_mock,
-        to_dict_mock,
         get_gca_resource_mock,
     ):
         test_reasoning_engine = reasoning_engines.ReasoningEngine(_TEST_RESOURCE_ID)
         get_reasoning_engine_mock.assert_called_with(
             name=_TEST_REASONING_ENGINE_RESOURCE_NAME,
             retry=_TEST_RETRY,
         )
-        test_reasoning_engine.query(query=_TEST_QUERY_PROMPT)
-        query_reasoning_engine_mock.assert_called_with(
-            request=_TEST_REASONING_ENGINE_QUERY_REQUEST
-        )
-        to_dict_mock.assert_called_once()
+        with mock.patch.object(_utils, "to_dict") as to_dict_mock:
+            to_dict_mock.return_value = {}
+            test_reasoning_engine.query(query=_TEST_QUERY_PROMPT)
+            query_reasoning_engine_mock.assert_called_with(
+                request=_TEST_REASONING_ENGINE_QUERY_REQUEST
+            )
+            to_dict_mock.assert_called_once()
 
     def test_operation_schemas(
         self,
diff --git a/vertexai/reasoning_engines/_reasoning_engines.py b/vertexai/reasoning_engines/_reasoning_engines.py
@@ -19,8 +19,9 @@
 import os
 import sys
 import tarfile
+import types
 import typing
-from typing import Any, Dict, List, Optional, Protocol, Sequence, Union
+from typing import Any, Callable, Dict, List, Optional, Protocol, Sequence, Union
 
 import proto
 
@@ -29,7 +30,7 @@
 from google.cloud.aiplatform import base
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils as aip_utils
-from google.cloud.aiplatform_v1beta1 import types
+from google.cloud.aiplatform_v1beta1 import types as aip_types
 from google.cloud.aiplatform_v1beta1.types import reasoning_engine_service
 from vertexai.reasoning_engines import _utils
 from google.protobuf import field_mask_pb2
@@ -44,6 +45,19 @@
 _STANDARD_API_MODE = ""
 _MODE_KEY_IN_SCHEMA = "api_mode"
 _DEFAULT_METHOD_NAME = "query"
+_DEFAULT_METHOD_DOCSTRING = """
+    Runs the Reasoning Engine to serve the user query.
+
+    This will be based on the `.query(...)` method of the python object that
+    was passed in when creating the Reasoning Engine.
+
+    Args:
+        **kwargs:
+            Optional. The arguments of the `.query(...)` method.
+
+    Returns:
+        dict[str, Any]: The response from serving the user query.
+"""
 
 
 @typing.runtime_checkable
@@ -73,7 +87,7 @@ def register_operations(self, **kwargs):
         """Register the user provided operations (modes and methods)."""
 
 
-class ReasoningEngine(base.VertexAiResourceNounWithFutureManager, Queryable):
+class ReasoningEngine(base.VertexAiResourceNounWithFutureManager):
     """Represents a Vertex AI Reasoning Engine resource."""
 
     client_class = aip_utils.ReasoningEngineClientWithOverride
@@ -98,6 +112,7 @@ def __init__(self, reasoning_engine_name: str):
             client_class=aip_utils.ReasoningEngineExecutionClientWithOverride,
         )
         self._gca_resource = self._get_gca_resource(resource_name=reasoning_engine_name)
+        _register_api_method(self)
         self._operation_schemas = None
 
     @property
@@ -233,7 +248,7 @@ def create(
             extra_packages=extra_packages,
         )
         # Update the package spec.
-        package_spec = types.ReasoningEngineSpec.PackageSpec(
+        package_spec = aip_types.ReasoningEngineSpec.PackageSpec(
             python_version=sys_version,
             pickle_object_gcs_uri="{}/{}/{}".format(
                 staging_bucket,
@@ -253,7 +268,7 @@ def create(
                 gcs_dir_name,
                 _REQUIREMENTS_FILE,
             )
-        reasoning_engine_spec = types.ReasoningEngineSpec(
+        reasoning_engine_spec = aip_types.ReasoningEngineSpec(
             package_spec=package_spec,
         )
         class_methods_spec = _generate_class_methods_spec_or_raise(
@@ -264,7 +279,7 @@ def create(
             parent=initializer.global_config.common_location_path(
                 project=sdk_resource.project, location=sdk_resource.location
             ),
-            reasoning_engine=types.ReasoningEngine(
+            reasoning_engine=aip_types.ReasoningEngine(
                 name=reasoning_engine_name,
                 display_name=display_name,
                 description=description,
@@ -289,6 +304,7 @@ def create(
             credentials=sdk_resource.credentials,
             location_override=sdk_resource.location,
         )
+        _register_api_method(sdk_resource)
         sdk_resource._operation_schemas = None
         return sdk_resource
 
@@ -431,30 +447,6 @@ def operation_schemas(self) -> Sequence[_utils.JsonDict]:
             self._operation_schemas = spec.get("class_methods", [])
         return self._operation_schemas
 
-    def query(self, **kwargs) -> _utils.JsonDict:
-        """Runs the Reasoning Engine to serve the user query.
-
-        This will be based on the `.query(...)` method of the python object that
-        was passed in when creating the Reasoning Engine.
-
-        Args:
-            **kwargs:
-                Optional. The arguments of the `.query(...)` method.
-
-        Returns:
-            dict[str, Any]: The response from serving the user query.
-        """
-        response = self.execution_api_client.query_reasoning_engine(
-            request=types.QueryReasoningEngineRequest(
-                name=self.resource_name,
-                input=kwargs,
-            ),
-        )
-        output = _utils.to_dict(response)
-        if "output" in output:
-            return output.get("output")
-        return output
-
 
 def _validate_sys_version_or_raise(sys_version: str) -> None:
     """Tries to validate the python system version."""
@@ -630,8 +622,8 @@ def _generate_update_request_or_raise(
     """Tries to generates the update request for the reasoning engine."""
     is_spec_update = False
     update_masks: List[str] = []
-    reasoning_engine_spec = types.ReasoningEngineSpec()
-    package_spec = types.ReasoningEngineSpec.PackageSpec()
+    reasoning_engine_spec = aip_types.ReasoningEngineSpec()
+    package_spec = aip_types.ReasoningEngineSpec.PackageSpec()
     if requirements is not None:
         is_spec_update = True
         update_masks.append("spec.package_spec.requirements_gcs_uri")
@@ -662,7 +654,7 @@ def _generate_update_request_or_raise(
         reasoning_engine_spec.class_methods.extend(class_methods_spec)
         update_masks.append("spec.class_methods")
 
-    reasoning_engine_message = types.ReasoningEngine(name=resource_name)
+    reasoning_engine_message = aip_types.ReasoningEngine(name=resource_name)
     if is_spec_update:
         reasoning_engine_spec.package_spec = package_spec
         reasoning_engine_message.spec = reasoning_engine_spec
@@ -684,6 +676,54 @@ def _generate_update_request_or_raise(
     )
 
 
+def _wrap_query_operation(method_name: str, doc: str) -> Callable[..., _utils.JsonDict]:
+    """Wraps a Reasoning Engine method, creating a callable for `query` API.
+
+    This function creates a callable object that executes the specified
+    Reasoning Engine method using the `query` API.  It handles the creation of
+    the API request and the processing of the API response.
+
+    Args:
+        method_name: The name of the Reasoning Engine method to call.
+        doc: Documentation string for the method.
+
+    Returns:
+        A callable object that executes the method on the Reasoning Engine via
+        the `query` API.
+    """
+
+    def _method(self, **kwargs) -> _utils.JsonDict:
+        response = self.execution_api_client.query_reasoning_engine(
+            request=aip_types.QueryReasoningEngineRequest(
+                name=self.resource_name,
+                input=kwargs,
+            ),
+        )
+        output = _utils.to_dict(response)
+        return output.get("output", output)
+
+    _method.__name__ = method_name
+    _method.__doc__ = doc
+
+    return _method
+
+
+def _register_api_method(obj: "ReasoningEngine"):
+    """Registers Reasoning Engine API methods based on operation schemas.
+
+    This function registers `query` method on the ReasoningEngine object
+    to handle API calls based on the specified API mode.
+
+    Args:
+        obj: The ReasoningEngine object to augment with API methods.
+    """
+    query_method = _wrap_query_operation(
+        method_name=_DEFAULT_METHOD_NAME, doc=_DEFAULT_METHOD_DOCSTRING
+    )
+    # Binds the method to the object.
+    setattr(obj, _DEFAULT_METHOD_NAME, types.MethodType(query_method, obj))
+
+
 def _get_registered_operations(reasoning_engine: Any) -> Dict[str, List[str]]:
     """Retrieves registered operations for a ReasoningEngine."""
     if isinstance(reasoning_engine, OperationRegistrable):
