From 532a00518cb4000b958229f83a46c7c0b3b855b8 Mon Sep 17 00:00:00 2001
From: LXW <1264174498@qq.com>
Date: Thu, 29 Feb 2024 12:45:11 +0800
Subject: [PATCH] (improvement)(chat) Fix agent filter rule query (#777)

Co-authored-by: jolunoluo
---
 .../parser/sql/rule/AgentCheckParser.java     | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/core/parser/sql/rule/AgentCheckParser.java b/chat/core/src/main/java/com/tencent/supersonic/chat/core/parser/sql/rule/AgentCheckParser.java
index 9031c6fb5..3fd10dfdd 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/core/parser/sql/rule/AgentCheckParser.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/core/parser/sql/rule/AgentCheckParser.java
@@ -23,20 +23,22 @@ public class AgentCheckParser implements SemanticParser {
     @Override
     public void parse(QueryContext queryContext, ChatContext chatContext) {
         List<SemanticQuery> queries = queryContext.getCandidateQueries();
-        agentCanSupport(queryContext, queries);
+        log.info("query size before agent filter:{}", queryContext.getCandidateQueries().size());
+        filterQueries(queryContext, queries);
+        log.info("query size after agent filter: {}", queryContext.getCandidateQueries().size());
     }
 
-    private void agentCanSupport(QueryContext queryContext, List<SemanticQuery> queries) {
+    private void filterQueries(QueryContext queryContext, List<SemanticQuery> queries) {
         Agent agent = queryContext.getAgent();
         if (agent == null) {
             return;
         }
         List<RuleParserTool> queryTools = getRuleTools(agent);
         if (CollectionUtils.isEmpty(queryTools)) {
-            queries.clear();
+            queryContext.setCandidateQueries(Lists.newArrayList());
             return;
         }
-        log.info("queries resolved:{} {}", agent.getName(),
+        log.info("agent name :{}, queries resolved: {}", agent.getName(),
                 queries.stream().map(SemanticQuery::getQueryMode).collect(Collectors.toList()));
         queries.removeIf(query -> {
             for (RuleParserTool tool : queryTools) {
@@ -46,10 +48,14 @@ private void agentCanSupport(QueryContext queryContext, List<SemanticQuery> quer
                 }
                 if (CollectionUtils.isNotEmpty(tool.getQueryTypes())) {
                     if (QueryManager.isTagQuery(query.getQueryMode())) {
-                        return !tool.getQueryTypes().contains(QueryType.TAG.name());
+                        if (!tool.getQueryTypes().contains(QueryType.TAG.name())) {
+                            return true;
+                        }
                     }
                     if (QueryManager.isMetricQuery(query.getQueryMode())) {
-                        return !tool.getQueryTypes().contains(QueryType.METRIC.name());
+                        if (!tool.getQueryTypes().contains(QueryType.METRIC.name())) {
+                            return true;
+                        }
                     }
                 }
                 if (CollectionUtils.isEmpty(tool.getViewIds())) {
@@ -62,7 +68,8 @@ private void agentCanSupport(QueryContext queryContext, List<SemanticQuery> quer
             }
             return true;
         });
-        log.info("rule queries witch can be supported by agent :{} {}", agent.getName(),
+        queryContext.setCandidateQueries(queries);
+        log.info("agent name :{}, rule queries witch can be supported by agent :{}", agent.getName(),
                 queries.stream().map(SemanticQuery::getQueryMode).collect(Collectors.toList()));
     }