Merge pull request #23 from PennChopMicrobiomeProgram/codex/plan-implementation-for-natural-language-query-feature

Add natural language query endpoint

Author: Ulthran

app/app.py

@@ -26,6 +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 werkzeug.middleware.proxy_fix import ProxyFix
 
 
@@ -218,6 +219,19 @@ def api_query():
         return {"error": "Query execution failed"}, 500
 
 
+@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")
+    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
+
+
 @app.route("/api", methods=["POST"])
 def api():
     try:

app/nl_query.py

@@ -0,0 +1,28 @@
+"""Utility functions for natural language to SQL queries."""
+
+from typing import Dict, List
+
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+
+def _to_sql(question: str) -> str:
+    """Naively convert a natural language question to SQL.
+
+    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")
+
+
+def run_nl_query(session: Session, question: str) -> Dict[str, List[Dict[str, int]]]:
+    """Execute a natural language query against the database."""
+
+    sql = _to_sql(question)
+    result = session.execute(text(sql)).fetchall()
+    return {"sql": sql, "result": [dict(row._mapping) for row in result]}

tests/test_api.py

@@ -64,3 +64,11 @@ def test_query_api(client):
     data = resp.get_json()
     assert data["recordsTotal"] == 1
     assert data["data"][0]["one"] == 1
+
+
+def test_nl_query_api(client):
+    resp = client.post("/api/nl_query", data={"query": "How many isolates are there?"})
+    assert resp.status_code == 200
+    data = resp.get_json()
+    assert data["sql"] == "SELECT COUNT(*) AS count FROM isolates"
+    assert data["result"][0]["count"] == 0