✨ Stateful conversation processing

Author: MrGraversen

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

@@ -0,0 +1,81 @@
+package io.graversen.replicate.facade;
+
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import io.graversen.replicate.llama3.Llama3Tokenizer;
+import io.graversen.replicate.service.Conversation;
+import io.graversen.replicate.service.ConversationService;
+import io.graversen.replicate.service.CreateConversation;
+import io.graversen.replicate.service.CreateTextPrediction;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import org.springframework.stereotype.Component;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
+
+@Component
+@RequiredArgsConstructor
+public class ConversationFacade {
+    private final ConcurrentMap<String, ConversationOptions> conversationOptions = new ConcurrentHashMap<>();
+
+    private final @NonNull ReplicateFacade replicateFacade;
+    private final @NonNull ConversationService conversationService;
+
+    public Conversation create(@NonNull CreateConversation createConversation, @NonNull ConversationOptions options) {
+        final var conversation = conversationService.create(createConversation);
+        conversationOptions.put(conversation.getId(), options);
+        return conversation;
+    }
+
+    public CompletableFuture<Conversation> chat(@NonNull String id, @NonNull TextMessage message) {
+        var conversation = conversationService.appendMessage(id, message);
+
+        if (conversation.isEmpty()) {
+            return CompletableFuture.failedFuture(new IllegalArgumentException("Could not find Conversation: " + id));
+        }
+
+        final var textConversation = conversation.get().getConversation();
+        final var model = conversation.get().getModel();
+        final var options = conversationOptions.get(conversation.get().getId());
+
+        final var createTextPrediction = new CreateTextPrediction(
+                textConversation,
+                options.getTemperature(),
+                null,
+                null,
+                null,
+                null,
+                null
+        );
+
+        final var pendingPrediction = replicateFacade.createPrediction(model, createTextPrediction);
+        return pendingPrediction
+                .thenApply(parseTextConversation())
+                .thenApply(updateConversation(conversation.get().getId()));
+    }
+
+    Function<TextConversation, Conversation> updateConversation(@NonNull String id) {
+        return conversation -> conversationService.update(id, conversation)
+                .orElseThrow(() -> new IllegalArgumentException("Could not find Conversation: " + id));
+    }
+
+    Function<PredictionResponseAndModel, TextConversation> parseTextConversation() {
+        return predictionResponseAndModel -> {
+            // For now, assume Llama3 is used and apply its tokenization parsing
+            // In the future, rely on a strategy pattern implementation selector for multi-model-family support
+
+            final var textInput = predictionResponseAndModel.getPredictionResponse().getInputKey("prompt")
+                    .orElseThrow(() -> new IllegalStateException("No text input found in prediction"));
+
+            var conversation = Llama3Tokenizer.parseTextCompletion(textInput);
+
+            final var textOutput = predictionResponseAndModel.getPredictionResponse().getTextOutput()
+                    .orElseThrow(() -> new IllegalStateException("No text output found in prediction"));
+
+            return conversation.append(TextMessage.assistant(textOutput));
+        };
+    }
+}

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

@@ -0,0 +1,11 @@
+package io.graversen.replicate.facade;
+
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+
+@Getter
+@RequiredArgsConstructor
+public class ConversationOptions {
+    private final @NonNull Double temperature;
+}

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

@@ -0,0 +1,51 @@
+package io.graversen.replicate.service;
+
+import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import lombok.AccessLevel;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
+
+import java.util.UUID;
+
+@Value
+@RequiredArgsConstructor(access = AccessLevel.PACKAGE)
+public class Conversation {
+    @NonNull String id;
+    @NonNull String userId;
+    @NonNull ReplicateModel model;
+    @NonNull TextConversation conversation;
+
+    public static Conversation createDefault(@NonNull String systemMessage, @NonNull ReplicateModel model) {
+        return new Conversation(
+                createId(),
+                "default",
+                model,
+                TextConversation.of(systemMessage)
+        );
+    }
+
+    public Conversation appendMessage(@NonNull TextMessage message) {
+        return new Conversation(
+                getId(),
+                getUserId(),
+                getModel(),
+                getConversation().append(message)
+        );
+    }
+
+    public Conversation update(@NonNull TextConversation conversation) {
+        return new Conversation(
+                getId(),
+                getUserId(),
+                getModel(),
+                conversation
+        );
+    }
+
+    private static String createId() {
+        return String.format("c_%s", UUID.randomUUID().toString());
+    }
+}

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

@@ -0,0 +1,48 @@
+package io.graversen.replicate.service;
+
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+import org.springframework.stereotype.Service;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+@Service
+@RequiredArgsConstructor
+public class ConversationService {
+    private final ConcurrentMap<String, Conversation> conversations = new ConcurrentHashMap<>();
+
+    public Conversation create(@NonNull CreateConversation createConversation) {
+        final var conversation = Conversation.createDefault(createConversation.getSystemMessage(), createConversation.getReplicateModel());
+        conversations.put(conversation.getId(), conversation);
+        return conversation;
+    }
+
+    public Optional<Conversation> appendMessage(@NonNull String id, @NonNull TextMessage message) {
+        final var conversationOrNull = conversations.computeIfPresent(id, (key, conversation) -> conversation.appendMessage(message));
+        return Optional.ofNullable(conversationOrNull);
+    }
+
+    public Optional<Conversation> update(@NonNull String id, @NonNull TextConversation conversation) {
+        final var conversationOrNull = conversations.computeIfPresent(id, (key, value) -> value.update(conversation));
+        return Optional.ofNullable(conversationOrNull);
+    }
+
+    public Optional<Conversation> getById(@NonNull String id) {
+        return Optional.ofNullable(conversations.get(id));
+    }
+
+    public List<Conversation> getByUser(@NonNull String userId) {
+        return conversations.values().stream()
+                .filter(conversation -> conversation.getUserId().equals(userId))
+                .toList();
+    }
+
+    public List<Conversation> getAll() {
+        return List.copyOf(conversations.values());
+    }
+}

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

@@ -0,0 +1,13 @@
+package io.graversen.replicate.service;
+
+import io.graversen.replicate.common.ReplicateModel;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+
+@Getter
+@RequiredArgsConstructor
+public class CreateConversation {
+    private final @NonNull String systemMessage;
+    private final @NonNull ReplicateModel replicateModel;
+}

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

@@ -8,6 +8,7 @@
 import java.time.Duration;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
@@ -73,6 +74,15 @@ public Optional<String> getTextOutput() {
         }
     }
 
+    public Optional<String> getInputKey(@NonNull String key) {
+        try {
+            final var input = (LinkedHashMap<String, Object>) getInput();
+            return Optional.of(((String) input.get(key)));
+        } catch (Exception e) {
+            return Optional.empty();
+        }
+    }
+
     private Function<List<String>, String> composeTextResponse() {
         return strings -> strings.stream().filter(string -> !string.isBlank()).collect(Collectors.joining());
     }