feat: ask users to confirm tables before query generation (pinterest#…

…1339) * feat: ask users to confirm tables before text2sql * comments * hidePitItButton
arrowtail-precision · Jan 3, 2024 · 28de77e · 28de77e
1 parent 2edc691
commit 28de77e
Show file tree

Hide file tree

Showing 7 changed files with 167 additions and 78 deletions.
diff --git 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
@@ -199,21 +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)
 
-        # 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/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
@@ -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,44 @@ 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) =>
+                                setTables((oldTables) =>
+                                    checked
+                                        ? uniq([...oldTables, table])
+                                        : oldTables.filter((t) => t !== table)
+                                )
+                            }
+                        />
+                        <TableTag
+                            metastoreId={queryEngineById[engineId].metastore_id}
+                            tableName={table}
+                        />
+                    </div>
+                ))}
+            </div>
+            <div className="mt8">
+                <Button
+                    title="Confirm & Generate SQL"
+                    onClick={onGenerate}
+                    color="confirm"
+                    disabled={tables.length === 0}
+                />
+            </div>
+        </div>
+    );
+
     const bottomDOM = newQuery && !generating && (
         <div className="right-align mb16">
             <Button
@@ -278,13 +325,13 @@ export const QueryGenerationModal = ({
                                     autoFocus: true,
                                 }}
                                 clearAfterSelect
-                                showTablePopoverTooltip
                             />
                         </div>
                     </div>
                 </div>
 
                 {questionBarDOM}
+                {tablesDOM}
                 {(explanation || data) && (
                     <div className="mt12">{explanation || data}</div>
                 )}

diff --git 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<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}

diff --git a/querybook/webapp/components/AIAssistant/TableTag.tsx 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<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>
+);