Assistant2 (#617)

* Keyboard shortcut to show schema hints (cmd+S / ctrl+S -- note that is a capital "S" so the full kbd commands is cmd+shift+s)
* DB-managed LLM prompts (editable in django admin)
* Versioned .js bundles (for cache busting)
* Automatically populate assistant responses that contain code into the editor
* `#616`_: Update schema/assistant tables/autocomplete on connection drop-down change

HISTORY.rst

@@ -5,6 +5,14 @@ Change Log
 This document records all notable changes to `django-sql-explorer <https://github.com/chrisclark/django-sql-explorer>`_.
 This project adheres to `Semantic Versioning <https://semver.org/>`_.
 
+`vNext`_ (unreleased)
+===========================
+* Keyboard shortcut to show schema hints (cmd+S / ctrl+S -- note that is a capital "S" so the full kbd commands is cmd+shift+s)
+* DB-managed LLM prompts (editable in django admin)
+* Versioned .js bundles (for cache busting)
+* Automatically populate assistant responses that contain code into the editor
+* `#616`_: Update schema/assistant tables/autocomplete on connection drop-down change
+
 `4.2.0`_ (2024-04-26)
 ===========================
 * `#609`_: Tracking should be opt-in and not use the SECRET_KEY

explorer/admin.py

@@ -1,7 +1,7 @@
 from django.contrib import admin
 
 from explorer.actions import generate_report_action
-from explorer.models import Query
+from explorer.models import Query, ExplorerValue
 
 
 @admin.register(Query)
@@ -10,3 +10,18 @@ class QueryAdmin(admin.ModelAdmin):
     list_filter = ("title",)
     raw_id_fields = ("created_by_user",)
     actions = [generate_report_action()]
+
+
+@admin.register(ExplorerValue)
+class ExplorerValueAdmin(admin.ModelAdmin):
+    list_display = ("key", "value", "display_key")
+    list_filter = ("key",)
+    readonly_fields = ("key",)
+    search_fields = ("key", "value")
+
+    def display_key(self, obj):
+        # Human-readable name for the key
+        return dict(ExplorerValue.EXPLORER_SETTINGS_CHOICES).get(obj.key, "")
+
+    display_key.short_description = "Setting Name"
+

explorer/assistant/views.py

@@ -5,8 +5,8 @@
 
 from explorer.telemetry import Stat, StatNames
 from explorer.utils import get_valid_connection
+from explorer.models import ExplorerValue
 from explorer.assistant.models import PromptLog
-from explorer.assistant.prompts import primary_prompt
 from explorer.assistant.utils import (
     do_req, extract_response, tables_from_schema_info,
     get_table_names_from_query, sample_rows_from_tables,
@@ -44,8 +44,10 @@ def run_assistant(request_data, user):
 
     user_prompt += f"## User's Request to Assistant ##\n\n{request_data['assistant_request']}\n\n"
 
-    prompt = primary_prompt.copy()
-    prompt["user"] = user_prompt
+    prompt = {
+        "system": ExplorerValue.objects.get_item(ExplorerValue.ASSISTANT_SYSTEM_PROMPT).value,
+        "user": user_prompt
+    }
 
     start = timezone.now()
     pl = PromptLog(

explorer/migrations/0016_alter_explorervalue_key.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.8 on 2024-04-26 13:05
+
+from django.db import migrations, models
+
+
+def insert_assistant_prompt(apps, schema_editor):
+
+    ExplorerValue = apps.get_model('explorer', 'ExplorerValue')
+    ExplorerValue.objects.get_or_create(
+        key="ASP",
+        value="""You are a data analyst's assistant and will be asked write or modify a SQL query to assist a business
+user with their analysis. The user will provide a prompt of what they are looking for help with, and may also
+provide SQL they have written so far, relevant table schema, and sample rows from the tables they are querying.
+
+For complex requests, you may use Common Table Expressions (CTEs) to break down the problem into smaller parts.
+CTEs are not needed for simpler requests.
+"""
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('explorer', '0015_explorervalue'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='explorervalue',
+            name='key',
+            field=models.CharField(choices=[('UUID', 'Install Unique ID'), ('SMLS', 'Startup metric last send'), ('ASP', 'System prompt for SQL Assistant')], max_length=5, unique=True),
+        ),
+        migrations.RunPython(insert_assistant_prompt),
+    ]

explorer/models.py

@@ -428,16 +428,21 @@ def set_startup_last_send(self, ts):
             obj.value = str(ts)
             obj.save()
 
+    def get_item(self, key):
+        return self.filter(key=key).first()
+
 
 class ExplorerValue(models.Model):
     INSTALL_UUID = "UUID"
     STARTUP_METRIC_LAST_SEND = "SMLS"
+    ASSISTANT_SYSTEM_PROMPT = "ASP"
     EXPLORER_SETTINGS_CHOICES = [
         (INSTALL_UUID, "Install Unique ID"),
         (STARTUP_METRIC_LAST_SEND, "Startup metric last send"),
+        (ASSISTANT_SYSTEM_PROMPT, "System prompt for SQL Assistant")
     ]
 
-    key = models.CharField(max_length=5, choices=EXPLORER_SETTINGS_CHOICES)
+    key = models.CharField(max_length=5, choices=EXPLORER_SETTINGS_CHOICES, unique=True)
     value = models.TextField(null=True, blank=True)
 
     objects = ExplorerValueManager()

explorer/src/js/assistant.js

@@ -2,38 +2,18 @@ import {getCsrfToken} from "./csrf";
 import { marked } from "marked";
 import DOMPurify from "dompurify";
 import * as bootstrap from 'bootstrap';
-import $ from "jquery";
 import List from "list.js";
+import { SchemaSvc } from "./schemaService"
 
 function getErrorMessage() {
     const errorElement = document.querySelector('.alert-danger.db-error');
     return errorElement ? errorElement.textContent.trim() : null;
 }
 
-export function setUpAssistant(expand = false) {
-
-    const error = getErrorMessage();
-
-    if(expand || error) {
-        const myCollapseElement = document.getElementById('assistant_collapse');
-        const bsCollapse = new bootstrap.Collapse(myCollapseElement, {
-          toggle: false
-        });
-        bsCollapse.show();
-        if(error) {
-            document.getElementById('id_error_help_message').classList.remove('d-none');
-        }
-    }
-
-    const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]');
-    [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl));
-
-    fetch('../schema.json/' + $("#id_connection").val())
-    .then(response => {
-        return response.json();
-    })
-    .then(data => {
-        const keys = Object.keys(data);
+function setupTableList() {
+    const conn = document.querySelector('#id_connection').value;
+    SchemaSvc.get(conn).then(schema => {
+        const keys = Object.keys(schema);
         const tableList = document.getElementById('table-list');
         tableList.innerHTML = '';
 
@@ -66,6 +46,29 @@ export function setUpAssistant(expand = false) {
     .catch(error => {
         console.error('Error retrieving JSON schema:', error);
     });
+}
+
+export function setUpAssistant(expand = false) {
+
+    const connEl = document.querySelector('#id_connection');
+    connEl.addEventListener('change', setupTableList);
+    setupTableList();
+
+    const error = getErrorMessage();
+
+    if(expand || error) {
+        const myCollapseElement = document.getElementById('assistant_collapse');
+        const bsCollapse = new bootstrap.Collapse(myCollapseElement, {
+          toggle: false
+        });
+        bsCollapse.show();
+        if(error) {
+            document.getElementById('id_error_help_message').classList.remove('d-none');
+        }
+    }
+
+    const tooltipTriggerList = document.querySelectorAll('[data-bs-toggle="tooltip"]');
+    [...tooltipTriggerList].map(tooltipTriggerEl => new bootstrap.Tooltip(tooltipTriggerEl));
 
     document.getElementById('id_assistant_input').addEventListener('keydown', function(event) {
         if ((event.ctrlKey || event.metaKey) && (event.key === 'Enter')) {
@@ -113,6 +116,19 @@ function submitAssistantAsk() {
         const output = DOMPurify.sanitize(marked.parse(data.message));
         document.getElementById("assistant_response").innerHTML = output;
         document.getElementById("assistant_spinner").classList.add('d-none');
+
+        // If there is exactly one code block in the response and the SQL editor is empty
+        // then copy the code directly into the editor
+        const preElements = document.querySelectorAll('#assistant_response pre');
+        if (preElements.length === 1 && window.editor?.state.doc.toString().trim() === "") {
+            window.editor.dispatch({
+                changes: {
+                    from: 0,
+                    insert: preElements[0].textContent
+                }
+            });
+        }
+
         setUpCopyButtons();
     })
         .catch(error => {

explorer/src/js/codemirror-config.js

@@ -12,13 +12,76 @@ import {autocompletion, completionKeymap, closeBrackets, closeBracketsKeymap, ac
 import {lintKeymap} from "@codemirror/lint"
 import { Prec } from "@codemirror/state";
 import {sql} from "@codemirror/lang-sql";
+import { SchemaSvc } from "./schemaService"
 
 let updateListenerExtension = EditorView.updateListener.of((update) => {
   if (update.docChanged) {
     document.dispatchEvent(new CustomEvent('docChanged', {}));
   }
 });
 
+const hideTooltipOnEsc = EditorView.domEventHandlers({
+    keydown(event, view) {
+        if (event.code === 'Escape') {
+            const tooltip = document.getElementById('schema_tooltip');
+            if (tooltip) {
+                tooltip.classList.add('d-none');
+                tooltip.classList.remove('d-block');
+            }
+            return true;
+        }
+        return false;
+    }
+});
+
+function displaySchemaTooltip(editor, content) {
+    let tooltip = document.getElementById('schema_tooltip');
+    if (tooltip) {
+        tooltip.classList.remove('d-none');
+        tooltip.classList.add('d-block');
+        tooltip.textContent = content;
+    }
+}
+
+function fetchAndShowSchema(view) {
+    const { state } = view;
+    const pos = state.selection.main.head;
+    const wordRange = state.wordAt(pos);
+
+    if (wordRange) {
+        const tableName = state.doc.sliceString(wordRange.from, wordRange.to);
+        const conn = document.querySelector('#id_connection').value;
+        SchemaSvc.get(conn).then(schema => {
+            let formattedSchema;
+            if (schema.hasOwnProperty(tableName)) {
+                formattedSchema = JSON.stringify(schema[tableName], null, 2);
+            } else {
+                formattedSchema = `Table '${tableName}' not found in schema for connection '${conn}'`;
+            }
+            displaySchemaTooltip(view, formattedSchema);
+        });
+    }
+    return true;
+}
+
+const schemaKeymap = [
+    {
+        key: "Ctrl-S",
+        mac: "Cmd-S",
+        run: (editor) => {
+            fetchAndShowSchema(editor);
+            return true;
+        }
+    },
+    {
+        key: "Cmd-S",
+        run: (editor) => {
+            fetchAndShowSchema(editor);
+            return true;
+        }
+    }
+];
+
 const submitEventFromCM = new CustomEvent('submitEventFromCM', {});
 const submitKeymapArr = [
     {
@@ -47,9 +110,7 @@ const autocompleteKeymap = [{key: "Tab", run: acceptCompletion}]
 
 
 export const explorerSetup = (() => [
-    sql({
-        //schema: window.schema_json
-    }),
+    sql({}),
     lineNumbers(),
     highlightActiveLineGutter(),
     highlightSpecialChars(),
@@ -66,6 +127,7 @@ export const explorerSetup = (() => [
     highlightSelectionMatches(),
     submitKeymap,
     updateListenerExtension,
+    hideTooltipOnEsc,
     keymap.of([
         ...closeBracketsKeymap,
         ...defaultKeymap,
@@ -74,6 +136,7 @@ export const explorerSetup = (() => [
         ...foldKeymap,
         ...completionKeymap,
         ...lintKeymap,
-        ...autocompleteKeymap
+        ...autocompleteKeymap,
+        ...schemaKeymap
     ])
 ])()

explorer/src/js/explorer.js

@@ -11,6 +11,25 @@ import { toggleFavorite } from "./favorites";
 
 import {schemaCompletionSource, StandardSQL} from "@codemirror/lang-sql";
 import {StateEffect} from "@codemirror/state";
+import {SchemaSvc} from "./schemaService";
+
+
+function updateSchema() {
+
+    const conn = document.querySelector('#id_connection').value;
+
+    SchemaSvc.get(conn).then(schema => {
+        window.editor.dispatch({
+            effects: StateEffect.appendConfig.of(
+                StandardSQL.language.data.of({
+                  autocomplete: schemaCompletionSource({schema: schema})
+                })
+            )
+        });
+    });
+
+    $("#schema_frame").attr("src", `../schema/{conn}`);
+}
 
 
 function editorFromTextArea(textarea) {
@@ -155,7 +174,6 @@ export class ExplorerEditor {
     }
 
     showSchema(noAutofocus) {
-        $("#schema_frame").attr("src", "../schema/" + $("#id_connection").val());
         if (noAutofocus === true) {
             $("#schema_frame").addClass("no-autofocus");
         }
@@ -346,22 +364,9 @@ export class ExplorerEditor {
             if(event.keyCode === 13){ this.showRows(); }
         }.bind(this));
 
-        fetch('../schema.json/' + $("#id_connection").val())
-            .then(response => {
-                return response.json();
-            })
-            .then(data => {
-                this.editor.dispatch({
-                    effects: StateEffect.appendConfig.of(
-                        StandardSQL.language.data.of({
-                          autocomplete: schemaCompletionSource({schema: data})
-                        })
-                    )
-                })
-                return data;
-            })
-            .catch(error => {
-                console.error('Error retrieving JSON schema:', error);
-            });
+        // Set up schema autocomplete in the editor. When the connection changes, load new schema.
+        const connEl = document.querySelector('#id_connection');
+        connEl.addEventListener('change', updateSchema);
+        updateSchema();
     }
 }