From f67bf3eeac8e9d90d4a2c33c93f0492359f4f202 Mon Sep 17 00:00:00 2001
From: jerryjzhang <jerryjzhang@tencent.com>
Date: Wed, 12 Mar 2025 16:47:40 +0800
Subject: [PATCH] (fix)(chat)Fix bug in creating chat model.

---
 .../service/impl/ChatModelServiceImpl.java    |  1 +
 .../chat/corrector/AggCorrectorTest.java      | 27 ++++++-------------
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/service/impl/ChatModelServiceImpl.java b/common/src/main/java/com/tencent/supersonic/common/service/impl/ChatModelServiceImpl.java
index dd289a3db..3259be432 100644
--- a/common/src/main/java/com/tencent/supersonic/common/service/impl/ChatModelServiceImpl.java
+++ b/common/src/main/java/com/tencent/supersonic/common/service/impl/ChatModelServiceImpl.java
@@ -57,6 +57,7 @@ public class ChatModelServiceImpl extends ServiceImpl<ChatModelMapper, ChatModel
             chatModelDO.setViewer(JsonUtil.toString(chatModel.getViewers()));
         }
         save(chatModelDO);
+        chatModel.setId(chatModelDO.getId());
         return chatModel;
     }
 
diff --git a/headless/chat/src/test/java/com/tencent/supersonic/headless/chat/corrector/AggCorrectorTest.java b/headless/chat/src/test/java/com/tencent/supersonic/headless/chat/corrector/AggCorrectorTest.java
index 2f5328773..a0dde6340 100644
--- a/headless/chat/src/test/java/com/tencent/supersonic/headless/chat/corrector/AggCorrectorTest.java
+++ b/headless/chat/src/test/java/com/tencent/supersonic/headless/chat/corrector/AggCorrectorTest.java
@@ -51,14 +51,15 @@ class AggCorrectorTest {
         dataSet.setDataSetId(dataSetId);
         semanticParseInfo.setDataSet(dataSet);
         SqlInfo sqlInfo = new SqlInfo();
-        String sql = "WITH 总停留时长 AS (SELECT 用户, SUM(停留时长) AS _总停留时长_ FROM 超音数数据集 WHERE 用户 IN ('alice', 'lucy') AND 数据日期 >= '2025-03-01' AND 数据日期 <= '2025-03-12' GROUP BY 用户) SELECT 用户, _总停留时长_ FROM 总停留时长";
+        String sql =
+                "WITH 总停留时长 AS (SELECT 用户, SUM(停留时长) AS _总停留时长_ FROM 超音数数据集 WHERE 用户 IN ('alice', 'lucy') AND 数据日期 >= '2025-03-01' AND 数据日期 <= '2025-03-12' GROUP BY 用户) SELECT 用户, _总停留时长_ FROM 总停留时长";
         sqlInfo.setParsedS2SQL(sql);
         sqlInfo.setCorrectedS2SQL(sql);
         semanticParseInfo.setSqlInfo(sqlInfo);
         corrector.correct(chatQueryContext, semanticParseInfo);
         Assert.assertEquals(
-                "WITH 总停留时长 AS (SELECT 用户名, SUM(停留时长) AS _总停留时长_ FROM 超音数数据集 WHERE 用户名 IN ('alice', 'lucy') AND 数据日期 " +
-                        ">= '2025-03-01' AND 数据日期 <= '2025-03-12' GROUP BY 用户名) SELECT 用户名, _总停留时长_ FROM 总停留时长",
+                "WITH 总停留时长 AS (SELECT 用户名, SUM(停留时长) AS _总停留时长_ FROM 超音数数据集 WHERE 用户名 IN ('alice', 'lucy') AND 数据日期 "
+                        + ">= '2025-03-01' AND 数据日期 <= '2025-03-12' GROUP BY 用户名) SELECT 用户名, _总停留时长_ FROM 总停留时长",
                 semanticParseInfo.getSqlInfo().getCorrectedS2SQL());
     }
 
@@ -73,29 +74,17 @@ class AggCorrectorTest {
         dataSetSchema.setDataSet(schemaElement);
         Set<SchemaElement> dimensions = new HashSet<>();
 
-        dimensions.add(SchemaElement.builder()
-                .dataSetId(1L)
-                .name("部门")
-                .build());
+        dimensions.add(SchemaElement.builder().dataSetId(1L).name("部门").build());
 
-        dimensions.add(SchemaElement.builder()
-                        .dataSetId(1L)
-                        .name("用户名")
-                        .build());
+        dimensions.add(SchemaElement.builder().dataSetId(1L).name("用户名").build());
 
         dataSetSchema.setDimensions(dimensions);
 
         Set<SchemaElement> metrics = new HashSet<>();
 
-        metrics.add(SchemaElement.builder()
-                .dataSetId(1L)
-                .name("访问次数")
-                .build());
+        metrics.add(SchemaElement.builder().dataSetId(1L).name("访问次数").build());
 
-        metrics.add(SchemaElement.builder()
-                .dataSetId(1L)
-                .name("停留时长")
-                .build());
+        metrics.add(SchemaElement.builder().dataSetId(1L).name("停留时长").build());
 
         dataSetSchema.setMetrics(metrics);
         dataSetSchemaList.add(dataSetSchema);