⚡ Emit PredictionCreatedEvent, response contains input object

Author: MrGraversen

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/CheckAndEmitPredictionCreationTask.java

@@ -0,0 +1,40 @@
+package io.graversen.replicate.facade;
+
+import jakarta.annotation.Nullable;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.SneakyThrows;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.stereotype.Component;
+
+import java.util.Optional;
+import java.util.concurrent.CompletionException;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+@Component
+@RequiredArgsConstructor
+public class CheckAndEmitPredictionCreationTask implements Function<Optional<PredictionResponseAndModel>, PredictionResponseAndModel> {
+    private final @NonNull ApplicationEventPublisher eventPublisher;
+
+    @Override
+    @SneakyThrows
+    public PredictionResponseAndModel apply(@Nullable Optional<PredictionResponseAndModel> predictionResponse) {
+        final var response = predictionResponse.orElseThrow(predictionFailedError());
+        final var event = mapPredictionCreatedEvent().apply(response);
+        eventPublisher.publishEvent(event);
+        return response;
+    }
+
+    Function<PredictionResponseAndModel, PredictionCreatedEvent> mapPredictionCreatedEvent() {
+        return predictionResponse -> new PredictionCreatedEvent(
+                predictionResponse.getPredictionResponse().getId(),
+                predictionResponse.getModel(),
+                predictionResponse.getPredictionResponse()
+        );
+    }
+
+    private Supplier<Exception> predictionFailedError() {
+        return () -> new CompletionException(new RuntimeException("Could not initiate prediction"));
+    }
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/CheckPredictionStateTask.java

@@ -0,0 +1,16 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.service.PredictionResponse;
+import lombok.AccessLevel;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+
+@Value
+@RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+public class PredictionCreatedEvent {
+    @NonNull String id;
+    @NonNull ReplicateModel model;
+    @NonNull PredictionResponse prediction;
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/PredictionCreatedEvent.java

@@ -19,7 +19,7 @@
 public class ReplicateFacade {
     private final @NonNull ReplicateService replicateService;
     private final @NonNull ExecutorService executorService;
-    private final @NonNull CheckPredictionStateTask checkPredictionStateTask;
+    private final @NonNull CheckAndEmitPredictionCreationTask checkAndEmitPredictionCreationTask;
     private final @NonNull PollPredictionStatusTask pollPredictionStatusTask;
     private final @NonNull EmitPredictionResponseTask emitPredictionResponseTask;
 
@@ -29,7 +29,7 @@ public CompletableFuture<PredictionResponseAndModel> createPrediction(
     ) {
         return CompletableFuture
                 .supplyAsync(doCreatePrediction(model, createPrediction), executorService)
-                .thenApplyAsync(checkPredictionStateTask, executorService)
+                .thenApplyAsync(checkAndEmitPredictionCreationTask, executorService)
                 .thenApplyAsync(pollPredictionStatusTask, executorService)
                 .whenCompleteAsync(emitPredictionResponseTask, executorService);
     }

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/ReplicateFacade.java

@@ -23,6 +23,7 @@ public class PredictionResponse {
     OffsetDateTime completedAt;
     String error;
     String status;
+    Object input;
     Object output;
     PredictionUrls urls;
 

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/service/PredictionResponse.java

@@ -24,6 +24,7 @@ public class ReplicateUtils {
     private static final String ATTRIBUTE_ERROR = "error";
     private static final String ATTRIBUTE_STATUS = "status";
     private static final String ATTRIBUTE_OUTPUT = "output";
+    private static final String ATTRIBUTE_INPUT = "input";
     private static final String ATTRIBUTE_URLS = "urls";
     private static final String ATTRIBUTE_CANCEL_URL = "cancel";
     private static final String ATTRIBUTE_GET_URL = "get";
@@ -37,6 +38,7 @@ public static Optional<PredictionResponse> mapPredictionResponse(@NonNull Linked
             final var completedAt = parseOffsetDateTime((String) responseMap.get(ATTRIBUTE_COMPLETED_AT));
             final var error = (String) responseMap.get(ATTRIBUTE_ERROR);
             final var status = (String) responseMap.get(ATTRIBUTE_STATUS);
+            final var input = responseMap.get(ATTRIBUTE_INPUT);
             final var output = responseMap.get(ATTRIBUTE_OUTPUT);
 
             final var urlsMap = (LinkedHashMap<String, Object>) responseMap.get(ATTRIBUTE_URLS);
@@ -48,7 +50,7 @@ public static Optional<PredictionResponse> mapPredictionResponse(@NonNull Linked
                 urls = new PredictionUrls(cancelUrl, getUrl);
             }
 
-            final var predictionResponse = new PredictionResponse(id, version, createdAt, startedAt, completedAt, error, status, output, urls);
+            final var predictionResponse = new PredictionResponse(id, version, createdAt, startedAt, completedAt, error, status, input, output, urls);
             return Optional.of(predictionResponse);
         } catch (Exception e) {
             log.error(e.getMessage(), e);