From 13fd67d62a38d8a6a8cf2f042a459acc2038b3b0 Mon Sep 17 00:00:00 2001
From: Chris Clark <chris@untrod.com>
Date: Tue, 9 Apr 2024 10:07:55 -0400
Subject: [PATCH] Sql assistant (#569)

* LLM-powered SQL Assistant
---
 HISTORY.rst                                |   2 +
 docs/features.rst                          |  11 +-
 docs/install.rst                           |  32 +++++
 explorer/__init__.py                       |   6 +-
 explorer/app_settings.py                   |   9 ++
 explorer/assistant/__init__.py             |   0
 explorer/assistant/models.py               |  21 +++
 explorer/assistant/prompts.py              |  11 ++
 explorer/assistant/tests.py                |  90 +++++++++++++
 explorer/assistant/utils.py                |  97 ++++++++++++++
 explorer/assistant/views.py                |  87 ++++++++++++
 explorer/migrations/0014_promptlog.py      |  29 ++++
 explorer/src/js/assistant.js               | 149 +++++++++++++++++++++
 explorer/src/js/explorer.js                |  14 +-
 explorer/src/js/main.js                    |  11 +-
 explorer/src/scss/assistant.scss           |  35 +++++
 explorer/src/scss/explorer.scss            |  12 +-
 explorer/src/scss/styles.scss              |   2 +-
 explorer/templates/explorer/assistant.html |  49 +++++++
 explorer/templates/explorer/play.html      |   3 +
 explorer/templates/explorer/query.html     |   3 +
 explorer/tests/settings.py                 |   1 +
 explorer/tracker.py                        |   4 +-
 explorer/urls.py                           |   2 +
 explorer/views/mixins.py                   |   1 +
 package-lock.json                          |  18 +++
 package.json                               |   4 +-
 requirements/base.txt                      |   1 +
 requirements/optional.txt                  |   2 +
 test_project/settings.py                   |   2 +
 30 files changed, 687 insertions(+), 21 deletions(-)
 create mode 100644 explorer/assistant/__init__.py
 create mode 100644 explorer/assistant/models.py
 create mode 100644 explorer/assistant/prompts.py
 create mode 100644 explorer/assistant/tests.py
 create mode 100644 explorer/assistant/utils.py
 create mode 100644 explorer/assistant/views.py
 create mode 100644 explorer/migrations/0014_promptlog.py
 create mode 100644 explorer/src/js/assistant.js
 create mode 100644 explorer/src/scss/assistant.scss
 create mode 100644 explorer/templates/explorer/assistant.html

diff --git a/HISTORY.rst b/HISTORY.rst
index d8197a76..444d7367 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -7,6 +7,8 @@ This project adheres to `Semantic Versioning <https://semver.org/>`_.
 
 `4.1.0`_ (TBD)
 ===========================
+* SQL Assistant: Built in query help via OpenAI (or LLM of choice), with relevant schema
+  automatically injected into the prompt. Enable via setting EXPLORER_AI_API_KEY.
 * Anonymous usage telemetry. Can be disabled by setting EXPLORER_ENABLE_ANONYMOUS_STATS to False
 * `#594`_: Eliminate <script> tags to prevent potential Content Security Policy issues
 
diff --git a/docs/features.rst b/docs/features.rst
index be6bedb0..98ffb7e1 100644
--- a/docs/features.rst
+++ b/docs/features.rst
@@ -25,13 +25,18 @@ Security
 
 Easy to get started
 -------------------
-- Built on Django's ORM, so works with Postgresql, Mysql, and
-  Sqlite. And, between you and me, it works fine on RedShift as
-  well.
+- Built on Django's ORM, so works with MySQL, Postgres, Oracle,
+  SQLite, Snowflake, MS SQL Server, RedShift, and MariaDB.
 - Small number of dependencies.
 - Just want to get in and write some ad-hoc queries? Go nuts with
   the Playground area.
 
+SQL Assistant
+-------------
+- Built in integration with OpenAI (or the LLM of your choosing)
+  to quickly get help with your query, with relevant schema
+  automatically injected into the prompt.
+
 Snapshots
 ---------
 - Tick the 'snapshot' box on a query, and Explorer will upload a
diff --git a/docs/install.rst b/docs/install.rst
index d85bf223..08d0e9a1 100644
--- a/docs/install.rst
+++ b/docs/install.rst
@@ -84,6 +84,38 @@ And run the server:
 
 You can now browse to http://127.0.0.1:8000/explorer/ and get exploring!
 
+Note that Explorer expects STATIC_URL to be set appropriately. This isn't a problem
+with vanilla Django setups, but if you are using e.g. Django Storages with S3, you
+must set your STATIC_URL to point to your S3 bucket (e.g. s3_bucket_url + '/static/')
+
+AI SQL Assistant
+----------------
+To enable AI features, you must install the OpenAI SDK and Tiktoken library from
+requirements/optional.txt. By default the Assistant is configured to use OpenAI and
+the `gpt-4-0125-preview` model. To use those settings, set an OpenAI API token in
+your project's settings.py file:
+
+``EXPLORER_AI_API_KEY = 'your_openai_api_key'``
+
+Or, more likely:
+
+``EXPLORER_AI_API_KEY = os.environ.get("OPENAI_API_KEY")``
+
+If you would prefer to use a different provider and/or different model, you can
+also override the AI API URL root and default model. For example, this would configure
+the Assistant to use OpenRouter and Mixtral 8x7B Instruct:
+
+..  code-block:: python
+    :emphasize-lines: 5
+
+    EXPLORER_ASSISTANT_MODEL = {"name": "mistralai/mixtral-8x7b-instruct:nitro",
+                                "max_tokens": 32768})
+    EXPLORER_ASSISTANT_BASE_URL = "https://openrouter.ai/api/v1"
+    EXPLORER_AI_API_KEY = os.environ.get("OPENROUTER_API_KEY")
+
+Other Parameters
+----------------
+
 The default behavior when viewing a parameterized query is to autorun the associated
 SQL with the default parameter values. This may perform poorly and you may want
 a chance for your users to review the parameters before running. If so you may add
diff --git a/explorer/__init__.py b/explorer/__init__.py
index 562f3d90..ee38d1c5 100644
--- a/explorer/__init__.py
+++ b/explorer/__init__.py
@@ -1,8 +1,8 @@
 __version_info__ = {
     "major": 4,
-    "minor": 0,
-    "patch": 2,
-    "releaselevel": "final",
+    "minor": 1,
+    "patch": 0,
+    "releaselevel": "beta",
     "serial": 0
 }
 
diff --git a/explorer/app_settings.py b/explorer/app_settings.py
index f7669459..5f52b97a 100644
--- a/explorer/app_settings.py
+++ b/explorer/app_settings.py
@@ -142,3 +142,12 @@
 # - user will need to run by clicking the Save & Run Button to execute
 EXPLORER_AUTORUN_QUERY_WITH_PARAMS = getattr(settings, "EXPLORER_AUTORUN_QUERY_WITH_PARAMS", True)
 VITE_DEV_MODE = getattr(settings, "VITE_DEV_MODE", False)
+
+
+# AI Assistant settings. Setting the first to an OpenAI key is the simplest way to enable the assistant
+EXPLORER_AI_API_KEY = getattr(settings, "EXPLORER_AI_API_KEY", None)
+EXPLORER_ASSISTANT_BASE_URL = getattr(settings, "EXPLORER_ASSISTANT_BASE_URL", "https://api.openai.com/v1")
+EXPLORER_ASSISTANT_MODEL = getattr(settings, "EXPLORER_ASSISTANT_MODEL",
+                                   # Return the model name and max_tokens it supports
+                                   {"name": "gpt-4-0125-preview",
+                                    "max_tokens": 128000})
diff --git a/explorer/assistant/__init__.py b/explorer/assistant/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/explorer/assistant/models.py b/explorer/assistant/models.py
new file mode 100644
index 00000000..21c6db6d
--- /dev/null
+++ b/explorer/assistant/models.py
@@ -0,0 +1,21 @@
+from django.db import models
+from django.conf import settings
+
+
+class PromptLog(models.Model):
+
+    class Meta:
+        app_label = "explorer"
+
+    prompt = models.TextField(blank=True)
+    response = models.TextField(blank=True)
+    run_by_user = models.ForeignKey(
+        settings.AUTH_USER_MODEL,
+        null=True,
+        blank=True,
+        on_delete=models.CASCADE
+    )
+    run_at = models.DateTimeField(auto_now_add=True)
+    duration = models.FloatField(blank=True, null=True)  # seconds
+    model = models.CharField(blank=True, max_length=128, default="")
+    error = models.TextField(blank=True, null=True)
diff --git a/explorer/assistant/prompts.py b/explorer/assistant/prompts.py
new file mode 100644
index 00000000..4e11786c
--- /dev/null
+++ b/explorer/assistant/prompts.py
@@ -0,0 +1,11 @@
+primary_prompt = {
+    "system": """
+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.
+""",
+    "user": ""
+}
diff --git a/explorer/assistant/tests.py b/explorer/assistant/tests.py
new file mode 100644
index 00000000..bd3aed9e
--- /dev/null
+++ b/explorer/assistant/tests.py
@@ -0,0 +1,90 @@
+from explorer.tests.factories import SimpleQueryFactory
+from unittest.mock import patch, Mock
+import unittest
+
+import json
+from django.test import TestCase
+from django.urls import reverse
+from django.contrib.auth.models import User
+from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN
+from explorer import app_settings
+
+
+class TestAssistantViews(TestCase):
+
+    def setUp(self):
+        self.user = User.objects.create_superuser(
+            "admin", "admin@admin.com", "pwd"
+        )
+        self.client.login(username="admin", password="pwd")
+        self.request_data = {
+           "sql": "SELECT * FROM explorer_query",
+           "connection": CONN,
+           "assistant_request": "Test Request"
+        }
+
+    @unittest.skipIf(not app_settings.EXPLORER_AI_API_KEY, "assistant not enabled")
+    @patch("explorer.assistant.utils.openai_client")
+    def test_do_modify_query(self, mocked_openai_client):
+        from explorer.assistant.views import run_assistant
+
+        # create.return_value should match: resp.choices[0].message
+        mocked_openai_client.return_value.chat.completions.create.return_value = Mock(
+            choices=[Mock(message=Mock(content="smart computer"))])
+        resp = run_assistant(self.request_data, None)
+        self.assertEqual(resp, "smart computer")
+
+    @unittest.skipIf(not app_settings.EXPLORER_AI_API_KEY, "assistant not enabled")
+    def test_assistant_help(self):
+        resp = self.client.post(reverse("assistant"),
+                                data=json.dumps(self.request_data),
+                                content_type="application/json")
+        self.assertIsNone(json.loads(resp.content)["message"])
+
+
+class TestPromptContext(TestCase):
+
+    def test_retrieves_sample_rows(self):
+        from explorer.assistant.utils import sample_rows_from_table, ROW_SAMPLE_SIZE
+        SimpleQueryFactory(title="First Query")
+        SimpleQueryFactory(title="Second Query")
+        SimpleQueryFactory(title="Third Query")
+        SimpleQueryFactory(title="Fourth Query")
+        ret = sample_rows_from_table(CONN, "explorer_query")
+        self.assertEqual(len(ret), ROW_SAMPLE_SIZE+1)  # includes header row
+
+    def test_format_rows_from_table(self):
+        from explorer.assistant.utils import format_rows_from_table
+        d = [
+            ["col1", "col2"],
+            ["val1", "val2"],
+        ]
+        ret = format_rows_from_table(d)
+        self.assertEqual(ret, "col1 | col2\n" + "-" * 50 + "\nval1 | val2\n")
+
+    def test_parsing_tables_from_query(self):
+        from explorer.assistant.utils import get_table_names_from_query
+        sql = "SELECT * FROM explorer_query"
+        ret = get_table_names_from_query(sql)
+        self.assertEqual(ret, ["explorer_query"])
+
+    def test_parsing_tables_from_query_bad_sql(self):
+        from explorer.assistant.utils import get_table_names_from_query
+        sql = "foo"
+        ret = get_table_names_from_query(sql)
+        self.assertEqual(ret, [])
+
+    def test_schema_info_from_table_names(self):
+        from explorer.assistant.utils import tables_from_schema_info
+        ret = tables_from_schema_info(CONN, ["explorer_query"])
+        expected = [("explorer_query", [
+            ("id", "AutoField"),
+            ("title", "CharField"),
+            ("sql", "TextField"),
+            ("description", "TextField"),
+            ("created_at", "DateTimeField"),
+            ("last_run_date", "DateTimeField"),
+            ("created_by_user_id", "IntegerField"),
+            ("snapshot", "BooleanField"),
+            ("connection", "CharField")])]
+        self.assertEqual(ret, expected)
diff --git a/explorer/assistant/utils.py b/explorer/assistant/utils.py
new file mode 100644
index 00000000..1a37b08d
--- /dev/null
+++ b/explorer/assistant/utils.py
@@ -0,0 +1,97 @@
+import logging
+from explorer import app_settings
+from explorer.schema import schema_info
+from explorer.utils import get_valid_connection
+from sql_metadata import Parser
+from django.db.utils import OperationalError
+
+if app_settings.EXPLORER_AI_API_KEY:
+    import tiktoken
+    from openai import OpenAI
+
+OPENAI_MODEL = app_settings.EXPLORER_ASSISTANT_MODEL["name"]
+ROW_SAMPLE_SIZE = 2
+
+
+def openai_client():
+    return OpenAI(
+        api_key=app_settings.EXPLORER_AI_API_KEY,
+        base_url=app_settings.EXPLORER_ASSISTANT_BASE_URL
+    )
+
+
+def do_req(prompt):
+    messages = [
+        {"role": "system", "content": prompt["system"]},
+        {"role": "user", "content": prompt["user"]},
+    ]
+    resp = openai_client().chat.completions.create(
+      model=OPENAI_MODEL,
+      messages=messages
+    )
+    messages.append(resp.choices[0].message)
+    logging.info(f"Response: {messages}")
+    return messages
+
+
+def extract_response(r):
+    return r[-1].content
+
+
+def tables_from_schema_info(connection, table_names):
+    schema = schema_info(connection)
+    return [table for table in schema if table[0] in set(table_names)]
+
+
+def sample_rows_from_tables(connection, table_names):
+    ret = None
+    for table_name in table_names:
+        ret = f"SAMPLE FROM TABLE {table_name}:\n"
+        ret = ret + format_rows_from_table(
+            sample_rows_from_table(connection, table_name)
+        ) + "\n\n"
+    return ret
+
+
+def sample_rows_from_table(connection, table_name):
+    conn = get_valid_connection(connection)
+    cursor = conn.cursor()
+    try:
+        cursor.execute(f"SELECT * FROM {table_name} LIMIT {ROW_SAMPLE_SIZE}")
+        ret = [[header[0] for header in cursor.description]]
+        ret = ret + cursor.fetchall()
+        return ret
+    except OperationalError as e:
+        return [[str(e)]]
+
+
+def format_rows_from_table(rows):
+    column_headers = list(rows[0])
+    ret = " | ".join(column_headers) + "\n" + "-" * 50 + "\n"
+    for row in rows[1:]:
+        row_str = " | ".join(str(item) for item in row)
+        ret += row_str + "\n"
+    return ret
+
+
+def get_table_names_from_query(sql):
+    if sql:
+        try:
+            parsed = Parser(sql)
+            return parsed.tables
+        except ValueError:
+            return []
+    return []
+
+
+def num_tokens_from_string(string: str) -> int:
+    """Returns the number of tokens in a text string."""
+    encoding = tiktoken.encoding_for_model(OPENAI_MODEL)
+    num_tokens = len(encoding.encode(string))
+    return num_tokens
+
+
+def fits_in_window(string: str) -> bool:
+    # Ratchet down by 5% to account for other boilerplate and system prompt
+    # TODO make this better by actually looking at the token count of the system prompt
+    return num_tokens_from_string(string) < (app_settings.EXPLORER_ASSISTANT_MODEL["max_tokens"] * 0.95)
diff --git a/explorer/assistant/views.py b/explorer/assistant/views.py
new file mode 100644
index 00000000..feb34f1f
--- /dev/null
+++ b/explorer/assistant/views.py
@@ -0,0 +1,87 @@
+from django.http import JsonResponse
+from django.utils import timezone
+from django.views.decorators.http import require_POST
+import json
+
+from explorer.tracker import Stat, StatNames
+from explorer.utils import get_valid_connection
+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,
+    fits_in_window
+)
+
+
+def run_assistant(request_data, user):
+
+    user_prompt = ""
+
+    db_vendor = get_valid_connection(request_data["connection"]).vendor
+    user_prompt += f"## Database Vendor / SQL Flavor is {db_vendor}\n\n"
+
+    sql = request_data.get("sql")
+    if sql:
+        user_prompt += f"## Existing SQL ##\n\n{sql}\n\n"
+
+    extra_tables = request_data.get("selected_tables", [])
+    included_tables = get_table_names_from_query(sql) + extra_tables
+    if included_tables:
+        table_struct = tables_from_schema_info(request_data["connection"],
+                                               included_tables)
+        user_prompt += f"## Relevant Tables' Structure ##\n\n{table_struct}\n\n"
+
+    if fits_in_window(user_prompt):
+        results_sample = sample_rows_from_tables(request_data["connection"],
+                                                 included_tables)
+        if fits_in_window(user_prompt + results_sample):
+            user_prompt += f"## Sample Results from Tables ##\n\n{results_sample}\n\n"
+
+    user_prompt += f"## User's Request to Assistant ##\n\n{request_data['assistant_request']}\n\n"
+
+    prompt = primary_prompt.copy()
+    prompt["user"] = user_prompt
+
+    start = timezone.now()
+    pl = PromptLog(
+        prompt=prompt,
+        run_by_user=user,
+        run_at=timezone.now(),
+    )
+    response_text = None
+    try:
+        resp = do_req(prompt)
+        response_text = extract_response(resp)
+        pl.response = response_text
+    except Exception as e:
+        pl.error = str(e)
+    finally:
+        stop = timezone.now()
+        pl.duration = (stop - start).total_seconds()
+        pl.save()
+        Stat(StatNames.ASSISTANT_RUN, {
+            "included_table_count": len(included_tables),
+            "extra_table_count": len(extra_tables),
+            "has_sql": bool(sql),
+            "duration": pl.duration,
+        }).track()
+    return response_text
+
+
+@require_POST
+def assistant_help(request):
+    try:
+        data = json.loads(request.body)
+
+        resp = run_assistant(data, request.user)
+
+        response_data = {
+            "status": "success",
+            "message": resp
+        }
+
+        return JsonResponse(response_data)
+
+    except json.JSONDecodeError:
+        return JsonResponse({"status": "error", "message": "Invalid JSON"}, status=400)
diff --git a/explorer/migrations/0014_promptlog.py b/explorer/migrations/0014_promptlog.py
new file mode 100644
index 00000000..38795e34
--- /dev/null
+++ b/explorer/migrations/0014_promptlog.py
@@ -0,0 +1,29 @@
+# Generated by Django 4.2.8 on 2024-01-11 08:22
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('explorer', '0013_querylog_error_querylog_success'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='PromptLog',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('prompt', models.TextField(blank=True)),
+                ('response', models.TextField(blank=True)),
+                ('run_at', models.DateTimeField(auto_now_add=True)),
+                ('duration', models.FloatField(blank=True, null=True)),
+                ('model', models.CharField(blank=True, default='', max_length=128)),
+                ('error', models.TextField(blank=True, null=True)),
+                ('run_by_user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)),
+            ],
+        ),
+    ]
diff --git a/explorer/src/js/assistant.js b/explorer/src/js/assistant.js
new file mode 100644
index 00000000..94f0853f
--- /dev/null
+++ b/explorer/src/js/assistant.js
@@ -0,0 +1,149 @@
+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";
+
+const spinner = "<div class=\"spinner-border text-primary\" role=\"status\"><span class=\"visually-hidden\">Loading...</span></div>";
+
+export function setUpAssistant(expand = false) {
+
+    if(expand) {
+        const myCollapseElement = document.getElementById('assistant_collapse');
+        const bsCollapse = new bootstrap.Collapse(myCollapseElement, {
+          toggle: false
+        });
+        bsCollapse.show();
+    }
+
+    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);
+        const tableList = document.getElementById('table-list');
+        tableList.innerHTML = '';
+
+        keys.forEach((key, index) => {
+            const div = document.createElement('div');
+            div.className = 'form-check';
+
+            const input = document.createElement('input');
+            input.className = 'form-check-input table-checkbox';
+            input.type = 'checkbox';
+            input.value = key;
+            input.id = 'flexCheckDefault' + index;
+
+            const label = document.createElement('label');
+            label.className = 'form-check-label';
+            label.setAttribute('for', input.id);
+            label.textContent = key;
+
+            div.appendChild(input);
+            div.appendChild(label);
+            tableList.appendChild(div);
+        });
+
+        let options = {
+            valueNames: ['form-check-label'],
+        };
+
+        new List('additional_table_container', options);
+    })
+    .catch(error => {
+        console.error('Error retrieving JSON schema:', error);
+    });
+
+    const checkbox = document.getElementById('include_other_tables');
+    const additionalTableContainer = document.getElementById('additional_table_container');
+    const assistantInputWrapper = document.getElementById('assistant_input_wrapper');
+
+    function showHideExtraTables(checked) {
+        if (checked) {
+            additionalTableContainer.classList.remove('d-none');
+            assistantInputWrapper.classList.remove('col-12');
+            assistantInputWrapper.classList.add('col-9');
+        } else {
+            additionalTableContainer.classList.add('d-none');
+            assistantInputWrapper.classList.remove('col-9');
+            assistantInputWrapper.classList.add('col-12');
+        }
+    }
+    checkbox.addEventListener('change', function() {
+        showHideExtraTables(this.checked);
+    });
+    showHideExtraTables(checkbox.checked);
+
+    document.getElementById('ask_assistant_btn').addEventListener('click', function() {
+
+        const selectedTables = Array.from(
+            document.querySelectorAll('.table-checkbox:checked')
+        ).map(cb => cb.value);
+
+        const data = {
+            sql: window.editor?.state.doc.toString() ?? null,
+            connection: document.getElementById("id_connection")?.value ?? null,
+            assistant_request: document.getElementById("id_assistant_input")?.value ?? null,
+            selected_tables: selectedTables
+        };
+
+        document.getElementById("response_block").style.display = "block";
+        document.getElementById("assistant_response").innerHTML = spinner;
+
+        fetch('/assistant/', {
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+                'X-CSRFToken': getCsrfToken()
+            },
+            body: JSON.stringify(data)
+        })
+        .then(response => {
+            if (!response.ok) {
+                throw new Error('Network response was not ok');
+            }
+            return response.json();
+        })
+        .then(data => {
+            const output = DOMPurify.sanitize(marked.parse(data.message));
+            document.getElementById("response_block").style.display = "block";
+            document.getElementById("assistant_response").innerHTML = output;
+            setUpCopyButtons();
+        })
+            .catch(error => {
+            console.error('Error:', error);
+        });
+    });
+}
+
+function setUpCopyButtons(){
+    document.querySelectorAll('#assistant_response pre').forEach(pre => {
+
+        const btn = document.createElement('i');
+        btn.classList.add('copy-btn');
+        btn.classList.add('bi-copy');
+        const msg = document.createElement('span');
+        msg.textContent = 'Copied!';
+        msg.style.display = 'none';
+        msg.style.marginLeft = '8px';
+        btn.appendChild(msg);
+        pre.appendChild(btn);
+
+        btn.addEventListener('click', function() {
+            const code = this.parentNode.firstElementChild.innerText;
+            navigator.clipboard.writeText(code).then(() => {
+                msg.style.display = 'inline';
+                setTimeout(() => {
+                    msg.style.display = 'none';
+                }, 2000);
+            }).catch(err => {
+                console.error('Error in copying text: ', err);
+            });
+        });
+    });
+}
diff --git a/explorer/src/js/explorer.js b/explorer/src/js/explorer.js
index d9c891da..3c1d80bd 100644
--- a/explorer/src/js/explorer.js
+++ b/explorer/src/js/explorer.js
@@ -1,7 +1,7 @@
 import $ from 'jquery';
 import { EditorView } from "codemirror";
-
 import { explorerSetup } from "./codemirror-config";
+import { setUpAssistant } from "./assistant";
 
 import cookie from 'cookiejs';
 import List from 'list.js'
@@ -29,6 +29,16 @@ function editorFromTextArea(textarea) {
 
 export class ExplorerEditor {
     constructor(queryId) {
+
+        const aa = document.getElementById('assistant_accordion');
+        const pa = document.getElementById('nav-preview');
+        if (aa) {
+            // Expand the assistant only if a new query is being created
+            // and no results are yet being shown
+            const expand = !pa && queryId === 'new';
+            setUpAssistant(expand);
+        }
+
         this.queryId = queryId;
         this.$table = $("#preview");
         this.$rows = $("#rows");
@@ -43,6 +53,8 @@ export class ExplorerEditor {
 
         this.editor = editorFromTextArea(document.getElementById("id_sql"));
 
+        window.editor = this.editor;
+
         document.addEventListener('submitEventFromCM', (e) => {
             this.$submit.click();
         });
diff --git a/explorer/src/js/main.js b/explorer/src/js/main.js
index c9851315..2ab7188b 100644
--- a/explorer/src/js/main.js
+++ b/explorer/src/js/main.js
@@ -7,18 +7,15 @@ the page. The dynamic import() allows Vite to chunk the JS and only load what's
 necessary for each page. Concretely, this matters because the pages with SQL
 Editors require fairly heavy JS (CodeMirror).
 */
-
-import * as bootstrap from 'bootstrap'; // eslint-disable-line no-unused-vars
+import * as bootstrap from 'bootstrap';
 
 const route_initializers = {
     explorer_index:      () => import('./query-list').then(({setupQueryList}) => setupQueryList()),
-    query_detail:        () => import('./explorer').then(({ExplorerEditor}) => {
-        const queryId = document.getElementById('queryIdGlobal').value;
-        new ExplorerEditor(queryId);
-    }),
+    query_detail:        () => import('./explorer').then(({ExplorerEditor}) =>
+        new ExplorerEditor(document.getElementById('queryIdGlobal').value)),
     query_create:        () => import('./explorer').then(({ExplorerEditor}) => new ExplorerEditor('new')),
     explorer_playground: () => import('./explorer').then(({ExplorerEditor}) => new ExplorerEditor('new')),
-    explorer_schema:     () => import('./schema').then(({setupSchema}) => setupSchema()),
+    explorer_schema:     () => import('./schema').then(({setupSchema}) => setupSchema())
 };
 
 document.addEventListener('DOMContentLoaded', function() {
diff --git a/explorer/src/scss/assistant.scss b/explorer/src/scss/assistant.scss
new file mode 100644
index 00000000..4e0a6e99
--- /dev/null
+++ b/explorer/src/scss/assistant.scss
@@ -0,0 +1,35 @@
+#assistant_description_label {
+    white-space: pre-wrap;
+}
+
+#assistant_response pre {
+    position: relative;
+    background: antiquewhite;
+}
+
+#assistant_response .copy-btn {
+    position: absolute;
+    top: 5px;
+    right: 5px;
+    cursor: pointer;
+}
+
+#additional_table_container {
+    overflow-y: auto;
+    max-height: 120px;
+}
+
+#assistant_input_parent {
+    max-height: 120px;
+    overflow: hidden;
+}
+
+#response_block::after {
+  content: "";
+  position: absolute;
+  top: -20px;
+  right: 10px;
+  border-width: 10px;
+  border-style: solid;
+  border-color: transparent transparent var(--bs-border-color) transparent;
+}
diff --git a/explorer/src/scss/explorer.scss b/explorer/src/scss/explorer.scss
index 22c4f472..a55e6333 100644
--- a/explorer/src/scss/explorer.scss
+++ b/explorer/src/scss/explorer.scss
@@ -28,7 +28,7 @@
 .data-headers {
     background: white;
     position: sticky;
-    top: 0px;
+    top: 0;
 }
 
 .log-sql {
@@ -37,7 +37,7 @@
 }
 
 .list {
-    -webkit-padding-start: 0px;
+    -webkit-padding-start: 0;
     list-style: none;
 }
 
@@ -58,12 +58,12 @@ td.name.indented {
 }
 
 .stats-th .counter{
-    border: 0px solid white;
+    border: 0 solid white;
 }
 
 .table>thead>tr>th.preview-header {
     white-space: nowrap;
-    border: 0px solid white;
+    border: 0 solid white;
 }
 
 div.sort {
@@ -124,3 +124,7 @@ nav.navbar {
         color: var(--bs-primary-bg-subtle) !important;
     }
 }
+
+.dropdown-toggle::after {
+    margin-left: 0 !important;
+}
diff --git a/explorer/src/scss/styles.scss b/explorer/src/scss/styles.scss
index 4881de73..d19f6c8b 100644
--- a/explorer/src/scss/styles.scss
+++ b/explorer/src/scss/styles.scss
@@ -1,10 +1,10 @@
 @import "~bootstrap/scss/bootstrap";
 
-//$bootstrap-icons-font-file: "~bootstrap-icons/font/fonts/bootstrap-icons";
 $bootstrap-icons-font-dir: "../../../node_modules/bootstrap-icons/font/fonts";
 @import "~bootstrap-icons/font/bootstrap-icons.css";
 
 @import "variables";
 @import "explorer";
+@import "assistant";
 
 @import "pivot.css";
diff --git a/explorer/templates/explorer/assistant.html b/explorer/templates/explorer/assistant.html
new file mode 100644
index 00000000..bfb2ad0f
--- /dev/null
+++ b/explorer/templates/explorer/assistant.html
@@ -0,0 +1,49 @@
+{% load i18n %}
+<div class ="accordion accordion-flush my-3 border" id="assistant_accordion">
+    <div class="accordion-item">
+        <div class="accordion-header" id="assistant_accordion_header">
+            <button class="accordion-button bg-light collapsed" type="button" data-bs-toggle="collapse" data-bs-target="#assistant_collapse" aria-expanded="false" aria-controls="assistant_collapse">
+                <label for="id_assistant_input">SQL Assistant</label>
+            </button>
+        </div>
+    </div>
+    <div id="assistant_collapse" class="accordion-collapse collapse" aria-labelledby="assistant_accordion_header" data-bs-parent="#assistant_accordion">
+        <div class="accordion-body">
+            <div class="row assistant_input_parent">
+                <div class="mt-3 col-12" id="assistant_input_wrapper">
+                    <textarea
+                        class="form-control" id="id_assistant_input"
+                        name="sql_assistant" rows="3" placeholder="What do you need help with?"></textarea>
+                    <label for="id_assistant_input" class="form-label" id="id_assistant_input_label"></label>
+                </div>
+                <div id="additional_table_container" class="d-none col-3">
+                    <input id="search_assistant_tables" class="search" placeholder="{% trans "Search Tables" %}" />
+                    <label for="search_assistant_tables" class="d-none" id="search_assistant_tables_label">Search tables</label>
+                    <div id="table-list" class="list"></div>
+                </div>
+            </div>
+
+            <div class="row">
+                <div class="col-6">
+                    <div class="form-check form-check-inline">
+                        <input type="checkbox" class="form-check-input" name="assistant-include-other-tables" id="include_other_tables" autocomplete="off">
+                        <label class="form-check-label" for="include_other_tables">Include Other Tables (<a class="small text-secondary" style="cursor: pointer;"
+                            data-bs-toggle="tooltip" data-bs-placement="top"
+                            data-bs-title="SQL Assistant builds a prompt with your query, your request, relevant schema (tables referenced in your query) and sample data from those tables. You can optionally include other tables (schema + data sample) that you'd like SQL Assistant to know about.">What is this?</a>)</label>
+                    </div>
+                </div>
+                <div class="col-6 text-end">
+                    <div class="btn-group" role="group">
+                        <button type="button" class="btn btn-outline-primary" id="ask_assistant_btn">Ask Assistant</button>
+                    </div>
+                </div>
+            </div>
+
+
+            <div id="response_block" style="display: none" class="position-relative">
+                <div class="mt-3 p-2 pt-3 rounded-2 border bg-light" id="assistant_response"></div>
+            </div>
+        </div>
+    </div>
+</div>
+
diff --git a/explorer/templates/explorer/play.html b/explorer/templates/explorer/play.html
index f3d0ddba..c555a0b0 100644
--- a/explorer/templates/explorer/play.html
+++ b/explorer/templates/explorer/play.html
@@ -64,6 +64,9 @@ <h2>{% trans "Playground" %}</h2>
                     </div>
                 </div>
                 <input type="hidden" value="{% trans 'Playground Query' %}" name="title" />
+                {% if assistant_enabled %}
+                    {% include 'explorer/assistant.html' %}
+                {% endif %}
             </form>
         </div>
         <div id="schema" style="display: none;">
diff --git a/explorer/templates/explorer/query.html b/explorer/templates/explorer/query.html
index ab5d8d64..63f36b56 100644
--- a/explorer/templates/explorer/query.html
+++ b/explorer/templates/explorer/query.html
@@ -126,6 +126,9 @@ <h2>
                         {% endif %}
                     </div>
                 </div>
+                {% if assistant_enabled %}
+                    {% include 'explorer/assistant.html' %}
+                {% endif %}
                 </form>
             </div>
         </div>
diff --git a/explorer/tests/settings.py b/explorer/tests/settings.py
index 9cff1378..d956466d 100644
--- a/explorer/tests/settings.py
+++ b/explorer/tests/settings.py
@@ -2,6 +2,7 @@
 
 EXPLORER_ENABLE_ANONYMOUS_STATS = False
 EXPLORER_TASKS_ENABLED = False  # set to true to test async tasks
+EXPLORER_AI_API_KEY = None  # set to any value to enable assistant
 CELERY_BROKER_URL = "redis://localhost:6379/0"
 CELERY_TASK_ALWAYS_EAGER = True
 
diff --git a/explorer/tracker.py b/explorer/tracker.py
index 9c2f9cb7..81d58c67 100644
--- a/explorer/tracker.py
+++ b/explorer/tracker.py
@@ -56,6 +56,7 @@ class StatNames(SelfNamedEnum):
     QUERY_RUN = auto()
     QUERY_STREAM = auto()
     STARTUP_STATS = auto()
+    ASSISTANT_RUN = auto()
 
 
 class Stat:
@@ -109,7 +110,8 @@ def gather_summary_stats():
             "debug": settings.DEBUG,
             "tasks_enabled": app_settings.ENABLE_TASKS,
             "unsafe_rendering": app_settings.UNSAFE_RENDERING,
-            "transform_count": len(app_settings.EXPLORER_TRANSFORMS)
+            "transform_count": len(app_settings.EXPLORER_TRANSFORMS),
+            "assistant_enabled": app_settings.EXPLORER_AI_API_KEY is not None,
         }
     except Exception as e:
         return {"error": "error gathering stats: %s" % e}
diff --git a/explorer/urls.py b/explorer/urls.py
index ae0947b8..68a57464 100644
--- a/explorer/urls.py
+++ b/explorer/urls.py
@@ -6,6 +6,7 @@
     StreamQueryView, format_sql,
 )
 
+from assistant.views import assistant_help
 
 urlpatterns = [
     path(
@@ -42,4 +43,5 @@
     path("favorites/", QueryFavoritesView.as_view(), name="query_favorites"),
     path("favorite/<int:query_id>", QueryFavoriteView.as_view(), name="query_favorite"),
     path("", ListQueryView.as_view(), name="explorer_index"),
+    path("assistant/", assistant_help, name="assistant"),
 ]
diff --git a/explorer/views/mixins.py b/explorer/views/mixins.py
index 6689c234..f973acbe 100644
--- a/explorer/views/mixins.py
+++ b/explorer/views/mixins.py
@@ -14,6 +14,7 @@ def gen_ctx(self):
             "can_change": app_settings.EXPLORER_PERMISSION_CHANGE(
                 self.request
             ),
+            "assistant_enabled": app_settings.EXPLORER_AI_API_KEY is not None,
             "csrf_cookie_name": settings.CSRF_COOKIE_NAME,
             "csrf_cookie_httponly": settings.CSRF_COOKIE_HTTPONLY,
             "view_name": self.request.resolver_match.view_name,
diff --git a/package-lock.json b/package-lock.json
index 10e19af9..994f19d0 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -13,8 +13,10 @@
         "bootstrap-icons": "^1.11.2",
         "codemirror": "^6.0.1",
         "cookiejs": "^2.1.3",
+        "dompurify": "^3.0.7",
         "jquery": "^3.7.1",
         "list.js": "^2.3.1",
+        "marked": "^11.1.1",
         "sortablejs": "^1.15.2"
       },
       "devDependencies": {
@@ -1139,6 +1141,11 @@
       "resolved": "https://registry.npmjs.org/crelt/-/crelt-1.0.6.tgz",
       "integrity": "sha512-VQ2MBenTq1fWZUH9DJNGti7kKv6EeAuYr3cLwxUWhIu1baTaXh4Ib5W2CqHVqib4/MqbYGJqiL3Zb8GJZr3l4g=="
     },
+    "node_modules/dompurify": {
+      "version": "3.0.7",
+      "resolved": "https://registry.npmjs.org/dompurify/-/dompurify-3.0.7.tgz",
+      "integrity": "sha512-BViYTZoqP3ak/ULKOc101y+CtHDUvBsVgSxIF1ku0HmK6BRf+C03MC+tArMvOPtVtZp83DDh5puywKDu4sbVjQ=="
+    },
     "node_modules/esbuild": {
       "version": "0.19.10",
       "resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.19.10.tgz",
@@ -1267,6 +1274,17 @@
         "node": "^6.0 || ^8.0 || ^10.0 || ^12.0 || >=14"
       }
     },
+    "node_modules/marked": {
+      "version": "11.1.1",
+      "resolved": "https://registry.npmjs.org/marked/-/marked-11.1.1.tgz",
+      "integrity": "sha512-EgxRjgK9axsQuUa/oKMx5DEY8oXpKJfk61rT5iY3aRlgU6QJtUcxU5OAymdhCvWvhYcd9FKmO5eQoX8m9VGJXg==",
+      "bin": {
+        "marked": "bin/marked.js"
+      },
+      "engines": {
+        "node": ">= 18"
+      }
+    },
     "node_modules/nanoid": {
       "version": "3.3.7",
       "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.7.tgz",
diff --git a/package.json b/package.json
index f5c09a68..22fc81c8 100644
--- a/package.json
+++ b/package.json
@@ -29,6 +29,8 @@
     "cookiejs": "^2.1.3",
     "jquery": "^3.7.1",
     "list.js": "^2.3.1",
-    "sortablejs": "^1.15.2"
+    "sortablejs": "^1.15.2",
+    "dompurify": "^3.0.7",
+    "marked": "^11.1.1"
   }
 }
diff --git a/requirements/base.txt b/requirements/base.txt
index 463ba439..05366204 100644
--- a/requirements/base.txt
+++ b/requirements/base.txt
@@ -2,3 +2,4 @@ sqlparse>=0.4.0
 coverage
 factory-boy>=3.1.0
 requests~=2.31.0
+sql_metadata>=2.10
diff --git a/requirements/optional.txt b/requirements/optional.txt
index e22b30ff..afb9d1cb 100644
--- a/requirements/optional.txt
+++ b/requirements/optional.txt
@@ -6,3 +6,5 @@ boto3>=1.20.0
 xlsxwriter>=1.3.6
 matplotlib<4
 seaborn<0.12
+openai>=1.6.1
+tiktoken>=0.6
diff --git a/test_project/settings.py b/test_project/settings.py
index 84f66a2d..f506c155 100644
--- a/test_project/settings.py
+++ b/test_project/settings.py
@@ -115,3 +115,5 @@
 CELERY_BROKER_URL = 'redis://localhost:6379/0'
 
 EXPLORER_S3_BUCKET = "thisismybucket.therearemanylikeit.butthisoneismine"
+EXPLORER_AI_API_KEY = os.environ.get("AI_API_KEY")
+EXPLORER_ASSISTANT_BASE_URL = os.environ.get("AI_BASE_URL")