modify user setting for query suggestions

querybook/config/user_setting.yaml

@@ -50,6 +50,14 @@ show_full_view:
         - disabled
     helper: Instead of modal, show full view when opening table/execution/snippet
 
+query_suggestions:
+    default: disabled
+    tab: editor
+    options:
+        - enabled
+        - disabled
+    helper: Enable to receive inline LLM generated query suggestions as you type from within the editor.
+
 editor_font_size:
     default: medium
     tab: editor

querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx

@@ -49,8 +49,14 @@ export const BoundQueryEditor = React.forwardRef<
         const editorRef = useForwardedRef<IQueryEditorHandles>(ref);
 
         // Code Editor related Props
-        const { codeEditorTheme, keyMap, options, fontSize, autoCompleteType } =
-            useUserQueryEditorConfig(searchContext);
+        const {
+            codeEditorTheme,
+            keyMap,
+            options,
+            fontSize,
+            autoCompleteType,
+            querySuggestionsEnabled,
+        } = useUserQueryEditorConfig(searchContext);
         const combinedOptions = useMemo(
             () => ({
                 ...options,
@@ -103,6 +109,7 @@ export const BoundQueryEditor = React.forwardRef<
                 theme={codeEditorTheme}
                 autoCompleteType={autoCompleteType}
                 fontSize={fontSize}
+                querySuggestionsEnabled={querySuggestionsEnabled}
                 getTableByName={fetchDataTable}
                 functionDocumentationByNameByLanguage={
                     functionDocumentationByNameByLanguage

querybook/webapp/components/QueryEditor/QueryEditor.tsx

@@ -23,6 +23,8 @@ import {
 import { useAutoComplete } from 'hooks/queryEditor/useAutoComplete';
 import { useCodeAnalysis } from 'hooks/queryEditor/useCodeAnalysis';
 import { useLint } from 'hooks/queryEditor/useLint';
+import { useUserQueryEditorConfig } from 'hooks/redux/useUserQueryEditorConfig';
+import { useUserQueryEditorConfig } from 'hooks/redux/useUserQueryEditorConfig';
 import { useDebouncedFn } from 'hooks/useDebouncedFn';
 import CodeMirror, { CodeMirrorKeyMap } from 'lib/codemirror';
 import { SQL_JINJA_MODE } from 'lib/codemirror/codemirror-mode';
@@ -57,6 +59,8 @@ export interface IQueryEditorProps extends IStyledQueryEditorProps {
     keyMap?: CodeMirrorKeyMap;
     className?: string;
     autoCompleteType?: AutoCompleteType;
+    querySuggestionsEnabled?: boolean;
+    querySuggestionsEnabled?: boolean;
 
     /**
      * If provided, then the container component will handle the fullscreen logic
@@ -115,6 +119,8 @@ export const QueryEditor: React.FC<
             keyMap = {},
             className,
             autoCompleteType = 'all',
+            querySuggestionsEnabled,
+            querySuggestionsEnabled,
             onFullScreen,
 
             onChange,
@@ -536,6 +542,16 @@ export const QueryEditor: React.FC<
             });
         }, [codeAnalysis]);
 
+        useEffect(() => {
+            if (querySuggestionsEnabled) {
+                // Enable copilot suggestion feature
+                enableCopilotSuggestions();
+            } else {
+                // Disable copilot suggestion feature
+                disableCopilotSuggestions();
+            }
+        }, [querySuggestionsEnabled]);
+
         useImperativeHandle(
             ref,
             () => ({
@@ -617,16 +633,24 @@ export const QueryEditor: React.FC<
             onTextHover,
         ]);
 
-        const editorDidMount = useCallback((editor: CodeMirror.Editor) => {
-            editorRef.current = editor;
+        const editorDidMount = useCallback(
+            (editor: CodeMirror.Editor) => {
+                editorRef.current = editor;
+
+                if (querySuggestionsEnabled) {
+                    // Enable copilot suggestion feature
+                    editor.querySuggestions();
+                }
 
-            // There is a strange bug where codemirror would start with the wrong height (on Execs tab)
-            // which can only be solved by clicking on it
-            // The current work around is to add refresh on mount
-            setTimeout(() => {
-                editor.refresh();
-            }, 50);
-        }, []);
+                // There is a strange bug where codemirror would start with the wrong height (on Execs tab)
+                // which can only be solved by clicking on it
+                // The current work around is to add refresh on mount
+                setTimeout(() => {
+                    editor.refresh();
+                }, 50);
+            },
+            [querySuggestionsEnabled]
+        );
 
         const onBeforeChange = useCallback(
             (editor: CodeMirror.Editor, data, value: string) => {

querybook/webapp/hooks/redux/useUserQueryEditorConfig.ts

@@ -17,6 +17,7 @@ export function useUserQueryEditorConfig(
     keyMap: CodeMirrorKeyMap;
     options: CodeMirror.EditorConfiguration;
     autoCompleteType: AutoCompleteType;
+    querySuggestionsEnabled: boolean;
 } {
     const editorSettings = useShallowSelector((state: IStoreState) => ({
         theme: getCodeEditorTheme(state.user.computedSettings['theme']),
@@ -26,6 +27,8 @@ export function useUserQueryEditorConfig(
             ],
         autoComplete: state.user.computedSettings['auto_complete'],
         tab: state.user.computedSettings['tab'],
+        querySuggestionsEnabled:
+            state.user.computedSettings['query_suggestions'] === 'enabled',
     }));
     const indentWithTabs = editorSettings.tab === 'tab';
     const tabSize =
@@ -106,5 +109,6 @@ export function useUserQueryEditorConfig(
         // From: https://github.com/codemirror/CodeMirror/issues/988
         keyMap,
         options,
+        querySuggestionsEnabled: editorSettings.querySuggestionsEnabled,
     };
 }

querybook/webapp/lib/codemirror/codemirror-copilot.ts

@@ -0,0 +1,9 @@
+import CodeMirror from 'codemirror';
+
+CodeMirror.defineExtension('querySuggestions', function () {
+    this.on('keyup', async (editor, event) => {
+        if (event.code === 'Space') {
+            console.log('CodeMirror Query Suggestions Extension');
+        }
+    });
+});

querybook/webapp/lib/codemirror/index.ts

@@ -18,21 +18,25 @@ import 'codemirror/addon/runmode/runmode';
 // Search highlighting
 import 'codemirror/addon/search/match-highlighter.js';
 import 'codemirror/lib/codemirror.css';
+import 'codemirror/mode/jinja2/jinja2';
 // From codemirror non-react package:
 import 'codemirror/mode/sql/sql';
-import 'codemirror/mode/jinja2/jinja2';
-
 import 'codemirror/theme/duotone-light.css';
 import 'codemirror/theme/material-palenight.css';
 import 'codemirror/theme/monokai.css';
 import 'codemirror/theme/solarized.css';
+// Query Suggestions
+import 'lib/codemirror/codemirror-copilot';
 // This should apply the hover option to codemirror
 import 'lib/codemirror/codemirror-hover';
 
 // Local styling
 import './editor_styles.scss';
 
 declare module 'codemirror' {
+    interface Editor {
+        querySuggestions(): void;
+    }
     // This is copied from runmode.d.ts. Not sure how to import it :(
     function runMode(
         text: string,
@@ -44,6 +48,10 @@ declare module 'codemirror' {
     function normalizeKeyMap(
         keyMap: Record<string, string | (() => any)>
     ): Record<string, string | (() => any)>;
+
+    interface Editor {
+        querySuggestions(): void;
+    }
 }
 
 attachCustomCommand(CodeMirror.commands);