fix: Improve retry mechanisms for split pdf logic

Author: jordan-homan

src/hooks/custom/SplitPdfHook.ts

@@ -29,6 +29,7 @@ import {
   PARTITION_FORM_FILES_KEY,
   PARTITION_FORM_SPLIT_PDF_PAGE_KEY,
 } from "./common.js";
+import {retry, RetryConfig} from "../../lib/retries";
 
 /**
  * Represents a hook for splitting and sending PDF files as per page requests.
@@ -178,9 +179,11 @@ export class SplitPdfHook
         file.name,
         firstPageNumber
       );
+      const timeoutInMs = 60 * 10 * 1000;
       const req = new Request(requestClone, {
         headers,
         body,
+        signal: AbortSignal.timeout(timeoutInMs)
       });
       requests.push(req);
       setIndex+=1;
@@ -190,11 +193,19 @@ export class SplitPdfHook
 
     const allowFailed = this.allowFailed;
 
+    const retryConfig = { strategy: "backoff" } as RetryConfig;
+    const retryCodes = ["502", "503", "504"];
+
     this.partitionRequests[operationID] = async.parallelLimit(
       requests.slice(0, -1).map((req, pageIndex) => async () => {
         const pageNumber = pageIndex + startingPageNumber;
         try {
-          const response = await this.client!.request(req);
+         const response = await retry(
+              async () => {
+                return await this.client!.request(req.clone());
+              },
+              { config: retryConfig, statusCodes: retryCodes }
+          );
           if (response.status === 200) {
             (this.partitionSuccessfulResponses[operationID] as Response[])[pageIndex] =
               response.clone();
@@ -206,7 +217,7 @@ export class SplitPdfHook
               }
           }
         } catch (e) {
-          console.error(`Failed to send request for page ${pageNumber}.`);
+          console.error(`Failed to send request for page ${pageNumber}.`, e);
           if (!allowFailed) {
             throw e;
           }
@@ -238,13 +249,13 @@ export class SplitPdfHook
        const failedResponse = failedResponses[0]?.clone();
        if (failedResponse) {
             responseBody = await failedResponse.text();
-            responseStatus = failedResponse.status;
             responseStatusText = failedResponse.statusText;
         } else {
             responseBody = JSON.stringify({"details:": "Unknown error"});
-            responseStatus = 503
             responseStatusText = "Unknown error"
         }
+        // if the response status is unknown or was 502, 503, 504, set back to 500 to ensure we don't cause more retries
+        responseStatus = 500;
         console.warn(
             `${numFailedResponses} requests failed. The partition operation is cancelled.`
           );

test/integration/SplitPdfHook.test.ts

@@ -4,7 +4,8 @@ import { UnstructuredClient } from "../../src";
 import { PartitionResponse } from "../../src/sdk/models/operations";
 import { PartitionParameters, Strategy } from "../../src/sdk/models/shared";
 
-const localServer = "http://localhost:8000"
+const localServer = "http://localhost:8000";
+const SECONDS = 1000;
 
 describe("SplitPdfHook integration tests check splitted file is same as not splitted", () => {
   const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
@@ -354,3 +355,42 @@ describe("SplitPdfHook integration tests page range parameter", () => {
     300000
   );
 });
+
+
+describe("SplitPDF succeeds for large PDF with high concurrency", () => {
+    const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
+
+    it.each([
+      `${localServer}/general/v0/general`,
+    ])("succeed", async (serverURL) => {
+        const client = new UnstructuredClient({
+            serverURL: serverURL,
+            security: {
+                apiKeyAuth: FAKE_API_KEY,
+            },
+        });
+
+        const file = {
+            content: readFileSync("test/data/layout-parser-paper.pdf"),
+            fileName: "test/data/layout-parser-paper.pdf"
+        };
+
+        const requestParams: PartitionParameters = {
+            files: file,
+            splitPdfPage: true,
+            strategy: Strategy.HiRes,
+            splitPdfAllowFailed: false,
+            splitPdfConcurrencyLevel: 15
+        };
+
+        const res: PartitionResponse = await client.general.partition({
+            partitionParameters: {
+                ...requestParams
+            },
+        });
+
+        expect(res.statusCode).toEqual(200);
+        expect(res.elements?.length).toBeGreaterThan(0);
+    },
+    300000);
+});