feat: add AI Explain Query and Fix Error via Ollama

Extract shared call_ollama_chat helper from generate_sql to reuse
settings loading and Ollama API call logic. Add two new AI commands:
- explain_sql: explains what a SQL query does in plain language
- fix_sql_error: suggests corrected SQL based on the error and schema

UI additions: "AI Explain" toolbar button, "Explain" and "Fix with AI"
action buttons on query errors, inline explanation display in results.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-02-13 18:48:39 +03:00
parent c001ad597f
commit 3ad0ee5cc3
6 changed files with 268 additions and 33 deletions

View File

@@ -73,16 +73,13 @@ pub async fn list_ollama_models(ollama_url: String) -> TuskResult<Vec<OllamaMode
Ok(tags.models) Ok(tags.models)
} }
#[tauri::command] async fn call_ollama_chat(
pub async fn generate_sql( app: &AppHandle,
app: AppHandle, system_prompt: String,
state: State<'_, Arc<AppState>>, user_content: String,
connection_id: String,
prompt: String,
) -> TuskResult<String> { ) -> TuskResult<String> {
// Load AI settings
let settings = { let settings = {
let path = get_ai_settings_path(&app)?; let path = get_ai_settings_path(app)?;
if !path.exists() { if !path.exists() {
return Err(TuskError::Ai( return Err(TuskError::Ai(
"No AI model selected. Open AI settings to choose a model.".to_string(), "No AI model selected. Open AI settings to choose a model.".to_string(),
@@ -98,24 +95,6 @@ pub async fn generate_sql(
)); ));
} }
// Build schema context
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are a PostgreSQL SQL generator. Given the database schema below and a natural language request, \
output ONLY a valid PostgreSQL SQL query. Do not include any explanation, markdown formatting, \
or code fences. Output raw SQL only.\n\n\
RULES:\n\
- Use FK relationships for correct JOIN conditions.\n\
- timestamp - timestamp = interval. To get a number use EXTRACT(EPOCH FROM (ts1 - ts2)).\n\
- interval cannot be cast to numeric directly.\n\
- When using UNION/UNION ALL, ensure matching column types; cast enums to text if they differ.\n\
- Use COALESCE for nullable columns in aggregations when appropriate.\n\
- Prefer LEFT JOIN when the related row may not exist.\n\n\
DATABASE SCHEMA:\n{}",
schema_text
);
let request = OllamaChatRequest { let request = OllamaChatRequest {
model: settings.model, model: settings.model,
messages: vec![ messages: vec![
@@ -125,7 +104,7 @@ pub async fn generate_sql(
}, },
OllamaChatMessage { OllamaChatMessage {
role: "user".to_string(), role: "user".to_string(),
content: prompt, content: user_content,
}, },
], ],
stream: false, stream: false,
@@ -162,8 +141,89 @@ pub async fn generate_sql(
.await .await
.map_err(|e| TuskError::Ai(format!("Failed to parse Ollama response: {}", e)))?; .map_err(|e| TuskError::Ai(format!("Failed to parse Ollama response: {}", e)))?;
let sql = clean_sql_response(&chat_resp.message.content); Ok(chat_resp.message.content)
Ok(sql) }
#[tauri::command]
pub async fn generate_sql(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
prompt: String,
) -> TuskResult<String> {
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are a PostgreSQL SQL generator. Given the database schema below and a natural language request, \
output ONLY a valid PostgreSQL SQL query. Do not include any explanation, markdown formatting, \
or code fences. Output raw SQL only.\n\n\
RULES:\n\
- Use FK relationships for correct JOIN conditions.\n\
- timestamp - timestamp = interval. To get a number use EXTRACT(EPOCH FROM (ts1 - ts2)).\n\
- interval cannot be cast to numeric directly.\n\
- When using UNION/UNION ALL, ensure matching column types; cast enums to text if they differ.\n\
- Use COALESCE for nullable columns in aggregations when appropriate.\n\
- Prefer LEFT JOIN when the related row may not exist.\n\n\
DATABASE SCHEMA:\n{}",
schema_text
);
let raw = call_ollama_chat(&app, system_prompt, prompt).await?;
Ok(clean_sql_response(&raw))
}
#[tauri::command]
pub async fn explain_sql(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
sql: String,
) -> TuskResult<String> {
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are a PostgreSQL expert. Explain what this SQL query does in clear, concise language. \
Focus on the business logic, mention the tables, joins, and filters used. \
Use short paragraphs or bullet points.\n\n\
DATABASE SCHEMA:\n{}",
schema_text
);
call_ollama_chat(&app, system_prompt, sql).await
}
#[tauri::command]
pub async fn fix_sql_error(
app: AppHandle,
state: State<'_, Arc<AppState>>,
connection_id: String,
sql: String,
error_message: String,
) -> TuskResult<String> {
let schema_text = build_schema_context(&state, &connection_id).await?;
let system_prompt = format!(
"You are a PostgreSQL expert. Fix the SQL query based on the error message. \
Output ONLY the corrected valid PostgreSQL SQL. Do not include any explanation, \
markdown formatting, or code fences. Output raw SQL only.\n\n\
RULES:\n\
- Use FK relationships for correct JOIN conditions.\n\
- timestamp - timestamp = interval. To get a number use EXTRACT(EPOCH FROM (ts1 - ts2)).\n\
- interval cannot be cast to numeric directly.\n\
- When using UNION/UNION ALL, ensure matching column types; cast enums to text if they differ.\n\
- Use COALESCE for nullable columns in aggregations when appropriate.\n\
- Prefer LEFT JOIN when the related row may not exist.\n\n\
DATABASE SCHEMA:\n{}",
schema_text
);
let user_content = format!(
"Original SQL:\n{}\n\nError:\n{}",
sql, error_message
);
let raw = call_ollama_chat(&app, system_prompt, user_content).await?;
Ok(clean_sql_response(&raw))
} }
async fn build_schema_context( async fn build_schema_context(

View File

@@ -94,6 +94,8 @@ pub fn run() {
commands::ai::save_ai_settings, commands::ai::save_ai_settings,
commands::ai::list_ollama_models, commands::ai::list_ollama_models,
commands::ai::generate_sql, commands::ai::generate_sql,
commands::ai::explain_sql,
commands::ai::fix_sql_error,
// lookup // lookup
commands::lookup::entity_lookup, commands::lookup::entity_lookup,
]) ])

View File

@@ -1,7 +1,8 @@
import { ResultsTable } from "./ResultsTable"; import { ResultsTable } from "./ResultsTable";
import { ResultsJsonView } from "./ResultsJsonView"; import { ResultsJsonView } from "./ResultsJsonView";
import type { QueryResult } from "@/types"; import type { QueryResult } from "@/types";
import { Loader2, AlertCircle } from "lucide-react"; import { Loader2, AlertCircle, Sparkles, Wand2 } from "lucide-react";
import { Button } from "@/components/ui/button";
interface Props { interface Props {
result?: QueryResult | null; result?: QueryResult | null;
@@ -14,6 +15,10 @@ interface Props {
value: unknown value: unknown
) => void; ) => void;
highlightedCells?: Set<string>; highlightedCells?: Set<string>;
aiExplanation?: string | null;
isAiLoading?: boolean;
onExplainError?: () => void;
onFixError?: () => void;
} }
export function ResultsPanel({ export function ResultsPanel({
@@ -23,6 +28,10 @@ export function ResultsPanel({
viewMode = "table", viewMode = "table",
onCellDoubleClick, onCellDoubleClick,
highlightedCells, highlightedCells,
aiExplanation,
isAiLoading,
onExplainError,
onFixError,
}: Props) { }: Props) {
if (isLoading) { if (isLoading) {
return ( return (
@@ -33,13 +42,65 @@ export function ResultsPanel({
); );
} }
if (aiExplanation) {
return (
<div className="h-full overflow-auto p-4">
<div className="rounded-md border bg-muted/30 p-4">
<div className="mb-2 flex items-center gap-2 text-xs font-medium text-muted-foreground">
<Sparkles className="h-3.5 w-3.5" />
AI Explanation
</div>
<pre className="whitespace-pre-wrap font-sans text-sm leading-relaxed text-foreground">
{aiExplanation}
</pre>
</div>
</div>
);
}
if (error) { if (error) {
return ( return (
<div className="flex h-full items-center justify-center p-4"> <div className="flex h-full flex-col items-center justify-center gap-3 p-4">
<div className="flex items-start gap-2 rounded-md border border-destructive/50 bg-destructive/10 p-3 text-sm text-destructive"> <div className="flex items-start gap-2 rounded-md border border-destructive/50 bg-destructive/10 p-3 text-sm text-destructive">
<AlertCircle className="mt-0.5 h-4 w-4 shrink-0" /> <AlertCircle className="mt-0.5 h-4 w-4 shrink-0" />
<pre className="whitespace-pre-wrap font-mono text-xs">{error}</pre> <pre className="whitespace-pre-wrap font-mono text-xs">{error}</pre>
</div> </div>
{(onExplainError || onFixError) && (
<div className="flex items-center gap-2">
{onExplainError && (
<Button
size="sm"
variant="outline"
className="h-7 gap-1.5 text-xs"
onClick={onExplainError}
disabled={isAiLoading}
>
{isAiLoading ? (
<Loader2 className="h-3 w-3 animate-spin" />
) : (
<Sparkles className="h-3 w-3" />
)}
Explain
</Button>
)}
{onFixError && (
<Button
size="sm"
variant="outline"
className="h-7 gap-1.5 text-xs"
onClick={onFixError}
disabled={isAiLoading}
>
{isAiLoading ? (
<Loader2 className="h-3 w-3 animate-spin" />
) : (
<Wand2 className="h-3 w-3" />
)}
Fix with AI
</Button>
)}
</div>
)}
</div> </div>
); );
} }

View File

@@ -13,7 +13,7 @@ import { useCompletionSchema } from "@/hooks/use-completion-schema";
import { useConnections } from "@/hooks/use-connections"; import { useConnections } from "@/hooks/use-connections";
import { useAppStore } from "@/stores/app-store"; import { useAppStore } from "@/stores/app-store";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { Play, Loader2, Lock, BarChart3, Download, AlignLeft, Bookmark, Table2, Braces, Sparkles } from "lucide-react"; import { Play, Loader2, Lock, BarChart3, Download, AlignLeft, Bookmark, Table2, Braces, Sparkles, BrainCircuit } from "lucide-react";
import { format as formatSql } from "sql-formatter"; import { format as formatSql } from "sql-formatter";
import { SaveQueryDialog } from "@/components/saved-queries/SaveQueryDialog"; import { SaveQueryDialog } from "@/components/saved-queries/SaveQueryDialog";
import { import {
@@ -26,6 +26,7 @@ import { exportCsv, exportJson } from "@/lib/tauri";
import { save } from "@tauri-apps/plugin-dialog"; import { save } from "@tauri-apps/plugin-dialog";
import { toast } from "sonner"; import { toast } from "sonner";
import { AiBar } from "@/components/ai/AiBar"; import { AiBar } from "@/components/ai/AiBar";
import { useExplainSql, useFixSqlError } from "@/hooks/use-ai";
import type { QueryResult, ExplainResult } from "@/types"; import type { QueryResult, ExplainResult } from "@/types";
interface Props { interface Props {
@@ -53,8 +54,11 @@ export function WorkspacePanel({
const [resultViewMode, setResultViewMode] = useState<"table" | "json">("table"); const [resultViewMode, setResultViewMode] = useState<"table" | "json">("table");
const [saveDialogOpen, setSaveDialogOpen] = useState(false); const [saveDialogOpen, setSaveDialogOpen] = useState(false);
const [aiBarOpen, setAiBarOpen] = useState(false); const [aiBarOpen, setAiBarOpen] = useState(false);
const [aiExplanation, setAiExplanation] = useState<string | null>(null);
const queryMutation = useQueryExecution(); const queryMutation = useQueryExecution();
const explainMutation = useExplainSql();
const fixMutation = useFixSqlError();
const addHistoryMutation = useAddHistory(); const addHistoryMutation = useAddHistory();
const { data: connections } = useConnections(); const { data: connections } = useConnections();
const { data: completionSchema } = useCompletionSchema(connectionId); const { data: completionSchema } = useCompletionSchema(connectionId);
@@ -98,6 +102,7 @@ export function WorkspacePanel({
if (!sqlValue.trim() || !connectionId) return; if (!sqlValue.trim() || !connectionId) return;
setError(null); setError(null);
setExplainData(null); setExplainData(null);
setAiExplanation(null);
setResultView("results"); setResultView("results");
queryMutation.mutate( queryMutation.mutate(
{ connectionId, sql: sqlValue }, { connectionId, sql: sqlValue },
@@ -191,6 +196,60 @@ export function WorkspacePanel({
[result] [result]
); );
const isAiLoading = explainMutation.isPending || fixMutation.isPending;
const handleAiExplain = useCallback(() => {
if (!sqlValue.trim() || !connectionId) return;
setAiExplanation(null);
setResultView("results");
explainMutation.mutate(
{ connectionId, sql: sqlValue },
{
onSuccess: (explanation) => {
setAiExplanation(explanation);
},
onError: (err) => {
toast.error("AI Explain failed", { description: String(err) });
},
}
);
}, [connectionId, sqlValue, explainMutation]);
const handleExplainError = useCallback(() => {
if (!sqlValue.trim() || !connectionId || !error) return;
setAiExplanation(null);
explainMutation.mutate(
{ connectionId, sql: `${sqlValue}\n\n-- Error: ${error}` },
{
onSuccess: (explanation) => {
setAiExplanation(explanation);
},
onError: (err) => {
toast.error("AI Explain failed", { description: String(err) });
},
}
);
}, [connectionId, sqlValue, error, explainMutation]);
const handleFixError = useCallback(() => {
if (!sqlValue.trim() || !connectionId || !error) return;
fixMutation.mutate(
{ connectionId, sql: sqlValue, errorMessage: error },
{
onSuccess: (fixedSql) => {
setSqlValue(fixedSql);
onSqlChange?.(fixedSql);
setError(null);
setAiExplanation(null);
toast.success("SQL replaced by AI suggestion");
},
onError: (err) => {
toast.error("AI Fix failed", { description: String(err) });
},
}
);
}, [connectionId, sqlValue, error, fixMutation, onSqlChange]);
return ( return (
<> <>
<ResizablePanelGroup orientation="vertical"> <ResizablePanelGroup orientation="vertical">
@@ -257,6 +316,21 @@ export function WorkspacePanel({
<Sparkles className="h-3 w-3" /> <Sparkles className="h-3 w-3" />
AI AI
</Button> </Button>
<Button
size="sm"
variant="ghost"
className="h-6 gap-1 text-xs"
onClick={handleAiExplain}
disabled={isAiLoading || !sqlValue.trim()}
title="Explain query with AI"
>
{isAiLoading ? (
<Loader2 className="h-3 w-3 animate-spin" />
) : (
<BrainCircuit className="h-3 w-3" />
)}
AI Explain
</Button>
{result && result.columns.length > 0 && ( {result && result.columns.length > 0 && (
<DropdownMenu> <DropdownMenu>
<DropdownMenuTrigger asChild> <DropdownMenuTrigger asChild>
@@ -314,7 +388,7 @@ export function WorkspacePanel({
<ResizableHandle withHandle /> <ResizableHandle withHandle />
<ResizablePanel id="results" defaultSize="60%" minSize="15%"> <ResizablePanel id="results" defaultSize="60%" minSize="15%">
<div className="flex h-full flex-col overflow-hidden"> <div className="flex h-full flex-col overflow-hidden">
{(explainData || result || error) && ( {(explainData || result || error || aiExplanation) && (
<div className="flex shrink-0 items-center border-b text-xs"> <div className="flex shrink-0 items-center border-b text-xs">
<button <button
className={`px-3 py-1 font-medium ${ className={`px-3 py-1 font-medium ${
@@ -377,6 +451,10 @@ export function WorkspacePanel({
error={error} error={error}
isLoading={queryMutation.isPending && resultView === "results"} isLoading={queryMutation.isPending && resultView === "results"}
viewMode={resultViewMode} viewMode={resultViewMode}
aiExplanation={aiExplanation}
isAiLoading={isAiLoading}
onExplainError={error ? handleExplainError : undefined}
onFixError={error ? handleFixError : undefined}
/> />
)} )}
</div> </div>

View File

@@ -4,6 +4,8 @@ import {
saveAiSettings, saveAiSettings,
listOllamaModels, listOllamaModels,
generateSql, generateSql,
explainSql,
fixSqlError,
} from "@/lib/tauri"; } from "@/lib/tauri";
import type { AiSettings } from "@/types"; import type { AiSettings } from "@/types";
@@ -45,3 +47,29 @@ export function useGenerateSql() {
}) => generateSql(connectionId, prompt), }) => generateSql(connectionId, prompt),
}); });
} }
export function useExplainSql() {
return useMutation({
mutationFn: ({
connectionId,
sql,
}: {
connectionId: string;
sql: string;
}) => explainSql(connectionId, sql),
});
}
export function useFixSqlError() {
return useMutation({
mutationFn: ({
connectionId,
sql,
errorMessage,
}: {
connectionId: string;
sql: string;
errorMessage: string;
}) => fixSqlError(connectionId, sql, errorMessage),
});
}

View File

@@ -254,6 +254,12 @@ export const listOllamaModels = (ollamaUrl: string) =>
export const generateSql = (connectionId: string, prompt: string) => export const generateSql = (connectionId: string, prompt: string) =>
invoke<string>("generate_sql", { connectionId, prompt }); invoke<string>("generate_sql", { connectionId, prompt });
export const explainSql = (connectionId: string, sql: string) =>
invoke<string>("explain_sql", { connectionId, sql });
export const fixSqlError = (connectionId: string, sql: string, errorMessage: string) =>
invoke<string>("fix_sql_error", { connectionId, sql, errorMessage });
// Entity Lookup // Entity Lookup
export const entityLookup = ( export const entityLookup = (
config: ConnectionConfig, config: ConnectionConfig,