From 72d01bb4198cd1792851f37af1f9885c13e3196f Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Fri, 11 Oct 2024 23:09:13 +0800
Subject: [PATCH] [improvement][chat] Optimize the MemoryReviewTask task by
 adding individual exception handling (#1788)

---
 .../chat/server/memory/MemoryReviewTask.java        | 13 ++++++++-----
 .../chat/server/service/impl/AgentServiceImpl.java  |  3 +--
 .../supersonic/headless/api/pojo/ModelSchema.java   |  3 ++-
 .../headless/chat/corrector/LLMSqlCorrector.java    |  3 ++-
 4 files changed, 13 insertions(+), 9 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 8014a9256..7578b35ae 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
@@ -49,11 +49,14 @@ public class MemoryReviewTask {
 
     @Scheduled(fixedDelay = 60 * 1000)
     public void review() {
-        try {
-            memoryService.getMemoriesForLlmReview().stream().forEach(this::processMemory);
-        } catch (Exception e) {
-            log.error("Exception occurred during memory review task", e);
-        }
+        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);
+            }
+        });
     }
 
     private void processMemory(ChatMemoryDO m) {
diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
index a183a94e0..702c473e6 100644
--- a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
+++ b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
@@ -105,8 +105,7 @@ public class AgentServiceImpl extends ServiceImpl<AgentDOMapper, AgentDO> implem
     }
 
     private synchronized void doExecuteAgentExamples(Agent agent) {
-        if (!agent.containsDatasetTool()
-                || !agent.enableMemoryReview()
+        if (!agent.containsDatasetTool() || !agent.enableMemoryReview()
                 || !ModelConfigHelper.testConnection(
                         ModelConfigHelper.getChatModelConfig(agent, ChatModelType.TEXT_TO_SQL))
                 || CollectionUtils.isEmpty(agent.getExamples())) {
diff --git a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/ModelSchema.java b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/ModelSchema.java
index fdf7750b6..1eafb8a44 100644
--- a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/ModelSchema.java
+++ b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/ModelSchema.java
@@ -1,8 +1,9 @@
 package com.tencent.supersonic.headless.api.pojo;
 
-import java.util.List;
 import lombok.Data;
 
+import java.util.List;
+
 @Data
 public class ModelSchema {
 
diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMSqlCorrector.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMSqlCorrector.java
index 31d1467e3..f1ba1c49b 100644
--- a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMSqlCorrector.java
+++ b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMSqlCorrector.java
@@ -56,7 +56,8 @@ public class LLMSqlCorrector extends BaseSemanticCorrector {
             return;
         }
 
-        ChatLanguageModel chatLanguageModel = ModelProvider.getChatModel(chatQueryContext.getModelConfig());
+        ChatLanguageModel chatLanguageModel =
+                ModelProvider.getChatModel(chatQueryContext.getModelConfig());
         SemanticSqlExtractor extractor =
                 AiServices.create(SemanticSqlExtractor.class, chatLanguageModel);
         Prompt prompt = generatePrompt(chatQueryContext.getQueryText(), semanticParseInfo);