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(() => { useEffect(() => {
const tableList = schemaList.find((schema) => schema.name === selectedSchemaName)?.tables || []; const tableList = schemaList.find((schema) => schema.name === selectedSchemaName)?.tables || [];
updateTableList(tableList); 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]); }, [selectedSchemaName, schemaList]);
useEffect(() => { useEffect(() => {
@ -144,15 +158,15 @@ const ConnectionSidebar = () => {
const handleTableCheckboxChange = async (tableName: string, value: boolean) => { const handleTableCheckboxChange = async (tableName: string, value: boolean) => {
if (value) { if (value) {
conversationStore.updateSelectedTablesName([...selectedTableNameList, tableName]); conversationStore.updateSelectedTablesNameList([...selectedTableNameList, tableName]);
} else { } else {
conversationStore.updateSelectedTablesName(selectedTableNameList.filter((name) => name !== tableName)); conversationStore.updateSelectedTablesNameList(selectedTableNameList.filter((name) => name !== tableName));
} }
}; };
const handleSchemaNameSelect = async (schemaName: string) => { const handleSchemaNameSelect = async (schemaName: string) => {
// need to empty selectedTableNameList when schemaName changed. because selectedTableNameList may not exist in new schema. // need to empty selectedTableNameList when schemaName changed. because selectedTableNameList may not exist in new schema.
conversationStore.updateSelectedTablesName([]); conversationStore.updateSelectedTablesNameList([]);
conversationStore.updateSelectedSchemaName(schemaName); conversationStore.updateSelectedSchemaName(schemaName);
}; };

View File

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