(improvement)(Headless) Add Text2SQLType to control whether rules and large models are passed (#891)

Co-authored-by: jolunoluo
This commit is contained in:
LXW
2024-04-07 14:33:17 +08:00
committed by GitHub
parent 12e25c0c50
commit 5f6e9ae194
8 changed files with 48 additions and 17 deletions

View File

@@ -70,6 +70,10 @@ public class Agent extends RecordInfo {
return !CollectionUtils.isEmpty(getParserTools(AgentToolType.NL2SQL_LLM));
}
public boolean containsRuleTool() {
return !CollectionUtils.isEmpty(getParserTools(AgentToolType.NL2SQL_RULE));
}
public boolean containsNL2SQLTool() {
return !CollectionUtils.isEmpty(getParserTools(AgentToolType.NL2SQL_LLM))
|| !CollectionUtils.isEmpty(getParserTools(AgentToolType.NL2SQL_RULE));

View File

@@ -2,6 +2,7 @@ package com.tencent.supersonic.chat.server.util;
import com.tencent.supersonic.chat.server.agent.Agent;
import com.tencent.supersonic.chat.server.pojo.ChatParseContext;
import com.tencent.supersonic.common.pojo.enums.Text2SQLType;
import com.tencent.supersonic.common.util.BeanMapper;
import com.tencent.supersonic.headless.api.pojo.request.QueryReq;
import org.apache.commons.collections.MapUtils;
@@ -17,8 +18,12 @@ public class QueryReqConverter {
if (agent == null) {
return queryReq;
}
if (agent.containsLLMParserTool()) {
queryReq.setEnableLLM(true);
if (agent.containsLLMParserTool() && agent.containsRuleTool()) {
queryReq.setText2SQLType(Text2SQLType.RULE_AND_LLM);
} else if (agent.containsLLMParserTool()) {
queryReq.setText2SQLType(Text2SQLType.ONLY_LLM);
} else if (agent.containsRuleTool()) {
queryReq.setText2SQLType(Text2SQLType.ONLY_RULE);
}
queryReq.setDataSetIds(agent.getDataSetIds());
if (Objects.nonNull(queryReq.getMapInfo())