feat: add toolType as useAIChat feature, and add sql chat which can help user generate sql which easy to query

Author: moonrailgun

src/client/components/CodeEditor/sql.tsx

@@ -32,6 +32,7 @@ interface SQLEditorProps {
   onExecuteLine?: (lineNumber: number, sql: string) => void | Promise<void>;
   enableRunButton?: boolean;
   executingLine?: number | null;
+  onSelectSqlChange?: (lineNumber: number, sql: string) => void;
 }
 
 export const SQLEditor: React.FC<SQLEditorProps> = React.memo((props) => {
@@ -41,11 +42,13 @@ export const SQLEditor: React.FC<SQLEditorProps> = React.memo((props) => {
     enableRunButton = false,
     onExecuteLine,
     executingLine = null,
+    onSelectSqlChange,
   } = props;
   const colorScheme = useTheme();
   const theme = colorScheme === 'dark' ? 'vs-dark' : 'light';
   const editorRef = useRef<editor.IStandaloneCodeEditor | null>(null);
   const monacoRef = useRef<Monaco | null>(null);
+  const lastSelectedLineRef = useRef<number | null>(null);
 
   const registerSQLCompletions = useEvent((monaco: Monaco) => {
     return createSQLCompletionProvider(monaco, tables);
@@ -78,6 +81,36 @@ export const SQLEditor: React.FC<SQLEditorProps> = React.memo((props) => {
         executingLine
       );
     }
+
+    // Listen to cursor position changes
+    if (onSelectSqlChange) {
+      editor.onDidChangeCursorPosition((e) => {
+        const lineNumber = e.position.lineNumber;
+
+        // Only trigger if line changed
+        if (lastSelectedLineRef.current === lineNumber) {
+          return;
+        }
+        lastSelectedLineRef.current = lineNumber;
+
+        const model = editor.getModel();
+        if (!model) {
+          return;
+        }
+
+        const content = model.getValue();
+        const statements = parseSQLStatements(content);
+
+        // Find the statement that contains the current line
+        const statement = statements.find(
+          (stmt) => stmt.startLine <= lineNumber && stmt.endLine >= lineNumber
+        );
+
+        if (statement) {
+          onSelectSqlChange(lineNumber, statement.sql);
+        }
+      });
+    }
   });
 
   const handleEditorWillMount = useEvent((monaco: Monaco) => {

src/client/components/ai/AIChatbot.tsx

@@ -1,5 +1,5 @@
 import React from 'react';
-import { LuCircleAlert } from 'react-icons/lu';
+import { LuCircleAlert, LuSparkles } from 'react-icons/lu';
 import { SimpleContextUsage } from '../ai-elements/context';
 import {
   Conversation,
@@ -30,7 +30,7 @@ interface AIChatbotProps {
   input: string;
   setInput: (input: string) => void;
   placeholder?: string;
-  usage?: LanguageModelUsage;
+  usage?: LanguageModelUsage | null;
   isDisabled?: boolean;
   suggestions?: string[];
   alert?: React.ReactNode;
@@ -62,6 +62,24 @@ export const AIChatbot: React.FC<AIChatbotProps> = React.memo((props) => {
         <ConversationScrollButton />
       </Conversation>
 
+      {props.messages.length === 0 && (
+        <div className="flex flex-1 flex-col items-center justify-center gap-4 p-8 text-center">
+          <div className="bg-primary/10 text-primary flex h-16 w-16 items-center justify-center rounded-full">
+            <LuSparkles className="h-8 w-8" />
+          </div>
+          <div className="space-y-2">
+            <h4 className="text-lg font-semibold">
+              {t('Use AI to improve your work')}
+            </h4>
+            <p className="text-muted-foreground max-w-sm text-sm">
+              {t(
+                'Leverage AI to generate SQL queries, gain insights, and enhance your productivity. Describe your data analysis needs and let AI assist you effortlessly.'
+              )}
+            </p>
+          </div>
+        </div>
+      )}
+
       {props.error && (
         <div className="p-3">
           <Alert variant="destructive">

src/client/components/ai/AIResponseMessages.tsx

@@ -1,5 +1,5 @@
 import { AbstractChat, ChatStatus, UIDataTypes, UIMessage, UITools } from 'ai';
-import React from 'react';
+import React, { Fragment } from 'react';
 import {
   Message,
   MessageAction,
@@ -31,8 +31,8 @@ export const AIResponseMessages: React.FC<AIResponseMessagesProps> = React.memo(
           const isLastMessage = messageIndex === messages.length - 1;
 
           return (
-            <>
-              <Message from={message.role} key={message.id}>
+            <Fragment key={message.id}>
+              <Message from={message.role}>
                 <MessageContent
                   className={message.role === 'assistant' ? 'w-full' : 'w-auto'}
                 >
@@ -71,7 +71,7 @@ export const AIResponseMessages: React.FC<AIResponseMessagesProps> = React.memo(
                     </MessageAction>
                   </MessageActions>
                 )}
-            </>
+            </Fragment>
           );
         })}
 

src/client/components/insights/warehouse/WarehouseAIChatPanel.tsx

@@ -1,41 +1,67 @@
 import React from 'react';
 import { useTranslation } from '@i18next-toolkit/react';
-import { LuSparkles } from 'react-icons/lu';
+import { AIChatbot } from '@/components/ai/AIChatbot';
+import { useWarehouseSqlChat } from '@/hooks/useWarehouseSqlChat';
+import { AppRouterOutput } from '@/api/trpc';
+
+type DatabaseType =
+  AppRouterOutput['insights']['warehouse']['database']['list'][number];
 
 interface WarehouseAIChatPanelProps {
   workspaceId: string;
-  databaseId: string;
+  database: DatabaseType;
+  currentSelectedSql: string | null;
   onSQLGenerated?: (sql: string) => void;
 }
 
 export const WarehouseAIChatPanel: React.FC<WarehouseAIChatPanelProps> =
-  React.memo(() => {
+  React.memo((props) => {
     const { t } = useTranslation();
 
+    const {
+      messages,
+      status,
+      input,
+      setInput,
+      usage,
+      addToolResult,
+      handleReset,
+      handleSend,
+      handleRegenerate,
+    } = useWarehouseSqlChat({
+      workspaceId: props.workspaceId,
+      selectedScopes: [
+        {
+          type: 'database',
+          id: props.database.id,
+          name: props.database.name,
+        },
+      ],
+      currentSelectedSql: props.currentSelectedSql,
+      onSQLGenerated: props.onSQLGenerated,
+    });
+
     return (
       <div className="flex h-full flex-col">
-        {/* Header */}
         <div className="border-b px-4 py-3">
           <h3 className="font-semibold">{t('AI Assistant')}</h3>
           <p className="text-muted-foreground text-xs">
             {t('AI-powered SQL query generation')}
           </p>
         </div>
 
-        {/* Coming Soon Content */}
-        <div className="flex flex-1 flex-col items-center justify-center gap-4 p-8 text-center">
-          <div className="bg-primary/10 text-primary flex h-16 w-16 items-center justify-center rounded-full">
-            <LuSparkles className="h-8 w-8" />
-          </div>
-          <div className="space-y-2">
-            <h4 className="text-lg font-semibold">{t('Coming Soon')}</h4>
-            <p className="text-muted-foreground max-w-sm text-sm">
-              {t(
-                'AI-powered SQL query generation and assistance will be available soon. Stay tuned!'
-              )}
-            </p>
-          </div>
-        </div>
+        <AIChatbot
+          className="flex-1 overflow-hidden"
+          messages={messages}
+          status={status}
+          input={input}
+          setInput={setInput}
+          usage={usage}
+          onAddToolResult={addToolResult}
+          onSubmit={handleSend}
+          onReset={handleReset}
+          onRegenerate={handleRegenerate}
+        />
       </div>
     );
   });

src/client/hooks/useAIChat.ts

@@ -2,6 +2,7 @@ import { useChat } from '@ai-sdk/react';
 import {
   ChatInit,
   DefaultChatTransport,
+  HttpChatTransportInitOptions,
   LanguageModelUsage,
   lastAssistantMessageIsCompleteWithToolCalls,
   UIMessage,
@@ -20,6 +21,8 @@ interface UseAIChatOptions {
    */
   apiEndpoint: string;
 
+  transportOptions?: Omit<HttpChatTransportInitOptions<UIMessage>, 'api'>;
+
   /**
    * Custom tool call handler
    */
@@ -53,6 +56,7 @@ interface UseAIChatOptions {
  */
 export function useAIChat({
   apiEndpoint,
+  transportOptions,
   onToolCall,
   initialMessages,
   onMessagesUpdate,
@@ -73,9 +77,10 @@ export function useAIChat({
   const transport = useMemo(
     () =>
       new DefaultChatTransport({
+        ...transportOptions,
         api: apiEndpoint,
       }),
-    [apiEndpoint]
+    [apiEndpoint, transportOptions]
   );
 
   // Initialize chat with AI SDK

src/client/hooks/useWarehouseInsightChat.ts

@@ -21,7 +21,7 @@ type ChartBlock = InsightsTimeEventChartProps & {
   title: string;
 };
 
-interface UseWarehouseAIChatOptions {
+interface UseWarehouseInsightChatOptions {
   workspaceId: string;
   selectedScopes: WarehouseScope[];
   isDisabled?: boolean;
@@ -36,7 +36,7 @@ export function useWarehouseInsightChat({
   workspaceId,
   selectedScopes,
   isDisabled = false,
-}: UseWarehouseAIChatOptions) {
+}: UseWarehouseInsightChatOptions) {
   const { t } = useTranslation();
   const [input, setInput] = useState('');
 
@@ -82,6 +82,11 @@ export function useWarehouseInsightChat({
     regenerate: regenerateChat,
   } = useAIChat({
     apiEndpoint: `/api/insights/${workspaceId}/chat`,
+    transportOptions: {
+      body: {
+        toolType: 'chart',
+      },
+    },
     initialMessages: session?.messages ?? [],
     onMessagesUpdate: (messages) => {
       // Save messages when AI response finishes