feat: add support for Document AI Layout Parser as a RAG import option

PiperOrigin-RevId: 702358388

Author: ilai-deutel

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -22,6 +22,7 @@
  EmbeddingModelConfig,
  Filter,
  HybridSearch,
+ LayoutParserConfig,
  LlmRanker,
  Pinecone,
  RagCorpus,
@@ -612,6 +613,16 @@
  ],
 )
 
+TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG = LayoutParserConfig(
+ processor_name="projects/test-project/locations/us/processors/abc123",
+ max_parsing_requests_per_min=100,
+)
+
+TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG = LayoutParserConfig(
+ processor_name="projects/test-project/locations/us/processors/abc123/processorVersions/pretrained-layout-parser-v0.0-2020-01-0",
+ max_parsing_requests_per_min=100,
+)
+
 TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE_NO_FOLDERS = ImportRagFilesConfig(
  rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
  share_point_sources=GapicSharePointSources(
@@ -634,6 +645,38 @@
  import_rag_files_config=TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE,
 )
 
+TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_PATH = ImportRagFilesConfig(
+ TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER
+)
+TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_PATH.rag_file_parsing_config = (
+ RagFileParsingConfig(
+ layout_parser=RagFileParsingConfig.LayoutParser(
+ processor_name="projects/test-project/locations/us/processors/abc123",
+ max_parsing_requests_per_min=100,
+ )
+ )
+)
+
+TEST_IMPORT_REQUEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH = ImportRagFilesRequest(
+ parent=TEST_RAG_CORPUS_RESOURCE_NAME,
+ import_rag_files_config=TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_PATH,
+)
+
+TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH = (
+ ImportRagFilesConfig(TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER)
+)
+TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH.rag_file_parsing_config = RagFileParsingConfig(
+ layout_parser=RagFileParsingConfig.LayoutParser(
+ processor_name="projects/test-project/locations/us/processors/abc123/processorVersions/pretrained-layout-parser-v0.0-2020-01-0",
+ max_parsing_requests_per_min=100,
+ )
+)
+
+TEST_IMPORT_REQUEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH = ImportRagFilesRequest(
+ parent=TEST_RAG_CORPUS_RESOURCE_NAME,
+ import_rag_files_config=TEST_IMPORT_FILES_CONFIG_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH,
+)
+
 # Retrieval
 TEST_QUERY_TEXT = "What happen to the fox and the dog?"
 TEST_CONTEXTS = RagContexts(

tests/unit/vertex_rag/test_rag_data_preview.py

@@ -1088,6 +1088,74 @@ def test_prepare_import_files_request_sharepoint_source_no_folders(self):
  test_rag_constants_preview.TEST_IMPORT_REQUEST_SHARE_POINT_SOURCE_NO_FOLDERS,
  )
 
+ def test_prepare_import_files_request_valid_layout_parser_with_processor_path(self):
+ request = prepare_import_files_request(
+ corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants_preview.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ layout_parser=test_rag_constants_preview.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ )
+ import_files_request_eq(
+ request,
+ test_rag_constants_preview.TEST_IMPORT_REQUEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH,
+ )
+
+ def test_prepare_import_files_request_valid_layout_parser_with_processor_version_path(
+ self,
+ ):
+ request = prepare_import_files_request(
+ corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants_preview.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ layout_parser=test_rag_constants_preview.TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG,
+ )
+ import_files_request_eq(
+ request,
+ test_rag_constants_preview.TEST_IMPORT_REQUEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH,
+ )
+
+ def test_prepare_import_files_request_invalid_layout_parser_name(self):
+ layout_parser = rag.LayoutParserConfig(
+ processor_name="projects/test-project/locations/us/processorTypes/LAYOUT_PARSER",
+ )
+ with pytest.raises(ValueError) as e:
+ prepare_import_files_request(
+ corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants_preview.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ layout_parser=layout_parser,
+ )
+ e.match("processor_name must be of the format")
+
+ def test_advanced_pdf_parsing_and_layout_parser_both_set_error(self):
+ with pytest.raises(ValueError) as e:
+ rag.import_files(
+ corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants_preview.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ use_advanced_pdf_parsing=True,
+ layout_parser=test_rag_constants_preview.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ )
+ e.match(
+ "Only one of use_advanced_pdf_parsing or layout_parser may be "
+ "passed in at a time"
+ )
+
+ @pytest.mark.asyncio
+ async def test_advanced_pdf_parsing_and_layout_parser_both_set_error_async(self):
+ with pytest.raises(ValueError) as e:
+ await rag.import_files_async(
+ corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants_preview.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ use_advanced_pdf_parsing=True,
+ layout_parser=test_rag_constants_preview.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ )
+ e.match(
+ "Only one of use_advanced_pdf_parsing or layout_parser may be "
+ "passed in at a time"
+ )
+
  def test_set_embedding_model_config_set_both_error(self):
  embedding_model_config = rag.EmbeddingModelConfig(
  publisher_model="whatever",

vertexai/preview/rag/__init__.py

@@ -40,6 +40,7 @@
  HybridSearch,
  JiraQuery,
  JiraSource,
+ LayoutParserConfig,
  LlmRanker,
  Pinecone,
  RagCorpus,
@@ -70,6 +71,7 @@
  "HybridSearch",
  "JiraQuery",
  "JiraSource",
+ "LayoutParserConfig",
  "LlmRanker",
  "Pinecone",
  "RagCorpus",

vertexai/preview/rag/rag_data.py

@@ -45,6 +45,7 @@
 from vertexai.preview.rag.utils.resources import (
  EmbeddingModelConfig,
  JiraSource,
+ LayoutParserConfig,
  Pinecone,
  RagCorpus,
  RagFile,
@@ -466,6 +467,7 @@ def import_files(
  max_embedding_requests_per_min: int = 1000,
  use_advanced_pdf_parsing: Optional[bool] = False,
  partial_failures_sink: Optional[str] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
 ) -> ImportRagFilesResponse:
  """
  Import files to an existing RagCorpus, wait until completion.
@@ -581,13 +583,21 @@ def import_files(
  exist - if it does not exist, it will be created. If it does exist,
  the schema will be checked and the partial failures will be appended
  to the table.
+ layout_parser: Configuration for the Document AI Layout Parser Processor
+ to use for document parsing. Optional.
+ If not None,`use_advanced_pdf_parsing` must be False.
  Returns:
  ImportRagFilesResponse.
  """
  if source is not None and paths is not None:
  raise ValueError("Only one of source or paths must be passed in at a time")
  if source is None and paths is None:
  raise ValueError("One of source or paths must be passed in")
+ if use_advanced_pdf_parsing and layout_parser is not None:
+ raise ValueError(
+ "Only one of use_advanced_pdf_parsing or layout_parser may be "
+ "passed in at a time"
+ )
  corpus_name = _gapic_utils.get_corpus_name(corpus_name)
  request = _gapic_utils.prepare_import_files_request(
  corpus_name=corpus_name,
@@ -599,6 +609,7 @@ def import_files(
  max_embedding_requests_per_min=max_embedding_requests_per_min,
  use_advanced_pdf_parsing=use_advanced_pdf_parsing,
  partial_failures_sink=partial_failures_sink,
+ layout_parser=layout_parser,
  )
  client = _gapic_utils.create_rag_data_service_client()
  try:
@@ -619,6 +630,7 @@ async def import_files_async(
  max_embedding_requests_per_min: int = 1000,
  use_advanced_pdf_parsing: Optional[bool] = False,
  partial_failures_sink: Optional[str] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
 ) -> operation_async.AsyncOperation:
  """
  Import files to an existing RagCorpus asynchronously.
@@ -734,13 +746,21 @@ async def import_files_async(
  exist - if it does not exist, it will be created. If it does exist,
  the schema will be checked and the partial failures will be appended
  to the table.
+ layout_parser: Configuration for the Document AI Layout Parser Processor
+ to use for document parsing. Optional.
+ If not None,`use_advanced_pdf_parsing` must be False.
  Returns:
  operation_async.AsyncOperation.
  """
  if source is not None and paths is not None:
  raise ValueError("Only one of source or paths must be passed in at a time")
  if source is None and paths is None:
  raise ValueError("One of source or paths must be passed in")
+ if use_advanced_pdf_parsing and layout_parser is not None:
+ raise ValueError(
+ "Only one of use_advanced_pdf_parsing or layout_parser may be "
+ "passed in at a time"
+ )
  corpus_name = _gapic_utils.get_corpus_name(corpus_name)
  request = _gapic_utils.prepare_import_files_request(
  corpus_name=corpus_name,
@@ -752,6 +772,7 @@ async def import_files_async(
  max_embedding_requests_per_min=max_embedding_requests_per_min,
  use_advanced_pdf_parsing=use_advanced_pdf_parsing,
  partial_failures_sink=partial_failures_sink,
+ layout_parser=layout_parser,
  )
  async_client = _gapic_utils.create_rag_data_service_async_client()
  try:

vertexai/preview/rag/utils/_gapic_utils.py

@@ -42,6 +42,7 @@
 from vertexai.preview.rag.utils.resources import (
  EmbeddingModelConfig,
  VertexPredictionEndpoint,
+ LayoutParserConfig,
  Pinecone,
  RagCorpus,
  RagFile,
@@ -60,6 +61,9 @@
 
 
 _VALID_RESOURCE_NAME_REGEX = "[a-z][a-zA-Z0-9._-]{0,127}"
+_VALID_DOCUMENT_AI_PROCESSOR_NAME_REGEX = (
+ r"projects/[^/]+/locations/[^/]+/processors/[^/]+(?:/processorVersions/[^/]+)?"
+)
 
 
 def create_rag_data_service_client():
@@ -445,6 +449,7 @@ def prepare_import_files_request(
  max_embedding_requests_per_min: int = 1000,
  use_advanced_pdf_parsing: bool = False,
  partial_failures_sink: Optional[str] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
 ) -> ImportRagFilesRequest:
  if len(corpus_name.split("/")) != 6:
  raise ValueError(
@@ -456,6 +461,24 @@ def prepare_import_files_request(
  use_advanced_pdf_parsing=use_advanced_pdf_parsing,
  ),
  )
+ if layout_parser is not None:
+ if (
+ re.fullmatch(
+ _VALID_DOCUMENT_AI_PROCESSOR_NAME_REGEX, layout_parser.processor_name
+ )
+ is None
+ ):
+ raise ValueError(
+ "processor_name must be of the format "
+ "`projects/{project_id}/locations/{location}/processors/{processor_id}`"
+ "or "
+ "`projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`, "
+ f"got {layout_parser.processor_name!r}"
+ )
+ rag_file_parsing_config.layout_parser = RagFileParsingConfig.LayoutParser(
+ processor_name=layout_parser.processor_name,
+ max_parsing_requests_per_min=layout_parser.max_parsing_requests_per_min,
+ )
  local_chunk_size = chunk_size
  local_chunk_overlap = chunk_overlap
  if transformation_config and transformation_config.chunking_config:

vertexai/preview/rag/utils/resources.py

@@ -477,3 +477,26 @@ class TransformationConfig:
  """
 
  chunking_config: Optional[ChunkingConfig] = None
+
+
+@dataclasses.dataclass
+class LayoutParserConfig:
+ """Configuration for the Document AI Layout Parser Processor.
+
+ Attributes:
+ processor_name (str):
+ The full resource name of a Document AI processor or processor
+ version. The processor must have type `LAYOUT_PARSER_PROCESSOR`.
+ Format:
+ - `projects/{project_id}/locations/{location}/processors/{processor_id}`
+ - `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`
+ max_parsing_requests_per_min (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 will be used.
+ """
+
+ processor_name: str
+ max_parsing_requests_per_min: Optional[int] = None