From 524ec38edcc9d56737c04d450c1f7117dbd9f1f9 Mon Sep 17 00:00:00 2001
From: yudong <whyuds@163.com>
Date: Fri, 8 Nov 2024 23:06:43 +0800
Subject: [PATCH] =?UTF-8?q?[improvement][chat]=E8=AE=B0=E5=BF=86=E8=AF=84?=
 =?UTF-8?q?=E4=BC=B0=E6=80=A7=E8=83=BD=E4=BC=98=E5=8C=96=20(#1887)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../chat/server/memory/MemoryReviewTask.java  | 32 ++++++++++++-------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/memory/MemoryReviewTask.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/memory/MemoryReviewTask.java
index 21be35ba2..28cbda888 100644
--- a/chat/server/src/main/java/com/tencent/supersonic/chat/server/memory/MemoryReviewTask.java
+++ b/chat/server/src/main/java/com/tencent/supersonic/chat/server/memory/MemoryReviewTask.java
@@ -1,6 +1,7 @@
 package com.tencent.supersonic.chat.server.memory;
 
 import com.tencent.supersonic.chat.api.pojo.enums.MemoryReviewResult;
+import com.tencent.supersonic.chat.api.pojo.request.ChatMemoryFilter;
 import com.tencent.supersonic.chat.server.agent.Agent;
 import com.tencent.supersonic.chat.server.persistence.dataobject.ChatMemoryDO;
 import com.tencent.supersonic.chat.server.service.AgentService;
@@ -24,6 +25,7 @@ import java.util.Collections;
 import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.List;
 
 @Component
 @Slf4j
@@ -57,24 +59,30 @@ public class MemoryReviewTask {
 
     @Scheduled(fixedDelay = 60 * 1000)
     public void review() {
-        memoryService.getMemoriesForLlmReview().stream().forEach(memory -> {
-            try {
-                processMemory(memory);
-            } catch (Exception e) {
-                log.error("Exception occurred while processing memory with id {}: {}",
-                        memory.getId(), e.getMessage(), e);
+        List<Agent> agentList = agentService.getAgents();
+        for (Agent agent : agentList) {
+            if(!agent.enableMemoryReview()){
+                continue;
             }
-        });
+            ChatMemoryFilter chatMemoryFilter = ChatMemoryFilter.builder().agentId(agent.getId()).build();
+            memoryService.getMemories(chatMemoryFilter).stream().forEach(memory -> {
+                try {
+                    processMemory(memory, agent);
+                } catch (Exception e) {
+                    log.error("Exception occurred while processing memory with id {}: {}",
+                            memory.getId(), e.getMessage(), e);
+                }
+            });
+        }
     }
 
-    private void processMemory(ChatMemoryDO m) {
-        Agent chatAgent = agentService.getAgent(m.getAgentId());
-        if (Objects.isNull(chatAgent)) {
+    private void processMemory(ChatMemoryDO m, Agent agent) {
+        if (Objects.isNull(agent)) {
             log.warn("Agent id {} not found or memory review disabled", m.getAgentId());
             return;
         }
 
-        ChatApp chatApp = chatAgent.getChatAppConfig().get(APP_KEY);
+        ChatApp chatApp = agent.getChatAppConfig().get(APP_KEY);
         if (Objects.isNull(chatApp) || !chatApp.isEnable()) {
             return;
         }
@@ -90,7 +98,7 @@ public class MemoryReviewTask {
                     response);
             processResponse(response, m);
         } else {
-            log.debug("ChatLanguageModel not found for agent:{}", chatAgent.getId());
+            log.debug("ChatLanguageModel not found for agent:{}", agent.getId());
         }
     }