fix: Propagating import result sink correctly in the vertexai sdk.

PiperOrigin-RevId: 741317609

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -209,6 +209,8 @@
  ),
  ),
 )
+TEST_IMPORT_RESULT_GCS_SINK = "gs://test-bucket/test-object.ndjson"
+TEST_IMPORT_RESULT_BIGQUERY_SINK = "bq://test-project.test_dataset.test_table"
 # GCS
 TEST_IMPORT_FILES_CONFIG_GCS = ImportRagFilesConfig(
  rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,

tests/unit/vertex_rag/test_rag_data.py

@@ -276,6 +276,14 @@ def import_files_request_eq(returned_request, expected_request):
  returned_request.import_rag_files_config.rag_file_transformation_config
  == expected_request.import_rag_files_config.rag_file_transformation_config
  )
+ assert (
+ returned_request.import_rag_files_config.import_result_gcs_sink
+ == expected_request.import_rag_files_config.import_result_gcs_sink
+ )
+ assert (
+ returned_request.import_rag_files_config.import_result_bigquery_sink
+ == expected_request.import_rag_files_config.import_result_bigquery_sink
+ )
 
 
 @pytest.mark.usefixtures("google_auth_mock")
@@ -517,6 +525,26 @@ def test_import_files(self, import_files_mock):
 
  assert response.imported_rag_files_count == 2
 
+ def test_import_files_with_import_result_gcs_sink(self, import_files_mock):
+ response = rag.import_files(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_GCS_PATH],
+ import_result_sink=test_rag_constants.TEST_IMPORT_RESULT_GCS_SINK,
+ )
+ import_files_mock.assert_called_once()
+
+ assert response.imported_rag_files_count == 2
+
+ def test_import_files_with_import_result_bigquery_sink(self, import_files_mock):
+ response = rag.import_files(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_GCS_PATH],
+ import_result_sink=test_rag_constants.TEST_IMPORT_RESULT_BIGQUERY_SINK,
+ )
+ import_files_mock.assert_called_once()
+
+ assert response.imported_rag_files_count == 2
+
  @pytest.mark.usefixtures("rag_data_client_mock_exception")
  def test_import_files_failure(self):
  with pytest.raises(RuntimeError) as e:
@@ -536,6 +564,32 @@ async def test_import_files_async(self, import_files_async_mock):
 
  assert response.result().imported_rag_files_count == 2
 
+ @pytest.mark.asyncio
+ async def test_import_files_with_import_result_gcs_sink_async(
+ self, import_files_async_mock
+ ):
+ response = await rag.import_files_async(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_GCS_PATH],
+ import_result_sink=test_rag_constants.TEST_IMPORT_RESULT_GCS_SINK,
+ )
+ import_files_async_mock.assert_called_once()
+
+ assert response.result().imported_rag_files_count == 2
+
+ @pytest.mark.asyncio
+ async def test_import_files_with_import_result_bigquery_sink_async(
+ self, import_files_async_mock
+ ):
+ response = await rag.import_files_async(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_GCS_PATH],
+ import_result_sink=test_rag_constants.TEST_IMPORT_RESULT_BIGQUERY_SINK,
+ )
+ import_files_async_mock.assert_called_once()
+
+ assert response.result().imported_rag_files_count == 2
+
  @pytest.mark.asyncio
  @pytest.mark.usefixtures("rag_data_async_client_mock_exception")
  async def test_import_files_async_failure(self):

vertexai/rag/rag_data.py

@@ -395,6 +395,7 @@ def import_files(
  transformation_config: Optional[TransformationConfig] = None,
  timeout: int = 600,
  max_embedding_requests_per_min: int = 1000,
+ import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
  parser: Optional[LayoutParserConfig] = None,
 ) -> ImportRagFilesResponse:
@@ -509,8 +510,17 @@ def import_files(
  here. If unspecified, a default value of 1,000
  QPM would be used.
  timeout: Default is 600 seconds.
- partial_failures_sink: Either a GCS path to store partial failures or a
- BigQuery table to store partial failures. The format is
+ import_result_sink: Either a GCS path to store import results or a
+ BigQuery table to store import results. The format is
+ "gs://my-bucket/my/object.ndjson" for GCS or
+ "bq://my-project.my-dataset.my-table" for BigQuery. An existing GCS
+ object cannot be used. However, the BigQuery table may or may not
+ exist - if it does not exist, it will be created. If it does exist,
+ the schema will be checked and the import results will be appended
+ to the table.
+ partial_failures_sink: Deprecated. Prefer to use `import_result_sink`.
+ Either a GCS path to store partial failures or a BigQuery table to
+ store partial failures. The format is
  "gs://my-bucket/my/object.ndjson" for GCS or
  "bq://my-project.my-dataset.my-table" for BigQuery. An existing GCS
  object cannot be used. However, the BigQuery table may or may not
@@ -534,6 +544,7 @@ def import_files(
  source=source,
  transformation_config=transformation_config,
  max_embedding_requests_per_min=max_embedding_requests_per_min,
+ import_result_sink=import_result_sink,
  partial_failures_sink=partial_failures_sink,
  parser=parser,
  )
@@ -552,6 +563,7 @@ async def import_files_async(
  source: Optional[Union[SlackChannelsSource, JiraSource, SharePointSources]] = None,
  transformation_config: Optional[TransformationConfig] = None,
  max_embedding_requests_per_min: int = 1000,
+ import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
  parser: Optional[LayoutParserConfig] = None,
 ) -> operation_async.AsyncOperation:
@@ -666,8 +678,17 @@ async def import_files_async(
  page on the project to set an appropriate value
  here. If unspecified, a default value of 1,000
  QPM would be used.
- partial_failures_sink: Either a GCS path to store partial failures or a
- BigQuery table to store partial failures. The format is
+ import_result_sink: Either a GCS path to store import results or a
+ BigQuery table to store import results. The format is
+ "gs://my-bucket/my/object.ndjson" for GCS or
+ "bq://my-project.my-dataset.my-table" for BigQuery. An existing GCS
+ object cannot be used. However, the BigQuery table may or may not
+ exist - if it does not exist, it will be created. If it does exist,
+ the schema will be checked and the import results will be appended
+ to the table.
+ partial_failures_sink: Deprecated. Prefer to use `import_result_sink`.
+ Either a GCS path to store partial failures or a BigQuery table to
+ store partial failures. The format is
  "gs://my-bucket/my/object.ndjson" for GCS or
  "bq://my-project.my-dataset.my-table" for BigQuery. An existing GCS
  object cannot be used. However, the BigQuery table may or may not
@@ -691,6 +712,7 @@ async def import_files_async(
  source=source,
  transformation_config=transformation_config,
  max_embedding_requests_per_min=max_embedding_requests_per_min,
+ import_result_sink=import_result_sink,
  partial_failures_sink=partial_failures_sink,
  parser=parser,
  )

vertexai/rag/utils/_gapic_utils.py

@@ -360,6 +360,7 @@ def prepare_import_files_request(
  source: Optional[Union[SlackChannelsSource, JiraSource, SharePointSources]] = None,
  transformation_config: Optional[TransformationConfig] = None,
  max_embedding_requests_per_min: int = 1000,
+ import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
  parser: Optional[LayoutParserConfig] = None,
 ) -> ImportRagFilesRequest:
@@ -407,6 +408,22 @@ def prepare_import_files_request(
  max_embedding_requests_per_min=max_embedding_requests_per_min,
  )
 
+ import_result_sink = import_result_sink or partial_failures_sink
+
+ if import_result_sink is not None:
+ if import_result_sink.startswith("gs://"):
+ import_rag_files_config.partial_failure_gcs_sink.output_uri_prefix = (
+ import_result_sink
+ )
+ elif import_result_sink.startswith("bq://"):
+ import_rag_files_config.partial_failure_bigquery_sink.output_uri = (
+ import_result_sink
+ )
+ else:
+ raise ValueError(
+ "import_result_sink must be a GCS path or a BigQuery table."
+ )
+
  if source is not None:
  gapic_source = convert_source_for_rag_import(source)
  if isinstance(gapic_source, GapicSlackSource):