Add an option to AiTranslate only a specific list of text units

box · Nov 22, 2024 · 5b47da8 · 5b47da8
1 parent 719ec38
commit 5b47da8
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 5 deletions.
diff --git a/cli/src/main/java/com/box/l10n/mojito/cli/command/RepositoryAiTranslationCommand.java b/cli/src/main/java/com/box/l10n/mojito/cli/command/RepositoryAiTranslationCommand.java
@@ -55,6 +55,13 @@ public class RepositoryAiTranslationCommand extends Command {
               + "sending too many strings to MT)")
   int sourceTextMaxCount = 100;
 
+  @Parameter(
+    names = {"--text-unit-ids"},
+    arity = 1,
+    description =
+      "The list of TmTextUnitIds to translate")
+  List<Long> textUnitIds;
+
   @Parameter(
       names = {"--use-batch"},
       arity = 1,
@@ -87,7 +94,7 @@ public void execute() throws CommandException {
     ProtoAiTranslateResponse protoAiTranslateResponse =
         repositoryAiTranslateClient.translateRepository(
             new RepositoryAiTranslateClient.ProtoAiTranslateRequest(
-                repositoryParam, locales, sourceTextMaxCount, useBatch));
+                repositoryParam, locales, sourceTextMaxCount, textUnitIds, useBatch));
 
     PollableTask pollableTask = protoAiTranslateResponse.pollableTask();
     commandHelper.waitForPollableTask(pollableTask.getId());

diff --git a/restclient/src/main/java/com/box/l10n/mojito/rest/client/RepositoryAiTranslateClient.java b/restclient/src/main/java/com/box/l10n/mojito/rest/client/RepositoryAiTranslateClient.java
@@ -32,6 +32,7 @@ public record ProtoAiTranslateRequest(
       String repositoryName,
       List<String> targetBcp47tags,
       int sourceTextMaxCountPerLocale,
+      List<Long> tmTextUnitIds,
       boolean useBatch) {}
 
   public record ProtoAiTranslateResponse(PollableTask pollableTask) {}

diff --git a/webapp/src/main/java/com/box/l10n/mojito/rest/textunit/AiTranslateWS.java b/webapp/src/main/java/com/box/l10n/mojito/rest/textunit/AiTranslateWS.java
@@ -43,6 +43,7 @@ public ProtoAiTranslateResponse aiTranslate(
                 protoAiTranslateRequest.repositoryName(),
                 protoAiTranslateRequest.targetBcp47tags(),
                 protoAiTranslateRequest.sourceTextMaxCountPerLocale(),
+                protoAiTranslateRequest.tmTextUnitIds(),
                 protoAiTranslateRequest.useBatch()));
 
     return new ProtoAiTranslateResponse(pollableFuture.getPollableTask());
@@ -53,6 +54,7 @@ public record ProtoAiTranslateRequest(
       List<String> targetBcp47tags,
       int sourceTextMaxCountPerLocale,
       boolean useBatch,
+      List<Long> tmTextUnitIds,
       boolean allLocales) {}
 
   public record ProtoAiTranslateResponse(PollableTask pollableTask) {}

diff --git a/webapp/src/main/java/com/box/l10n/mojito/service/oaitranslate/AiTranslateService.java b/webapp/src/main/java/com/box/l10n/mojito/service/oaitranslate/AiTranslateService.java
@@ -126,6 +126,7 @@ public record AiTranslateInput(
       String repositoryName,
       List<String> targetBcp47tags,
       int sourceTextMaxCountPerLocale,
+      List<Long> tmTextUnitIds,
       boolean useBatch) {}
 
   public PollableFuture<Void> aiTranslateAsync(AiTranslateInput aiTranslateInput) {
@@ -161,7 +162,7 @@ public void aiTranslateNoBatch(AiTranslateInput aiTranslateInput) {
         .flatMap(
             rl ->
                 asyncProcessLocale(
-                    rl, aiTranslateInput.sourceTextMaxCountPerLocale(), openAIClientPool),
+                    rl, aiTranslateInput.sourceTextMaxCountPerLocale(), aiTranslateInput.tmTextUnitIds(), openAIClientPool),
             10)
         .then()
         .doOnTerminate(
@@ -174,6 +175,7 @@ public void aiTranslateNoBatch(AiTranslateInput aiTranslateInput) {
   Mono<Void> asyncProcessLocale(
       RepositoryLocale repositoryLocale,
       int sourceTextMaxCountPerLocale,
+      List<Long> tmTextUnitIds,
       OpenAIClientPool openAIClientPool) {
 
     Repository repository = repositoryLocale.getRepository();
@@ -187,7 +189,12 @@ Mono<Void> asyncProcessLocale(
     textUnitSearcherParameters.setRepositoryIds(repository.getId());
     textUnitSearcherParameters.setStatusFilter(StatusFilter.FOR_TRANSLATION);
     textUnitSearcherParameters.setLocaleId(repositoryLocale.getLocale().getId());
-    textUnitSearcherParameters.setLimit(sourceTextMaxCountPerLocale);
+    if (tmTextUnitIds != null) {
+      logger.debug("Using tmTextUnitIds: {} for ai translate repository: {}", tmTextUnitIds, repository.getName());
+      textUnitSearcherParameters.setTmTextUnitIds(tmTextUnitIds);
+    } else {
+      textUnitSearcherParameters.setLimit(sourceTextMaxCountPerLocale);
+    }
 
     List<TextUnitDTO> textUnitDTOS = textUnitSearcher.search(textUnitSearcherParameters);
 

diff --git a/webapp/src/test/java/com/box/l10n/mojito/service/oaitranslate/AiTranslateServiceTest.java b/webapp/src/test/java/com/box/l10n/mojito/service/oaitranslate/AiTranslateServiceTest.java
@@ -33,7 +33,7 @@ public void aiTranslateBatch() throws ExecutionException, InterruptedException {
     aiTranslateService
         .aiTranslateAsync(
             new AiTranslateService.AiTranslateInput(
-                tmTestData.repository.getName(), null, 100, true))
+                tmTestData.repository.getName(), null, 100, null, true))
         .get();
   }
 
@@ -47,7 +47,7 @@ public void aiTranslateNoBatch()
     aiTranslateService
         .aiTranslateAsync(
             new AiTranslateService.AiTranslateInput(
-                tmTestData.repository.getName(), null, 100, false))
+                tmTestData.repository.getName(), null, 100, null, false))
         .get();
   }
 }