Working NLQ demo

app/app.py

@@ -26,7 +26,7 @@
 from sqlalchemy import text
 from sqlalchemy import select
 from app.datatables import datatables_response, init_app, query_columns
-from app.nl_query import run_nl_query
+from app.nl_query import generate_sql, generate_sql_modification
 from werkzeug.middleware.proxy_fix import ProxyFix
 
 
@@ -222,14 +222,24 @@ def api_query():
 @app.route("/api/nl_query", methods=["POST"])
 def api_nl_query():
     """Translate a natural language question into SQL and execute it."""
-    question = request.form.get("query")
+    question = request.form.get("prompt")
+    starting_query = request.form.get("query")
+
     if not question:
         return {"error": "No query provided"}, 400
-    try:
-        return run_nl_query(db.session, question), 200
-    except Exception as e:
-        print(f"Error executing NL query: {e}")
-        return {"error": "Query execution failed"}, 500
+
+    if not starting_query:
+        try:
+            return generate_sql(question), 200
+        except Exception as e:
+            print(f"Error creating NL query: {e}")
+            return {"error": "Query generation failed"}, 500
+    else:
+        try:
+            return generate_sql_modification(question, starting_query), 200
+        except Exception as e:
+            print(f"Error creating NL query with starting query: {e}")
+            return {"error": "Query modification failed"}, 500
 
 
 @app.route("/api", methods=["POST"])

app/nl_query.py

@@ -1,28 +1,94 @@
-"""Utility functions for natural language to SQL queries."""
+from typing_extensions import Annotated, TypedDict
+from langchain import hub
+from langgraph.graph import START, StateGraph
+from langchain_openai import ChatOpenAI
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.schema import CreateTable
 
-from typing import Dict, List
+# Import SQLAlchemy models installed via requirements
+from marc_db.models import Base
 
-from sqlalchemy import text
-from sqlalchemy.orm import Session
 
+class QueryOutput(TypedDict):
+    """Generated SQL query."""
 
-def _to_sql(question: str) -> str:
-    """Naively convert a natural language question to SQL.
+    query: Annotated[str, ..., "Syntactically valid SQL query."]
 
-    This is a placeholder implementation meant to demonstrate how the
-    translation could be structured. Only a very small set of questions is
-    supported.
-    """
 
-    q = question.lower()
-    if "how many isolates" in q:
-        return "SELECT COUNT(*) AS count FROM isolates"
-    raise ValueError("Unsupported question")
+class State(TypedDict):
+    question: str
+    query: str
 
 
-def run_nl_query(session: Session, question: str) -> Dict[str, List[Dict[str, int]]]:
-    """Execute a natural language query against the database."""
+# Build schema from SQLAlchemy models
+SCHEMA = "\n\n".join(
+    str(CreateTable(table).compile(dialect=sqlite.dialect()))
+    for table in Base.metadata.sorted_tables
+)
 
-    sql = _to_sql(question)
-    result = session.execute(text(sql)).fetchall()
-    return {"sql": sql, "result": [dict(row._mapping) for row in result]}
+llm = ChatOpenAI(
+    model="gpt-4o",
+    temperature=0,
+    max_tokens=None,
+    timeout=None,
+    max_retries=2,
+)
+query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")
+
+
+def write_query(state: State) -> State:
+    prompt = query_prompt_template.invoke(
+        {
+            "dialect": "sqlite",
+            "top_k": 10,
+            "table_info": SCHEMA,
+            "input": state["question"],
+        }
+    )
+    structured_llm = llm.with_structured_output(QueryOutput)
+    result = structured_llm.invoke(prompt)
+    return {"query": result["query"]}
+
+
+def modify_query(state: State) -> State:
+    prompt = query_prompt_template.invoke(
+        {
+            "dialect": "sqlite",
+            "top_k": 10,
+            "table_info": SCHEMA,
+            "input": state["question"],
+            "query": state["query"],
+        }
+    )
+    structured_llm = llm.with_structured_output(QueryOutput)
+    result = structured_llm.invoke(prompt)
+    return {"query": result["query"]}
+
+
+def generate_sql(question: str) -> str:
+    graph_builder = StateGraph(State)
+    graph_builder.add_node("write_query", write_query)
+    graph_builder.add_edge(START, "write_query")
+    graph = graph_builder.compile()
+    result = graph.invoke({"question": question})
+    return result["query"]
+
+
+def generate_sql_modification(question: str, starting_query: str) -> str:
+    graph_builder = StateGraph(State)
+    graph_builder.add_node("modify_query", write_query)
+    graph_builder.add_edge(START, "modify_query")
+    graph = graph_builder.compile()
+    result = graph.invoke({"question": question, "query": starting_query})
+    return result["query"]
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Generate an SQL query for a natural language question."
+    )
+    parser.add_argument("question", help="User question to answer with SQL.")
+    args = parser.parse_args()
+    print(generate_sql(args.question))

app/static/query_editor.js

@@ -7,12 +7,15 @@ function initQueryEditor(containerId, hiddenInputId, tables) {
       keywords.push(f);
     }
   }
+  console.log("Keywords:", keywords);
+  console.log("Tables:", tables);
   const editor = CodeMirror(document.getElementById(containerId), {
     value: input.value || "",
     mode: "text/x-sql",
     lineNumbers: true,
     extraKeys: { "Ctrl-Space": "autocomplete" },
     hintOptions: { tables: tables, keywords: keywords },
+    lineWrapping: true,
   });
 
   editor.on("inputRead", function (cm, change) {

app/templates/query.html

@@ -13,12 +13,25 @@
 {% block body %}
 <div class="container py-2">
 <div class="d-flex flex-row">
-  <form class="d-flex flex-column w-100" method="POST" action="{{ url_for('query') }}">
-    <label class="p-2">Custom Query</label>
-    <input type="hidden" id="query-input" name="query" value="{{ query }}">
-    <div id="query-editor" class="w-100 border" style="height: 300px;"></div>
-    <button type="submit" id="submit" class="my-2 p-2">Submit</button>
-  </form>
+  <div class="d-flex flex-column">
+    <form class="d-flex flex-column w-100" method="POST" action="{{ url_for('query') }}">
+      <label class="p-2">Custom Query</label>
+      <input type="hidden" id="query-input" name="query" value="{{ query }}">
+      <div id="query-editor" class="w-100 border" style="height: 300px; max-width: 300;"></div>
+      <button type="submit" id="submit" class="my-2 p-2">Submit</button>
+    </form>
+
+    <form class="d-flex flex-column w-50 p-5 m-5" method="POST" action="{{ url_for('api_nl_query') }}">
+      <label for="prompt" class="form-label">
+        Natural Language Query
+        <i class="bi bi-question-circle" tabindex="0" data-bs-toggle="popover"
+          data-bs-content="Ask a question about the data. The system will generate an SQL query to answer it."></i>
+      </label>
+      <textarea id="prompt" name="prompt" class="form-control mb-3" rows="3">{{ prompt }}</textarea>
+      <input type="hidden" name="query" value="{{ query }}">
+      <button type="submit" class="btn btn-primary">Submit</button>
+    </form>
+  </div>
 
   <div class="d-flex flex-column w-50 p-5 m-5">
     <label for="model-select" class="form-label">