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 () => {
+      requests.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;
           }
@@ -215,7 +226,9 @@ export class SplitPdfHook
       concurrencyLevel
     );
 
-    return requests.at(-1) as Request;
+    const dummyRequest = new Request("https://no-op/");
+    await this.client!.request(dummyRequest);
+    return dummyRequest;
   }
 
     /**
@@ -230,28 +243,39 @@ export class SplitPdfHook
                     successfulResponses: Response[],
                     failedResponses: Response[]
   ): Promise<Response>  {
+    let realResponse = response.clone();
+    const firstSuccessfulResponse = successfulResponses.at(0);
+    const isFakeResponse = response.headers.has("fake-response");
+    if (firstSuccessfulResponse !== undefined && isFakeResponse) {
+      realResponse = firstSuccessfulResponse.clone();
+    }
+
     let responseBody, responseStatus, responseStatusText;
     const numFailedResponses = failedResponses?.length ?? 0;
-    const headers = prepareResponseHeaders(response);
+    const headers = prepareResponseHeaders(realResponse);
 
     if (!this.allowFailed && failedResponses && failedResponses.length > 0) {
        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.`
           );
       } else {
-        responseBody = await prepareResponseBody([...successfulResponses, response]);
-        responseStatus = response.status
-        responseStatusText = response.statusText
+        if (isFakeResponse) {
+          responseBody = await prepareResponseBody([...successfulResponses]);
+        } else {
+          responseBody = await prepareResponseBody([...successfulResponses, response]);
+        }
+        responseStatus = realResponse.status
+        responseStatusText = realResponse.statusText
         if (numFailedResponses > 0) {
           console.warn(
             `${numFailedResponses} requests failed. The results might miss some pages.`

src/lib/http.ts

@@ -52,6 +52,16 @@ export class HTTPClient {
 
   async request(request: Request): Promise<Response> {
     let req = request;
+    if (req.url === "https://no-op/") {
+      return new Response('{}', {
+        headers: [
+            ["fake-response", "fake-response"]
+        ],
+        status: 200,
+        statusText: 'OK_NO_OP'
+      });
+    }
+
     for (const hook of this.requestHooks) {
       const nextRequest = await hook(req);
       if (nextRequest) {

test/integration/SplitPdfHook.test.ts

@@ -4,7 +4,7 @@ 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";
 
 describe("SplitPdfHook integration tests check splitted file is same as not splitted", () => {
   const FAKE_API_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
@@ -354,3 +354,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);
+});