feat: ask users to confirm tables before query generation

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.28.1",
+    "version": "3.28.2",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

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:

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}

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<string>('');
     const [streamData, setStreamData] = useState<{ [key: string]: string }>({});
+    const [foundTables, setFoundTables] = useState<string[]>([]);
 
     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,31 +125,36 @@ 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 (
                 !generating &&
                 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 = ({
         </div>
     );
 
+    const tablesDOM = foundTables.length > 0 && (
+        <div className="mt12">
+            <div>
+                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.
+            </div>
+            <div className="mt8">
+                {foundTables.map((table) => (
+                    <div key={table} className="flex-row">
+                        <Checkbox
+                            value={tables.includes(table)}
+                            onChange={(checked) => {
+                                if (checked) {
+                                    setTables(uniq([...tables, table]));
+                                } else {
+                                    setTables(
+                                        tables.filter((t) => t !== table)
+                                    );
+                                }
+                            }}
+                        />
+                        <TableTag
+                            metastoreId={queryEngineById[engineId].metastore_id}
+                            tableName={table}
+                        />
+                    </div>
+                ))}
+            </div>
+            <div className="mt8">
+                <Button title="Continue" onClick={onGenerate} color="confirm" />
+            </div>
+        </div>
+    );
+
     const bottomDOM = newQuery && !generating && (
         <div className="right-align mb16">
             <Button
@@ -278,13 +322,13 @@ export const QueryGenerationModal = ({
                                     autoFocus: true,
                                 }}
                                 clearAfterSelect
-                                showTablePopoverTooltip
                             />
                         </div>
                     </div>
                 </div>
 
                 {questionBarDOM}
+                {tablesDOM}
                 {(explanation || data) && (
                     <div className="mt12">{explanation || data}</div>
                 )}

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<ITableSelectProps> = ({
@@ -32,7 +29,6 @@ export const TableSelector: React.FunctionComponent<ITableSelectProps> = ({
     usePortalMenu = true,
     selectProps = {},
     clearAfterSelect = false,
-    showTablePopoverTooltip = false,
 }) => {
     const [searchText, setSearchText] = useState('');
     const asyncSelectProps: Partial<AsyncProps<any, false>> = {};
@@ -68,41 +64,6 @@ export const TableSelector: React.FunctionComponent<ITableSelectProps> = ({
         [metastoreId, tableNames]
     );
 
-    const getTableTagDOM = (tableName) => (
-        <PopoverHoverWrapper>
-            {(showPopover, anchorElement) => (
-                <>
-                    <HoverIconTag
-                        name={tableName}
-                        iconOnHover="X"
-                        onIconHoverClick={() => {
-                            const newTableNames = tableNames.filter(
-                                (name) => name !== tableName
-                            );
-                            onTableNamesChange(newTableNames);
-                        }}
-                        tooltip={showTablePopoverTooltip ? null : tableName}
-                        tooltipPos="right"
-                        mini
-                        highlighted
-                        light
-                    />
-                    {showTablePopoverTooltip && showPopover && (
-                        <Popover
-                            onHide={() => null}
-                            anchor={anchorElement}
-                            layout={['right']}
-                        >
-                            <TableTooltipByName
-                                metastoreId={metastoreId}
-                                tableFullName={tableName}
-                            />
-                        </Popover>
-                    )}
-                </>
-            )}
-        </PopoverHoverWrapper>
-    );
     return (
         <div className="TableSelect">
             <AsyncSelect
@@ -128,7 +89,18 @@ export const TableSelector: React.FunctionComponent<ITableSelectProps> = ({
             {tableNames.length ? (
                 <div className="flex-row flex-wrap mt8 gap8">
                     {tableNames.map((tableName) => (
-                        <div key={tableName}>{getTableTagDOM(tableName)}</div>
+                        <TableTag
+                            key={tableName}
+                            metastoreId={metastoreId}
+                            tableName={tableName}
+                            onIconClick={() => {
+                                const newTableNames = tableNames.filter(
+                                    (name) => name !== tableName
+                                );
+                                onTableNamesChange(newTableNames);
+                            }}
+                            highlighted
+                        />
                     ))}
                 </div>
             ) : null}

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<ITableTagProps> = ({
+    metastoreId,
+    tableName,
+    onIconClick,
+    highlighted = false,
+}) => (
+    <PopoverHoverWrapper>
+        {(showPopover, anchorElement) => (
+            <>
+                <HoverIconTag
+                    name={tableName}
+                    iconOnHover={onIconClick ? 'X' : undefined}
+                    onIconHoverClick={onIconClick}
+                    mini
+                    highlighted={highlighted}
+                    light
+                />
+                {showPopover && (
+                    <Popover
+                        onHide={() => null}
+                        anchor={anchorElement}
+                        layout={['right']}
+                    >
+                        <TableTooltipByName
+                            metastoreId={metastoreId}
+                            tableFullName={tableName}
+                            showDetails={true}
+                        />
+                    </Popover>
+                )}
+            </>
+        )}
+    </PopoverHoverWrapper>
+);

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,13 +16,15 @@ interface IProps {
     table: IDataTable;
     columns: IDataColumn[];
     schema: IDataSchema;
+    showPinItButton?: boolean;
     openTableModal?: () => any;
 }
 
 export const TableTooltip: React.FunctionComponent<IProps> = ({
     table,
     columns,
     schema,
+    showPinItButton = true,
     openTableModal,
 }) => {
     const tableName =
@@ -47,7 +50,7 @@ export const TableTooltip: React.FunctionComponent<IProps> = ({
             className="ml4"
         />
     );
-    const pinToSidebarButton = (
+    const pinToSidebarButton = showPinItButton && (
         <IconButton
             noPadding
             size={18}
@@ -113,11 +116,23 @@ export const TableTooltip: React.FunctionComponent<IProps> = ({
 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
+            }
         />
     );
 };