feat: add @mention to select a table for the command input (#1432)

* feat: add @mention support for the command input

* address comments

* fix table selector disappear issue

* comments

* remove id

* use suggest

* add comment

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.33.1",
+    "version": "3.33.2",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {
@@ -84,6 +84,7 @@
         "react-hot-toast": "^2.4.0",
         "react-json-view": "^1.21.3",
         "react-lazyload": "^3.2.0",
+        "react-mentions": "^4.4.10",
         "react-redux": "7.2.4",
         "react-router-dom": "5.1.2",
         "react-select": "4.3.1",

querybook/config/elasticsearch.yaml

@@ -212,7 +212,7 @@ tables:
                 full_name_ngram:
                     type: text
                     analyzer: edge_ngram_lowercase
-                completion_name:
+                completion_name: # to be deprecated
                     type: completion
                     analyzer: keyword
                     contexts:

querybook/server/datasources/search.py

@@ -130,20 +130,14 @@ def vector_search_tables(
 
 
 @register("/suggest/<int:metastore_id>/tables/", methods=["GET"])
-def suggest_tables(metastore_id, prefix, limit=10):
+def suggest_tables(metastore_id, keyword, limit=10):
     api_assert(limit is None or limit <= 100, "Requesting too many tables")
     verify_metastore_permission(metastore_id)
 
-    query = construct_suggest_table_query(prefix, limit, metastore_id)
+    query = construct_suggest_table_query(keyword, limit, metastore_id)
     options = get_matching_suggestions(query, ES_CONFIG["tables"]["index_name"])
-    texts = [
-        "{}.{}".format(
-            option.get("_source", {}).get("schema", ""),
-            option.get("_source", {}).get("name", ""),
-        )
-        for option in options
-    ]
-    return texts
+    tables = [option.get("_source", {}).get("full_name") for option in options]
+    return tables
 
 
 # /search/ but it is actually suggest

querybook/server/lib/elasticsearch/search_utils.py

@@ -79,8 +79,6 @@ def get_matching_suggestions(query: Union[str, Dict], index_name: str):
         if result is None:
             result = {}
 
-    options = next(iter(result.get("suggest", {}).get("suggest", [])), {}).get(
-        "options", []
-    )
+    options = result.get("hits", {}).get("hits", [])
 
     return options

querybook/server/lib/elasticsearch/suggest_table.py

@@ -1,17 +1,18 @@
 def construct_suggest_table_query(
-    prefix: str,
+    keyword: str,
     limit: int,
     metastore_id: int,
 ):
     return {
-        "suggest": {
-            "suggest": {
-                "text": prefix,
-                "completion": {
-                    "field": "completion_name",
-                    "size": limit,
-                    "contexts": {"metastore_id": metastore_id},
+        "from": 0,
+        "size": limit,
+        "query": {
+            "bool": {
+                "must": {
+                    "match": {"full_name_ngram": {"query": keyword, "operator": "and"}}
                 },
+                "filter": {"match": {"metastore_id": metastore_id}},
             }
         },
+        "_source": ["id", "full_name"],
     }

querybook/server/logic/elasticsearch.py

@@ -601,6 +601,7 @@ def get_completion_name():
         "name": table_name,
         "full_name": full_name,
         "full_name_ngram": full_name,
+        # To be deprecated: completion_name is not used anymore
         "completion_name": get_completion_name,
         "description": get_table_description,
         "created_at": lambda: DATETIME_TO_UTC(table.created_at),

querybook/webapp/components/AIAssistant/AICommandBar.tsx

@@ -1,5 +1,11 @@
-import { set, uniq } from 'lodash';
-import React, { forwardRef, useCallback, useEffect, useState } from 'react';
+import { uniq } from 'lodash';
+import React, {
+    forwardRef,
+    useCallback,
+    useEffect,
+    useMemo,
+    useState,
+} from 'react';
 
 import { AICommandInput } from 'components/AIAssistant/AICommandInput';
 import { AICommandResultView } from 'components/AIAssistant/AICommandResultView';
@@ -11,14 +17,11 @@ import {
 } from 'const/command';
 import { IQueryEngine } from 'const/queryEngine';
 import { CommandRunner, useCommand } from 'hooks/useCommand';
-import { useEvent } from 'hooks/useEvent';
 import { useForwardedRef } from 'hooks/useForwardedRef';
 import { trackClick } from 'lib/analytics';
 import { TableToken } from 'lib/sql-helper/sql-lexer';
-import { matchKeyPress } from 'lib/utils/keyboard';
 import { analyzeCode } from 'lib/web-worker';
-import { Button, TextButton } from 'ui/Button/Button';
-import { IconButton } from 'ui/Button/IconButton';
+import { Button } from 'ui/Button/Button';
 import { Message } from 'ui/Message/Message';
 import { IResizableTextareaHandles } from 'ui/ResizableTextArea/ResizableTextArea';
 import { StyledText } from 'ui/StyledText/StyledText';
@@ -65,7 +68,8 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
             (cmd) => cmd.name === (query ? 'edit' : 'generate')
         );
         const tablesInQuery = useTablesInQuery(query, queryEngine.language);
-        const [tables, setTables] = useState(tablesInQuery);
+        const [tables, setTables] = useState<string[]>([]);
+        const [mentionedTables, setMentionedTables] = useState<string[]>([]);
         const commandInputRef = useForwardedRef<IResizableTextareaHandles>(ref);
         const [showPopupView, setShowPopupView] = useState(false);
         const [command, setCommand] =
@@ -77,6 +81,14 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
         );
         const [showConfirm, setShowConfirm] = useState(false);
 
+        const finalTablesToUse = useMemo(() => {
+            if (mentionedTables.length > 0) {
+                return mentionedTables;
+            } else {
+                return uniq([...tablesInQuery, ...tables]);
+            }
+        }, [tablesInQuery, mentionedTables, tables]);
+
         const {
             runCommand,
             isRunning,
@@ -86,8 +98,14 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
         } = useCommand(command, commandRunner);
 
         useEffect(() => {
-            setTables((tables) => uniq([...tablesInQuery, ...tables]));
-        }, [tablesInQuery]);
+            if (!query && mentionedTables.length === 1) {
+                onUpdateQuery(
+                    `SELECT * FROM ${mentionedTables[0]} LIMIT 10`,
+                    false
+                );
+                onFormatQuery();
+            }
+        }, [mentionedTables]);
 
         useEffect(() => {
             if (command.name === 'format') {
@@ -96,7 +114,7 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
             } else if (command.name === 'generate' || command.name === 'edit') {
                 setCommandKwargs({
                     query_engine_id: queryEngine.id,
-                    tables: tables,
+                    tables: finalTablesToUse,
                     question: commandInputValue,
                     original_query: command.name === 'generate' ? '' : query,
                 });
@@ -122,25 +140,22 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
                     aux: {
                         mode: command.name,
                         question: commandKwargs.question,
-                        tables,
+                        tables: finalTablesToUse,
                     },
                 });
             }
         }, [command, runCommand, setShowPopupView, commandKwargs]);
 
         const getCommandResultView = () => {
-            if (!commandResult) {
-                return null;
-            }
-
             if (command.name === 'generate' || command.name === 'edit') {
                 return (
                     <AICommandResultView
                         command={command}
                         commandKwargs={commandKwargs}
                         metastoreId={queryEngine.metastore_id}
                         commandResult={commandResult}
-                        tables={tables}
+                        tables={finalTablesToUse}
+                        hasMentionedTables={mentionedTables.length > 0}
                         originalQuery={query}
                         isStreaming={isRunning}
                         onContinue={handleCommand}
@@ -226,6 +241,9 @@ export const AICommandBar: React.FC<IQueryCellCommandBarProps> = forwardRef(
                             );
                             setCommandInputValue(inputValue);
                         }}
+                        mentionedTables={mentionedTables}
+                        onMentionedTablesChange={setMentionedTables}
+                        metastoreId={queryEngine.metastore_id}
                         onSubmit={handleCommand}
                         running={isRunning}
                         cancelGeneration={cancelCommand}

querybook/webapp/components/AIAssistant/AICommandInput.scss

@@ -6,10 +6,6 @@
     background-color: var(--bg-lightest);
     border-radius: var(--border-radius-sm);
 
-    &:hover {
-        background-color: var(--bg-hover);
-    }
-
     .stars-icon {
         position: relative;
         color: var(--icon);
@@ -41,15 +37,6 @@
         color: var(--color-pink-dark);
     }
 
-    .question-text-area {
-        flex: 1;
-        min-width: auto;
-        line-height: 24px;
-        min-height: 40px;
-        overflow: hidden;
-        padding: 8px 8px;
-    }
-
     .button {
         height: 28px;
         margin: 6px;
@@ -59,6 +46,7 @@
 .AICommandInput-popover {
     border-radius: var(--border-radius-sm);
     overflow: hidden;
+    box-shadow: 0 0px 4px var(--bg-dark);
 
     .command-item {
         display: flex;