diff --git a/agentplatform/_genai/rag.py b/agentplatform/_genai/rag.py index f00db97a5e..26e14942c6 100644 --- a/agentplatform/_genai/rag.py +++ b/agentplatform/_genai/rag.py @@ -135,6 +135,19 @@ def _GetCorpusOperationParameters_to_vertex( return to_object +def _GetImportFilesOperationParameters_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, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operation_name"], getv(from_object, ["operation_name"]) + ) + + return to_object + + def _GetRagConfigRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -174,6 +187,148 @@ def _GetRagFileRequestParameters_to_vertex( return to_object +def _ImportRagFilesConfig_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, ["gcs_source"]) is not None: + setv(to_object, ["gcsSource"], getv(from_object, ["gcs_source"])) + + if getv(from_object, ["global_max_embedding_requests_per_min"]) is not None: + setv( + to_object, + ["globalMaxEmbeddingRequestsPerMin"], + getv(from_object, ["global_max_embedding_requests_per_min"]), + ) + + if getv(from_object, ["google_drive_source"]) is not None: + setv( + to_object, ["googleDriveSource"], getv(from_object, ["google_drive_source"]) + ) + + if getv(from_object, ["import_result_bigquery_sink"]) is not None: + setv( + to_object, + ["importResultBigquerySink"], + getv(from_object, ["import_result_bigquery_sink"]), + ) + + if getv(from_object, ["import_result_gcs_sink"]) is not None: + setv( + to_object, + ["importResultGcsSink"], + getv(from_object, ["import_result_gcs_sink"]), + ) + + if getv(from_object, ["jira_source"]) is not None: + setv(to_object, ["jiraSource"], getv(from_object, ["jira_source"])) + + if getv(from_object, ["max_embedding_requests_per_min"]) is not None: + setv( + to_object, + ["maxEmbeddingRequestsPerMin"], + getv(from_object, ["max_embedding_requests_per_min"]), + ) + + if getv(from_object, ["partial_failure_bigquery_sink"]) is not None: + setv( + to_object, + ["partialFailureBigquerySink"], + getv(from_object, ["partial_failure_bigquery_sink"]), + ) + + if getv(from_object, ["partial_failure_gcs_sink"]) is not None: + setv( + to_object, + ["partialFailureGcsSink"], + getv(from_object, ["partial_failure_gcs_sink"]), + ) + + if getv(from_object, ["rag_file_chunking_config"]) is not None: + setv( + to_object, + ["ragFileChunkingConfig"], + getv(from_object, ["rag_file_chunking_config"]), + ) + + if getv(from_object, ["rag_file_metadata_config"]) is not None: + setv( + to_object, + ["ragFileMetadataConfig"], + getv(from_object, ["rag_file_metadata_config"]), + ) + + if getv(from_object, ["rag_file_parsing_config"]) is not None: + setv( + to_object, + ["ragFileParsingConfig"], + _RagFileParsingConfig_to_vertex( + getv(from_object, ["rag_file_parsing_config"]), to_object + ), + ) + + if getv(from_object, ["rag_file_transformation_config"]) is not None: + setv( + to_object, + ["ragFileTransformationConfig"], + getv(from_object, ["rag_file_transformation_config"]), + ) + + if getv(from_object, ["rebuild_ann_index"]) is not None: + setv(to_object, ["rebuildAnnIndex"], getv(from_object, ["rebuild_ann_index"])) + + if getv(from_object, ["share_point_sources"]) is not None: + setv( + to_object, ["sharePointSources"], getv(from_object, ["share_point_sources"]) + ) + + if getv(from_object, ["slack_source"]) is not None: + setv(to_object, ["slackSource"], getv(from_object, ["slack_source"])) + + return to_object + + +def _ImportRagFilesRequestParameters_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, ["import_rag_files_request"]) is not None: + setv( + to_object, + ["_self"], + _ImportRagFilesRequest_to_vertex( + getv(from_object, ["import_rag_files_request"]), 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, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _ImportRagFilesRequest_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, ["import_rag_files_config"]) is not None: + setv( + to_object, + ["importRagFilesConfig"], + _ImportRagFilesConfig_to_vertex( + getv(from_object, ["import_rag_files_config"]), to_object + ), + ) + + return to_object + + def _ListRagCorporaConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -812,6 +967,36 @@ def _RagFileParsingConfigLlmParser_to_vertex( return to_object +def _RagFileParsingConfig_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, ["advanced_parser"]) is not None: + setv(to_object, ["advancedParser"], getv(from_object, ["advanced_parser"])) + + if getv(from_object, ["layout_parser"]) is not None: + setv(to_object, ["layoutParser"], getv(from_object, ["layout_parser"])) + + if getv(from_object, ["llm_parser"]) is not None: + setv( + to_object, + ["llmParser"], + _RagFileParsingConfigLlmParser_to_vertex( + getv(from_object, ["llm_parser"]), to_object + ), + ) + + if getv(from_object, ["use_advanced_pdf_parsing"]) is not None: + setv( + to_object, + ["useAdvancedPdfParsing"], + getv(from_object, ["use_advanced_pdf_parsing"]), + ) + + return to_object + + def _RagManagedDbConfigSpanner_from_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -2300,6 +2485,148 @@ def retrieve_contexts( self._api_client._verify_response(return_value) return return_value + def _import_files( + self, + *, + import_rag_files_request: types.ImportRagFilesRequestOrDict, + config: Optional[types.ImportRagFilesRequestConfigOrDict] = None, + name: str, + ) -> types.ImportRagFilesOperation: + """ + Imports files into a RAG Corpus. + """ + + parameter_model = types._ImportRagFilesRequestParameters( + import_rag_files_request=import_rag_files_request, + 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 = _ImportRagFilesRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/ragFiles:import".format_map(request_url_dict) + else: + path = "{name}/ragFiles:import" + + 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 = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportRagFilesOperation._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 + + def _get_import_files_operation( + self, + *, + operation_name: str, + config: Optional[types.GetImportFilesOperationConfigOrDict] = None, + ) -> types.ImportRagFilesOperation: + parameter_model = types._GetImportFilesOperationParameters( + operation_name=operation_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 = _GetImportFilesOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operation_name}".format_map(request_url_dict) + else: + path = "{operation_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 = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportRagFilesOperation._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 + def create_corpus( self, *, @@ -2372,6 +2699,58 @@ def delete_file( return None + def import_files( + self, + *, + name: str, + import_config: types.ImportRagFilesConfigOrDict, + config: Optional[types.ImportRagFilesRequestConfigOrDict] = None, + ) -> types.ImportRagFilesResponse: + """ + Imports files into a Rag Corpus. To get a list of all files on the corpus after calling import_files, call list_files. + + Example usage: + + ``` + corpus = client.rag.import_files( + name="projects/test-project/locations/us-central1/ragCorpora/123456789", + import_config=types.ImportRagFilesConfig( + gcs_source=types.GcsSource(uris=["gs://test-bucket/test-file.pdf"]), + ), + ) + + file_list = client.rag.list_files(name=corpus.name) + ``` + + Args: + name: The name of the Rag Corpus to import files into, format: + `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}` + import_config: The configuration for importing files, including source files and processing options. + config: The configuration to use for the import request. + + Returns: + The ImportRagFilesResponse from the import files operation. + """ + operation = self._import_files( + name=name, + import_rag_files_request=types.ImportRagFilesRequest( + import_rag_files_config=import_config + ), + config=config, + ) + + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self._get_import_files_operation, + ) + + if operation.error: + raise RuntimeError( + f"Failed to import files into RagCorpus: {operation.error}" + ) + + return operation.response + class AsyncRag(_api_module.BaseModule): @@ -3333,6 +3712,152 @@ async def retrieve_contexts( self._api_client._verify_response(return_value) return return_value + async def _import_files( + self, + *, + import_rag_files_request: types.ImportRagFilesRequestOrDict, + config: Optional[types.ImportRagFilesRequestConfigOrDict] = None, + name: str, + ) -> types.ImportRagFilesOperation: + """ + Imports files into a RAG Corpus. + """ + + parameter_model = types._ImportRagFilesRequestParameters( + import_rag_files_request=import_rag_files_request, + 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 = _ImportRagFilesRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}/ragFiles:import".format_map(request_url_dict) + else: + path = "{name}/ragFiles:import" + + 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( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportRagFilesOperation._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_import_files_operation( + self, + *, + operation_name: str, + config: Optional[types.GetImportFilesOperationConfigOrDict] = None, + ) -> types.ImportRagFilesOperation: + parameter_model = types._GetImportFilesOperationParameters( + operation_name=operation_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 = _GetImportFilesOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operation_name}".format_map(request_url_dict) + else: + path = "{operation_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( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportRagFilesOperation._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 create_corpus( self, *, @@ -3404,3 +3929,55 @@ async def delete_file( ) return None + + async def import_files( + self, + *, + name: str, + import_config: types.ImportRagFilesConfigOrDict, + config: Optional[types.ImportRagFilesRequestConfigOrDict] = None, + ) -> types.ImportRagFilesResponse: + """ + Imports files into a Rag Corpus. To get a list of all files on the corpus after calling import_files, call list_files. + + Example usage: + + ``` + corpus = client.rag.import_files( + name="projects/test-project/locations/us-central1/ragCorpora/123456789", + import_config=types.ImportRagFilesConfig( + gcs_source=types.GcsSource(uris=["gs://test-bucket/test-file.pdf"]), + ), + ) + + file_list = client.rag.list_files(name=corpus.name) + ``` + + Args: + name: The name of the Rag Corpus to import files into, format: + `projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}` + import_config: The configuration for importing files, including source files and processing options. + config: The configuration to use for the import request. + + Returns: + The ImportRagFilesResponse from the import files operation. + """ + operation = await self._import_files( + name=name, + import_rag_files_request=types.ImportRagFilesRequest( + import_rag_files_config=import_config + ), + config=config, + ) + + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self._get_import_files_operation, + ) + + if operation.error: + raise RuntimeError( + f"Failed to import files into RagCorpus: {operation.error}" + ) + + return operation.response diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 8f4787820f..7887de83cb 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -92,6 +92,7 @@ from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters +from .common import _GetImportFilesOperationParameters from .common import _GetMultimodalDatasetOperationParameters from .common import _GetMultimodalDatasetParameters from .common import _GetRagConfigRequestParameters @@ -103,6 +104,7 @@ from .common import _GetSkillOperationParameters from .common import _GetSkillRequestParameters from .common import _GetSkillRevisionRequestParameters +from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters from .common import _ListAgentEngineMemoryRequestParameters from .common import _ListAgentEngineMemoryRevisionsRequestParameters @@ -245,6 +247,9 @@ from .common import BatchPredictionValidationAssessmentResult from .common import BatchPredictionValidationAssessmentResultDict from .common import BatchPredictionValidationAssessmentResultOrDict +from .common import BigQueryDestination +from .common import BigQueryDestinationDict +from .common import BigQueryDestinationOrDict from .common import BigQueryRequestSet from .common import BigQueryRequestSetDict from .common import BigQueryRequestSetOrDict @@ -584,6 +589,9 @@ from .common import FileStatusDict from .common import FileStatusOrDict from .common import Framework +from .common import GcsDestination +from .common import GcsDestinationDict +from .common import GcsDestinationOrDict from .common import GcsSource from .common import GcsSourceDict from .common import GcsSourceOrDict @@ -684,6 +692,9 @@ from .common import GetEvaluationSetConfig from .common import GetEvaluationSetConfigDict from .common import GetEvaluationSetConfigOrDict +from .common import GetImportFilesOperationConfig +from .common import GetImportFilesOperationConfigDict +from .common import GetImportFilesOperationConfigOrDict from .common import GetMultimodalDatasetOperationConfig from .common import GetMultimodalDatasetOperationConfigDict from .common import GetMultimodalDatasetOperationConfigOrDict @@ -722,6 +733,21 @@ from .common import GoogleDriveSourceResourceIdOrDict from .common import IdentityType from .common import Importance +from .common import ImportRagFilesConfig +from .common import ImportRagFilesConfigDict +from .common import ImportRagFilesConfigOrDict +from .common import ImportRagFilesOperation +from .common import ImportRagFilesOperationDict +from .common import ImportRagFilesOperationOrDict +from .common import ImportRagFilesRequest +from .common import ImportRagFilesRequestConfig +from .common import ImportRagFilesRequestConfigDict +from .common import ImportRagFilesRequestConfigOrDict +from .common import ImportRagFilesRequestDict +from .common import ImportRagFilesRequestOrDict +from .common import ImportRagFilesResponse +from .common import ImportRagFilesResponseDict +from .common import ImportRagFilesResponseOrDict from .common import IngestEventsConfig from .common import IngestEventsConfigDict from .common import IngestEventsConfigOrDict @@ -1110,11 +1136,32 @@ from .common import RagEngineConfigDict from .common import RagEngineConfigOrDict from .common import RagFile +from .common import RagFileChunkingConfig +from .common import RagFileChunkingConfigDict +from .common import RagFileChunkingConfigFixedLengthChunking +from .common import RagFileChunkingConfigFixedLengthChunkingDict +from .common import RagFileChunkingConfigFixedLengthChunkingOrDict +from .common import RagFileChunkingConfigOrDict from .common import RagFileDict +from .common import RagFileMetadataConfig +from .common import RagFileMetadataConfigDict +from .common import RagFileMetadataConfigOrDict from .common import RagFileOrDict +from .common import RagFileParsingConfig +from .common import RagFileParsingConfigAdvancedParser +from .common import RagFileParsingConfigAdvancedParserDict +from .common import RagFileParsingConfigAdvancedParserOrDict +from .common import RagFileParsingConfigDict +from .common import RagFileParsingConfigLayoutParser +from .common import RagFileParsingConfigLayoutParserDict +from .common import RagFileParsingConfigLayoutParserOrDict from .common import RagFileParsingConfigLlmParser from .common import RagFileParsingConfigLlmParserDict from .common import RagFileParsingConfigLlmParserOrDict +from .common import RagFileParsingConfigOrDict +from .common import RagFileTransformationConfig +from .common import RagFileTransformationConfigDict +from .common import RagFileTransformationConfigOrDict from .common import RagFileType from .common import RagManagedDbConfig from .common import RagManagedDbConfigBasic @@ -2731,6 +2778,51 @@ "RetrieveContextsResponse", "RetrieveContextsResponseDict", "RetrieveContextsResponseOrDict", + "ImportRagFilesRequestConfig", + "ImportRagFilesRequestConfigDict", + "ImportRagFilesRequestConfigOrDict", + "BigQueryDestination", + "BigQueryDestinationDict", + "BigQueryDestinationOrDict", + "GcsDestination", + "GcsDestinationDict", + "GcsDestinationOrDict", + "RagFileChunkingConfigFixedLengthChunking", + "RagFileChunkingConfigFixedLengthChunkingDict", + "RagFileChunkingConfigFixedLengthChunkingOrDict", + "RagFileChunkingConfig", + "RagFileChunkingConfigDict", + "RagFileChunkingConfigOrDict", + "RagFileMetadataConfig", + "RagFileMetadataConfigDict", + "RagFileMetadataConfigOrDict", + "RagFileParsingConfigAdvancedParser", + "RagFileParsingConfigAdvancedParserDict", + "RagFileParsingConfigAdvancedParserOrDict", + "RagFileParsingConfigLayoutParser", + "RagFileParsingConfigLayoutParserDict", + "RagFileParsingConfigLayoutParserOrDict", + "RagFileParsingConfig", + "RagFileParsingConfigDict", + "RagFileParsingConfigOrDict", + "RagFileTransformationConfig", + "RagFileTransformationConfigDict", + "RagFileTransformationConfigOrDict", + "ImportRagFilesConfig", + "ImportRagFilesConfigDict", + "ImportRagFilesConfigOrDict", + "ImportRagFilesRequest", + "ImportRagFilesRequestDict", + "ImportRagFilesRequestOrDict", + "ImportRagFilesResponse", + "ImportRagFilesResponseDict", + "ImportRagFilesResponseOrDict", + "ImportRagFilesOperation", + "ImportRagFilesOperationDict", + "ImportRagFilesOperationOrDict", + "GetImportFilesOperationConfig", + "GetImportFilesOperationConfigDict", + "GetImportFilesOperationConfigOrDict", "GetAgentEngineRuntimeRevisionConfig", "GetAgentEngineRuntimeRevisionConfigDict", "GetAgentEngineRuntimeRevisionConfigOrDict", @@ -3334,6 +3426,8 @@ "_DeleteRagFileRequestParameters", "_UpdateRagConfigRequestParameters", "_RetrieveRagContextsRequestParameters", + "_ImportRagFilesRequestParameters", + "_GetImportFilesOperationParameters", "_GetAgentEngineRuntimeRevisionRequestParameters", "_ListAgentEngineRuntimeRevisionsRequestParameters", "_DeleteAgentEngineRuntimeRevisionRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index e5638c98ef..69292288a4 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -14654,6 +14654,598 @@ class RetrieveContextsResponseDict(TypedDict, total=False): ] +class ImportRagFilesRequestConfig(_common.BaseModel): + """Config for importing RAG files.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class ImportRagFilesRequestConfigDict(TypedDict, total=False): + """Config for importing RAG files.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +ImportRagFilesRequestConfigOrDict = Union[ + ImportRagFilesRequestConfig, ImportRagFilesRequestConfigDict +] + + +class BigQueryDestination(_common.BaseModel): + """The BigQuery location for the output content.""" + + output_uri: Optional[str] = Field( + default=None, + description="""Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""", + ) + + +class BigQueryDestinationDict(TypedDict, total=False): + """The BigQuery location for the output content.""" + + output_uri: Optional[str] + """Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""" + + +BigQueryDestinationOrDict = Union[BigQueryDestination, BigQueryDestinationDict] + + +class GcsDestination(_common.BaseModel): + """The Google Cloud Storage location where the output is to be written to.""" + + output_uri_prefix: Optional[str] = Field( + default=None, + description="""Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist.""", + ) + + +class GcsDestinationDict(TypedDict, total=False): + """The Google Cloud Storage location where the output is to be written to.""" + + output_uri_prefix: Optional[str] + """Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist.""" + + +GcsDestinationOrDict = Union[GcsDestination, GcsDestinationDict] + + +class RagFileChunkingConfigFixedLengthChunking(_common.BaseModel): + """Specifies the fixed length chunking config.""" + + chunk_overlap: Optional[int] = Field( + default=None, description="""The overlap between chunks.""" + ) + chunk_size: Optional[int] = Field( + default=None, description="""The size of the chunks.""" + ) + + +class RagFileChunkingConfigFixedLengthChunkingDict(TypedDict, total=False): + """Specifies the fixed length chunking config.""" + + chunk_overlap: Optional[int] + """The overlap between chunks.""" + + chunk_size: Optional[int] + """The size of the chunks.""" + + +RagFileChunkingConfigFixedLengthChunkingOrDict = Union[ + RagFileChunkingConfigFixedLengthChunking, + RagFileChunkingConfigFixedLengthChunkingDict, +] + + +class RagFileChunkingConfig(_common.BaseModel): + """Specifies the size and overlap of chunks for RagFiles.""" + + chunk_overlap: Optional[int] = Field( + default=None, description="""The overlap between chunks.""" + ) + chunk_size: Optional[int] = Field( + default=None, description="""The size of the chunks.""" + ) + fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunking] = Field( + default=None, description="""Specifies the fixed length chunking config.""" + ) + + +class RagFileChunkingConfigDict(TypedDict, total=False): + """Specifies the size and overlap of chunks for RagFiles.""" + + chunk_overlap: Optional[int] + """The overlap between chunks.""" + + chunk_size: Optional[int] + """The size of the chunks.""" + + fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunkingDict] + """Specifies the fixed length chunking config.""" + + +RagFileChunkingConfigOrDict = Union[RagFileChunkingConfig, RagFileChunkingConfigDict] + + +class RagFileMetadataConfig(_common.BaseModel): + """Metadata config for RagFile.""" + + gcs_metadata_schema_source: Optional[GcsSource] = Field( + default=None, + description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", + ) + gcs_metadata_source: Optional[GcsSource] = Field( + default=None, + description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", + ) + google_drive_metadata_schema_source: Optional[GoogleDriveSource] = Field( + default=None, + description="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", + ) + google_drive_metadata_source: Optional[GoogleDriveSource] = Field( + default=None, + description="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", + ) + inline_metadata_schema_source: Optional[str] = Field( + default=None, + description="""Inline metadata schema source. Must be a JSON string.""", + ) + inline_metadata_source: Optional[str] = Field( + default=None, description="""Inline metadata source. Must be a JSON string.""" + ) + + +class RagFileMetadataConfigDict(TypedDict, total=False): + """Metadata config for RagFile.""" + + gcs_metadata_schema_source: Optional[GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" + + gcs_metadata_source: Optional[GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" + + google_drive_metadata_schema_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" + + google_drive_metadata_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" + + inline_metadata_schema_source: Optional[str] + """Inline metadata schema source. Must be a JSON string.""" + + inline_metadata_source: Optional[str] + """Inline metadata source. Must be a JSON string.""" + + +RagFileMetadataConfigOrDict = Union[RagFileMetadataConfig, RagFileMetadataConfigDict] + + +class RagFileParsingConfigAdvancedParser(_common.BaseModel): + """Specifies the advanced parsing for RagFiles.""" + + use_advanced_pdf_parsing: Optional[bool] = Field( + default=None, description="""Whether to use advanced PDF parsing.""" + ) + + +class RagFileParsingConfigAdvancedParserDict(TypedDict, total=False): + """Specifies the advanced parsing for RagFiles.""" + + use_advanced_pdf_parsing: Optional[bool] + """Whether to use advanced PDF parsing.""" + + +RagFileParsingConfigAdvancedParserOrDict = Union[ + RagFileParsingConfigAdvancedParser, RagFileParsingConfigAdvancedParserDict +] + + +class RagFileParsingConfigLayoutParser(_common.BaseModel): + """Document AI Layout Parser config.""" + + global_max_parsing_requests_per_min: Optional[int] = Field( + default=None, + description="""The maximum number of requests the job is allowed to make to the Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""", + ) + max_parsing_requests_per_min: Optional[int] = Field( + default=None, + description="""The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""", + ) + processor_name: Optional[str] = Field( + default=None, + description="""The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""", + ) + + +class RagFileParsingConfigLayoutParserDict(TypedDict, total=False): + """Document AI Layout Parser config.""" + + global_max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""" + + max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""" + + processor_name: Optional[str] + """The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""" + + +RagFileParsingConfigLayoutParserOrDict = Union[ + RagFileParsingConfigLayoutParser, RagFileParsingConfigLayoutParserDict +] + + +class RagFileParsingConfig(_common.BaseModel): + """Specifies the parsing config for RagFiles.""" + + advanced_parser: Optional[RagFileParsingConfigAdvancedParser] = Field( + default=None, description="""The Advanced Parser to use for RagFiles.""" + ) + layout_parser: Optional[RagFileParsingConfigLayoutParser] = Field( + default=None, description="""The Layout Parser to use for RagFiles.""" + ) + llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( + default=None, description="""The LLM Parser to use for RagFiles.""" + ) + use_advanced_pdf_parsing: Optional[bool] = Field( + default=None, description="""Whether to use advanced PDF parsing.""" + ) + + +class RagFileParsingConfigDict(TypedDict, total=False): + """Specifies the parsing config for RagFiles.""" + + advanced_parser: Optional[RagFileParsingConfigAdvancedParserDict] + """The Advanced Parser to use for RagFiles.""" + + layout_parser: Optional[RagFileParsingConfigLayoutParserDict] + """The Layout Parser to use for RagFiles.""" + + llm_parser: Optional[RagFileParsingConfigLlmParserDict] + """The LLM Parser to use for RagFiles.""" + + use_advanced_pdf_parsing: Optional[bool] + """Whether to use advanced PDF parsing.""" + + +RagFileParsingConfigOrDict = Union[RagFileParsingConfig, RagFileParsingConfigDict] + + +class RagFileTransformationConfig(_common.BaseModel): + """Specifies the transformation config for RagFiles.""" + + rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( + default=None, description="""Specifies the chunking config for RagFiles.""" + ) + + +class RagFileTransformationConfigDict(TypedDict, total=False): + """Specifies the transformation config for RagFiles.""" + + rag_file_chunking_config: Optional[RagFileChunkingConfigDict] + """Specifies the chunking config for RagFiles.""" + + +RagFileTransformationConfigOrDict = Union[ + RagFileTransformationConfig, RagFileTransformationConfigDict +] + + +class ImportRagFilesConfig(_common.BaseModel): + """Config for importing RagFiles.""" + + gcs_source: Optional[GcsSource] = Field( + default=None, + description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""", + ) + global_max_embedding_requests_per_min: Optional[int] = Field( + default=None, + description="""Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""", + ) + google_drive_source: Optional[GoogleDriveSource] = Field( + default=None, + description="""Google Drive location. Supports importing individual files as well as Google Drive folders.""", + ) + import_result_bigquery_sink: Optional[BigQueryDestination] = Field( + default=None, + description="""The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""", + ) + import_result_gcs_sink: Optional[GcsDestination] = Field( + default=None, + description="""The Cloud Storage path to write import result to.""", + ) + jira_source: Optional[JiraSource] = Field( + default=None, + description="""Jira queries with their corresponding authentication.""", + ) + max_embedding_requests_per_min: Optional[int] = Field( + default=None, + description="""Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""", + ) + partial_failure_bigquery_sink: Optional[BigQueryDestination] = Field( + default=None, + description="""The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""", + ) + partial_failure_gcs_sink: Optional[GcsDestination] = Field( + default=None, + description="""The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""", + ) + rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( + default=None, + description="""Specifies the size and overlap of chunks after importing RagFiles.""", + ) + rag_file_metadata_config: Optional[RagFileMetadataConfig] = Field( + default=None, + description="""Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""", + ) + rag_file_parsing_config: Optional[RagFileParsingConfig] = Field( + default=None, + description="""Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""", + ) + rag_file_transformation_config: Optional[RagFileTransformationConfig] = Field( + default=None, + description="""Specifies the transformation config for RagFiles.""", + ) + rebuild_ann_index: Optional[bool] = Field( + default=None, + description="""Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""", + ) + share_point_sources: Optional[SharePointSources] = Field( + default=None, description="""SharePoint sources.""" + ) + slack_source: Optional[SlackSource] = Field( + default=None, + description="""Slack channels with their corresponding access tokens.""", + ) + + +class ImportRagFilesConfigDict(TypedDict, total=False): + """Config for importing RagFiles.""" + + gcs_source: Optional[GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""" + + global_max_embedding_requests_per_min: Optional[int] + """Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""" + + google_drive_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders.""" + + import_result_bigquery_sink: Optional[BigQueryDestinationDict] + """The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""" + + import_result_gcs_sink: Optional[GcsDestinationDict] + """The Cloud Storage path to write import result to.""" + + jira_source: Optional[JiraSourceDict] + """Jira queries with their corresponding authentication.""" + + max_embedding_requests_per_min: Optional[int] + """Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""" + + partial_failure_bigquery_sink: Optional[BigQueryDestinationDict] + """The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""" + + partial_failure_gcs_sink: Optional[GcsDestinationDict] + """The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""" + + rag_file_chunking_config: Optional[RagFileChunkingConfigDict] + """Specifies the size and overlap of chunks after importing RagFiles.""" + + rag_file_metadata_config: Optional[RagFileMetadataConfigDict] + """Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""" + + rag_file_parsing_config: Optional[RagFileParsingConfigDict] + """Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""" + + rag_file_transformation_config: Optional[RagFileTransformationConfigDict] + """Specifies the transformation config for RagFiles.""" + + rebuild_ann_index: Optional[bool] + """Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""" + + share_point_sources: Optional[SharePointSourcesDict] + """SharePoint sources.""" + + slack_source: Optional[SlackSourceDict] + """Slack channels with their corresponding access tokens.""" + + +ImportRagFilesConfigOrDict = Union[ImportRagFilesConfig, ImportRagFilesConfigDict] + + +class ImportRagFilesRequest(_common.BaseModel): + """Request message for VertexRagDataService.ImportRagFiles.""" + + import_rag_files_config: Optional[ImportRagFilesConfig] = Field( + default=None, + description="""Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""", + ) + + +class ImportRagFilesRequestDict(TypedDict, total=False): + """Request message for VertexRagDataService.ImportRagFiles.""" + + import_rag_files_config: Optional[ImportRagFilesConfigDict] + """Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""" + + +ImportRagFilesRequestOrDict = Union[ImportRagFilesRequest, ImportRagFilesRequestDict] + + +class _ImportRagFilesRequestParameters(_common.BaseModel): + """Parameters for importing RAG files.""" + + import_rag_files_request: Optional[ImportRagFilesRequest] = Field( + default=None, description="""""" + ) + config: Optional[ImportRagFilesRequestConfig] = Field( + default=None, description="""""" + ) + name: Optional[str] = Field(default=None, description="""""") + + +class _ImportRagFilesRequestParametersDict(TypedDict, total=False): + """Parameters for importing RAG files.""" + + import_rag_files_request: Optional[ImportRagFilesRequestDict] + """""" + + config: Optional[ImportRagFilesRequestConfigDict] + """""" + + name: Optional[str] + """""" + + +_ImportRagFilesRequestParametersOrDict = Union[ + _ImportRagFilesRequestParameters, _ImportRagFilesRequestParametersDict +] + + +class ImportRagFilesResponse(_common.BaseModel): + """Response from importing RAG files.""" + + imported_rag_files_count: Optional[int] = Field( + default=None, + description="""The number of RagFiles imported into the RagCorpus.""", + ) + failed_rag_files_count: Optional[int] = Field( + default=None, + description="""The number of RagFiles that failed to import into the RagCorpus.""", + ) + skipped_rag_files_count: Optional[int] = Field( + default=None, + description="""The number of RagFiles skipped while importing into the RagCorpus.""", + ) + partial_failures_gcs_path: Optional[str] = Field( + default=None, + description="""The Google Cloud Storage path with the partial failures.""", + ) + partial_failures_bigquery_table: Optional[str] = Field( + default=None, + description="""The BigQuery table where the partial failures were written.""", + ) + + +class ImportRagFilesResponseDict(TypedDict, total=False): + """Response from importing RAG files.""" + + imported_rag_files_count: Optional[int] + """The number of RagFiles imported into the RagCorpus.""" + + failed_rag_files_count: Optional[int] + """The number of RagFiles that failed to import into the RagCorpus.""" + + skipped_rag_files_count: Optional[int] + """The number of RagFiles skipped while importing into the RagCorpus.""" + + partial_failures_gcs_path: Optional[str] + """The Google Cloud Storage path with the partial failures.""" + + partial_failures_bigquery_table: Optional[str] + """The BigQuery table where the partial failures were written.""" + + +ImportRagFilesResponseOrDict = Union[ImportRagFilesResponse, ImportRagFilesResponseDict] + + +class ImportRagFilesOperation(_common.BaseModel): + """Operation for importing RAG files.""" + + 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.""", + ) + response: Optional[ImportRagFilesResponse] = Field( + default=None, description="""The response from the import operation.""" + ) + + +class ImportRagFilesOperationDict(TypedDict, total=False): + """Operation for importing RAG files.""" + + 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.""" + + response: Optional[ImportRagFilesResponseDict] + """The response from the import operation.""" + + +ImportRagFilesOperationOrDict = Union[ + ImportRagFilesOperation, ImportRagFilesOperationDict +] + + +class GetImportFilesOperationConfig(_common.BaseModel): + """Config for getting an import files operation.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetImportFilesOperationConfigDict(TypedDict, total=False): + """Config for getting an import files operation.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +GetImportFilesOperationConfigOrDict = Union[ + GetImportFilesOperationConfig, GetImportFilesOperationConfigDict +] + + +class _GetImportFilesOperationParameters(_common.BaseModel): + """Parameters for getting an import files operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetImportFilesOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetImportFilesOperationParametersDict(TypedDict, total=False): + """Parameters for getting an import files operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetImportFilesOperationConfigDict] + """Used to override the default configuration.""" + + +_GetImportFilesOperationParametersOrDict = Union[ + _GetImportFilesOperationParameters, _GetImportFilesOperationParametersDict +] + + class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): """Config for getting an Agent Engine Runtime Revision.""" diff --git a/tests/unit/agentplatform/genai/replays/test_import_rag_files.py b/tests/unit/agentplatform/genai/replays/test_import_rag_files.py new file mode 100644 index 0000000000..94e06b8003 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_import_rag_files.py @@ -0,0 +1,77 @@ +# 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 import files 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_rag_import_private(client): + + response = client.rag._import_files( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/3671559596213796864", + import_rag_files_request=types.ImportRagFilesRequest( + import_rag_files_config=types.ImportRagFilesConfig( + gcs_source=types.GcsSource( + uris=["gs://sararob_test/test-rag-file.pdf"] + ), + ), + ), + ) + + assert isinstance(response, types.ImportRagFilesOperation) + + +def test_rag_import(client): + + gcs_filepath = "gs://sararob_test/test-rag-file.pdf" + + import_file_response = client.rag.import_files( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/6301661778598166528", + import_config=types.ImportRagFilesConfig( + gcs_source=types.GcsSource(uris=[gcs_filepath]), + ), + ) + + assert isinstance(import_file_response, types.ImportRagFilesResponse) + + assert import_file_response.imported_rag_files_count == 1 + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_import_files_async(client): + + gcs_filepath = "gs://sararob_test/test-rag-file.pdf" + + import_file_response = await client.aio.rag.import_files( + name="projects/vertex-sdk-dev/locations/us-central1/ragCorpora/3671559596213796864", + import_config=types.ImportRagFilesConfig( + gcs_source=types.GcsSource(uris=[gcs_filepath]), + ), + ) + + assert isinstance(import_file_response, types.ImportRagFilesResponse) + # This file is already imported in this corpus, so it should be skipped. + assert import_file_response.skipped_rag_files_count == 1