profit/lakehouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix SQL generation: stricter prompt + auto-retry on schema errors

Browse Source 
- Prompt now says "CRITICAL: ONLY use columns from schema, do NOT invent"
- Strips markdown backticks from model output
- Auto-retry: if SQL fails with "Schema error" or "No field named",
  feeds the error + schema back to the model for a corrected query
- Both button click and Enter key paths have retry logic

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root 2026-03-27 20:32:25 -05:00
parent 399fc81ab5
commit 2c5aeaeada
1 changed files with 63 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
crates/ui/src/main.rs
View File

@ -260,21 +260,54 @@ fn AskPanel(datasets: Vec<Dataset>) -> Element {
// Step 2: Generate SQL
step.set("writing SQL...".into());
let prompt = format!(
"You are a SQL assistant for a data lakehouse using Apache DataFusion (PostgreSQL-compatible SQL).\n\n\
"You are a SQL assistant. You write Apache DataFusion SQL (PostgreSQL-compatible).\n\n\
CRITICAL: You MUST only use column names that appear in the schema below. Do NOT invent or guess column names.\n\n\
{schema_ctx}\n\
User question: {q}\n\n\
Write a SQL query that answers this question. Output ONLY the SQL query, nothing else. No markdown, no explanation, no backticks."
Rules:\n\
- ONLY use table and column names listed above\n\
- If unsure about a column name, pick the closest match from the schema\n\
- Output ONLY the SQL query. No markdown, no explanation, no backticks."
);
match ai_generate(&prompt, 512).await {
Ok(resp) => {
let sql = resp.text.trim().to_string();
// Clean markdown backticks if model adds them
let sql = sql.trim_start_matches("```sql").trim_start_matches("```").trim_end_matches("```").trim().to_string();
generated_sql.set(Some(sql.clone()));
// Step 3: Execute
step.set("running query...".into());
let query_result = run_sql(&sql).await;
result.set(Some(query_result));
// Step 3b: If schema error, retry with the error as feedback
if let Err(ref err) = query_result {
if err.contains("Schema error") || err.contains("No field named") {
step.set("fixing SQL...".into());
let retry_prompt = format!(
"The SQL you wrote had an error:\n{err}\n\n\
{schema_ctx}\n\
Original question: {q}\n\n\
Write a CORRECTED SQL query using ONLY the columns listed in the schema. Output ONLY SQL."
);
if let Ok(retry_resp) = ai_generate(&retry_prompt, 512).await {
let retry_sql = retry_resp.text.trim()
.trim_start_matches("```sql").trim_start_matches("```")
.trim_end_matches("```").trim().to_string();
generated_sql.set(Some(retry_sql.clone()));
step.set("running corrected query...".into());
let retry_result = run_sql(&retry_sql).await;
result.set(Some(retry_result));
} else {
result.set(Some(query_result));
}
} else {
result.set(Some(query_result));
}
} else {
result.set(Some(query_result));
}
}
Err(e) => {
result.set(Some(Err(format!("AI error: {e}"))));
@ -312,18 +345,41 @@ fn AskPanel(datasets: Vec<Dataset>) -> Element {
let schema_ctx = get_schema_context(&ds).await;
step.set("writing SQL...".into());
let prompt = format!(
"You are a SQL assistant for a data lakehouse using Apache DataFusion (PostgreSQL-compatible SQL).\n\n\
"You are a SQL assistant. You write Apache DataFusion SQL (PostgreSQL-compatible).\n\n\
CRITICAL: You MUST only use column names that appear in the schema below. Do NOT invent or guess column names.\n\n\
{schema_ctx}\n\
User question: {q}\n\n\
Write a SQL query that answers this question. Output ONLY the SQL query, nothing else. No markdown, no explanation, no backticks."
Rules:\n\
- ONLY use table and column names listed above\n\
- If unsure about a column name, pick the closest match from the schema\n\
- Output ONLY the SQL query. No markdown, no explanation, no backticks."
);
match ai_generate(&prompt, 512).await {
Ok(resp) => {
let sql = resp.text.trim().to_string();
let sql = resp.text.trim().trim_start_matches("```sql").trim_start_matches("```").trim_end_matches("```").trim().to_string();
generated_sql.set(Some(sql.clone()));
step.set("running query...".into());
let query_result = run_sql(&sql).await;
result.set(Some(query_result));
if let Err(ref err) = query_result {
if err.contains("Schema error") || err.contains("No field named") {
step.set("fixing SQL...".into());
let retry_prompt = format!(
"The SQL you wrote had an error:\n{err}\n\n{schema_ctx}\n\nOriginal question: {q}\n\nWrite a CORRECTED SQL query using ONLY the columns listed. Output ONLY SQL."
);
if let Ok(rr) = ai_generate(&retry_prompt, 512).await {
let rsql = rr.text.trim().trim_start_matches("```sql").trim_start_matches("```").trim_end_matches("```").trim().to_string();
generated_sql.set(Some(rsql.clone()));
step.set("running corrected query...".into());
result.set(Some(run_sql(&rsql).await));
} else {
result.set(Some(query_result));
}
} else {
result.set(Some(query_result));
}
} else {
result.set(Some(query_result));
}
}
Err(e) => {
result.set(Some(Err(format!("AI error: {e}"))));