feat: Parameter to send custom page range when splitting pdf

To match the python feature: Unstructured-IO/unstructured-python-client#125

Add a client-side param called `splitPdfPageRange` which takes a list of two integers, `[start, end]`. If `splitPdfPage` is `true` and a range is set, slice the doc from `start` up to and including `end`. Only this page range will be sent to the API. The subset of pages is still split up as needed. If `[start, end]` is out of bounds, throw an error to the user.

Author: awalker4

overlay_client.yaml

@@ -21,3 +21,15 @@ actions:
           "description": "Number of maximum concurrent requests made when splitting PDF. Ignored on backend.",
           "default": 5,
         }
+  - target: $["components"]["schemas"]["partition_parameters"]["properties"]
+    update:
+      "split_pdf_page_range":
+        {
+          "type": "array",
+          "title": "Split Pdf Page Range",
+          "description": "When `split_pdf_page is set to `True`, this parameter selects a subset of the pdf to send to the API. The parameter is a list of 2 integers within the range [1, length_of_pdf]. TODO A ValueError is thrown if the given range is invalid. It's an internal parameter for the Python client and is not sent to the backend.",
+          "items": {"type": "integer"},
+          "minItems": 2,
+          "maxItems": 2,
+          "example": [1, 10],
+        }

src/hooks/custom/SplitPdfHook.ts

@@ -15,6 +15,7 @@ import {
   getOptimalSplitSize,
   getSplitPdfConcurrencyLevel,
   getStartingPageNumber,
+  getSplitPdfPageRange,
   loadPdf,
   prepareRequestBody,
   prepareRequestHeaders,
@@ -96,12 +97,16 @@ export class SplitPdfHook
       return request;
     }
 
-    const [error, pdf, pagesCount] = await loadPdf(file);
+    const [error, pdf, totalPages] = await loadPdf(file);
     if (file === null || pdf === null || error) {
       console.info("Partitioning without split.")
       return request;
     }
 
+    let [pageRangeStart, pageRangeEnd] = getSplitPdfPageRange(formData, totalPages);
+
+    let pagesCount = pageRangeEnd - pageRangeStart + 1;
+
     if (pagesCount < MIN_PAGES_PER_THREAD) {
       console.info(
         `PDF has less than ${MIN_PAGES_PER_THREAD} pages. Partitioning without split.`
@@ -118,15 +123,17 @@ export class SplitPdfHook
     const splitSize = await getOptimalSplitSize(pagesCount, concurrencyLevel);
     console.info("Determined optimal split size of %d pages.", splitSize)
 
-    if (splitSize >= pagesCount) {
+    // If the doc is small enough, and we aren't slicing it with a page range:
+    // do not split, just continue with the original request
+    if (splitSize >= pagesCount && pagesCount == totalPages) {
       console.info(
           "Document has too few pages (%d) to be split efficiently. Partitioning without split.",
           pagesCount,
       )
       return request;
     }
 
-    const splits = await splitPdf(pdf, splitSize);
+    const splits = await splitPdf(pdf, splitSize, pageRangeStart, pageRangeEnd);
     const numberOfSplits = splits.length
     console.info(
         "Document split into %d, %d-paged sets.",

src/hooks/custom/common.ts

@@ -11,6 +11,7 @@ export const BASE_PROTOCOL = "https:";
 export const PARTITION_FORM_FILES_KEY = "files";
 export const PARTITION_FORM_SPLIT_PDF_PAGE_KEY = "split_pdf_page";
 export const PARTITION_FORM_STARTING_PAGE_NUMBER_KEY = "starting_page_number";
+export const PARTITION_FORM_SPLIT_PDF_PAGE_RANGE_KEY = "split_pdf_page_range";
 export const PARTITION_FORM_SPLIT_PDF_CONCURRENCY_LEVEL =
   "split_pdf_concurrency_level";
 

src/hooks/custom/utils/form.ts

@@ -3,6 +3,7 @@ import {
   DEFAULT_STARTING_PAGE_NUMBER,
   MAX_NUMBER_OF_PARALLEL_REQUESTS,
   PARTITION_FORM_SPLIT_PDF_CONCURRENCY_LEVEL,
+  PARTITION_FORM_SPLIT_PDF_PAGE_RANGE_KEY,
   PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
 } from "../common.js";
 
@@ -41,6 +42,31 @@ function getIntegerParameter(
   return numberParameter;
 }
 
+/**
+ * Retrieves and validates a page range from FormData, ensuring that the start and end values are defined and within bounds.
+ *
+ * @param formData - The FormData object containing the page range parameter.
+ * @param maxPages - The maximum number of pages in the document.
+ * @returns {[number, number]} - A tuple containing the validated start and end page numbers.
+ *
+ * @throws Will throw an error if the page range is invalid or out of bounds.
+ */
+export function getSplitPdfPageRange(formData: FormData, maxPages: number): [number, number] {
+  const formDataParameter = formData.get(PARTITION_FORM_SPLIT_PDF_PAGE_RANGE_KEY);
+  const pageRange = String(formDataParameter).split(",").map(Number)
+
+  let start = pageRange[0] || 1;
+  let end = pageRange[1] || maxPages;
+
+  if (!(start > 0 && start <= maxPages) || !(end > 0 && end <= maxPages) || !(start <= end)) {
+    const msg = `Page range (${start} to ${end}) is out of bounds. Values should be between 1 and ${maxPages}.`;
+    console.error(msg);
+    throw new Error(msg);
+  }
+
+  return [start, end];
+}
+
 /**
  * Gets the number of maximum requests that can be made when splitting PDF.
  * - The number of maximum requests is determined by the value of the request parameter

src/hooks/custom/utils/pdf.ts

@@ -66,24 +66,32 @@ export async function getOptimalSplitSize(
  *
  * @param pdf - The PDF file to extract pages from.
  * @param splitSize - The number of pages per split.
+ * @param [pageRangeStart=1] - The starting page of the range to be split (1-based index). Defaults to the first page of the document.
+ * @param [pageRangeEnd=pdf.getPageCount()] - The ending page of the range to be split (1-based index). Defaults to the last page of the document.
  * @returns A promise that resolves to an array of objects containing Blob files and
  * start and end page numbers from the original document.
  */
 export async function splitPdf(
   pdf: PDFDocument,
-  splitSize: number
+  splitSize: number,
+  pageRangeStart?: number,
+  pageRangeEnd?: number
 ): Promise<PdfSplit[]> {
   const pdfSplits: PdfSplit[] = [];
-  const pagesCount = pdf.getPages().length;
+
+  const startPage = pageRangeStart || 1;
+  const endPage = pageRangeEnd || pdf.getPageCount();
+  const pagesCount = endPage - startPage + 1
+
   const numberOfSplits = Math.ceil(pagesCount / splitSize);
 
   for (let i = 0; i < numberOfSplits; ++i) {
     const offset = i * splitSize;
-    const startPage = offset + 1;
-    // If it's the last split, take the rest of the pages
-    const endPage = Math.min(pagesCount, offset + splitSize);
-    const pdfSplit = await pdfPagesToBlob(pdf, startPage, endPage);
-    pdfSplits.push({ content: pdfSplit, startPage, endPage });
+    const splitStartPage = offset + startPage;
+    const splitEndPage = Math.min(endPage, splitStartPage + splitSize - 1);
+
+    const pdfSplit = await pdfPagesToBlob(pdf, splitStartPage, splitEndPage);
+    pdfSplits.push({ content: pdfSplit, startPage: splitStartPage, endPage: splitEndPage });
   }
 
   return pdfSplits;

src/sdk/general.ts

@@ -174,6 +174,12 @@ export class General extends ClientSDK {
         if (payload$.partition_parameters.split_pdf_page !== undefined) {
             body$.append("split_pdf_page", String(payload$.partition_parameters.split_pdf_page));
         }
+        if (payload$.partition_parameters.split_pdf_page_range !== undefined) {
+            body$.append(
+                "split_pdf_page_range",
+                String(payload$.partition_parameters.split_pdf_page_range)
+            );
+        }
         if (payload$.partition_parameters.starting_page_number !== undefined) {
             body$.append(
                 "starting_page_number",

src/sdk/models/shared/partitionparameters.ts

@@ -140,6 +140,10 @@ export type PartitionParameters = {
      * Should the pdf file be split at client. Ignored on backend.
      */
     splitPdfPage?: boolean | undefined;
+    /**
+     * When `split_pdf_page is set to `True`, this parameter selects a subset of the pdf to send to the API. The parameter is a list of 2 integers within the range [1, length_of_pdf]. TODO A ValueError is thrown if the given range is invalid. It's an internal parameter for the Python client and is not sent to the backend.
+     */
+    splitPdfPageRange?: Array<number> | undefined;
     /**
      * When PDF is split into pages before sending it into the API, providing this information will allow the page number to be assigned correctly. Introduced in 1.0.27.
      */
@@ -288,6 +292,7 @@ export const PartitionParameters$inboundSchema: z.ZodType<
         skip_infer_table_types: z.array(z.string()).optional(),
         split_pdf_concurrency_level: z.number().int().default(5),
         split_pdf_page: z.boolean().default(true),
+        split_pdf_page_range: z.array(z.number().int()).optional(),
         starting_page_number: z.nullable(z.number().int()).optional(),
         strategy: Strategy$inboundSchema.default(Strategy.Auto),
         unique_element_ids: z.boolean().default(false),
@@ -313,6 +318,7 @@ export const PartitionParameters$inboundSchema: z.ZodType<
             skip_infer_table_types: "skipInferTableTypes",
             split_pdf_concurrency_level: "splitPdfConcurrencyLevel",
             split_pdf_page: "splitPdfPage",
+            split_pdf_page_range: "splitPdfPageRange",
             starting_page_number: "startingPageNumber",
             unique_element_ids: "uniqueElementIds",
             xml_keep_tags: "xmlKeepTags",
@@ -344,6 +350,7 @@ export type PartitionParameters$Outbound = {
     skip_infer_table_types?: Array<string> | undefined;
     split_pdf_concurrency_level: number;
     split_pdf_page: boolean;
+    split_pdf_page_range?: Array<number> | undefined;
     starting_page_number?: number | null | undefined;
     strategy: string;
     unique_element_ids: boolean;
@@ -380,6 +387,7 @@ export const PartitionParameters$outboundSchema: z.ZodType<
         skipInferTableTypes: z.array(z.string()).optional(),
         splitPdfConcurrencyLevel: z.number().int().default(5),
         splitPdfPage: z.boolean().default(true),
+        splitPdfPageRange: z.array(z.number().int()).optional(),
         startingPageNumber: z.nullable(z.number().int()).optional(),
         strategy: Strategy$outboundSchema.default(Strategy.Auto),
         uniqueElementIds: z.boolean().default(false),
@@ -405,6 +413,7 @@ export const PartitionParameters$outboundSchema: z.ZodType<
             skipInferTableTypes: "skip_infer_table_types",
             splitPdfConcurrencyLevel: "split_pdf_concurrency_level",
             splitPdfPage: "split_pdf_page",
+            splitPdfPageRange: "split_pdf_page_range",
             startingPageNumber: "starting_page_number",
             uniqueElementIds: "unique_element_ids",
             xmlKeepTags: "xml_keep_tags",

test/integration/SplitPdfHook.test.ts

@@ -243,3 +243,87 @@ describe("SplitPdfHook integration tests check splitted file is same as not spli
     }).rejects.toThrow(/.*File type None is not supported.*/);
   });
 });
+
+describe("SplitPdfHook integration tests page range parameter", () => {
+  const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
+  const consoleInfoSpy = jest.spyOn(console, "info");
+  const consoleWarnSpy = jest.spyOn(console, "warn");
+  const consoleErrorSpy = jest.spyOn(console, "error");
+
+  let client = new UnstructuredClient({
+    serverURL: localServer,
+    security: {
+      apiKeyAuth: FAKE_API_KEY,
+    },
+  });
+
+  beforeEach(async () => {
+    try {
+      const res = await fetch(`${localServer}/general/docs`);
+      expect(res.status).toEqual(200);
+    } catch {
+      throw Error(`The unstructured-api is not running on ${localServer}`);
+    }
+
+    client = new UnstructuredClient({
+      serverURL: localServer,
+      security: {
+        apiKeyAuth: FAKE_API_KEY,
+      },
+    });
+  });
+
+  afterEach(() => {
+    consoleInfoSpy.mockClear();
+    consoleWarnSpy.mockClear();
+    consoleErrorSpy.mockClear();
+  });
+
+  it.each([
+    { pageRange: [1, 14], expectedOk: true, expectedPages: [1, 14] }, // Valid range, start on boundary
+    { pageRange: [4, 16], expectedOk: true, expectedPages: [4, 16] }, // Valid range, end on boundary
+    { pageRange: [2, 5], expectedOk: true, expectedPages: [2, 5] },   // Valid range within boundary
+    { pageRange: [6, 6], expectedOk: true, expectedPages: [6, 6] },   // Single page range
+    { pageRange: [2, 100], expectedOk: false, expectedPages: null },  // End page too high
+    { pageRange: [50, 100], expectedOk: false, expectedPages: null }, // Range too high
+    { pageRange: [-50, 5], expectedOk: false, expectedPages: null },  // Start page too low
+    { pageRange: [-50, -2], expectedOk: false, expectedPages: null }, // Range too low
+    { pageRange: [10, 2], expectedOk: false, expectedPages: null },   // Backwards range
+  ])(
+    "for page range $pageRange",
+    async ({ pageRange, expectedOk, expectedPages }) => {
+      const filename = "test/data/layout-parser-paper.pdf";
+      const file = { content: readFileSync(filename), fileName: filename };
+
+      let startingPageNumber = 1;
+      try {
+        let response = await client.general.partition({
+          partitionParameters: {
+            files: file,
+            strategy: Strategy.Fast,
+            splitPdfPage: true,
+            splitPdfPageRange: pageRange
+          },
+        });
+
+        // Grab the set of page numbers in the result
+        // Assert that all returned elements are in the expected page range
+        const pageNumbers = new Set(response?.elements?.map((element: any) => element.metadata.page_number));
+        const minPageNumber = expectedPages?.[0] ?? 0 + startingPageNumber - 1;
+        const maxPageNumber = expectedPages?.[1] ?? 0 + startingPageNumber - 1;
+
+        expect(Math.min(...pageNumbers)).toBe(minPageNumber);
+        expect(Math.max(...pageNumbers)).toBe(maxPageNumber);
+      } catch (e) {
+        if (!expectedOk) {
+          expect(consoleErrorSpy).toHaveBeenCalledWith(expect.stringContaining("is out of bounds"));
+          return;
+        } else {
+          throw e;
+        }
+      }
+
+    },
+    300000
+  );
+});