From 9cf98dbaf70491ea738586c3e72c97f1ca0aaccc Mon Sep 17 00:00:00 2001 From: Mahesh Kommareddi Date: Sat, 20 Jul 2024 16:25:20 -0400 Subject: [PATCH] Agents now have memory in addition to context with Tree of thought --- src/main/java/com/ioa/agent/AgentInfo.java | 25 +++++++++- src/main/java/com/ioa/agent/Memory.java | 47 +++++++++++++++++++ .../com/ioa/service/WebSocketService.java | 18 +++++++ src/main/java/com/ioa/util/TreeOfThought.java | 8 +--- 4 files changed, 90 insertions(+), 8 deletions(-) create mode 100644 src/main/java/com/ioa/agent/Memory.java diff --git a/src/main/java/com/ioa/agent/AgentInfo.java b/src/main/java/com/ioa/agent/AgentInfo.java index 472635c..4b7a676 100644 --- a/src/main/java/com/ioa/agent/AgentInfo.java +++ b/src/main/java/com/ioa/agent/AgentInfo.java @@ -15,6 +15,7 @@ public class AgentInfo { private String name; private List capabilities; private List tools; + private Memory memory; private TreeOfThought treeOfThought; private WebSocketService webSocketService; private ToolRegistry toolRegistry; @@ -31,16 +32,36 @@ public class AgentInfo { this.webSocketService = webSocketService; this.toolRegistry = toolRegistry; this.model = model; + this.memory = new Memory(); } public void receiveMessage(Message message) { - // Process the received message + memory.addToHistory(message.getContent()); String prompt = "You are " + name + " with capabilities: " + capabilities + "\nYou received a message: " + message.getContent() + - "\nHow would you respond or what actions would you take based on this message?"; + "\nBased on your memory and context, how would you respond or what actions would you take?" + + "\n\nMemory:\n" + memory.getFormattedMemory(); String response = model.generate(prompt, null); System.out.println("DEBUG: " + name + " processed message: " + message.getContent()); System.out.println("DEBUG: " + name + " response: " + response); + + // Add the response to memory + memory.addToHistory("My response: " + response); + } + + public void performTreeOfThought(String task) { + String prompt = "You are " + name + " with capabilities: " + capabilities + + "\nTask: " + task + + "\nBased on your memory and context, perform a tree of thought reasoning to approach this task." + + "\n\nMemory:\n" + memory.getFormattedMemory(); + + Map totResult = treeOfThought.reason(prompt, 3, 2); + String reasoning = (String) totResult.get("reasoning"); + + // Add the reasoning to memory + memory.addContextualFact("Tree of Thought for task '" + task + "': " + reasoning); + + System.out.println("DEBUG: " + name + " Tree of Thought reasoning: " + reasoning); } public void notifyTurn(ConversationFSM conversation) { diff --git a/src/main/java/com/ioa/agent/Memory.java b/src/main/java/com/ioa/agent/Memory.java new file mode 100644 index 0000000..385f09b --- /dev/null +++ b/src/main/java/com/ioa/agent/Memory.java @@ -0,0 +1,47 @@ +package com.ioa.agent; + +import java.util.ArrayList; +import java.util.List; + +public class Memory { + private List conversationHistory; + private List contextualFacts; + private int maxHistorySize = 10; // Adjust as needed + + public Memory() { + this.conversationHistory = new ArrayList<>(); + this.contextualFacts = new ArrayList<>(); + } + + public void addToHistory(String message) { + conversationHistory.add(message); + if (conversationHistory.size() > maxHistorySize) { + conversationHistory.remove(0); + } + } + + public void addContextualFact(String fact) { + contextualFacts.add(fact); + } + + public List getConversationHistory() { + return new ArrayList<>(conversationHistory); + } + + public List getContextualFacts() { + return new ArrayList<>(contextualFacts); + } + + public String getFormattedMemory() { + StringBuilder sb = new StringBuilder(); + sb.append("Conversation History:\n"); + for (String message : conversationHistory) { + sb.append("- ").append(message).append("\n"); + } + sb.append("\nContextual Facts:\n"); + for (String fact : contextualFacts) { + sb.append("- ").append(fact).append("\n"); + } + return sb.toString(); + } +} \ No newline at end of file diff --git a/src/main/java/com/ioa/service/WebSocketService.java b/src/main/java/com/ioa/service/WebSocketService.java index 38ce9c8..60fe75f 100644 --- a/src/main/java/com/ioa/service/WebSocketService.java +++ b/src/main/java/com/ioa/service/WebSocketService.java @@ -8,6 +8,8 @@ import org.springframework.beans.factory.annotation.Autowired; import org.springframework.messaging.simp.SimpMessagingTemplate; import org.springframework.stereotype.Service; import org.springframework.context.annotation.Lazy; +import org.springframework.web.socket.TextMessage; +import org.springframework.web.socket.WebSocketSession; @Service public class WebSocketService { @@ -22,6 +24,22 @@ public class WebSocketService { this.objectMapper = new ObjectMapper(); } + public void handleMessage(WebSocketSession session, TextMessage message) { + try { + String[] messageArray = objectMapper.readValue(message.getPayload(), String[].class); + if (messageArray.length > 0) { + String content = messageArray[0]; + // Assume we're using a default conversation ID for simplicity + String conversationId = "default"; + // Assume we're using a default user ID for simplicity + String userId = "user"; + conversationManager.postMessage(conversationId, userId, content); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + public void sendUpdate(String topic, Object payload) { messagingTemplate.convertAndSend("/topic/" + topic, payload); } diff --git a/src/main/java/com/ioa/util/TreeOfThought.java b/src/main/java/com/ioa/util/TreeOfThought.java index 164a86a..a87e5dc 100644 --- a/src/main/java/com/ioa/util/TreeOfThought.java +++ b/src/main/java/com/ioa/util/TreeOfThought.java @@ -45,12 +45,8 @@ public class TreeOfThought { "\nExplore a new branch of thought (branch " + (i+1) + "/" + branches + "):"; String thought = model.generate(branchPrompt, null); - if (!thought.equals("No response generated") && !thought.startsWith("Error:")) { - Map childNode = exploreThought(task, depth - 1, branches, thought); - children.add(childNode); - } else { - System.out.println("WARNING: Failed to generate thought. Result: " + thought); - } + Map childNode = exploreThought(task, depth - 1, branches, thought); + children.add(childNode); } node.put("children", children);