From 868cdb3b7886712d54a6e6ba31ac3beb23234777 Mon Sep 17 00:00:00 2001
From: Georgios Andrianakis <geoand@gmail.com>
Date: Thu, 14 Dec 2023 16:51:42 +0200
Subject: [PATCH] Make sure ChatMemoryStore beans are unremovable

Fixes: #158
---
 .../langchain4j/deployment/ChatMemoryProcessor.java            | 3 +++
 .../runtime/aiservice/InMemoryChatMemoryStoreProducer.java     | 2 --
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ChatMemoryProcessor.java b/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ChatMemoryProcessor.java
index 05d56a82d..a4f9d3d72 100644
--- a/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ChatMemoryProcessor.java
+++ b/core/deployment/src/main/java/io/quarkiverse/langchain4j/deployment/ChatMemoryProcessor.java
@@ -13,6 +13,7 @@
 import io.quarkiverse.langchain4j.runtime.aiservice.ChatMemoryConfig;
 import io.quarkus.arc.SyntheticCreationalContext;
 import io.quarkus.arc.deployment.SyntheticBeanBuildItem;
+import io.quarkus.arc.deployment.UnremovableBeanBuildItem;
 import io.quarkus.deployment.annotations.BuildProducer;
 import io.quarkus.deployment.annotations.BuildStep;
 import io.quarkus.deployment.annotations.ExecutionTime;
@@ -24,6 +25,7 @@ public class ChatMemoryProcessor {
     @Record(ExecutionTime.RUNTIME_INIT)
     void setupBeans(ChatMemoryBuildConfig buildConfig, ChatMemoryConfig runtimeConfig,
             ChatMemoryRecorder recorder,
+            BuildProducer<UnremovableBeanBuildItem> unremovableProducer,
             BuildProducer<SyntheticBeanBuildItem> syntheticBeanProducer) {
 
         Function<SyntheticCreationalContext<ChatMemoryProvider>, ChatMemoryProvider> fun;
@@ -47,5 +49,6 @@ void setupBeans(ChatMemoryBuildConfig buildConfig, ChatMemoryConfig runtimeConfi
         configurator.createWith(fun);
 
         syntheticBeanProducer.produce(configurator.done());
+        unremovableProducer.produce(UnremovableBeanBuildItem.beanTypes(ChatMemoryStore.class));
     }
 }
diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/InMemoryChatMemoryStoreProducer.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/InMemoryChatMemoryStoreProducer.java
index 500b31a4d..694c2de76 100644
--- a/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/InMemoryChatMemoryStoreProducer.java
+++ b/core/runtime/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/InMemoryChatMemoryStoreProducer.java
@@ -6,7 +6,6 @@
 import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
 import io.quarkiverse.langchain4j.RegisterAiService;
 import io.quarkus.arc.DefaultBean;
-import io.quarkus.arc.Unremovable;
 
 /**
  * Creates the default {@link InMemoryChatMemoryStore} store to be used by classes annotated with {@link RegisterAiService}
@@ -16,7 +15,6 @@ public class InMemoryChatMemoryStoreProducer {
     @Produces
     @Singleton
     @DefaultBean
-    @Unremovable
     public InMemoryChatMemoryStore chatMemoryStore() {
         return new InMemoryChatMemoryStore();
     }