Use @RegisterAiService in AssistantWithToolsResource

quarkiverse · Nov 15, 2023 · 1e97b07 · 1e97b07
1 parent 8126e4c
commit 1e97b07
Showing 1 changed file with 32 additions and 11 deletions.
diff --git a/...s/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java b/...s/openai/src/main/java/org/acme/example/openai/aiservices/AssistantWithToolsResource.java
@@ -1,35 +1,42 @@
 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;
 
 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
     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();
+        }
     }
 }