Merge pull request #119 from agnosticeng/feat/ask-fix-ai

feat: implement fix with AI button

src/lib/components/Ai/Chat.svelte

@@ -8,6 +8,7 @@
 	import Stop from '$lib/icons/Stop.svelte';
 	import { getTextFromElement, transform } from '$lib/markdown';
 	import type { Table } from '$lib/olap-engine';
+	import { onMount } from 'svelte';
 	import ChangeModelBox from './ChangeModelBox.svelte';
 	import DatasetsBox from './DatasetsBox.svelte';
 	import Loader from './Loader.svelte';
@@ -42,8 +43,8 @@
 	let datasetSelectbox = $state<ReturnType<typeof Select>>();
 	let textarea = $state<HTMLTextAreaElement>();
 	let abortController: AbortController | undefined;
-	let chatMessages = $derived(messages.filter((m) => m.role === 'user' || m.role === 'assistant'));
 	let modelSelectbox = $state<ReturnType<typeof Select>>();
+	let form = $state<HTMLFormElement>();
 	const uid = $props.id();
 
 	function getContextFromTable(table: Table): string {
@@ -55,6 +56,8 @@
 		dataset ??= datasets?.at(0);
 	});
 
+	onMount(() => form?.dispatchEvent(new SubmitEvent('submit')));
+
 	const client = $derived(new OpenAIClient(selectedModel.baseURL));
 
 	async function handleSubmit(
@@ -62,14 +65,18 @@
 	) {
 		event.preventDefault();
 
-		const data = new FormData(event.currentTarget);
-		let content = data.get('message');
-		if (!content || typeof content !== 'string') return;
-		content = content.trim();
-		if (!content.length) return;
+		const lastMessage = messages.at(-1);
+		if (lastMessage?.role !== 'user') {
+			const data = new FormData(event.currentTarget);
+			let content = data.get('message');
+			if (!content || typeof content !== 'string') return;
+			content = content.trim();
+			if (!content.length) return;
+
+			message = '';
+			messages = messages.concat({ role: 'user', content });
+		}
 
-		message = '';
-		messages = messages.concat({ content, role: 'user' });
 		loading = true;
 
 		try {
@@ -139,7 +146,7 @@
 		role="presentation"
 		onclick={(e) => e.target === e.currentTarget && textarea?.focus()}
 	>
-		{#each chatMessages as { role, content }, index}
+		{#each messages as { role, content }, index}
 			<article data-role={role}>
 				<h2>
 					{#if role === 'user'}
@@ -162,8 +169,8 @@
 		{:else}
 			<article>
 				<h2>You</h2>
-				{#if chatMessages.length === 0 && dataset}{@render context(dataset)}{/if}
-				<form id="{uid}-user-message" method="POST" onsubmit={handleSubmit}>
+				{#if messages.length === 0 && dataset}{@render context(dataset)}{/if}
+				<form id="{uid}-user-message" method="POST" onsubmit={handleSubmit} bind:this={form}>
 					<textarea
 						name="message"
 						tabindex="0"
@@ -239,7 +246,13 @@
 				<Stop size="11" />
 			</button>
 		{:else}
-			<button form="{uid}-user-message" type="submit" bind:this={submitter} title="Send ⌘⏎">
+			<button
+				form="{uid}-user-message"
+				type="submit"
+				bind:this={submitter}
+				title="Send ⌘⏎"
+				disabled={!message}
+			>
 				Send ⌘⏎
 			</button>
 		{/if}

src/lib/components/Ai/fix_query.md

@@ -0,0 +1,11 @@
+Please fix the SQL query so that it works correctly in ClickHouse. Keep the logic as close as possible to the original intent. Only return the corrected SQL.
+
+```sql
+{{sql}}
+```
+
+Error:
+
+```
+{{error}}
+```

src/lib/components/Ai/index.ts

@@ -1,6 +1,6 @@
 import type { Table } from '$lib/olap-engine';
+import fix_query_template from './fix_query.md?raw';
 import type { ChatInput, Model } from './types';
-
 export { OpenAIClient } from './OpenAI';
 export { default as AiPanel } from './Panel.svelte';
 
@@ -30,3 +30,7 @@ export function deserializeModel(model: string): Model | null {
 		return null;
 	}
 }
+
+export function generateFixMessage(sql: string, error: string) {
+	return fix_query_template.replace('{{sql}}', sql).replace('{{error}}', error);
+}

src/lib/components/Console.svelte

@@ -3,15 +3,19 @@
 		level: 'error';
 		timestamp: Date;
 		data: string;
+		query: string;
 	}
 </script>
 
 <script lang="ts">
+	import Sparkles from '$lib/icons/Sparkles.svelte';
+
 	interface Props {
 		logs: Log[];
+		onClickFixWithAi?: (log: Log) => void;
 	}
 
-	let { logs }: Props = $props();
+	let { logs, onClickFixWithAi }: Props = $props();
 </script>
 
 <div class="container">
@@ -19,6 +23,11 @@
 		<div class="line">
 			<span class="timestamp">{log.timestamp.toUTCString()}</span>
 			<span class={log.level}>{log.data}</span>
+			{#if onClickFixWithAi}
+				<button title="Fix with the AI Assistant" onclick={() => onClickFixWithAi(log)}>
+					<Sparkles size="12" />
+				</button>
+			{/if}
 		</div>
 	{/each}
 </div>
@@ -52,6 +61,29 @@
 			&:not(:last-of-type) {
 				margin-bottom: 12px;
 			}
+
+			& > button {
+				height: 18px;
+				aspect-ratio: 1;
+				border-radius: 4px;
+				padding-top: 2px;
+
+				&:hover {
+					background-color: hsl(0deg 0% 10%);
+				}
+			}
+		}
+	}
+
+	button {
+		appearance: none;
+		outline: none;
+		border: none;
+		background: none;
+		padding: 0;
+
+		&:not(:disabled):hover {
+			cursor: pointer;
 		}
 	}
 </style>

src/lib/components/Result.svelte

@@ -14,6 +14,7 @@
 		tab?: 'data' | 'chart' | 'logs';
 		onClearLogs?: () => void;
 		warnings?: string[];
+		onClickFixWithAi?: (log: Log) => void;
 	}
 
 	let {
@@ -21,7 +22,8 @@
 		logs = [],
 		tab = $bindable('data'),
 		onClearLogs,
-		warnings = []
+		warnings = [],
+		onClickFixWithAi
 	}: Props = $props();
 </script>
 
@@ -68,7 +70,7 @@
 		</div>
 
 		<div class="tab-content" style:visibility={tab === 'logs' ? 'visible' : 'hidden'}>
-			<Console {logs} />
+			<Console {logs} {onClickFixWithAi} />
 		</div>
 	</div>
 </section>

src/lib/olap-engine/engine-local.ts

@@ -1,4 +1,4 @@
-import { invoke, Channel } from '@tauri-apps/api/core';
+import { Channel, invoke } from '@tauri-apps/api/core';
 import { InternalEventEmitter } from './EventListener';
 import type { Events, ExecOptions, OLAPEngine, OLAPResponse, Table } from './index';
 
@@ -68,7 +68,7 @@ export class LocalEngine extends InternalEventEmitter<Events> implements OLAPEng
 			return data;
 		} catch (e) {
 			console.error(body);
-			if (_emit) this.emit('error', new Error(body));
+			if (_emit) this.emit('error', query, new Error(body));
 		}
 	}
 

src/lib/olap-engine/engine-remote.ts

@@ -26,7 +26,7 @@ export class RemoteEngine extends InternalEventEmitter<Events> implements OLAPEn
 			return data;
 		} catch (e) {
 			console.error(e);
-			if (_emit) this.emit('error', e);
+			if (_emit) this.emit('error', query, e);
 		}
 	}
 

src/routes/+page.svelte

@@ -1,5 +1,11 @@
 <script lang="ts">
-	import { AiPanel, deserializeModel, serializeModel, type Chat } from '$lib/components/Ai';
+	import {
+		AiPanel,
+		deserializeModel,
+		generateFixMessage,
+		serializeModel,
+		type Chat
+	} from '$lib/components/Ai';
 	import type { Log } from '$lib/components/Console.svelte';
 	import { ContextMenuState } from '$lib/components/ContextMenu';
 	import ContextMenu from '$lib/components/ContextMenu/ContextMenu.svelte';
@@ -303,10 +309,10 @@
 
 	let bottomPanelTab = $state<'data' | 'chart' | 'logs'>('data');
 	let errors = $state.raw<Log[]>([]);
-	engine.on('error', (e) => {
+	engine.on('error', (query: string, e: unknown) => {
 		if (e instanceof Error) {
 			if (e.message === 'Canceled') return;
-			errors = errors.concat({ level: 'error', timestamp: new Date(), data: e.message });
+			errors = errors.concat({ level: 'error', timestamp: new Date(), data: e.message, query });
 			bottomPanel.open = true;
 			bottomPanelTab = 'logs';
 		}
@@ -389,6 +395,18 @@ LIMIT 100;`;
 			) ?? fallback
 		);
 	});
+
+	function handleFixQuery({ data, query }: Log) {
+		focusedChat =
+			chats.push({
+				id: crypto.randomUUID(),
+				name: 'New Chat',
+				messages: [{ role: 'user', content: generateFixMessage(query, data) }]
+			}) - 1;
+
+		if (isMobile) rightDrawerOpened = true;
+		else rightPanel.open = true;
+	}
 </script>
 
 <svelte:window onkeydown={handleKeyDown} bind:innerWidth={screenWidth} />
@@ -566,6 +584,7 @@ LIMIT 100;`;
 									logs={errors}
 									bind:tab={bottomPanelTab}
 									onClearLogs={() => (errors = [])}
+									onClickFixWithAi={handleFixQuery}
 									{warnings}
 								/>
 							{/snippet}