Starting page number

Author: mpolomdeepsense

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -29,12 +29,16 @@
 )
 from unstructured_client.models import shared
 
+# TODO: (Marek Połom) - Update documentation before merging
+
 logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
 
 PARTITION_FORM_FILES_KEY = "files"
 PARTITION_FORM_SPLIT_PDF_PAGE_KEY = "split_pdf_page"
+PARTITION_FORM_STARTING_PAGE_NUMBER_KEY = "starting_page_number"
 
 SUBSTITUTE_FILENAME = "file_for_partition.pdf"
+DEFAULT_STARTING_PAGE_NUMBER = 1
 
 
 FormData = dict[str, Union[str, shared.Files]]
@@ -88,6 +92,10 @@ def before_request(
             Union[requests.PreparedRequest, Exception]: If `splitPdfPage` is set to `true`,
             the last page request; otherwise, the original request.
         """
+        if self.client is None:
+            logger.warning("HTTP client not accessible! Continuing without splitting.")
+            return request
+
         operation_id = hook_ctx.operation_id
         content_type = request.headers.get("Content-Type")
         body = request.body
@@ -104,9 +112,7 @@ def before_request(
         if file is None or not isinstance(file, shared.Files) or not self._is_pdf(file):
             return request
 
-        if self.client is None:
-            logger.warning("HTTP client not accessible! Continuing without splitting.")
-            return request
+        starting_page_number = self._get_starting_page_number(form_data)
 
         pages = self._get_pdf_pages(file.content)
         call_api_partial = functools.partial(
@@ -119,10 +125,12 @@ def before_request(
         self.partition_requests[operation_id] = []
         last_page_content = io.BytesIO()
         with ThreadPoolExecutor(max_workers=call_threads) as executor:
-            for page_content, page_number, all_pages_number in pages:
+            for page_content, page_index, all_pages_number in pages:
+                page_number = page_index + starting_page_number
                 # Check if this page is the last one
-                if page_number == all_pages_number:
+                if page_index == all_pages_number:
                     last_page_content = page_content
+                    last_page_number = page_number
                     break
                 self.partition_requests[operation_id].append(
                     executor.submit(call_api_partial, (page_content, page_number))
@@ -131,7 +139,7 @@ def before_request(
         # `before_request` method needs to return a request so we skip sending the last page in parallel
         # and return that last page at the end of this method
         last_page_request = self._create_request(
-            request, form_data, last_page_content, file.file_name
+            request, form_data, last_page_content, file.file_name, last_page_number
         )
         last_page_prepared_request = self.client.prepare_request(last_page_request)
         return last_page_prepared_request
@@ -270,8 +278,7 @@ def _get_pdf_pages(
             new_pdf.write(pdf_buffer)
             pdf_buffer.seek(0)
 
-            # 1-index the page numbers
-            yield pdf_buffer, offset + 1, offset_end
+            yield pdf_buffer, offset, offset_end
             offset += split_size
 
     def _parse_form_data(self, decoded_data: MultipartDecoder) -> FormData:
@@ -362,7 +369,9 @@ def _call_api(
             raise RuntimeError("HTTP client not accessible!")
         page_content, page_number = page
 
-        new_request = self._create_request(request, form_data, page_content, filename)
+        new_request = self._create_request(
+            request, form_data, page_content, filename, page_number
+        )
         prepared_request = self.client.prepare_request(new_request)
 
         try:
@@ -377,6 +386,7 @@ def _create_request(
         form_data: FormData,
         page_content: io.BytesIO,
         filename: str,
+        page_number: int,
     ) -> requests.Request:
         """
         Creates a request object for a part of a splitted PDF file.
@@ -392,7 +402,7 @@ def _create_request(
             original file.
         """
         headers = self._prepare_request_headers(request.headers)
-        payload = self._prepare_request_payload(form_data)
+        payload = self._prepare_request_payload(form_data, page_number)
         body = MultipartEncoder(
             fields={
                 **payload,
@@ -428,7 +438,9 @@ def _prepare_request_headers(
         headers.pop("Content-Length", None)
         return headers
 
-    def _prepare_request_payload(self, form_data: FormData) -> FormData:
+    def _prepare_request_payload(
+        self, form_data: FormData, page_number: int
+    ) -> FormData:
         """
         Prepares the request payload by removing unnecessary keys and updating the
         file.
@@ -442,7 +454,12 @@ def _prepare_request_payload(self, form_data: FormData) -> FormData:
         payload = copy.deepcopy(form_data)
         payload.pop(PARTITION_FORM_SPLIT_PDF_PAGE_KEY, None)
         payload.pop(PARTITION_FORM_FILES_KEY, None)
-        payload.update({PARTITION_FORM_SPLIT_PDF_PAGE_KEY: "false"})
+        payload.pop(PARTITION_FORM_STARTING_PAGE_NUMBER_KEY, None)
+        updated_parameters = {
+            PARTITION_FORM_SPLIT_PDF_PAGE_KEY: "false",
+            PARTITION_FORM_STARTING_PAGE_NUMBER_KEY: str(page_number),
+        }
+        payload.update(updated_parameters)
         return payload
 
     def _create_response(
@@ -540,3 +557,28 @@ def _clear_operation(self, operation_id: str) -> None:
         """
         self.partition_responses.pop(operation_id, None)
         self.partition_requests.pop(operation_id, None)
+
+    def _get_starting_page_number(self, form_data: FormData) -> int:
+        starting_page_number = DEFAULT_STARTING_PAGE_NUMBER
+        try:
+            _starting_page_number = (
+                form_data.get(PARTITION_FORM_STARTING_PAGE_NUMBER_KEY)
+                or DEFAULT_STARTING_PAGE_NUMBER
+            )
+            starting_page_number = int(_starting_page_number)  # type: ignore
+        except ValueError:
+            logger.warning(
+                "'%s' is not a valid integer. Using default value '%d'.",
+                PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
+                DEFAULT_STARTING_PAGE_NUMBER,
+            )
+
+        if starting_page_number < 1:
+            logger.warning(
+                "'%s' is less than 1. Using default value '%d'.",
+                PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
+                DEFAULT_STARTING_PAGE_NUMBER,
+            )
+            starting_page_number = DEFAULT_STARTING_PAGE_NUMBER
+
+        return starting_page_number