From c2ce3a75b7422e7f567959bf156bc790c077d74b Mon Sep 17 00:00:00 2001
From: luowanghaoyun <luowanghaoyun@qiyi.com>
Date: Tue, 12 Aug 2025 11:21:11 +0800
Subject: [PATCH] Improve the performance of the chat page retrieve function

---
 .../headless/server/service/impl/RetrieveServiceImpl.java      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/RetrieveServiceImpl.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/RetrieveServiceImpl.java
index 2d87bd34c..a47275277 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/RetrieveServiceImpl.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/RetrieveServiceImpl.java
@@ -1,5 +1,6 @@
 package com.tencent.supersonic.headless.server.service.impl;
 
+import com.google.common.collect.Lists;
 import com.tencent.supersonic.common.pojo.User;
 import com.tencent.supersonic.common.pojo.enums.DictWordType;
 import com.tencent.supersonic.headless.api.pojo.SchemaElement;
@@ -79,7 +80,7 @@ public class RetrieveServiceImpl implements RetrieveService {
         Set<Long> dataSetIds = queryNLReq.getDataSetIds();
 
         ChatQueryContext chatQueryContext = new ChatQueryContext(queryNLReq);
-        chatQueryContext.setModelIdToDataSetIds(dataSetService.getModelIdToDataSetIds());
+        chatQueryContext.setModelIdToDataSetIds(modelIdToDataSetIds);
 
         Map<MatchText, List<HanlpMapResult>> regTextMap =
                 searchMatchStrategy.match(chatQueryContext, originals, dataSetIds);