From cc66ebd684c060148015f3c816c141ee46fb7598 Mon Sep 17 00:00:00 2001
From: jerryjzhang <jerryjzhang@tencent.com>
Date: Wed, 5 Feb 2025 17:24:48 +0800
Subject: [PATCH] [improvement][chat]Add queryId to QueryNLReq.

---
 .../chat/server/service/impl/ChatQueryServiceImpl.java           | 1 +
 .../tencent/supersonic/headless/api/pojo/request/QueryNLReq.java | 1 +
 .../tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java | 1 +
 3 files changed, 3 insertions(+)

diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java
index 52d49b930..b1b37897e 100644
--- a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java
+++ b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java
@@ -91,6 +91,7 @@ public class ChatQueryServiceImpl implements ChatQueryService {
         Long queryId = chatParseReq.getQueryId();
         if (Objects.isNull(queryId)) {
             queryId = chatManageService.createChatQuery(chatParseReq);
+            chatParseReq.setQueryId(queryId);
         }
 
         ParseContext parseContext = buildParseContext(chatParseReq, new ChatParseResp(queryId));
diff --git a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/request/QueryNLReq.java b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/request/QueryNLReq.java
index a6fef7d7a..ea25f5388 100644
--- a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/request/QueryNLReq.java
+++ b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/request/QueryNLReq.java
@@ -18,6 +18,7 @@ import java.util.Set;
 
 @Data
 public class QueryNLReq extends SemanticQueryReq implements Serializable {
+    private Long queryId;
     private String queryText;
     private Set<Long> dataSetIds = Sets.newHashSet();
     private User user;
diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java
index 944ef25fd..c3a8c4228 100644
--- a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java
+++ b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java
@@ -18,6 +18,7 @@ import java.util.stream.Collectors;
 
 @Data
 public class LLMReq {
+    private Long queryId;
     private String queryText;
     private LLMSchema schema;
     private List<Term> terms;