bringing ai with tweaks to db:studio

packages/core/src/components/chat/chat-sidebar.tsx

@@ -2,8 +2,9 @@
 
 import { fetchServerSentEvents, useChat } from "@tanstack/ai-react";
 import { Plus } from "lucide-react";
-import { useState } from "react";
-import { CHAT_SUGGESTIONS, DEFAULTS } from "shared/constants";
+import { useEffect, useMemo, useState } from "react";
+import { CHAT_SUGGESTIONS, DEFAULTS, MODEL_LIST } from "shared/constants";
+import type { ExecuteQueryResult } from "shared/types";
 import {
 	Conversation,
 	ConversationContent,
@@ -33,21 +34,120 @@ import {
 import { Suggestion, Suggestions } from "@/components/ai-elements/suggestion";
 import { SheetSidebar } from "@/components/sheet-sidebar";
 import { Button } from "@/components/ui/button";
+import {
+	Select,
+	SelectContent,
+	SelectItem,
+	SelectTrigger,
+	SelectValue,
+} from "@/components/ui/select";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+import { useExecuteQuery } from "@/hooks/use-execute-query";
 import { useRateLimit } from "@/hooks/use-rate-limit";
+import { getDbType } from "@/lib/api";
+import { useAiPrefillStore } from "@/stores/ai-prefill.store";
+import { useAiSettingsStore } from "@/stores/ai-settings.store";
+import { useDatabaseStore } from "@/stores/database.store";
+import { useInsertSqlStore } from "@/stores/insert-sql.store";
 import { useSheetStore } from "@/stores/sheet.store";
 
 export const ChatSidebar = () => {
 	const { rateLimit, refetchRateLimit } = useRateLimit();
 	const [text, setText] = useState("");
+	const [chatResults, setChatResults] = useState<
+		Record<string, { isLoading?: boolean; error?: string; result?: ExecuteQueryResult }>
+	>({});
 	const { isSheetOpen, closeSheet } = useSheetStore();
+	const { selectedDatabase } = useDatabaseStore();
+	const {
+		includeSchemaInAiContext,
+		useByocProxy,
+		byocProxyUrl,
+		provider,
+		model,
+		apiKeys,
+		setModel,
+	} = useAiSettingsStore();
+	const { executeQuery: executeSandboxQuery } = useExecuteQuery("sandbox");
+
+	const chatUrl = useMemo(() => {
+		const dbType = getDbType() ?? "pg";
+		return `${DEFAULTS.BASE_URL}/${dbType}/chat`;
+	}, []);
+
+	const chatBody = useMemo(() => {
+		const body: Record<string, string | boolean | undefined> = {
+			db: selectedDatabase ?? undefined,
+			includeSchemaInAiContext,
+			provider,
+			model,
+			apiKey: apiKeys[provider] || undefined,
+		};
+		if (useByocProxy && byocProxyUrl?.trim()) {
+			body.proxyUrl = byocProxyUrl.trim();
+		}
+		return body;
+	}, [
+		selectedDatabase,
+		includeSchemaInAiContext,
+		useByocProxy,
+		byocProxyUrl,
+		provider,
+		model,
+		apiKeys,
+	]);
+
+	const { prefillMessage, setPrefillMessage } = useAiPrefillStore();
+	const { setPendingSql } = useInsertSqlStore();
+
+	function extractSqlBlock(text: string): string | null {
+		const match = text.match(/```sql\s*([\s\S]*?)```/);
+		return match ? match[1].trim() : null;
+	}
+
+	useEffect(() => {
+		if (isSheetOpen("ai-assistant") && prefillMessage) {
+			setText(prefillMessage);
+			setPrefillMessage(null);
+		}
+	}, [isSheetOpen("ai-assistant"), prefillMessage, setPrefillMessage]);
 
 	const { messages, sendMessage, isLoading, clear, stop } = useChat({
-		connection: fetchServerSentEvents(`${DEFAULTS.BASE_URL}/chat`),
+		connection: fetchServerSentEvents(chatUrl, { body: chatBody }),
 		onError: (error) => console.error("Error:", error.message),
 		onResponse: (response) => console.log("Response:", response),
 		onFinish: (message) => {
 			console.log("Finish:", message);
+			const messageAny = message as {
+				parts?: { type: string; content?: string }[];
+				content?: string;
+			};
+			const textParts = Array.isArray(messageAny.parts)
+				? messageAny.parts
+						.filter((part) => part.type === "text")
+						.map((part) => part.content ?? "")
+						.join("")
+				: (messageAny.content ?? "");
+			const sqlBlock = extractSqlBlock(textParts);
+			if (sqlBlock && message.id && selectedDatabase) {
+				setChatResults((prev) => ({
+					...prev,
+					[message.id]: { isLoading: true },
+				}));
+				executeSandboxQuery({ query: sqlBlock })
+					.then((result) => {
+						setChatResults((prev) => ({
+							...prev,
+							[message.id]: { isLoading: false, result },
+						}));
+					})
+					.catch((error: Error) => {
+						setChatResults((prev) => ({
+							...prev,
+							[message.id]: { isLoading: false, error: error.message },
+						}));
+					});
+			}
 			refetchRateLimit();
 		},
 	});
@@ -67,6 +167,7 @@ export const ChatSidebar = () => {
 	const handleNewChat = () => {
 		clear();
 		setText("");
+		setChatResults({});
 	};
 
 	const handleSuggestionClick = (suggestion: string) => {
@@ -86,6 +187,24 @@ export const ChatSidebar = () => {
 			title="AI Assistant"
 			cta={
 				<div className="flex items-center gap-2">
+					<Select
+						value={model}
+						onValueChange={setModel}
+					>
+						<SelectTrigger className="h-8 w-[210px] text-xs">
+							<SelectValue placeholder="Model" />
+						</SelectTrigger>
+						<SelectContent>
+							{MODEL_LIST.filter((item) => item.provider === provider).map((item) => (
+								<SelectItem
+									key={item.id}
+									value={item.id}
+								>
+									{item.name}
+								</SelectItem>
+							))}
+						</SelectContent>
+					</Select>
 					{rateLimit && (
 						<Tooltip>
 							<TooltipTrigger asChild>
@@ -139,10 +258,13 @@ export const ChatSidebar = () => {
 									);
 									const textContent = message.parts
 										.filter((part) => part.type === "text")
-										.map((part) => part.content)
+										.map((part) => ("content" in part ? part.content : ""))
 										.join("");
 
 									const hasThinking = thinkingParts.length > 0;
+									const sqlBlock =
+										message.role === "assistant" ? extractSqlBlock(textContent) : null;
+									const chatResult = message.id ? chatResults[message.id] : undefined;
 
 									return (
 										<MessageBranch
@@ -151,18 +273,63 @@ export const ChatSidebar = () => {
 										>
 											<MessageBranchContent>
 												<Message from={message.role === "user" ? "user" : "assistant"}>
-													<div>
+													<div className="space-y-2">
 														{hasThinking && message.role === "assistant" && (
 															<Reasoning duration={0}>
 																<ReasoningTrigger />
 																<ReasoningContent>
-																	{thinkingParts.map((part) => part.content).join("\n")}
+																	{thinkingParts
+																		.map((part) => ("content" in part ? part.content : ""))
+																		.join("\n")}
 																</ReasoningContent>
 															</Reasoning>
 														)}
 														<MessageContent>
 															<MessageResponse>{textContent}</MessageResponse>
 														</MessageContent>
+														{sqlBlock && (
+															<div className="space-y-2">
+																<Button
+																	type="button"
+																	variant="outline"
+																	size="sm"
+																	className="mt-2"
+																	onClick={() => setPendingSql(sqlBlock)}
+																>
+																	Insert into editor
+																</Button>
+																{chatResult?.isLoading && (
+																	<div className="text-xs text-muted-foreground">
+																		Running in sandbox...
+																	</div>
+																)}
+																{chatResult?.error && (
+																	<div className="text-xs text-destructive">
+																		Error: {chatResult.error}
+																	</div>
+																)}
+																{chatResult?.result && (
+																	<div className="rounded-md border border-zinc-800 p-2 text-xs space-y-1">
+																		<div className="text-amber-400">
+																			Sandbox result — no changes saved
+																		</div>
+																		<div className="text-muted-foreground">
+																			{chatResult.result.rowCount} rows •{" "}
+																			{chatResult.result.duration.toFixed(2)}ms
+																		</div>
+																		{chatResult.result.rows.length > 0 && (
+																			<pre className="whitespace-pre-wrap">
+																				{JSON.stringify(
+																					chatResult.result.rows.slice(0, 5),
+																					null,
+																					2,
+																				)}
+																			</pre>
+																		)}
+																	</div>
+																)}
+															</div>
+														)}
 													</div>
 												</Message>
 											</MessageBranchContent>
@@ -222,7 +389,7 @@ export const ChatSidebar = () => {
 									className="h-8!"
 									status={status}
 									onClick={isLoading ? handleStop : undefined}
-									disabled={rateLimit && rateLimit.remaining === 0}
+									disabled={rateLimit?.remaining === 0 || !selectedDatabase}
 								/>
 							</PromptInputFooter>
 						</PromptInput>

packages/core/src/components/components/command-palette.tsx

@@ -682,13 +682,7 @@ export function CommandPalette() {
 								</div>
 							</CommandItem>
 							<CommandItem
-								onSelect={() =>
-									handleAction(() => {
-										toast.info("Settings - Coming Soon!", {
-											description: "Customize your DB Studio experience",
-										});
-									})
-								}
+								onSelect={() => handleAction(() => openSheet("settings"), "Opening settings")}
 							>
 								<Settings className="mr-2 h-4 w-4" />
 								<div className="flex flex-col">