Ability to iterate on SQL queries with follow-up prompts

[simonw]

I ran this:

For llm_prices calculate for each model the price if I have 10000 input tokens and 500 output tokens

And got this:

Now I want to send a follow-up prompt saying "return cost in cents"

[simonw]

This relates to:

• 4

It's another way of prompting from the query page, only here it's not just for errors.

[simonw]

One option could be a query action menu item - but that's not visible, so it would be several clicks to activate.

Or it could use the new top_query() hook from Datasette 1.0a8.

[simonw]

First prototype:

diff --git a/datasette_query_assistant/__init__.py b/datasette_query_assistant/__init__.py
index 5228bb3..008c515 100644
--- a/datasette_query_assistant/__init__.py
+++ b/datasette_query_assistant/__init__.py
@@ -2,7 +2,7 @@ from anthropic import AsyncAnthropic
 from datasette import hookimpl, Response, Forbidden
 import os
 import urllib
-import re
+from markupsafe import escape
 from datasette.utils import sqlite3
 from typing import List, Set

@@ -67,7 +67,11 @@ async def generate_sql(client, messages, prefix=""):
     return prefix + message.content[0].text

-async def generate_sql_with_retries(client, db, question, schema, max_retries=3):
+async def generate_sql_with_retries(
+    client, db, question, schema, sql=None, max_retries=3
+):
+    if sql:
+        question = "Previous query:\n" + sql + "\n\n" + question
     messages = [
         {"role": "user", "content": "The table schema is:\n" + schema},
         {"role": "assistant", "content": "Ask questions to generate SQL"},
@@ -122,6 +126,7 @@ async def assistant(request, datasette):
     if request.method == "POST":
         post_vars = await request.post_vars()
         question = (post_vars.get("question") or "").strip()
+        sql = post_vars.get("sql") or None
         table = post_vars.get("table") or None
         if not question:
             datasette.add_message(request, "Question is required", datasette.ERROR)
@@ -132,7 +137,7 @@ async def assistant(request, datasette):

         client = AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

-        sql = await generate_sql_with_retries(client, db, question, schema)
+        sql = await generate_sql_with_retries(client, db, question, schema, sql=sql)
         return Response.redirect(
             datasette.urls.database(database)
             + "?"
@@ -183,6 +188,23 @@ def database_actions(datasette, actor, database):
     return inner

+@hookimpl
+def top_query(request, datasette, database, sql):
+    return """
+    <details><summary>AI query assistant</summary>
+    <form action="{}/-/assistant" method="post">
+    <p><textarea placeholder="Describe a change to make to this query" name="question" style="width: 80%; height: 3em"></textarea></p>
+    <p>
+      <input type="submit" value="Update SQL">
+      <input type="hidden" name="sql" value="{}">
+      <input type="hidden" name="csrftoken" value="{}">
+    </p>
+    </form></details>
+    """.format(
+        datasette.urls.database(database), escape(sql), request.scope["csrftoken"]()
+    )
+
+
 @hookimpl
 def register_routes():
     return [

[simonw]

I'm not sold on the UI here. It also loses track of which tables were in the original schema and goes back to sending the full schema - and it doesn't include the previous conversation, just the previous SQL query.

[simonw]

I think if you came in from the SQL query assistant the UI for a follow-up should be expanded by default. If you ran a query yourself by hand it can be collapsed.

It should also show below the existing SQL query textarea - though that's not where the plugin hook puts it, so I'd have to re-arrange the page with JavaScript a bit.

[simonw]

It could expand by default if there's an error message, and prefill the textarea with a request to fix that error - addressing #4