feat: select as many tables up to the maximum token by default

This commit is contained in:
Tianzhou Chen
2023-06-08 01:23:34 +08:00
parent e774078613
commit 88f80aaefe
2 changed files with 19 additions and 5 deletions

View File

@ -102,6 +102,20 @@ const ConnectionSidebar = () => {
useEffect(() => {
const tableList = schemaList.find((schema) => schema.name === selectedSchemaName)?.tables || [];
updateTableList(tableList);
// By default, select as many tables up to the maximum token.
const defaultCheckedTableList = [];
if (selectedTableNameList.length === 0) {
let tokenCount = 0;
for (const table of tableList) {
tokenCount += countTextTokens(table?.structure || "");
if (tokenCount > maxToken) {
break;
}
defaultCheckedTableList.push(table.name);
}
conversationStore.updateSelectedTablesNameList(defaultCheckedTableList);
}
}, [selectedSchemaName, schemaList]);
useEffect(() => {
@ -144,15 +158,15 @@ const ConnectionSidebar = () => {
const handleTableCheckboxChange = async (tableName: string, value: boolean) => {
if (value) {
conversationStore.updateSelectedTablesName([...selectedTableNameList, tableName]);
conversationStore.updateSelectedTablesNameList([...selectedTableNameList, tableName]);
} else {
conversationStore.updateSelectedTablesName(selectedTableNameList.filter((name) => name !== tableName));
conversationStore.updateSelectedTablesNameList(selectedTableNameList.filter((name) => name !== tableName));
}
};
const handleSchemaNameSelect = async (schemaName: string) => {
// need to empty selectedTableNameList when schemaName changed. because selectedTableNameList may not exist in new schema.
conversationStore.updateSelectedTablesName([]);
conversationStore.updateSelectedTablesNameList([]);
conversationStore.updateSelectedSchemaName(schemaName);
};

View File

@ -23,7 +23,7 @@ interface ConversationState {
getConversationById: (conversationId: Id | undefined) => Conversation | undefined;
updateConversation: (conversationId: Id, conversation: Partial<Conversation>) => void;
clearConversation: (filter: (conversation: Conversation) => boolean) => void;
updateSelectedTablesName: (selectedTableNameList: string[]) => void;
updateSelectedTablesNameList: (selectedTableNameList: string[]) => void;
updateSelectedSchemaName: (selectedSchemaName: string) => void;
}
@ -65,7 +65,7 @@ export const useConversationStore = create<ConversationState>()(
conversationList: state.conversationList.filter(filter),
}));
},
updateSelectedTablesName: (selectedTableNameList: string[]) => {
updateSelectedTablesNameList: (selectedTableNameList: string[]) => {
const currentConversation = get().getConversationById(get().currentConversationId);
if (currentConversation) {
get().updateConversation(currentConversation.id, {