[improvement][chat]记忆评估性能优化

tencentmusic · Nov 5, 2024 · fdfdc49 · fdfdc49
1 parent 9edcb9f
commit fdfdc49
Showing 1 changed file with 20 additions and 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
@@ -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.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.List;
 
 @Component
 @Slf4j
@@ -57,24 +59,30 @@ public 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 @@ private void processMemory(ChatMemoryDO m) {
                     response);
             processResponse(response, m);
         } else {
-            log.debug("ChatLanguageModel not found for agent:{}", chatAgent.getId());
+            log.debug("ChatLanguageModel not found for agent:{}", agent.getId());
         }
     }