# -*- coding: utf-8 -*- # Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import abc import inspect import io import logging import os import sys import tarfile import types import typing from typing import ( Any, Callable, Dict, Iterable, List, Optional, Protocol, Sequence, Union, ) import proto from google.api_core import exceptions from google.cloud import storage 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 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 _LOGGER = base.Logger(__name__) _SUPPORTED_PYTHON_VERSIONS = ("3.9", "3.10", "3.11", "3.12", "3.13") _DEFAULT_GCS_DIR_NAME = "reasoning_engine" _BLOB_FILENAME = "reasoning_engine.pkl" _REQUIREMENTS_FILE = "requirements.txt" _EXTRA_PACKAGES_FILE = "dependencies.tar.gz" _STANDARD_API_MODE = "" _STREAM_API_MODE = "stream" _MODE_KEY_IN_SCHEMA = "api_mode" _METHOD_NAME_KEY_IN_SCHEMA = "name" _DEFAULT_METHOD_NAME = "query" _DEFAULT_STREAM_METHOD_NAME = "stream_query" _DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]" _DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]" _DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ Runs the Reasoning Engine to serve the user request. This will be based on the `.{method_name}(...)` of the python object that was passed in when creating the Reasoning Engine. The method will invoke the `{default_method_name}` API client of the python object. Args: **kwargs: Optional. The arguments of the `.{method_name}(...)` method. Returns: {return_type}: The response from serving the user request. """ @typing.runtime_checkable class Queryable(Protocol): """Protocol for Reasoning Engine applications that can be queried.""" @abc.abstractmethod def query(self, **kwargs): """Runs the Reasoning Engine to serve the user query.""" @typing.runtime_checkable class StreamQueryable(Protocol): """Protocol for Reasoning Engine applications that can stream responses.""" @abc.abstractmethod def stream_query(self, **kwargs): """Stream responses to serve the user query.""" @typing.runtime_checkable class Cloneable(Protocol): """Protocol for Reasoning Engine applications that can be cloned.""" @abc.abstractmethod def clone(self): """Return a clone of the object.""" @typing.runtime_checkable class OperationRegistrable(Protocol): """Protocol for applications that has registered operations.""" @abc.abstractmethod def register_operations(self, **kwargs): """Register the user provided operations (modes and methods).""" class ReasoningEngine(base.VertexAiResourceNounWithFutureManager): """Represents a Vertex AI Reasoning Engine resource.""" client_class = aip_utils.ReasoningEngineClientWithOverride _resource_noun = "reasoning_engine" _getter_method = "get_reasoning_engine" _list_method = "list_reasoning_engines" _delete_method = "delete_reasoning_engine" _parse_resource_name_method = "parse_reasoning_engine_path" _format_resource_name_method = "reasoning_engine_path" def __init__(self, reasoning_engine_name: str): """Retrieves a Reasoning Engine resource. Args: reasoning_engine_name (str): Required. A fully-qualified resource name or ID such as "projects/123/locations/us-central1/reasoningEngines/456" or "456" when project and location are initialized or passed. """ super().__init__(resource_name=reasoning_engine_name) self.execution_api_client = initializer.global_config.create_client( client_class=aip_utils.ReasoningEngineExecutionClientWithOverride, ) self._gca_resource = self._get_gca_resource(resource_name=reasoning_engine_name) try: _register_api_methods_or_raise(self) except Exception as e: logging.warning("Failed to register API methods: {%s}", e) self._operation_schemas = None @property def resource_name(self) -> str: """Fully-qualified resource name.""" return self._gca_resource.name @classmethod def create( cls, reasoning_engine: Union[Queryable, OperationRegistrable], *, requirements: Optional[Union[str, Sequence[str]]] = None, reasoning_engine_name: Optional[str] = None, display_name: Optional[str] = None, description: Optional[str] = None, gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, sys_version: Optional[str] = None, extra_packages: Optional[Sequence[str]] = None, ) -> "ReasoningEngine": """Creates a new ReasoningEngine. The Reasoning Engine will be an instance of the `reasoning_engine` that was passed in, running remotely on Vertex AI. Sample ``src_dir`` contents (e.g. ``./user_src_dir``): .. code-block:: python user_src_dir/ |-- main.py |-- requirements.txt |-- user_code/ | |-- utils.py | |-- ... |-- ... To build a Reasoning Engine: .. code-block:: python remote_app = ReasoningEngine.create( local_app, requirements=[ # I.e. the PyPI dependencies listed in requirements.txt "google-cloud-aiplatform==1.25.0", "langchain==0.0.242", ... ], extra_packages=[ "./user_src_dir/main.py", # a single file "./user_src_dir/user_code", # a directory ... ], ) Args: reasoning_engine (ReasoningEngineInterface): Required. The Reasoning Engine to be created. requirements (Union[str, Sequence[str]]): Optional. The set of PyPI dependencies needed. It can either be the path to a single file (requirements.txt), or an ordered list of strings corresponding to each line of the requirements file. reasoning_engine_name (str): Optional. A fully-qualified resource name or ID such as "projects/123/locations/us-central1/reasoningEngines/456" or "456" when project and location are initialized or passed. If specifying the ID, it should be 4-63 characters. Valid characters are lowercase letters, numbers and hyphens ("-"), and it should start with a number or a lower-case letter. If not provided, Vertex AI will generate a value for this ID. display_name (str): Optional. The user-defined name of the Reasoning Engine. The name can be up to 128 characters long and can comprise any UTF-8 character. description (str): Optional. The description of the Reasoning Engine. gcs_dir_name (CreateReasoningEngineOptions): Optional. The GCS bucket directory under `staging_bucket` to use for staging the artifacts needed. sys_version (str): Optional. The Python system version used. Currently supports any of "3.9", "3.10", "3.11", "3.12", "3.13". If not specified, it defaults to the "{major}.{minor}" attributes of sys.version_info. extra_packages (Sequence[str]): Optional. The set of extra user-provided packages (if any). Returns: ReasoningEngine: The Reasoning Engine that was created. Raises: ValueError: If `sys.version` is not supported by ReasoningEngine. ValueError: If the `project` was not set using `vertexai.init`. ValueError: If the `location` was not set using `vertexai.init`. ValueError: If the `staging_bucket` was not set using vertexai.init. ValueError: If the `staging_bucket` does not start with "gs://". FileNotFoundError: If `extra_packages` includes a file or directory that does not exist. IOError: If requirements is a string that corresponds to a nonexistent file. """ if not sys_version: sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" _validate_sys_version_or_raise(sys_version) reasoning_engine = _validate_reasoning_engine_or_raise(reasoning_engine) requirements = _validate_requirements_or_raise(requirements) extra_packages = _validate_extra_packages_or_raise(extra_packages) if reasoning_engine_name: _LOGGER.warning( "ReasoningEngine does not support user-defined resource IDs at " f"the moment. Therefore {reasoning_engine_name=} would be " "ignored and a random ID will be generated instead." ) sdk_resource = cls.__new__(cls) base.VertexAiResourceNounWithFutureManager.__init__( sdk_resource, resource_name=reasoning_engine_name, ) staging_bucket = initializer.global_config.staging_bucket _validate_staging_bucket_or_raise(staging_bucket) # Prepares the Reasoning Engine for creation in Vertex AI. # This involves packaging and uploading the artifacts for # reasoning_engine, requirements and extra_packages to # `staging_bucket/gcs_dir_name`. _prepare( reasoning_engine=reasoning_engine, requirements=requirements, project=sdk_resource.project, location=sdk_resource.location, staging_bucket=staging_bucket, gcs_dir_name=gcs_dir_name, extra_packages=extra_packages, ) # Update the package spec. package_spec = aip_types.ReasoningEngineSpec.PackageSpec( python_version=sys_version, pickle_object_gcs_uri="{}/{}/{}".format( staging_bucket, gcs_dir_name, _BLOB_FILENAME, ), ) if extra_packages: package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( staging_bucket, gcs_dir_name, _EXTRA_PACKAGES_FILE, ) if requirements: package_spec.requirements_gcs_uri = "{}/{}/{}".format( staging_bucket, gcs_dir_name, _REQUIREMENTS_FILE, ) reasoning_engine_spec = aip_types.ReasoningEngineSpec( package_spec=package_spec, ) class_methods_spec = _generate_class_methods_spec_or_raise( reasoning_engine, _get_registered_operations(reasoning_engine) ) reasoning_engine_spec.class_methods.extend(class_methods_spec) operation_future = sdk_resource.api_client.create_reasoning_engine( parent=initializer.global_config.common_location_path( project=sdk_resource.project, location=sdk_resource.location ), reasoning_engine=aip_types.ReasoningEngine( name=reasoning_engine_name, display_name=display_name, description=description, spec=reasoning_engine_spec, ), ) _LOGGER.log_create_with_lro(cls, operation_future) created_resource = operation_future.result() _LOGGER.log_create_complete( cls, created_resource, cls._resource_noun, module_name="vertexai.preview.reasoning_engines", ) # We use `._get_gca_resource(...)` instead of `created_resource` to # fully instantiate the attributes of the reasoning engine. sdk_resource._gca_resource = sdk_resource._get_gca_resource( resource_name=created_resource.name ) sdk_resource.execution_api_client = initializer.global_config.create_client( client_class=aip_utils.ReasoningEngineExecutionClientWithOverride, credentials=sdk_resource.credentials, location_override=sdk_resource.location, ) try: _register_api_methods_or_raise(sdk_resource) except Exception as e: logging.warning("Failed to register API methods: {%s}", e) sdk_resource._operation_schemas = None return sdk_resource def update( self, *, reasoning_engine: Optional[Union[Queryable, OperationRegistrable]] = None, requirements: Optional[Union[str, Sequence[str]]] = None, display_name: Optional[str] = None, description: Optional[str] = None, gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, sys_version: Optional[str] = None, extra_packages: Optional[Sequence[str]] = None, ) -> "ReasoningEngine": """Updates an existing ReasoningEngine. This method updates the configuration of an existing ReasoningEngine running remotely, which is identified by its resource name. Unlike the `create` function which requires a `reasoning_engine` object, all arguments in this method are optional. This method allows you to modify individual aspects of the configuration by providing any of the optional arguments. Note that you must provide at least one argument (except `sys_version`). Args: reasoning_engine (ReasoningEngineInterface): Optional. The Reasoning Engine to be replaced. If it is not specified, the existing Reasoning Engine will be used. requirements (Union[str, Sequence[str]]): Optional. The set of PyPI dependencies needed. It can either be the path to a single file (requirements.txt), or an ordered list of strings corresponding to each line of the requirements file. If it is not specified, the existing requirements will be used. If it is set to an empty string or list, the existing requirements will be removed. display_name (str): Optional. The user-defined name of the Reasoning Engine. The name can be up to 128 characters long and can comprise any UTF-8 character. description (str): Optional. The description of the Reasoning Engine. gcs_dir_name (CreateReasoningEngineOptions): Optional. The GCS bucket directory under `staging_bucket` to use for staging the artifacts needed. sys_version (str): Optional. The Python system version used. Currently updating sys version is not supported. extra_packages (Sequence[str]): Optional. The set of extra user-provided packages (if any). If it is not specified, the existing extra packages will be used. If it is set to an empty list, the existing extra packages will be removed. Returns: ReasoningEngine: The Reasoning Engine that was updated. Raises: ValueError: If `sys.version` is updated. ValueError: If the `staging_bucket` was not set using vertexai.init. ValueError: If the `staging_bucket` does not start with "gs://". FileNotFoundError: If `extra_packages` includes a file or directory that does not exist. ValueError: if none of `display_name`, `description`, `requirements`, `extra_packages`, or `reasoning_engine` were specified. IOError: If requirements is a string that corresponds to a nonexistent file. """ staging_bucket = initializer.global_config.staging_bucket _validate_staging_bucket_or_raise(staging_bucket) historical_operation_schemas = self.operation_schemas() # Validate the arguments. if not any( [ reasoning_engine, requirements, extra_packages, display_name, description, ] ): raise ValueError( "At least one of `reasoning_engine`, `requirements`, " "`extra_packages`, `display_name`, or `description` must be " "specified." ) if sys_version: _LOGGER.warning("Updated sys_version is not supported.") if requirements is not None: requirements = _validate_requirements_or_raise(requirements) if extra_packages is not None: extra_packages = _validate_extra_packages_or_raise(extra_packages) if reasoning_engine is not None: reasoning_engine = _validate_reasoning_engine_or_raise(reasoning_engine) # Prepares the Reasoning Engine for update in Vertex AI. # This involves packaging and uploading the artifacts for # reasoning_engine, requirements and extra_packages to # `staging_bucket/gcs_dir_name`. _prepare( reasoning_engine=reasoning_engine, requirements=requirements, project=self.project, location=self.location, staging_bucket=staging_bucket, gcs_dir_name=gcs_dir_name, extra_packages=extra_packages, ) update_request = _generate_update_request_or_raise( resource_name=self.resource_name, staging_bucket=staging_bucket, gcs_dir_name=gcs_dir_name, reasoning_engine=reasoning_engine, requirements=requirements, extra_packages=extra_packages, display_name=display_name, description=description, ) operation_future = self.api_client.update_reasoning_engine( request=update_request ) _LOGGER.info( f"Update ReasoningEngine backing LRO: {operation_future.operation.name}" ) created_resource = operation_future.result() _LOGGER.info(f"ReasoningEngine updated. Resource name: {created_resource.name}") self._operation_schemas = None self.execution_api_client = initializer.global_config.create_client( client_class=aip_utils.ReasoningEngineExecutionClientWithOverride, ) # We use `._get_gca_resource(...)` instead of `created_resource` to # fully instantiate the attributes of the reasoning engine. self._gca_resource = self._get_gca_resource(resource_name=self.resource_name) if ( reasoning_engine is None or historical_operation_schemas == self.operation_schemas() ): # As the API/operations of the reasoning engine are unchanged, we # can return it here. return self # If the reasoning engine has changed and the historical operation # schemas are different from the current operation schemas, we need to # unregister the historical operation schemas and register the current # operation schemas. _unregister_api_methods(self, historical_operation_schemas) try: _register_api_methods_or_raise(self) except Exception as e: logging.warning("Failed to register API methods: {%s}", e) return self def operation_schemas(self) -> Sequence[_utils.JsonDict]: """Returns the (Open)API schemas for the Reasoning Engine.""" spec = _utils.to_dict(self._gca_resource.spec) if not hasattr(self, "_operation_schemas") or self._operation_schemas is None: self._operation_schemas = spec.get("classMethods", []) return self._operation_schemas def _validate_sys_version_or_raise(sys_version: str) -> None: """Tries to validate the python system version.""" if sys_version not in _SUPPORTED_PYTHON_VERSIONS: raise ValueError( f"Unsupported python version: {sys_version}. ReasoningEngine " f"only supports {_SUPPORTED_PYTHON_VERSIONS} at the moment." ) if sys_version != f"{sys.version_info.major}.{sys.version_info.minor}": _LOGGER.warning( f"{sys_version=} is inconsistent with {sys.version_info=}. " "This might result in issues with deployment, and should only " "be used as a workaround for advanced cases." ) def _validate_staging_bucket_or_raise(staging_bucket: str) -> str: """Tries to validate the staging bucket.""" if not staging_bucket: raise ValueError("Please provide a `staging_bucket` in `vertexai.init(...)`") if not staging_bucket.startswith("gs://"): raise ValueError(f"{staging_bucket=} must start with `gs://`") def _validate_reasoning_engine_or_raise( reasoning_engine: Union[Queryable, OperationRegistrable, StreamQueryable], ) -> Union[Queryable, OperationRegistrable, StreamQueryable]: """Tries to validate the reasoning engine. The reasoning engine must have one of the following: * a callable method named `query` * a callable method named `stream_query` * a callable method named `register_operations` Args: reasoning_engine: The reasoning engine to be validated. Returns: The validated reasoning engine. Raises: TypeError: If the reasoning engine has no callable method named `query`, `stream_query` or `register_operations`. ValueError: If the reasoning engine has an invalid `query`, `stream_query` or `register_operations` signature. """ is_queryable = isinstance(reasoning_engine, Queryable) and callable( reasoning_engine.query ) is_stream_queryable = isinstance(reasoning_engine, StreamQueryable) and callable( reasoning_engine.stream_query ) is_operation_registrable = isinstance( reasoning_engine, OperationRegistrable ) and callable(reasoning_engine.register_operations) if not (is_queryable or is_stream_queryable or is_operation_registrable): raise TypeError( "reasoning_engine has neither a callable method named `query`" " nor a callable method named `register_operations`." ) if is_queryable: try: inspect.signature(getattr(reasoning_engine, "query")) except ValueError as err: raise ValueError( "Invalid query signature. This might be due to a missing " "`self` argument in the reasoning_engine.query method." ) from err if is_stream_queryable: try: inspect.signature(getattr(reasoning_engine, "stream_query")) except ValueError as err: raise ValueError( "Invalid stream_query signature. This might be due to a missing" " `self` argument in the reasoning_engine.stream_query method." ) from err if is_operation_registrable: try: inspect.signature(getattr(reasoning_engine, "register_operations")) except ValueError as err: raise ValueError( "Invalid register_operations signature. This might be due to a " "missing `self` argument in the " "reasoning_engine.register_operations method." ) from err if isinstance(reasoning_engine, Cloneable): # Avoid undeployable ReasoningChain states. reasoning_engine = reasoning_engine.clone() return reasoning_engine def _validate_requirements_or_raise(requirements: Sequence[str]) -> Sequence[str]: """Tries to validate the requirements.""" if isinstance(requirements, str): try: _LOGGER.info(f"Reading requirements from {requirements=}") with open(requirements) as f: requirements = f.read().splitlines() _LOGGER.info(f"Read the following lines: {requirements}") except IOError as err: raise IOError(f"Failed to read requirements from {requirements=}") from err return requirements or [] def _validate_extra_packages_or_raise(extra_packages: Sequence[str]) -> Sequence[str]: """Tries to validates the extra packages.""" extra_packages = extra_packages or [] for extra_package in extra_packages: if not os.path.exists(extra_package): raise FileNotFoundError( f"Extra package specified but not found: {extra_package=}" ) return extra_packages def _get_gcs_bucket(project: str, location: str, staging_bucket: str) -> storage.Bucket: """Gets or creates the GCS bucket.""" storage = _utils._import_cloud_storage_or_raise() storage_client = storage.Client(project=project) staging_bucket = staging_bucket.replace("gs://", "") try: gcs_bucket = storage_client.get_bucket(staging_bucket) _LOGGER.info(f"Using bucket {staging_bucket}") except exceptions.NotFound: new_bucket = storage_client.bucket(staging_bucket) gcs_bucket = storage_client.create_bucket(new_bucket, location=location) _LOGGER.info(f"Creating bucket {staging_bucket} in {location=}") return gcs_bucket def _upload_reasoning_engine( reasoning_engine: Union[Queryable, OperationRegistrable], gcs_bucket: storage.Bucket, gcs_dir_name: str, ) -> None: """Uploads the reasoning engine to GCS.""" cloudpickle = _utils._import_cloudpickle_or_raise() blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") with blob.open("wb") as f: cloudpickle.dump(reasoning_engine, f) dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" _LOGGER.info(f"Writing to {dir_name}/{_BLOB_FILENAME}") def _upload_requirements( requirements: Sequence[str], gcs_bucket: storage.Bucket, gcs_dir_name: str, ) -> None: """Uploads the requirements file to GCS.""" blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}") blob.upload_from_string("\n".join(requirements)) dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" _LOGGER.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}") def _upload_extra_packages( extra_packages: Sequence[str], gcs_bucket: storage.Bucket, gcs_dir_name: str, ) -> None: """Uploads extra packages to GCS.""" _LOGGER.info("Creating in-memory tarfile of extra_packages") tar_fileobj = io.BytesIO() with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar: for file in extra_packages: tar.add(file) tar_fileobj.seek(0) blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}") blob.upload_from_string(tar_fileobj.read()) dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" _LOGGER.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}") def _prepare( reasoning_engine: Optional[Union[Queryable, OperationRegistrable]], requirements: Optional[Sequence[str]], extra_packages: Optional[Sequence[str]], project: str, location: str, staging_bucket: str, gcs_dir_name: str, ) -> None: """Prepares the reasoning engine for creation or updates in Vertex AI. This involves packaging and uploading artifacts to Cloud Storage. Note that 1. This does not actually update the Reasoning Engine in Vertex AI. 2. This will only generate and upload a pickled object if specified. 3. This will only generate and upload the dependencies.tar.gz file if extra_packages is non-empty. Args: reasoning_engine: The reasoning engine to be prepared. requirements (Sequence[str]): The set of PyPI dependencies needed. extra_packages (Sequence[str]): The set of extra user-provided packages. project (str): The project for the staging bucket. location (str): The location for the staging bucket. staging_bucket (str): The staging bucket name in the form "gs://...". gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to use for staging the artifacts needed. """ gcs_bucket = _get_gcs_bucket(project, location, staging_bucket) if reasoning_engine is not None: _upload_reasoning_engine(reasoning_engine, gcs_bucket, gcs_dir_name) if requirements is not None: _upload_requirements(requirements, gcs_bucket, gcs_dir_name) if extra_packages is not None: _upload_extra_packages(extra_packages, gcs_bucket, gcs_dir_name) def _generate_update_request_or_raise( resource_name: str, staging_bucket: str, gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, reasoning_engine: Optional[Union[Queryable, OperationRegistrable]] = None, requirements: Optional[Union[str, Sequence[str]]] = None, extra_packages: Optional[Sequence[str]] = None, display_name: Optional[str] = None, description: Optional[str] = None, ) -> reasoning_engine_service.UpdateReasoningEngineRequest: """Tries to generates the update request for the reasoning engine.""" is_spec_update = False update_masks: List[str] = [] 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") package_spec.requirements_gcs_uri = "{}/{}/{}".format( staging_bucket, gcs_dir_name, _REQUIREMENTS_FILE, ) if extra_packages is not None: is_spec_update = True update_masks.append("spec.package_spec.dependency_files_gcs_uri") package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( staging_bucket, gcs_dir_name, _EXTRA_PACKAGES_FILE, ) if reasoning_engine is not None: is_spec_update = True update_masks.append("spec.package_spec.pickle_object_gcs_uri") package_spec.pickle_object_gcs_uri = "{}/{}/{}".format( staging_bucket, gcs_dir_name, _BLOB_FILENAME, ) class_methods_spec = _generate_class_methods_spec_or_raise( reasoning_engine, _get_registered_operations(reasoning_engine) ) reasoning_engine_spec.class_methods.extend(class_methods_spec) update_masks.append("spec.class_methods") 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 if display_name: reasoning_engine_message.display_name = display_name update_masks.append("display_name") if description: reasoning_engine_message.description = description update_masks.append("description") if not update_masks: raise ValueError( "At least one of `reasoning_engine`, `requirements`, " "`extra_packages`, `display_name`, or `description` must be " "specified." ) return reasoning_engine_service.UpdateReasoningEngineRequest( reasoning_engine=reasoning_engine_message, update_mask=field_mask_pb2.FieldMask(paths=update_masks), ) 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, class_method=method_name, ), ) output = _utils.to_dict(response) return output.get("output", output) _method.__name__ = method_name _method.__doc__ = doc return _method def _wrap_stream_query_operation( method_name: str, doc: str ) -> Callable[..., Iterable[Any]]: """Wraps a Reasoning Engine method, creating a callable for `stream_query` API. This function creates a callable object that executes the specified Reasoning Engine method using the `stream_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 `stream_query` API. """ def _method(self, **kwargs) -> Iterable[Any]: response = self.execution_api_client.stream_query_reasoning_engine( request=aip_types.StreamQueryReasoningEngineRequest( name=self.resource_name, input=kwargs, class_method=method_name, ), ) for chunk in response: for parsed_json in _utils.yield_parsed_json(chunk): if parsed_json is not None: yield parsed_json _method.__name__ = method_name _method.__doc__ = doc return _method def _unregister_api_methods( obj: "ReasoningEngine", operation_schemas: Sequence[_utils.JsonDict] ): """Unregisters Reasoning Engine API methods based on operation schemas. This function iterates through operation schemas provided by the ReasoningEngine object. Each schema defines an API mode and method name. It dynamically unregisters methods on the ReasoningEngine object. This should only be used when updating the object. Args: obj: The ReasoningEngine object to augment with API methods. operation_schemas: The operation schemas to use for method unregistration. """ for operation_schema in operation_schemas: if "name" in operation_schema: method_name = operation_schema.get("name") if hasattr(obj, method_name): delattr(obj, method_name) def _register_api_methods_or_raise(obj: "ReasoningEngine"): """Registers Reasoning Engine API methods based on operation schemas. This function iterates through operation schemas provided by the ReasoningEngine object. Each schema defines an API mode and method name. It dynamically creates and registers methods on the ReasoningEngine object to handle API calls based on the specified API mode. Currently, only standard API mode `` is supported. Args: obj: The ReasoningEngine object to augment with API methods. Raises: ValueError: If the API mode is not supported or if the operation schema missing required `api_mode` or `name` fields. """ for operation_schema in obj.operation_schemas(): if _MODE_KEY_IN_SCHEMA not in operation_schema: raise ValueError( f"Operation schema {operation_schema} does not" " contain an `api_mode` field." ) api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA) if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema: raise ValueError( f"Operation schema {operation_schema} does not" " contain a `name` field." ) method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) method_description = operation_schema.get("description") if api_mode == _STANDARD_API_MODE: method_description = ( method_description or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( method_name=method_name, default_method_name=_DEFAULT_METHOD_NAME, return_type=_DEFAULT_METHOD_RETURN_TYPE, ) ) method = _wrap_query_operation( method_name=method_name, doc=method_description, ) elif api_mode == _STREAM_API_MODE: method_description = ( method_description or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( method_name=method_name, default_method_name=_DEFAULT_STREAM_METHOD_NAME, return_type=_DEFAULT_STREAM_METHOD_RETURN_TYPE, ) ) method = _wrap_stream_query_operation( method_name=method_name, doc=method_description, ) else: raise ValueError( f"Unsupported api mode: `{api_mode}`," f" Supported modes are: `{_STANDARD_API_MODE}`" f" and `{_STREAM_API_MODE}`." ) # Binds the method to the object. setattr(obj, method_name, types.MethodType(method, obj)) def _get_registered_operations(reasoning_engine: Any) -> Dict[str, List[str]]: """Retrieves registered operations for a ReasoningEngine.""" if isinstance(reasoning_engine, OperationRegistrable): return reasoning_engine.register_operations() operations = {} if isinstance(reasoning_engine, Queryable): operations[_STANDARD_API_MODE] = [_DEFAULT_METHOD_NAME] if isinstance(reasoning_engine, StreamQueryable): operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME] return operations def _generate_class_methods_spec_or_raise( reasoning_engine: Any, operations: Dict[str, List[str]] ) -> List[proto.Message]: """Generates a ReasoningEngineSpec based on the registered operations. Args: reasoning_engine: The ReasoningEngine instance. operations: A dictionary of API modes and method names. Returns: A list of ReasoningEngineSpec.ClassMethod messages. Raises: ValueError: If a method defined in `register_operations` is not found on the ReasoningEngine. """ class_methods_spec = [] for mode, method_names in operations.items(): for method_name in method_names: if not hasattr(reasoning_engine, method_name): raise ValueError( f"Method `{method_name}` defined in `register_operations`" " not found on ReasoningEngine." ) method = getattr(reasoning_engine, method_name) try: schema_dict = _utils.generate_schema(method, schema_name=method_name) except Exception as e: _LOGGER.warning(f"failed to generate schema for {method_name}: {e}") continue class_method = _utils.to_proto(schema_dict) class_method[_MODE_KEY_IN_SCHEMA] = mode class_methods_spec.append(class_method) return class_methods_spec