diff --git a/src/components/ConnectionSidebar.tsx b/src/components/ConnectionSidebar.tsx index fbeff56..4db1464 100644 --- a/src/components/ConnectionSidebar.tsx +++ b/src/components/ConnectionSidebar.tsx @@ -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); }; diff --git a/src/store/conversation.ts b/src/store/conversation.ts index 415a6e3..59f782f 100644 --- a/src/store/conversation.ts +++ b/src/store/conversation.ts @@ -23,7 +23,7 @@ interface ConversationState { getConversationById: (conversationId: Id | undefined) => Conversation | undefined; updateConversation: (conversationId: Id, conversation: Partial) => 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()( conversationList: state.conversationList.filter(filter), })); }, - updateSelectedTablesName: (selectedTableNameList: string[]) => { + updateSelectedTablesNameList: (selectedTableNameList: string[]) => { const currentConversation = get().getConversationById(get().currentConversationId); if (currentConversation) { get().updateConversation(currentConversation.id, {