From 1e97b0729b689fdb3de4ea52c737bbfc7c56e1aa Mon Sep 17 00:00:00 2001
From: Georgios Andrianakis <geoand@gmail.com>
Date: Wed, 15 Nov 2023 11:19:15 +0200
Subject: [PATCH] Use @RegisterAiService in AssistantWithToolsResource

---
 .../AssistantWithToolsResource.java           | 43 ++++++++++++++-----
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/integration-tests/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java b/integration-tests/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java
index b5e7961e5..813f87e7a 100644
--- a/integration-tests/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java
+++ b/integration-tests/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java
@@ -1,5 +1,10 @@
 package org.acme.example.openai.aiservices;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import jakarta.annotation.PreDestroy;
+import jakarta.enterprise.context.RequestScoped;
 import jakarta.inject.Singleton;
 import jakarta.ws.rs.GET;
 import jakarta.ws.rs.Path;
@@ -7,22 +12,18 @@
 import org.jboss.resteasy.reactive.RestQuery;
 
 import dev.langchain4j.agent.tool.Tool;
+import dev.langchain4j.memory.ChatMemory;
+import dev.langchain4j.memory.chat.ChatMemoryProvider;
 import dev.langchain4j.memory.chat.MessageWindowChatMemory;
-import dev.langchain4j.model.chat.ChatLanguageModel;
-import dev.langchain4j.service.AiServices;
+import io.quarkiverse.langchain4j.RegisterAiService;
 
 @Path("assistant-with-tool")
 public class AssistantWithToolsResource {
 
     private final Assistant assistant;
 
-    public AssistantWithToolsResource(ChatLanguageModel chatLanguageModel,
-            Calculator calculator) {
-        this.assistant = AiServices.builder(Assistant.class)
-                .chatLanguageModel(chatLanguageModel)
-                .tools(calculator)
-                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
-                .build();
+    public AssistantWithToolsResource(Assistant assistant) {
+        this.assistant = assistant;
     }
 
     @GET
@@ -30,6 +31,12 @@ public String get(@RestQuery String message) {
         return assistant.chat(message);
     }
 
+    @RegisterAiService(tools = Calculator.class, chatMemoryProviderSupplier = RegisterAiService.BeanChatMemoryProviderSupplier.class)
+    public interface Assistant {
+
+        String chat(String userMessage);
+    }
+
     @Singleton
     public static class Calculator {
 
@@ -49,8 +56,22 @@ int add(int a, int b) {
         }
     }
 
-    interface Assistant {
+    @RequestScoped
+    public static class ChatMemoryBean implements ChatMemoryProvider {
 
-        String chat(String userMessage);
+        private final Map<Object, ChatMemory> memories = new ConcurrentHashMap<>();
+
+        @Override
+        public ChatMemory get(Object memoryId) {
+            return memories.computeIfAbsent(memoryId, id -> MessageWindowChatMemory.builder()
+                    .maxMessages(20)
+                    .id(memoryId)
+                    .build());
+        }
+
+        @PreDestroy
+        public void close() {
+            memories.clear();
+        }
     }
 }