✨ add support for workflow polling (#238)

Author: ianardee

README.md

@@ -74,6 +74,7 @@ import com.mindee.product.us.bankcheck.BankCheckV1;
 ### Custom Documents (docTI & Custom APIs)
 ```java
 import com.mindee.MindeeClient;
+import com.mindee.PredictOptions;
 import com.mindee.input.LocalInputSource;
 import com.mindee.parsing.common.PredictResponse;
 import com.mindee.product.generated.GeneratedV1;
@@ -98,6 +99,7 @@ public class SimpleMindeeClient {
     Document<GeneratedV1> customDocument = mindeeClient.enqueueAndParse(
         localInputSource,
         endpoint
+        // PredictOptions.builder().build(),
     );
   }
 }
@@ -116,6 +118,7 @@ This is the easiest way to get started.
 
 ```java
 import com.mindee.MindeeClient;
+import com.mindee.PredictOptions;
 import com.mindee.input.LocalInputSource;
 import com.mindee.parsing.common.AsyncPredictResponse;
 import com.mindee.product.internationalid.InternationalIdV2;
@@ -138,6 +141,7 @@ public class SimpleMindeeClient {
     AsyncPredictResponse<InternationalIdV2> response = mindeeClient.enqueueAndParse(
         InternationalIdV2.class,
         inputSource
+        // PredictOptions.builder().build(),
     );
 
     // Print a summary of the response

docs/code_samples/workflow_execution.txt

@@ -24,7 +24,6 @@ public class SimpleMindeeClient {
         inputSource
     );
 
-
     // Alternatively: give an alias to the document
     //    WorkflowResponse response = mindeeClient.executeWorkflow(
     //        workflowId,

src/main/java/com/mindee/MindeeClient.java

@@ -146,6 +146,30 @@ public <T extends Inference> AsyncPredictResponse<T> enqueue(
     );
   }
 
+  /**
+   * Send a local file to an async queue.
+   * @param <T> Type of inference.
+   * @param type Type of inference.
+   * @param localInputSource A local input source file.
+   * @param predictOptions Prediction options for the enqueuing.
+   * @return an instance of {@link AsyncPredictResponse}.
+   * @throws IOException Throws if the file can't be accessed.
+   */
+  public <T extends Inference> AsyncPredictResponse<T> enqueue(
+      Class<T> type,
+      LocalInputSource localInputSource,
+      PredictOptions predictOptions
+  ) throws IOException {
+    return this.enqueue(
+        type,
+        new Endpoint(type),
+        localInputSource.getFile(),
+        localInputSource.getFilename(),
+        predictOptions,
+        null
+    );
+  }
+
   /**
    * Send a remote file to an async queue.
    * @param <T> Type of inference.
@@ -290,6 +314,60 @@ public <T extends Inference> AsyncPredictResponse<T> enqueueAndParse(
     );
   }
 
+  /**
+   * Send a local file to an async queue, poll, and parse when complete.
+   * @param <T> Type of inference.
+   * @param type Type of inference.
+   * @param localInputSource A local input source file.
+   * @param predictOptions Prediction options for the enqueuing.
+   * @param pollingOptions Options for async call parameters.
+   * @return an instance of {@link AsyncPredictResponse}.
+   * @throws IOException Throws if the file can't be accessed.
+   * @throws InterruptedException Throws in the event of a timeout.
+   */
+  public <T extends Inference> AsyncPredictResponse<T> enqueueAndParse(
+      Class<T> type,
+      LocalInputSource localInputSource,
+      PredictOptions predictOptions,
+      AsyncPollingOptions pollingOptions
+  ) throws IOException, InterruptedException {
+    return this.enqueueAndParse(
+        type,
+        new Endpoint(type),
+        pollingOptions,
+        localInputSource.getFile(),
+        localInputSource.getFilename(),
+        predictOptions,
+        null
+    );
+  }
+
+  /**
+   * Send a local file to an async queue, poll, and parse when complete.
+   * @param <T> Type of inference.
+   * @param type Type of inference.
+   * @param localInputSource A local input source file.
+   * @param predictOptions Prediction options for the enqueuing.
+   * @return an instance of {@link AsyncPredictResponse}.
+   * @throws IOException Throws if the file can't be accessed.
+   * @throws InterruptedException Throws in the event of a timeout.
+   */
+  public <T extends Inference> AsyncPredictResponse<T> enqueueAndParse(
+      Class<T> type,
+      LocalInputSource localInputSource,
+      PredictOptions predictOptions
+  ) throws IOException, InterruptedException {
+    return this.enqueueAndParse(
+        type,
+        new Endpoint(type),
+        null,
+        localInputSource.getFile(),
+        localInputSource.getFilename(),
+        predictOptions,
+        null
+    );
+  }
+
   /**
    * Send a remote file to an async queue, poll, and parse when complete.
    * @param <T> Type of inference.

src/main/java/com/mindee/PredictOptions.java

@@ -24,15 +24,28 @@ public class PredictOptions {
    * size.
    */
   Boolean fullText;
+  /**
+   * If set, will enqueue to a workflow queue instead of a product's endpoint.
+   */
+  String workflowId;
+  /**
+   * If set, will enable Retrieval-Augmented Generation.
+   * Only works if a valid workflowId is set.
+   */
+  Boolean rag;
 
   @Builder
   private PredictOptions(
       Boolean allWords,
       Boolean fullText,
-      Boolean cropper
+      Boolean cropper,
+      String workflowId,
+      Boolean rag
   ) {
     this.allWords = allWords == null ? Boolean.FALSE : allWords;
     this.fullText = fullText == null ? Boolean.FALSE : fullText;
     this.cropper = cropper == null ? Boolean.FALSE : cropper;
+    this.workflowId = workflowId;
+    this.rag = rag == null ? Boolean.FALSE : rag;
   }
 }

src/main/java/com/mindee/http/MindeeHttpApi.java

@@ -41,8 +41,9 @@
 public final class MindeeHttpApi extends MindeeApi {
 
   private static final ObjectMapper mapper = new ObjectMapper();
-  private final Function<Endpoint, String> buildBaseUrl = this::buildProductUrl;
-  private final Function<String, String> buildWorkflowBaseUrl = this::buildWorkflowUrl;
+  private final Function<Endpoint, String> buildProductPredicBasetUrl = this::buildProductPredictBaseUrl;
+  private final Function<String, String> buildWorkflowPredictBaseUrl = this::buildWorkflowPredictBaseUrl;
+  private final Function<String, String> buildWorkflowExecutionBaseUrl = this::buildWorkflowExecutionUrl;
   /**
    * The MindeeSetting needed to make the api call.
    */
@@ -53,24 +54,27 @@ public final class MindeeHttpApi extends MindeeApi {
    */
   private final HttpClientBuilder httpClientBuilder;
   /**
-   * The function used to generate the API endpoint URL.
+   * The function used to generate the synchronous API endpoint URL.
    * Only needs to be set if the api calls need to be directed through internal URLs.
    */
   private final Function<Endpoint, String> urlFromEndpoint;
-
   /**
-   * The function used to generate the API endpoint URL for workflow execution calls.
+   * The function used to generate the asynchronous API endpoint URL for a product.
    * Only needs to be set if the api calls need to be directed through internal URLs.
    */
   private final Function<Endpoint, String> asyncUrlFromEndpoint;
+  /**
+   * The function used to generate the asynchronous API endpoint URL for a workflow.
+   * Only needs to be set if the api calls need to be directed through internal URLs.
+   */
+  private final Function<String, String> asyncUrlFromWorkflow;
   /**
    * The function used to generate the Job status URL for Async calls.
    * Only needs to be set if the api calls need to be directed through internal URLs.
    */
   private final Function<Endpoint, String> documentUrlFromEndpoint;
-
   /**
-   * The function used to generate the Job status URL for Async calls.
+   * The function used to generate the Job status URL for workflow execution calls.
    * Only needs to be set if the api calls need to be directed through internal URLs.
    */
   private final Function<String, String> workflowUrlFromId;
@@ -82,6 +86,7 @@ public MindeeHttpApi(MindeeSettings mindeeSettings) {
         null,
         null,
         null,
+        null,
         null
     );
   }
@@ -93,7 +98,8 @@ private MindeeHttpApi(
       Function<Endpoint, String> urlFromEndpoint,
       Function<Endpoint, String> asyncUrlFromEndpoint,
       Function<Endpoint, String> documentUrlFromEndpoint,
-      Function<String, String> workflowUrlFromEndpoint
+      Function<String, String> workflowUrlFromEndpoint,
+      Function<String, String> asyncUrlFromWorkflow
   ) {
     this.mindeeSettings = mindeeSettings;
 
@@ -106,26 +112,35 @@ private MindeeHttpApi(
     if (urlFromEndpoint != null) {
       this.urlFromEndpoint = urlFromEndpoint;
     } else {
-      this.urlFromEndpoint = buildBaseUrl.andThen((url) -> url.concat("/predict"));
+      this.urlFromEndpoint = buildProductPredicBasetUrl.andThen(
+          (url) -> url.concat("/predict"));
+    }
+
+    if (asyncUrlFromWorkflow != null) {
+      this.asyncUrlFromWorkflow = asyncUrlFromWorkflow;
+    } else {
+      this.asyncUrlFromWorkflow = this.buildWorkflowPredictBaseUrl.andThen(
+          (url) -> url.concat("/predict_async"));
     }
 
     if (asyncUrlFromEndpoint != null) {
       this.asyncUrlFromEndpoint = asyncUrlFromEndpoint;
     } else {
-      this.asyncUrlFromEndpoint = this.urlFromEndpoint.andThen((url) -> url.concat("_async"));
+      this.asyncUrlFromEndpoint = this.buildProductPredicBasetUrl.andThen(
+          (url) -> url.concat("/predict_async"));
     }
 
     if (documentUrlFromEndpoint != null) {
       this.documentUrlFromEndpoint = documentUrlFromEndpoint;
     } else {
-      this.documentUrlFromEndpoint = this.buildBaseUrl.andThen(
+      this.documentUrlFromEndpoint = this.buildProductPredicBasetUrl.andThen(
           (url) -> url.concat("/documents/queue/"));
     }
 
     if (workflowUrlFromEndpoint != null) {
       this.workflowUrlFromId = workflowUrlFromEndpoint;
     } else {
-      this.workflowUrlFromId = this.buildWorkflowBaseUrl;
+      this.workflowUrlFromId = this.buildWorkflowExecutionBaseUrl;
     }
   }
 
@@ -233,7 +248,12 @@ public <DocT extends Inference> AsyncPredictResponse<DocT> predictAsyncPost(
       RequestParameters requestParameters
   ) throws IOException {
 
-    String url = asyncUrlFromEndpoint.apply(endpoint);
+    String url;
+    if (requestParameters.getPredictOptions().getWorkflowId() != null) {
+      url = asyncUrlFromWorkflow.apply(requestParameters.getPredictOptions().getWorkflowId());
+    } else {
+      url = asyncUrlFromEndpoint.apply(endpoint);
+    }
     HttpPost post = buildHttpPost(url, requestParameters);
 
     // required to register jackson date module format to deserialize
@@ -340,7 +360,7 @@ private <ResponseT extends ApiResponse> MindeeHttpException getHttpError(
     return new MindeeHttpException(statusCode, message, details, errorCode);
   }
 
-  private String buildProductUrl(Endpoint endpoint) {
+  private String buildProductPredictBaseUrl(Endpoint endpoint) {
     return this.mindeeSettings.getBaseUrl()
         + "/products/"
         + endpoint.getAccountName()
@@ -350,7 +370,11 @@ private String buildProductUrl(Endpoint endpoint) {
         + endpoint.getVersion();
   }
 
-  private String buildWorkflowUrl(String workflowId) {
+  private String buildWorkflowPredictBaseUrl(String workflowId) {
+    return this.mindeeSettings.getBaseUrl() + "/workflows/" + workflowId;
+  }
+
+  private String buildWorkflowExecutionUrl(String workflowId) {
     return this.mindeeSettings.getBaseUrl() + "/workflows/" + workflowId + "/executions";
   }
 
@@ -388,7 +412,9 @@ private List<NameValuePair> buildPostParams(
     if (Boolean.TRUE.equals(requestParameters.getPredictOptions().getFullText())) {
       params.add(new BasicNameValuePair("full_text_ocr", "true"));
     }
-    if (Boolean.TRUE.equals(requestParameters.getWorkflowOptions().getRag())) {
+    if (Boolean.TRUE.equals(requestParameters.getWorkflowOptions().getRag())
+        || Boolean.TRUE.equals(requestParameters.getPredictOptions().getRag())
+    ) {
       params.add(new BasicNameValuePair("rag", "true"));
     }
     return params;

src/main/java/com/mindee/parsing/common/InferenceExtras.java

@@ -1,6 +1,7 @@
 package com.mindee.parsing.common;
 
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
@@ -17,4 +18,9 @@ public class InferenceExtras {
    * Full Text OCR result.
    */
   private String fullTextOcr;
+  /**
+   * Retrieval-Augmented Generation results.
+   */
+  @JsonProperty("rag")
+  private Rag rag;
 }

src/main/java/com/mindee/parsing/common/Rag.java

@@ -0,0 +1,22 @@
+package com.mindee.parsing.common;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+
+/**
+ * Retrieval-Augmented Generation info class.
+ */
+@Getter
+@EqualsAndHashCode
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class Rag {
+  /**
+   * The document ID that was matched.
+   */
+  @Setter
+  @JsonProperty("matching_document_id")
+  private String matchingDocumentId;
+}