feat: add sql completions powered by LLM (#1514)

* feat: add sql completions powered by LLM

* fix linter

* fix linter

* address comment

containers/bundled_querybook_config.yaml

@@ -28,6 +28,10 @@ ELASTICSEARCH_HOST: http://elasticsearch:9200
 #         model_args:
 #             model_name: gpt-4o-mini
 #             temperature: 0
+#     sql_complete:
+#         model_args:
+#             model_name: gpt-4o-mini
+#             temperature: 0
 
 # Uncomment below to enable vector store to support embedding based table search.
 # Please check langchain doc for the configs of each provider.

package.json

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

querybook/config/querybook_public_config.yaml

@@ -15,6 +15,9 @@ ai_assistant:
     table_vector_search:
         enabled: false
 
+    sql_complete:
+        enabled: false
+
 survey:
     global_response_cooldown: 2592000 # 30 days
     global_trigger_cooldown: 600 # 10 minutes

querybook/config/user_setting.yaml

@@ -87,3 +87,11 @@ tab:
         - 'tab space 4'
         - tab
     helper: The spaces setting for code editor, this does not modify the existing code.
+
+sql_complete:
+    default: disabled
+    tab: editor
+    options:
+        - enabled
+        - disabled
+    helper: (Experimental) Enable it to receive inline AI-generated SQL completions as you type within the editor.

querybook/server/const/ai_assistant.py

@@ -9,6 +9,7 @@ class AICommandType(Enum):
     SQL_SUMMARY = "sql_summary"
     TABLE_SUMMARY = "table_summary"
     TABLE_SELECT = "table_select"
+    SQL_COMPLETE = "sql_complete"
 
 
 AI_ASSISTANT_NAMESPACE = "/ai_assistant"

querybook/server/datasources_socketio/ai_assistant.py

@@ -30,3 +30,14 @@ def sql_fix(payload={}):
     ai_assistant.query_auto_fix(
         query_execution_id=query_execution_id,
     )
+
+
+@register_socket(AICommandType.SQL_COMPLETE.value, namespace=AI_ASSISTANT_NAMESPACE)
+def sql_complete(payload={}):
+    prefix = payload["prefix"]
+    suffix = payload["suffix"]
+    query_engine_id = payload["query_engine_id"]
+    tables = payload.get("tables", [])
+    ai_assistant.get_sql_completion(
+        query_engine_id=query_engine_id, tables=tables, prefix=prefix, suffix=suffix
+    )

querybook/server/lib/ai_assistant/base_ai_assistant.py

@@ -29,6 +29,7 @@
 from .prompts.table_select_prompt import TABLE_SELECT_PROMPT
 from .prompts.table_summary_prompt import TABLE_SUMMARY_PROMPT
 from .prompts.text_to_sql_prompt import TEXT_TO_SQL_PROMPT
+from .prompts.sql_complete_prompt import SQL_AUTOCOMPLETE_PROMPT
 from .tools.table_schema import (
     get_slimmed_table_schemas,
     get_table_schema_by_name,
@@ -160,6 +161,14 @@ def _get_table_select_prompt(self, top_n, question, table_schemas):
             table_schemas=table_schemas,
         )
 
+    def _get_sql_complete_prompt(self, dialect, table_schemas, prefix, suffix):
+        return SQL_AUTOCOMPLETE_PROMPT.format(
+            dialect=dialect,
+            table_schemas=table_schemas,
+            prefix=prefix,
+            suffix=suffix,
+        )
+
     def _get_error_msg(self, error) -> str:
         """Override this method to return specific error messages for your own assistant."""
         if isinstance(error, ValidationError):
@@ -458,3 +467,45 @@ def find_tables(self, metastore_id, question, session=None):
         except Exception as e:
             LOG.error(e, exc_info=True)
             return []
+
+    @catch_error
+    @with_session
+    @with_ai_socket(command_type=AICommandType.SQL_COMPLETE)
+    def get_sql_completion(
+        self,
+        query_engine_id: int,
+        tables: list[str],
+        prefix: str,
+        suffix: str,
+        socket=None,
+        session=None,
+    ):
+        """
+        Generate SQL completions based on the given context.
+        """
+        query_engine = admin_logic.get_query_engine_by_id(
+            query_engine_id, session=session
+        )
+        table_schemas = get_table_schemas_by_names(
+            metastore_id=query_engine.metastore_id,
+            full_table_names=tables,
+            should_skip_column=self._should_skip_column,
+            session=session,
+        )
+        prompt = self._get_sql_complete_prompt(
+            dialect=query_engine.language,
+            table_schemas=table_schemas,
+            prefix=prefix,
+            suffix=suffix,
+        )
+        llm = self._get_llm(
+            ai_command=AICommandType.SQL_COMPLETE.value,
+            prompt_length=self._get_token_count(
+                AICommandType.SQL_COMPLETE.value, prompt
+            ),
+        )
+
+        chain = llm | JsonOutputParser()
+        response = chain.invoke(prompt)
+        socket.send_data(response)
+        socket.close()

querybook/server/lib/ai_assistant/prompts/sql_complete_prompt.py

@@ -0,0 +1,42 @@
+from langchain.prompts import PromptTemplate
+
+prompt_template = """You are an expert in the {dialect} SQL dialect, skilled in providing precise SQL code completions.
+Your task is to complete the SQL query based on the given context.
+
+<Prefix><FILL_ME><Suffix>
+
+===Table Schemas
+{table_schemas}
+
+===Response Guidelines:
+1. Analyze the partial query and table schemas to understand the context and determine the query's goal.
+2. Identify the relevant tables and columns necessary for the query.
+3. Replace <FILL_ME> with appropriate SQL code, or leave it empty if no completion is needed.
+4. Make sure the completion does not overlap with the prefix or suffix.
+5. Respond in JSON format
+
+===Response Format:
+{{
+    "completion": "the SQL code to replace <FILL_ME>, if any"
+}}
+
+
+===Example
+Input:
+sele<FILL_ME> from some_table
+
+Reasoning:
+The prefix "sele" suggests that the query is likely a SELECT statement. The table schemas indicate the available columns. The completion should be a list of columns to select from the table "some_table".
+As it already has a partial query, the completion should be starting from "ct" to complete the word "select", and then followed by the columns to select.
+
+Output:
+{{
+    "completion": "ct *"
+}}
+
+===Input
+{prefix}<FILL_ME>{suffix}
+"""
+
+
+SQL_AUTOCOMPLETE_PROMPT = PromptTemplate.from_template(prompt_template)

querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx

@@ -38,8 +38,13 @@ export const BoundQueryEditor = React.forwardRef<
     const editorRef = useForwardedRef<IQueryEditorHandles>(ref);
 
     // Code Editor related Props
-    const { codeEditorTheme, options, fontSize, autoCompleteType } =
-        useUserQueryEditorConfig();
+    const {
+        codeEditorTheme,
+        options,
+        fontSize,
+        autoCompleteType,
+        sqlCompleteEnabled,
+    } = useUserQueryEditorConfig();
     const combinedOptions = useMemo(
         () => ({
             ...options,
@@ -79,6 +84,7 @@ export const BoundQueryEditor = React.forwardRef<
             searchContext={searchContext}
             cellId={cellId}
             engineId={engine?.id}
+            sqlCompleteEnabled={sqlCompleteEnabled}
         />
     );
 });

querybook/webapp/components/QueryEditor/QueryEditor.tsx

@@ -22,6 +22,7 @@ import { useKeyMapExtension } from 'hooks/queryEditor/extensions/useKeyMapExtens
 import { useLintExtension } from 'hooks/queryEditor/extensions/useLintExtension';
 import { useOptionsExtension } from 'hooks/queryEditor/extensions/useOptionsExtension';
 import { useSearchExtension } from 'hooks/queryEditor/extensions/useSearchExtension';
+import { useSqlCompleteExtension } from 'hooks/queryEditor/extensions/useSqlCompleteExtension';
 import { useStatusBarExtension } from 'hooks/queryEditor/extensions/useStatusBarExtension';
 import { useAutoComplete } from 'hooks/queryEditor/useAutoComplete';
 import { useCodeAnalysis } from 'hooks/queryEditor/useCodeAnalysis';
@@ -49,6 +50,7 @@ export interface IQueryEditorProps {
     keyMap?: CodeMirrorKeyMap;
     className?: string;
     autoCompleteType?: AutoCompleteType;
+    sqlCompleteEnabled?: boolean;
 
     engineId: number;
     templatedVariables?: TDataDocMetaVariables;
@@ -103,6 +105,7 @@ export const QueryEditor: React.FC<
             keyMap = {},
             className,
             autoCompleteType = 'all',
+            sqlCompleteEnabled = false,
             engineId,
             cellId,
             templatedVariables = [],
@@ -355,6 +358,12 @@ export const QueryEditor: React.FC<
             [onSelection]
         );
 
+        const sqlCompleteExtension = useSqlCompleteExtension({
+            enabled: sqlCompleteEnabled,
+            engineId,
+            tables: tableNamesSet,
+        });
+
         const extensions = useMemo(
             () => [
                 mixedSQL(),
@@ -368,6 +377,7 @@ export const QueryEditor: React.FC<
                 optionsExtension,
                 searchExtension,
                 selectionExtension,
+                sqlCompleteExtension,
             ],
             [
                 keyMapExtention,
@@ -379,6 +389,7 @@ export const QueryEditor: React.FC<
                 optionsExtension,
                 searchExtension,
                 selectionExtension,
+                sqlCompleteExtension,
             ]
         );
 

querybook/webapp/components/UserSettingsMenu/UserSettingsMenu.tsx

@@ -2,6 +2,7 @@ import React, { useMemo } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
 import { UserSettingsTab } from 'components/EnvironmentAppRouter/modalRoute/UserSettingsMenuRoute';
+import PublicConfig from 'config/querybook_public_config.yaml';
 import userSettingConfig from 'config/user_setting.yaml';
 import { titleize } from 'lib/utils';
 import { availableEnvironmentsSelector } from 'redux/environment/selector';
@@ -13,6 +14,8 @@ import { makeSelectOptions, Select } from 'ui/Select/Select';
 
 import './UserSettingsMenu.scss';
 
+const AIAssistantConfig = PublicConfig.ai_assistant;
+
 export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
     tab,
 }) => {
@@ -35,9 +38,16 @@ export const UserSettingsMenu: React.FC<{ tab: UserSettingsTab }> = ({
 
     const settingsToShow = useMemo(
         () =>
-            Object.entries(userSettingConfig).filter(
-                ([key, value]) => value.tab === tab
-            ),
+            Object.entries(userSettingConfig).filter(([key, value]) => {
+                if (key === 'sql_complete') {
+                    return (
+                        AIAssistantConfig.enabled &&
+                        AIAssistantConfig.sql_complete.enabled &&
+                        value.tab === tab
+                    );
+                }
+                return value.tab === tab;
+            }),
         [tab]
     );
 

querybook/webapp/config.d.ts

@@ -102,6 +102,10 @@ declare module 'config/querybook_public_config.yaml' {
             table_vector_search: {
                 enabled: boolean;
             };
+
+            sql_complete: {
+                enabled: boolean;
+            };
         };
         survey?: {
             global_response_cooldown?: number;

querybook/webapp/const/aiAssistant.ts

@@ -5,6 +5,7 @@ export enum AICommandType {
     TEXT_TO_SQL = 'text_to_sql',
     TABLE_SUMMARY = 'table_summary',
     TABLE_SELECT = 'table_select',
+    SQL_COMPLETE = 'sql_complete',
 }
 
 export enum AISocketEvent {