Skip to content

Commit 61df9a7

Browse files
ja-openaija-openai
committed
Refactor AiTranslate batch to use Quartz
lots of copy/pasta from aiReview, not trivial to share code though, we'll check later
1 parent 70f2cbd commit 61df9a7

File tree

7 files changed

+612
-123
lines changed

7 files changed

+612
-123
lines changed

‎cli/src/main/java/com/box/l10n/mojito/cli/command/RepositoryAiTranslationCommand.java‎

Lines changed: 260 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,21 @@
44
import com.beust.jcommander.Parameters;
55
import com.box.l10n.mojito.cli.command.param.Param;
66
import com.box.l10n.mojito.cli.console.ConsoleWriter;
7+
import com.box.l10n.mojito.json.ObjectMapper;
8+
import com.box.l10n.mojito.openai.OpenAIClient;
9+
import com.box.l10n.mojito.rest.client.PollableTaskClient;
710
import com.box.l10n.mojito.rest.client.RepositoryAiTranslateClient;
11+
import com.box.l10n.mojito.rest.client.RepositoryAiTranslateClient.ProtoAiTranslateRequest;
812
import com.box.l10n.mojito.rest.client.RepositoryAiTranslateClient.ProtoAiTranslateResponse;
13+
import com.box.l10n.mojito.rest.client.exception.PollableTaskException;
914
import com.box.l10n.mojito.rest.entity.PollableTask;
15+
import java.util.Comparator;
1016
import java.util.List;
17+
import java.util.Map;
18+
import java.util.Optional;
19+
import java.util.concurrent.atomic.AtomicBoolean;
1120
import java.util.stream.Collectors;
21+
import org.fusesource.jansi.Ansi;
1222
import org.fusesource.jansi.Ansi.Color;
1323
import org.slf4j.Logger;
1424
import org.slf4j.LoggerFactory;
@@ -79,10 +89,28 @@ public class RepositoryAiTranslationCommand extends Command {
7989
description = "Text to append to the end of the base prompt")
8090
String promptSuffix;
8191

92+
@Parameter(
93+
names = "--attach-job-id",
94+
arity = 1,
95+
description =
96+
"ID of an existing job to re-attach to; the CLI will only poll its status and will not start any new work.")
97+
Long attachJobId;
98+
99+
@Parameter(
100+
names = "--retry-import-job-id",
101+
arity = 1,
102+
description =
103+
"ID of an existing job to try to re-import; If a job stopped because a transient error, try to import remaining data")
104+
Long retryImportJobId;
105+
82106
@Autowired CommandHelper commandHelper;
83107

84108
@Autowired RepositoryAiTranslateClient repositoryAiTranslateClient;
85109

110+
@Autowired PollableTaskClient pollableTaskClient;
111+
112+
@Autowired ObjectMapper objectMapper;
113+
86114
@Override
87115
public boolean shouldShowInCommandList() {
88116
return false;
@@ -91,29 +119,239 @@ public boolean shouldShowInCommandList() {
91119
@Override
92120
public void execute() throws CommandException {
93121

122+
if (retryImportJobId != null) {
123+
consoleWriter.a("Retry importing task id: ").fg(Color.MAGENTA).a(retryImportJobId).println();
124+
125+
Optional<PollableTask> lastForReimport =
126+
pollableTaskClient.getPollableTask(retryImportJobId).getSubTasks().stream()
127+
.filter(t -> t.getCreatedDate() != null)
128+
.sorted(Comparator.comparing(PollableTask::getCreatedDate).reversed())
129+
.filter(PollableTask::isAllFinished)
130+
.filter(pt -> pt.getErrorMessage() != null)
131+
.findFirst();
132+
133+
if (lastForReimport.isPresent()) {
134+
long pollableTaskId =
135+
repositoryAiTranslateClient.retryImport(lastForReimport.get().getId());
136+
waitForPollable(pollableTaskId);
137+
} else {
138+
consoleWriter
139+
.fg(Color.YELLOW)
140+
.a("Last task did not finish with an error, don't retry")
141+
.println();
142+
}
143+
144+
} else if (attachJobId != null) {
145+
consoleWriter.a("Attaching, task id: ").fg(Color.MAGENTA).a(attachJobId).println();
146+
waitForPollable(attachJobId);
147+
} else {
148+
consoleWriter
149+
.newLine()
150+
.a("Ai review repository: ")
151+
.fg(Color.CYAN)
152+
.a(repositoryParam)
153+
.reset()
154+
.a(", model: ")
155+
.fg(Color.CYAN)
156+
.a(useModel)
157+
.reset()
158+
.a(" for locales: ")
159+
.fg(Color.CYAN)
160+
.a(
161+
locales == null
162+
? "<all>"
163+
: locales.stream().collect(Collectors.joining(", ", "[", "]")))
164+
.println(2);
165+
166+
ProtoAiTranslateResponse protoAiTranslateResponse =
167+
repositoryAiTranslateClient.translateRepository(
168+
new ProtoAiTranslateRequest(
169+
repositoryParam,
170+
locales,
171+
sourceTextMaxCount,
172+
textUnitIds,
173+
useBatch,
174+
useModel,
175+
promptSuffix));
176+
177+
PollableTask pollableTask = protoAiTranslateResponse.pollableTask();
178+
consoleWriter.a("Running, task id: ").fg(Color.MAGENTA).a(pollableTask.getId()).println();
179+
waitForPollable(pollableTask.getId());
180+
}
181+
182+
consoleWriter.fg(Ansi.Color.GREEN).newLine().a("Finished").println(2);
183+
}
184+
185+
void waitForPollable(Long pollableTaskId) {
186+
try {
187+
final AtomicBoolean firstRender = new AtomicBoolean(true);
188+
189+
pollableTaskClient.waitForPollableTask(
190+
pollableTaskId,
191+
PollableTaskClient.NO_TIMEOUT,
192+
pollableTask -> {
193+
Optional<PollableTask> lastFinishedForOutput =
194+
pollableTask.getSubTasks().stream()
195+
.filter(t -> t.getCreatedDate() != null)
196+
.sorted(Comparator.comparing(PollableTask::getCreatedDate).reversed())
197+
.filter(PollableTask::isAllFinished)
198+
.findFirst();
199+
200+
if (lastFinishedForOutput.isPresent()) {
201+
if (!firstRender.get()) {
202+
consoleWriter.erasePreviouslyPrintedLines();
203+
} else {
204+
firstRender.set(false);
205+
}
206+
207+
Long lastFinishedTaskId = lastFinishedForOutput.get().getId();
208+
consoleWriter
209+
.a("Running, task id: ")
210+
.fg(Color.MAGENTA)
211+
.a(pollableTaskId)
212+
.reset()
213+
.a(", child task id: ")
214+
.fg(Color.MAGENTA)
215+
.a(lastFinishedTaskId)
216+
.newLine();
217+
String pollableTaskOutput =
218+
pollableTaskClient.getPollableTaskOutput(lastFinishedTaskId);
219+
try {
220+
renderAiReviewBatchesImportOutput(
221+
objectMapper.readValueUnchecked(
222+
pollableTaskOutput, AiTranslateBatchesImportOutput.class));
223+
} catch (Exception e) {
224+
logger.error("Can't render", e);
225+
consoleWriter
226+
.reset()
227+
.a("Can't render:" + e.getMessage())
228+
.newLine()
229+
.a(pollableTaskOutput)
230+
.newLine();
231+
}
232+
}
233+
});
234+
235+
} catch (PollableTaskException e) {
236+
throw new CommandException(e.getMessage(), e.getCause());
237+
}
238+
}
239+
240+
void renderAiReviewBatchesImportOutput(
241+
AiTranslateBatchesImportOutput aiTranslateBatchesImportOutput) {
242+
243+
if (!aiTranslateBatchesImportOutput.batchCreationErrors().isEmpty()) {
244+
consoleWriter
245+
.fg(Color.RED)
246+
.a(
247+
"Some batches failed to be created. The following locales will not be processed and will need to be retried:")
248+
.newLine();
249+
for (String batchCreationError : aiTranslateBatchesImportOutput.batchCreationErrors()) {
250+
consoleWriter.a("- " + batchCreationError).newLine();
251+
}
252+
consoleWriter.newLine();
253+
}
254+
255+
if (!aiTranslateBatchesImportOutput.skippedLocales().isEmpty()) {
256+
consoleWriter
257+
.reset()
258+
.a("No content to review for the following locales; skipping: ")
259+
.fg(Color.MAGENTA)
260+
.a(String.join(",", aiTranslateBatchesImportOutput.skippedLocales()))
261+
.reset()
262+
.newLine();
263+
consoleWriter.newLine();
264+
}
265+
266+
aiTranslateBatchesImportOutput
267+
.retrieveBatchResponses()
268+
.forEach(
269+
r ->
270+
renderBatch(
271+
r,
272+
aiTranslateBatchesImportOutput.failedToImport.get(r.id()),
273+
aiTranslateBatchesImportOutput.processed.contains(r.id())));
274+
consoleWriter.println();
275+
}
276+
277+
void renderBatch(
278+
OpenAIClient.RetrieveBatchResponse retrieveBatchResponse,
279+
String importError,
280+
boolean processed) {
281+
consoleWriter.a("- ").fg(Color.CYAN).a(retrieveBatchResponse.id()).a(" ");
282+
283+
consoleWriter.reset().a("[import: ");
284+
if (importError != null) {
285+
consoleWriter.fg(Color.RED).a("failed");
286+
} else {
287+
if (processed) {
288+
if ("completed".equals(retrieveBatchResponse.status())) {
289+
consoleWriter.fg(Color.GREEN).a("success");
290+
} else {
291+
consoleWriter.fg(Color.YELLOW).a(" - ");
292+
}
293+
} else {
294+
consoleWriter.fg(Color.YELLOW).a("waiting");
295+
}
296+
}
297+
consoleWriter.reset().a("]");
298+
299+
Color batchStatusColor =
300+
switch (retrieveBatchResponse.status()) {
301+
case "completed" -> Color.GREEN;
302+
case "failed" -> Color.RED;
303+
case "running", "queued", "in_progress" -> Color.YELLOW;
304+
default -> Color.DEFAULT;
305+
};
306+
307+
OpenAIClient.RetrieveBatchResponse.RequestCounts c = retrieveBatchResponse.requestCounts();
308+
94309
consoleWriter
95-
.newLine()
96-
.a("Ai translate repository: ")
97-
.fg(Color.CYAN)
98-
.a(repositoryParam)
99310
.reset()
100-
.a(" for locales: ")
101-
.fg(Color.CYAN)
102-
.a(locales == null ? "<all>" : locales.stream().collect(Collectors.joining(", ", "[", "]")))
103-
.println(2);
104-
105-
ProtoAiTranslateResponse protoAiTranslateResponse =
106-
repositoryAiTranslateClient.translateRepository(
107-
new RepositoryAiTranslateClient.ProtoAiTranslateRequest(
108-
repositoryParam,
109-
locales,
110-
sourceTextMaxCount,
111-
textUnitIds,
112-
useBatch,
113-
useModel,
114-
promptSuffix));
115-
116-
PollableTask pollableTask = protoAiTranslateResponse.pollableTask();
117-
commandHelper.waitForPollableTask(pollableTask.getId());
311+
.a(" [batch: ")
312+
.fg(batchStatusColor)
313+
.a(retrieveBatchResponse.status())
314+
.reset()
315+
.a(" ; total=")
316+
.a(c.total())
317+
.a(", completed=")
318+
.a(c.completed())
319+
.a(", ");
320+
321+
if (c.failed() > 0) {
322+
consoleWriter.fg(Color.RED);
323+
}
324+
325+
consoleWriter.a("failed=").a(c.failed()).reset();
326+
consoleWriter.a("]").newLine();
327+
328+
if (importError != null) {
329+
consoleWriter.newLine().a("Import error: ").newLine().fg(Color.RED).a(importError).newLine();
330+
}
331+
332+
if (retrieveBatchResponse.errors() != null
333+
&& retrieveBatchResponse.errors().data() != null
334+
&& !retrieveBatchResponse.errors().data().isEmpty()) {
335+
consoleWriter.fg(Color.RED).a(" Errors:").reset().newLine();
336+
retrieveBatchResponse
337+
.errors()
338+
.data()
339+
.forEach(
340+
e ->
341+
consoleWriter
342+
.a(" - ")
343+
.a(
344+
"[%s] %s (param=%s, line=%s)"
345+
.formatted(e.code(), e.message(), e.param(), e.line()))
346+
.newLine());
347+
}
118348
}
349+
350+
public record AiTranslateBatchesImportOutput(
351+
List<OpenAIClient.RetrieveBatchResponse> retrieveBatchResponses,
352+
List<String> skippedLocales,
353+
List<String> batchCreationErrors,
354+
List<String> processed,
355+
Map<String, String> failedToImport,
356+
Long nextJob) {}
119357
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,15 @@ public ProtoAiTranslateResponse translateRepository(
2828
getBasePathForEntity(), protoAiTranslateRequest, ProtoAiTranslateResponse.class);
2929
}
3030

31+
public long retryImport(Long childPollableTaskId) {
32+
ProtoAiTranslateRetryImportResponse protoAiReviewRetryImportResponse =
33+
authenticatedRestTemplate.postForObject(
34+
getBasePathForEntity() + "/retry-import",
35+
new ProtoAiTranslateRetryImportRequest(childPollableTaskId),
36+
ProtoAiTranslateRetryImportResponse.class);
37+
return protoAiReviewRetryImportResponse.pollableTaskId();
38+
}
39+
3140
public record ProtoAiTranslateRequest(
3241
String repositoryName,
3342
List<String> targetBcp47tags,
@@ -38,4 +47,8 @@ public record ProtoAiTranslateRequest(
3847
String promptSuffix) {}
3948

4049
public record ProtoAiTranslateResponse(PollableTask pollableTask) {}
50+
51+
public record ProtoAiTranslateRetryImportRequest(long childPollableTaskId) {}
52+
53+
public record ProtoAiTranslateRetryImportResponse(long pollableTaskId) {}
4154
}

‎webapp/src/main/java/com/box/l10n/mojito/rest/textunit/AiTranslateWS.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,4 +61,19 @@ public record ProtoAiTranslateRequest(
6161
String promptSuffix) {}
6262

6363
public record ProtoAiTranslateResponse(PollableTask pollableTask) {}
64+
65+
@RequestMapping(method = RequestMethod.POST, value = "/api/proto-ai-translate/retry-import")
66+
@ResponseStatus(HttpStatus.OK)
67+
public ProtoAiReviewRetryImportResponse aiReviewRetryImport(
68+
@RequestBody ProtoAiReviewRetryImportRequest protoAiReviewRetryImportRequest) {
69+
PollableFuture<Void> pollableFuture =
70+
aiTranslateService.retryImport(
71+
protoAiReviewRetryImportRequest.childPollableTaskId(),
72+
PollableTask.INJECT_CURRENT_TASK);
73+
return new ProtoAiReviewRetryImportResponse(pollableFuture.getPollableTask().getId());
74+
}
75+
76+
public record ProtoAiReviewRetryImportRequest(long childPollableTaskId) {}
77+
78+
public record ProtoAiReviewRetryImportResponse(long pollableTaskId) {}
6479
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -779,7 +779,7 @@ String generateBatchFileContent(
779779
.collect(joining("\n"));
780780
}
781781

782-
// this is not review specific move
782+
// TODO this is not review specific move
783783
OpenAIClient.RetrieveBatchResponse retrieveBatchWithRetry(
784784
OpenAIClient.CreateBatchResponse batch) {
785785

0 commit comments

Comments
 (0)