From 266f7f0b45faef527bd19a46d183316685df72fb Mon Sep 17 00:00:00 2001
From: Georgios Andrianakis <geoand@gmail.com>
Date: Wed, 13 Dec 2023 17:02:34 +0200
Subject: [PATCH] Make @ApplicationScoped beans work with ChatMemoryRemover

---
 .../java/io/quarkiverse/langchain4j/ChatMemoryRemover.java | 7 +++++--
 .../examples/aiservices/CustomChatMemoryStoreTest.java     | 3 ++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/core/runtime/src/main/java/io/quarkiverse/langchain4j/ChatMemoryRemover.java b/core/runtime/src/main/java/io/quarkiverse/langchain4j/ChatMemoryRemover.java
index 6edf201d3..4c7f49fb2 100644
--- a/core/runtime/src/main/java/io/quarkiverse/langchain4j/ChatMemoryRemover.java
+++ b/core/runtime/src/main/java/io/quarkiverse/langchain4j/ChatMemoryRemover.java
@@ -5,6 +5,7 @@
 import dev.langchain4j.memory.ChatMemory;
 import dev.langchain4j.store.memory.chat.ChatMemoryStore;
 import io.quarkiverse.langchain4j.runtime.aiservice.ChatMemoryRemovable;
+import io.quarkus.arc.ClientProxy;
 
 /**
  * Allows the application to manually control when a {@link ChatMemory} should be removed from the underlying
@@ -22,7 +23,8 @@ private ChatMemoryRemover() {
      * @param memoryId The object used as memory IDs for which the corresponding {@link ChatMemory} should be removed
      */
     public static void remove(Object aiService, Object memoryId) {
-        if (aiService instanceof ChatMemoryRemovable r) {
+        var obj = ClientProxy.unwrap(aiService);
+        if (obj instanceof ChatMemoryRemovable r) {
             r.remove(memoryId);
         }
     }
@@ -32,7 +34,8 @@ public static void remove(Object aiService, Object memoryId) {
      * @param memoryIds The objects used as memory IDs for which the corresponding {@link ChatMemory} should be removed
      */
     public static void remove(Object aiService, List<Object> memoryIds) {
-        if (aiService instanceof ChatMemoryRemovable r) {
+        var obj = ClientProxy.unwrap(aiService);
+        if (obj instanceof ChatMemoryRemovable r) {
             r.remove(memoryIds.toArray(EMPTY_OBJECT_ARRAY));
         }
     }
diff --git a/openai/openai-vanilla/deployment/src/test/java/org/acme/examples/aiservices/CustomChatMemoryStoreTest.java b/openai/openai-vanilla/deployment/src/test/java/org/acme/examples/aiservices/CustomChatMemoryStoreTest.java
index 01e9b9181..dccf6f8c7 100644
--- a/openai/openai-vanilla/deployment/src/test/java/org/acme/examples/aiservices/CustomChatMemoryStoreTest.java
+++ b/openai/openai-vanilla/deployment/src/test/java/org/acme/examples/aiservices/CustomChatMemoryStoreTest.java
@@ -14,6 +14,7 @@
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.enterprise.inject.Produces;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
@@ -81,7 +82,7 @@ void setup() {
     }
 
     @RegisterAiService
-    @Singleton
+    @ApplicationScoped
     interface ChatWithSeparateMemoryForEachUser {
 
         String chat(@MemoryId int memoryId, @UserMessage String userMessage);