Merge pull request #4 from muhrifqii/dev/chat-memory

muhrifqii · web-flow · commit 76a1cc5b71a3 · 2024-09-04T23:54:43.000+07:00
Chat Memory
diff --git a/README.md b/README.md
@@ -11,6 +11,17 @@ Sample project to showcase a java spring boot application with Generative AI pow
 - Flyway
 - [Ollama](https://ollama.com)
 
+## Prerequisites - Ollama
+
+- Install Ollama on your local machine
+- Run `ollama run llama3.1`
+
+## Running the project
+
+```bash
+make up
+```
+
 ## Video on The Stream Chat Response Feature
 
 
diff --git a/llm/api/src/main/java/com/muhrifqii/llm/api/usecases/MessageStoreUsecase.java b/llm/api/src/main/java/com/muhrifqii/llm/api/usecases/MessageStoreUsecase.java
@@ -1,15 +1,12 @@
 package com.muhrifqii.llm.api.usecases;
 
 import com.muhrifqii.llm.api.datamodels.conversations.Message;
-import com.muhrifqii.llm.api.datamodels.conversations.UserMessage;
 
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 public interface MessageStoreUsecase {
     Flux<Message> getMessages(String conversationId, String cursor);
 
-    Mono<Message> saveUserMessage(UserMessage userMessage);
-
-    Mono<Message> saveAssistantMessage(Message message);
+    Mono<Message> saveMessage(Message message, boolean newEntity);
 }
diff --git a/llm/api/src/main/java/com/muhrifqii/llm/api/utils/DateUtils.java b/llm/api/src/main/java/com/muhrifqii/llm/api/utils/DateUtils.java
@@ -30,4 +30,11 @@ public static String toIsoString(LocalDateTime date) {
                 .map(DateTimeFormatter.ISO_DATE_TIME::format)
                 .orElse("");
     }
+
+    public static LocalDateTime fromIsoString(String date) {
+        return Optional.ofNullable(date)
+                .map(DateTimeFormatter.ISO_DATE_TIME::parse)
+                .map(LocalDateTime::from)
+                .orElse(null);
+    }
 }
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/configurations/ChatModelConfig.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/configurations/ChatModelConfig.java
@@ -1,26 +1,37 @@
 package com.muhrifqii.llm.configurations;
 
+import java.time.Duration;
+
 import org.springframework.ai.chat.client.ChatClient;
-import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
 import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
-import org.springframework.ai.chat.memory.ChatMemory;
+import org.springframework.boot.web.client.ClientHttpRequestFactories;
+import org.springframework.boot.web.client.ClientHttpRequestFactorySettings;
+import org.springframework.boot.web.client.RestClientCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import com.muhrifqii.llm.annotations.MemCachedChatMemory;
+import com.muhrifqii.llm.services.MessageStoreAdvisor;
 
 @Configuration
 public class ChatModelConfig {
     @Bean
     ChatClient chatClient(
             ChatClient.Builder builder,
-            @MemCachedChatMemory ChatMemory chatMemory) {
+            MessageStoreAdvisor messageStoreAdvisor) {
         return builder
                 .defaultSystem(
                         "You are a Pokédex AI named Slaking.AI, a highly advanced AI that specializes in providing detailed and accurate information about Pokémon. You have access to all known data about Pokémon species, including their types, abilities, evolutions, habitat, and more. Your responses should be concise, factual, and directly related to the Pokémon in question. Ensure to offer relevant insights based on the user's query, and avoid speculation. Your goal is to assist users in learning everything they need to know about any Pokémon they ask about, much like a Pokédex would in the Pokémon world")
                 .defaultAdvisors(
-                        new SimpleLoggerAdvisor(),
-                        new MessageChatMemoryAdvisor(chatMemory))
+                        messageStoreAdvisor,
+                        new SimpleLoggerAdvisor())
                 .build();
     }
+
+    @Bean
+    RestClientCustomizer restClientCustomizer() {
+        return restClientBuilder -> restClientBuilder
+                .requestFactory(ClientHttpRequestFactories.get(ClientHttpRequestFactorySettings.DEFAULTS
+                        .withConnectTimeout(Duration.ofSeconds(180))
+                        .withReadTimeout(Duration.ofSeconds(180))));
+    }
 }
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/repositories/MessageEntity.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/repositories/MessageEntity.java
@@ -7,21 +7,29 @@
 import org.springframework.ai.chat.messages.MessageType;
 import org.springframework.data.annotation.CreatedDate;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.Transient;
 import org.springframework.data.domain.Persistable;
 import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.lang.Nullable;
 
-import lombok.Builder;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.experimental.Accessors;
 
 @Table("ai_messages")
-@Builder
-public record MessageEntity(
-        @Id String id,
-        String coversationId,
-        String content,
-        String messageType,
-        @CreatedDate LocalDateTime createdAt)
-        implements Persistable<String>, Message {
+@Getter
+@Setter
+@Accessors(fluent = true)
+public class MessageEntity implements Persistable<String>, Message {
+    @Id
+    private String id;
+    private String conversationId;
+    private String content;
+    private String messageType;
+    @CreatedDate
+    private LocalDateTime createdAt;
+    @Transient
+    private boolean newEntity;
 
     @Override
     @Nullable
@@ -31,7 +39,7 @@ public String getId() {
 
     @Override
     public boolean isNew() {
-        return createdAt == null;
+        return createdAt == null || newEntity;
     }
 
     @Override
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ChatHelper.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ChatHelper.java
@@ -35,7 +35,6 @@ public static Message newMessage(String conversationID) {
                 .id(generateId())
                 .conversationId(conversationID)
                 .messageType(MessageType.USER.getValue())
-                .createdAt(DateUtils.nowIsoString())
                 .build();
     }
 
@@ -81,20 +80,40 @@ public static com.muhrifqii.llm.repositories.Conversation updateConversation(
     public static Message mapMessage(MessageEntity source) {
         return Message.builder()
                 .id(source.id())
-                .conversationId(source.coversationId())
+                .conversationId(source.conversationId())
                 .content(source.content())
                 .messageType(source.messageType())
                 .createdAt(DateUtils.toIsoString(source.createdAt()))
                 .build();
     }
 
-    public static MessageEntity mapFromAi(String conversationId, org.springframework.ai.chat.messages.Message source) {
-        return MessageEntity.builder()
+    public static MessageEntity mapMessage(Message source, boolean newEntity) {
+        return new MessageEntity()
+                .id(source.id())
+                .conversationId(source.conversationId())
+                .content(source.content())
+                .messageType(source.messageType())
+                .createdAt(DateUtils.fromIsoString(source.createdAt()))
+                .newEntity(newEntity);
+    }
+
+    public static Message mapFromAi(String conversationId, org.springframework.ai.chat.messages.Message source) {
+        return Message.builder()
+                .id(generateId())
+                .conversationId(conversationId)
+                .content(source.getContent())
+                .messageType(source.getMessageType().getValue())
+                .createdAt(DateUtils.nowIsoString())
+                .build();
+    }
+
+    public static Message mapFromAi(String conversationId, org.springframework.ai.chat.messages.UserMessage source) {
+        return Message.builder()
                 .id(generateId())
-                .coversationId(conversationId)
+                .conversationId(conversationId)
                 .content(source.getContent())
                 .messageType(source.getMessageType().getValue())
-                .createdAt(DateUtils.now())
+                .createdAt(DateUtils.nowIsoString())
                 .build();
     }
 
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ChatService.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ChatService.java
@@ -27,7 +27,7 @@ public class ChatService implements ChatServiceTrait {
 
     @Override
     public Mono<Message> chat(String conversationID, UserMessage message) {
-        return Mono.fromCallable(() -> chatClient
+        return Mono.fromSupplier(() -> chatClient
                 .prompt()
                 .user(message.content())
                 .advisors(advSpec -> chatMemoryAdvisorSpec(advSpec, conversationID))
@@ -71,7 +71,7 @@ private Flux<Message> chatPromptStream(Message source, UserMessage userMessage)
     }
 
     private void chatMemoryAdvisorSpec(AdvisorSpec advisorSpec, String conversationId) {
-        if (Constants.EMPTY_SLUG.equals(conversationId)) {
+        if (Constants.EMPTY_SLUG.equals(conversationId) || conversationId == null) {
             return;
         }
         advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId);
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ConversationService.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/ConversationService.java
@@ -10,7 +10,6 @@
 import com.muhrifqii.llm.api.datamodels.HandledError;
 import com.muhrifqii.llm.api.datamodels.conversations.Conversation;
 import com.muhrifqii.llm.api.datamodels.conversations.Message;
-import com.muhrifqii.llm.api.datamodels.conversations.UserMessage;
 import com.muhrifqii.llm.api.traits.ConversationServiceTrait;
 import com.muhrifqii.llm.repositories.ConversationRepository;
 import com.muhrifqii.llm.repositories.MessageRepository;
@@ -74,13 +73,9 @@ public Flux<Message> getMessages(String conversationId, String cursor) {
     }
 
     @Override
-    public Mono<Message> saveAssistantMessage(Message message) {
-        return null;
-    }
-
-    @Override
-    public Mono<Message> saveUserMessage(UserMessage userMessage) {
-        return null;
+    public Mono<Message> saveMessage(Message message, boolean newEntity) {
+        return messageRepository.save(ChatHelper.mapMessage(message, newEntity))
+                .map(ChatHelper::mapMessage);
     }
 
     private Mono<Conversation> throwIfConversationIdNotExist(String id) {
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/DbPersistedChatMemory.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/DbPersistedChatMemory.java
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/MemcachedChatMemory.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/MemcachedChatMemory.java
@@ -5,6 +5,7 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
+import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
 import org.springframework.ai.chat.memory.ChatMemory;
 import org.springframework.ai.chat.messages.Message;
 import org.springframework.stereotype.Component;
@@ -22,22 +23,40 @@ public class MemcachedChatMemory implements ChatMemory {
 
     @Override
     public void add(String conversationId, List<Message> messages) {
+        if (!isConversationIdValid(conversationId)) {
+            return;
+        }
+
         log.debug("add:{} with {} messages", conversationId, messages.size());
         this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
         this.conversationHistory.get(conversationId).addAll(messages);
     }
 
     @Override
     public List<Message> get(String conversationId, int lastN) {
+        if (!isConversationIdValid(conversationId)) {
+            return List.of();
+        }
+
         log.debug("get:{}:{}", lastN, conversationId);
         List<Message> all = this.conversationHistory.get(conversationId);
-        return all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).toList() : List.of();
+        return all != null ? all.stream()
+                .skip(Math.max(0, all.size() - lastN))
+                .toList() : List.of();
     }
 
     @Override
     public void clear(String conversationId) {
+        if (!isConversationIdValid(conversationId)) {
+            return;
+        }
+
         log.debug("clear:{}", conversationId);
         this.conversationHistory.remove(conversationId);
     }
 
+    private boolean isConversationIdValid(String conversationId) {
+        return !AbstractChatMemoryAdvisor.DEFAULT_CHAT_MEMORY_CONVERSATION_ID.equals(conversationId);
+    }
+
 }
diff --git a/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/MessageStoreAdvisor.java b/llm/ollama-provider/src/main/java/com/muhrifqii/llm/services/MessageStoreAdvisor.java
diff --git a/llm/ollama-provider/src/main/resources/application.properties b/llm/ollama-provider/src/main/resources/application.properties

Original file line number	Diff line number	Diff line change
`@@ -30,4 +30,11 @@ public static String toIsoString(LocalDateTime date) {`
`30`	`30`	`.map(DateTimeFormatter.ISO_DATE_TIME::format)`
`31`	`31`	`.orElse("");`
`32`	`32`	`}`
	`33`	`+`
	`34`	`+ public static LocalDateTime fromIsoString(String date) {`
	`35`	`+ return Optional.ofNullable(date)`
	`36`	`+ .map(DateTimeFormatter.ISO_DATE_TIME::parse)`
	`37`	`+ .map(LocalDateTime::from)`
	`38`	`+ .orElse(null);`
	`39`	`+ }`
`33`	`40`	`}`