From 2c5aeaeada104deab9dab2d1176abd106daa31e5 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 27 Mar 2026 20:32:25 -0500 Subject: [PATCH] Fix SQL generation: stricter prompt + auto-retry on schema errors - 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) --- crates/ui/src/main.rs | 70 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 63 insertions(+), 7 deletions(-) diff --git a/crates/ui/src/main.rs b/crates/ui/src/main.rs index 38def9e..3bbd6dc 100644 --- a/crates/ui/src/main.rs +++ b/crates/ui/src/main.rs @@ -260,21 +260,54 @@ fn AskPanel(datasets: Vec) -> 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) -> 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}"))));