From 7dbd23d75170ecd8575c5aded6a2b6a61fe65352 Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Thu, 5 Sep 2024 16:10:45 +0800
Subject: [PATCH] (improvement)(chat) Fixed the issue where setting an
 incorrect embeddingStore address in the system settings caused the service to
 fail to start (#1631)

---
 .../chat/server/memory/MemoryReviewTask.java  | 71 +++++++++++--------
 .../service/impl/ExemplarServiceImpl.java     |  2 +-
 2 files changed, 42 insertions(+), 31 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 f9802a5d6..ad3bea715 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
@@ -2,6 +2,7 @@ package com.tencent.supersonic.chat.server.memory;
 
 import com.tencent.supersonic.chat.api.pojo.enums.MemoryReviewResult;
 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;
 import com.tencent.supersonic.chat.server.service.MemoryService;
 import dev.langchain4j.model.chat.ChatLanguageModel;
@@ -50,37 +51,47 @@ public class MemoryReviewTask {
 
     @Scheduled(fixedDelay = 60 * 1000)
     public void review() {
-        memoryService.getMemoriesForLlmReview().stream()
-                .forEach(m -> {
-                    Agent chatAgent = agentService.getAgent(m.getAgentId());
-                    if (Objects.nonNull(chatAgent) && chatAgent.enableMemoryReview()) {
-                        String promptStr = String.format(INSTRUCTION, m.getQuestion(), m.getDbSchema(),
-                                m.getSideInfo(), m.getS2sql());
-                        Prompt prompt = PromptTemplate.from(promptStr).apply(Collections.EMPTY_MAP);
+        try {
+            memoryService.getMemoriesForLlmReview().stream().forEach(this::processMemory);
+        } catch (Exception e) {
+            log.error("Exception occurred during memory review task", e);
+        }
+    }
 
-                        keyPipelineLog.info("MemoryReviewTask reqPrompt:\n{}", promptStr);
-                        ChatLanguageModel chatLanguageModel = ModelProvider.getChatModel(
-                                chatAgent.getModelConfig());
-                        if (Objects.nonNull(chatLanguageModel)) {
-                            String response = chatLanguageModel.generate(prompt.toUserMessage()).content().text();
-                            keyPipelineLog.info("MemoryReviewTask modelResp:\n{}", response);
+    private void processMemory(ChatMemoryDO m) {
+        Agent chatAgent = agentService.getAgent(m.getAgentId());
+        if (Objects.isNull(chatAgent) || !chatAgent.enableMemoryReview()) {
+            log.debug("Agent id {} not found or memory review disabled", m.getAgentId());
+            return;
+        }
+        String promptStr = createPromptString(m);
+        Prompt prompt = PromptTemplate.from(promptStr).apply(Collections.EMPTY_MAP);
 
-                            Matcher matcher = OUTPUT_PATTERN.matcher(response);
-                            if (matcher.find()) {
-                                m.setLlmReviewRet(MemoryReviewResult.getMemoryReviewResult(matcher.group(1)));
-                                m.setLlmReviewCmt(matcher.group(2));
-                                // directly enable memory if the LLM determines it positive
-                                if (MemoryReviewResult.POSITIVE.equals(m.getLlmReviewRet())) {
-                                    memoryService.enableMemory(m);
-                                }
-                                memoryService.updateMemory(m);
-                            }
-                        } else {
-                            log.debug("ChatLanguageModel not found for agent:{}", chatAgent.getId());
-                        }
-                    } else {
-                        log.debug("Agent id {} not found or memory review disabled", m.getAgentId());
-                    }
-                });
+        keyPipelineLog.info("MemoryReviewTask reqPrompt:\n{}", promptStr);
+        ChatLanguageModel chatLanguageModel = ModelProvider.getChatModel(chatAgent.getModelConfig());
+        if (Objects.nonNull(chatLanguageModel)) {
+            String response = chatLanguageModel.generate(prompt.toUserMessage()).content().text();
+            keyPipelineLog.info("MemoryReviewTask modelResp:\n{}", response);
+            processResponse(response, m);
+        } else {
+            log.debug("ChatLanguageModel not found for agent:{}", chatAgent.getId());
+        }
+    }
+
+    private String createPromptString(ChatMemoryDO m) {
+        return String.format(INSTRUCTION, m.getQuestion(), m.getDbSchema(), m.getSideInfo(), m.getS2sql());
+    }
+
+    private void processResponse(String response, ChatMemoryDO m) {
+        Matcher matcher = OUTPUT_PATTERN.matcher(response);
+        if (matcher.find()) {
+            m.setLlmReviewRet(MemoryReviewResult.getMemoryReviewResult(matcher.group(1)));
+            m.setLlmReviewCmt(matcher.group(2));
+            // directly enable memory if the LLM determines it positive
+            if (MemoryReviewResult.POSITIVE.equals(m.getLlmReviewRet())) {
+                memoryService.enableMemory(m);
+            }
+            memoryService.updateMemory(m);
+        }
     }
 }
diff --git a/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java b/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
index 5bac8bf5e..8483a8ff8 100644
--- a/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
+++ b/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
@@ -83,7 +83,7 @@ public class ExemplarServiceImpl implements ExemplarService, CommandLineRunner {
     public void run(String... args) {
         try {
             loadSysExemplars();
-        } catch (IOException e) {
+        } catch (Exception e) {
             log.error("Failed to load system exemplars", e);
         }
     }