From ad640414678eb89da43ab45fc83b6bdfa31ec6ce Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Thu, 28 May 2026 10:01:17 -0700 Subject: [PATCH] chore: GenAI client - add generated private methods for RAG GenAI module PiperOrigin-RevId: 922833969 --- agentplatform/_genai/rag.py | 879 ++++++++++++++++-- agentplatform/_genai/types/__init__.py | 56 ++ agentplatform/_genai/types/common.py | 342 ++++++- .../genai/replays/test_rag_delete.py | 34 + .../genai/replays/test_rag_delete_file.py | 35 + .../genai/replays/test_rag_update.py | 63 ++ .../genai/replays/test_rag_update_config.py | 56 ++ 7 files changed, 1382 insertions(+), 83 deletions(-) create mode 100644 tests/unit/agentplatform/genai/replays/test_rag_delete.py create mode 100644 tests/unit/agentplatform/genai/replays/test_rag_delete_file.py create mode 100644 tests/unit/agentplatform/genai/replays/test_rag_update.py create mode 100644 tests/unit/agentplatform/genai/replays/test_rag_update_config.py diff --git a/agentplatform/_genai/rag.py b/agentplatform/_genai/rag.py index 313a56109e..e7aa5a3a52 100644 --- a/agentplatform/_genai/rag.py +++ b/agentplatform/_genai/rag.py @@ -48,6 +48,37 @@ def _CreateRagCorpusRequestParameters_to_vertex( return to_object +def _DeleteRagCorpusRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + +def _DeleteRagFileRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["corpus_id"]) is not None: + setv(to_object, ["_url", "corpus_id"], getv(from_object, ["corpus_id"])) + + if getv(from_object, ["file_id"]) is not None: + setv(to_object, ["_url", "file_id"], getv(from_object, ["file_id"])) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + def _GetRagConfigRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -313,6 +344,86 @@ def _RagCorpus_to_vertex( return to_object +def _RagEngineConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(parent_object, ["name"]) is not None: + setv(to_object, ["name"], getv(parent_object, ["name"])) + + if getv(parent_object, ["ragManagedDbConfig"]) is not None: + setv( + to_object, + ["rag_managed_db_config"], + getv(parent_object, ["ragManagedDbConfig"]), + ) + + return to_object + + +def _RagEngineConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(parent_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["rag_managed_db_config"]) is not None: + setv( + parent_object, + ["ragManagedDbConfig"], + getv(from_object, ["rag_managed_db_config"]), + ) + + return to_object + + +def _UpdateRagConfigRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["updated_config"]) is not None: + setv( + to_object, + ["_self"], + _RagEngineConfig_to_vertex( + getv(from_object, ["updated_config"]), to_object + ), + ) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + +def _UpdateRagCorpusRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["corpus_id"]) is not None: + setv(to_object, ["_url", "corpus_id"], getv(from_object, ["corpus_id"])) + + if getv(from_object, ["rag_corpus"]) is not None: + setv( + to_object, + ["_self"], + _RagCorpus_to_vertex(getv(from_object, ["rag_corpus"]), to_object), + ) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + return to_object + + class Rag(_api_module.BaseModule): def _create_corpus( @@ -719,6 +830,9 @@ def get_config( response_dict = {} if not response.body else json.loads(response.body) + if self._api_client.vertexai: + response_dict = _RagEngineConfig_from_vertex(response_dict) + return_value = types.RagEngineConfig._from_response( response=response_dict, kwargs=( @@ -743,22 +857,23 @@ def get_config( self._api_client._verify_response(return_value) return return_value - -class AsyncRag(_api_module.BaseModule): - - async def _create_corpus( + def _update_corpus( self, *, + corpus_id: str, rag_corpus: types.RagCorpusOrDict, - config: Optional[types.CreateRagCorpusConfigOrDict] = None, - ) -> types.CreateRagCorpusOperation: + config: Optional[types.UpdateRagCorpusConfigOrDict] = None, + name: Optional[str] = None, + ) -> types.UpdateRagCorpusOperation: """ - Creates a new Rag Corpus. + Updates an existing Rag Corpus. """ - parameter_model = types._CreateRagCorpusRequestParameters( + parameter_model = types._UpdateRagCorpusRequestParameters( + corpus_id=corpus_id, rag_corpus=rag_corpus, config=config, + name=name, ) request_url_dict: Optional[dict[str, str]] @@ -767,12 +882,12 @@ async def _create_corpus( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateRagCorpusRequestParameters_to_vertex(parameter_model) + request_dict = _UpdateRagCorpusRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "ragCorpora".format_map(request_url_dict) + path = "ragCorpora/{corpus_id}".format_map(request_url_dict) else: - path = "ragCorpora" + path = "ragCorpora/{corpus_id}" query_params = request_dict.get("_query") if query_params: @@ -790,13 +905,11 @@ async def _create_corpus( request_dict = _common.convert_to_dict(request_dict) request_dict = _common.encode_unserializable_types(request_dict) - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) + response = self._api_client.request("patch", path, request_dict, http_options) response_dict = {} if not response.body else json.loads(response.body) - return_value = types.CreateRagCorpusOperation._from_response( + return_value = types.UpdateRagCorpusOperation._from_response( response=response_dict, kwargs=( { @@ -820,14 +933,17 @@ async def _create_corpus( self._api_client._verify_response(return_value) return return_value - async def get_corpus( - self, *, name: str, config: Optional[types.GetRagCorpusConfigOrDict] = None - ) -> types.RagCorpus: + def _delete_corpus( + self, + *, + name: Optional[str] = None, + config: Optional[types.DeleteRagCorpusConfigOrDict] = None, + ) -> types.DeleteRagCorpusOperation: """ - Gets a RAG Corpus. + Deletes a RAG Corpus. """ - parameter_model = types._GetRagCorpusRequestParameters( + parameter_model = types._DeleteRagCorpusRequestParameters( name=name, config=config, ) @@ -838,7 +954,7 @@ async def get_corpus( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetRagCorpusRequestParameters_to_vertex(parameter_model) + request_dict = _DeleteRagCorpusRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -861,16 +977,11 @@ async def get_corpus( request_dict = _common.convert_to_dict(request_dict) request_dict = _common.encode_unserializable_types(request_dict) - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) + response = self._api_client.request("delete", path, request_dict, http_options) response_dict = {} if not response.body else json.loads(response.body) - if self._api_client.vertexai: - response_dict = _RagCorpus_from_vertex(response_dict) - - return_value = types.RagCorpus._from_response( + return_value = types.DeleteRagCorpusOperation._from_response( response=response_dict, kwargs=( { @@ -894,14 +1005,20 @@ async def get_corpus( self._api_client._verify_response(return_value) return return_value - async def list_corpora( - self, *, config: Optional[types.ListRagCorporaConfigOrDict] = None - ) -> types.ListRagCorporaResponse: + def _delete_file( + self, + *, + corpus_id: str, + file_id: str, + config: Optional[types.DeleteRagFileConfigOrDict] = None, + ) -> types.DeleteRagFileOperation: """ - Lists RagCorpora. + Deletes a RAG File. """ - parameter_model = types._ListRagCorporaRequestParameters( + parameter_model = types._DeleteRagFileRequestParameters( + corpus_id=corpus_id, + file_id=file_id, config=config, ) @@ -911,12 +1028,14 @@ async def list_corpora( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListRagCorporaRequestParameters_to_vertex(parameter_model) + request_dict = _DeleteRagFileRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "ragCorpora".format_map(request_url_dict) + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}".format_map( + request_url_dict + ) else: - path = "ragCorpora" + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}" query_params = request_dict.get("_query") if query_params: @@ -934,16 +1053,11 @@ async def list_corpora( request_dict = _common.convert_to_dict(request_dict) request_dict = _common.encode_unserializable_types(request_dict) - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) + response = self._api_client.request("delete", path, request_dict, http_options) response_dict = {} if not response.body else json.loads(response.body) - if self._api_client.vertexai: - response_dict = _ListRagCorporaResponse_from_vertex(response_dict) - - return_value = types.ListRagCorporaResponse._from_response( + return_value = types.DeleteRagFileOperation._from_response( response=response_dict, kwargs=( { @@ -967,20 +1081,18 @@ async def list_corpora( self._api_client._verify_response(return_value) return return_value - async def get_file( + def _update_config( self, *, - corpus_id: str, - file_id: str, - config: Optional[types.GetRagFileConfigOrDict] = None, - ) -> types.RagFile: + updated_config: types.RagEngineConfigOrDict, + config: Optional[types.UpdateRagConfigOrDict] = None, + ) -> types.UpdateRagConfigOperation: """ - Gets a RagFile. + Updates a RAG Engine Config. """ - parameter_model = types._GetRagFileRequestParameters( - corpus_id=corpus_id, - file_id=file_id, + parameter_model = types._UpdateRagConfigRequestParameters( + updated_config=updated_config, config=config, ) @@ -990,14 +1102,12 @@ async def get_file( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetRagFileRequestParameters_to_vertex(parameter_model) + request_dict = _UpdateRagConfigRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "ragCorpora/{corpus_id}/ragFiles/{file_id}".format_map( - request_url_dict - ) + path = "ragEngineConfig".format_map(request_url_dict) else: - path = "ragCorpora/{corpus_id}/ragFiles/{file_id}" + path = "ragEngineConfig" query_params = request_dict.get("_query") if query_params: @@ -1015,13 +1125,11 @@ async def get_file( request_dict = _common.convert_to_dict(request_dict) request_dict = _common.encode_unserializable_types(request_dict) - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) + response = self._api_client.request("patch", path, request_dict, http_options) response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RagFile._from_response( + return_value = types.UpdateRagConfigOperation._from_response( response=response_dict, kwargs=( { @@ -1045,16 +1153,22 @@ async def get_file( self._api_client._verify_response(return_value) return return_value - async def list_files( - self, *, config: Optional[types.ListRagFilesConfigOrDict] = None, corpus_id: str - ) -> types.ListRagFilesResponse: + +class AsyncRag(_api_module.BaseModule): + + async def _create_corpus( + self, + *, + rag_corpus: types.RagCorpusOrDict, + config: Optional[types.CreateRagCorpusConfigOrDict] = None, + ) -> types.CreateRagCorpusOperation: """ - Lists RagFile instances within a RagCorpus. + Creates a new Rag Corpus. """ - parameter_model = types._ListRagFilesRequestParameters( + parameter_model = types._CreateRagCorpusRequestParameters( + rag_corpus=rag_corpus, config=config, - corpus_id=corpus_id, ) request_url_dict: Optional[dict[str, str]] @@ -1063,12 +1177,12 @@ async def list_files( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListRagFilesRequestParameters_to_vertex(parameter_model) + request_dict = _CreateRagCorpusRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "ragCorpora/{corpus_id}/ragFiles".format_map(request_url_dict) + path = "ragCorpora".format_map(request_url_dict) else: - path = "ragCorpora/{corpus_id}/ragFiles" + path = "ragCorpora" query_params = request_dict.get("_query") if query_params: @@ -1087,12 +1201,12 @@ async def list_files( request_dict = _common.encode_unserializable_types(request_dict) response = await self._api_client.async_request( - "get", path, request_dict, http_options + "post", path, request_dict, http_options ) response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListRagFilesResponse._from_response( + return_value = types.CreateRagCorpusOperation._from_response( response=response_dict, kwargs=( { @@ -1116,14 +1230,15 @@ async def list_files( self._api_client._verify_response(return_value) return return_value - async def get_config( - self, *, config: Optional[types.GetRagConfigOrDict] = None - ) -> types.RagEngineConfig: + async def get_corpus( + self, *, name: str, config: Optional[types.GetRagCorpusConfigOrDict] = None + ) -> types.RagCorpus: """ - Gets a RAG Engine Config. + Gets a RAG Corpus. """ - parameter_model = types._GetRagConfigRequestParameters( + parameter_model = types._GetRagCorpusRequestParameters( + name=name, config=config, ) @@ -1133,12 +1248,12 @@ async def get_config( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetRagConfigRequestParameters_to_vertex(parameter_model) + request_dict = _GetRagCorpusRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: - path = "ragEngineConfig".format_map(request_url_dict) + path = "{name}".format_map(request_url_dict) else: - path = "ragEngineConfig" + path = "{name}" query_params = request_dict.get("_query") if query_params: @@ -1162,7 +1277,609 @@ async def get_config( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RagEngineConfig._from_response( + if self._api_client.vertexai: + response_dict = _RagCorpus_from_vertex(response_dict) + + return_value = types.RagCorpus._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def list_corpora( + self, *, config: Optional[types.ListRagCorporaConfigOrDict] = None + ) -> types.ListRagCorporaResponse: + """ + Lists RagCorpora. + """ + + parameter_model = types._ListRagCorporaRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRagCorporaRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragCorpora".format_map(request_url_dict) + else: + path = "ragCorpora" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListRagCorporaResponse_from_vertex(response_dict) + + return_value = types.ListRagCorporaResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def get_file( + self, + *, + corpus_id: str, + file_id: str, + config: Optional[types.GetRagFileConfigOrDict] = None, + ) -> types.RagFile: + """ + Gets a RagFile. + """ + + parameter_model = types._GetRagFileRequestParameters( + corpus_id=corpus_id, + file_id=file_id, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRagFileRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}".format_map( + request_url_dict + ) + else: + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RagFile._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def list_files( + self, *, config: Optional[types.ListRagFilesConfigOrDict] = None, corpus_id: str + ) -> types.ListRagFilesResponse: + """ + Lists RagFile instances within a RagCorpus. + """ + + parameter_model = types._ListRagFilesRequestParameters( + config=config, + corpus_id=corpus_id, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRagFilesRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragCorpora/{corpus_id}/ragFiles".format_map(request_url_dict) + else: + path = "ragCorpora/{corpus_id}/ragFiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ListRagFilesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def get_config( + self, *, config: Optional[types.GetRagConfigOrDict] = None + ) -> types.RagEngineConfig: + """ + Gets a RAG Engine Config. + """ + + parameter_model = types._GetRagConfigRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRagConfigRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragEngineConfig".format_map(request_url_dict) + else: + path = "ragEngineConfig" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RagEngineConfig_from_vertex(response_dict) + + return_value = types.RagEngineConfig._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update_corpus( + self, + *, + corpus_id: str, + rag_corpus: types.RagCorpusOrDict, + config: Optional[types.UpdateRagCorpusConfigOrDict] = None, + name: Optional[str] = None, + ) -> types.UpdateRagCorpusOperation: + """ + Updates an existing Rag Corpus. + """ + + parameter_model = types._UpdateRagCorpusRequestParameters( + corpus_id=corpus_id, + rag_corpus=rag_corpus, + config=config, + name=name, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRagCorpusRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragCorpora/{corpus_id}".format_map(request_url_dict) + else: + path = "ragCorpora/{corpus_id}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.UpdateRagCorpusOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete_corpus( + self, + *, + name: Optional[str] = None, + config: Optional[types.DeleteRagCorpusConfigOrDict] = None, + ) -> types.DeleteRagCorpusOperation: + """ + Deletes a RAG Corpus. + """ + + parameter_model = types._DeleteRagCorpusRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRagCorpusRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRagCorpusOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete_file( + self, + *, + corpus_id: str, + file_id: str, + config: Optional[types.DeleteRagFileConfigOrDict] = None, + ) -> types.DeleteRagFileOperation: + """ + Deletes a RAG File. + """ + + parameter_model = types._DeleteRagFileRequestParameters( + corpus_id=corpus_id, + file_id=file_id, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRagFileRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}".format_map( + request_url_dict + ) + else: + path = "ragCorpora/{corpus_id}/ragFiles/{file_id}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRagFileOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update_config( + self, + *, + updated_config: types.RagEngineConfigOrDict, + config: Optional[types.UpdateRagConfigOrDict] = None, + ) -> types.UpdateRagConfigOperation: + """ + Updates a RAG Engine Config. + """ + + parameter_model = types._UpdateRagConfigRequestParameters( + updated_config=updated_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRagConfigRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "ragEngineConfig".format_map(request_url_dict) + else: + path = "ragEngineConfig" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.UpdateRagConfigOperation._from_response( response=response_dict, kwargs=( { diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index d6eebed8e5..7161ab4e7c 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -56,6 +56,8 @@ from .common import _DeleteEvaluationMetricParameters from .common import _DeleteMultimodalDatasetRequestParameters from .common import _DeletePromptVersionRequestParameters +from .common import _DeleteRagCorpusRequestParameters +from .common import _DeleteRagFileRequestParameters from .common import _DeleteSandboxEnvironmentSnapshotRequestParameters from .common import _DeleteSandboxEnvironmentTemplateRequestParameters from .common import _DeleteSkillRequestParameters @@ -138,6 +140,8 @@ from .common import _UpdateAgentEngineSessionRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateMultimodalDatasetParameters +from .common import _UpdateRagConfigRequestParameters +from .common import _UpdateRagCorpusRequestParameters from .common import _UpdateSkillRequestParameters from .common import A2aTask from .common import A2aTaskDict @@ -404,6 +408,18 @@ from .common import DeletePromptVersionOperation from .common import DeletePromptVersionOperationDict from .common import DeletePromptVersionOperationOrDict +from .common import DeleteRagCorpusConfig +from .common import DeleteRagCorpusConfigDict +from .common import DeleteRagCorpusConfigOrDict +from .common import DeleteRagCorpusOperation +from .common import DeleteRagCorpusOperationDict +from .common import DeleteRagCorpusOperationOrDict +from .common import DeleteRagFileConfig +from .common import DeleteRagFileConfigDict +from .common import DeleteRagFileConfigOrDict +from .common import DeleteRagFileOperation +from .common import DeleteRagFileOperationDict +from .common import DeleteRagFileOperationOrDict from .common import DeleteSandboxEnvironmentSnapshotConfig from .common import DeleteSandboxEnvironmentSnapshotConfigDict from .common import DeleteSandboxEnvironmentSnapshotConfigOrDict @@ -1615,6 +1631,18 @@ from .common import UpdatePromptConfig from .common import UpdatePromptConfigDict from .common import UpdatePromptConfigOrDict +from .common import UpdateRagConfig +from .common import UpdateRagConfigDict +from .common import UpdateRagConfigOperation +from .common import UpdateRagConfigOperationDict +from .common import UpdateRagConfigOperationOrDict +from .common import UpdateRagConfigOrDict +from .common import UpdateRagCorpusConfig +from .common import UpdateRagCorpusConfigDict +from .common import UpdateRagCorpusConfigOrDict +from .common import UpdateRagCorpusOperation +from .common import UpdateRagCorpusOperationDict +from .common import UpdateRagCorpusOperationOrDict from .common import UpdateSkillConfig from .common import UpdateSkillConfigDict from .common import UpdateSkillConfigOrDict @@ -2556,6 +2584,30 @@ "RagEngineConfig", "RagEngineConfigDict", "RagEngineConfigOrDict", + "UpdateRagCorpusConfig", + "UpdateRagCorpusConfigDict", + "UpdateRagCorpusConfigOrDict", + "UpdateRagCorpusOperation", + "UpdateRagCorpusOperationDict", + "UpdateRagCorpusOperationOrDict", + "DeleteRagCorpusConfig", + "DeleteRagCorpusConfigDict", + "DeleteRagCorpusConfigOrDict", + "DeleteRagCorpusOperation", + "DeleteRagCorpusOperationDict", + "DeleteRagCorpusOperationOrDict", + "DeleteRagFileConfig", + "DeleteRagFileConfigDict", + "DeleteRagFileConfigOrDict", + "DeleteRagFileOperation", + "DeleteRagFileOperationDict", + "DeleteRagFileOperationOrDict", + "UpdateRagConfig", + "UpdateRagConfigDict", + "UpdateRagConfigOrDict", + "UpdateRagConfigOperation", + "UpdateRagConfigOperationDict", + "UpdateRagConfigOperationOrDict", "GetAgentEngineRuntimeRevisionConfig", "GetAgentEngineRuntimeRevisionConfigDict", "GetAgentEngineRuntimeRevisionConfigOrDict", @@ -3152,6 +3204,10 @@ "_GetRagFileRequestParameters", "_ListRagFilesRequestParameters", "_GetRagConfigRequestParameters", + "_UpdateRagCorpusRequestParameters", + "_DeleteRagCorpusRequestParameters", + "_DeleteRagFileRequestParameters", + "_UpdateRagConfigRequestParameters", "_GetAgentEngineRuntimeRevisionRequestParameters", "_ListAgentEngineRuntimeRevisionsRequestParameters", "_DeleteAgentEngineRuntimeRevisionRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 3a2e17b8b6..561921fd15 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -13601,7 +13601,7 @@ class RagManagedDbConfigSpannerDict(TypedDict, total=False): class RagManagedDbConfig(_common.BaseModel): - """Configuration message for RagManagedDb used by RagEngine.""" + """The backend config of the RagEngineConfig.""" basic: Optional[RagManagedDbConfigBasic] = Field( default=None, @@ -13629,7 +13629,7 @@ class RagManagedDbConfig(_common.BaseModel): class RagManagedDbConfigDict(TypedDict, total=False): - """Configuration message for RagManagedDb used by RagEngine.""" + """The backend config of the RagEngineConfig.""" basic: Optional[RagManagedDbConfigBasicDict] """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Basic tier.""" @@ -13677,6 +13677,344 @@ class RagEngineConfigDict(TypedDict, total=False): RagEngineConfigOrDict = Union[RagEngineConfig, RagEngineConfigDict] +class UpdateRagCorpusConfig(_common.BaseModel): + """Config for updating a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class UpdateRagCorpusConfigDict(TypedDict, total=False): + """Config for updating a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +UpdateRagCorpusConfigOrDict = Union[UpdateRagCorpusConfig, UpdateRagCorpusConfigDict] + + +class _UpdateRagCorpusRequestParameters(_common.BaseModel): + """Parameters for updating a RAG corpus.""" + + corpus_id: Optional[str] = Field(default=None, description="""""") + rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") + config: Optional[UpdateRagCorpusConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") + + +class _UpdateRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for updating a RAG corpus.""" + + corpus_id: Optional[str] + """""" + + rag_corpus: Optional[RagCorpusDict] + """""" + + config: Optional[UpdateRagCorpusConfigDict] + """""" + + name: Optional[str] + """""" + + +_UpdateRagCorpusRequestParametersOrDict = Union[ + _UpdateRagCorpusRequestParameters, _UpdateRagCorpusRequestParametersDict +] + + +class UpdateRagCorpusOperation(_common.BaseModel): + """Operation for updating a RAG corpus.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class UpdateRagCorpusOperationDict(TypedDict, total=False): + """Operation for updating a RAG corpus.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +UpdateRagCorpusOperationOrDict = Union[ + UpdateRagCorpusOperation, UpdateRagCorpusOperationDict +] + + +class DeleteRagCorpusConfig(_common.BaseModel): + """Config for deleting a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class DeleteRagCorpusConfigDict(TypedDict, total=False): + """Config for deleting a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +DeleteRagCorpusConfigOrDict = Union[DeleteRagCorpusConfig, DeleteRagCorpusConfigDict] + + +class _DeleteRagCorpusRequestParameters(_common.BaseModel): + """Parameters for deleting a RAG corpus.""" + + name: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteRagCorpusConfig] = Field(default=None, description="""""") + + +class _DeleteRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a RAG corpus.""" + + name: Optional[str] + """""" + + config: Optional[DeleteRagCorpusConfigDict] + """""" + + +_DeleteRagCorpusRequestParametersOrDict = Union[ + _DeleteRagCorpusRequestParameters, _DeleteRagCorpusRequestParametersDict +] + + +class DeleteRagCorpusOperation(_common.BaseModel): + """Operation for deleting a RAG corpus.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class DeleteRagCorpusOperationDict(TypedDict, total=False): + """Operation for deleting a RAG corpus.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteRagCorpusOperationOrDict = Union[ + DeleteRagCorpusOperation, DeleteRagCorpusOperationDict +] + + +class DeleteRagFileConfig(_common.BaseModel): + """Config for deleting a RAG File.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class DeleteRagFileConfigDict(TypedDict, total=False): + """Config for deleting a RAG File.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +DeleteRagFileConfigOrDict = Union[DeleteRagFileConfig, DeleteRagFileConfigDict] + + +class _DeleteRagFileRequestParameters(_common.BaseModel): + """Parameters for deleting a RAG File.""" + + corpus_id: Optional[str] = Field(default=None, description="""""") + file_id: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteRagFileConfig] = Field(default=None, description="""""") + + +class _DeleteRagFileRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a RAG File.""" + + corpus_id: Optional[str] + """""" + + file_id: Optional[str] + """""" + + config: Optional[DeleteRagFileConfigDict] + """""" + + +_DeleteRagFileRequestParametersOrDict = Union[ + _DeleteRagFileRequestParameters, _DeleteRagFileRequestParametersDict +] + + +class DeleteRagFileOperation(_common.BaseModel): + """Operation for deleting a RAG File.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class DeleteRagFileOperationDict(TypedDict, total=False): + """Operation for deleting a RAG File.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteRagFileOperationOrDict = Union[DeleteRagFileOperation, DeleteRagFileOperationDict] + + +class UpdateRagConfig(_common.BaseModel): + """Config for updating a RAG Config.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class UpdateRagConfigDict(TypedDict, total=False): + """Config for updating a RAG Config.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +UpdateRagConfigOrDict = Union[UpdateRagConfig, UpdateRagConfigDict] + + +class _UpdateRagConfigRequestParameters(_common.BaseModel): + """Parameters for updating a RAG Config.""" + + updated_config: Optional[RagEngineConfig] = Field(default=None, description="""""") + config: Optional[UpdateRagConfig] = Field(default=None, description="""""") + + +class _UpdateRagConfigRequestParametersDict(TypedDict, total=False): + """Parameters for updating a RAG Config.""" + + updated_config: Optional[RagEngineConfigDict] + """""" + + config: Optional[UpdateRagConfigDict] + """""" + + +_UpdateRagConfigRequestParametersOrDict = Union[ + _UpdateRagConfigRequestParameters, _UpdateRagConfigRequestParametersDict +] + + +class UpdateRagConfigOperation(_common.BaseModel): + """Operation for updating a RAG Config.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class UpdateRagConfigOperationDict(TypedDict, total=False): + """Operation for updating a RAG Config.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +UpdateRagConfigOperationOrDict = Union[ + UpdateRagConfigOperation, UpdateRagConfigOperationDict +] + + class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): """Config for getting an Agent Engine Runtime Revision.""" diff --git a/tests/unit/agentplatform/genai/replays/test_rag_delete.py b/tests/unit/agentplatform/genai/replays/test_rag_delete.py new file mode 100644 index 0000000000..afeadc5663 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_rag_delete.py @@ -0,0 +1,34 @@ +# Copyright 2026 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. +# +"""Tests the rag._delete_corpus() method against the Agent Platform endpoint using replays.""" + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_delete_rag_corpus_private(client): + + corpus_op = client.rag._delete_corpus( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/1113515007867355136", + corpus_id="1113515007867355136", + ) + + assert isinstance(corpus_op, types.DeleteRagCorpusOperation) diff --git a/tests/unit/agentplatform/genai/replays/test_rag_delete_file.py b/tests/unit/agentplatform/genai/replays/test_rag_delete_file.py new file mode 100644 index 0000000000..2a4ef22d09 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_rag_delete_file.py @@ -0,0 +1,35 @@ +# Copyright 2026 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. +# +"""Tests the rag._delete_file() method against the Agent Platform endpoint using replays.""" + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_delete_rag_file_private(client): + + file_op = client.rag._delete_file( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/6301661778598166528/ragFiles/5708249684881115624", + corpus_id="6301661778598166528", + file_id="5708249684881115624", + ) + + assert isinstance(file_op, types.DeleteRagFileOperation) diff --git a/tests/unit/agentplatform/genai/replays/test_rag_update.py b/tests/unit/agentplatform/genai/replays/test_rag_update.py new file mode 100644 index 0000000000..fe562093c0 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_rag_update.py @@ -0,0 +1,63 @@ +# Copyright 2026 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. +# +"""Tests the rag._update_corpus() method against the Agent Platform endpoint using replays.""" + +import pytest + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_update_rag_corpus_private(client): + + corpus_op = client.rag._update_corpus( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/5685794529555251200", + corpus_id="5685794529555251200", + rag_corpus=types.RagCorpus( + display_name="My Updated Vertex AI Search Test Corpus", + description="My Updated Test Corpus Description", + vertex_ai_search_config=types.VertexAiSearchConfig( + serving_config="projects/vertex-sdk-dev/locations/us-central1/collections/default_collection/engines/test-engine/servingConfigs/default_serving_config" + ), + ), + ) + + assert isinstance(corpus_op, types.UpdateRagCorpusOperation) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_update_rag_corpus_private_async(client): + corpus_op = await client.aio.rag._update_corpus( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/5685794529555251200", + corpus_id="5685794529555251200", + rag_corpus=types.RagCorpus( + display_name="My Updated Vertex AI Search Test Corpus", + description="My Updated Test Corpus Description", + vertex_ai_search_config=types.VertexAiSearchConfig( + serving_config="projects/vertex-sdk-dev/locations/us-central1/collections/default_collection/engines/test-engine/servingConfigs/default_serving_config" + ), + ), + ) + + assert isinstance(corpus_op, types.UpdateRagCorpusOperation) diff --git a/tests/unit/agentplatform/genai/replays/test_rag_update_config.py b/tests/unit/agentplatform/genai/replays/test_rag_update_config.py new file mode 100644 index 0000000000..4d81241440 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_rag_update_config.py @@ -0,0 +1,56 @@ +# Copyright 2026 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. +# +"""Tests the rag._update_config() method against the Agent Platform endpoint using replays.""" + +import pytest + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform._genai import types + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_update_rag_config_private(client): + config_op = client.rag._update_config( + updated_config=types.RagEngineConfig( + name="projects/vertex-sdk-dev/locations/us-central1/ragEngineConfig/test_rag_config", + rag_managed_db_config=types.RagManagedDbConfig( + serverless=types.RagManagedDbConfigServerless() + ), + ), + ) + + assert isinstance(config_op, types.UpdateRagConfigOperation) + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_update_rag_config_private_async(client): + config_op = await client.aio.rag._update_config( + updated_config=types.RagEngineConfig( + name="projects/vertex-sdk-dev/locations/us-central1/ragEngineConfig/test_rag_config", + rag_managed_db_config=types.RagManagedDbConfig( + serverless=types.RagManagedDbConfigServerless() + ), + ), + ) + + assert isinstance(config_op, types.UpdateRagConfigOperation)