Implement Batch Mode for Ai Review

Allow to specify the model to use in the CLI.

Author: ja-openai

cli/src/main/java/com/box/l10n/mojito/cli/command/RepositoryAiReviewCommand.java

@@ -67,9 +67,9 @@ public class RepositoryAiReviewCommand extends Command {
   boolean useBatch = false;
 
   @Parameter(
-    names = {"--use-model"},
-    arity = 1,
-    description = "Use a specific model for the review")
+      names = {"--use-model"},
+      arity = 1,
+      description = "Use a specific model for the review")
   String useModel;
 
   @Autowired CommandHelper commandHelper;

restclient/src/main/java/com/box/l10n/mojito/rest/client/RepositoryAiReviewClient.java

@@ -28,12 +28,12 @@ public ProtoAiReviewResponse reviewRepository(ProtoAiReviewRequest protoAiReview
   }
 
   public record ProtoAiReviewRequest(
-    String repositoryName,
-    List<String> targetBcp47tags,
-    int sourceTextMaxCountPerLocale,
-    List<Long> tmTextUnitIds,
-    boolean useBatch,
-    String useModel) {}
+      String repositoryName,
+      List<String> targetBcp47tags,
+      int sourceTextMaxCountPerLocale,
+      List<Long> tmTextUnitIds,
+      boolean useBatch,
+      String useModel) {}
 
   public record ProtoAiReviewResponse(PollableTask pollableTask) {}
 }

webapp/src/main/java/com/box/l10n/mojito/service/oaireview/AiReviewBatchesImportJob.java

@@ -0,0 +1,98 @@
+package com.box.l10n.mojito.service.oaireview;
+
+import com.box.l10n.mojito.openai.OpenAIClient.CreateBatchResponse;
+import com.box.l10n.mojito.openai.OpenAIClient.RetrieveBatchResponse;
+import com.box.l10n.mojito.quartz.QuartzPollableJob;
+import com.box.l10n.mojito.service.pollableTask.PollableFuture;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+
+@Component
+public class AiReviewBatchesImportJob
+    extends QuartzPollableJob<
+        AiReviewBatchesImportJob.AiReviewBatchesImportInput,
+        AiReviewBatchesImportJob.AiReviewBatchesImportOutput> {
+
+  static Logger logger = LoggerFactory.getLogger(AiReviewBatchesImportJob.class);
+
+  @Autowired AiReviewService aiReviewService;
+
+  public record AiReviewBatchesImportInput(
+      List<CreateBatchResponse> createBatchResponses, List<String> processed, int attempt) {}
+
+  public record AiReviewBatchesImportOutput(
+      List<RetrieveBatchResponse> retrieveBatchResponses, List<String> processed, Long nextJob) {}
+
+  @Override
+  public AiReviewBatchesImportOutput call(AiReviewBatchesImportInput aiReviewBatchesImportInput)
+      throws Exception {
+
+    List<RetrieveBatchResponse> retrieveBatchResponses = new ArrayList<>();
+    Set<String> processed = new HashSet<>(aiReviewBatchesImportInput.processed());
+
+    List<CreateBatchResponse> toImport =
+        aiReviewBatchesImportInput.createBatchResponses().stream()
+            .filter(b -> !processed.contains(b.id()))
+            .toList();
+
+    logger.info(
+        "{} batches to import, already processed: {} and total: {}",
+        toImport.size(),
+        processed.size(),
+        aiReviewBatchesImportInput.createBatchResponses.size());
+
+    for (CreateBatchResponse createBatchResponse : toImport) {
+
+      logger.debug("Retrieve current status of batch: {}", createBatchResponse.id());
+      RetrieveBatchResponse retrieveBatchResponse =
+          aiReviewService.retrieveBatchWithRetry(createBatchResponse);
+      retrieveBatchResponses.add(retrieveBatchResponse);
+
+      if ("completed".equals(retrieveBatchResponse.status())) {
+        logger.info("Completed batch: {}", retrieveBatchResponse.id());
+        aiReviewService.importBatch(retrieveBatchResponse);
+        processed.add(createBatchResponse.id());
+      } else if ("failed".equals(retrieveBatchResponse.status())) {
+        logger.error("Batch failed, skipping it: {}", retrieveBatchResponse);
+        processed.add(createBatchResponse.id());
+      } else if ("expired".equals(retrieveBatchResponse.status())) {
+        logger.info("Batch expired, skipping it: {}", retrieveBatchResponse);
+        processed.add(createBatchResponse.id());
+      } else if ("cancelled".equals(retrieveBatchResponse.status())) {
+        logger.info("Batch cancelled, skipping it: {}", retrieveBatchResponse);
+        processed.add(createBatchResponse.id());
+      } else {
+        logger.info("Batch is still processing will process later: {}", retrieveBatchResponse);
+      }
+    }
+
+    PollableFuture<AiReviewBatchesImportOutput> pollableFuture = null;
+
+    if (processed.size() >= aiReviewBatchesImportInput.createBatchResponses().size()) {
+      logger.info(
+          "Everything has been processed ({}/{}), don't reschedule",
+          processed.size(),
+          aiReviewBatchesImportInput.createBatchResponses().size());
+    } else {
+      logger.info("Schedule new job to process remaining batches, processed: {}", processed);
+      pollableFuture =
+          aiReviewService.aiReviewBatchesImportAsync(
+              new AiReviewBatchesImportInput(
+                  aiReviewBatchesImportInput.createBatchResponses(),
+                  processed.stream().toList(),
+                  aiReviewBatchesImportInput.attempt() + 1),
+              getCurrentPollableTask().getParentTask());
+    }
+
+    return new AiReviewBatchesImportOutput(
+        retrieveBatchResponses,
+        processed.stream().toList(),
+        pollableFuture != null ? pollableFuture.getPollableTask().getId() : null);
+  }
+}

webapp/src/main/java/com/box/l10n/mojito/service/oaireview/AiReviewJob.java

@@ -16,7 +16,7 @@ public class AiReviewJob extends QuartzPollableJob<AiReviewInput, Void> {
 
   @Override
   public Void call(AiReviewInput aiReviewJobInput) throws Exception {
-    aiReviewService.aiReview(aiReviewJobInput);
+    aiReviewService.aiReview(aiReviewJobInput, getCurrentPollableTask());
     return null;
   }
 }