feat: Update v1 sdk to support llmparser in import file functions

PiperOrigin-RevId: 756617457

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -21,6 +21,7 @@
 from vertexai.rag import (
  Filter,
  LayoutParserConfig,
+ LlmParserConfig,
  LlmRanker,
  Pinecone,
  RagCorpus,
@@ -629,3 +630,26 @@
  llm_ranker=LlmRanker(model_name="test-llm-ranker"),
  ),
 )
+TEST_LLM_PARSER_CONFIG = LlmParserConfig(
+ model_name="gemini-1.5-pro-002",
+ max_parsing_requests_per_min=500,
+ custom_parsing_prompt="test-custom-parsing-prompt",
+)
+
+
+TEST_IMPORT_FILES_CONFIG_LLM_PARSER = ImportRagFilesConfig(
+ TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER
+)
+
+TEST_IMPORT_FILES_CONFIG_LLM_PARSER.rag_file_parsing_config = RagFileParsingConfig(
+ llm_parser=RagFileParsingConfig.LlmParser(
+ model_name="gemini-1.5-pro-002",
+ max_parsing_requests_per_min=500,
+ custom_parsing_prompt="test-custom-parsing-prompt",
+ )
+)
+
+TEST_IMPORT_REQUEST_LLM_PARSER = ImportRagFilesRequest(
+ parent=TEST_RAG_CORPUS_RESOURCE_NAME,
+ import_rag_files_config=TEST_IMPORT_FILES_CONFIG_LLM_PARSER,
+)

tests/unit/vertex_rag/test_rag_data.py

@@ -931,7 +931,7 @@ def test_prepare_import_files_request_valid_layout_parser_with_processor_path(se
  corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
  paths=[test_rag_constants.TEST_DRIVE_FOLDER],
  transformation_config=create_transformation_config(),
- parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ layout_parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
  )
  import_files_request_eq(
  request,
@@ -945,7 +945,7 @@ def test_prepare_import_files_request_valid_layout_parser_with_processor_version
  corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
  paths=[test_rag_constants.TEST_DRIVE_FOLDER],
  transformation_config=create_transformation_config(),
- parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG,
+ layout_parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG,
  )
  import_files_request_eq(
  request,
@@ -961,10 +961,45 @@ def test_prepare_import_files_request_invalid_layout_parser_name(self):
  corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
  paths=[test_rag_constants.TEST_DRIVE_FOLDER],
  transformation_config=create_transformation_config(),
- parser=layout_parser,
+ layout_parser=layout_parser,
  )
  e.match("processor_name must be of the format")
 
+ def test_prepare_import_files_request_llm_parser(self):
+ request = prepare_import_files_request(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ llm_parser=test_rag_constants.TEST_LLM_PARSER_CONFIG,
+ )
+ import_files_request_eq(
+ request,
+ test_rag_constants.TEST_IMPORT_REQUEST_LLM_PARSER,
+ )
+
+ def test_layout_parser_and_llm_parser_both_set_error(self):
+ with pytest.raises(ValueError) as e:
+ rag.import_files(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ layout_parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ llm_parser=test_rag_constants.TEST_LLM_PARSER_CONFIG,
+ )
+ e.match("Only one of layout_parser or llm_parser may be passed in at a time")
+
+ @pytest.mark.asyncio
+ async def test_layout_parser_and_llm_parser_both_set_error_async(self):
+ with pytest.raises(ValueError) as e:
+ await rag.import_files_async(
+ corpus_name=test_rag_constants.TEST_RAG_CORPUS_RESOURCE_NAME,
+ paths=[test_rag_constants.TEST_DRIVE_FOLDER],
+ transformation_config=create_transformation_config(),
+ layout_parser=test_rag_constants.TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG,
+ llm_parser=test_rag_constants.TEST_LLM_PARSER_CONFIG,
+ )
+ e.match("Only one of layout_parser or llm_parser may be passed in at a time")
+
  def test_set_embedding_model_config_set_both_error(self):
  embedding_model_config = rag.RagEmbeddingModelConfig(
  vertex_prediction_endpoint=rag.VertexPredictionEndpoint(

vertexai/rag/__init__.py

@@ -43,6 +43,7 @@
  JiraQuery,
  JiraSource,
  LayoutParserConfig,
+ LlmParserConfig,
  LlmRanker,
  Pinecone,
  RagCorpus,
@@ -71,6 +72,7 @@
  "JiraQuery",
  "JiraSource",
  "LayoutParserConfig",
+ "LlmParserConfig",
  "LlmRanker",
  "Pinecone",
  "RagCorpus",

vertexai/rag/rag_data.py

@@ -45,6 +45,7 @@
 from vertexai.rag.utils.resources import (
  JiraSource,
  LayoutParserConfig,
+ LlmParserConfig,
  RagCorpus,
  RagFile,
  RagVectorDbConfig,
@@ -433,7 +434,8 @@ def import_files(
  max_embedding_requests_per_min: int = 1000,
  import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
- parser: Optional[LayoutParserConfig] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
+ llm_parser: Optional[LlmParserConfig] = None,
 ) -> ImportRagFilesResponse:
  """
  Import files to an existing RagCorpus, wait until completion.
@@ -573,6 +575,10 @@ def import_files(
  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 layout_parser is not None and llm_parser is not None:
+ raise ValueError(
+ "Only one of layout_parser or llm_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,
@@ -582,7 +588,8 @@ def import_files(
  max_embedding_requests_per_min=max_embedding_requests_per_min,
  import_result_sink=import_result_sink,
  partial_failures_sink=partial_failures_sink,
- parser=parser,
+ layout_parser=layout_parser,
+ llm_parser=llm_parser,
  )
  client = _gapic_utils.create_rag_data_service_client()
  try:
@@ -601,7 +608,8 @@ async def import_files_async(
  max_embedding_requests_per_min: int = 1000,
  import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
- parser: Optional[LayoutParserConfig] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
+ llm_parser: Optional[LlmParserConfig] = None,
 ) -> operation_async.AsyncOperation:
  """
  Import files to an existing RagCorpus asynchronously.
@@ -741,6 +749,10 @@ async def import_files_async(
  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 layout_parser is not None and llm_parser is not None:
+ raise ValueError(
+ "Only one of layout_parser or llm_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,
@@ -750,7 +762,8 @@ async def import_files_async(
  max_embedding_requests_per_min=max_embedding_requests_per_min,
  import_result_sink=import_result_sink,
  partial_failures_sink=partial_failures_sink,
- parser=parser,
+ layout_parser=layout_parser,
+ llm_parser=llm_parser,
  )
  async_client = _gapic_utils.create_rag_data_service_async_client()
  try:

vertexai/rag/utils/_gapic_utils.py

@@ -41,6 +41,7 @@
 )
 from vertexai.rag.utils.resources import (
  LayoutParserConfig,
+ LlmParserConfig,
  Pinecone,
  RagCorpus,
  RagEmbeddingModelConfig,
@@ -381,30 +382,45 @@ def prepare_import_files_request(
  max_embedding_requests_per_min: int = 1000,
  import_result_sink: Optional[str] = None,
  partial_failures_sink: Optional[str] = None,
- parser: Optional[LayoutParserConfig] = None,
+ layout_parser: Optional[LayoutParserConfig] = None,
+ llm_parser: Optional[LlmParserConfig] = None,
 ) -> ImportRagFilesRequest:
  if len(corpus_name.split("/")) != 6:
  raise ValueError(
  "corpus_name must be of the format `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`"
  )
 
  rag_file_parsing_config = RagFileParsingConfig()
- if parser is not None:
+ if layout_parser is not None:
  if (
- re.fullmatch(_VALID_DOCUMENT_AI_PROCESSOR_NAME_REGEX, parser.processor_name)
+ 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 {parser.processor_name!r}"
+ "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=parser.processor_name,
- max_parsing_requests_per_min=parser.max_parsing_requests_per_min,
+ processor_name=layout_parser.processor_name,
+ max_parsing_requests_per_min=layout_parser.max_parsing_requests_per_min,
+ )
+ if llm_parser is not None:
+ rag_file_parsing_config.llm_parser = RagFileParsingConfig.LlmParser(
+ model_name=llm_parser.model_name
  )
+ if llm_parser.max_parsing_requests_per_min is not None:
+ rag_file_parsing_config.llm_parser.max_parsing_requests_per_min = (
+ llm_parser.max_parsing_requests_per_min
+ )
+ if llm_parser.custom_parsing_prompt is not None:
+ rag_file_parsing_config.llm_parser.custom_parsing_prompt = (
+ llm_parser.custom_parsing_prompt
+ )
 
  chunk_size = 1024
  chunk_overlap = 200

vertexai/rag/utils/resources.py

@@ -445,3 +445,27 @@ class LayoutParserConfig:
 
  processor_name: str
  max_parsing_requests_per_min: Optional[int] = None
+
+
+@dataclasses.dataclass
+class LlmParserConfig:
+ """Configuration for the Document AI Layout Parser Processor.
+
+ Attributes:
+ model_name (str):
+ The full resource name of a Vertex AI model. Format:
+ - `projects/{project_id}/locations/{location}/publishers/google/models/{model_id}`
+ - `projects/{project_id}/locations/{location}/models/{model_id}`
+ max_parsing_requests_per_min (int):
+ The maximum number of requests the job is allowed to make to the
+ Vertex AI model per minute. Consult
+ https://cloud.google.com/vertex-ai/generative-ai/docs/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.
+ custom_parsing_prompt (str):
+ A custom prompt to use for parsing.
+ """
+
+ model_name: str
+ max_parsing_requests_per_min: Optional[int] = None
+ custom_parsing_prompt: Optional[str] = None