From 759c3373029dba9f47a32115ecddf7964f072378 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Tue, 26 Sep 2023 17:35:35 +0000 Subject: [PATCH 1/3] feat: ask users to confirm tables before text2sql --- package.json | 2 +- .../lib/ai_assistant/base_ai_assistant.py | 2 + .../prompts/table_select_prompt.py | 10 ++- .../AIAssistant/QueryGenerationModal.tsx | 80 ++++++++++++++----- .../components/AIAssistant/TableSelector.tsx | 56 ++++--------- .../components/AIAssistant/TableTag.tsx | 48 +++++++++++ .../CodeMirrorTooltip/TableTooltip.tsx | 28 +++++-- 7 files changed, 157 insertions(+), 69 deletions(-) create mode 100644 querybook/webapp/components/AIAssistant/TableTag.tsx diff --git a/package.json b/package.json index de66d712e..841de2ea6 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "querybook", - "version": "3.28.1", + "version": "3.28.2", "description": "A Big Data Webapp", "private": true, "scripts": { diff --git a/querybook/server/lib/ai_assistant/base_ai_assistant.py b/querybook/server/lib/ai_assistant/base_ai_assistant.py index 22078efd3..e120df877 100644 --- a/querybook/server/lib/ai_assistant/base_ai_assistant.py +++ b/querybook/server/lib/ai_assistant/base_ai_assistant.py @@ -207,6 +207,8 @@ def generate_sql_query( ) if tables: socket.send_tables_for_sql_gen(tables) + socket.close() + return # not finding any relevant tables if not tables: diff --git a/querybook/server/lib/ai_assistant/prompts/table_select_prompt.py b/querybook/server/lib/ai_assistant/prompts/table_select_prompt.py index 4d07bbcc5..e1a8febc0 100644 --- a/querybook/server/lib/ai_assistant/prompts/table_select_prompt.py +++ b/querybook/server/lib/ai_assistant/prompts/table_select_prompt.py @@ -2,10 +2,14 @@ prompt_template = """ -You are a data scientist that can help select the most suitable tables for SQL query tasks. +You are a data scientist that can help select the most relevant tables for SQL query tasks. -Please select at most top {top_n} tables from tables provided to answer the question below. -Please response in a valid JSON array format with table names which can be parsed by Python json.loads(). +Please select the most relevant table(s) that can be used to generate SQL query for the question. + +===Response Guidelines +- Only return the most relevant table(s). +- Return at most {top_n} tables. +- Response should be a valid JSON array of table names which can be parsed by Python json.loads(). For a single table, the format should be ["table_name"]. ===Tables {table_schemas} diff --git a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx index 2f91302b2..70169735b 100644 --- a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx +++ b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx @@ -13,6 +13,7 @@ import { trimSQLQuery } from 'lib/stream'; import { matchKeyPress } from 'lib/utils/keyboard'; import { analyzeCode } from 'lib/web-worker'; import { Button } from 'ui/Button/Button'; +import { Checkbox } from 'ui/Checkbox/Checkbox'; import { Icon } from 'ui/Icon/Icon'; import { Message } from 'ui/Message/Message'; import { Modal } from 'ui/Modal/Modal'; @@ -21,6 +22,7 @@ import { StyledText } from 'ui/StyledText/StyledText'; import { Tag } from 'ui/Tag/Tag'; import { TableSelector } from './TableSelector'; +import { TableTag } from './TableTag'; import { TextToSQLMode, TextToSQLModeSelector } from './TextToSQLModeSelector'; import './QueryGenerationModal.scss'; @@ -97,10 +99,12 @@ export const QueryGenerationModal = ({ ); const [newQuery, setNewQuery] = useState(''); const [streamData, setStreamData] = useState<{ [key: string]: string }>({}); + const [foundTables, setFoundTables] = useState([]); const onData = useCallback(({ type, data }) => { if (type === 'tables') { - setTables(uniq([...tables, ...data])); + setTables([...data.slice(0, 1)]); // select the first table by default + setFoundTables(data); } else { setStreamData(data); } @@ -121,6 +125,25 @@ export const QueryGenerationModal = ({ setNewQuery(trimSQLQuery(rawNewQuery)); }, [rawNewQuery]); + const onGenerate = useCallback(() => { + setFoundTables([]); + generateSQL({ + query_engine_id: engineId, + tables: tables, + question: question, + original_query: query, + }); + trackClick({ + component: ComponentType.AI_ASSISTANT, + element: ElementType.QUERY_GENERATION_BUTTON, + aux: { + mode: textToSQLMode, + question, + tables, + }, + }); + }, [engineId, question, tables, query, generateSQL, trackClick]); + const onKeyDown = useCallback( (event: React.KeyboardEvent) => { if ( @@ -128,24 +151,10 @@ export const QueryGenerationModal = ({ matchKeyPress(event, 'Enter') && !event.shiftKey ) { - generateSQL({ - query_engine_id: engineId, - tables: tables, - question: question, - original_query: query, - }); - trackClick({ - component: ComponentType.AI_ASSISTANT, - element: ElementType.QUERY_GENERATION_BUTTON, - aux: { - mode: textToSQLMode, - question, - tables, - }, - }); + onGenerate(); } }, - [engineId, question, tables, query, generateSQL, generating] + [onGenerate] ); const questionBarDOM = ( @@ -188,6 +197,41 @@ export const QueryGenerationModal = ({ ); + const tablesDOM = foundTables.length > 0 && ( +
+
+ Please review the tables below that I found for your question. + Select the tables you would like to use or you can also search + for tables above. +
+
+ {foundTables.map((table) => ( +
+ { + if (checked) { + setTables(uniq([...tables, table])); + } else { + setTables( + tables.filter((t) => t !== table) + ); + } + }} + /> + +
+ ))} +
+
+
+
+ ); + const bottomDOM = newQuery && !generating && (
{questionBarDOM} + {tablesDOM} {(explanation || data) && (
{explanation || data}
)} diff --git a/querybook/webapp/components/AIAssistant/TableSelector.tsx b/querybook/webapp/components/AIAssistant/TableSelector.tsx index c3337e9dc..ed7f879ec 100644 --- a/querybook/webapp/components/AIAssistant/TableSelector.tsx +++ b/querybook/webapp/components/AIAssistant/TableSelector.tsx @@ -1,16 +1,14 @@ import React, { useCallback, useState } from 'react'; import AsyncSelect, { Props as AsyncProps } from 'react-select/async'; -import { TableTooltipByName } from 'components/CodeMirrorTooltip/TableTooltip'; import { asyncReactSelectStyles, makeReactSelectStyle, } from 'lib/utils/react-select'; import { SearchTableResource } from 'resource/search'; import { overlayRoot } from 'ui/Overlay/Overlay'; -import { Popover } from 'ui/Popover/Popover'; -import { PopoverHoverWrapper } from 'ui/Popover/PopoverHoverWrapper'; -import { HoverIconTag } from 'ui/Tag/HoverIconTag'; + +import { TableTag } from './TableTag'; interface ITableSelectProps { metastoreId: number; @@ -22,7 +20,6 @@ interface ITableSelectProps { // remove the selected table name after select clearAfterSelect?: boolean; - showTablePopoverTooltip?: boolean; } export const TableSelector: React.FunctionComponent = ({ @@ -32,7 +29,6 @@ export const TableSelector: React.FunctionComponent = ({ usePortalMenu = true, selectProps = {}, clearAfterSelect = false, - showTablePopoverTooltip = false, }) => { const [searchText, setSearchText] = useState(''); const asyncSelectProps: Partial> = {}; @@ -68,41 +64,6 @@ export const TableSelector: React.FunctionComponent = ({ [metastoreId, tableNames] ); - const getTableTagDOM = (tableName) => ( - - {(showPopover, anchorElement) => ( - <> - { - const newTableNames = tableNames.filter( - (name) => name !== tableName - ); - onTableNamesChange(newTableNames); - }} - tooltip={showTablePopoverTooltip ? null : tableName} - tooltipPos="right" - mini - highlighted - light - /> - {showTablePopoverTooltip && showPopover && ( - null} - anchor={anchorElement} - layout={['right']} - > - - - )} - - )} - - ); return (
= ({ {tableNames.length ? (
{tableNames.map((tableName) => ( -
{getTableTagDOM(tableName)}
+ { + const newTableNames = tableNames.filter( + (name) => name !== tableName + ); + onTableNamesChange(newTableNames); + }} + highlighted + /> ))}
) : null} diff --git a/querybook/webapp/components/AIAssistant/TableTag.tsx b/querybook/webapp/components/AIAssistant/TableTag.tsx new file mode 100644 index 000000000..fd34acd26 --- /dev/null +++ b/querybook/webapp/components/AIAssistant/TableTag.tsx @@ -0,0 +1,48 @@ +import React from 'react'; + +import { TableTooltipByName } from 'components/CodeMirrorTooltip/TableTooltip'; +import { Popover } from 'ui/Popover/Popover'; +import { PopoverHoverWrapper } from 'ui/Popover/PopoverHoverWrapper'; +import { HoverIconTag } from 'ui/Tag/HoverIconTag'; + +interface ITableTagProps { + metastoreId: number; + tableName: string; + onIconClick?: () => void; + highlighted?: boolean; +} + +export const TableTag: React.FunctionComponent = ({ + metastoreId, + tableName, + onIconClick, + highlighted = false, +}) => ( + + {(showPopover, anchorElement) => ( + <> + + {showPopover && ( + null} + anchor={anchorElement} + layout={['right']} + > + + + )} + + )} + +); diff --git a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx index 67025e0e8..c48052b16 100644 --- a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx +++ b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx @@ -1,11 +1,12 @@ import { ContentState } from 'draft-js'; -import React, { useEffect, useState } from 'react'; +import React, { useCallback, useEffect, useState } from 'react'; import { useDispatch } from 'react-redux'; import { DataTableTags } from 'components/DataTableTags/DataTableTags'; import { IDataColumn, IDataSchema, IDataTable } from 'const/metastore'; import { useShallowSelector } from 'hooks/redux/useShallowSelector'; import { setSidebarTableId } from 'lib/querybookUI'; +import { navigateWithinEnv } from 'lib/utils/query-string'; import * as dataSourcesActions from 'redux/dataSources/action'; import { IStoreState } from 'redux/store/types'; import { IconButton } from 'ui/Button/IconButton'; @@ -15,6 +16,7 @@ interface IProps { table: IDataTable; columns: IDataColumn[]; schema: IDataSchema; + showPinItButton?: boolean; openTableModal?: () => any; } @@ -22,6 +24,7 @@ export const TableTooltip: React.FunctionComponent = ({ table, columns, schema, + showPinItButton = true, openTableModal, }) => { const tableName = @@ -47,7 +50,7 @@ export const TableTooltip: React.FunctionComponent = ({ className="ml4" /> ); - const pinToSidebarButton = ( + const pinToSidebarButton = showPinItButton && ( = ({ export const TableTooltipByName: React.FunctionComponent<{ metastoreId: number; tableFullName: string; - openTableModal?: () => any; -}> = ({ metastoreId, tableFullName, openTableModal }) => { + showPinItButton?: boolean; + showDetails?: boolean; +}> = ({ + metastoreId, + tableFullName, + showPinItButton = false, + showDetails = true, +}) => { const dispatch = useDispatch(); const [tableId, setTableId] = useState(null); + const openTableModal = useCallback((tableId: number) => { + navigateWithinEnv(`/table/${tableId}/`, { + isModal: true, + }); + }, []); + useEffect(() => { const fetchTable = async () => { try { @@ -167,7 +182,10 @@ export const TableTooltipByName: React.FunctionComponent<{ table={table} schema={schema} columns={columns} - openTableModal={openTableModal} + showPinItButton={showPinItButton} + openTableModal={ + showDetails ? () => openTableModal(tableId) : undefined + } /> ); }; From 461f32bd93c8a559119ac50d6bce40f3cf639331 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Tue, 26 Sep 2023 20:23:37 +0000 Subject: [PATCH 2/3] comments --- .../lib/ai_assistant/base_ai_assistant.py | 22 +++++++++--------- .../AIAssistant/QueryGenerationModal.tsx | 23 +++++++++++-------- .../CodeMirrorTooltip/TableTooltip.tsx | 8 +++---- 3 files changed, 27 insertions(+), 26 deletions(-) diff --git a/querybook/server/lib/ai_assistant/base_ai_assistant.py b/querybook/server/lib/ai_assistant/base_ai_assistant.py index e120df877..fae8af071 100644 --- a/querybook/server/lib/ai_assistant/base_ai_assistant.py +++ b/querybook/server/lib/ai_assistant/base_ai_assistant.py @@ -199,23 +199,23 @@ def generate_sql_query( query_engine = admin_logic.get_query_engine_by_id( query_engine_id, session=session ) + if not tables: - tables = self.find_tables( + suggested_tables = self.find_tables( metastore_id=query_engine.metastore_id, question=question, session=session, ) - if tables: - socket.send_tables_for_sql_gen(tables) - socket.close() - return - # not finding any relevant tables - if not tables: - # ask user to provide table names - socket.send_data( - "Sorry, I can't find any relevant tables by the given context. Please provide table names above." - ) + if suggested_tables: + socket.send_tables_for_sql_gen(suggested_tables) + else: + # not finding any relevant tables + # ask user to provide table names + socket.send_data( + "Sorry, I can't find any relevant tables by the given context. Please provide table names above." + ) + socket.close() return diff --git a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx index 70169735b..8136e1b03 100644 --- a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx +++ b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx @@ -209,15 +209,13 @@ export const QueryGenerationModal = ({
{ - if (checked) { - setTables(uniq([...tables, table])); - } else { - setTables( - tables.filter((t) => t !== table) - ); - } - }} + onChange={(checked) => + setTables((oldTables) => + checked + ? uniq([...oldTables, table]) + : oldTables.filter((t) => t !== table) + ) + } />
-
); diff --git a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx index c48052b16..d3a751d63 100644 --- a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx +++ b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx @@ -127,11 +127,11 @@ export const TableTooltipByName: React.FunctionComponent<{ const dispatch = useDispatch(); const [tableId, setTableId] = useState(null); - const openTableModal = useCallback((tableId: number) => { + const openTableModal = useCallback(() => { navigateWithinEnv(`/table/${tableId}/`, { isModal: true, }); - }, []); + }, [tableId]); useEffect(() => { const fetchTable = async () => { @@ -183,9 +183,7 @@ export const TableTooltipByName: React.FunctionComponent<{ schema={schema} columns={columns} showPinItButton={showPinItButton} - openTableModal={ - showDetails ? () => openTableModal(tableId) : undefined - } + openTableModal={showDetails ? openTableModal : undefined} /> ); }; From 96f7ed04b0c4bde7e203aecc8b21da60e4900121 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Tue, 26 Sep 2023 23:20:13 +0000 Subject: [PATCH 3/3] hidePitItButton --- .../components/CodeMirrorTooltip/TableTooltip.tsx | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx index d3a751d63..b062b079a 100644 --- a/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx +++ b/querybook/webapp/components/CodeMirrorTooltip/TableTooltip.tsx @@ -16,7 +16,7 @@ interface IProps { table: IDataTable; columns: IDataColumn[]; schema: IDataSchema; - showPinItButton?: boolean; + hidePinItButton?: boolean; openTableModal?: () => any; } @@ -24,7 +24,7 @@ export const TableTooltip: React.FunctionComponent = ({ table, columns, schema, - showPinItButton = true, + hidePinItButton = false, openTableModal, }) => { const tableName = @@ -50,7 +50,7 @@ export const TableTooltip: React.FunctionComponent = ({ className="ml4" /> ); - const pinToSidebarButton = showPinItButton && ( + const pinToSidebarButton = !hidePinItButton && ( = ({ export const TableTooltipByName: React.FunctionComponent<{ metastoreId: number; tableFullName: string; - showPinItButton?: boolean; + hidePinItButton?: boolean; showDetails?: boolean; }> = ({ metastoreId, tableFullName, - showPinItButton = false, + hidePinItButton = true, showDetails = true, }) => { const dispatch = useDispatch(); @@ -182,7 +182,7 @@ export const TableTooltipByName: React.FunctionComponent<{ table={table} schema={schema} columns={columns} - showPinItButton={showPinItButton} + hidePinItButton={hidePinItButton} openTableModal={showDetails ? openTableModal : undefined} /> );