From 2801b27adee548a637c2fb25c3d154474842bf8b Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Sun, 18 Aug 2024 00:16:55 +0800
Subject: [PATCH] (improvement)(chat) Fix the issue where the query gets stuck
 when two terms appear simultaneously. (#1581)

---
 .../headless/chat/mapper/BaseMapper.java      | 58 +++++++++----------
 .../server/utils/ChatWorkflowEngine.java      |  2 -
 2 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/BaseMapper.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/BaseMapper.java
index 2fc0713c8..0a1b6fa9e 100644
--- a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/BaseMapper.java
+++ b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/BaseMapper.java
@@ -95,43 +95,41 @@ public abstract class BaseMapper implements SchemaMapper {
 
     private static void filterByQueryDataType(ChatQueryContext chatQueryContext,
                                               Predicate<SchemaElement> needRemovePredicate) {
-        chatQueryContext.getMapInfo().getDataSetElementMatches().values().stream().forEach(
-                schemaElementMatches -> schemaElementMatches.removeIf(
-                        schemaElementMatch -> {
-                            SchemaElement element = schemaElementMatch.getElement();
-                            SchemaElementType type = element.getType();
-                            if (SchemaElementType.ENTITY.equals(type) || SchemaElementType.DATASET.equals(type)
-                                    || SchemaElementType.ID.equals(type)) {
-                                return false;
-                            }
-                            return needRemovePredicate.test(element);
-                        }
-                ));
+        chatQueryContext.getMapInfo().getDataSetElementMatches().values().forEach(schemaElementMatches -> {
+            schemaElementMatches.removeIf(schemaElementMatch -> {
+                SchemaElement element = schemaElementMatch.getElement();
+                SchemaElementType type = element.getType();
+
+                boolean isEntityOrDatasetOrId = SchemaElementType.ENTITY.equals(type)
+                        || SchemaElementType.DATASET.equals(type)
+                        || SchemaElementType.ID.equals(type);
+
+                return !isEntityOrDatasetOrId && needRemovePredicate.test(element);
+            });
+        });
     }
 
     public abstract void doMap(ChatQueryContext chatQueryContext);
 
     public void addToSchemaMap(SchemaMapInfo schemaMap, Long dataSetId, SchemaElementMatch newElementMatch) {
         Map<Long, List<SchemaElementMatch>> dataSetElementMatches = schemaMap.getDataSetElementMatches();
-        List<SchemaElementMatch> schemaElementMatches = dataSetElementMatches.putIfAbsent(dataSetId, new ArrayList<>());
-        if (schemaElementMatches == null) {
-            schemaElementMatches = dataSetElementMatches.get(dataSetId);
-        }
-        //remove duplication
-        AtomicBoolean needAddNew = new AtomicBoolean(true);
-        schemaElementMatches.removeIf(
-                existElementMatch -> {
-                    if (isEquals(existElementMatch, newElementMatch)) {
-                        if (newElementMatch.getSimilarity() > existElementMatch.getSimilarity()) {
-                            return true;
-                        } else {
-                            needAddNew.set(false);
-                        }
-                    }
-                    return false;
+        List<SchemaElementMatch> schemaElementMatches = dataSetElementMatches.computeIfAbsent(dataSetId,
+                k -> new ArrayList<>());
+
+        AtomicBoolean shouldAddNew = new AtomicBoolean(true);
+
+        schemaElementMatches.removeIf(existingElementMatch -> {
+            if (isEquals(existingElementMatch, newElementMatch)) {
+                if (newElementMatch.getSimilarity() > existingElementMatch.getSimilarity()) {
+                    return true;
+                } else {
+                    shouldAddNew.set(false);
                 }
-        );
-        if (needAddNew.get()) {
+            }
+            return false;
+        });
+
+        if (shouldAddNew.get()) {
             schemaElementMatches.add(newElementMatch);
         }
     }
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java
index 46ca010d2..f5ecc607d 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java
@@ -47,8 +47,6 @@ public class ChatWorkflowEngine {
                         parseResult.setState(ParseResp.ParseState.FAILED);
                         parseResult.setErrorMsg("No semantic entities can be mapped against user question.");
                         queryCtx.setChatWorkflowState(ChatWorkflowState.FINISHED);
-                    } else if (queryCtx.getMapInfo().needContinueMap()) {
-                        queryCtx.setChatWorkflowState(ChatWorkflowState.MAPPING);
                     } else {
                         queryCtx.setChatWorkflowState(ChatWorkflowState.PARSING);
                     }